Update app.py

app.py

@@ -28,28 +28,24 @@ def generate(prompt, option, progress=gr.Progress()):
 
     # Main pipeline.
     unet = UNet2DConditionModel.from_config(base, subfolder="unet")
-    unet.load_state_dict(load_file(hf_hub_download(repo, opts["4 Steps"][0]))).to(device, dtype)
     pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=dtype, variant="fp16")
-    pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing").to(device, dtype)
+    pipe.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device))
+    pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample" if step == 1 else "epsilon")
 
     # Safety checker.
     safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker").to(device, dtype)
     feature_extractor=CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32")
     image_processor = VaeImageProcessor(vae_scale_factor=8)
 
-    progress((0, step))
-
-    if step != step_loaded:
-        print(f"Switching checkpoint from {step_loaded} to {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))
-        step_loaded = step
     def inference_callback(p, i, t, kwargs):
         progress((i+1, step))
         return kwargs
+    
+    # Inference loop.
+    progress((0, step))
     results = pipe(prompt, num_inference_steps=step, guidance_scale=0, callback_on_step_end=inference_callback, output_type="pt")
 
-    # Safety check
+    # Safety check.
     feature_extractor_input = image_processor.postprocess(results.images, output_type="pil")
     safety_checker_input = feature_extractor(feature_extractor_input, return_tensors="pt")
     images, has_nsfw_concept = safety_checker(