HM-RunningHub-ComfyUI_RH_FlashHead/nodes.py

import comfy.utils

import folder_paths
try:
    from comfy_api.input_impl.video_types import VideoFromFile
except ImportError:
    VideoFromFile = None
from pathlib import Path
from optimum.quanto import freeze, qint8, quantize 
import uuid

import os
import numpy as np
import time
import torch
import torch.distributed as dist
import subprocess
import imageio
import librosa
import numpy as np
from loguru import logger
from collections import deque
from datetime import datetime
from PIL import Image
import io
import av
import yaml

# from .flash_head.inference import get_pipeline, get_base_data, get_infer_params, get_audio_embedding, run_pipeline
from .flash_head.src.pipeline.flash_head_pipeline import FlashHeadPipeline
from .flash_head.inference import get_audio_embedding, run_pipeline

class RunningHub_FlashHead_Loader:

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model_type": (["pro", "lite"], {"default": "lite"}),
            }
        }

    RETURN_TYPES = ('RH_FlashHead_Pipeline', )
    RETURN_NAMES = ('FlashHead Pipeline', )
    FUNCTION = "load"
    CATEGORY = "RunningHub/FlashHead"

    def load(self, model_type):

        ckpt_dir = os.path.join(folder_paths.models_dir, 'Soul-AILab', 'SoulX-FlashHead-1_3B')
        wav2vec_dir = os.path.join(folder_paths.models_dir, 'wav2vec', '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 RunningHub_FlashHead_Sampler:

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "pipeline": ("RH_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 = "RunningHub/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))

        # Flush encoder
        output_container.mux(out_stream.encode(None))

        # Close containers
        output_container.close()

        # Write the output to file
        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. * 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]):
                    frame = frames[i, :, :, :]
                    writer.append_data(frame)
        
        # merge video and audio
        # cmd = ['ffmpeg', '-i', temp_video_path, '-i', audio_path, '-c', 'copy', '-shortest', video_path, '-y']
        cmd = ['ffmpeg', '-i', temp_video_path, '-i', audio_path, '-c:v', 'copy', '-c:a', 'aac', '-shortest', video_path, '-y']
        subprocess.run(cmd)
        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_avatar_image_path = os.path.join(folder_paths.get_temp_directory(), f"flashtalk_avatar_image_{uuid.uuid4()}.png")
        self.save_avatar_image(avatar_image, tmp_avatar_image_path)

        infer_params_path = Path(__file__).resolve().parent / "flash_head" / "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') ^ (2 ** 32)
        width = kwargs.get('width', infer_params['width'])
        height = kwargs.get('height', infer_params['height'])
        infer_params['width'] = width
        infer_params['height'] = height

        #kiki
        use_face_crop = False

        # TODO: move to args
        if pipeline.model_type == "pretrained":
            infer_params['sample_steps'] = 20
        else:
            infer_params['sample_steps'] = 4

        pipeline.prepare_params(
            cond_image_path_or_dir=tmp_avatar_image_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=use_face_crop,
        )

        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=infer_params['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
        human_speech_array_slices = human_speech_array_all[:(len(human_speech_array_all)//(human_speech_array_slice_len))*human_speech_array_slice_len].reshape(-1, human_speech_array_slice_len)

        generated_list = []
        chunk_num = len(human_speech_array_slices)
        self.pbar = comfy.utils.ProgressBar(chunk_num)
        for chunk_idx, human_speech_array in enumerate(human_speech_array_slices):
            torch.cuda.synchronize()
            start_time = time.time()

            # streaming encode audio chunks
            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, infer_params=infer_params)

            # inference
            video = run_pipeline(pipeline, audio_embedding)

            torch.cuda.synchronize()
            end_time = time.time()
            logger.info(f"Generate video chunk-{chunk_idx} done, cost time: {(end_time - start_time):.3f}s")

            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):
        """Create ComfyUI VIDEO object"""
        if VideoFromFile is not None:
            return VideoFromFile(video_path)
        else:
            # Fallback: return file path as string
            return video_path

    def update(self):
        self.pbar.update(1)

NODE_CLASS_MAPPINGS = {
    "RunningHub SoulX-FlashHead Loader": RunningHub_FlashHead_Loader,
    "RunningHub SoulX-FlashHead Sampler": RunningHub_FlashHead_Sampler,
}