Update app.py

app.py

@@ -112,6 +112,7 @@ class AnimateController:
                 motion_module_dropdown, map_location="cpu")
             missing, unexpected = self.unet.load_state_dict(
                 motion_module_state_dict, strict=False)
+            del motion_module_state_dict
             assert len(unexpected) == 0
             return gr.Dropdown.update()
 
@@ -134,6 +135,9 @@ class AnimateController:
             converted_unet_checkpoint = convert_ldm_unet_checkpoint(
                 base_model_state_dict, self.unet.config)
             self.unet.load_state_dict(converted_unet_checkpoint, strict=False)
+            del converted_unet_checkpoint
+            del converted_vae_checkpoint
+            del base_model_state_dict
 
             # self.text_encoder = convert_ldm_clip_checkpoint(base_model_state_dict)
             return gr.Dropdown.update()
@@ -174,7 +178,7 @@ class AnimateController:
                 **OmegaConf.to_container(self.inference_config.noise_scheduler_kwargs))
         ).to("cuda")
 
-        original_state_dict = {k: v.cpu().clone() for k, v in pipeline.unet.state_dict().items()}
+        original_state_dict = {k: v.cpu().clone() for k, v in pipeline.unet.state_dict().items() if "motion_modules." not in k}
         pipeline.unet = convert_lcm_lora(pipeline.unet, self.lcm_lora_path, spatial_lora_slider)
 
         pipeline.to("cuda")
@@ -196,7 +200,7 @@ class AnimateController:
                 video_length=length_slider,
             ).videos
 
-        pipeline.unet.load_state_dict(original_state_dict)
+        pipeline.unet.load_state_dict(original_state_dict,strict=False)
         del original_state_dict
 
         save_sample_path = os.path.join(