Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import torch
 import spaces
-from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler
+from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 
@@ -9,26 +9,25 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 base = "stabilityai/stable-diffusion-xl-base-1.0"
 repo = "ByteDance/SDXL-Lightning"
 opts = {
-    "1 Step"  : ["sdxl_lightning_1step_unet_x0.safetensors", 1],
-    "2 Steps" : ["sdxl_lightning_2step_unet.safetensors", 2],
-    "4 Steps" : ["sdxl_lightning_4step_unet.safetensors", 4],
-    "8 Steps" : ["sdxl_lightning_8step_unet.safetensors", 8],
+    "1 Step"  : ("sdxl_lightning_1step_unet_x0.safetensors", 1),
+    "2 Steps" : ("sdxl_lightning_2step_unet.safetensors", 2),
+    "4 Steps" : ("sdxl_lightning_4step_unet.safetensors", 4),
+    "8 Steps" : ("sdxl_lightning_8step_unet.safetensors", 8),
 }
 
+step_loaded = 4
+unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
+unet.load_state_dict(load_file(hf_hub_download(repo, opts["4 Steps"][0]), device=device))
 pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to(device)
-last_step = None
 
-# Function 
 @spaces.GPU(enable_queue=True)
 def generate_image(prompt, option):
     ckpt, step = opts[option]
-    if last_step != step:
+    if step_loaded != step:
         pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample" if step == 1 else "epsilon")
         pipe.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device))
-        last_step = step
-    image = pipe(prompt, num_inference_steps=step, guidance_scale=0).images[0]
-    return image
-
+        step_loaded = step
+    return pipe(prompt, num_inference_steps=step, guidance_scale=0).images[0]
 
 with gr.Blocks() as demo:
     gr.HTML("<h1><center>SDXL-Lightning</center></h1>")