Update app.py

app.py

@@ -172,7 +172,7 @@ def main():
     global progress
         
     device = "cuda"
-    dtype = torch.bfloat16
+    dtype = torch.float16
     result=[]
     step = 2
 
@@ -183,8 +183,8 @@ def main():
     repo = "ByteDance/SDXL-Lightning"
     ckpt = f"sdxl_lightning_{step}step_unet.safetensors"
 
-    #unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(device, dtype)
-    #unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device))
+    unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(device, dtype)
+    unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device))
 
     repo = "ByteDance/AnimateDiff-Lightning"
     ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
@@ -192,7 +192,7 @@ def main():
     adapter = MotionAdapter().to(device, dtype)
     adapter.load_state_dict(load_file(hf_hub_download(repo ,ckpt), device=device))
 
-    pipe = AnimateDiffPipeline.from_pretrained(base, motion_adapter=adapter, torch_dtype=dtype, variant="fp16").to(device)
+    pipe = AnimateDiffPipeline.from_pretrained(base, unet=unet, motion_adapter=adapter, torch_dtype=dtype, variant="fp16").to(device, dtype)
     pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear")
     
     mp.set_start_method("spawn", force=True)