219 lines
8.3 KiB
Python
219 lines
8.3 KiB
Python
import comfy.utils
|
|
import folder_paths
|
|
try:
|
|
from comfy_api.input_impl.video_types import VideoFromFile
|
|
except ImportError:
|
|
VideoFromFile = None
|
|
|
|
import uuid
|
|
import os
|
|
import io
|
|
import av
|
|
import time
|
|
import yaml
|
|
import imageio
|
|
import subprocess
|
|
import numpy as np
|
|
import torch
|
|
import librosa
|
|
from pathlib import Path
|
|
from collections import deque
|
|
from PIL import Image
|
|
from loguru import logger
|
|
|
|
import flash_head
|
|
_FLASH_HEAD_ROOT = Path(flash_head.__path__[0])
|
|
|
|
from flash_head.src.pipeline.flash_head_pipeline import FlashHeadPipeline
|
|
from flash_head.inference import get_audio_embedding, run_pipeline
|
|
|
|
|
|
|
|
class iAskAsk_Local_FlashHead_Loader:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model_type": (["pro", "lite"], {"default": "lite"}),
|
|
},
|
|
"optional": {
|
|
"ckpt_dir": ("STRING", {"default": "", "multiline": False}),
|
|
"wav2vec_dir": ("STRING", {"default": "", "multiline": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("iAskAsk_FlashHead_Pipeline",)
|
|
RETURN_NAMES = ("FlashHead Pipeline",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "iAskAsk/FlashHead"
|
|
|
|
def load(self, model_type, ckpt_dir="", wav2vec_dir=""):
|
|
if not ckpt_dir:
|
|
ckpt_dir = os.path.join(folder_paths.models_dir, "Soul-AILab", "SoulX-FlashHead-1_3B")
|
|
if not wav2vec_dir:
|
|
wav2vec_dir = os.path.join(folder_paths.models_dir, "facebook", "wav2vec2-base-960h")
|
|
pipeline = FlashHeadPipeline(
|
|
checkpoint_dir=ckpt_dir,
|
|
model_type=model_type,
|
|
wav2vec_dir=wav2vec_dir,
|
|
device="cuda",
|
|
use_usp=False,
|
|
)
|
|
return (pipeline,)
|
|
|
|
|
|
class iAskAsk_Local_FlashHead_Sampler:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("iAskAsk_FlashHead_Pipeline",),
|
|
"ref_audio": ("AUDIO",),
|
|
"avatar_image": ("IMAGE",),
|
|
"seed": ("INT", {"default": 42, "min": 0, "max": 0xffffffffffffffff}),
|
|
},
|
|
"optional": {
|
|
"width": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 8}),
|
|
"height": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 8}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("VIDEO",)
|
|
RETURN_NAMES = ("video",)
|
|
FUNCTION = "sample"
|
|
CATEGORY = "iAskAsk/FlashHead"
|
|
|
|
def save_wav(self, ref_audio, save_path):
|
|
waveform = ref_audio["waveform"][0]
|
|
layout = "mono" if waveform.shape[0] == 1 else "stereo"
|
|
output_buffer = io.BytesIO()
|
|
output_container = av.open(output_buffer, mode="w", format="flac")
|
|
out_stream = output_container.add_stream("flac", rate=ref_audio["sample_rate"], layout=layout)
|
|
frame = av.AudioFrame.from_ndarray(
|
|
waveform.movedim(0, 1).reshape(1, -1).float().numpy(),
|
|
format="flt", layout=layout,
|
|
)
|
|
frame.sample_rate = ref_audio["sample_rate"]
|
|
frame.pts = 0
|
|
output_container.mux(out_stream.encode(frame))
|
|
output_container.mux(out_stream.encode(None))
|
|
output_container.close()
|
|
output_buffer.seek(0)
|
|
with open(save_path, "wb") as f:
|
|
f.write(output_buffer.getbuffer())
|
|
|
|
def save_avatar_image(self, avatar_image, save_path):
|
|
i = 255.0 * avatar_image.squeeze().cpu().numpy()
|
|
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
img.save(save_path)
|
|
|
|
def save_video(self, frames_list, video_path, audio_path, fps):
|
|
temp_video_path = video_path.replace("flashtalk_", "flashtalk_temp_")
|
|
with imageio.get_writer(
|
|
temp_video_path, format="mp4", mode="I",
|
|
fps=fps, codec="h264", ffmpeg_params=["-bf", "0"],
|
|
) as writer:
|
|
for frames in frames_list:
|
|
frames = frames.numpy().astype(np.uint8)
|
|
for i in range(frames.shape[0]):
|
|
writer.append_data(frames[i, :, :, :])
|
|
cmd = ["ffmpeg", "-i", temp_video_path, "-i", audio_path,
|
|
"-c:v", "copy", "-c:a", "aac", "-shortest", video_path, "-y"]
|
|
subprocess.run(cmd, check=True)
|
|
os.remove(temp_video_path)
|
|
|
|
def sample(self, **kwargs):
|
|
ref_audio = kwargs.get("ref_audio")
|
|
tmp_wav_path = os.path.join(
|
|
folder_paths.get_temp_directory(), f"flashtalk_audio_{uuid.uuid4()}.wav")
|
|
self.save_wav(ref_audio, tmp_wav_path)
|
|
|
|
avatar_image = kwargs.get("avatar_image")
|
|
tmp_img_path = os.path.join(
|
|
folder_paths.get_temp_directory(), f"flashtalk_avatar_{uuid.uuid4()}.png")
|
|
self.save_avatar_image(avatar_image, tmp_img_path)
|
|
|
|
infer_params_path = _FLASH_HEAD_ROOT / "configs" / "infer_params.yaml"
|
|
with open(infer_params_path, "r") as f:
|
|
infer_params = yaml.safe_load(f)
|
|
|
|
pipeline = kwargs.get("pipeline")
|
|
motion_frames_latent_num = infer_params["motion_frames_latent_num"]
|
|
motion_frames_num = (motion_frames_latent_num - 1) * pipeline.config.vae_stride[0] + 1
|
|
infer_params["motion_frames_num"] = motion_frames_num
|
|
|
|
base_seed = kwargs.get("seed")
|
|
width = kwargs.get("width", infer_params["width"])
|
|
height = kwargs.get("height", infer_params["height"])
|
|
infer_params["width"] = width
|
|
infer_params["height"] = height
|
|
|
|
infer_params["sample_steps"] = 4
|
|
|
|
pipeline.prepare_params(
|
|
cond_image_path_or_dir=tmp_img_path,
|
|
target_size=(infer_params["height"], infer_params["width"]),
|
|
frame_num=infer_params["frame_num"],
|
|
motion_frames_num=infer_params["motion_frames_num"],
|
|
sampling_steps=infer_params["sample_steps"],
|
|
seed=base_seed,
|
|
shift=infer_params["sample_shift"],
|
|
color_correction_strength=infer_params["color_correction_strength"],
|
|
use_face_crop=False,
|
|
)
|
|
|
|
sample_rate = infer_params["sample_rate"]
|
|
tgt_fps = infer_params["tgt_fps"]
|
|
cached_audio_duration = infer_params["cached_audio_duration"]
|
|
frame_num = infer_params["frame_num"]
|
|
motion_frames_num = infer_params["motion_frames_num"]
|
|
slice_len = frame_num - motion_frames_num
|
|
|
|
human_speech_array_all, _ = librosa.load(tmp_wav_path, sr=sample_rate, mono=True)
|
|
|
|
cached_audio_length_sum = sample_rate * cached_audio_duration
|
|
audio_end_idx = cached_audio_duration * tgt_fps
|
|
audio_start_idx = audio_end_idx - frame_num
|
|
|
|
audio_dq = deque([0.0] * cached_audio_length_sum, maxlen=cached_audio_length_sum)
|
|
|
|
human_speech_array_slice_len = slice_len * sample_rate // tgt_fps
|
|
aligned_len = (len(human_speech_array_all) // human_speech_array_slice_len) * human_speech_array_slice_len
|
|
human_speech_array_slices = human_speech_array_all[:aligned_len].reshape(-1, human_speech_array_slice_len)
|
|
|
|
generated_list = []
|
|
chunk_num = len(human_speech_array_slices)
|
|
self.pbar = comfy.utils.ProgressBar(chunk_num)
|
|
|
|
generator_ctx = torch.autocast('cuda', dtype=torch.bfloat16) if torch.cuda.get_device_capability()[0] >= 8 else torch.autocast('cuda', dtype=torch.float16)
|
|
|
|
for chunk_idx, human_speech_array in enumerate(human_speech_array_slices):
|
|
torch.cuda.synchronize()
|
|
audio_dq.extend(human_speech_array.tolist())
|
|
audio_array = np.array(audio_dq)
|
|
audio_embedding = get_audio_embedding(pipeline, audio_array, audio_start_idx, audio_end_idx)
|
|
with generator_ctx:
|
|
video = run_pipeline(pipeline, audio_embedding)
|
|
video = video[motion_frames_num:]
|
|
torch.cuda.synchronize()
|
|
generated_list.append(video.cpu())
|
|
self.pbar.update(1)
|
|
|
|
output_path = os.path.join(
|
|
folder_paths.get_output_directory(), f"flashtalk_video_{uuid.uuid4()}.mp4")
|
|
self.save_video(generated_list, output_path, tmp_wav_path, fps=tgt_fps)
|
|
return (self._create_video_object(output_path),)
|
|
|
|
def _create_video_object(self, video_path):
|
|
if VideoFromFile is not None:
|
|
return VideoFromFile(video_path)
|
|
return video_path
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"iAskAsk Local FlashHead Loader": iAskAsk_Local_FlashHead_Loader,
|
|
"iAskAsk Local FlashHead Sampler": iAskAsk_Local_FlashHead_Sampler,
|
|
}
|