update infer script

Author: Yuwei Guo

scripts/animate.py

@@ -15,14 +15,12 @@
 from animatediff.models.unet import UNet3DConditionModel
 from animatediff.pipelines.pipeline_animation import AnimationPipeline
 from animatediff.utils.util import save_videos_grid
-from animatediff.utils.convert_from_ckpt import convert_ldm_unet_checkpoint, convert_ldm_clip_checkpoint, convert_ldm_vae_checkpoint
-from animatediff.utils.convert_lora_safetensor_to_diffusers import convert_lora
+from animatediff.utils.util import load_weights
 from diffusers.utils.import_utils import is_xformers_available
 
 from einops import rearrange, repeat
 
 import csv, pdb, glob
-from safetensors import safe_open
 import math
 from pathlib import Path
 
@@ -60,50 +58,16 @@ def main(args):
                 scheduler=DDIMScheduler(**OmegaConf.to_container(inference_config.noise_scheduler_kwargs)),
             ).to("cuda")
 
-            # 1. unet ckpt
-            # 1.1 motion module
-            motion_module_state_dict = torch.load(motion_module, map_location="cpu")
-            if "global_step" in motion_module_state_dict: func_args.update({"global_step": motion_module_state_dict["global_step"]})
-            missing, unexpected = pipeline.unet.load_state_dict(motion_module_state_dict, strict=False)
-            assert len(unexpected) == 0
-            
-            # 1.2 T2I
-            if model_config.path != "":
-                if model_config.path.endswith(".ckpt"):
-                    state_dict = torch.load(model_config.path)
-                    pipeline.unet.load_state_dict(state_dict)
-                    
-                elif model_config.path.endswith(".safetensors"):
-                    state_dict = {}
-                    with safe_open(model_config.path, framework="pt", device="cpu") as f:
-                        for key in f.keys():
-                            state_dict[key] = f.get_tensor(key)
-                            
-                    is_lora = all("lora" in k for k in state_dict.keys())
-                    if not is_lora:
-                        base_state_dict = state_dict
-                    else:
-                        base_state_dict = {}
-                        with safe_open(model_config.base, framework="pt", device="cpu") as f:
-                            for key in f.keys():
-                                base_state_dict[key] = f.get_tensor(key)                
-                    
-                    # vae
-                    converted_vae_checkpoint = convert_ldm_vae_checkpoint(base_state_dict, pipeline.vae.config)
-                    pipeline.vae.load_state_dict(converted_vae_checkpoint)
-                    # unet
-                    converted_unet_checkpoint = convert_ldm_unet_checkpoint(base_state_dict, pipeline.unet.config)
-                    pipeline.unet.load_state_dict(converted_unet_checkpoint, strict=False)
-                    # text_model
-                    pipeline.text_encoder = convert_ldm_clip_checkpoint(base_state_dict)
-                    
-                    # import pdb
-                    # pdb.set_trace()
-                    if is_lora:
-                        pipeline = convert_lora(pipeline, state_dict, alpha=model_config.lora_alpha)
-
-            pipeline.to("cuda")
-            ### <<< create validation pipeline <<< ###
+            pipeline = load_weights(
+                pipeline,
+                # motion module
+                motion_module_path         = motion_module,
+                motion_module_lora_configs = model_config.get("motion_module_lora_configs", []),
+                # image layers
+                dreambooth_model_path      = model_config.get("dreambooth_path", ""),
+                lora_model_path            = model_config.get("lora_model_path", ""),
+                lora_alpha                 = model_config.get("lora_alpha", 0.8),
+            ).to("cuda")
 
             prompts      = model_config.prompt
             n_prompts    = list(model_config.n_prompt) * len(prompts) if len(model_config.n_prompt) == 1 else model_config.n_prompt