| 20 | |
| 21 | |
| 22 | class AnimateController: |
| 23 | def __init__( |
| 24 | self, |
| 25 | config_path="./configs/prompts/animation.yaml", |
| 26 | weight_dtype=torch.float16, |
| 27 | ): |
| 28 | # Read pretrained weights path from config |
| 29 | self.config = OmegaConf.load(config_path) |
| 30 | self.pipeline = None |
| 31 | self.weight_dtype = weight_dtype |
| 32 | |
| 33 | def animate( |
| 34 | self, |
| 35 | ref_image, |
| 36 | pose_video_path, |
| 37 | width=512, |
| 38 | height=768, |
| 39 | length=24, |
| 40 | num_inference_steps=25, |
| 41 | cfg=3.5, |
| 42 | seed=123, |
| 43 | ): |
| 44 | generator = torch.manual_seed(seed) |
| 45 | if isinstance(ref_image, np.ndarray): |
| 46 | ref_image = Image.fromarray(ref_image) |
| 47 | if self.pipeline is None: |
| 48 | vae = AutoencoderKL.from_pretrained( |
| 49 | self.config.pretrained_vae_path, |
| 50 | ).to("cuda", dtype=self.weight_dtype) |
| 51 | |
| 52 | reference_unet = UNet2DConditionModel.from_pretrained( |
| 53 | self.config.pretrained_base_model_path, |
| 54 | subfolder="unet", |
| 55 | ).to(dtype=self.weight_dtype, device="cuda") |
| 56 | |
| 57 | inference_config_path = self.config.inference_config |
| 58 | infer_config = OmegaConf.load(inference_config_path) |
| 59 | denoising_unet = UNet3DConditionModel.from_pretrained_2d( |
| 60 | self.config.pretrained_base_model_path, |
| 61 | self.config.motion_module_path, |
| 62 | subfolder="unet", |
| 63 | unet_additional_kwargs=infer_config.unet_additional_kwargs, |
| 64 | ).to(dtype=self.weight_dtype, device="cuda") |
| 65 | |
| 66 | pose_guider = PoseGuider(320, block_out_channels=(16, 32, 96, 256)).to( |
| 67 | dtype=self.weight_dtype, device="cuda" |
| 68 | ) |
| 69 | |
| 70 | image_enc = CLIPVisionModelWithProjection.from_pretrained( |
| 71 | self.config.image_encoder_path |
| 72 | ).to(dtype=self.weight_dtype, device="cuda") |
| 73 | sched_kwargs = OmegaConf.to_container(infer_config.noise_scheduler_kwargs) |
| 74 | scheduler = DDIMScheduler(**sched_kwargs) |
| 75 | |
| 76 | # load pretrained weights |
| 77 | denoising_unet.load_state_dict( |
| 78 | torch.load(self.config.denoising_unet_path, map_location="cpu"), |
| 79 | strict=False, |