updated

app.py

@@ -5,6 +5,7 @@ from huggingface_hub import hf_hub_download
 import numpy as np
 from PIL import Image
 import os
+import gc
 
 # Suppress symlink warnings
 os.environ['HF_HUB_DISABLE_SYMLINKS_WARNING'] = "1"
@@ -38,15 +39,26 @@ styles = {
     }
 }
 
+# Pre-generate example images
+example_images = {
+    "glitch": "examples/glitch_example.jpg",
+    "anime80s": "examples/anime80s_example.jpg",
+    "night": "examples/night_example.jpg"
+}
+
 def load_pipeline():
     """Load and prepare the pipeline with all style embeddings"""
     # Check if CUDA is available
     device = "cuda" if torch.cuda.is_available() else "cpu"
     dtype = torch.float16 if device == "cuda" else torch.float32
 
+    # Use smaller model for CPU
+    model_id = "runwayml/stable-diffusion-v1-5" if device == "cuda" else "CompVis/stable-diffusion-v1-4"
+    
     pipe = StableDiffusionPipeline.from_pretrained(
-        "CompVis/stable-diffusion-v1-4",
-        torch_dtype=dtype
+        model_id,
+        torch_dtype=dtype,
+        low_cpu_mem_usage=True
     ).to(device)
     
     # Load all embeddings
@@ -69,6 +81,28 @@ def apply_purple_guidance(image, strength=0.5):
 
 def generate_image(prompt, style, seed, apply_guidance, guidance_strength=0.5):
     """Generate an image with selected style and optional purple guidance"""
+    # Check if this is one of our examples with pre-generated images
+    if prompt == "A serene mountain landscape with a lake at sunset" and style == "glitch" and seed == 42:
+        if os.path.exists(example_images["glitch"]):
+            image = Image.open(example_images["glitch"])
+            if apply_guidance:
+                image = apply_purple_guidance(image, guidance_strength)
+            return image
+    
+    if prompt == "A magical forest at twilight" and style == "anime80s" and seed == 789:
+        if os.path.exists(example_images["anime80s"]):
+            image = Image.open(example_images["anime80s"])
+            if apply_guidance:
+                image = apply_purple_guidance(image, guidance_strength)
+            return image
+    
+    if prompt == "A cyberpunk city at night" and style == "night" and seed == 456:
+        if os.path.exists(example_images["night"]):
+            image = Image.open(example_images["night"])
+            if apply_guidance:
+                image = apply_purple_guidance(image, guidance_strength)
+            return image
+    
     if style not in styles:
         return None
     
@@ -82,20 +116,39 @@ def generate_image(prompt, style, seed, apply_guidance, guidance_strength=0.5):
     # Create styled prompt
     styled_prompt = f"{prompt} {style_info['token']}"
     
-    # Generate image
-    image = pipe(
-        styled_prompt,
-        generator=generator,
-        guidance_scale=7.5,
-        num_inference_steps=20
-    ).images[0]
+    # Generate image with reduced settings for CPU
+    if device == "cpu":
+        # Use much smaller image size and fewer steps on CPU
+        image = pipe(
+            styled_prompt,
+            generator=generator,
+            guidance_scale=7.5,
+            num_inference_steps=10,  # Reduced steps
+            height=256,  # Smaller height
+            width=256    # Smaller width
+        ).images[0]
+    else:
+        image = pipe(
+            styled_prompt,
+            generator=generator,
+            guidance_scale=7.5,
+            num_inference_steps=20
+        ).images[0]
     
     # Apply purple guidance if requested
     if apply_guidance:
         image = apply_purple_guidance(image, guidance_strength)
     
+    # Clean up memory
+    gc.collect()
+    if device == "cuda":
+        torch.cuda.empty_cache()
+    
     return image
 
+# Create examples directory
+os.makedirs("examples", exist_ok=True)
+
 # Initialize the pipeline globally
 print("Loading pipeline and embeddings...")
 pipe = load_pipeline()
@@ -113,14 +166,35 @@ demo = gr.Interface(
     outputs=gr.Image(label="Generated Image"),
     title="Style-Guided Image Generation with Purple Enhancement",
     description="""Generate images in different styles with optional purple color guidance.
-    Choose a style, enter a prompt, and optionally apply purple color enhancement.""",
+    Choose a style, enter a prompt, and optionally apply purple color enhancement.
+    Note: Generation may take a few minutes on CPU.""",
     examples=[
         ["A serene mountain landscape with a lake at sunset", "glitch", 42, True, 0.5],
         ["A magical forest at twilight", "anime80s", 789, True, 0.7],
         ["A cyberpunk city at night", "night", 456, False, 0.5],
     ],
-    cache_examples=True
+    cache_examples=True,
+    allow_flagging="never"  # Disable flagging to reduce overhead
 )
 
 if __name__ == "__main__":
-    demo.launch(share=False) 
+    # Generate and save example images if they don't exist
+    if not all(os.path.exists(path) for path in example_images.values()):
+        print("Pre-generating example images...")
+        # Example 1
+        if not os.path.exists(example_images["glitch"]):
+            img = generate_image("A serene mountain landscape with a lake at sunset", "glitch", 42, False, 0.5)
+            img.save(example_images["glitch"])
+        
+        # Example 2
+        if not os.path.exists(example_images["anime80s"]):
+            img = generate_image("A magical forest at twilight", "anime80s", 789, False, 0.7)
+            img.save(example_images["anime80s"])
+        
+        # Example 3
+        if not os.path.exists(example_images["night"]):
+            img = generate_image("A cyberpunk city at night", "night", 456, False, 0.5)
+            img.save(example_images["night"])
+    
+    # Launch the app
+    demo.launch(share=False, show_error=True)