diff --git a/__init__.py b/__init__.py index 13b1fc5..7acc01d 100644 --- a/__init__.py +++ b/__init__.py @@ -6,7 +6,7 @@ import cv2 from tqdm import tqdm import re import torch -from .nodes import PoseGenVideo, RefImagePath, GenerateRefPose, Audio2Video, AudioPath +from .nodes import PoseGenVideo, RefImagePath, Audio2Video, AudioPath #,GenerateRefPose from .src.utils.util import get_fps, read_frames, save_videos_from_pil, calculate_file_hash, get_sorted_dir_files_from_directory, get_audio, lazy_eval, hash_path, validate_path import numpy as np @@ -100,7 +100,7 @@ NODE_CLASS_MAPPINGS = { "AniPortrait_LoadVideoPath": LoadVideoPath, "AniPortrait_Pose_Gen_Video": PoseGenVideo, "AniPortrait_Ref_Image_Path": RefImagePath, - "AniPortrait_Generate_Ref_Pose": GenerateRefPose, +# "AniPortrait_Generate_Ref_Pose": GenerateRefPose, "AniPortrait_Audio2Video": Audio2Video, "AniPortrait_Audio_Path": AudioPath, } @@ -110,7 +110,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "AniPortrait_LoadVideoPath": "Load Video (Path) šŸŽ„AniPortrait", "AniPortrait_Pose_Gen_Video": "Pose Generate Video šŸŽ„AniPortrait", "AniPortrait_Ref_Image_Path": "Ref Image Path šŸŽ„AniPortrait", - "AniPortrait_Generate_Ref_Pose": "Generate Ref Pose šŸŽ„AniPortrait", +# "AniPortrait_Generate_Ref_Pose": "Generate Ref Pose šŸŽ„AniPortrait", "AniPortrait_Audio2Video": "Audio Gen Video šŸŽ„AniPortrait", "AniPortrait_Audio_Path": "Audio Path šŸŽ„AniPortrait", } diff --git a/configs/prompts/animation_audio.yaml b/configs/prompts/animation_audio.yaml index e280f6b..a949ed2 100644 --- a/configs/prompts/animation_audio.yaml +++ b/configs/prompts/animation_audio.yaml @@ -13,7 +13,6 @@ weight_dtype: 'fp16' #pose_temp: "./configs/inference/head_pose_temp/pose_temp.npy" - test_cases: # "./configs/inference/ref_images/lyl.png": # - "./configs/inference/audio/lyl.wav" diff --git a/nodes.py b/nodes.py index 28ef4cc..920d60f 100644 --- a/nodes.py +++ b/nodes.py @@ -12,6 +12,7 @@ import numpy as np import cv2 import torch import torchvision +import random from tqdm import tqdm from diffusers import AutoencoderKL, DDIMScheduler from diffusers.pipelines.stable_diffusion import StableDiffusionPipeline @@ -29,10 +30,11 @@ from .src.pipelines.pipeline_pose2vid_long import Pose2VideoPipeline from .src.utils.util import get_fps, read_frames, save_videos_grid, calculate_file_hash, get_sorted_dir_files_from_directory, get_audio, lazy_eval, hash_path, validate_path from .src.utils.frame_interpolation import init_frame_interpolation_model, batch_images_interpolation_tool from .src.audio_models.model import Audio2MeshModel +from .src.audio_models.pose_model import Audio2PoseModel from .src.utils.audio_util import prepare_audio_feature from .src.utils.mp_utils import LMKExtractor from .src.utils.draw_util import FaceMeshVisualizer -from .src.utils.pose_util import project_points, project_points_with_trans +from .src.utils.pose_util import project_points, project_points_with_trans, matrix_to_euler_and_translation, euler_and_translation_to_matrix, smooth_pose_seq from scipy.spatial.transform import Rotation as R from scipy.interpolate import interp1d @@ -338,6 +340,7 @@ class AudioPath: def load_reference_audio(audio_path: str): return (audio_path) # if return only one node,has to add comma +''' class GenerateRefPose: @classmethod def INPUT_TYPES(s): @@ -418,7 +421,7 @@ def smooth_pose_seq(pose_seq, window_size=5): smoothed_pose_seq[i] = np.mean(pose_seq[start:end], axis=0) return smoothed_pose_seq - +''' class Audio2Video: @classmethod @@ -435,6 +438,7 @@ class Audio2Video: "model": ([audio_config.pretrained_base_model_path],), "weight_dtype": (["fp16", "fp32"],), "accelerate": ("BOOLEAN", {"default": True}), + "length": ("INT", {"default": 0, "min":0, "max": 0xffffffffffffffff, "step": 1}), "fi_step": ("INT", {"default": 3}), "motion_module_path": ([audio_config.motion_module_path],), "image_encoder_path": ([audio_config.image_encoder_path],), @@ -445,7 +449,7 @@ class Audio2Video: "optional": { "images": ("IMAGE", ), "audio_path": ("Audio_Path",), - "ref_pose_path": ("FILENAMES", ), + #"ref_pose_path": ("FILENAMES", ), "fps": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}), }, } @@ -456,8 +460,8 @@ class Audio2Video: CATEGORY = "AniPortrait šŸŽ„Video" FUNCTION = "audio_2_video" - def audio_2_video(self, ref_image, height, width, seed, cfg, steps, vae_path, model, weight_dtype, accelerate, fi_step, motion_module_path, image_encoder_path, denoising_unet_path, reference_unet_path, pose_guider_path, fps=0, images=None, ref_pose_path=None, audio_path=None): - if ref_pose_path: + def audio_2_video(self, ref_image, height, width, seed, cfg, steps, vae_path, model, weight_dtype, accelerate, length, fi_step, motion_module_path, image_encoder_path, denoising_unet_path, reference_unet_path, pose_guider_path, fps=0, images=None, audio_path=None): + if audio_path: if weight_dtype == "fp16": weight_dtype = torch.float16 else: @@ -480,6 +484,10 @@ class Audio2Video: a2m_model.load_state_dict(torch.load(os.path.join(os.path.dirname(os.path.abspath(__file__)), audio_infer_config['pretrained_model']['a2m_ckpt'])), strict=False) a2m_model.cuda().eval() + a2p_model = Audio2PoseModel(audio_infer_config['a2p_model']) + a2p_model.load_state_dict(torch.load(os.path.join(os.path.dirname(os.path.abspath(__file__)),audio_infer_config['pretrained_model']['a2p_ckpt'])), strict=False) + a2p_model.cuda().eval() + vae = AutoencoderKL.from_pretrained(vae_path,).to(device, dtype=weight_dtype) reference_unet = UNet2DConditionModel.from_pretrained(model,subfolder="unet",).to(dtype=weight_dtype, device=device) @@ -536,13 +544,42 @@ class Audio2Video: pred = pred.squeeze().detach().cpu().numpy() pred = pred.reshape(pred.shape[0], -1, 3) pred = pred + face_result['lmks3d'] - - pose_seq = np.load(ref_pose_path) - mirrored_pose_seq = np.concatenate((pose_seq, pose_seq[-2:0:-1]), axis=0) - cycled_pose_seq = np.tile(mirrored_pose_seq, (sample['seq_len'] // len(mirrored_pose_seq) + 1, 1))[:sample['seq_len']] + + if 'pose_temp' in audio_config and audio_config['pose_temp'] is not None: + pose_seq = np.load(audio_config['pose_temp']) + mirrored_pose_seq = np.concatenate((pose_seq, pose_seq[-2:0:-1]), axis=0) + pose_seq = np.tile(mirrored_pose_seq, (sample['seq_len'] // len(mirrored_pose_seq) + 1, 1))[:sample['seq_len']] + else: + id_seed = random.randint(0, 99) + id_seed = torch.LongTensor([id_seed]).cuda() + # Currently, only inference up to a maximum length of 10 seconds is supported. + chunk_duration = 5 # 5 seconds + sr = 16000 + fps = 30 + chunk_size = sr * chunk_duration + + audio_chunks = list(sample['audio_feature'].split(chunk_size, dim=1)) + seq_len_list = [chunk_duration*fps] * (len(audio_chunks) - 1) + [sample['seq_len'] % (chunk_duration*fps)] # 30 fps + + audio_chunks[-2] = torch.cat((audio_chunks[-2], audio_chunks[-1]), dim=1) + seq_len_list[-2] = seq_len_list[-2] + seq_len_list[-1] + del audio_chunks[-1] + del seq_len_list[-1] + + pose_seq = [] + for audio, seq_len in zip(audio_chunks, seq_len_list): + pose_seq_chunk = a2p_model.infer(audio, seq_len, id_seed) + pose_seq_chunk = pose_seq_chunk.squeeze().detach().cpu().numpy() + pose_seq_chunk[:, :3] *= 0.5 + pose_seq.append(pose_seq_chunk) + + pose_seq = np.concatenate(pose_seq, 0) + pose_seq = smooth_pose_seq(pose_seq, 7) + # project 3D mesh to 2D landmark - projected_vertices = project_points(pred, face_result['trans_mat'], cycled_pose_seq, [height, width]) + #projected_vertices = project_points(pred, face_result['trans_mat'], cycled_pose_seq, [height, width]) + projected_vertices = project_points(pred, face_result['trans_mat'], pose_seq, [height, width]) pose_images = [] for i, verts in enumerate(projected_vertices): @@ -550,10 +587,10 @@ class Audio2Video: pose_images.append(lmk_img) pose_list = [] + ''' pose_tensor_list = [] print(f"pose video has {len(pose_images)} frames") pose_transform = transforms.Compose([transforms.Resize((height, width)), transforms.ToTensor()]) - frame_length = len(pose_images) for pose_image_np in pose_images[: frame_length]: pose_image_pil = Image.fromarray(cv2.cvtColor(pose_image_np, cv2.COLOR_BGR2RGB)) @@ -562,15 +599,21 @@ class Audio2Video: for pose_image_np in pose_images[: frame_length: sub_step]: pose_image_np = cv2.resize(pose_image_np, (width, height)) pose_list.append(pose_image_np) - + ''' + frame_length = len(pose_images) if length==0 or length > len(pose_images) else length + sub_step = fi_step if accelerate else 1 + for pose_image_np in pose_images[: frame_length: sub_step]: + pose_image_np = cv2.resize(pose_image_np, (width, height)) + pose_list.append(pose_image_np) + pose_list = np.array(pose_list) video_length = len(pose_list) - + ''' pose_tensor = torch.stack(pose_tensor_list, dim=0) # (f, c, h, w) pose_tensor = pose_tensor.transpose(0, 1) pose_tensor = pose_tensor.unsqueeze(0) - + ''' video = pipe( Image.fromarray(ref_image_pil), pose_list, @@ -689,13 +732,16 @@ class Audio2Video: pose_trans_list = [] verts_list = [] bs_list = [] + ''' src_tensor_list = [] frame_length = len(images) for src_image_pil in images[: frame_length: step]: src_image_pil = (src_image_pil.numpy() * 255).astype(np.uint8) src_tensor_list.append(pose_transform(Image.fromarray(src_image_pil))) - sub_step = step*fi_step if accelerate else step - for src_image_pil in images[: frame_length: sub_step]: + ''' + frame_length = len(images) if length==0 or length*step > len(images) else length*step + sub_step = fi_step if accelerate else step + for src_image_pil in images[: frame_length: step*sub_step]: src_image_pil = (src_image_pil.numpy() * 255).astype(np.uint8) src_img_np = cv2.cvtColor(np.array(src_image_pil), cv2.COLOR_RGB2BGR) frame_height, frame_width, _ = src_img_np.shape @@ -705,16 +751,32 @@ class Audio2Video: pose_trans_list.append(src_img_result['trans_mat']) verts_list.append(src_img_result['lmks3d']) bs_list.append(src_img_result['bs']) - - pose_arr = np.array(pose_trans_list) + + #pose_arr = np.array(pose_trans_list) + trans_mat_arr = np.array(pose_trans_list) verts_arr = np.array(verts_list) bs_arr = np.array(bs_list) min_bs_idx = np.argmin(bs_arr.sum(1)) + # compute delta pose + pose_arr = np.zeros([trans_mat_arr.shape[0], 6]) + for i in range(pose_arr.shape[0]): + euler_angles, translation_vector = matrix_to_euler_and_translation(trans_mat_arr[i]) # real pose of source + pose_arr[i, :3] = euler_angles + pose_arr[i, 3:6] = translation_vector + + init_tran_vec = face_result['trans_mat'][:3, 3] # init translation of tgt + pose_arr[:, 3:6] = pose_arr[:, 3:6] - pose_arr[0, 3:6] + init_tran_vec # (relative translation of source) + (init translation of tgt) + + pose_arr_smooth = smooth_pose_seq(pose_arr, window_size=3) + pose_mat_smooth = [euler_and_translation_to_matrix(pose_arr_smooth[i][:3], pose_arr_smooth[i][3:6]) for i in range(pose_arr_smooth.shape[0])] + pose_mat_smooth = np.array(pose_mat_smooth) + # face retarget verts_arr = verts_arr - verts_arr[min_bs_idx] + face_result['lmks3d'] # project 3D mesh to 2D landmark - projected_vertices = project_points_with_trans(verts_arr, pose_arr, [frame_height, frame_width]) + #projected_vertices = project_points_with_trans(verts_arr, pose_arr, [frame_height, frame_width]) + projected_vertices = project_points_with_trans(verts_arr, pose_mat_smooth, [frame_height, frame_width]) pose_list = [] for i, verts in enumerate(projected_vertices): @@ -726,11 +788,11 @@ class Audio2Video: video_length = len(pose_list) - + ''' src_tensor = torch.stack(src_tensor_list, dim=0) # (f, c, h, w) src_tensor = src_tensor.transpose(0, 1) src_tensor = src_tensor.unsqueeze(0) - + ''' video_gen = pipe( Image.fromarray(ref_image_pil), pose_list, diff --git a/src/audio_models/pose_model.py b/src/audio_models/pose_model.py new file mode 100644 index 0000000..fe28b6a --- /dev/null +++ b/src/audio_models/pose_model.py @@ -0,0 +1,126 @@ +import os +import math +import torch +import torch.nn as nn +from transformers import Wav2Vec2Config + +from .torch_utils import get_mask_from_lengths +from .wav2vec2 import Wav2Vec2Model + + +def init_biased_mask(n_head, max_seq_len, period): + def get_slopes(n): + def get_slopes_power_of_2(n): + start = (2**(-2**-(math.log2(n)-3))) + ratio = start + return [start*ratio**i for i in range(n)] + if math.log2(n).is_integer(): + return get_slopes_power_of_2(n) + else: + closest_power_of_2 = 2**math.floor(math.log2(n)) + return get_slopes_power_of_2(closest_power_of_2) + get_slopes(2*closest_power_of_2)[0::2][:n-closest_power_of_2] + slopes = torch.Tensor(get_slopes(n_head)) + bias = torch.arange(start=0, end=max_seq_len, step=period).unsqueeze(1).repeat(1,period).view(-1)//(period) + bias = - torch.flip(bias,dims=[0]) + alibi = torch.zeros(max_seq_len, max_seq_len) + for i in range(max_seq_len): + alibi[i, :i+1] = bias[-(i+1):] + alibi = slopes.unsqueeze(1).unsqueeze(1) * alibi.unsqueeze(0) + mask = (torch.triu(torch.ones(max_seq_len, max_seq_len)) == 1).transpose(0, 1) + mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0)) + mask = mask.unsqueeze(0) + alibi + return mask + + +def enc_dec_mask(device, T, S): + mask = torch.ones(T, S) + for i in range(T): + mask[i, i] = 0 + return (mask==1).to(device=device) + + +class PositionalEncoding(nn.Module): + def __init__(self, d_model, max_len=600): + super(PositionalEncoding, self).__init__() + pe = torch.zeros(max_len, d_model) + position = torch.arange(0, max_len).unsqueeze(1).float() + div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + pe = pe.unsqueeze(0) + self.register_buffer('pe', pe) + + def forward(self, x): + x = x + self.pe[:, :x.size(1)] + return x + + +class Audio2PoseModel(nn.Module): + def __init__( + self, + config + ): + + super().__init__() + + latent_dim = config['latent_dim'] + model_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), config['model_path']) + #model_path = config['model_path'] + only_last_fetures = config['only_last_fetures'] + from_pretrained = config['from_pretrained'] + out_dim = config['out_dim'] + + self.out_dim = out_dim + + self._only_last_features = only_last_fetures + + self.audio_encoder_config = Wav2Vec2Config.from_pretrained(model_path, local_files_only=True) + if from_pretrained: + self.audio_encoder = Wav2Vec2Model.from_pretrained(model_path, local_files_only=True) + else: + self.audio_encoder = Wav2Vec2Model(self.audio_encoder_config) + self.audio_encoder.feature_extractor._freeze_parameters() + + hidden_size = self.audio_encoder_config.hidden_size + + self.pose_map = nn.Linear(out_dim, latent_dim) + self.in_fn = nn.Linear(hidden_size, latent_dim) + + self.PPE = PositionalEncoding(latent_dim) + self.biased_mask = init_biased_mask(n_head = 8, max_seq_len = 600, period=1) + decoder_layer = nn.TransformerDecoderLayer(d_model=latent_dim, nhead=8, dim_feedforward=2*latent_dim, batch_first=True) + self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=8) + self.pose_map_r = nn.Linear(latent_dim, out_dim) + + self.id_embed = nn.Embedding(100, latent_dim) # 100 ids + + + def infer(self, input_value, seq_len, id_seed=None): + embeddings = self.audio_encoder(input_value, seq_len=seq_len, output_hidden_states=True) + + if self._only_last_features: + hidden_states = embeddings.last_hidden_state + else: + hidden_states = sum(embeddings.hidden_states) / len(embeddings.hidden_states) + + hidden_states = self.in_fn(hidden_states) + + id_embedding = self.id_embed(id_seed).unsqueeze(1) + + init_pose = torch.zeros([hidden_states.shape[0], 1, self.out_dim]).to(hidden_states.device) + for i in range(seq_len): + if i==0: + pose_emb = self.pose_map(init_pose) + pose_input = self.PPE(pose_emb) + else: + pose_input = self.PPE(pose_emb) + + pose_input = pose_input + id_embedding + tgt_mask = self.biased_mask[:, :pose_input.shape[1], :pose_input.shape[1]].clone().detach().to(hidden_states.device) + memory_mask = enc_dec_mask(hidden_states.device, pose_input.shape[1], hidden_states.shape[1]) + pose_out = self.transformer_decoder(pose_input, hidden_states, tgt_mask=tgt_mask, memory_mask=memory_mask) + pose_out = self.pose_map_r(pose_out) + new_output = self.pose_map(pose_out[:,-1,:]).unsqueeze(1) + pose_emb = torch.cat((pose_emb, new_output), 1) + return pose_out + diff --git a/src/utils/pose_util.py b/src/utils/pose_util.py index 611c690..4909ac1 100644 --- a/src/utils/pose_util.py +++ b/src/utils/pose_util.py @@ -67,4 +67,23 @@ def euler_and_translation_to_matrix(euler_angles, translation_vector): matrix[:3, :3] = rotation_matrix matrix[:3, 3] = translation_vector - return matrix \ No newline at end of file + return matrix + + +def matrix_to_euler_and_translation(matrix): + rotation_matrix = matrix[:3, :3] + translation_vector = matrix[:3, 3] + rotation = R.from_matrix(rotation_matrix) + euler_angles = rotation.as_euler('xyz', degrees=True) + return euler_angles, translation_vector + + +def smooth_pose_seq(pose_seq, window_size=5): + smoothed_pose_seq = np.zeros_like(pose_seq) + + for i in range(len(pose_seq)): + start = max(0, i - window_size // 2) + end = min(len(pose_seq), i + window_size // 2 + 1) + smoothed_pose_seq[i] = np.mean(pose_seq[start:end], axis=0) + + return smoothed_pose_seq