Update app.py

app.py

@@ -2,14 +2,29 @@ import gradio as gr
 import numpy as np
 import random
 import torch
-from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, FlowMatchEulerDiscreteScheduler
+from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler,  AutoencoderKL, UNet2DConditionModel
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
 import spaces
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16
 
-repo = "stabilityai/stable-cascade"
-pipe = StableDiffusion3Pipeline.from_pretrained(repo, torch_dtype=torch.float16).to(device)
+# Use the correct repo for SDXL
+repo = "stabilityai/sdxl-turbo"  # This is the correct repo for SDXL
+
+# Load the model components separately
+vae = AutoencoderKL.from_pretrained(repo, subfolder="vae", torch_dtype=torch.float16).to(device)
+text_encoder = SD3Transformer2DModel.from_pretrained(repo, subfolder="text_encoder", torch_dtype=torch.float16).to(device)
+unet = UNet2DConditionModel.from_pretrained(repo, subfolder="unet", torch_dtype=torch.float16).to(device)
+scheduler = EulerDiscreteScheduler.from_pretrained(repo, subfolder="scheduler", torch_dtype=torch.float16)
+
+# Construct the pipeline (this is how you work with SDXL)
+pipe = StableDiffusionPipeline(
+    vae=vae,
+    text_encoder=text_encoder,
+    unet=unet,
+    scheduler=scheduler
+).to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1344
@@ -22,6 +37,7 @@ def infer(prompts, negative_prompts, seeds, randomize_seeds, widths, heights, gu
 
         generator = torch.Generator().manual_seed(seeds[i])
 
+        # SDXL requires a slightly different call format:
         image = pipe(
             prompt=prompt,
             negative_prompt=negative_prompts[i],
@@ -53,7 +69,7 @@ with gr.Blocks(css=css) as demo:
 
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""
-        # Demo [Automated Stable Diffusion 3 Medium](https://huggingface.co/stabilityai/stable-diffusion-3-medium)
+        # Demo [Automated Stable Diffusion XL](https://huggingface.co/stabilityai/stablediffusion-xl)
         """)
 
         with gr.Row():