musetalk undetected face mismatch fix

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"
 }

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

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: