Update app.py

app.py

@@ -1,12 +1,36 @@
-import gradio as gr
 import torch
+import gradio as gr
+import spaces  # Import spaces for ZeroGPU support
+from functools import lru_cache
 from diffusers import DiffusionPipeline
 
-# Path for the LoRA model
+# LoRA model path on Hugging Face Hub
 color_book_lora_path = "artificialguybr/ColoringBookRedmond-V2"
 color_book_trigger = ", ColoringBookAF, Coloring Book"
 
-# Predefined styles for users to choose from
+# Load model on CPU initially
+@lru_cache(maxsize=1)
+def load_pipeline(use_lora: bool):
+    """Load Stable Diffusion pipeline and LoRA weights (if selected)."""
+    
+    # Load the base model
+    pipe = DiffusionPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-xl-refiner-1.0",
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        use_safetensors=True,
+        variant="fp16" if torch.cuda.is_available() else None
+    )
+
+    # Keep the model on CPU until GPU is requested
+    pipe.to("cpu")
+
+    # Load LoRA if selected
+    if use_lora:
+        pipe.load_lora_weights(color_book_lora_path)
+
+    return pipe
+
+# Define styles
 styles = {
     "Neonpunk": {
         "prompt": "neonpunk style, cyberpunk, vaporwave, neon, vibrant, stunningly beautiful, crisp, "
@@ -31,68 +55,40 @@ styles = {
     }
 }
 
-# Cache the pipeline to avoid reloading it multiple times
-from functools import lru_cache
-@lru_cache(maxsize=1)
-def load_pipeline(use_lora: bool):
-    """Load the diffusion pipeline with or without LoRA weights."""
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    # Load the base model
-    pipe = DiffusionPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-xl-refiner-1.0",
-        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-        use_safetensors=True,
-        variant="fp16" if device == "cuda" else None
-    )
-
-    # Apply LoRA weights if selected
-    if use_lora:
-        pipe.load_lora_weights(color_book_lora_path)
-
-    return pipe
-
+@spaces.GPU  # ZeroGPU: Allocate GPU only when generating images
 def generate_image(prompt: str, style_name: str, use_lora: bool):
-    """Generates an AI-generated image based on user input and selected style."""
-    try:
-        # Load the pre-trained pipeline
-        pipeline = load_pipeline(use_lora)
+    """Generate an image using Stable Diffusion with optional LoRA fine-tuning."""
+    
+    # Load the pipeline (cached)
+    pipeline = load_pipeline(use_lora)
+
+    # Move model to GPU only when needed
+    pipeline.to("cuda")
 
-        # Retrieve style prompts
-        style_prompt = styles.get(style_name, {}).get("prompt", "")
-        negative_prompt = styles.get(style_name, {}).get("negative_prompt", "")
+    # Get the selected style details
+    style_prompt = styles.get(style_name, {}).get("prompt", "")
+    negative_prompt = styles.get(style_name, {}).get("negative_prompt", "")
 
-        # Apply LoRA trigger phrase if selected
-        if use_lora:
-            prompt += color_book_trigger
+    # Apply LoRA trigger phrase if enabled
+    if use_lora:
+        prompt += color_book_trigger
 
-        # Generate the image
-        image = pipeline(
-            prompt=prompt + " " + style_prompt,
-            negative_prompt="blurred, ugly, watermark, low resolution, " + negative_prompt,
-            num_inference_steps=20,
-            guidance_scale=9.0
-        ).images[0]
+    # Generate image
+    image = pipeline(
+        prompt=prompt + " " + style_prompt,
+        negative_prompt="blurred, ugly, watermark, low resolution, " + negative_prompt,
+        num_inference_steps=20,
+        guidance_scale=9.0
+    ).images[0]
 
-        return image
+    # Move model back to CPU to free GPU resources
+    pipeline.to("cpu")
 
-    except Exception as e:
-        return f"Error generating image: {str(e)}"
+    return image
 
-# Gradio UI
+# Gradio Interface for Hugging Face Spaces (ZeroGPU-compatible)
 interface = gr.Interface(
     fn=generate_image,
     inputs=[
         gr.Textbox(label="Enter Your Prompt", placeholder="A cute lion"),
-        gr.Dropdown(label="Select a Style", choices=list(styles.keys()), value="None"),
-        gr.Checkbox(label="Use Coloring Book LoRA", value=False)
-    ],
-    outputs=gr.Image(label="Generated Image"),
-    title="🎨 AI Coloring Book & Style Generator",
-    description="Generate amazing AI-powered art using Stable Diffusion XL. "
-                "Choose a style or apply a Coloring Book LoRA for unique results."
-)
-
-# Run the Gradio interface
-if __name__ == "__main__":
-    interface.launch()
+        gr.Dropdown(label="Select a Style", choices=list(styles.keys()