Update app.py

app.py

@@ -1,70 +1,129 @@
 import gradio as gr
-from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
 import torch
+from datetime import datetime
 
-# Check if we have GPU
+# Optimization settings
+torch.backends.cuda.matmul.allow_tf32 = True  # Enable tf32 for faster matmuls on Ampere GPUs
+torch.backends.cudnn.allow_tf32 = True  # Enable tf32 for cuDNN
+
+# Check for GPU and set up the model
 device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32
 
-# Load the model (we'll use a smaller model for free tier compatibility)
+# Use a smaller, faster model (try these alternatives if this doesn't work well)
+# model_id = "prompthero/openjourney-v4"  # Midjourney style
+# model_id = "nitrosocke/redshift-diffusion"  # 3D render style
 model_id = "runwayml/stable-diffusion-v1-5"
-pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
-pipe = pipe.to(device)
 
-# Reduce memory usage
-pipe.enable_attention_slicing()
+# Load the pipeline with optimizations
+scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
+pipe = StableDiffusionPipeline.from_pretrained(
+    model_id,
+    scheduler=scheduler,
+    torch_dtype=dtype,
+    safety_checker=None,  # Disable safety checker for faster generation
+    requires_safety_checker=False
+).to(device)
+
+# Apply optimizations
+pipe.enable_xformers_memory_efficient_attention()  # Flash attention
+pipe.enable_attention_slicing(1)  # Minimal slicing for balance between speed and memory
 
-def generate_image(prompt, negative_prompt, steps, guidance_scale, seed):
+# Warm up the model (important for consistent speed)
+print("Warming up the model...")
+with torch.inference_mode():
+    _ = pipe("warmup", num_inference_steps=1, guidance_scale=1)
+
+def generate_image(
+    prompt: str,
+    negative_prompt: str = "",
+    steps: int = 15,
+    guidance_scale: float = 7.0,
+    seed: int = -1,
+    width: int = 512,
+    height: int = 512
+):
+    # Start timer
+    start_time = datetime.now()
+    
     # Set random seed if provided
+    generator = None
     if seed != -1:
         generator = torch.Generator(device=device).manual_seed(seed)
-    else:
-        generator = None
     
-    # Generate image
-    with torch.autocast(device):
+    # Generate image with optimizations
+    with torch.inference_mode(), torch.autocast(device):
         image = pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
             num_inference_steps=int(steps),
             guidance_scale=guidance_scale,
-            generator=generator
+            generator=generator,
+            width=width,
+            height=height
         ).images[0]
     
-    return image
+    # Calculate generation time
+    gen_time = (datetime.now() - start_time).total_seconds()
+    print(f"Generated image in {gen_time:.2f} seconds")
+    
+    return image, f"Generated in {gen_time:.2f}s | Steps: {steps} | CFG: {guidance_scale} | Seed: {seed if seed != -1 else 'random'}"
 
-# Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# 🎨 Diffusion Art Generator")
-    gr.Markdown("Generate images using Stable Diffusion")
+# Gradio interface with optimizations
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # ⚡ Ultra-Fast Diffusion Image Generator
+    *Optimized to generate images in under 5 seconds*
+    """)
     
     with gr.Row():
         with gr.Column():
-            prompt = gr.Textbox(label="Prompt", placeholder="A beautiful landscape with mountains and a lake")
-            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="blurry, low quality, distorted")
-            steps = gr.Slider(1, 50, value=25, label="Inference Steps")
-            guidance_scale = gr.Slider(1, 20, value=7.5, label="Guidance Scale")
-            seed = gr.Number(value=-1, label="Seed (-1 for random)")
-            generate_btn = gr.Button("Generate Image")
+            prompt = gr.Textbox(
+                label="Prompt",
+                placeholder="A beautiful landscape with mountains and a lake",
+                max_lines=3
+            )
+            negative_prompt = gr.Textbox(
+                label="Negative Prompt",
+                placeholder="blurry, low quality, distorted",
+                max_lines=2
+            )
+            
+            with gr.Accordion("Advanced Settings", open=False):
+                with gr.Row():
+                    steps = gr.Slider(8, 30, value=15, step=1, label="Inference Steps")
+                    guidance_scale = gr.Slider(1, 10, value=7.0, step=0.5, label="Guidance Scale")
+                with gr.Row():
+                    seed = gr.Number(value=-1, label="Seed (-1 for random)")
+                    width = gr.Slider(256, 768, value=512, step=64, label="Width")
+                    height = gr.Slider(256, 768, value=512, step=64, label="Height")
+            
+            generate_btn = gr.Button("Generate Image", variant="primary")
+            info_output = gr.Textbox(label="Generation Info", interactive=False)
         
         with gr.Column():
-            output_image = gr.Image(label="Generated Image", type="pil")
-
-    generate_btn.click(
-        fn=generate_image,
-        inputs=[prompt, negative_prompt, steps, guidance_scale, seed],
-        outputs=output_image
-    )
+            output_image = gr.Image(label="Generated Image", type="pil", show_download_button=True)
     
+    # Examples for quick testing
     gr.Examples(
         examples=[
-            ["A futuristic cityscape at sunset, digital art", "blurry, low quality", 25, 7.5, 42],
-            ["A cute corgi puppy wearing a superhero cape", "ugly, deformed", 30, 8, -1],
-            ["An astronaut riding a horse on Mars, photorealistic", "cartoon, drawing", 40, 9, 123]
+            ["A futuristic cyberpunk city at night, neon lights, rain reflections", "blurry, low quality", 15, 7.0, 42],
+            ["A cute robot cat, digital art, vibrant colors", "ugly, deformed", 12, 7.5, -1],
+            ["A majestic dragon flying over mountains at sunset, fantasy art", "cartoon, sketch", 20, 8.0, 123]
         ],
         inputs=[prompt, negative_prompt, steps, guidance_scale, seed],
-        outputs=output_image,
+        outputs=[output_image, info_output],
         fn=generate_image,
         cache_examples=True
     )
+    
+    generate_btn.click(
+        fn=generate_image,
+        inputs=[prompt, negative_prompt, steps, guidance_scale, seed, width, height],
+        outputs=[output_image, info_output]
+    )
 
-demo.launch()
+# For Hugging Face Spaces
+demo.queue(max_size=4)  # Limit concurrent requests to prevent OOM errors
+demo.launch(debug=False)