Update app.py

app.py

@@ -149,7 +149,7 @@ class AnimateController:
                 for key in f.keys():
                     self.lora_model_state_dict[key] = f.get_tensor(key)
         return gr.Dropdown.update()
-
+    @torch.no_grad()
     def animate(
         self,
         lora_alpha_slider,
@@ -174,8 +174,8 @@ class AnimateController:
                 **OmegaConf.to_container(self.inference_config.noise_scheduler_kwargs))
         ).to("cuda")
 
-        pipeline.unet = convert_lcm_lora(copy.deepcopy(
-            self.unet), self.lcm_lora_path, spatial_lora_slider)
+        original_state_dict = {k: v.cpu().clone() for k, v in model.state_dict().items()}
+        pipeline.unet = convert_lcm_lora(pipeline.unet, self.lcm_lora_path, spatial_lora_slider)
 
         pipeline.to("cuda")
 
@@ -185,15 +185,19 @@ class AnimateController:
             torch.seed()
         seed = torch.initial_seed()
 
-        sample = pipeline(
-            prompt_textbox,
-            negative_prompt=negative_prompt_textbox,
-            num_inference_steps=sample_step_slider,
-            guidance_scale=cfg_scale_slider,
-            width=width_slider,
-            height=height_slider,
-            video_length=length_slider,
-        ).videos
+        with torch.autocast("cuda"):
+            sample = pipeline(
+                prompt_textbox,
+                negative_prompt=negative_prompt_textbox,
+                num_inference_steps=sample_step_slider,
+                guidance_scale=cfg_scale_slider,
+                width=width_slider,
+                height=height_slider,
+                video_length=length_slider,
+            ).videos
+
+        pipeline.unet.load(original_state_dict)
+        del original_state_dict
 
         save_sample_path = os.path.join(
             self.savedir_sample, f"{sample_idx}.mp4")