Update app.py

app.py

@@ -4,22 +4,23 @@ import random
 from diffusers import DiffusionPipeline, LCMScheduler
 import time
 
-# Load model with specific CPU optimizations to prevent black images
+# Load model
 model_id = "runwayml/stable-diffusion-v1-5"
 adapter_id = "latent-consistency/lcm-lora-sdv1-5"
 
+# Use float32 explicitly for CPU to avoid black images (NaN errors)
 pipe = DiffusionPipeline.from_pretrained(
     model_id, 
-    safety_checker=None, # Safety checkers sometimes cause black images on CPU
-    requires_safety_checker=False
+    safety_checker=None, 
+    torch_dtype=torch.float32 
 )
 pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
 pipe.load_lora_weights(adapter_id)
 pipe.to("cpu")
 
-# Estimation logic: Approx 8 seconds per step on 2-vCPUs for 512x512
 def update_estimate(steps, size):
-    base_time_per_step = 8 # estimated seconds
+    # 2-vCPU logic: ~8s per step at 512, scaled by total pixels
+    base_time_per_step = 8 
     resolution_factor = (int(size) / 512) ** 2
     estimate = round(steps * base_time_per_step * resolution_factor, 1)
     return f"⏳ Estimated generation time: ~{estimate} seconds"
@@ -27,16 +28,14 @@ def update_estimate(steps, size):
 def generate(prompt, size, steps, is_random, manual_seed):
     start_time = time.time()
     
-    # Handle Seed
     seed = random.randint(0, 1000000) if is_random else int(manual_seed)
-    generator = torch.Generator().manual_seed(seed)
+    generator = torch.Generator(device="cpu").manual_seed(seed)
     
-    # Generate image
-    # Note: 1024x1024 on 2-vCPU may take 5+ minutes and risk crashing RAM
-    image = pipe(
+    # Run inference
+    output = pipe(
         prompt=prompt, 
         num_inference_steps=int(steps), 
-        guidance_scale=1.0, # LCM standard
+        guidance_scale=1.0,
         width=int(size),
         height=int(size),
         generator=generator
@@ -45,19 +44,21 @@ def generate(prompt, size, steps, is_random, manual_seed):
     end_time = time.time()
     duration = round(end_time - start_time, 2)
     
-    return image, f"✅ Done in {duration}s | Seed: {seed}"
+    return output, f"✅ Done in {duration}s | Seed: {seed}"
 
-# UI Layout
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
+# Custom CSS to make it look clean
+css = ".gradio-container {max-width: 850px !important}"
+
+with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     gr.Markdown("# ⚡ Fast CPU Image Gen (SD 1.5 + LCM)")
     
     with gr.Row():
-        with gr.Column(scale=1):
-            prompt = gr.Textbox(label="Prompt", placeholder="A sunset over a futuristic city...")
+        with gr.Column():
+            prompt_input = gr.Textbox(label="Prompt", placeholder="A cute robot eating a taco...", lines=3)
             
             with gr.Accordion("Advanced Settings", open=False):
-                size = gr.Radio([128, 512, 1024], label="Image Size", value=512)
-                steps = gr.Slider(1, 6, value=4, step=1, label="Steps")
+                size_choice = gr.Radio([128, 256, 512, 768, 1024], label="Image Size", value=512)
+                step_slider = gr.Slider(1, 6, value=4, step=1, label="Steps (Keep low for CPU)")
                 
                 with gr.Row():
                     random_seed = gr.Checkbox(label="Randomize Seed", value=True)
@@ -65,22 +66,22 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 
                 estimate_display = gr.Markdown("⏳ Estimated generation time: ~32 seconds")
 
-            generate_btn = gr.Button("Generate", variant="primary")
+            generate_btn = gr.Button("Generate Image", variant="primary")
         
-        with gr.Column(scale=1):
+        with gr.Column():
             status_label = gr.Label(value="Ready", label="Status")
-            image_out = gr.Image(label="Output")
+            image_out = gr.Image(label="Output", type="pil")
 
-    # Dynamic visibility for Seed
+    # Visibility Logic
     random_seed.change(lambda x: gr.update(visible=not x), inputs=random_seed, outputs=seed_val)
     
-    # Update estimate whenever settings change
-    size.change(update_estimate, [steps, size], estimate_display)
-    steps.change(update_estimate, [steps, size], estimate_display)
+    # Estimation Logic
+    size_choice.change(update_estimate, [step_slider, size_choice], estimate_display)
+    step_slider.change(update_estimate, [step_slider, size_choice], estimate_display)
 
     generate_btn.click(
         fn=generate, 
-        inputs=[prompt, size, steps, random_seed, seed_val], 
+        inputs=[prompt_input, size_choice, step_slider, random_seed, seed_val], 
         outputs=[image_out, status_label]
     )