343 lines
15 KiB
Python
343 lines
15 KiB
Python
import os,sys
|
|
import uuid
|
|
from .data_preparation import CirculateVideo,work_dir,now_dir,ouput_dir
|
|
sys.path.append(now_dir)
|
|
import basicsr
|
|
sys.modules["basicsr"] = basicsr
|
|
import cv2
|
|
import shutil
|
|
import torchaudio
|
|
from tqdm import tqdm
|
|
from PIL import Image
|
|
import numpy as np
|
|
from talkingface.render_model import RenderModel
|
|
from talkingface.audio_model import AudioModel
|
|
|
|
|
|
checkpoints_dir = os.path.join(now_dir,"checkpoint")
|
|
|
|
class StaticVideo:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required":{
|
|
"image":("IMAGE",),
|
|
"length":("INT",{
|
|
"default":3
|
|
})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VIDEO",)
|
|
#RETURN_NAMES = ("image_output_name",)
|
|
|
|
FUNCTION = "generate"
|
|
|
|
#OUTPUT_NODE = False
|
|
|
|
CATEGORY = "AIFSH_DHLive"
|
|
|
|
def comfyimage2Image(self,comfyimage):
|
|
comfyimage = comfyimage.numpy()[0] * 255
|
|
image_np = comfyimage.astype(np.uint8)
|
|
image = Image.fromarray(image_np)
|
|
return image
|
|
|
|
def generate(self,image,length):
|
|
image = self.comfyimage2Image(image)
|
|
image_path = os.path.join(work_dir,"source.jpg")
|
|
image.save(image_path)
|
|
video_path = os.path.join(work_dir,"video.mp4")
|
|
cmd = f"ffmpeg -r 25 -f image2 -loop 1 -i {image_path} -vcodec libx264 -pix_fmt yuv420p -r 25 -t {length} -y {video_path}"
|
|
os.system(cmd)
|
|
return (video_path,)
|
|
|
|
import torch
|
|
def rgbnp2tensor(rgbnplist):
|
|
rgbnps = np.array(rgbnplist).copy()
|
|
lqinput = np.array(np.array(rgbnps) / 255.0, np.float32) # [t,h,w,c]
|
|
lqinput = torch.from_numpy(lqinput).permute(0, 3, 1, 2).cuda()
|
|
return lqinput
|
|
|
|
def apply_net_to_frames(frames, model, w=1.0):
|
|
lqinput = rgbnp2tensor(frames)
|
|
with torch.no_grad():
|
|
restored_faces = model(lqinput, w=w)[0][1]
|
|
restored_faces = torch.clamp(restored_faces, 0, 1)
|
|
restored_face = restored_faces.detach().cpu().permute(1, 2, 0).numpy() * 255
|
|
npface = np.array(restored_face, np.uint8)
|
|
return npface
|
|
|
|
import yaml
|
|
from collections import OrderedDict
|
|
|
|
def ordered_yaml():
|
|
"""Support OrderedDict for yaml.
|
|
|
|
Returns:
|
|
yaml Loader and Dumper.
|
|
"""
|
|
try:
|
|
from yaml import CDumper as Dumper
|
|
from yaml import CLoader as Loader
|
|
except ImportError:
|
|
from yaml import Dumper, Loader
|
|
|
|
_mapping_tag = yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG
|
|
|
|
def dict_representer(dumper, data):
|
|
return dumper.represent_dict(data.items())
|
|
|
|
def dict_constructor(loader, node):
|
|
return OrderedDict(loader.construct_pairs(node))
|
|
|
|
Dumper.add_representer(OrderedDict, dict_representer)
|
|
Loader.add_constructor(_mapping_tag, dict_constructor)
|
|
return Loader, Dumper
|
|
|
|
import torch
|
|
from PIL import Image
|
|
from torchvision.transforms.functional import normalize
|
|
from basicsr.utils import imwrite, img2tensor, tensor2img
|
|
from basicsr.utils.download_util import load_file_from_url
|
|
from basicsr.utils.registry import ARCH_REGISTRY
|
|
|
|
prompt_sr = 16000
|
|
class DHLiveNode:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required":{
|
|
"audio":("AUDIO",),
|
|
"video":("VIDEO",),
|
|
"if_data_preparation":("BOOLEAN",{
|
|
"default":True
|
|
}),
|
|
"upscale":(["None","GFPGANv1.4","CodeFormer",
|
|
"RestoreFormer++","PGTFormer","KEEP"],),
|
|
"padding":("INT",{
|
|
"default": 512,
|
|
}),
|
|
"if_RIFE":("BOOLEAN",{
|
|
"default":True
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VIDEO",)
|
|
#RETURN_NAMES = ("image_output_name",)
|
|
|
|
FUNCTION = "generate"
|
|
|
|
#OUTPUT_NODE = False
|
|
|
|
CATEGORY = "AIFSH_DHLive"
|
|
|
|
def generate(self,audio,video,if_data_preparation,upscale,padding,if_RIFE):
|
|
## 1. data preparation
|
|
video_data_path = os.path.join(work_dir,"video_data")
|
|
if if_data_preparation:
|
|
shutil.rmtree(video_data_path,ignore_errors=True)
|
|
os.makedirs(video_data_path,exist_ok=True)
|
|
video_out_path = os.path.join(video_data_path,"circle.mp4")
|
|
CirculateVideo(video_in_path=video,video_out_path=video_out_path)
|
|
|
|
audioModel = AudioModel()
|
|
audioModel.loadModel(os.path.join(checkpoints_dir,"audio.pkl"))
|
|
|
|
renderModel = RenderModel()
|
|
renderModel.loadModel(os.path.join(checkpoints_dir,"render.pth"))
|
|
|
|
pkl_path = "{}/keypoint_rotate.pkl".format(video_data_path)
|
|
video_path = "{}/circle.mp4".format(video_data_path)
|
|
renderModel.reset_charactor(video_path, pkl_path)
|
|
|
|
waveform = audio['waveform'].squeeze(0)
|
|
source_sr = audio['sample_rate']
|
|
speech = waveform.mean(dim=0,keepdim=True)
|
|
if source_sr != prompt_sr:
|
|
speech = torchaudio.transforms.Resample(orig_freq=source_sr, new_freq=prompt_sr)(speech)
|
|
wavpath = os.path.join(ouput_dir,"tmp.wav")
|
|
torchaudio.save(wavpath,speech,prompt_sr,format="wav")
|
|
mouth_frame = audioModel.interface_wav(speech.numpy()[0])
|
|
|
|
cap_input = cv2.VideoCapture(video_path)
|
|
vid_width = cap_input.get(cv2.CAP_PROP_FRAME_WIDTH) # 宽度
|
|
vid_height = cap_input.get(cv2.CAP_PROP_FRAME_HEIGHT) # 高度
|
|
if if_RIFE:
|
|
import math
|
|
def get_64x_num(num):
|
|
return math.ceil(num / 64) * 64
|
|
height,width = (1024,get_64x_num(1024*vid_width/vid_height)) if vid_height > vid_width else (get_64x_num(1024*vid_height/vid_width),1024)
|
|
cap_input.release()
|
|
task_id = str(uuid.uuid1())
|
|
output_video_name = os.path.join(ouput_dir,task_id+".mp4")
|
|
task_id_path = os.path.join(ouput_dir,"DHLive","output",task_id)
|
|
os.makedirs(task_id_path, exist_ok=True)
|
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
|
save_path = os.path.join(task_id_path,"silence.mp4")
|
|
videoWriter = cv2.VideoWriter(save_path, fourcc, 25, (int(vid_width) * 1, int(vid_height)))
|
|
image_frame_list = []
|
|
if upscale == "PGTFormer":
|
|
import yaml
|
|
with open(os.path.join(checkpoints_dir,"upscale","PGTFormer.yml"), mode='r') as f:
|
|
opt = yaml.load(f, Loader=ordered_yaml()[0])
|
|
ooo = opt['network_g']
|
|
print(ooo)
|
|
# from talkingface.basicsr.archs.pgtformer_arch import PGTFormer
|
|
# upscale_model = PGTFormer(**ooo).cuda()
|
|
upscale_model = ARCH_REGISTRY.get("PGTFormer")(**ooo).cuda()
|
|
state_dict = torch.load(os.path.join(checkpoints_dir,"upscale","PGTFormer.pth"))
|
|
upscale_model.load_state_dict(state_dict=state_dict['params_ema'])
|
|
upscale_model.eval()
|
|
upscale_model.requires_grad_(False)
|
|
# Read the first frame
|
|
frame_buffer = []
|
|
pre_frame = None
|
|
pre_crop_coords = None
|
|
for i , raw_frame in tqdm(enumerate(mouth_frame)):
|
|
frame, face_numpy, crop_coords = renderModel.interface(raw_frame,upscale="None",padding=padding)
|
|
face_frame = cv2.resize(face_numpy, (512, 512), interpolation=cv2.INTER_LINEAR)
|
|
if i == 0:
|
|
# print(face_frame.shape)
|
|
frame_buffer.append(face_frame)
|
|
frame_buffer.append(face_frame)
|
|
else:
|
|
frame_buffer.append(face_frame)
|
|
if len(frame_buffer) == 3:
|
|
processed_frame = apply_net_to_frames(frame_buffer, upscale_model)
|
|
# processed_face_numpy = cv2.resize(processed_frame, (256,256), interpolation=cv2.INTER_LINEAR)
|
|
x_min, y_min, x_max, y_max = pre_crop_coords
|
|
|
|
img_face = cv2.resize(processed_frame, (x_max - x_min, y_max - y_min),interpolation=cv2.INTER_LINEAR)
|
|
|
|
pre_frame[y_min:y_max, x_min:x_max] = img_face
|
|
videoWriter.write(pre_frame)
|
|
if if_RIFE:
|
|
pre_frame = cv2.resize(pre_frame, (width, height),interpolation=cv2.INTER_LINEAR)
|
|
image_frame_list.append(Image.fromarray(cv2.cvtColor(pre_frame,cv2.COLOR_BGR2RGB)))
|
|
# Remove the first frame from the buffer and continue
|
|
frame_buffer.pop(0)
|
|
pre_crop_coords = crop_coords
|
|
pre_frame = frame
|
|
|
|
# Process the last two frames (when padding a frame is needed)
|
|
if len(frame_buffer) == 2:
|
|
frame_buffer.append(frame_buffer[-1]) # Pad the last frame
|
|
processed_frame = apply_net_to_frames(frame_buffer, upscale_model)
|
|
# processed_face_numpy = cv2.resize(processed_frame, (256,256), interpolation=cv2.INTER_LINEAR)
|
|
x_min, y_min, x_max, y_max = pre_crop_coords
|
|
img_face = cv2.resize(processed_frame, (x_max - x_min, y_max - y_min),interpolation=cv2.INTER_LINEAR)
|
|
pre_frame[y_min:y_max, x_min:x_max] = img_face
|
|
videoWriter.write(pre_frame)
|
|
if if_RIFE:
|
|
pre_frame = cv2.resize(pre_frame, (width, height),interpolation=cv2.INTER_LINEAR)
|
|
image_frame_list.append(Image.fromarray(cv2.cvtColor(pre_frame,cv2.COLOR_BGR2RGB)))
|
|
|
|
else:
|
|
cropped_faces = []
|
|
crop_coords = []
|
|
frames_list = []
|
|
for frame in tqdm(mouth_frame):
|
|
if upscale == "KEEP":
|
|
frame,face_img,coords = renderModel.interface(frame,upscale="None",padding=padding)
|
|
face_img = cv2.resize(face_img, (512, 512), interpolation=cv2.INTER_LINEAR)
|
|
cropped_face_t = img2tensor(
|
|
face_img / 255., bgr2rgb=True, float32=True)
|
|
normalize(cropped_face_t, (0.5, 0.5, 0.5),
|
|
(0.5, 0.5, 0.5), inplace=True)
|
|
cropped_faces.append(cropped_face_t)
|
|
crop_coords.append(coords)
|
|
frames_list.append(frame)
|
|
else:
|
|
frame,face_img,coords = renderModel.interface(frame,upscale=upscale,padding=padding)
|
|
videoWriter.write(frame)
|
|
if if_RIFE:
|
|
frame = cv2.resize(frame, (width, height),interpolation=cv2.INTER_LINEAR)
|
|
image_frame_list.append(Image.fromarray(cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)))
|
|
|
|
if upscale == "KEEP":
|
|
# ------------------ set up restorer -------------------
|
|
net = ARCH_REGISTRY.get('KEEP')(img_size=512, emb_dim=256, dim_embd=512,
|
|
n_head=8, n_layers=9, codebook_size=1024, connect_list=['16', '32', '64'],
|
|
flow_type='gmflow', flownet_path=None,
|
|
kalman_attn_head_dim=48, num_uncertainty_layers=3, cross_fuse_list=['16', '32'],
|
|
cross_fuse_nhead=4, cross_fuse_dim=256).to("cuda")
|
|
|
|
# ckpt_path = 'weights/KEEP/KEEP-a1a14d46.pth'
|
|
# checkpoint = torch.load(ckpt_path)['params_ema']
|
|
|
|
ckpt_path = load_file_from_url(
|
|
url='https://github.com/jnjaby/KEEP/releases/download/v0.1.0/KEEP-a1a14d46.pth',
|
|
model_dir=os.path.join(checkpoints_dir,"upscale"), progress=True, file_name=None)
|
|
checkpoint = torch.load(ckpt_path)
|
|
net.load_state_dict(checkpoint)
|
|
net.eval()
|
|
|
|
cropped_faces = torch.stack(
|
|
cropped_faces, dim=0).unsqueeze(0).to("cuda")
|
|
print(cropped_faces.shape)
|
|
with torch.no_grad():
|
|
video_length = cropped_faces.shape[1]
|
|
output = []
|
|
for start_idx in range(0, video_length, 20):
|
|
end_idx = min(start_idx + 20, video_length)
|
|
if end_idx - start_idx == 1:
|
|
output.append(net(
|
|
cropped_faces[:, [start_idx, start_idx], ...], need_upscale=False)[:, 0:1, ...])
|
|
else:
|
|
output.append(net(
|
|
cropped_faces[:, start_idx:end_idx, ...], need_upscale=False))
|
|
output = torch.cat(output, dim=1).squeeze(0)
|
|
assert output.shape[0] == video_length, "Differer number of frames"
|
|
|
|
restored_faces = [tensor2img(
|
|
x, rgb2bgr=True, min_max=(-1, 1)) for x in output]
|
|
del output
|
|
torch.cuda.empty_cache()
|
|
for frame,face,coords in tqdm(zip(frames_list,restored_faces,crop_coords)):
|
|
|
|
x_min, y_min, x_max, y_max = coords
|
|
img_face = cv2.resize(face, (x_max - x_min, y_max - y_min),interpolation=cv2.INTER_LINEAR)
|
|
frame[y_min:y_max, x_min:x_max] = img_face
|
|
videoWriter.write(frame)
|
|
if if_RIFE:
|
|
frame = cv2.resize(frame, (width, height),interpolation=cv2.INTER_LINEAR)
|
|
image_frame_list.append(Image.fromarray(cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)))
|
|
|
|
videoWriter.release()
|
|
if if_RIFE:
|
|
import imageio
|
|
from .rife import IFNet,RIFESmoother
|
|
model = IFNet().eval()
|
|
state_dict = torch.load(os.path.join(checkpoints_dir,"upscale","flownet.pkl"),map_location="cpu")
|
|
model.load_state_dict(model.state_dict_converter().from_civitai(state_dict))
|
|
model.to(torch.float32).to("cuda")
|
|
smoother = RIFESmoother.load_model(model)
|
|
out_video = smoother(image_frame_list)
|
|
print("use RIFE flow")
|
|
def save_video(frames, save_path, fps=25, quality=9):
|
|
writer = imageio.get_writer(save_path, fps=fps, quality=quality)
|
|
for frame in tqdm(frames, desc="Saving video"):
|
|
frame = np.array(frame)
|
|
writer.append_data(frame)
|
|
writer.close()
|
|
save_video(out_video,save_path)
|
|
|
|
os.system(
|
|
"ffmpeg -i {} -i {} -c:v libx264 -pix_fmt yuv420p -loglevel quiet {}".format(save_path, wavpath, output_video_name))
|
|
shutil.rmtree(task_id_path)
|
|
return (output_video_name,)
|
|
|
|
WEB_DIRECTORY = "./web"
|
|
|
|
from .util_nodes import PreViewVideo, LoadVideo,CombineVideo
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"CombineVideo":CombineVideo,
|
|
"PreViewVideo":PreViewVideo,
|
|
"DHLiveNode": DHLiveNode,
|
|
"StaticVideo":StaticVideo,
|
|
"DHLIVELoadVideo":LoadVideo,
|
|
} |