magicwang1111-ComfyUI-Keysync-wrapper/keysync_wrapper/media_utils.py

import subprocess
from pathlib import Path

import cv2
import numpy as np
import soundfile as sf
import torch

from .ffmpeg_utils import check_ffmpeg
from .file_utils import resolve_path


def _as_image_tensor(images):
    if isinstance(images, list):
        images = torch.stack(images)
    if not isinstance(images, torch.Tensor):
        images = torch.as_tensor(images)
    if images.dim() != 4:
        raise ValueError(f"Expected IMAGE tensor with shape [frames, height, width, channels], got {tuple(images.shape)}")
    if images.shape[-1] < 3:
        raise ValueError(f"Expected IMAGE tensor with at least 3 channels, got {tuple(images.shape)}")
    return images.detach().cpu()


def _frame_to_uint8(frame):
    frame = frame[..., :3]
    if frame.dtype != torch.uint8:
        frame = frame.float()
        if float(frame.max()) <= 1.0:
            frame = frame * 255.0
        frame = frame.clamp(0, 255).byte()
    return np.ascontiguousarray(frame.numpy())


def save_images_to_video(images, output_path, fps):
    image_tensor = _as_image_tensor(images)
    output_file = resolve_path(output_path)
    output_file.parent.mkdir(parents=True, exist_ok=True)

    frame_count, height, width, _ = image_tensor.shape
    if frame_count <= 0:
        raise ValueError("No frames were provided to save_images_to_video.")

    ffmpeg = check_ffmpeg()
    command = [
        ffmpeg,
        "-y",
        "-f",
        "rawvideo",
        "-vcodec",
        "rawvideo",
        "-s",
        f"{width}x{height}",
        "-pix_fmt",
        "rgb24",
        "-r",
        str(fps),
        "-i",
        "-",
        "-an",
        "-c:v",
        "libx264",
        "-pix_fmt",
        "yuv420p",
        "-crf",
        "18",
        str(output_file),
    ]

    process = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    try:
        for frame in image_tensor:
            process.stdin.write(_frame_to_uint8(frame).tobytes())
        stdout, stderr = process.communicate()
    except BrokenPipeError:
        stdout, stderr = process.communicate()
    finally:
        if process.stdin and not process.stdin.closed:
            process.stdin.close()

    if process.returncode != 0:
        message = stderr.decode("utf-8", errors="replace").strip()
        raise RuntimeError(f"Failed to encode IMAGE frames to mp4: {message}")

    return output_file


def _audio_waveform_to_tensor(audio):
    if not isinstance(audio, dict) or "waveform" not in audio or "sample_rate" not in audio:
        raise ValueError("AUDIO input must be a dict with 'waveform' and 'sample_rate'.")

    waveform = audio["waveform"]
    if not isinstance(waveform, torch.Tensor):
        waveform = torch.as_tensor(waveform)
    waveform = waveform.detach().cpu().float()

    if waveform.dim() == 3:
        waveform = waveform[0]
    elif waveform.dim() == 1:
        waveform = waveform.unsqueeze(0)
    elif waveform.dim() != 2:
        raise ValueError(f"Expected AUDIO waveform with shape [batch, channels, samples] or [channels, samples], got {tuple(waveform.shape)}")

    sample_rate = int(audio["sample_rate"])
    return waveform, sample_rate


def trim_or_pad_audio(audio, duration_sec=None, pad_to_duration=False):
    waveform, sample_rate = _audio_waveform_to_tensor(audio)

    if duration_sec is not None and duration_sec > 0:
        target_samples = max(1, int(round(float(duration_sec) * sample_rate)))
        if waveform.shape[-1] > target_samples:
            waveform = waveform[:, :target_samples]
        elif pad_to_duration and waveform.shape[-1] < target_samples:
            pad = torch.zeros((waveform.shape[0], target_samples - waveform.shape[-1]), dtype=waveform.dtype)
            waveform = torch.cat([waveform, pad], dim=-1)

    return {"waveform": waveform.unsqueeze(0).contiguous(), "sample_rate": sample_rate}


def save_audio_to_wav(audio, output_path, duration_sec=None, pad_to_duration=False):
    adjusted = trim_or_pad_audio(audio, duration_sec=duration_sec, pad_to_duration=pad_to_duration)
    waveform = adjusted["waveform"].squeeze(0).clamp(-1.0, 1.0)
    sample_rate = int(adjusted["sample_rate"])

    output_file = resolve_path(output_path)
    output_file.parent.mkdir(parents=True, exist_ok=True)
    audio_np = waveform.transpose(0, 1).numpy()
    sf.write(str(output_file), audio_np, sample_rate, subtype="PCM_16")
    return output_file


def load_video_to_images(video_path):
    video_file = resolve_path(video_path)
    capture = cv2.VideoCapture(str(video_file))
    if not capture.isOpened():
        raise RuntimeError(f"Failed to open output video: {video_file}")

    fps = capture.get(cv2.CAP_PROP_FPS)
    frames = []
    try:
        while True:
            ok, frame_bgr = capture.read()
            if not ok:
                break
            frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
            frames.append(frame_rgb)
    finally:
        capture.release()

    if not frames:
        raise RuntimeError(f"No frames were decoded from output video: {video_file}")

    array = np.stack(frames).astype(np.float32) / 255.0
    return torch.from_numpy(array), int(round(fps or 25))