Update app.py

app.py

@@ -5,7 +5,7 @@ from random import randint
 from all_models import models
 from io import BytesIO
 from PIL import Image
-from fastapi import FastAPI, Request, Response
+from fastapi import FastAPI, Request
 from deep_translator import GoogleTranslator
 import tempfile
 import uuid
@@ -37,17 +37,6 @@ def convert_to_png(image):
     png_buffer.seek(0)
     return Image.open(png_buffer)
 
-# Dictionary to store temporary files in memory
-temp_images = {}
-
-@app.get("/temp_images/{image_id}.png")
-async def get_temp_image(image_id: str):
-    """Endpoint to serve temporary images directly"""
-    if image_id in temp_images:
-        image_bytes = temp_images[image_id]
-        return Response(content=image_bytes, media_type="image/png")
-    return Response(status_code=404)
-
 def gen_image(model_str, prompt):
     if model_str == 'NA':
         return None, None
@@ -69,18 +58,15 @@ def gen_image(model_str, prompt):
     elif not isinstance(generated_image, Image.Image):
         generated_image = Image.fromarray(generated_image)
     
-    # Convert to PNG and store in memory
-    png_buffer = BytesIO()
-    generated_image.save(png_buffer, format="PNG")
-    png_bytes = png_buffer.getvalue()
+    # Create temp directory
+    temp_dir = os.path.join(tempfile.gettempdir(), "gradio_images")
+    os.makedirs(temp_dir, exist_ok=True)
     
-    # Generate unique ID for the image
-    image_id = str(uuid.uuid4())
-    temp_images[image_id] = png_bytes
+    # Save to temporary file
+    temp_path = os.path.join(temp_dir, f"{uuid.uuid4()}.png")
+    generated_image.save(temp_path, format="PNG")
     
-    # Return direct URL without proxy
-    image_url = f"/temp_images/{image_id}.png"
-    return image_url, seed
+    return temp_path, seed
 
 # Gradio Interface
 def make_me():
@@ -100,41 +86,37 @@ def make_me():
                                     elem_id="custom_stop_button")
                 
                 def on_generate_click():
-                    return gr.Button('Generating...', interactive=False, elem_id="custom_gen_button"), gr.Button('Stop', variant='secondary', interactive=True, elem_id="custom_stop_button")
+                    return gr.Button('Generating...', interactive=False), gr.Button('Stop', interactive=True)
                 
                 def on_stop_click():
-                    return gr.Button('Generate', variant='primary', elem_id="custom_gen_button"), gr.Button('Stop', variant='secondary', interactive=False, elem_id="custom_stop_button")
+                    return gr.Button('Generate', interactive=True), gr.Button('Stop', interactive=False)
                 
-                gen_button.click(on_generate_click, inputs=None, outputs=[gen_button, stop_button])
-                stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button])
+                gen_button.click(on_generate_click, None, [gen_button, stop_button])
+                stop_button.click(on_stop_click, None, [gen_button, stop_button])
             
             model_dropdown = gr.Dropdown(
                 models, 
                 label="Select Model", 
-                value=models[0] if models else None,
-                interactive=True
+                value=models[0] if models else None
             )
         
         with gr.Column(scale=2):
-            with gr.Row():
-                output_gallery = gr.Gallery(
-                    label="Generated PNG Images",
-                    columns=2,
-                    height="auto",
-                    elem_id="gallery"
-                )
+            output_gallery = gr.Gallery(
+                label="Generated PNG Images",
+                columns=2,
+                height="auto",
+                elem_id="gallery"
+            )
             seed_output = gr.Textbox(
                 label="Seed used",
                 interactive=False
             )
             
             def generate_wrapper(model_str, prompt):
-                image_url, seed = gen_image(model_str, prompt)
-                if image_url is None:
+                image_path, seed = gen_image(model_str, prompt)
+                if image_path is None:
                     return None, ""
-                
-                # Return the direct image URL
-                return [image_url], str(seed)
+                return [image_path], str(seed)
                 
             gen_event = gen_button.click(
                 generate_wrapper, 
@@ -142,10 +124,9 @@ def make_me():
                 [output_gallery, seed_output]
             )
             stop_button.click(
-                fn=lambda: [gr.Button('Generate', variant='primary', interactive=True), 
-                          gr.Button('Stop', variant='secondary', interactive=False)],
-                inputs=None,
-                outputs=[gen_button, stop_button],
+                on_stop_click,
+                None,
+                [gen_button, stop_button],
                 cancels=[gen_event]
             )