infer.py

# !/usr/bin/env python
# -*- coding: UTF-8 -*-
import warnings
import os
import numpy as np
import torch
import torch.utils.checkpoint
from PIL import Image
from diffusers import AutoencoderKLTemporalDecoder
from diffusers.schedulers import EulerDiscreteScheduler
from transformers import CLIPVisionModelWithProjection
import torchvision.transforms as transforms
import torch.nn.functional as F
from torchvision.utils import save_image
import random
import cv2
# pipeline 
from .src.pipelines.pipeline import LQ2VideoLongSVDPipeline
from .src.utils.util import save_videos_grid, seed_everything
from .src.models.id_proj import IDProjConvModel
from .src.models import model_insightface_360k
from .src.dataset.face_align.align import AlignImage
from .src.models.svfr_adapter.unet_3d_svd_condition_ip import UNet3DConditionSVDModel
from .src.dataset.dataset import get_affine_transform, mean_face_lm5p_256

warnings.filterwarnings("ignore")


def main_loader(weight_dtype, repo, unet_path, det_path, id_path, face_path,device):
    vae = AutoencoderKLTemporalDecoder.from_pretrained(repo, subfolder="vae", variant="fp16")
    val_noise_scheduler = EulerDiscreteScheduler.from_pretrained(repo, subfolder="scheduler")
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(repo, subfolder="image_encoder", variant="fp16")
    unet = UNet3DConditionSVDModel.from_pretrained(repo, subfolder="unet", variant="fp16")
    
    align_instance = AlignImage(device, det_path=det_path)
    
    import torch.nn as nn
    class InflatedConv3d(nn.Conv2d):
        def forward(self, x):
            x = super().forward(x)
            return x
    
    # Add ref channel
    old_weights = unet.conv_in.weight
    old_bias = unet.conv_in.bias
    new_conv1 = InflatedConv3d(
        12,
        old_weights.shape[0],
        kernel_size=unet.conv_in.kernel_size,
        stride=unet.conv_in.stride,
        padding=unet.conv_in.padding,
        bias=True if old_bias is not None else False,
    )
    param = torch.zeros((320, 4, 3, 3), requires_grad=True)
    new_conv1.weight = torch.nn.Parameter(torch.cat((old_weights, param), dim=1))
    if old_bias is not None:
        new_conv1.bias = old_bias
    unet.conv_in = new_conv1
    unet.config["in_channels"] = 12
    unet.config.in_channels = 12
    
    id_linear = IDProjConvModel(in_channels=512, out_channels=1024).to(device=device)
    
    # load pretrained weights
    unet.load_state_dict(
        torch.load(unet_path, map_location="cpu"),
        strict=True,
    )
    
    id_linear.load_state_dict(
        torch.load(id_path, map_location="cpu"),
        strict=True,
    )
    
    net_arcface = model_insightface_360k.getarcface(face_path).eval().to(device=device)
    
    image_encoder.to(weight_dtype)
    vae.to(weight_dtype)
    unet.to(weight_dtype)
    id_linear.to(weight_dtype)
    net_arcface.requires_grad_(False).to(weight_dtype)
    
    pipe = LQ2VideoLongSVDPipeline(
        unet=unet,
        image_encoder=image_encoder,
        vae=vae,
        scheduler=val_noise_scheduler,
        feature_extractor=None
    
    )
    pipe = pipe.to(device, dtype=unet.dtype)
    
    return pipe, id_linear, net_arcface, align_instance


def main_sampler(pipe, align_instance, net_arcface, id_linear, save_dir, weight_dtype, seed, input_frames_pil, task_ids,
                 mask_array,
                 save_video, decode_chunk_size, noise_aug_strength, min_appearance_guidance_scale,
                 max_appearance_guidance_scale,
                 overlap, i2i_noise_strength, steps, n_sample_frames,device):
    to_tensor = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
    
    seed_everything(seed)
    
    if 2 in task_ids and isinstance(mask_array, np.ndarray):
        white_positions = mask_array == 255
    
    print('task_ids:', task_ids)
    task_prompt = [0, 0, 0]
    for i in range(3):
        if i in task_ids:
            task_prompt[i] = 1
    print("task_prompt:", task_prompt)
    
    files_prefix = ''.join(random.choice("0123456789") for _ in range(5))
    video_name = f"infer_{files_prefix}"
    # print(video_name)
    
    if os.path.exists(os.path.join(save_dir, "result_frames", video_name[:-4])):
        print(os.path.join(save_dir, "result_frames", video_name[:-4]))
        # continue
    
    # cap = decord.VideoReader(input_frames_pil, fault_tol=1)
    
    total_frames = len(input_frames_pil)
    T = total_frames  #
    print("total_frames:", total_frames)
    step = 1
    drive_idx_start = 0
    drive_idx_list = list(range(drive_idx_start, drive_idx_start + T * step, step))
    assert len(drive_idx_list) == T
    
    imSameIDs = []
    vid_gt = []
    for i, drive_idx in enumerate(drive_idx_list):
        imSameID = input_frames_pil[drive_idx]
        image_array = np.array(imSameID)
        if 2 in task_ids and isinstance(mask_array, np.ndarray):
            image_array[white_positions] = [255, 255, 255]  # mask for inpainting task
        vid_gt.append(np.float32(image_array / 255.))
        imSameIDs.append(imSameID)
    
    vid_lq = [(torch.from_numpy(frame).permute(2, 0, 1) - 0.5) / 0.5 for frame in vid_gt]  # torch.Size([3, 512, 512])
    
    val_data = dict(
        pixel_values_vid_lq=torch.stack(vid_lq, dim=0),
        # pixel_values_ref_img=self.to_tensor(target_image),
        # pixel_values_ref_concat_img=self.to_tensor(imSrc2),
        task_ids=task_ids,
        task_id_input=torch.tensor(task_prompt),
        total_frames=total_frames,
    )
    
    window_overlap = 0
    inter_frame_list = get_overlap_slide_window_indices(val_data["total_frames"], n_sample_frames, window_overlap)
    
    lq_frames = val_data["pixel_values_vid_lq"]
    task_ids = val_data["task_ids"]
    task_id_input = val_data["task_id_input"]
    height, width = val_data["pixel_values_vid_lq"].shape[-2:]
    
    print("Generating the first clip...")
    output = pipe(
        lq_frames[inter_frame_list[0]].to(device).to(weight_dtype),  # lq
        None,  # ref concat
        torch.zeros((1, len(inter_frame_list[0]), 49, 1024)).to(device).to(weight_dtype),  # encoder_hidden_states
        task_id_input.to(device).to(weight_dtype),
        height=height,
        width=width,
        num_frames=len(inter_frame_list[0]),
        decode_chunk_size=decode_chunk_size,
        noise_aug_strength=noise_aug_strength,
        min_guidance_scale=min_appearance_guidance_scale,
        max_guidance_scale=max_appearance_guidance_scale,
        overlap=overlap,
        frames_per_batch=len(inter_frame_list[0]),
        num_inference_steps=steps,
        i2i_noise_strength=i2i_noise_strength,
    )
    video = output.frames
    
    ref_img_tensor = video[0][:, -1]
    ref_img = (video[0][:, -1] * 0.5 + 0.5).clamp(0, 1) * 255.
    ref_img = ref_img.permute(1, 2, 0).cpu().numpy().astype(np.uint8)
    
    pts5 = align_instance(ref_img[:, :, [2, 1, 0]], maxface=True)[0][0]
    
    warp_mat = get_affine_transform(pts5, mean_face_lm5p_256 * height / 256)
    ref_img = cv2.warpAffine(np.array(Image.fromarray(ref_img)), warp_mat, (height, width), flags=cv2.INTER_CUBIC)
    ref_img = to_tensor(ref_img).to(device).to(weight_dtype)
    
    save_image(ref_img * 0.5 + 0.5, f"{save_dir}/ref_img_align.png")
    
    ref_img = F.interpolate(ref_img.unsqueeze(0)[:, :, 0:224, 16:240], size=[112, 112], mode='bilinear')
    _, id_feature_conv = net_arcface(ref_img)
    id_embedding = id_linear(id_feature_conv)
    
    print('Generating all video clips...')
    video = pipe(
        lq_frames.to(device).to(weight_dtype),  # lq
        ref_img_tensor.to(device).to(weight_dtype),
        id_embedding.unsqueeze(1).repeat(1, len(lq_frames), 1, 1).to("cuda").to(weight_dtype),  # encoder_hidden_states
        task_id_input.to(device).to(weight_dtype),
        height=height,
        width=width,
        num_frames=val_data["total_frames"],  # frame_num,
        decode_chunk_size=decode_chunk_size,
        noise_aug_strength=noise_aug_strength,
        min_guidance_scale=min_appearance_guidance_scale,
        max_guidance_scale=max_appearance_guidance_scale,
        overlap=overlap,
        frames_per_batch=n_sample_frames,
        num_inference_steps=steps,
        i2i_noise_strength=i2i_noise_strength,
    ).frames
    
    video = (video * 0.5 + 0.5).clamp(0, 1)
    video = torch.cat([video.to(device=device)], dim=0).cpu()  # torch.Size([1, 3, 160, 512, 512])
    
    if save_video:
        save_videos_grid(video, f"{save_dir}/{video_name[:-4]}_{seed}.mp4", n_rows=1, fps=25)
    
    # if restore_frames:
    #     video = video.squeeze(0)
    #     os.makedirs(os.path.join(save_dir, "result_frames", f"{video_name[:-4]}_{seed}"),exist_ok=True)
    #     print(os.path.join(save_dir, "result_frames", video_name[:-4]))
    #     for i in range(video.shape[1]):
    #         save_frames_path = os.path.join(f"{save_dir}/result_frames", f"{video_name[:-4]}_{seed}", f'{i:08d}.png')
    #         save_image(video[:,i], save_frames_path)
    
    return video.squeeze(0).permute(1, 2, 3, 0)  # bcthw to B,H,W,C


def get_overlap_slide_window_indices(video_length, window_size, window_overlap):
    inter_frame_list = []
    for j in range(0, video_length, window_size - window_overlap):
        inter_frame_list.append([e % video_length for e in range(j, min(j + window_size, video_length))])
    
    return inter_frame_list