Add ComfyUI HighSync nodes

highsync_wrapper/__init__.py

@@ -0,0 +1 @@
+"""ComfyUI adapter utilities for HighSync."""

highsync_wrapper/download_utils.py

@@ -0,0 +1,79 @@
+from pathlib import Path
+
+from .file_utils import normalize_path, resolve_path
+
+
+CORE_FILES = (
+    "denoising_unet-500.pth",
+    "reference_unet-500.pth",
+    "audio_processor/whisper_tiny.pt",
+)
+
+ALLOW_PATTERNS = (
+    "denoising_unet-500.pth",
+    "reference_unet-500.pth",
+    "audio_processor/whisper_tiny.pt",
+    "sd-vae-ft-mse/**",
+    "sd-image-variations-diffusers/**",
+)
+
+
+def _directory_has_model_files(path):
+    if not path.is_dir():
+        return False
+    return any(candidate.is_file() for candidate in path.rglob("*"))
+
+
+def validate_required_models(models_dir):
+    root = resolve_path(models_dir)
+    missing = []
+
+    for rel_path in CORE_FILES:
+        if not (root / rel_path).is_file():
+            missing.append(rel_path)
+
+    vae_dir = root / "sd-vae-ft-mse"
+    if not _directory_has_model_files(vae_dir) or not (vae_dir / "config.json").is_file():
+        missing.append("sd-vae-ft-mse/config.json")
+
+    base_dir = root / "sd-image-variations-diffusers"
+    if not _directory_has_model_files(base_dir) or not (base_dir / "unet" / "config.json").is_file():
+        missing.append("sd-image-variations-diffusers/unet/config.json")
+
+    if missing:
+        missing_text = ", ".join(missing)
+        raise FileNotFoundError(
+            f"HighSync model files are missing in {root}: {missing_text}. "
+            "Run the HighSync Download Models node with download=True, or download "
+            "https://huggingface.co/saeed-5959/high_sync into this folder."
+        )
+
+    return root
+
+
+def download_or_check_models(models_dir, repo_id="saeed-5959/high_sync", download=False):
+    root = Path(normalize_path(models_dir))
+
+    if download:
+        try:
+            from huggingface_hub import snapshot_download
+        except ImportError as exc:
+            raise RuntimeError(
+                "huggingface_hub is required for HighSync Download Models. "
+                "Install this wrapper's requirements.txt first."
+            ) from exc
+
+        root.mkdir(parents=True, exist_ok=True)
+        snapshot_download(
+            repo_id=repo_id,
+            local_dir=str(root),
+            allow_patterns=list(ALLOW_PATTERNS),
+            repo_type="model",
+        )
+    elif not root.exists():
+        raise FileNotFoundError(
+            f"models_dir does not exist: {root}. "
+            "Run HighSync Download Models with download=True or download the models manually."
+        )
+
+    return str(validate_required_models(root))

highsync_wrapper/ffmpeg_utils.py

@@ -0,0 +1,45 @@
+import shutil
+import subprocess
+
+from .file_utils import resolve_path
+
+
+def check_ffmpeg():
+    ffmpeg = shutil.which("ffmpeg") or shutil.which("ffmpeg.exe")
+    if ffmpeg is None:
+        raise RuntimeError("ffmpeg not found. Please install ffmpeg and add it to PATH.")
+    return ffmpeg
+
+
+def run_ffmpeg(args, error_label):
+    command = [check_ffmpeg()] + [str(arg) for arg in args]
+    result = subprocess.run(command, capture_output=True, text=True)
+    if result.returncode != 0:
+        stderr = result.stderr.strip()
+        raise RuntimeError(f"{error_label} failed with exit code {result.returncode}: {stderr}")
+
+
+def mux_video_audio(video_path, audio_path, output_path):
+    output_file = resolve_path(output_path)
+    output_file.parent.mkdir(parents=True, exist_ok=True)
+    run_ffmpeg(
+        [
+            "-y",
+            "-i",
+            resolve_path(video_path),
+            "-i",
+            resolve_path(audio_path),
+            "-map",
+            "0:v:0",
+            "-map",
+            "1:a:0",
+            "-c:v",
+            "copy",
+            "-c:a",
+            "aac",
+            "-shortest",
+            output_file,
+        ],
+        "audio/video mux",
+    )
+    return output_file

highsync_wrapper/file_utils.py

@@ -0,0 +1,92 @@
+import os
+import re
+import shutil
+import uuid
+from datetime import datetime
+from pathlib import Path
+
+
+WINDOWS_DRIVE_RE = re.compile(r"^([A-Za-z]):[\\/](.*)$")
+
+
+def is_windows_drive_path(path):
+    return isinstance(path, str) and WINDOWS_DRIVE_RE.match(path.strip()) is not None
+
+
+def normalize_path(path):
+    if path is None:
+        return path
+
+    raw_path = str(path).strip().strip('"')
+    if raw_path == "":
+        return raw_path
+
+    match = WINDOWS_DRIVE_RE.match(raw_path)
+    if match and os.name != "nt":
+        drive = match.group(1).lower()
+        rest = match.group(2).replace("\\", "/")
+        return f"/mnt/{drive}/{rest}"
+
+    return os.path.abspath(os.path.expanduser(raw_path))
+
+
+def resolve_path(path):
+    return Path(normalize_path(path)).expanduser().resolve()
+
+
+def ensure_exists(path, label="path", is_dir=None):
+    if path is None or str(path).strip() == "":
+        raise FileNotFoundError(f"{label} is empty.")
+
+    resolved = resolve_path(path)
+    if is_dir is True and not resolved.is_dir():
+        raise FileNotFoundError(f"{label} does not exist: {resolved}")
+    if is_dir is False and not resolved.is_file():
+        raise FileNotFoundError(f"{label} does not exist: {resolved}")
+    if is_dir is None and not resolved.exists():
+        raise FileNotFoundError(f"{label} does not exist: {resolved}")
+    return resolved
+
+
+def safe_copy(src, dst):
+    src_path = ensure_exists(src, "source file", is_dir=False)
+    dst_path = resolve_path(dst)
+    dst_path.parent.mkdir(parents=True, exist_ok=True)
+    if src_path.resolve() != dst_path.resolve():
+        shutil.copy2(src_path, dst_path)
+    return dst_path
+
+
+def get_comfy_output_dir():
+    try:
+        import folder_paths
+
+        return Path(folder_paths.get_output_directory()).resolve()
+    except Exception:
+        current = Path(__file__).resolve()
+        for parent in current.parents:
+            if (parent / "custom_nodes").exists() and (parent / "output").exists():
+                return (parent / "output").resolve()
+        return Path("/mnt/d/ComfyUI/output").resolve()
+
+
+def create_task_dir(prefix="highsync"):
+    output_root = get_comfy_output_dir() / "highsync"
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    suffix = uuid.uuid4().hex[:8]
+    task_id = f"{prefix}_{timestamp}_{suffix}"
+    task_dir = output_root / task_id
+
+    paths = {
+        "task_id": task_id,
+        "task_dir": task_dir,
+        "media_dir": task_dir / "media",
+        "logs_dir": task_dir / "logs",
+        "frames_dir": task_dir / "frames",
+    }
+
+    for key, value in paths.items():
+        if key.endswith("_dir"):
+            value.mkdir(parents=True, exist_ok=True)
+
+    return paths

highsync_wrapper/media_utils.py

@@ -0,0 +1,224 @@
+import subprocess
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import torch
+
+from .ffmpeg_utils import check_ffmpeg
+from .file_utils import resolve_path
+
+
+TARGET_FPS = 25
+TARGET_AUDIO_SR = 16000
+
+
+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 images_to_uint8_rgb(images):
+    image_tensor = _as_image_tensor(images)[..., :3]
+    if image_tensor.dtype != torch.uint8:
+        image_tensor = image_tensor.float()
+        if image_tensor.numel() > 0 and float(image_tensor.max()) <= 1.0:
+            image_tensor = image_tensor * 255.0
+        image_tensor = image_tensor.clamp(0, 255).byte()
+    return np.ascontiguousarray(image_tensor.numpy())
+
+
+def frames_uint8_rgb_to_image_tensor(frames):
+    if len(frames) == 0:
+        raise ValueError("No output frames were generated.")
+    array = np.ascontiguousarray(np.stack(frames).astype(np.float32) / 255.0)
+    return torch.from_numpy(array)
+
+
+def resample_frames_to_fps(frames_rgb, input_fps, output_fps=TARGET_FPS, max_frames=0):
+    input_fps = int(input_fps)
+    output_fps = int(output_fps)
+    if input_fps <= 0 or output_fps <= 0:
+        raise ValueError("input_fps and output_fps must be positive.")
+    if len(frames_rgb) == 0:
+        raise ValueError("No IMAGE frames were provided.")
+
+    if input_fps == output_fps:
+        sampled = frames_rgb.copy()
+    else:
+        duration = len(frames_rgb) / float(input_fps)
+        target_count = max(1, int(round(duration * output_fps)))
+        indices = np.floor(np.arange(target_count) * input_fps / output_fps).astype(np.int64)
+        indices = np.clip(indices, 0, len(frames_rgb) - 1)
+        sampled = frames_rgb[indices]
+
+    max_frames = int(max_frames or 0)
+    if max_frames > 0:
+        sampled = sampled[:max_frames]
+
+    return np.ascontiguousarray(sampled)
+
+
+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(
+            "Expected AUDIO waveform with shape [batch, channels, samples], "
+            f"[channels, samples], or [samples], got {tuple(waveform.shape)}"
+        )
+
+    return waveform.contiguous(), int(audio["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 _resample_waveform(waveform, sample_rate, target_sample_rate):
+    if int(sample_rate) == int(target_sample_rate):
+        return waveform
+
+    try:
+        import torchaudio
+
+        resampler = torchaudio.transforms.Resample(
+            orig_freq=int(sample_rate),
+            new_freq=int(target_sample_rate),
+        )
+        return resampler(waveform)
+    except Exception as exc:
+        raise RuntimeError(
+            "Failed to resample audio with torchaudio. Install torchaudio or provide 16kHz audio."
+        ) from exc
+
+
+def save_audio_to_wav(audio, output_path, duration_sec=None, pad_to_duration=False, mono=False, sample_rate=None):
+    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)
+    source_sample_rate = int(adjusted["sample_rate"])
+
+    if mono and waveform.shape[0] > 1:
+        waveform = waveform.mean(dim=0, keepdim=True)
+
+    target_sample_rate = int(sample_rate or source_sample_rate)
+    waveform = _resample_waveform(waveform, source_sample_rate, target_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, target_sample_rate, subtype="PCM_16")
+    return output_file
+
+
+def encode_rgb_frames_to_mp4(frames_rgb, output_path, fps=TARGET_FPS):
+    output_file = resolve_path(output_path)
+    output_file.parent.mkdir(parents=True, exist_ok=True)
+
+    frames = np.ascontiguousarray(frames_rgb)
+    if frames.ndim != 4 or frames.shape[-1] != 3:
+        raise ValueError(f"Expected RGB frames with shape [frames, height, width, 3], got {frames.shape}")
+
+    frame_count, height, width, _ = frames.shape
+    if frame_count <= 0:
+        raise ValueError("No frames were provided to encode.")
+
+    command = [
+        check_ffmpeg(),
+        "-y",
+        "-f",
+        "rawvideo",
+        "-vcodec",
+        "rawvideo",
+        "-s",
+        f"{width}x{height}",
+        "-pix_fmt",
+        "rgb24",
+        "-r",
+        str(int(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 frames:
+            process.stdin.write(np.ascontiguousarray(frame).tobytes())
+        _, stderr = process.communicate()
+    except BrokenPipeError:
+        _, 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 HighSync output video: {message}")
+
+    return output_file
+
+
+def denoise_audio_file(input_audio_path, output_audio_path, device, cache):
+    try:
+        from denoiser import pretrained
+        from denoiser.dsp import convert_audio
+    except ImportError as exc:
+        raise RuntimeError(
+            "denoiser and torchaudio are required when denoise_audio=True. "
+            "Install this wrapper's requirements.txt or disable denoise_audio."
+        ) from exc
+
+    model = cache.get("denoiser_model")
+    model_device = torch.device(device)
+    if model is None:
+        model = pretrained.dns64().to(model_device)
+        model.eval()
+        cache["denoiser_model"] = model
+
+    audio_np, sample_rate = sf.read(str(Path(input_audio_path)), dtype="float32", always_2d=True)
+    wav = torch.from_numpy(audio_np.T.copy())
+    wav = convert_audio(wav.to(model_device), sample_rate, model.sample_rate, model.chin)
+
+    with torch.no_grad():
+        denoised = model(wav[None])[0]
+
+    output_file = resolve_path(output_audio_path)
+    output_file.parent.mkdir(parents=True, exist_ok=True)
+    audio_np = denoised.detach().cpu().clamp(-1.0, 1.0).numpy()[0]
+    sf.write(str(output_file), audio_np, int(model.sample_rate), subtype="PCM_16")
+    return output_file

highsync_wrapper/model_utils.py

@@ -0,0 +1,148 @@
+from pathlib import Path
+
+from .download_utils import validate_required_models
+
+
+_MODEL_CACHE = {}
+
+
+def _get_repo_root():
+    return Path(__file__).resolve().parents[1]
+
+
+def _torch_load(path, torch):
+    try:
+        return torch.load(path, map_location="cpu", weights_only=True)
+    except TypeError:
+        return torch.load(path, map_location="cpu")
+
+
+def _get_torch_device(torch):
+    try:
+        import comfy.model_management as model_management
+
+        return model_management.get_torch_device()
+    except Exception:
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def _weight_dtype(torch, weight_dtype):
+    dtype_text = str(weight_dtype or "fp16").lower()
+    if dtype_text == "fp16":
+        return torch.float16
+    if dtype_text == "bf16":
+        return torch.bfloat16
+    if dtype_text == "fp32":
+        return torch.float32
+    raise ValueError(f"Unsupported HighSync weight_dtype: {weight_dtype}")
+
+
+def _load_config(models_root, weight_dtype):
+    try:
+        from omegaconf import OmegaConf
+    except ImportError as exc:
+        raise RuntimeError("omegaconf is required to load HighSync config. Install requirements.txt first.") from exc
+
+    config_path = _get_repo_root() / "high_sync" / "configs" / "inference" / "default.yaml"
+    config = OmegaConf.load(str(config_path))
+
+    config.weight_dtype = str(weight_dtype)
+    config.denoising_unet_path = str(models_root / "denoising_unet-500.pth")
+    config.reference_unet_path = str(models_root / "reference_unet-500.pth")
+    config.base_model_path = str(models_root / "sd-image-variations-diffusers")
+    config.audio_model_path = str(models_root / "audio_processor" / "whisper_tiny.pt")
+    config.vae.model_path = str(models_root / "sd-vae-ft-mse")
+    config.mask_path = str(_get_repo_root() / "high_sync" / "masks" / "mask_64.png")
+    config.mask_512_path = str(_get_repo_root() / "high_sync" / "masks" / "mask_512.png")
+
+    return config
+
+
+def _build_scheduler(config):
+    try:
+        from diffusers import DDIMScheduler
+        from omegaconf import OmegaConf
+    except ImportError as exc:
+        raise RuntimeError("diffusers and omegaconf are required for HighSync. Install requirements.txt first.") from exc
+
+    sched_kwargs = OmegaConf.to_container(config.noise_scheduler_kwargs, resolve=True)
+    if bool(config.enable_zero_snr):
+        sched_kwargs.update(
+            rescale_betas_zero_snr=True,
+            timestep_spacing="trailing",
+            prediction_type="v_prediction",
+        )
+    return DDIMScheduler(**sched_kwargs)
+
+
+def load_highsync_model(models_dir, weight_dtype="fp16"):
+    try:
+        import torch
+        from diffusers import AutoencoderKL
+    except ImportError as exc:
+        raise RuntimeError(
+            "torch and diffusers are required to load HighSync. Install this wrapper's requirements.txt first."
+        ) from exc
+
+    models_root = validate_required_models(models_dir)
+    device = _get_torch_device(torch)
+    dtype = _weight_dtype(torch, weight_dtype)
+    cache_key = (str(models_root), str(weight_dtype).lower(), str(device))
+
+    if cache_key in _MODEL_CACHE:
+        return _MODEL_CACHE[cache_key]
+
+    config = _load_config(models_root, weight_dtype)
+
+    from omegaconf import OmegaConf
+
+    from high_sync.image_processor import ImageProcessor
+    from high_sync.src.models.unet_2d_condition import UNet2DConditionModel
+    from high_sync.src.models.unet_3d_echo import EchoUNet3DConditionModel
+    from high_sync.src.models.whisper.audio2feature import load_audio_model
+    from high_sync.src.pipelines.face_animate import FaceAnimatePipeline
+
+    scheduler = _build_scheduler(config)
+
+    vae = AutoencoderKL.from_pretrained(str(config.vae.model_path))
+    reference_unet = UNet2DConditionModel.from_pretrained(str(config.base_model_path), subfolder="unet")
+    denoising_unet = EchoUNet3DConditionModel.from_pretrained_2d(
+        str(config.base_model_path),
+        "",
+        subfolder="unet",
+        unet_additional_kwargs=OmegaConf.to_container(config.unet_additional_kwargs, resolve=True),
+    )
+
+    reference_unet.load_state_dict(_torch_load(str(config.reference_unet_path), torch))
+    denoising_unet.load_state_dict(_torch_load(str(config.denoising_unet_path), torch))
+
+    vae.requires_grad_(False)
+    reference_unet.requires_grad_(False)
+    denoising_unet.requires_grad_(False)
+
+    vae.eval()
+    reference_unet.eval()
+    denoising_unet.eval()
+
+    pipeline = FaceAnimatePipeline(
+        vae=vae,
+        reference_unet=reference_unet,
+        denoising_unet=denoising_unet,
+        scheduler=scheduler,
+    )
+    pipeline.to(device=device, dtype=dtype)
+
+    audio_processor = load_audio_model(model_path=str(config.audio_model_path), device=device)
+
+    model = {
+        "models_dir": str(models_root),
+        "device": device,
+        "weight_dtype": dtype,
+        "config": config,
+        "pipeline": pipeline,
+        "audio_processor": audio_processor,
+        "image_processor_cls": ImageProcessor,
+        "denoiser_model": None,
+    }
+    _MODEL_CACHE[cache_key] = model
+    return model

highsync_wrapper/runner.py

@@ -0,0 +1,323 @@
+import shutil
+
+import cv2
+import numpy as np
+import torch
+
+from .ffmpeg_utils import mux_video_audio
+from .file_utils import create_task_dir
+from .media_utils import (
+    TARGET_AUDIO_SR,
+    TARGET_FPS,
+    denoise_audio_file,
+    encode_rgb_frames_to_mp4,
+    frames_uint8_rgb_to_image_tensor,
+    images_to_uint8_rgb,
+    resample_frames_to_fps,
+    save_audio_to_wav,
+    trim_or_pad_audio,
+)
+
+
+def _maybe_throw_if_interrupted():
+    try:
+        import comfy.model_management as model_management
+    except Exception:
+        return
+    model_management.throw_exception_if_processing_interrupted()
+
+
+def _empty_cache(device):
+    if isinstance(device, torch.device) and device.type == "cuda" and torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    try:
+        import comfy.model_management as model_management
+
+        model_management.soft_empty_cache()
+    except Exception:
+        return
+
+
+def create_mask(mask_path):
+    raw_mask = cv2.imread(str(mask_path))
+    if raw_mask is None:
+        raise FileNotFoundError(f"HighSync mask could not be read: {mask_path}")
+    mask = 1 - (raw_mask / 255.0)
+
+    h, w, c = mask.shape
+    new_mask = np.zeros((h, w, c + 1), dtype=np.float32)
+    new_mask[:, :, :c] = mask
+    new_mask[:, :, c] = mask[:, :, c - 1]
+    return torch.FloatTensor(new_mask).permute(2, 0, 1).unsqueeze(1)
+
+
+def create_mask_alpha(width, height, coordinate):
+    left, top, right, bottom = coordinate
+    mask = np.ones((height, width, 3), dtype=np.float32)
+    mask[top:bottom, left:right] = 0
+
+    pad = min(int((right - left) / 15), int((bottom - top) / 15))
+    if pad <= 0:
+        return mask
+
+    for idx in range(left, right):
+        mask[bottom - pad:bottom, idx] = np.array([np.linspace(0, 1, pad)] * 3).transpose(1, 0)
+        mask[top:top + pad, idx] = np.array([np.linspace(1, 0, pad)] * 3).transpose(1, 0)
+
+    for idx in range(top, bottom):
+        mask[idx, right - pad:right] = np.array([np.linspace(0, 1, pad)] * 3).transpose(1, 0)
+        mask[idx, left:left + pad] = np.array([np.linspace(1, 0, pad)] * 3).transpose(1, 0)
+
+    for idx, start, end in zip(range(top, top + pad), np.linspace(1, 0.7, pad), np.linspace(0.7, 0, pad)):
+        mask[idx, left:left + pad] = np.array([np.linspace(start, end, pad)] * 3).transpose(1, 0)
+        mask[idx, right - pad:right] = np.array([np.linspace(end, start, pad)] * 3).transpose(1, 0)
+
+    for idx, start, end in zip(range(bottom - pad, bottom), np.linspace(0.7, 1, pad), np.linspace(0, 0.7, pad)):
+        mask[idx, left:left + pad] = np.array([np.linspace(start, end, pad)] * 3).transpose(1, 0)
+        mask[idx, right - pad:right] = np.array([np.linspace(end, start, pad)] * 3).transpose(1, 0)
+
+    return mask
+
+
+def _prepare_audio_files(task, audio, duration_sec, trim_audio_to_video, denoise_audio, device, model_cache):
+    audio_for_video = trim_or_pad_audio(
+        audio,
+        duration_sec=duration_sec if trim_audio_to_video else None,
+        pad_to_duration=bool(trim_audio_to_video),
+    )
+
+    mux_audio_path = save_audio_to_wav(
+        audio_for_video,
+        task["media_dir"] / "mux_audio.wav",
+    )
+    feature_audio_path = save_audio_to_wav(
+        audio_for_video,
+        task["media_dir"] / "feature_audio_16k.wav",
+        mono=True,
+        sample_rate=TARGET_AUDIO_SR,
+    )
+
+    if denoise_audio:
+        feature_audio_path = denoise_audio_file(
+            feature_audio_path,
+            task["media_dir"] / "feature_audio_16k_denoised.wav",
+            device=device,
+            cache=model_cache,
+        )
+
+    return audio_for_video, mux_audio_path, feature_audio_path
+
+
+def _ensure_highsync_model(highsync_model):
+    required = ("config", "pipeline", "audio_processor", "image_processor_cls", "device", "weight_dtype")
+    if not isinstance(highsync_model, dict) or any(key not in highsync_model for key in required):
+        raise ValueError("highsync_model must come from the HighSync Model Loader node.")
+
+
+def _append_original_frame(output_frames_rgb, frame_bgr):
+    output_frames_rgb.append(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB))
+
+
+def _run_generation_chunk(
+    pipeline,
+    image_processor,
+    config,
+    audio_fea_final,
+    frames_video,
+    output_frames_rgb,
+    count_img,
+    generator,
+    device,
+    weight_dtype,
+    inference_steps,
+    cfg_scale,
+):
+    source_image_pixels_all, bbox_all, frame_list_all = image_processor.preprocess_frames_jump(frames_video)
+
+    img_size = (int(config.data.source_image.width), int(config.data.source_image.height))
+    sample_frames = int(config.data.n_sample_frames)
+    mask_path = str(config.mask_path)
+
+    for source_image_pixels_list, bbox_list, frame_list in zip(source_image_pixels_all, bbox_all, frame_list_all):
+        _maybe_throw_if_interrupted()
+
+        if not frame_list:
+            continue
+
+        if bbox_list[0] != (0, 0, 0, 0):
+            audio_tensor = audio_fea_final[:, count_img: count_img + len(frame_list)]
+            main_length = len(frame_list)
+            diff = sample_frames - main_length
+
+            if diff > 0:
+                source_image_pixels_list.extend([source_image_pixels_list[-1]] * diff)
+                audio_tensor = torch.cat((audio_tensor, audio_tensor[:, -1:].repeat(1, diff, 1, 1)), dim=1)
+
+            source_image_pixels = torch.cat(source_image_pixels_list, dim=0)
+            mask_main = create_mask(mask_path).unsqueeze(0).to(device=device, dtype=weight_dtype)
+            _empty_cache(device)
+
+            with torch.no_grad():
+                pipeline_output = pipeline(
+                    gt_encode=source_image_pixels,
+                    audio_tensor=audio_tensor,
+                    mask_main=mask_main,
+                    width=img_size[0],
+                    height=img_size[1],
+                    video_length=len(source_image_pixels_list),
+                    num_inference_steps=int(inference_steps),
+                    guidance_scale=float(cfg_scale),
+                    generator=generator,
+                )
+
+            _empty_cache(device)
+
+            images_out = pipeline_output.videos.squeeze(0).permute(1, 2, 3, 0).cpu().numpy()[:main_length]
+            images_out = np.clip(images_out * 255, 0, 255).astype(np.uint8)
+
+            for image_out, bbox, frame in zip(images_out, bbox_list, frame_list):
+                frame_h, frame_w, _ = frame.shape
+                if bbox == (0, 0, 0, 0):
+                    _append_original_frame(output_frames_rgb, frame)
+                else:
+                    image_out_bgr = cv2.cvtColor(image_out, cv2.COLOR_RGB2BGR)
+                    frame_in_out = frame.copy()
+                    frame_in_out[bbox[1]:bbox[3], bbox[0]:bbox[2]] = cv2.resize(
+                        image_out_bgr,
+                        (bbox[2] - bbox[0], bbox[3] - bbox[1]),
+                    )
+                    mask = create_mask_alpha(frame_w, frame_h, bbox)
+                    blended = (frame * mask + frame_in_out * (1 - mask)).astype(np.uint8)
+                    _append_original_frame(output_frames_rgb, blended)
+
+                count_img += 1
+        else:
+            for frame in frame_list:
+                _append_original_frame(output_frames_rgb, frame)
+                count_img += 1
+
+    return count_img
+
+
+def run_highsync(
+    highsync_model,
+    images,
+    audio,
+    input_fps=25,
+    max_frames=0,
+    inference_steps=None,
+    cfg_scale=None,
+    seed=42,
+    preprocess_chunk_frames=1500,
+    denoise_audio=True,
+    trim_audio_to_video=True,
+    keep_intermediate_outputs=True,
+):
+    _ensure_highsync_model(highsync_model)
+
+    config = highsync_model["config"]
+    pipeline = highsync_model["pipeline"]
+    audio_processor = highsync_model["audio_processor"]
+    image_processor = highsync_model["image_processor_cls"]((int(config.data.source_image.width), int(config.data.source_image.height)))
+    device = torch.device(highsync_model["device"])
+    weight_dtype = highsync_model["weight_dtype"]
+
+    inference_steps = int(inference_steps if inference_steps is not None else config.inference_steps)
+    cfg_scale = float(cfg_scale if cfg_scale is not None else config.cfg_scale)
+    preprocess_chunk_frames = max(12, int(preprocess_chunk_frames))
+
+    task = create_task_dir()
+
+    frames_rgb = images_to_uint8_rgb(images)
+    frames_rgb = resample_frames_to_fps(frames_rgb, input_fps=input_fps, output_fps=TARGET_FPS, max_frames=max_frames)
+    frames_bgr = [cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) for frame in frames_rgb]
+    initial_duration_sec = len(frames_bgr) / float(TARGET_FPS)
+
+    _, _, feature_audio_path = _prepare_audio_files(
+        task=task,
+        audio=audio,
+        duration_sec=initial_duration_sec,
+        trim_audio_to_video=trim_audio_to_video,
+        denoise_audio=denoise_audio,
+        device=device,
+        model_cache=highsync_model,
+    )
+
+    whisper_feature = audio_processor.audio2feat(str(feature_audio_path))
+    whisper_chunks = audio_processor.feature2chunks(feature_array=whisper_feature, fps=TARGET_FPS)
+    audio_frame_num = int(whisper_chunks.shape[0])
+    video_length = min(len(frames_bgr), audio_frame_num)
+    if video_length <= 0:
+        raise RuntimeError("HighSync could not align any video frames with the provided audio.")
+
+    frames_bgr = frames_bgr[:video_length]
+    audio_fea_final = torch.as_tensor(whisper_chunks[:video_length], dtype=pipeline.vae.dtype, device=pipeline.vae.device)
+    audio_fea_final = audio_fea_final.unsqueeze(0)
+
+    final_duration_sec = video_length / float(TARGET_FPS)
+    output_audio, mux_audio_path, _ = _prepare_audio_files(
+        task=task,
+        audio=audio,
+        duration_sec=final_duration_sec,
+        trim_audio_to_video=trim_audio_to_video,
+        denoise_audio=False,
+        device=device,
+        model_cache=highsync_model,
+    )
+
+    torch.manual_seed(int(seed))
+    generator = torch.manual_seed(int(seed))
+
+    output_frames_rgb = []
+    frames_video = []
+    count_motion = 0
+    count_img = 0
+
+    while True:
+        _maybe_throw_if_interrupted()
+
+        if count_motion >= video_length:
+            finish = True
+        else:
+            frames_video.append(frames_bgr[count_motion])
+            count_motion += 1
+            finish = False
+
+        if len(frames_video) == 0:
+            break
+
+        if (count_motion % preprocess_chunk_frames == 0) or finish:
+            print(f"[HighSync] processing frames {count_img}-{count_motion} / {video_length}")
+            count_img = _run_generation_chunk(
+                pipeline=pipeline,
+                image_processor=image_processor,
+                config=config,
+                audio_fea_final=audio_fea_final,
+                frames_video=frames_video,
+                output_frames_rgb=output_frames_rgb,
+                count_img=count_img,
+                generator=generator,
+                device=device,
+                weight_dtype=weight_dtype,
+                inference_steps=inference_steps,
+                cfg_scale=cfg_scale,
+            )
+            frames_video = []
+
+        if finish:
+            break
+
+    if len(output_frames_rgb) == 0:
+        raise RuntimeError("HighSync finished without producing output frames.")
+
+    output_frames = np.ascontiguousarray(np.stack(output_frames_rgb))
+    no_audio_path = encode_rgb_frames_to_mp4(output_frames, task["media_dir"] / "highsync_no_audio.mp4", fps=TARGET_FPS)
+    final_video_path = mux_video_audio(no_audio_path, mux_audio_path, task["task_dir"] / "final.mp4")
+
+    output_images = frames_uint8_rgb_to_image_tensor(output_frames)
+
+    if not keep_intermediate_outputs:
+        for path in (task["media_dir"], task["frames_dir"], task["logs_dir"]):
+            shutil.rmtree(path, ignore_errors=True)
+
+    return output_images, output_audio, TARGET_FPS, str(final_video_path.resolve())