difpoint/inference.py

# -*- coding: UTF-8 -*-
'''
@File    ：inference.py
@Author  ：Chaolong Yang
@Date    ：2024/5/29 19:26 
'''
import glob

import os
os.environ['HYDRA_FULL_ERROR']='1'


import os
import time
import shutil
import uuid
import os
import cv2
import tyro

import numpy as np
from tqdm import tqdm
import cv2
from rich.progress import track

from difpoint.croper import Croper
from PIL import Image
import time


import torch
import torch.nn.functional as F
from torch import nn
import imageio
from pydub import AudioSegment
from pykalman import KalmanFilter
import scipy
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')

from difpoint.dataset_process import audio
import os
import argparse
import pdb
import ffmpeg
import cv2
import time
import numpy as np
import os
import datetime
import platform
from omegaconf import OmegaConf
#from difpoint.src.pipelines.faster_live_portrait_pipeline import FasterLivePortraitPipeline
from difpoint.src.live_portrait_pipeline import LivePortraitPipeline
from difpoint.src.config.argument_config import ArgumentConfig
from difpoint.src.config.inference_config import InferenceConfig
from difpoint.src.config.crop_config import CropConfig
from difpoint.src.live_portrait_pipeline import LivePortraitPipeline
from difpoint.src.utils.retargeting_utils import calc_eye_close_ratio, calc_lip_close_ratio
from difpoint.src.utils.camera import get_rotation_matrix
from difpoint.src.utils.video import images2video, concat_frames, get_fps, add_audio_to_video, has_audio_stream


FFMPEG = "ffmpeg"

def parse_audio_length(audio_length, sr, fps):
    bit_per_frames = sr / fps
    num_frames = int(audio_length / bit_per_frames)
    audio_length = int(num_frames * bit_per_frames)
    return audio_length, num_frames

def crop_pad_audio(wav, audio_length):
    if len(wav) > audio_length:
        wav = wav[:audio_length]
    elif len(wav) < audio_length:
        wav = np.pad(wav, [0, audio_length - len(wav)], mode='constant', constant_values=0)
    return wav

class Conv2d(nn.Module):
    def __init__(self, cin, cout, kernel_size, stride, padding, residual=False, use_act=True, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.conv_block = nn.Sequential(
            nn.Conv2d(cin, cout, kernel_size, stride, padding),
            nn.BatchNorm2d(cout)
        )
        self.act = nn.ReLU()
        self.residual = residual
        self.use_act = use_act

    def forward(self, x):
        out = self.conv_block(x)
        if self.residual:
            out += x

        if self.use_act:
            return self.act(out)
        else:
            return out

class AudioEncoder(nn.Module):
    def __init__(self, wav2lip_checkpoint, device):
        super(AudioEncoder, self).__init__()

        self.audio_encoder = nn.Sequential(
            Conv2d(1, 32, kernel_size=3, stride=1, padding=1),
            Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),
            Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),

            Conv2d(32, 64, kernel_size=3, stride=(3, 1), padding=1),
            Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
            Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),

            Conv2d(64, 128, kernel_size=3, stride=3, padding=1),
            Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
            Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),

            Conv2d(128, 256, kernel_size=3, stride=(3, 2), padding=1),
            Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True),

            Conv2d(256, 512, kernel_size=3, stride=1, padding=0),
            Conv2d(512, 512, kernel_size=1, stride=1, padding=0),)

        #### load the pre-trained audio_encoder
        wav2lip_state_dict = torch.load(wav2lip_checkpoint, map_location=torch.device(device))['state_dict']
        state_dict = self.audio_encoder.state_dict()

        for k,v in wav2lip_state_dict.items():
            if 'audio_encoder' in k:
                state_dict[k.replace('module.audio_encoder.', '')] = v
        self.audio_encoder.load_state_dict(state_dict)

    def forward(self, audio_sequences):
        # audio_sequences = (B, T, 1, 80, 16)
        B = audio_sequences.size(0)

        audio_sequences = torch.cat([audio_sequences[:, i] for i in range(audio_sequences.size(1))], dim=0)

        audio_embedding = self.audio_encoder(audio_sequences) # B, 512, 1, 1
        dim = audio_embedding.shape[1]
        audio_embedding = audio_embedding.reshape((B, -1, dim, 1, 1))

        return audio_embedding.squeeze(-1).squeeze(-1) #B seq_len+1 512

def partial_fields(target_class, kwargs):
    return target_class(**{k: v for k, v in kwargs.items() if hasattr(target_class, k)})

def dct2device(dct: dict, device):
    for key in dct:
        dct[key] = torch.tensor(dct[key]).to(device)
    return dct

def save_video_with_watermark(video, audio, save_path, watermark=False):
    temp_file = str(uuid.uuid4())+'.mp4'
    cmd = r'ffmpeg -y -i "%s" -i "%s" -vcodec copy "%s"' % (video, audio, temp_file)
    os.system(cmd)
    shutil.move(temp_file, save_path)
    

    
class Inferencer(object):
    def __init__(self):
        
        st=time.time()
        print('#'*25+'Start initialization'+'#'*25)
        self.device = 'cuda'
        from difpoint.model import get_model
        self.point_diffusion = get_model()
        ckpt = torch.load('./downloaded_repo/ckpts/KDTalker.pth', weights_only=False)

        self.point_diffusion.load_state_dict(ckpt['model'])
        print('model', self.point_diffusion.children())
        self.point_diffusion.eval()
        self.point_diffusion.to(self.device)

        lm_croper_checkpoint = './downloaded_repo/ckpts/shape_predictor_68_face_landmarks.dat'
        self.croper = Croper(lm_croper_checkpoint)

        self.norm_info = dict(np.load(r'difpoint/datasets/norm_info_d6.5_c8.5_vox1_train.npz'))

        wav2lip_checkpoint = './downloaded_repo/ckpts/wav2lip.pth'
        self.wav2lip_model = AudioEncoder(wav2lip_checkpoint, 'cuda')
        self.wav2lip_model.cuda()
        self.wav2lip_model.eval()

        args = tyro.cli(ArgumentConfig)

        self.inf_cfg = partial_fields(InferenceConfig, args.__dict__)  # use attribute of args to initial InferenceConfig
        self.crop_cfg = partial_fields(CropConfig, args.__dict__)  # use attribute of args to initial CropConfig

        self.live_portrait_pipeline = LivePortraitPipeline(inference_cfg=self.inf_cfg, crop_cfg=self.crop_cfg)
        print('#'*25+f'End initialization, cost time {time.time()-st}'+'#'*25)

    def _norm(self, data_dict):
        for k in data_dict.keys():
            if k in ['yaw', 'pitch', 'roll', 't', 'exp', 'scale', 'kp', 'c_lip', 'c_eye']:
                v=data_dict[k]
                data_dict[k] = (v - self.norm_info[k+'_mean'])/self.norm_info[k+'_std']
        return data_dict

    def _denorm(self, data_dict):
        for k in data_dict.keys():
            if k in ['yaw', 'pitch', 'roll', 't', 'exp', 'scale', 'kp', 'c_lip', 'c_eye']:
                v=data_dict[k]
                data_dict[k] = v * self.norm_info[k+'_std'] + self.norm_info[k+'_mean']
        return data_dict


    def output_to_dict(self, data):
        output = {}

        output['scale'] = data[:, 0]
        output['yaw'] = data[:, 1, None]
        output['pitch'] = data[:, 2, None]
        output['roll'] = data[:, 3, None]
        output['t'] = data[:, 4:7]
        output['exp'] = data[:, 7:]

        return output

    def extract_mel_from_audio(self, audio_file_path):
        syncnet_mel_step_size = 16
        fps = 25
        wav = audio.load_wav(audio_file_path, 16000)
        wav_length, num_frames = parse_audio_length(len(wav), 16000, 25)
        wav = crop_pad_audio(wav, wav_length)
        orig_mel = audio.melspectrogram(wav).T
        spec = orig_mel.copy()
        indiv_mels = []

        for i in tqdm(range(num_frames), 'mel:'):
            start_frame_num = i - 2
            start_idx = int(80. * (start_frame_num / float(fps)))
            end_idx = start_idx + syncnet_mel_step_size
            seq = list(range(start_idx, end_idx))
            seq = [min(max(item, 0), orig_mel.shape[0] - 1) for item in seq]
            m = spec[seq, :]
            indiv_mels.append(m.T)
        indiv_mels = np.asarray(indiv_mels)  # T 80 16
        return indiv_mels

    def extract_wav2lip_from_audio(self, audio_file_path):
        asd_mel = self.extract_mel_from_audio(audio_file_path)
        asd_mel = torch.FloatTensor(asd_mel).cuda().unsqueeze(0).unsqueeze(2)
        with torch.no_grad():
            hidden = self.wav2lip_model(asd_mel)
        return hidden[0].cpu().detach().numpy()

    def headpose_pred_to_degree(self, pred):
        device = pred.device
        idx_tensor = [idx for idx in range(66)]
        idx_tensor = torch.FloatTensor(idx_tensor).to(device)
        pred = F.softmax(pred)
        degree = torch.sum(pred * idx_tensor, 1) * 3 - 99
        return degree

    def calc_combined_eye_ratio(self, c_d_eyes_i, c_s_eyes):
        c_s_eyes_tensor = torch.from_numpy(c_s_eyes).float().to(self.device)
        c_d_eyes_i_tensor = c_d_eyes_i[0].reshape(1, 1).to(self.device)
        # [c_s,eyes, c_d,eyes,i]
        combined_eye_ratio_tensor = torch.cat([c_s_eyes_tensor, c_d_eyes_i_tensor], dim=1)
        return combined_eye_ratio_tensor

    def calc_combined_lip_ratio(self, c_d_lip_i, c_s_lip):
        c_s_lip_tensor = torch.from_numpy(c_s_lip).float().to(self.device)
        c_d_lip_i_tensor = c_d_lip_i[0].to(self.device).reshape(1, 1) # 1x1
        # [c_s,lip, c_d,lip,i]
        combined_lip_ratio_tensor = torch.cat([c_s_lip_tensor, c_d_lip_i_tensor], dim=1) # 1x2
        return combined_lip_ratio_tensor

    # 2024.06.26
    @torch.no_grad()
    def generate_with_audio_img(self, upload_audio_path, tts_audio_path, audio_type, image_path, smoothed_pitch, smoothed_yaw, smoothed_roll, smoothed_t,  save_path='./downloaded_repo/'):
        print(audio_type)
        if audio_type == 'upload':
            audio_path = upload_audio_path
        elif audio_type == 'tts':
            audio_path = tts_audio_path
        save_path = os.path.join(save_path, "output.mp4")
        image = [np.array(Image.open(image_path).convert('RGB'))]
        if image[0].shape[0] != 256 or image[0].shape[1] != 256:
            cropped_image, crop, quad = self.croper.crop(image, still=False, xsize=512)
            input_image = cv2.resize(cropped_image[0], (256, 256))
        else:
            input_image = image[0]

        I_s = torch.FloatTensor(input_image.transpose((2, 0, 1))).unsqueeze(0).cuda() / 255

        x_s_info = self.live_portrait_pipeline.live_portrait_wrapper.get_kp_info(I_s)
        x_c_s = x_s_info['kp'].reshape(1, 21, -1)
        R_s = get_rotation_matrix(x_s_info['pitch'], x_s_info['yaw'], x_s_info['roll'])
        f_s = self.live_portrait_pipeline.live_portrait_wrapper.extract_feature_3d(I_s)
        x_s = self.live_portrait_pipeline.live_portrait_wrapper.transform_keypoint(x_s_info)

        flag_lip_zero = self.inf_cfg.flag_lip_zero  # not overwrite 


        ######## process driving info ########
        kp_info = {}
        for k in x_s_info.keys():
            kp_info[k] = x_s_info[k].cpu().numpy()
        # kp_info['c_lip'] = c_s_lip
        # kp_info['c_eye'] = c_s_eye

        kp_info = self._norm(kp_info)

        ori_kp = torch.cat([torch.zeros([1, 7]), torch.Tensor(kp_info['kp'])], -1).cuda()

        input_x = np.concatenate([kp_info[k] for k in ['scale', 'yaw', 'pitch', 'roll', 't']], 1)
        input_x = np.concatenate((input_x, kp_info['exp'].reshape(1, 63)), axis=1)
        input_x = np.expand_dims(input_x, -1)
        input_x = np.expand_dims(input_x, 0)
        input_x = np.concatenate([input_x, input_x, input_x], -1)

        aud_feat = self.extract_wav2lip_from_audio(audio_path)

        outputs = [input_x]

        st = time.time()
        print('#' * 25 + 'Start Inference' + '#' * 25)
        sample_frame = 64    # 32 aud_feat.shape[0]
        
        for i in range(0, aud_feat.shape[0] - 1, sample_frame):
            input_mel = torch.Tensor(aud_feat[i: i + sample_frame]).unsqueeze(0).cuda()
            kp0 = torch.Tensor(outputs[-1])[:, -1].cuda()
            pred_kp = self.point_diffusion.forward_sample(70, ref_kps=kp0, ori_kps=ori_kp, aud_feat=input_mel,
                                                          scheduler='ddim', num_inference_steps=50)
            outputs.append(pred_kp.cpu().numpy())


        outputs = np.mean(np.concatenate(outputs, 1)[0], -1)[1:, ]
        output_dict = self.output_to_dict(outputs)
        output_dict = self._denorm(output_dict)

        num_frame = output_dict['yaw'].shape[0]
        x_d_info = {}
        for key in output_dict:
            x_d_info[key] = torch.tensor(output_dict[key]).cuda()

        # smooth
        def smooth(sequence, n_dim_state=1):
            kf = KalmanFilter(initial_state_mean=sequence[0],
                              transition_covariance=0.05 * np.eye(n_dim_state),  # 较小的过程噪声
                              observation_covariance=0.001 * np.eye(n_dim_state))  # 可以增大观测噪声，减少敏感性
            state_means, _ = kf.smooth(sequence)
            return state_means

        # scale_data = x_d_info['scale'].cpu().numpy()
        yaw_data = x_d_info['yaw'].cpu().numpy()
        pitch_data = x_d_info['pitch'].cpu().numpy()
        roll_data = x_d_info['roll'].cpu().numpy()
        t_data = x_d_info['t'].cpu().numpy()
        exp_data = x_d_info['exp'].cpu().numpy()

        smoothed_pitch = smooth(pitch_data, n_dim_state=1) * smoothed_pitch
        smoothed_yaw = smooth(yaw_data, n_dim_state=1) * smoothed_yaw
        smoothed_roll = smooth(roll_data, n_dim_state=1) * smoothed_roll
        # smoothed_scale = smooth(scale_data, n_dim_state=1)
        smoothed_t = smooth(t_data, n_dim_state=3) * smoothed_t
        smoothed_exp = smooth(exp_data, n_dim_state=63)

        # x_d_info['scale'] = torch.Tensor(smoothed_scale).cuda()
        x_d_info['pitch'] = torch.Tensor(smoothed_pitch).cuda()
        x_d_info['yaw'] = torch.Tensor(smoothed_yaw).cuda()
        x_d_info['roll'] = torch.Tensor(smoothed_roll).cuda()
        x_d_info['t'] = torch.Tensor(smoothed_t).cuda()
        x_d_info['exp'] = torch.Tensor(smoothed_exp).cuda()



        template_dct = {'motion': [], 'c_d_eyes_lst': [], 'c_d_lip_lst': []}
        for i in track(range(num_frame), description='Making motion templates...', total=num_frame):
            # collect s_d, R_d, δ_d and t_d for inference
            x_d_i_info = x_d_info
            R_d_i = get_rotation_matrix(x_d_i_info['pitch'][i], x_d_i_info['yaw'][i], x_d_i_info['roll'][i])

            item_dct = {
                'scale': x_d_i_info['scale'][i].cpu().numpy().astype(np.float32),
                'R_d': R_d_i.cpu().numpy().astype(np.float32),
                'exp': x_d_i_info['exp'][i].reshape(1, 21, -1).cpu().numpy().astype(np.float32),
                't': x_d_i_info['t'][i].cpu().numpy().astype(np.float32),
            }

            template_dct['motion'].append(item_dct)
            # template_dct['c_d_eyes_lst'].append(x_d_i_info['c_eye'][i])
            # template_dct['c_d_lip_lst'].append(x_d_i_info['c_lip'][i])

        I_p_lst = []
        R_d_0, x_d_0_info = None, None

        for i in track(range(num_frame), description='Animating...', total=num_frame):
            x_d_i_info = template_dct['motion'][i]

            for key in x_d_i_info:
                x_d_i_info[key] = torch.tensor(x_d_i_info[key]).cuda()

            R_d_i = x_d_i_info['R_d']

            if i == 0:
                R_d_0 = R_d_i
                x_d_0_info = x_d_i_info
            

            if self.inf_cfg.flag_relative_motion:
                R_new = (R_d_i @ R_d_0.permute(0, 2, 1)) @ R_s
                delta_new = x_s_info['exp'].reshape(1, 21, -1) + (x_d_i_info['exp'] - x_d_0_info['exp'])
                scale_new = x_s_info['scale'] * (x_d_i_info['scale'] / x_d_0_info['scale'])
                t_new = x_s_info['t'] + (x_d_i_info['t'] - x_d_0_info['t'])
            else:
                R_new = R_d_i
                delta_new = x_d_i_info['exp']
                scale_new = x_s_info['scale']
                t_new = x_d_i_info['t']
            t_new[..., 2] = 0  # zero tz
            
            x_d_i_new = scale_new * (x_c_s @ R_new + delta_new) + t_new
    

            # Algorithm 1:
            if not self.inf_cfg.flag_stitching and not self.inf_cfg.flag_eye_retargeting and not self.inf_cfg.flag_lip_retargeting:
                # without stitching or retargeting
                if flag_lip_zero:
                    x_d_i_new += lip_delta_before_animation.reshape(-1, x_s.shape[1], 3)
                else:
                    pass
            elif self.inf_cfg.flag_stitching and not self.inf_cfg.flag_eye_retargeting and not self.inf_cfg.flag_lip_retargeting:
                # with stitching and without retargeting
                if flag_lip_zero:
                    x_d_i_new = self.live_portrait_pipeline.live_portrait_wrapper.stitching(x_s, x_d_i_new) + lip_delta_before_animation.reshape(-1, x_s.shape[1], 3)
                else:
                    x_d_i_new = self.live_portrait_pipeline.live_portrait_wrapper.stitching(x_s, x_d_i_new)
            else:
                eyes_delta, lip_delta = None, None

                if self.inf_cfg.flag_relative_motion:  # use x_s
                    x_d_i_new = x_s + \
                                (eyes_delta.reshape(-1, x_s.shape[1], 3) if eyes_delta is not None else 0) + \
                                (lip_delta.reshape(-1, x_s.shape[1], 3) if lip_delta is not None else 0)
                else:  # use x_d,i
                    x_d_i_new = x_d_i_new + \
                                (eyes_delta.reshape(-1, x_s.shape[1], 3) if eyes_delta is not None else 0) + \
                                (lip_delta.reshape(-1, x_s.shape[1], 3) if lip_delta is not None else 0)

                if self.inf_cfg.flag_stitching:
                    x_d_i_new = self.live_portrait_pipeline.live_portrait_wrapper.stitching(x_s, x_d_i_new)


            out = self.live_portrait_pipeline.live_portrait_wrapper.warp_decode(f_s, x_s, x_d_i_new)
            I_p_i = self.live_portrait_pipeline.live_portrait_wrapper.parse_output(out['out'])[0]
            I_p_lst.append(I_p_i)

        video_name = os.path.basename(save_path)
        video_save_dir = os.path.dirname(save_path)
        path = os.path.join(video_save_dir, video_name)
                  
        imageio.mimsave(path, I_p_lst, fps=float(25))

        audio_name = audio_path.split('/')[-1]
        new_audio_path = os.path.join(video_save_dir, audio_name)
        start_time = 0
        # cog will not keep the .mp3 filename
        sound = AudioSegment.from_file(audio_path)
        end_time = start_time + num_frame * 1 / 25 * 1000
        word1 = sound.set_frame_rate(16000)
        word = word1[start_time:end_time]
        word.export(new_audio_path, format="wav")

        save_video_with_watermark(path, new_audio_path, save_path, watermark=False)
        print(f'The generated video is named {video_save_dir}/{video_name}')

        print('#' * 25 + f'End Inference, cost time {time.time() - st}' + '#' * 25)
        return save_path