Update app.py

app.py

@@ -5,41 +5,62 @@ from random import randint
 from all_models import models
 from io import BytesIO
 from PIL import Image
+from fastapi import FastAPI, Request
+import json
 
 css_code = os.getenv("DazDinGo_CSS")
 
-def load_fn(models):
-    global models_load
-    models_load = {}
-    for model in models:
-        if model not in models_load.keys():
-            try:
-                m = gr.load(f'models/{model}')
-            except Exception as error:
-                m = gr.Interface(lambda txt: None, ['text'], ['image'])
-            models_load.update({model: m})
+# Load models
+models_load = {}
+for model in models:
+    try:
+        models_load[model] = gr.load(f'models/{model}')
+    except Exception as error:
+        models_load[model] = gr.Interface(lambda txt: None, ['text'], ['image'])
 
-load_fn(models)
+app = FastAPI()
 
-def gen_fn(model_str, prompt):
+def gen_image(model_str, prompt):
     if model_str == 'NA':
-        return None, ""
+        return None
     noise = str(randint(0, 4294967296))
     klir = '| ultra detail, ultra elaboration, ultra quality, perfect'
-    image = models_load[model_str](f'{prompt} {klir} {noise}')
-    
-    # Convert to base64
+    return models_load[model_str](f'{prompt} {klir} {noise}')
+
+def image_to_base64(image):
     buffered = BytesIO()
     if isinstance(image, str):  # if it's a file path
         img = Image.open(image)
         img.save(buffered, format="JPEG")
     else:  # if it's a PIL Image
         image.save(buffered, format="JPEG")
-    img_str = base64.b64encode(buffered.getvalue()).decode()
-    base64_code = f"data:image/jpeg;base64,{img_str}"
+    return base64.b64encode(buffered.getvalue()).decode()
+
+# API endpoint
+@app.post("/generate")
+async def api_generate(request: Request):
+    data = await request.json()
+    model = data.get('model', models[0])
+    prompt = data.get('prompt', '')
+    
+    if model not in models:
+        return {"error": "Model not found"}
+    
+    image = gen_image(model, prompt)
+    if image is None:
+        return {"error": "Image generation failed"}
     
-    return image, base64_code
+    base64_str = image_to_base64(image)
+    
+    return {
+        "status": "success",
+        "model": model,
+        "prompt": prompt,
+        "image_base64": base64_str,
+        "image_format": "jpeg"
+    }
 
+# Gradio Interface
 def make_me():
     with gr.Row():
         with gr.Column(scale=4):
@@ -66,24 +87,38 @@ def make_me():
                                         value=models[0] if models else None)
             output_image = gr.Image(label="Generated Image", width=512, 
                                   max_height=768, elem_id="custom_image", 
-                                  show_label=True, interactive=False, 
-                                  show_share_button=False)
+                                  show_label=True, interactive=False)
+            
+            # JSON output
+            json_output = gr.JSON(label="API Response")
             
-            # Base64 output
-            base64_output = gr.Textbox(label="Base64 Code", interactive=True, 
-                                     placeholder="Base64 code will appear here after generation",
-                                     lines=4)
+            def generate_wrapper(model_str, prompt):
+                image = gen_image(model_str, prompt)
+                if image is None:
+                    return None, {"error": "Generation failed"}
+                
+                base64_str = image_to_base64(image)
+                response = {
+                    "status": "success",
+                    "model": model_str,
+                    "prompt": prompt,
+                    "image_base64": base64_str,
+                    "image_format": "jpeg"
+                }
+                return image, response
             
-            gen_event = gen_button.click(gen_fn, [model_dropdown, txt_input], 
-                                       [output_image, base64_output])
+            gen_event = gen_button.click(generate_wrapper, [model_dropdown, txt_input], 
+                                       [output_image, json_output])
             stop_button.click(on_stop_click, inputs=None, 
                             outputs=[gen_button, stop_button], cancels=[gen_event])
-    
-    with gr.Row():
-        gr.HTML("")
 
+# Create Gradio app
 with gr.Blocks(css=css_code) as demo: 
     make_me()
 
-demo.queue(concurrency_count=50)
-demo.launch()
+# Mount Gradio app to FastAPI
+app = gr.mount_gradio_app(app, demo, path="/")
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)