inference.py

import argparse
import os
import warnings
import torch
from diffusers import AutoencoderKL, DDIMScheduler
from omegaconf import OmegaConf
from torch import nn
import cv2
import numpy as np
from typing import List
from einops import rearrange
import librosa
import soundfile as sf
import torchaudio
import imageio
import subprocess
from denoiser import pretrained
from denoiser.dsp import convert_audio

from src.pipelines.face_animate import FaceAnimatePipeline
from src.models.unet_2d_condition import UNet2DConditionModel
from src.models.unet_3d_echo import EchoUNet3DConditionModel
from src.models.whisper.audio2feature import load_audio_model
from image_processor import ImageProcessor
from src.utils.util import save_video, filter_non_none


warnings.filterwarnings("ignore")

class Net(nn.Module):
    """
    The Net class combines all the necessary modules for the inference process.
    
    Args:
        reference_unet (UNet2DConditionModel): The UNet2DConditionModel used as a reference for inference.
        denoising_unet (UNet3DConditionModel): The UNet3DConditionModel used for denoising the input audio.
    """
    def __init__(
        self,
        reference_unet: UNet2DConditionModel,
        denoising_unet: EchoUNet3DConditionModel,
    ):
        super().__init__()
        self.reference_unet = reference_unet
        self.denoising_unet = denoising_unet

    def forward(self,):
        """
        empty function to override abstract function of nn Module
        """

    def get_modules(self):
        """
        Simple method to avoid too-few-public-methods pylint error
        """
        return {
            "reference_unet": self.reference_unet,
            "denoising_unet": self.denoising_unet,
        }


def create_mask(mask_path, mask_512_path, clip_length):
        mask = 1 - (cv2.imread(mask_path) / 255.)
        (h, w, c) = mask.shape
        new_mask = np.zeros((h, w, c+1))
        new_mask[:, :, :c] = mask
        new_mask[:, :, c] = mask[:, :, c-1]

        mask = torch.FloatTensor(new_mask).permute(2,0,1).unsqueeze(1)

        return mask


def create_mask_alpha(w, h, coordinate):
    (left, top, right, bottom) = coordinate

    mask = np.ones((h,w,3))
    mask[top:bottom, left:right] = 0
    
    pad = min(int((right-left)/15), int((bottom-top)/15)) # 100/15=6.6% for padding
    
    for idx in range(left, right):
        mask[bottom-pad:bottom, idx] = np.array([np.linspace(0,1,pad)]*3).transpose(1,0)
        mask[top:top+pad, idx] = np.array([np.linspace(1,0,pad)]*3).transpose(1,0)
    
    for idx in range(top, bottom):
        mask[idx, right-pad:right] = np.array([np.linspace(0,1,pad)]*3).transpose(1,0)
        mask[idx, left:left+pad] = np.array([np.linspace(1,0,pad)]*3).transpose(1,0)
    
    for idx, start, end in zip(range(top, top+pad), np.linspace(1,0.7, pad), np.linspace(0.7,0, pad)):
        mask[idx, left:left+pad] = np.array([np.linspace(start,end,pad)]*3).transpose(1,0)
        mask[idx, right-pad:right] = np.array([np.linspace(end,start,pad)]*3).transpose(1,0)
    
    for idx, start, end in zip(range(bottom-pad, bottom), np.linspace(0.7,1, pad), np.linspace(0,0.7, pad)):
        mask[idx, left:left+pad] = np.array([np.linspace(start,end,pad)]*3).transpose(1,0)
        mask[idx, right-pad:right] = np.array([np.linspace(end,start,pad)]*3).transpose(1,0)
    
    return mask

def process_audio(input_audio_path, output_audio_path):  
    wav,_ = librosa.load(input_audio_path, sr=16000)
    tmp_path = output_audio_path[:-4]+"_16000.wav"
    sf.write(tmp_path, wav, 16000)

    model = pretrained.dns64().cuda()
    wav, sr = torchaudio.load(input_audio_path)
    wav = convert_audio(wav.cuda(), sr, model.sample_rate, model.chin)

    with torch.no_grad():
        denoised = model(wav[None])[0]
        
    sf.write(output_audio_path, denoised.data.cpu().numpy()[0], 16000)

    return

img_size = (512,512)
    

def mem_gpu():
    # Total memory
    total = torch.cuda.get_device_properties(0).total_memory
    # Allocated memory
    allocated = torch.cuda.memory_allocated(0)
    # Cached memory (used for optimization)
    cached = torch.cuda.memory_reserved(0)

    print(f"Total: {total / 1e9:.2f} GB")
    print(f"Allocated: {allocated / 1e9:.2f} GB")
    print(f"Cached: {cached / 1e9:.2f} GB")


def inference_process(args: argparse.Namespace):
    """
    Perform inference processing.

    Args:
        args (argparse.Namespace): Command-line arguments.

    This function initializes the configuration for the inference process. It sets up the necessary
    modules and variables to prepare for the upcoming inference steps.
    """
    # 1. init config
    cli_args = filter_non_none(vars(args))
    config = OmegaConf.load(args.config)
    config = OmegaConf.merge(config, cli_args)
    save_path = config.save_path
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    # 2. runtime variables
    device = torch.device(
        "cuda") if torch.cuda.is_available() else torch.device("cpu")
    if config.weight_dtype == "fp16":
        weight_dtype = torch.float16
    elif config.weight_dtype == "bf16":
        weight_dtype = torch.bfloat16
    elif config.weight_dtype == "fp32":
        weight_dtype = torch.float32
    else:
        weight_dtype = torch.float32

    # 3. prepare inference data
    # 3.1 prepare source image, face mask
    img_size = (config.data.source_image.width,
                config.data.source_image.height)
    clip_length = config.data.n_sample_frames
    mask_path = config.mask_path
    mask_512_path = config.mask_512_path

    temp_dir = os.path.join(save_path, "img")
    os.makedirs(temp_dir, exist_ok=True)
    os.system(f"rm -r {temp_dir}/*")

    # 3.2 prepare audio embeddings
    sample_rate = config.data.driving_audio.sample_rate
    assert sample_rate == 16000, "audio sample rate must be 16000"
    fps = config.data.export_video.fps

    ### load audio processor params
    audio_processor = load_audio_model(model_path=config.audio_model_path, device=device)

    audio_temp_path = os.path.join(save_path, "audio.wav")
    process_audio(config.driving_audio, audio_temp_path)

    # 4. build modules
    sched_kwargs = OmegaConf.to_container(config.noise_scheduler_kwargs)
    if config.enable_zero_snr:
        sched_kwargs.update(
            rescale_betas_zero_snr=True,
            timestep_spacing="trailing",
            prediction_type="v_prediction",
        )
    val_noise_scheduler = DDIMScheduler(**sched_kwargs)
    sched_kwargs.update({"beta_schedule": "scaled_linear"})

    vae = AutoencoderKL.from_pretrained(config.vae.model_path)
    reference_unet = UNet2DConditionModel.from_pretrained(
        config.base_model_path, subfolder="unet")
    denoising_unet = EchoUNet3DConditionModel.from_pretrained_2d(
        config.base_model_path,
        "",
        subfolder="unet",
        unet_additional_kwargs=OmegaConf.to_container(
            config.unet_additional_kwargs)
    )

    reference_unet.load_state_dict(
        torch.load(config.reference_unet_path, map_location="cpu"),
    )

    denoising_unet.load_state_dict(
        torch.load(config.denoising_unet_path, map_location="cpu"),
    )

    # Freeze
    vae.requires_grad_(False)
    reference_unet.requires_grad_(False)
    denoising_unet.requires_grad_(False)

    # reference_unet.enable_gradient_checkpointing()
    # denoising_unet.enable_gradient_checkpointing()

    vae.eval()
    reference_unet.eval()
    denoising_unet.eval()

    net = Net(
        reference_unet,
        denoising_unet
    )


    # 5. inference
    pipeline = FaceAnimatePipeline(
        vae=vae,
        reference_unet=net.reference_unet,
        denoising_unet=net.denoising_unet,
        scheduler=val_noise_scheduler
    )


    pipeline.to(device=device, dtype=weight_dtype)

    image_processor = ImageProcessor(img_size)

    generator = torch.manual_seed(42)

    video_stream = cv2.VideoCapture(args.source_video)
    
    temp_video_path = config.output[:4] + "_no_audio.mp4"
    writer = imageio.get_writer(
    temp_video_path,
    fps=25,
    codec='libx264',
    quality=10,  # 0 (low) to 10 (high)
    ffmpeg_params=[
        '-pix_fmt', 'yuv420p',
        '-crf', '18'  # Lower CRF = better quality (17–23 is visually lossless range)
    ])
    
    whisper_feature = audio_processor.audio2feat(audio_temp_path)

    whisper_chunks = audio_processor.feature2chunks(feature_array=whisper_feature, fps=fps)
    # torch.save(whisper_chunks, "./cache/audio.pt")
    audio_frame_num = whisper_chunks.shape[0]
    audio_fea_final = torch.Tensor(whisper_chunks).to(dtype=vae.dtype, device=vae.device)
    # audio_fea_final = torch.tensor(torch.load("/workspace/preprocess/audio_emb_whisper/cM096XGOCG4_C1.pt")).to(dtype=vae.dtype, device=vae.device)
    audio_fea_final = audio_fea_final.unsqueeze(0)

    video_length = min(int(video_stream.get(cv2.CAP_PROP_FRAME_COUNT)), audio_frame_num)
    if video_length < audio_frame_num:
        audio_fea_final = audio_fea_final[:, :video_length, :, :]

    times = video_length // config.inference_frames_num

    frames_video = []
    count_motion = 0
    count_img = 0
    count = 0
    finish = False
    
    while 1:

        still_reading, frame = video_stream.read()

        if (not still_reading) or (count_motion >= video_length):
            video_stream.release()
            finish = True
        else:
            frames_video.append(frame)
            count_motion += 1

        if len(frames_video)==0:
            break
        
        if ((count_motion % config.inference_frames_num) == 0) or finish:

            print(f"processing {count}/{times}")
            source_image_pixels_all, bbox_all, frame_list_all = image_processor.preprocess_frames_jump(frames_video)            

            del frames_video
            
            for source_image_pixels_list, bbox_list, frame_list in zip(source_image_pixels_all, bbox_all, frame_list_all):

                print(f"generating {count}/{times}")

                if bbox_list[0]!=(0,0,0,0):
                    
                    audio_tensor = audio_fea_final[:, count_img : count_img + len(frame_list)]

                    main_length = len(frame_list)
                    diff = 12 - main_length 

                    if diff > 0:
                        source_image_pixels_list.extend([source_image_pixels_list[-1]]*diff)
                        audio_tensor = torch.cat((audio_tensor, audio_tensor[:,-1:].repeat(1,diff,1,1)), dim=1)

                    source_image_pixels = torch.cat(source_image_pixels_list, dim=0)
                    mask_main = create_mask(mask_path, mask_512_path, len(frame_list))
                    mask_main = mask_main.unsqueeze(0).to(device="cuda", dtype=weight_dtype)
                    torch.cuda.empty_cache()

                    with torch.no_grad():
                        pipeline_output = pipeline(
                            gt_encode=source_image_pixels,
                            audio_tensor=audio_tensor,
                            mask_main=mask_main,
                            width=img_size[0],
                            height=img_size[1],
                            video_length=len(source_image_pixels_list),
                            num_inference_steps=config.inference_steps,
                            guidance_scale=config.cfg_scale,
                            generator=generator
                        )

                    torch.cuda.empty_cache()

                    images_out = pipeline_output.videos.squeeze(0).permute(1, 2, 3, 0).cpu().numpy()[:main_length]  # convert to [f, h, w, c]
                    images_out = np.clip(images_out * 255, 0, 255).astype(np.uint8)  # to [0, 255]

                    for image_out, bbox, frame in zip(images_out, bbox_list, frame_list):
                        
                        frame_h, frame_w, _ = frame.shape
                        path = os.path.join(temp_dir, f"image_{count_img}.png")
                        
                        if bbox==(0,0,0,0):
                            writer.append_data(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
                            # cv2.imwrite(path, frame)
                        else:
                            image_out = cv2.cvtColor(image_out, cv2.COLOR_RGB2BGR)
                            frame_in_out = frame.copy()
                            frame_in_out[bbox[1]:bbox[3], bbox[0]:bbox[2]] = cv2.resize(image_out, (bbox[2] - bbox[0], bbox[3] - bbox[1]))
                            mask = create_mask_alpha(frame_w, frame_h, bbox)
                            blurred_image = (frame * mask + frame_in_out * (1-mask)).astype(np.uint8)
                            writer.append_data(cv2.cvtColor(blurred_image, cv2.COLOR_RGB2BGR))    
                            # cv2.imwrite(path, blurred_image)
                        
                        count_img += 1

                else:
                    for frame in frame_list:
                        path = os.path.join(temp_dir, f"image_{count_img}.png")
                        writer.append_data(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
                        # cv2.imwrite(path, frame)
                        count_img += 1

            count += 1

            frames_video = []

        if finish:
            break

    writer.close()

    command = f"ffmpeg -i {temp_video_path} -i {config.driving_audio} -map 0:v -map 1:a -c:v h264 -shortest -y {config.output} -loglevel quiet"
    subprocess.run(command, shell=True, check=True, text=True)

    # output_file = config.output
    # frames_path = os.path.join(temp_dir, f"image_%d.png")
    # save_video(frames_path, config.driving_audio, output_file, frame_w, frame_h)

    return 


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument("-c", "--config", default="configs/inference/default.yaml")
    parser.add_argument("--source_video", type=str, required=False, help="source image")
    parser.add_argument("--driving_audio", type=str, required=False, help="driving audio")
    parser.add_argument("--output", type=str, help="output video file name", default="./cache/output.mp4")
    parser.add_argument("--inference_frames_num", type=int, help="inference frames number", default=1500)


    command_line_args = parser.parse_args()

    inference_process(command_line_args)