Update app.py

app.py

@@ -1,35 +1,34 @@
 import torch
 import gradio as gr
-import spaces  # Import spaces for ZeroGPU support
+import spaces
 from functools import lru_cache
-from diffusers import StableDiffusionXLPipeline  # ✅ Correct pipeline for text-to-image
+from diffusers import StableDiffusionXLPipeline
 
-# LoRA model path on Hugging Face Hub
-color_book_lora_path = "artificialguybr/ColoringBookRedmond-V2"
-color_book_trigger = ", ColoringBookAF, Coloring Book"
+# ===============================
+# 🩹 FIX for Gradio bug (bool schema issue)
+# ===============================
+import gradio_client.utils as gu
 
-# Load model on CPU initially
-@lru_cache(maxsize=1)
-def load_pipeline(use_lora: bool):
-    """Load Stable Diffusion XL pipeline and LoRA weights (if selected)."""
-    
-    # ✅ Use StableDiffusionXLPipeline for text-to-image generation
-    pipe = StableDiffusionXLPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-xl-base-1.0",
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        use_safetensors=True
-    )
+# Monkey patch for "TypeError: argument of type 'bool' is not iterable"
+if not hasattr(gu, "_patched_json_schema_to_python_type"):
+    orig_get_type = gu.get_type
 
-    # Keep the model on CPU until GPU is requested
-    pipe.to("cpu")
+    def safe_get_type(schema):
+        # Ensure schema is always a dict before checking keys
+        if not isinstance(schema, dict):
+            return str(schema)
+        return orig_get_type(schema)
 
-    # Load LoRA if selected
-    if use_lora:
-        pipe.load_lora_weights(color_book_lora_path)
+    gu.get_type = safe_get_type
+    gu._patched_json_schema_to_python_type = True
 
-    return pipe
 
-# Define styles
+# ===============================
+# 🎨 Model and Styles Configuration
+# ===============================
+color_book_lora_path = "artificialguybr/ColoringBookRedmond-V2"
+color_book_trigger = ", ColoringBookAF, Coloring Book"
+
 styles = {
     "Neonpunk": {
         "prompt": "neonpunk style, cyberpunk, vaporwave, neon, vibrant, stunningly beautiful, crisp, "
@@ -54,25 +53,46 @@ styles = {
     }
 }
 
-@spaces.GPU  # ZeroGPU: Allocate GPU only when generating images
+
+# ===============================
+# 🚀 Pipeline Loader (with caching)
+# ===============================
+@lru_cache(maxsize=1)
+def load_pipeline(use_lora: bool):
+    """Load Stable Diffusion XL pipeline and optionally apply LoRA weights."""
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-xl-base-1.0",
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        use_safetensors=True
+    )
+    pipe.to("cpu")
+
+    if use_lora:
+        pipe.load_lora_weights(color_book_lora_path)
+
+    return pipe
+
+
+# ===============================
+# 🎨 Image Generation Function
+# ===============================
+@spaces.GPU  # ZeroGPU: allocate GPU only when generating
 def generate_image(prompt: str, style_name: str, use_lora: bool):
     """Generate an image using Stable Diffusion XL with optional LoRA fine-tuning."""
-    
-    # Load the pipeline (cached)
-    pipeline = load_pipeline(use_lora)
 
-    # Move model to GPU only when needed
+    # Load cached pipeline
+    pipeline = load_pipeline(use_lora)
     pipeline.to("cuda")
 
-    # Get the selected style details
+    # Retrieve style info
     style_prompt = styles.get(style_name, {}).get("prompt", "")
     negative_prompt = styles.get(style_name, {}).get("negative_prompt", "")
 
-    # Apply LoRA trigger phrase if enabled
+    # Add LoRA trigger if needed
     if use_lora:
         prompt += color_book_trigger
 
-    # ✅ Ensure text-to-image pipeline is used correctly
+    # Generate image
     image = pipeline(
         prompt=prompt + " " + style_prompt,
         negative_prompt="blurred, ugly, watermark, low resolution, " + negative_prompt,
@@ -80,12 +100,15 @@ def generate_image(prompt: str, style_name: str, use_lora: bool):
         guidance_scale=9.0
     ).images[0]
 
-    # Move model back to CPU to free GPU resources
+    # Move model back to CPU to release GPU
     pipeline.to("cpu")
 
     return image
 
-# Gradio Interface for Hugging Face Spaces (ZeroGPU-compatible)
+
+# ===============================
+# 🌐 Gradio Interface (for Spaces)
+# ===============================
 interface = gr.Interface(
     fn=generate_image,
     inputs=[
@@ -95,10 +118,16 @@ interface = gr.Interface(
     ],
     outputs=gr.Image(label="Generated Image"),
     title="🎨 AI Coloring Book & Style Generator",
-    description="Generate AI-powered art using Stable Diffusion XL on Hugging Face Spaces. "
-                "Choose a style or enable a LoRA fine-tuned coloring book effect."
+    description=(
+        "Generate AI-powered art using Stable Diffusion XL on Hugging Face Spaces. "
+        "Choose a style or enable a LoRA fine-tuned coloring book effect. "
+        "This app dynamically allocates GPU (ZeroGPU) only during generation."
+    )
 )
 
-# Run Gradio app for Hugging Face Spaces
+
+# ===============================
+# 🏁 Launch App
+# ===============================
 if __name__ == "__main__":
     interface.launch()