AnimateAnyone と OpenVINO による画像からビデオへの合成#

%%skip not $SHOULD_CONVERT 
if not VAE_ENCODER_PATH.exists(): 
    class VaeEncoder(torch.nn.Module): 
        def __init__(self, vae): 
            super().__init__() 
            self.vae = vae 

        def forward(self, x): 
            return self.vae.encode(x).latent_dist.mean 
    vae.eval() 
    with torch.no_grad(): 
        vae_encoder = ov.convert_model(VaeEncoder(vae), example_input=torch.zeros(1,3,512,448)) 
    vae_encoder = nncf.compress_weights(vae_encoder) 
    ov.save_model(vae_encoder, VAE_ENCODER_PATH) 
    del vae_encoder 
    cleanup_torchscript_cache()

%%skip not $SHOULD_CONVERT 
if not VAE_DECODER_PATH.exists(): 
    class VaeDecoder(torch.nn.Module): 
        def __init__(self, vae): 
            super().__init__() 
            self.vae = vae 

    def forward(self, z): 
        return self.vae.decode(z).sample 
    vae.eval() 
    with torch.no_grad(): 
        vae_decoder = ov.convert_model(VaeDecoder(vae), example_input=torch.zeros(1,4,HEIGHT//8,WIDTH//8)) 
    vae_decoder = nncf.compress_weights(vae_decoder) 
    ov.save_model(vae_decoder, VAE_DECODER_PATH) 
    del vae_decoder 
    cleanup_torchscript_cache() 
del vae 
gc.collect()

%%skip not $SHOULD_CONVERT 
if not REFERENCE_UNET_PATH.exists(): 
    class ReferenceUNetWrapper(torch.nn.Module): 
        def __init__(self, reference_unet): 
            super().__init__() 
            self.reference_unet = reference_unet 

        def forward(self, sample, timestep, encoder_hidden_states): 
            return self.reference_unet(sample, timestep, encoder_hidden_states, return_dict=False)[1] 

    sample = torch.zeros(2, 4, HEIGHT // 8, WIDTH // 8) 
    timestep = torch.tensor(0) 
    encoder_hidden_states = torch.zeros(2, 1, 768) 
    reference_unet.eval() 
    with torch.no_grad(): 
        wrapper = ReferenceUNetWrapper(reference_unet) 
        example_input = (sample, timestep, encoder_hidden_states) 
        ref_features_shapes = {k: v.shape for k, v in wrapper(*example_input).items()} 
        ov_reference_unet = ov.convert_model( 
            wrapper, 
            example_input=example_input, 
        ) 
    ov_reference_unet = nncf.compress_weights(ov_reference_unet) 
    ov.save_model(ov_reference_unet, REFERENCE_UNET_PATH) 
    del ov_reference_unet 
    del wrapper 
    cleanup_torchscript_cache() 
del reference_unet 
gc.collect()

%%skip not $SHOULD_CONVERT 
if not DENOISING_UNET_PATH.exists(): 
    class DenoisingUNetWrapper(torch.nn.Module): 
        def __init__(self, denoising_unet): 
            super().__init__() 
            self.denoising_unet = denoising_unet 

        def forward( 
            self, 
            sample, 
            timestep, 
            encoder_hidden_states, 
            pose_cond_fea, 
            ref_features 
        ): 
            return self.denoising_unet( 
                sample, 
                timestep, 
                encoder_hidden_states, 
                ref_features, pose_cond_fea=pose_cond_fea, 
                return_dict=False) 

    example_input = { 
        "sample": torch.zeros(2, 4, VIDEO_LENGTH, HEIGHT // 8, WIDTH // 8), 
        "timestep": torch.tensor(999), "encoder_hidden_states": torch.zeros(2,1,768), 
        "pose_cond_fea": torch.zeros(2, 320, VIDEO_LENGTH, HEIGHT // 8, WIDTH // 8), 
        "ref_features": {k: torch.zeros(shape) for k, shape in ref_features_shapes.items()} 
    } 

    denoising_unet.eval() 
    with torch.no_grad(): 
        ov_denoising_unet = ov.convert_model( 
            DenoisingUNetWrapper(denoising_unet), 
            example_input=tuple(example_input.values()) 
        ) 
  ov_denoising_unet.inputs[0].get_node().set_partial_shape(ov.PartialShape((2, 4, VIDEO_LENGTH, HEIGHT // 8, WIDTH // 8))) 
ov_denoising_unet.inputs[2].get_node().set_partial_shape(ov.PartialShape((2, 1, 768))) 
ov_denoising_unet.inputs[3].get_node().set_partial_shape(ov.PartialShape((2, 320, VIDEO_LENGTH, HEIGHT // 8, WIDTH // 8))) 
    for ov_input, shape in zip(ov_denoising_unet.inputs[4:], ref_features_shapes.values()): 
        ov_input.get_node().set_partial_shape(ov.PartialShape(shape)) 
        ov_input.get_node().set_element_type(ov.Type.f32) 
    ov_denoising_unet.validate_nodes_and_infer_types() 
    ov_denoising_unet = nncf.compress_weights(ov_denoising_unet) 
    ov.save_model(ov_denoising_unet, DENOISING_UNET_PATH) 
    del ov_denoising_unet 
    cleanup_torchscript_cache() 
del denoising_unet 
gc.collect()

%%skip not $SHOULD_CONVERT 
if not POSE_GUIDER_PATH.exists(): 
    pose_guider.eval() 
    with torch.no_grad(): 
        ov_pose_guider = ov.convert_model(pose_guider, example_input=torch.zeros(1, 3, VIDEO_LENGTH, HEIGHT, WIDTH)) 
    ov_pose_guider = nncf.compress_weights(ov_pose_guider) 
    ov.save_model(ov_pose_guider, POSE_GUIDER_PATH) 
    del ov_pose_guider 
    cleanup_torchscript_cache() 
del pose_guider 
gc.collect()

%%skip not $SHOULD_CONVERT 
if not IMAGE_ENCODER_PATH.exists(): 
    image_enc.eval() 
    with torch.no_grad(): 
        ov_image_encoder = ov.convert_model(image_enc, example_input=torch.zeros(1, 3, 224, 224), input=(1, 3, 224, 224)) 
    ov_image_encoder = nncf.compress_weights(ov_image_encoder) 
    ov.save_model(ov_image_encoder, IMAGE_ENCODER_PATH) 
    del ov_image_encoder 
    cleanup_torchscript_cache() 
del image_enc 
gc.collect()

device = widgets.Dropdown( 
    options=core.available_devices + ["AUTO"], 
    value="AUTO", 
    description="Device:", 
    disabled=False, 
) 

device

class OVPose2VideoPipeline(Pose2VideoPipeline): 
    def __init__( 
        self, 
        vae_encoder_path=VAE_ENCODER_PATH, 
        vae_decoder_path=VAE_DECODER_PATH, 
        image_encoder_path=IMAGE_ENCODER_PATH, 
        reference_unet_path=REFERENCE_UNET_PATH, 
        denoising_unet_path=DENOISING_UNET_PATH, 
        pose_guider_path=POSE_GUIDER_PATH, 
        device=device.value, 
): 
    self.vae_encoder = core.compile_model(vae_encoder_path, device) 
    self.vae_decoder = core.compile_model(vae_decoder_path, device) 
    self.image_encoder = core.compile_model(image_encoder_path, device) 
    self.reference_unet = core.compile_model(reference_unet_path, device) 
    self.denoising_unet = core.compile_model(denoising_unet_path, device) 
    self.pose_guider = core.compile_model(pose_guider_path, device) 
    self.scheduler = DDIMScheduler(**OmegaConf.to_container(infer_config.noise_scheduler_kwargs)) 

    self.vae_scale_factor = 8 
    self.clip_image_processor = CLIPImageProcessor() 
    self.ref_image_processor = VaeImageProcessor(do_convert_rgb=True) 
    self.cond_image_processor = VaeImageProcessor(do_convert_rgb=True, do_normalize=False) 

    def decode_latents(self, latents): 
        video_length = latents.shape[2] 
        latents = 1 / 0.18215 * latents 
        latents = rearrange(latents, "b c f h w -> (b f) c h w") 
        # video = self.vae.decode(latents).sample 
        video = [] 
        for frame_idx in tqdm(range(latents.shape[0])): 
            video.append(torch.from_numpy(self.vae_decoder(latents[frame_idx : frame_idx + 1])[0])) 
        video = torch.cat(video) 
        video = rearrange(video, "(b f) c h w -> b c f h w", f=video_length) 
        video = (video / 2 + 0.5).clamp(0, 1) 
        # 常に float32 にキャストします。これは大きなオーバーヘッドを引き起こさない bfloa16 と互換性があるためです 
        video = video.cpu().float().numpy() 
        return video 

    def __call__( 
        self, 
        ref_image, 
        pose_images, 
        width, 
        height, 
        video_length, 
        num_inference_steps=30, 
        guidance_scale=3.5, 
        num_images_per_prompt=1, 
        eta: float = 0.0, 
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, 
        output_type: Optional[str] = "tensor", 
        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, 
        callback_steps: Optional[int] = 1, 
        context_schedule="uniform", 
        context_frames=24, 
        context_stride=1, 
        context_overlap=4, 
        context_batch_size=1, 
        interpolation_factor=1, 
        **kwargs, 
    ): 
        do_classifier_free_guidance = guidance_scale > 1.0 

        # タイムステップを準備 
        self.scheduler.set_timesteps(num_inference_steps) 
        timesteps = self.scheduler.timesteps 

        batch_size = 1 

        # クリップ画像の埋め込みを準備 
        clip_image = self.clip_image_processor.preprocess(ref_image.resize((224, 224)), return_tensors="pt").pixel_values 
        clip_image_embeds = self.image_encoder(clip_image)["image_embeds"] 
        clip_image_embeds = torch.from_numpy(clip_image_embeds) 
        encoder_hidden_states = clip_image_embeds.unsqueeze(1) 
        uncond_encoder_hidden_states = torch.zeros_like(encoder_hidden_states) 

        if do_classifier_free_guidance: 
            encoder_hidden_states = torch.cat([uncond_encoder_hidden_states, encoder_hidden_states], dim=0) 

        latents = self.prepare_latents( 
            batch_size * num_images_per_prompt, 
            4, 
            width, 
            height, 
            video_length, 
            clip_image_embeds.dtype, 
            torch.device("cpu"), 
            generator, 
    ) 

    # 追加のステップ kwargs を準備 
    extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) 

    # 参照画像の潜在変数を準備 
    ref_image_tensor = self.ref_image_processor.preprocess(ref_image, height=height, width=width) # (bs, c, width, height) 
    ref_image_latents = self.vae_encoder(ref_image_tensor)[0] 
    ref_image_latents = ref_image_latents * 0.18215 # (b, 4, h, w) 
    ref_image_latents = torch.from_numpy(ref_image_latents) 

    # ポーズ条件画像のリストを準備 
    pose_cond_tensor_list = [] 
    for pose_image in pose_images: 
        pose_cond_tensor = self.cond_image_processor.preprocess(pose_image, height=height, width=width) 
        pose_cond_tensor = pose_cond_tensor.unsqueeze(2) # (bs, c, 1, h, w) 
        pose_cond_tensor_list.append(pose_cond_tensor) 
    pose_cond_tensor = torch.cat(pose_cond_tensor_list, dim=2) # (bs, c, t, h, w) 
    pose_fea = self.pose_guider(pose_cond_tensor)[0] 
    pose_fea = torch.from_numpy(pose_fea) 

    context_scheduler = get_context_scheduler(context_schedule) 

    # ノイズ除去ループ 
    num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order 
    with self.progress_bar(total=num_inference_steps) as progress_bar: 
        for i, t in enumerate(timesteps): 
            noise_pred = torch.zeros( 
                ( 
                    latents.shape[0] * (2 if do_classifier_free_guidance else 1), 
                    *latents.shape[1:], 
                ), 
                device=latents.device, 
                dtype=latents.dtype, 
            ) 
            counter = torch.zeros( 
                (1, 1, latents.shape[2], 1, 1), 
                device=latents.device, 
                dtype=latents.dtype, 
            ) 

            # 1. Forward reference image 
            if i == 0: 
                ref_features = self.reference_unet( 
                    ( 
                        ref_image_latents.repeat((2 if do_classifier_free_guidance else 1), 1, 1, 1), 
                        torch.zeros_like(t), 
                        # t, 
                        encoder_hidden_states, 
                    ) 
                ).values() 

            context_queue = list( 
                context_scheduler( 
                    0, 
                    num_inference_steps, 
                    latents.shape[2], 
                    context_frames, 
                    context_stride, 
                    0,
                ) 
            ) 
            num_context_batches = math.ceil(len(context_queue) / context_batch_size) 

            context_queue = list( 
                context_scheduler( 
                    0, 
                    num_inference_steps, 
                    latents.shape[2], 
                    context_frames, 
                    context_stride, 
                    context_overlap, 
                ) 
            ) 

            num_context_batches = math.ceil(len(context_queue) / context_batch_size) 
            global_context = [] 
            for i in range(num_context_batches): 
                global_context.append(context_queue[i * context_batch_size : (i + 1) * context_batch_size]) 

            for context in global_context:
                # 3.1 分類器フリーガイダンスを行う場合は潜在変数を拡張 
                latent_model_input = torch.cat([latents[:, :, c] for c in context]).repeat(2 if do_classifier_free_guidance else 1, 1, 1, 1, 1) 
                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) 
                b, c, f, h, w = latent_model_input.shape 
                latent_pose_input = torch.cat([pose_fea[:, :, c] for c in context]).repeat(2 if do_classifier_free_guidance else 1, 1, 1, 1, 1) 

                pred = self.denoising_unet( 
                    ( 
                        latent_model_input, 
                        t, 
                        encoder_hidden_states[:b], 
                        latent_pose_input, 
                        *ref_features, 
                    ) 
                )[0] 

                for j, c in enumerate(context): 
                    noise_pred[:, :, c] = noise_pred[:, :, c] + pred 
                    counter[:, :, c] = counter[:, :, c] + 1 

            # ガイダンスを実行 
            if do_classifier_free_guidance: 
                noise_pred_uncond, noise_pred_text = (noise_pred / counter).chunk(2) 
                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) 

            latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample 

            if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): 
                progress_bar.update() 
                if callback is not None and i % callback_steps == 0: 
                    step_idx = i // getattr(self.scheduler, "order", 1) 
                    callback(step_idx, t, latents) 

        if interpolation_factor > 0: 
            latents = self.interpolate_latents(latents, interpolation_factor, latents.device) 
        # 後処理 
        images = self.decode_latents(latents) # (b, c, f, h, w) 

        # テンソルへ変換 
        if output_type == "tensor": 
            images = torch.from_numpy(images) 

        return images

pipe = OVPose2VideoPipeline()

pose_images = read_frames("Moore-AnimateAnyone/configs/inference/pose_videos/anyone-video-2_kps.mp4") 
src_fps = get_fps("Moore-AnimateAnyone/configs/inference/pose_videos/anyone-video-2_kps.mp4") 
ref_image = Image.open("Moore-AnimateAnyone/configs/inference/ref_images/anyone-5.png").convert("RGB") 
pose_list = [] 
for pose_image_pil in pose_images[:VIDEO_LENGTH]: 
    pose_list.append(pose_image_pil)

video = pipe( 
    ref_image, 
    pose_list, 
    width=WIDTH, 
    height=HEIGHT, 
    video_length=VIDEO_LENGTH, 
)