Update app.py

app.py

@@ -35,8 +35,11 @@ def clear_memory():
         torch.cuda.empty_cache()
     gc.collect()
 
-def add_metadata_and_save(image: Image.Image, filepath: str, prompt: str, negative_prompt: str, seed: int, steps: int, guidance: float, strength: Optional[float] = None):
+def add_metadata_and_save(image: Image.Image, prompt: str, negative_prompt: str, seed: int, steps: int, guidance: float, strength: Optional[float] = None):
     """Embed generation metadata into a PNG and save it."""
+    # Create temporary file with unique name
+    temp_path = tempfile.mktemp(suffix=".png")
+    
     meta = PngImagePlugin.PngInfo()
     meta.add_text("Prompt", prompt)
     meta.add_text("NegativePrompt", negative_prompt)
@@ -50,8 +53,8 @@ def add_metadata_and_save(image: Image.Image, filepath: str, prompt: str, negati
     meta.add_text("Seed", str(seed))
     meta.add_text("Date", datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
     
-    image.save(filepath, "PNG", pnginfo=meta)
-    return filepath
+    image.save(temp_path, "PNG", pnginfo=meta)
+    return temp_path
 
 def load_models():
     """Load both text2img and img2img pipelines optimized for Spaces"""
@@ -199,10 +202,9 @@ def generate_txt2img(prompt, negative_prompt, num_steps, guidance_scale, width,
         
         generation_time = time.time() - start_time
         
-        # Save with metadata
-        temp_path = tempfile.mktemp(suffix=".png")
-        add_metadata_and_save(
-            result.images[0], temp_path, enhanced_prompt, negative_prompt or "", 
+        # Save with metadata - returns file path
+        png_path = add_metadata_and_save(
+            result.images[0], enhanced_prompt, negative_prompt or "", 
             seed, num_steps, guidance_scale
         )
         
@@ -212,7 +214,7 @@ def generate_txt2img(prompt, negative_prompt, num_steps, guidance_scale, width,
             negative_prompt or "", seed, num_steps, guidance_scale
         )
         
-        return result.images[0], status
+        return png_path, status
         
     except Exception as e:
         return None, f"Generation failed: {str(e)}"
@@ -277,10 +279,9 @@ def generate_img2img(input_image, prompt, negative_prompt, num_steps, guidance_s
         
         generation_time = time.time() - start_time
         
-        # Save with metadata
-        temp_path = tempfile.mktemp(suffix=".png")
-        add_metadata_and_save(
-            result.images[0], temp_path, enhanced_prompt, negative_prompt or "", 
+        # Save with metadata - returns file path
+        png_path = add_metadata_and_save(
+            result.images[0], enhanced_prompt, negative_prompt or "", 
             seed, num_steps, guidance_scale, strength
         )
         
@@ -290,7 +291,7 @@ def generate_img2img(input_image, prompt, negative_prompt, num_steps, guidance_s
             negative_prompt or "", seed, num_steps, guidance_scale, strength
         )
         
-        return result.images[0], status
+        return png_path, status
         
     except Exception as e:
         return None, f"Transformation failed: {str(e)}"
@@ -315,8 +316,8 @@ def set_example_prompt():
 
 # Simplified negative prompt for better performance
 DEFAULT_NEGATIVE = """
-(low quality:1.3), (worst quality:1.3), (bad quality:1.2), (nsfw:1.5), (easynegative:1.3) (bad_prompt:1.3) badhandv4 bad-hands-5 (negative_hand-neg) (bad-picture-chill-75v), (worst quality:1.3), (low quality:1.3), (bad quality:1.3), (a shadow on skin:1.3), (a shaded skin:1.3), (a dark skin:1.3), (blush:1.3), (signature, watermark, username, letter, copyright name, 
-copyright, chinese text, artist name, name tag, company name, name tag, text, error:1.5), (bad anatomy:1.5), (low quality hand:1.5), (worst quality hand:1.5)
+(low quality:1.3), (worst quality:1.3), (bad quality:1.2), blurry, noisy, ugly, deformed, 
+(text, watermark:1.4), (extra limbs:1.3), (bad hands:1.3), (bad anatomy:1.2)
 """
 
 # Gradio interface optimized for Spaces
@@ -374,7 +375,7 @@ with gr.Blocks(
                     txt2img_btn = gr.Button("🎨 Generate", variant="primary", size="lg")
                 
                 with gr.Column():
-                    txt2img_output = gr.Image(label="Generated Image", height=400)
+                    txt2img_output = gr.File(label="Generated PNG with Metadata", file_types=[".png"])
                     txt2img_status = gr.Textbox(label="Generation Info", interactive=False, lines=6)
         
         with gr.TabItem("🖼️ Image to Image"):
@@ -420,7 +421,7 @@ with gr.Blocks(
                     img2img_btn = gr.Button("🖼️ Transform", variant="primary", size="lg")
                 
                 with gr.Column():
-                    img2img_output = gr.Image(label="Generated Image", height=400)
+                    img2img_output = gr.File(label="Generated PNG with Metadata", file_types=[".png"])
                     img2img_status = gr.Textbox(label="Generation Info", interactive=False, lines=6)
     
     # Event handlers