Update app.py

app.py

@@ -1,13 +1,16 @@
 import gradio as gr
 import torch
 from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
-from PIL import Image
+from PIL import Image, PngImagePlugin
+from datetime import datetime
 import os
 import gc
 import time
 import spaces
 from typing import Optional, Tuple
 from huggingface_hub import hf_hub_download
+import tempfile
+import random
 
 # Global pipeline variables
 txt2img_pipe = None
@@ -17,6 +20,14 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 # Hugging Face model configuration
 MODEL_REPO = "ajsbsd/CyberRealistic-Pony"
 MODEL_FILENAME = "cyberrealisticPony_v110.safetensors"
+model_id = f"{MODEL_REPO}/{MODEL_FILENAME}"
+
+# Generation configuration for metadata
+generation_config = {
+    "vae": "SDXL VAE",
+    "sampler": "DPM++ 2M Karras",
+    "steps": 20
+}
 
 def clear_memory():
     """Clear GPU memory"""
@@ -24,6 +35,24 @@ 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):
+    """Embed generation metadata into a PNG and save it."""
+    meta = PngImagePlugin.PngInfo()
+    meta.add_text("Prompt", prompt)
+    meta.add_text("NegativePrompt", negative_prompt)
+    meta.add_text("Model", model_id)
+    meta.add_text("VAE", generation_config["vae"])
+    meta.add_text("Sampler", generation_config["sampler"])
+    meta.add_text("Steps", str(steps))
+    meta.add_text("CFG_Scale", str(guidance))
+    if strength is not None:
+        meta.add_text("Strength", str(strength))
+    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
+
 def load_models():
     """Load both text2img and img2img pipelines optimized for Spaces"""
     global txt2img_pipe, img2img_pipe
@@ -108,6 +137,22 @@ def validate_dimensions(width: int, height: int) -> Tuple[int, int]:
     
     return width, height
 
+def format_status_with_metadata(generation_time: float, width: int, height: int, prompt: str, negative_prompt: str, seed: int, steps: int, guidance: float, strength: Optional[float] = None):
+    """Format status message with generation metadata"""
+    status_parts = [
+        f"✅ Generated in {generation_time:.1f}s ({width}×{height})",
+        f"🎯 Prompt: {prompt[:50]}..." if len(prompt) > 50 else f"🎯 Prompt: {prompt}",
+        f"🚫 Negative: {negative_prompt[:30]}..." if negative_prompt and len(negative_prompt) > 30 else f"🚫 Negative: {negative_prompt or 'None'}",
+        f"🎲 Seed: {seed}",
+        f"📏 Steps: {steps}",
+        f"🎛️ CFG: {guidance}"
+    ]
+    
+    if strength is not None:
+        status_parts.append(f"💪 Strength: {strength}")
+    
+    return "\n".join(status_parts)
+
 @spaces.GPU(duration=60)  # GPU decorator for Spaces
 def generate_txt2img(prompt, negative_prompt, num_steps, guidance_scale, width, height, seed, add_quality_tags):
     """Generate image from text prompt with Spaces GPU support"""
@@ -127,10 +172,12 @@ def generate_txt2img(prompt, negative_prompt, num_steps, guidance_scale, width,
         # Validate dimensions
         width, height = validate_dimensions(width, height)
         
+        # Handle seed
+        if seed == -1:
+            seed = random.randint(0, 2147483647)
+        
         # Set seed
-        generator = None
-        if seed != -1:
-            generator = torch.Generator(device=device).manual_seed(int(seed))
+        generator = torch.Generator(device=device).manual_seed(int(seed))
         
         # Enhance prompt
         enhanced_prompt = enhance_prompt(prompt, add_quality_tags)
@@ -151,7 +198,19 @@ def generate_txt2img(prompt, negative_prompt, num_steps, guidance_scale, width,
             )
         
         generation_time = time.time() - start_time
-        status = f"Generated in {generation_time:.1f}s ({width}x{height})"
+        
+        # Save with metadata
+        temp_path = tempfile.mktemp(suffix=".png")
+        add_metadata_and_save(
+            result.images[0], temp_path, enhanced_prompt, negative_prompt or "", 
+            seed, num_steps, guidance_scale
+        )
+        
+        # Format status with metadata
+        status = format_status_with_metadata(
+            generation_time, width, height, enhanced_prompt, 
+            negative_prompt or "", seed, num_steps, guidance_scale
+        )
         
         return result.images[0], status
         
@@ -179,10 +238,12 @@ def generate_img2img(input_image, prompt, negative_prompt, num_steps, guidance_s
     try:
         clear_memory()
         
+        # Handle seed
+        if seed == -1:
+            seed = random.randint(0, 2147483647)
+        
         # Set seed
-        generator = None
-        if seed != -1:
-            generator = torch.Generator(device=device).manual_seed(int(seed))
+        generator = torch.Generator(device=device).manual_seed(int(seed))
         
         # Enhance prompt
         enhanced_prompt = enhance_prompt(prompt, add_quality_tags)
@@ -215,7 +276,19 @@ def generate_img2img(input_image, prompt, negative_prompt, num_steps, guidance_s
             )
         
         generation_time = time.time() - start_time
-        status = f"Transformed in {generation_time:.1f}s (Strength: {strength})"
+        
+        # Save with metadata
+        temp_path = tempfile.mktemp(suffix=".png")
+        add_metadata_and_save(
+            result.images[0], temp_path, enhanced_prompt, negative_prompt or "", 
+            seed, num_steps, guidance_scale, strength
+        )
+        
+        # Format status with metadata
+        status = format_status_with_metadata(
+            generation_time, w, h, enhanced_prompt, 
+            negative_prompt or "", seed, num_steps, guidance_scale, strength
+        )
         
         return result.images[0], status
         
@@ -224,10 +297,26 @@ def generate_img2img(input_image, prompt, negative_prompt, num_steps, guidance_s
     finally:
         clear_memory()
 
+# Example prompts for inspiration
+EXAMPLE_PROMPTS = [
+    "beautiful anime girl with long flowing hair, cherry blossoms, soft lighting",
+    "cyberpunk cityscape at night, neon lights, rain reflections, detailed architecture",
+    "majestic dragon flying over mountains, fantasy landscape, dramatic clouds",
+    "cute anthropomorphic fox character, forest background, magical atmosphere",
+    "elegant woman in Victorian dress, portrait, ornate background, vintage style",
+    "futuristic robot with glowing eyes, metallic surface, sci-fi environment",
+    "mystical unicorn in enchanted forest, rainbow mane, sparkles, ethereal lighting",
+    "steampunk airship floating in sky, gears and brass, adventure scene"
+]
+
+def set_example_prompt():
+    """Return a random example prompt"""
+    return random.choice(EXAMPLE_PROMPTS)
+
 # Simplified negative prompt for better performance
 DEFAULT_NEGATIVE = """
-(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)
+(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)
 """
 
 # Gradio interface optimized for Spaces
@@ -241,17 +330,21 @@ with gr.Blocks(
     Generate high-quality images using the CyberRealistic Pony SDXL model.
     
     ⚠️ **Note**: First generation may take longer as the model loads. GPU time is limited on Spaces.
+    📋 **Metadata**: All generated images include embedded metadata (prompt, settings, seed, etc.)
     """)
     
     with gr.Tabs():
         with gr.TabItem("🎨 Text to Image"):
             with gr.Row():
                 with gr.Column():
-                    txt2img_prompt = gr.Textbox(
-                        label="Prompt",
-                        placeholder="beautiful landscape, mountains, sunset",
-                        lines=2
-                    )
+                    with gr.Row():
+                        txt2img_prompt = gr.Textbox(
+                            label="Prompt",
+                            placeholder="beautiful landscape, mountains, sunset",
+                            lines=2,
+                            scale=4
+                        )
+                        txt2img_example_btn = gr.Button("🎲 Random Example", scale=1)
                     
                     with gr.Accordion("Advanced Settings", open=False):
                         txt2img_negative = gr.Textbox(
@@ -273,24 +366,30 @@ with gr.Blocks(
                             txt2img_width = gr.Slider(512, 1024, 768, step=64, label="Width")
                             txt2img_height = gr.Slider(512, 1024, 768, step=64, label="Height")
                         
-                        txt2img_seed = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
+                        txt2img_seed = gr.Slider(
+                            minimum=-1, maximum=2147483647, value=-1, step=1,
+                            label="Seed (-1 for random)"
+                        )
                     
                     txt2img_btn = gr.Button("🎨 Generate", variant="primary", size="lg")
                 
                 with gr.Column():
                     txt2img_output = gr.Image(label="Generated Image", height=400)
-                    txt2img_status = gr.Textbox(label="Status", interactive=False)
+                    txt2img_status = gr.Textbox(label="Generation Info", interactive=False, lines=6)
         
         with gr.TabItem("🖼️ Image to Image"):
             with gr.Row():
                 with gr.Column():
                     img2img_input = gr.Image(label="Input Image", type="pil", height=250)
                     
-                    img2img_prompt = gr.Textbox(
-                        label="Prompt",
-                        placeholder="digital painting style, vibrant colors",
-                        lines=2
-                    )
+                    with gr.Row():
+                        img2img_prompt = gr.Textbox(
+                            label="Prompt",
+                            placeholder="digital painting style, vibrant colors",
+                            lines=2,
+                            scale=4
+                        )
+                        img2img_example_btn = gr.Button("🎲 Random Example", scale=1)
                     
                     with gr.Accordion("Advanced Settings", open=False):
                         img2img_negative = gr.Textbox(
@@ -313,13 +412,16 @@ with gr.Blocks(
                             label="Strength (Higher = more creative)"
                         )
                         
-                        img2img_seed = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
+                        img2img_seed = gr.Slider(
+                            minimum=-1, maximum=2147483647, value=-1, step=1,
+                            label="Seed (-1 for random)"
+                        )
                     
                     img2img_btn = gr.Button("🖼️ Transform", variant="primary", size="lg")
                 
                 with gr.Column():
                     img2img_output = gr.Image(label="Generated Image", height=400)
-                    img2img_status = gr.Textbox(label="Status", interactive=False)
+                    img2img_status = gr.Textbox(label="Generation Info", interactive=False, lines=6)
     
     # Event handlers
     txt2img_btn.click(
@@ -335,6 +437,17 @@ with gr.Blocks(
                 img2img_strength, img2img_seed, img2img_quality_tags],
         outputs=[img2img_output, img2img_status]
     )
+    
+    # Example prompt buttons
+    txt2img_example_btn.click(
+        fn=set_example_prompt,
+        outputs=[txt2img_prompt]
+    )
+    
+    img2img_example_btn.click(
+        fn=set_example_prompt,
+        outputs=[img2img_prompt]
+    )
 
 print(f"🚀 CyberRealistic Pony Generator initialized on {device}")