311 lines
13 KiB
Python
311 lines
13 KiB
Python
|
|
import json
|
|
from multiprocessing.pool import ThreadPool
|
|
import os
|
|
import shutil
|
|
import cv2
|
|
from tqdm import tqdm
|
|
|
|
from ..lib.audio import get_audio, remix_audio
|
|
from ..lib import BASE_CACHE_DIR
|
|
from ..lib.helper import gc_collect, get_hash, get_optimal_torch_device
|
|
from .helper import MultipleTypeProxy, model_downloader
|
|
from ..lib.musetalk.models.unet import PositionalEncoding, UNet
|
|
from ..lib.musetalk.models.vae import VAE
|
|
from ..lib.musetalk.whisper.audio2feature import Audio2Feature
|
|
import torch
|
|
import numpy as np
|
|
import folder_paths
|
|
from ..lib.musetalk.utils.face_parsing import FaceParsing
|
|
from ..lib.musetalk.utils.blending import get_image
|
|
from ..lib.musetalk.utils.utils import datagen
|
|
import torchvision.transforms as TT
|
|
import torchvision
|
|
from PIL import Image
|
|
|
|
CATEGORY = "🌺RVC-Studio/musetalk"
|
|
input_path = folder_paths.get_input_directory()
|
|
out_path = folder_paths.get_output_directory()
|
|
temp_path = folder_paths.get_temp_directory()
|
|
project_directory = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
def tensor_to_opencv(tensor_image: torch.Tensor):
|
|
numpy_image = tensor_image.detach().numpy()
|
|
if numpy_image.max()<=1: numpy_image *= 255
|
|
opencv_image = cv2.cvtColor(numpy_image.astype(np.uint8), cv2.COLOR_RGB2BGR)
|
|
opencv_image
|
|
return opencv_image
|
|
|
|
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))
|
|
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")
|
|
|
|
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!"
|
|
|
|
files = sorted(fname for fname in os.listdir(directory_path) if fname.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff')))
|
|
|
|
assert len(files)>0, f"{directory_path} is empty!"
|
|
return files
|
|
|
|
def load_images_from_directory(directory_path):
|
|
image_list = []
|
|
img_size = None
|
|
img_to_tensor = TT.ToTensor()
|
|
|
|
files = get_imagefiles(directory_path)
|
|
|
|
# Iterate over all files in the directory
|
|
for filename in tqdm(files):
|
|
# Load the image
|
|
img_path = os.path.join(directory_path, filename)
|
|
image = Image.open(img_path).convert('RGB')
|
|
|
|
# Infer image size from the first image
|
|
if img_size is None:
|
|
img_size = image.size # (width, height)
|
|
|
|
# Resize image if it's not the same size as the first image
|
|
if image.size != img_size:
|
|
image = image.resize(img_size)
|
|
|
|
# Convert image to tensor (HWC)
|
|
tensor_image = img_to_tensor(image).permute(1, 2, 0)
|
|
|
|
image_list.append(tensor_image)
|
|
|
|
# Stack all images to form a batch (B, H, W, C)
|
|
image_batch = torch.stack(image_list, dim=0)
|
|
|
|
return image_batch
|
|
|
|
class MuseAudioFeatureExtractionNode:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
|
|
return {
|
|
"required": {
|
|
"audio": (MultipleTypeProxy("AUDIO,VHS_AUDIO"),),
|
|
},
|
|
"optional": {
|
|
"fps": (MultipleTypeProxy("FLOAT,INT"),{"default": 1.}),
|
|
}
|
|
}
|
|
|
|
CATEGORY = CATEGORY
|
|
|
|
RETURN_TYPES = ("WHISPER_CHUNKS",)
|
|
RETURN_NAMES = ("whisper_chunks",)
|
|
FUNCTION = "process"
|
|
|
|
def process(self, audio, fps):
|
|
print("############################################## extract audio feature ##############################################")
|
|
model_path = model_downloader("musetalk/whisper/tiny.pt")
|
|
audio_processor = Audio2Feature(model_path=model_path)
|
|
input_audio = remix_audio(get_audio(audio),target_sr=16000,norm=True)
|
|
whisper_feature = audio_processor.audio2feat(np.array(input_audio[0],dtype=np.float32))
|
|
whisper_chunks = audio_processor.feature2chunks(feature_array=whisper_feature,fps=fps)
|
|
del audio_processor, whisper_feature, input_audio
|
|
gc_collect()
|
|
|
|
return (whisper_chunks,)
|
|
|
|
class MuseImageFeatureExtractionNode:
|
|
|
|
device = get_optimal_torch_device()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"bbox_detector": ("BBOX_DETECTOR", )
|
|
},
|
|
"optional": {
|
|
"threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"dilation": ("INT", {"default": 0, "min": -512, "max": 512, "step": 1}),
|
|
"use_cache": ("BOOLEAN",{"default": True})
|
|
}
|
|
}
|
|
|
|
CATEGORY = CATEGORY
|
|
|
|
RETURN_TYPES = ("COORDS","STRING")
|
|
RETURN_NAMES = ("coords","frames_dir")
|
|
FUNCTION = "process"
|
|
|
|
@staticmethod
|
|
def mask_to_bbox(mask):
|
|
# Ensure the mask is a binary tensor
|
|
if mask is None: return (0.0,0.0,0.0,0.0)
|
|
elif mask.dtype != torch.bool: mask = mask > 0
|
|
|
|
# Find the bounding box coordinates using torchvision's masks_to_boxes
|
|
boxes = torchvision.ops.masks_to_boxes(mask.unsqueeze(0))
|
|
|
|
# Convert the boxes tensor to a list of coordinates
|
|
x_min, y_min, x_max, y_max = boxes[0].tolist()
|
|
|
|
return x_min, y_min, x_max, y_max
|
|
|
|
def process(self, images, bbox_detector, threshold=.5, dilation=0, use_cache=True):
|
|
print("############################################## preprocess input image ##############################################")
|
|
cache_name = get_hash(images,bbox_detector,len(images),threshold,dilation)
|
|
frames_dir = os.path.join(temp_path,f"frames-{cache_name}")
|
|
crop_coord_save_path = os.path.join(BASE_CACHE_DIR,"musetalk",f"coords-{cache_name}.json")
|
|
|
|
if os.path.isfile(crop_coord_save_path) and os.path.isdir(frames_dir) and use_cache:
|
|
print("using extracted coordinates")
|
|
with open(crop_coord_save_path,'r') as f:
|
|
data = json.load(f)
|
|
coord_list = data["coord_list"]
|
|
else:
|
|
os.makedirs(frames_dir,exist_ok=True)
|
|
print("extracting landmarks...time consuming")
|
|
coord_list = []
|
|
has_face = False
|
|
for i,image in enumerate(tqdm(images.cpu())):
|
|
mask = bbox_detector.detect_combined(image.unsqueeze(0), threshold, dilation)
|
|
if mask is not None: has_face=True
|
|
coords = self.mask_to_bbox(mask)
|
|
coord_list.append(coords)
|
|
cv2.imwrite(os.path.join(frames_dir, f"{str(i).zfill(8)}.png"), tensor_to_opencv(image))
|
|
assert has_face, "No face detected!"
|
|
|
|
if use_cache:
|
|
os.makedirs(os.path.dirname(crop_coord_save_path),exist_ok=True)
|
|
with open(crop_coord_save_path, 'w') as f:
|
|
json.dump(dict(coord_list=coord_list), f)
|
|
|
|
gc_collect()
|
|
|
|
return (coord_list, frames_dir)
|
|
|
|
class MuseTalkNode:
|
|
|
|
device = get_optimal_torch_device()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
|
|
return {
|
|
"required": {
|
|
"coord_list": ("COORDS",),
|
|
"frames_dir": ("STRING",{"default": ""}),
|
|
"whisper_chunks": ("WHISPER_CHUNKS",)
|
|
},
|
|
"optional": {
|
|
"batch_size": ("INT",{"default": 1, "min": 1}),
|
|
}
|
|
}
|
|
|
|
CATEGORY = CATEGORY
|
|
|
|
RETURN_TYPES = ("IMAGE","STRING")
|
|
RETURN_NAMES = ("images","results_dir")
|
|
FUNCTION = "process"
|
|
|
|
def process(self, coord_list, frames_dir, whisper_chunks, batch_size=1):
|
|
frame_list = get_imagefiles(frames_dir)
|
|
results_dir = os.path.join(temp_path,f"results-{get_hash(coord_list, whisper_chunks, *frame_list)}")
|
|
|
|
if not os.path.isdir(results_dir) or len(os.listdir(results_dir))==0:
|
|
os.makedirs(results_dir,exist_ok=True)
|
|
|
|
print("############################################## process latents ##############################################")
|
|
vae_model = model_downloader("musetalk/sd-vae-ft-mse/diffusion_pytorch_model.safetensors")
|
|
_ = 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}")
|
|
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
|
|
frame_list_cycle = frame_list + frame_list[::-1]
|
|
coord_list_cycle = coord_list + coord_list[::-1]
|
|
input_latent_list_cycle = input_latent_list + input_latent_list[::-1]
|
|
print(f"{len(frame_list_cycle)=} {len(coord_list_cycle)=} {len(input_latent_list_cycle)=} {len(whisper_chunks)=}")
|
|
|
|
print("############################################## inference batch by batch ##############################################")
|
|
unet_config = model_downloader("musetalk/musetalk.json")
|
|
unet_model = model_downloader("musetalk/pytorch_model.bin")
|
|
unet = UNet(unet_config=unet_config,model_path=unet_model,use_float16=True)
|
|
pe = PositionalEncoding(d_model=384)
|
|
video_num = len(whisper_chunks)
|
|
dataset = datagen(whisper_chunks,input_latent_list_cycle,batch_size)
|
|
res_frame_list = []
|
|
timestep = torch.tensor([0], device=unet.device)
|
|
for (whisper_batch,latent_batch) in tqdm(dataset,total=int(np.ceil(float(video_num)/batch_size))):
|
|
tensor_list = [torch.FloatTensor(arr) for arr in whisper_batch]
|
|
audio_feature_batch = torch.stack(tensor_list) # torch, B, 5*N,384
|
|
audio_feature_batch = pe(audio_feature_batch)
|
|
|
|
pred_latents = unet.model(latent_batch.to(unet.device).half(), timestep, encoder_hidden_states=audio_feature_batch.to(unet.device).half()).sample
|
|
recon = vae.decode_latents(pred_latents)
|
|
for res_frame in recon:
|
|
res_frame_list.append(res_frame)
|
|
del pe, unet.model, vae, unet
|
|
gc_collect()
|
|
|
|
print("############################################## pad to full image ##############################################")
|
|
resnet_path = model_downloader('musetalk/face-parse-bisent/resnet18-5c106cde.pth')
|
|
face_model_pth = model_downloader('musetalk/face-parse-bisent/79999_iter.pth')
|
|
fp_model = FaceParsing(resnet_path,face_model_pth)
|
|
|
|
with ThreadPool(batch_size) as pool:
|
|
for i, res_frame in enumerate(tqdm(res_frame_list)):
|
|
pool.apply(process_frame, args=(i, res_frame, coord_list_cycle, frame_list_cycle, frames_dir, results_dir, fp_model))
|
|
|
|
del fp_model, frame_list_cycle, res_frame_list, coord_list_cycle, input_latent_list
|
|
gc_collect()
|
|
|
|
images = load_images_from_directory(results_dir)
|
|
|
|
print(f"{images.shape=} {images.max()=} {images.dtype=}")
|
|
|
|
return (images,results_dir)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"MuseTalkNode": MuseTalkNode,
|
|
"MuseAudioFeatureExtractionNode": MuseAudioFeatureExtractionNode,
|
|
"MuseImageFeatureExtractionNode": MuseImageFeatureExtractionNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"MuseTalkNode": "🌺MuseTalk Processor",
|
|
"MuseAudioFeatureExtractionNode": "🌺MuseTalk Feature Extractor",
|
|
"MuseImageFeatureExtractionNode": "🌺MuseTalk Image Preprocessing"
|
|
} |