Update app.py

app.py

@@ -19,11 +19,14 @@ opts = {
     "8 Steps" : ("sdxl_lightning_8step_unet.safetensors", 8),
 }
 
+# Default to load 4-step model.
 step_loaded = 4
-unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(device, dtype)
-unet.load_state_dict(load_file(hf_hub_download(repo, opts["4 Steps"][0]), device=device))
+unet = UNet2DConditionModel.from_config(base, subfolder="unet")
+unet.load_state_dict(load_file(hf_hub_download(repo, opts["4 Steps"][0])))
 pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=dtype, variant="fp16").to(device, dtype)
+pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
 
+# Inference function.
 @spaces.GPU(enable_queue=True)
 def generate_image(prompt, option):
     global step_loaded