Update app.py

app.py

@@ -10,8 +10,8 @@ import PIL
 base = "stabilityai/stable-diffusion-xl-base-1.0"
 repo = "tianweiy/DMD2"
 checkpoints = {
-    "1-Step" : ["dmd2_sdxl_1step_unet.bin", 1],
-    "4-Step" : ["dmd2_sdxl_4step_unet.bin", 4],
+    "1-Step" : ["dmd2_sdxl_1step_unet_fp16.bin", 1],
+    "4-Step" : ["dmd2_sdxl_4step_unet_fp16.bin", 4],
 }
 loaded = None
 
@@ -22,8 +22,8 @@ CSS = """
 """
 
 # Ensure model and scheduler are initialized in GPU-enabled function
-unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
 if torch.cuda.is_available():
+    unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
     pipe = DiffusionPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")
 
 
@@ -38,7 +38,7 @@ def generate_image(prompt, ckpt):
 
     if loaded != num_inference_steps:
         unet.load_state_dict(torch.load(hf_hub_download(repo, checkpoint), map_location="cuda"))
-        pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample" if num_inference_steps==1 else "epsilon")
+        pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
         loaded = num_inference_steps
         
     results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0)