95 lines
3.2 KiB
Python
95 lines
3.2 KiB
Python
import os
|
|
import sys
|
|
import folder_paths
|
|
import os.path as osp
|
|
now_dir = osp.dirname(__file__)
|
|
tmp_dir = osp.join(now_dir,"temp")
|
|
aifsh_dir = osp.join(folder_paths.models_dir,"AIFSH")
|
|
sys.path.append(now_dir)
|
|
minimates_models_dir = osp.join(aifsh_dir,"MiniMates")
|
|
output_dir = folder_paths.get_output_directory()
|
|
import cv2
|
|
import torch
|
|
import shutil
|
|
import torchaudio
|
|
import tempfile
|
|
import numpy as np
|
|
from PIL import Image
|
|
from pathlib import Path
|
|
|
|
class MiniMatesNode:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required":{
|
|
"avator":("IMAGE",),
|
|
"driving_audio":("AUDIO",),
|
|
"if_matting":("BOOLEAN",{
|
|
"default": False,
|
|
"tooltip":"if matting the person from avator image"
|
|
})
|
|
},
|
|
"optional":{
|
|
"driving_video":("VIDEO",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VIDEO",)
|
|
#RETURN_NAMES = ("image_output_name",)
|
|
|
|
FUNCTION = "gen_video"
|
|
|
|
#OUTPUT_NODE = False
|
|
|
|
CATEGORY = "AIFSH_MiniMates"
|
|
|
|
def gen_video(self,avator,driving_audio,if_matting,driving_video=None):
|
|
os.makedirs(tmp_dir,exist_ok=True)
|
|
os.environ['minimates_models_dir'] = minimates_models_dir
|
|
with tempfile.NamedTemporaryFile(suffix=".png",dir=tmp_dir,delete=False) as img:
|
|
avator_np = avator.numpy()[0] * 255
|
|
avator_np = avator_np.astype(np.uint8)
|
|
avator_cv2 = cv2.cvtColor(avator_np,cv2.COLOR_BGR2RGBA)
|
|
cv2.imwrite(img.name,avator_cv2)
|
|
avator_path = Path(img.name)
|
|
|
|
if if_matting:
|
|
from minimates.interface.matting import Matting
|
|
matting_model = Matting(osp.join(minimates_models_dir,"modnet.onnx"))
|
|
save_img_path = osp.join(avator_path.parent,avator_path.stem+"_rgba.png")
|
|
matting_model.predict_image_rgba(avator_path,save_img_path)
|
|
avator_path = Path(save_img_path)
|
|
|
|
py = sys.executable or "python"
|
|
if driving_video:
|
|
template_path = osp.join(output_dir,Path(driving_video).stem+".template")
|
|
py_path = osp.join(now_dir,"minimates/interface/generate_move_template.py")
|
|
cmd = f"""{py} {py_path} {driving_video} {template_path}"""
|
|
if not osp.exists(template_path):
|
|
print(cmd)
|
|
os.system(cmd)
|
|
else:
|
|
template_path = None
|
|
|
|
with tempfile.NamedTemporaryFile(suffix=".wav",dir=tmp_dir,delete=False) as aud:
|
|
waveform = driving_audio['waveform'][0]
|
|
sample_rate = driving_audio["sample_rate"]
|
|
torchaudio.save(aud.name,waveform,sample_rate)
|
|
audio_path = Path(aud.name)
|
|
py_path = osp.join(now_dir,"minimates/interface/interface_audio.py")
|
|
output_path = osp.join(output_dir,avator_path.stem+"_"+audio_path.stem+".mp4")
|
|
if template_path:
|
|
cmd = f"""{py} {py_path} {avator_path} {audio_path} {output_path} {template_path}"""
|
|
else:
|
|
cmd = f"""{py} {py_path} {avator_path} {audio_path} {output_path}"""
|
|
print(cmd)
|
|
os.system(cmd)
|
|
shutil.rmtree(tmp_dir)
|
|
return (output_path, )
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"MiniMatesNode": MiniMatesNode
|
|
}
|