Update app.py

app.py

@@ -3,16 +3,18 @@ import random
 import time
 import gradio as gr
 import torch
-from diffusers import DiffusionPipeline, LCMScheduler
+from diffusers import StableDiffusionPipeline, LCMScheduler
+from huggingface_hub import hf_hub_download
 
-# Configuration - Using optimized Flux model with LCM
-MODEL_ID = "ByteDance/Hyper-SD"
+# Configuration - Using optimized model
+MODEL_ID = "Lykon/dreamshaper-8-lcm"
 MODEL_CACHE = "model_cache"
 os.makedirs(MODEL_CACHE, exist_ok=True)
 
 # Load model with CPU optimizations
 def get_pipeline():
-    pipe = DiffusionPipeline.from_pretrained(
+    # Use traditional StableDiffusionPipeline instead of DiffusionPipeline
+    pipe = StableDiffusionPipeline.from_pretrained(
         MODEL_ID,
         torch_dtype=torch.float32,
         cache_dir=MODEL_CACHE,
@@ -36,18 +38,16 @@ print(f"Model loaded in {load_time:.2f} seconds")
 def generate_image(
     prompt: str,
     negative_prompt: str = "blurry, low quality, cartoon, drawing, text",
-    width: int = 768,
-    height: int = 768,
+    width: int = 512,  # Reduced for CPU performance
+    height: int = 512,  # Reduced for CPU performance
     seed: int = -1,
     guidance_scale: float = 2.0,
     num_inference_steps: int = 4
 ):
-    # Set seed if not provided
     if seed == -1:
         seed = random.randint(0, 2147483647)
     generator = torch.Generator(device="cpu").manual_seed(seed)
     
-    # Generate image with timing
     start_gen = time.time()
     with torch.no_grad():
         image = pipeline(
@@ -64,73 +64,33 @@ def generate_image(
     print(f"Generated {width}x{height} image in {gen_time:.2f} seconds")
     return image, seed, f"Generated in {gen_time:.2f}s | Loaded in {load_time:.2f}s"
 
-# Create optimized interface
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="teal")) as demo:
-    gr.Markdown("""
-    # ⚡ FLUX Turbo Generator
-    **Professional Quality Images · Lightning Fast CPU Generation**
-    """)
+# Create interface
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# ⚡ Turbo Image Generator")
     
     with gr.Row():
-        with gr.Column(scale=1):
-            prompt = gr.Textbox(
-                label="Prompt",
-                placeholder="Describe a professional, high-quality image...",
-                lines=3
-            )
-            negative_prompt = gr.Textbox(
-                label="Negative Prompt",
-                value="blurry, low quality, cartoon, drawing, text"
-            )
-            generate_btn = gr.Button("Generate Image", variant="primary")
+        with gr.Column():
+            prompt = gr.Textbox(label="Prompt", lines=3)
+            negative_prompt = gr.Textbox(label="Negative Prompt", value="blurry, low quality")
+            generate_btn = gr.Button("Generate", variant="primary")
             
-            with gr.Accordion("Advanced Settings", open=False):
-                with gr.Row():
-                    width = gr.Slider(512, 1024, value=768, step=64, label="Width")
-                    height = gr.Slider(512, 1024, value=768, step=64, label="Height")
-                with gr.Row():
-                    guidance = gr.Slider(1.0, 5.0, value=2.0, step=0.1, label="Guidance")
-                    steps = gr.Slider(1, 8, value=4, step=1, label="Steps")
+            with gr.Accordion("Advanced", open=False):
+                width = gr.Slider(384, 768, value=512, step=64, label="Width")
+                height = gr.Slider(384, 768, value=512, step=64, label="Height")
+                guidance = gr.Slider(1.0, 5.0, value=2.0, step=0.1, label="Guidance")
+                steps = gr.Slider(1, 8, value=4, step=1, label="Steps")
                 seed = gr.Number(label="Seed", value=-1)
         
-        with gr.Column(scale=1):
-            output_image = gr.Image(label="Result", type="pil", height=500)
+        with gr.Column():
+            output_image = gr.Image(label="Result", type="pil")
             used_seed = gr.Textbox(label="Used Seed")
             perf_info = gr.Textbox(label="Performance Info")
-    
-    # Generation handler
+
     generate_btn.click(
         generate_image,
         inputs=[prompt, negative_prompt, width, height, seed, guidance, steps],
         outputs=[output_image, used_seed, perf_info]
     )
-    
-    # Professional examples
-    gr.Examples(
-        examples=[
-            [
-                "Professional photograph of a futuristic city at golden hour, cinematic lighting, ultra-detailed",
-                "blurry, cartoon, drawing, text, watermark",
-                768,
-                768
-            ],
-            [
-                "Hyperrealistic portrait of a wise elderly man, detailed wrinkles, studio lighting, 8k resolution",
-                "anime, cartoon, deformed, ugly",
-                768,
-                1024
-            ],
-            [
-                "Majestic mountain landscape with crystal clear lake reflection, autumn colors, sharp focus",
-                "low quality, blurry, people, buildings",
-                1024,
-                768
-            ]
-        ],
-        inputs=[prompt, negative_prompt, width, height],
-        label="Professional Examples"
-    )
 
-# Launch the app
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)))