Update app.py

app.py

@@ -18,12 +18,13 @@ opts = {
 step_loaded = 4
 unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(device, 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)
+pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to(device)
 
 @spaces.GPU(enable_queue=True)
 def generate_image(prompt, option):
+    global step_loaded
     ckpt, step = opts[option]
-    if step_loaded != step:
+    if step != step_loaded:
         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))
         step_loaded = step