Update app.py

app.py

@@ -2,7 +2,7 @@ import torch
 import gradio as gr
 import spaces  # Import spaces for ZeroGPU support
 from functools import lru_cache
-from diffusers import DiffusionPipeline
+from diffusers import StableDiffusionXLPipeline  # ✅ Correct pipeline for text-to-image
 
 # LoRA model path on Hugging Face Hub
 color_book_lora_path = "artificialguybr/ColoringBookRedmond-V2"
@@ -11,14 +11,13 @@ color_book_trigger = ", ColoringBookAF, Coloring Book"
 # 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 Stable Diffusion XL pipeline and LoRA weights (if selected)."""
     
-    # Load the base model
-    pipe = DiffusionPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-xl-refiner-1.0",
+    # ✅ 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,
-        variant="fp16" if torch.cuda.is_available() else None
+        use_safetensors=True
     )
 
     # Keep the model on CPU until GPU is requested
@@ -57,7 +56,7 @@ styles = {
 
 @spaces.GPU  # ZeroGPU: Allocate GPU only when generating images
 def generate_image(prompt: str, style_name: str, use_lora: bool):
-    """Generate an image using Stable Diffusion with optional LoRA fine-tuning."""
+    """Generate an image using Stable Diffusion XL with optional LoRA fine-tuning."""
     
     # Load the pipeline (cached)
     pipeline = load_pipeline(use_lora)
@@ -73,7 +72,7 @@ def generate_image(prompt: str, style_name: str, use_lora: bool):
     if use_lora:
         prompt += color_book_trigger
 
-    # Generate image
+    # ✅ Ensure text-to-image pipeline is used correctly
     image = pipeline(
         prompt=prompt + " " + style_prompt,
         negative_prompt="blurred, ugly, watermark, low resolution, " + negative_prompt,