From 61dbf9d59ae8d5313e4b70f21f9efdf7c7f69ca7 Mon Sep 17 00:00:00 2001 From: SayanoAI <28544690+SayanoAI@users.noreply.github.com> Date: Mon, 16 Sep 2024 14:31:50 -0400 Subject: [PATCH] musetalk undetected face mismatch fix --- custom_nodes/audio_nodes.py | 27 ++++++++++++++++++++--- custom_nodes/musetalk_nodes.py | 40 ++++++++++++++++++++++------------ custom_nodes/utility_nodes.py | 6 ----- 3 files changed, 50 insertions(+), 23 deletions(-) diff --git a/custom_nodes/audio_nodes.py b/custom_nodes/audio_nodes.py index e1c2a8f..44933c2 100644 --- a/custom_nodes/audio_nodes.py +++ b/custom_nodes/audio_nodes.py @@ -54,7 +54,26 @@ class LoadAudio: def IS_CHANGED(cls, audio, sr): return get_hash(audio,sr) - +class AudioInfoNode: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "audio": (MultipleTypeProxy('AUDIO,VHS_AUDIO'),), + }} + + CATEGORY = CATEGORY + + RETURN_TYPES = ("VHS_AUDIO","AUDIO","FLOAT", "INT") + RETURN_NAMES = ("vhs_audio","audio","seconds","sr") + FUNCTION = "get_info" + + def get_info(self, audio): + input_audio = get_audio(audio) + seconds = len(input_audio[0])/input_audio[1] + return (lambda:audio_to_bytes(*input_audio),to_audio_dict(*input_audio),float(seconds),int(input_audio[1])) + + class DownloadAudio: @classmethod def INPUT_TYPES(cls): @@ -336,7 +355,8 @@ NODE_CLASS_MAPPINGS = { "MergeAudioNode": MergeAudioNode, "AudioBatchValueNode": AudioBatchValueNode, "DownloadAudio": DownloadAudio, - "ProcessAudioNode": ProcessAudioNode + "ProcessAudioNode": ProcessAudioNode, + "AudioInfoNode": AudioInfoNode } NODE_DISPLAY_NAME_MAPPINGS = { @@ -345,5 +365,6 @@ NODE_DISPLAY_NAME_MAPPINGS = { "RVC-Studio.PreviewAudio": "🌺Save Audio", "MergeAudioNode": "🌺Merge Audio", "AudioBatchValueNode": "🌺Audio RMS Batch Values", - "ProcessAudioNode": "🌺Audio Processor" + "ProcessAudioNode": "🌺Audio Processor", + "AudioInfoNode": "🌺Audio Info" } \ No newline at end of file diff --git a/custom_nodes/musetalk_nodes.py b/custom_nodes/musetalk_nodes.py index fa6e88c..aaa5f10 100644 --- a/custom_nodes/musetalk_nodes.py +++ b/custom_nodes/musetalk_nodes.py @@ -2,8 +2,8 @@ import json from multiprocessing.pool import ThreadPool import os +import shutil import cv2 -import copy from tqdm import tqdm from ..lib.audio import get_audio, remix_audio @@ -39,15 +39,20 @@ def tensor_to_opencv(tensor_image: torch.Tensor): def process_frame(i, res_frame, coord_list_cycle, frame_list_cycle, frames_dir, results_dir, fp_model): bbox = coord_list_cycle[i % (len(coord_list_cycle))] bbox = list(map(int, bbox)) - ori_frame = cv2.imread(os.path.join(frames_dir,frame_list_cycle[i % (len(frame_list_cycle))])) - x1, y1, x2, y2 = bbox + img_src = os.path.join(frames_dir,frame_list_cycle[i % (len(frame_list_cycle))]) + img_dest = os.path.join(results_dir, f"{str(i).zfill(8)}.png") - try: - res_frame = cv2.resize(res_frame.astype(np.uint8), (x2-x1, y2-y1)) - combined_frame = get_image(fp_model, ori_frame, res_frame, bbox) - cv2.imwrite(os.path.join(results_dir, f"{str(i).zfill(8)}.png"), combined_frame) - except Exception as error: - print(f"{error=}") + x1, y1, x2, y2 = bbox + if sum(bbox)>0: + try: + ori_frame = cv2.imread(img_src) + res_frame = cv2.resize(res_frame.astype(np.uint8), (x2-x1, y2-y1)) + if sum(bbox)==0: combined_frame = ori_frame.clone() #use original frame if no face detected + else: combined_frame = get_image(fp_model, ori_frame, res_frame, bbox) + return cv2.imwrite(img_dest, combined_frame) + except Exception as error: + print(f"{error=}") + return shutil.copy(img_src,img_dest) def get_imagefiles(directory_path): assert os.path.isdir(directory_path), f"{directory_path} is not a directory!" @@ -228,17 +233,24 @@ class MuseTalkNode: _ = model_downloader("musetalk/sd-vae-ft-mse/config.json") vae = VAE(model_path = os.path.dirname(vae_model), use_float16=True) + # init empty latent + empty_frame = cv2.imread(os.path.join(frames_dir,frame_list[0])) + empty_frame = cv2.resize(empty_frame,(256,256),interpolation = cv2.INTER_NEAREST) + empty_latent = torch.zeros_like(vae.get_latents_for_unet(empty_frame)) + input_latent_list = [] for bbox, frame in tqdm(zip(coord_list, frame_list)): bbox = list(map(int,bbox)) frame = cv2.imread(os.path.join(frames_dir,frame)) if sum(bbox) == 0: print(f"No face detected: {bbox}") - continue - x1, y1, x2, y2 = bbox - crop_frame = frame[y1:y2, x1:x2] - crop_frame = cv2.resize(crop_frame,(256,256),interpolation = cv2.INTER_LANCZOS4) - latents = vae.get_latents_for_unet(crop_frame) + latents = empty_latent.clone() + else: + x1, y1, x2, y2 = bbox + crop_frame = frame[y1:y2, x1:x2] + crop_frame = cv2.resize(crop_frame,(256,256),interpolation = cv2.INTER_LANCZOS4) + latents = vae.get_latents_for_unet(crop_frame) + input_latent_list.append(latents) # to smooth the first and the last frame diff --git a/custom_nodes/utility_nodes.py b/custom_nodes/utility_nodes.py index 3d6596f..426f3e7 100644 --- a/custom_nodes/utility_nodes.py +++ b/custom_nodes/utility_nodes.py @@ -75,12 +75,6 @@ class MergeImageBatches: for new_img in images[1:]: image = torch.cat([image,new_img]) gc_collect() - - # if self.__archive__ is None: - # with tempfile.NamedTemporaryFile(delete=False) as ntf: - # self.__archive__ = ntf.name - # images = self.iter_images(self.__archive__,images) - del images images = image gc_collect() else: