"""
Generate Audio for Video — multi-model Gradio app.
Supported models
----------------
TARO – video-conditioned diffusion via CAVP + onset features (16 kHz, 8.192 s window)
MMAudio – multimodal flow-matching with CLIP/Synchformer + text prompt (44 kHz, 8 s window)
HunyuanFoley – text-guided foley via SigLIP2 + Synchformer + CLAP (48 kHz, up to 15 s)
"""
import functools
import gc
import html as _html
import math
import os
import sys
import json
import shutil
import tempfile
import random
import threading
import time
import warnings
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
# HunyuanVideo-Foley's DAC VAE loads its weights via the pre-parametrization
# torch.nn.utils.weight_norm API — harmless, but noisy on every model load.
warnings.filterwarnings(
"ignore",
message="`torch.nn.utils.weight_norm` is deprecated",
category=FutureWarning,
)
import torch
import numpy as np
import torchaudio
import soundfile as sf
import ffmpeg
import spaces
import gradio as gr
from huggingface_hub import HfApi, hf_hub_download, snapshot_download
# ------------------------------------------------------------------ #
# Gradio's own Video component (processing_utils.convert_video_to_ #
# playable_mp4) writes its ffmpeg output to Path(video_path) #
# .with_suffix(".mp4") — a no-op rename whenever the upload is #
# already named *.mp4 (e.g. HEVC-in-mp4). That makes it overwrite the #
# exact file the browser is concurrently streaming for the live #
# preview widget, corrupting the response mid-flight. Since HF Spaces #
# multiplexes a whole browser tab over one HTTP/2 connection, that #
# corrupted stream takes the *entire* connection down with it — even #
# unrelated in-flight requests like the generation job's SSE queue. #
# Patch it to convert a same-named copy instead, so the file the #
# browser already has a handle on is never mutated in place. #
import gradio.processing_utils as _gr_processing_utils
_gr_convert_video_to_playable_mp4 = _gr_processing_utils.convert_video_to_playable_mp4
def _patched_convert_video_to_playable_mp4(video_path: str) -> str:
p = Path(video_path)
if p.suffix.lower() == ".mp4":
copy_path = str(p.with_name(p.stem + "_gr_conv.mp4"))
shutil.copy2(video_path, copy_path)
video_path = copy_path
return _gr_convert_video_to_playable_mp4(video_path)
_gr_processing_utils.convert_video_to_playable_mp4 = _patched_convert_video_to_playable_mp4
# ------------------------------------------------------------------ #
# Gradio 5.x/6.x's Blocks.get_api_info() always returns an empty #
# unnamed_endpoints dict (confirmed still true through 5.50.0 and #
# 6.20.0 — it's simply never populated). But the JS client dispatches #
# every internal (non-api_name) click/change/upload event by numeric #
# fn_index, and looks that fn_index up in exactly that always-empty #
# dict, so endpoint_info resolves to undefined and every such event #
# logs a harmless "Too many arguments provided for the endpoint." #
# console warning (it doesn't drop data or block the request — see #
# gradio's client/src/helpers/api_info.ts map_data_to_params, which #
# just falls back to the raw args array). Since every dependency #
# already gets an auto-generated api_name and a fully-computed entry #
# in named_endpoints regardless, just mirror each of those into #
# unnamed_endpoints keyed by its numeric fn_index so the JS client's #
# lookup actually finds it. #
_orig_get_api_info = gr.Blocks.get_api_info
def _patched_get_api_info(self, all_endpoints: bool = False):
api_info = _orig_get_api_info(self, all_endpoints=all_endpoints)
for fn_index, fn in self.fns.items():
if not fn.fn or fn.api_name is False:
continue
entry = api_info["named_endpoints"].get(f"/{fn.api_name}")
if entry is not None:
api_info["unnamed_endpoints"][str(fn_index)] = entry
return api_info
gr.Blocks.get_api_info = _patched_get_api_info
# ================================================================== #
# CHECKPOINT CONFIGURATION #
# ================================================================== #
CKPT_REPO_ID = "JackIsNotInTheBox/Generate_Audio_for_Video_Checkpoints"
CACHE_DIR = "/tmp/model_ckpts"
os.makedirs(CACHE_DIR, exist_ok=True)
# ---- Local directories that must exist before parallel downloads start ----
MMAUDIO_WEIGHTS_DIR = Path(CACHE_DIR) / "MMAudio" / "weights"
MMAUDIO_EXT_DIR = Path(CACHE_DIR) / "MMAudio" / "ext_weights"
HUNYUAN_MODEL_DIR = Path(CACHE_DIR) / "HunyuanFoley"
MMAUDIO_WEIGHTS_DIR.mkdir(parents=True, exist_ok=True)
MMAUDIO_EXT_DIR.mkdir(parents=True, exist_ok=True)
HUNYUAN_MODEL_DIR.mkdir(parents=True, exist_ok=True)
# ------------------------------------------------------------------ #
# Parallel checkpoint + model downloads #
# All downloads are I/O-bound (network), so running them in threads #
# cuts Space cold-start time roughly proportional to the number of #
# independent groups (previously sequential, now concurrent). #
# hf_hub_download / snapshot_download are thread-safe. #
# ------------------------------------------------------------------ #
def _dl_taro():
"""Download TARO .ckpt/.pt files and return their local paths."""
c = hf_hub_download(repo_id=CKPT_REPO_ID, filename="TARO/cavp_epoch66.ckpt", cache_dir=CACHE_DIR)
o = hf_hub_download(repo_id=CKPT_REPO_ID, filename="TARO/onset_model.ckpt", cache_dir=CACHE_DIR)
t = hf_hub_download(repo_id=CKPT_REPO_ID, filename="TARO/taro_ckpt.pt", cache_dir=CACHE_DIR)
print("TARO checkpoints downloaded.")
return c, o, t
def _dl_mmaudio():
"""Download MMAudio .pth files and return their local paths."""
m = hf_hub_download(repo_id=CKPT_REPO_ID, filename="MMAudio/mmaudio_large_44k_v2.pth",
cache_dir=CACHE_DIR, local_dir=str(MMAUDIO_WEIGHTS_DIR))
v = hf_hub_download(repo_id=CKPT_REPO_ID, filename="MMAudio/v1-44.pth",
cache_dir=CACHE_DIR, local_dir=str(MMAUDIO_EXT_DIR))
s = hf_hub_download(repo_id=CKPT_REPO_ID, filename="MMAudio/synchformer_state_dict.pth",
cache_dir=CACHE_DIR, local_dir=str(MMAUDIO_EXT_DIR))
print("MMAudio checkpoints downloaded.")
return m, v, s
def _dl_hunyuan():
"""Download HunyuanVideoFoley .pth files."""
hf_hub_download(repo_id=CKPT_REPO_ID, filename="HunyuanVideo-Foley/hunyuanvideo_foley.pth",
cache_dir=CACHE_DIR, local_dir=str(HUNYUAN_MODEL_DIR))
hf_hub_download(repo_id=CKPT_REPO_ID, filename="HunyuanVideo-Foley/vae_128d_48k.pth",
cache_dir=CACHE_DIR, local_dir=str(HUNYUAN_MODEL_DIR))
hf_hub_download(repo_id=CKPT_REPO_ID, filename="HunyuanVideo-Foley/synchformer_state_dict.pth",
cache_dir=CACHE_DIR, local_dir=str(HUNYUAN_MODEL_DIR))
print("HunyuanVideoFoley checkpoints downloaded.")
def _populate_hf_cache_from_mirror(orig_repo_id, mirror_subpath):
"""Download an encoder slice from CKPT_REPO_ID and wire it into the HF
cache under `orig_repo_id` so existing from_pretrained() calls keyed by
the upstream id resolve from the mirror without networking upstream."""
sha = "mirror-snapshot"
mirror_root = snapshot_download(
repo_id=CKPT_REPO_ID,
allow_patterns=[f"{mirror_subpath}/**"],
)
src = Path(mirror_root) / mirror_subpath
cache_root = Path(os.environ.get(
"HF_HUB_CACHE", os.path.expanduser("~/.cache/huggingface/hub")
))
repo_cache = cache_root / f"models--{orig_repo_id.replace('/', '--')}"
snap = repo_cache / "snapshots" / sha
refs = repo_cache / "refs"
snap.parent.mkdir(parents=True, exist_ok=True)
refs.mkdir(parents=True, exist_ok=True)
if snap.is_symlink() or snap.exists():
if snap.is_symlink():
snap.unlink()
else:
shutil.rmtree(snap)
snap.symlink_to(src, target_is_directory=True)
(refs / "main").write_text(sha)
return str(snap)
def _dl_siglip2():
p = _populate_hf_cache_from_mirror(
"google/siglip2-base-patch16-512",
"encoders/google/siglip2-base-patch16-512",
)
print(f"SigLIP2 pre-downloaded from mirror -> {p}")
def _dl_clap():
p = _populate_hf_cache_from_mirror(
"laion/larger_clap_general",
"encoders/laion/larger_clap_general",
)
print(f"CLAP pre-downloaded from mirror -> {p}")
def _dl_clip():
# Space code uses the legacy '-384' alias which upstream redirects to
# '-378'. We mirror under the canonical name and key the cache under
# the alias so the existing 'hf-hub:apple/...-384' calls resolve.
p = _populate_hf_cache_from_mirror(
"apple/DFN5B-CLIP-ViT-H-14-384",
"encoders/apple/DFN5B-CLIP-ViT-H-14-378",
)
print(f"MMAudio CLIP pre-downloaded from mirror -> {p}")
def _dl_audioldm2():
p = _populate_hf_cache_from_mirror(
"cvssp/audioldm2",
"encoders/cvssp/audioldm2",
)
print(f"AudioLDM2 pre-downloaded from mirror -> {p}")
def _dl_bigvgan():
"""Returns the local snapshot dir so MMAudio's BigVGANv2.from_pretrained()
receives a local path rather than reaching out to HF Hub."""
p = _populate_hf_cache_from_mirror(
"nvidia/bigvgan_v2_44khz_128band_512x",
"encoders/nvidia/bigvgan_v2_44khz_128band_512x",
)
print(f"BigVGAN pre-downloaded from mirror -> {p}")
return p
print("[startup] Starting parallel checkpoint + model downloads…")
_t_dl_start = time.perf_counter()
with ThreadPoolExecutor(max_workers=8) as _pool:
_fut_taro = _pool.submit(_dl_taro)
_fut_mmaudio = _pool.submit(_dl_mmaudio)
_fut_hunyuan = _pool.submit(_dl_hunyuan)
_fut_clap = _pool.submit(_dl_clap)
_fut_clip = _pool.submit(_dl_clip)
_fut_aldm2 = _pool.submit(_dl_audioldm2)
_fut_bigvgan = _pool.submit(_dl_bigvgan)
_fut_siglip2 = _pool.submit(_dl_siglip2)
# Raise any download exceptions immediately
for _fut in as_completed([_fut_taro, _fut_mmaudio, _fut_hunyuan,
_fut_clap, _fut_clip, _fut_aldm2, _fut_bigvgan,
_fut_siglip2]):
_fut.result()
cavp_ckpt_path, onset_ckpt_path, taro_ckpt_path = _fut_taro.result()
mmaudio_model_path, mmaudio_vae_path, mmaudio_synchformer_path = _fut_mmaudio.result()
bigvgan_local_dir = _fut_bigvgan.result()
print(f"[startup] All downloads done in {time.perf_counter() - _t_dl_start:.1f}s")
print(f"[startup] main process pid={os.getpid()}")
# Force any later HF Hub call to be cache-only. With every encoder cache
# pre-populated above from the user's mirror, downstream from_pretrained()
# calls (CLAP, AudioLDM2) resolve from the populated cache instead of
# networking out to upstream. Best-effort: the exact attribute name/shape of
# huggingface_hub's offline-mode internals has proven to vary by installed
# version (an is_offline_mode() debug print crashed live with AttributeError,
# meaning the deployed version's constants.py doesn't match what was tested
# against locally) — os.environ is the one thing every version reads at
# *some* point, so set it and best-effort mirror it onto the constant too,
# but don't rely on this alone (see the open_clip-specific patch below for
# the one confirmed, reproduced failure).
os.environ["HF_HUB_OFFLINE"] = "1"
import huggingface_hub
import huggingface_hub.constants as _hf_constants
print(f"[startup] huggingface_hub version={getattr(huggingface_hub, '__version__', '?')}")
if hasattr(_hf_constants, "HF_HUB_OFFLINE"):
_hf_constants.HF_HUB_OFFLINE = True
def _ensure_offline_in_worker():
"""Re-assert offline mode from inside a @spaces.GPU function's own body
(called as the first line of all six raw GPU functions) in case the
worker thread anyio dispatches @spaces.GPU calls onto doesn't share the
same already-mutated module state as the main thread."""
os.environ["HF_HUB_OFFLINE"] = "1"
if hasattr(_hf_constants, "HF_HUB_OFFLINE"):
_hf_constants.HF_HUB_OFFLINE = True
print(f"[_ensure_offline_in_worker] pid={os.getpid()} "
f"HF_HUB_OFFLINE={getattr(_hf_constants, 'HF_HUB_OFFLINE', '?')}")
# Settle, empirically, whether this call actually has real GPU access.
# ZeroGPU docs state @spaces.GPU is a no-op outside a ZeroGPU-hardware
# Space, silently falling back to CPU (or a CUDA-emulation shim) — which
# would perfectly explain a "hung" generation that's actually just very
# slow CPU-bound diffusion. Print hard facts instead of guessing.
try:
cuda_ok = torch.cuda.is_available()
print(f"[_ensure_offline_in_worker] torch.cuda.is_available()={cuda_ok} "
f"device_name={torch.cuda.get_device_name(0) if cuda_ok else 'N/A'} "
f"device_count={torch.cuda.device_count() if cuda_ok else 0}")
except Exception as e:
print(f"[_ensure_offline_in_worker] CUDA check raised: {e!r}")
# The one CONFIRMED, reproduced failure: MMAudio's open_clip CLIP loader
# (create_model_from_pretrained('hf-hub:apple/DFN5B-CLIP-ViT-H-14-384', ...))
# re-downloads its ~4GB checkpoint live during GPU inference regardless of the
# HF_HUB_OFFLINE patches above — verified via a direct HTTP repro against the
# deployed Space (queue/join straight to _run_mmaudio, bypassing the browser).
# open_clip.pretrained calls a *bound* local reference to hf_hub_download
# (`from huggingface_hub import hf_hub_download` at its own import time), so
# rather than continue chasing which huggingface_hub-internal offline-mode
# attribute the deployed version actually honors, force the one parameter
# that unambiguously controls this regardless of version: local_files_only.
import open_clip.pretrained as _open_clip_pretrained
_orig_open_clip_hf_hub_download = _open_clip_pretrained.hf_hub_download
def _patched_open_clip_hf_hub_download(*args, **kwargs):
kwargs["local_files_only"] = True
return _orig_open_clip_hf_hub_download(*args, **kwargs)
_open_clip_pretrained.hf_hub_download = _patched_open_clip_hf_hub_download
# ================================================================== #
# SHARED CONSTANTS / HELPERS #
# ================================================================== #
# CPU → GPU context passing via function-name-keyed global store.
#
# Problem: ZeroGPU runs @spaces.GPU functions on its own worker thread, so
# threading.local() is invisible to the GPU worker. Passing ctx as a
# function argument exposes it to Gradio's API endpoint, causing
# "Too many arguments" errors.
#
# Solution: store context in a plain global dict keyed by function name.
# A per-key Lock serialises concurrent callers for the same function
# (ZeroGPU is already synchronous — the wrapper blocks until the GPU fn
# returns — so in practice only one call per GPU fn is in-flight at a time).
# The global dict is readable from any thread.
_GPU_CTX: dict = {}
_GPU_CTX_LOCK = threading.Lock()
def _ctx_store(fn_name: str, data: dict) -> None:
"""Store *data* under *fn_name* key (overwrites previous)."""
with _GPU_CTX_LOCK:
_GPU_CTX[fn_name] = data
def _ctx_load(fn_name: str) -> dict:
"""Pop and return the context dict stored under *fn_name*."""
with _GPU_CTX_LOCK:
return _GPU_CTX.pop(fn_name, {})
MAX_SLOTS = 8 # max parallel generation slots shown in UI
MAX_SEGS = 8 # max segments per slot (same as MAX_SLOTS; video ≤ ~64 s at 8 s/seg)
# Segment overlay palette — shared between _build_waveform_html and _build_regen_pending_html
SEG_COLORS = [
"rgba(100,180,255,{a})", "rgba(255,160,100,{a})",
"rgba(120,220,140,{a})", "rgba(220,120,220,{a})",
"rgba(255,220,80,{a})", "rgba(80,220,220,{a})",
"rgba(255,100,100,{a})", "rgba(180,255,180,{a})",
]
# ------------------------------------------------------------------ #
# Micro-helpers that eliminate repeated boilerplate across the file #
# ------------------------------------------------------------------ #
def _ensure_syspath(subdir: str) -> str:
"""Add *subdir* (relative to app.py) to sys.path if not already present.
Returns the absolute path for convenience."""
p = os.path.join(os.path.dirname(os.path.abspath(__file__)), subdir)
if p not in sys.path:
sys.path.insert(0, p)
return p
# Ensure all model library subdirs are in sys.path for the MAIN PROCESS.
# ZeroGPU workers also call _ensure_syspath internally, but when the worker
# pickles its return value and the main process deserialises it via
# _ForkingPickler.loads(), the main process must already have these modules
# importable — otherwise ModuleNotFoundError is raised during unpickling.
_ensure_syspath("HunyuanVideo-Foley")
_ensure_syspath("MMAudio")
_ensure_syspath("TARO")
print("[startup] Model library sys.path entries registered in main process.")
def _get_device_and_dtype() -> tuple:
"""Return (device, weight_dtype) pair used by all GPU functions.
Returns a torch.device object so callers can use device.type safely."""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
return device, torch.bfloat16
def _extract_segment_clip(silent_video: str, seg_start: float, seg_dur: float,
output_path: str) -> str:
"""Stream-copy a segment from *silent_video* to *output_path*. Returns *output_path*."""
ffmpeg.input(silent_video, ss=seg_start, t=seg_dur).output(
output_path, vcodec="copy", an=None
).run(overwrite_output=True, quiet=True)
return output_path
# Per-slot reentrant locks — prevent concurrent regens on the same slot from
# producing a race condition where the second regen reads stale state
# (the shared seg_state textbox hasn't been updated yet by the first regen).
# Locks are keyed by slot_id string (e.g. "taro_0", "mma_2").
_SLOT_LOCKS: dict = {}
_SLOT_LOCKS_MUTEX = threading.Lock()
def _get_slot_lock(slot_id: str) -> threading.Lock:
with _SLOT_LOCKS_MUTEX:
if slot_id not in _SLOT_LOCKS:
_SLOT_LOCKS[slot_id] = threading.Lock()
return _SLOT_LOCKS[slot_id]
def set_global_seed(seed: int) -> None:
np.random.seed(seed % (2**32))
random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(seed)
def get_random_seed() -> int:
return random.randint(0, 2**32 - 1)
def _resolve_seed(seed_val) -> int:
"""Normalise seed_val to a non-negative int.
Negative values (UI default 'random') produce a fresh random seed."""
seed_val = int(seed_val)
return seed_val if seed_val >= 0 else get_random_seed()
def get_video_duration(video_path: str) -> float:
"""Return video duration in seconds (CPU only)."""
probe = ffmpeg.probe(video_path)
return float(probe["format"]["duration"])
def strip_audio_from_video(video_path: str, output_path: str) -> None:
"""Write a silent copy of *video_path* to *output_path*.
Tries a fast stream-copy first. Falls back to re-encoding when the
source codec can't be stream-copied into an mp4 container (e.g.
VP9/webm, ProRes/mov, or any upload that skipped/failed the
browser-preview H.264 transcode in _transcode_for_browser)."""
try:
ffmpeg.input(video_path).output(output_path, vcodec="copy", an=None).run(
overwrite_output=True, quiet=True
)
except ffmpeg.Error as e:
stderr = e.stderr.decode(errors="replace") if e.stderr else str(e)
print(f"[strip_audio_from_video] stream-copy failed, re-encoding instead: {stderr[-2000:]}")
ffmpeg.input(video_path).output(
output_path, vcodec="libx264", preset="fast", crf=18,
pix_fmt="yuv420p", an=None,
).run(overwrite_output=True, quiet=True)
def _transcode_for_browser(video_path: str):
"""Re-encode uploaded video to H.264/AAC MP4 so the browser preview widget can play it.
Returns a NEW path in a fresh /tmp/gradio/ subdirectory. Gradio probes the
returned path fresh, sees H.264, and serves it directly without its own
slow fallback converter. The in-place overwrite approach loses the race
because Gradio probes the original path at upload time before this callback runs.
Only called on upload — not during generation.
Cross-tab sync (_sync) mirrors whatever this returns onto the other two
tabs' Video components, and each of THEIR .upload() handlers can re-fire
on that reassignment — without the guards below, that reprocesses an
already-transcoded "_h264.mp4"/"_gr_conv.mp4" file, and if it has since
been cleaned up this turns into an infinite failed-retry loop (observed
in prod: 39 repeated "No such file" failures alternating between the two
names for one upload). Both guards return gr.update() (no-op, doesn't
touch the component's value) instead of re-assigning a path, so there's
nothing for a re-triggered handler to chase.
"""
if video_path is None:
return video_path
base = os.path.basename(video_path)
if base.endswith("_h264.mp4") or base.endswith("_gr_conv.mp4"):
print(f"[transcode_for_browser] already processed ({base}), skipping")
return gr.update()
if not os.path.exists(video_path):
print(f"[transcode_for_browser] input vanished before processing: {video_path!r}")
return gr.update()
try:
probe = ffmpeg.probe(video_path)
has_audio = any(s["codec_type"] == "audio" for s in probe.get("streams", []))
# Check if already H.264 — skip transcode if so
video_streams = [s for s in probe.get("streams", []) if s["codec_type"] == "video"]
if video_streams and video_streams[0].get("codec_name") == "h264":
print(f"[transcode_for_browser] already H.264, skipping")
return video_path
# Write the H.264 output into the SAME directory as the original upload.
# Gradio's file server only allows paths under dirs it registered — the
# upload dir is already allowed, so a sibling file there will serve fine.
import os as _os
upload_dir = _os.path.dirname(video_path)
stem = _os.path.splitext(_os.path.basename(video_path))[0]
out_path = _os.path.join(upload_dir, stem + "_h264.mp4")
kwargs = dict(
vcodec="libx264", preset="fast", crf=18,
pix_fmt="yuv420p", movflags="+faststart",
)
if has_audio:
kwargs["acodec"] = "aac"
kwargs["audio_bitrate"] = "128k"
else:
kwargs["an"] = None
# map 0:v:0 explicitly to skip non-video streams (e.g. data/timecode tracks)
ffmpeg.input(video_path).output(out_path, map="0:v:0", **kwargs).run(
overwrite_output=True, quiet=True
)
print(f"[transcode_for_browser] transcoded to H.264: {out_path}")
return out_path
except Exception as e:
stderr = e.stderr.decode(errors="replace") if isinstance(e, ffmpeg.Error) and e.stderr else str(e)
print(f"[transcode_for_browser] failed, using original: {stderr[-2000:]}")
return video_path
# ------------------------------------------------------------------ #
# Temp directory registry — tracks dirs for cleanup on new generation #
# ------------------------------------------------------------------ #
_TEMP_DIRS: list = [] # list of tmp_dir paths created by generate_*
_TEMP_DIRS_MAX = 10 # keep at most this many; older ones get cleaned up
def _register_tmp_dir(tmp_dir: str) -> str:
"""Register a temp dir so it can be cleaned up when newer ones replace it."""
_TEMP_DIRS.append(tmp_dir)
while len(_TEMP_DIRS) > _TEMP_DIRS_MAX:
old = _TEMP_DIRS.pop(0)
try:
shutil.rmtree(old, ignore_errors=True)
print(f"[cleanup] Removed old temp dir: {old}")
except Exception:
pass
return tmp_dir
def _to_stereo(w: np.ndarray) -> np.ndarray:
"""Ensure *w* is stereo (2, T). Squeezes (1,T) then duplicates mono."""
if w.ndim == 2 and w.shape[0] == 1:
w = w.squeeze(0) # (1, T) → (T,)
if w.ndim == 1:
w = np.stack([w, w], axis=0) # (T,) → (2, T)
return w
def _save_seg_wavs(wavs: list[np.ndarray], tmp_dir: str, prefix: str) -> list[str]:
"""Save a list of numpy wav arrays to .npy files as stereo (2, T).
Mono arrays are duplicated to stereo so the on-disk format is always
uniform — this avoids shape mismatches during cross-model regens."""
paths = []
for i, w in enumerate(wavs):
p = os.path.join(tmp_dir, f"{prefix}_seg{i}.npy")
np.save(p, _to_stereo(w))
paths.append(p)
return paths
def _load_seg_wavs(paths: list[str]) -> list[np.ndarray]:
"""Load segment wav arrays from .npy file paths.
All files on disk are expected to be stereo (2, T) — _save_seg_wavs
guarantees this. Raises AssertionError if any array has unexpected shape.
"""
wavs = []
for p in paths:
w = np.load(p)
assert w.ndim == 2 and w.shape[0] == 2, \
f"Expected stereo (2, T) in {p}, got shape {w.shape}"
wavs.append(w)
return wavs
# ------------------------------------------------------------------ #
# Shared model-loading helpers (deduplicate generate / regen code) #
# ------------------------------------------------------------------ #
def _load_taro_models(device, weight_dtype):
"""Load TARO MMDiT + AudioLDM2 VAE/vocoder. Returns (model_net, vae, vocoder, latents_scale)."""
from TARO.models import MMDiT
from diffusers import AutoencoderKL
from transformers import SpeechT5HifiGan
model_net = MMDiT(adm_in_channels=120, z_dims=[768], encoder_depth=4).to(device)
model_net.load_state_dict(torch.load(taro_ckpt_path, map_location=device, weights_only=False)["ema"])
model_net.eval().to(weight_dtype)
vae = AutoencoderKL.from_pretrained("cvssp/audioldm2", subfolder="vae",
local_files_only=True).to(device).eval()
vocoder = SpeechT5HifiGan.from_pretrained("cvssp/audioldm2", subfolder="vocoder",
local_files_only=True).to(device)
latents_scale = torch.tensor([0.18215] * 8).view(1, 8, 1, 1).to(device)
return model_net, vae, vocoder, latents_scale
def _load_taro_feature_extractors(device):
"""Load CAVP + onset extractors. Returns (extract_cavp, onset_model)."""
from TARO.cavp_util import Extract_CAVP_Features
from TARO.onset_util import VideoOnsetNet
extract_cavp = Extract_CAVP_Features(
device=device, config_path="TARO/cavp/cavp.yaml", ckpt_path=cavp_ckpt_path,
)
raw_sd = torch.load(onset_ckpt_path, map_location=device, weights_only=False)["state_dict"]
onset_sd = {}
for k, v in raw_sd.items():
if "model.net.model" in k: k = k.replace("model.net.model", "net.model")
elif "model.fc." in k: k = k.replace("model.fc", "fc")
onset_sd[k] = v
onset_model = VideoOnsetNet(pretrained=False).to(device)
onset_model.load_state_dict(onset_sd)
onset_model.eval()
return extract_cavp, onset_model
def _load_mmaudio_models(device, dtype):
"""Load MMAudio net + feature_utils. Returns (net, feature_utils, model_cfg, seq_cfg)."""
from mmaudio.eval_utils import all_model_cfg
from mmaudio.model.networks import get_my_mmaudio
from mmaudio.model.utils.features_utils import FeaturesUtils
model_cfg = all_model_cfg["large_44k_v2"]
model_cfg.model_path = Path(mmaudio_model_path)
model_cfg.vae_path = Path(mmaudio_vae_path)
model_cfg.synchformer_ckpt = Path(mmaudio_synchformer_path)
model_cfg.bigvgan_16k_path = None
seq_cfg = model_cfg.seq_cfg
net = get_my_mmaudio(model_cfg.model_name).to(device, dtype).eval()
net.load_weights(torch.load(model_cfg.model_path, map_location=device, weights_only=True))
feature_utils = FeaturesUtils(
tod_vae_ckpt=str(model_cfg.vae_path),
synchformer_ckpt=str(model_cfg.synchformer_ckpt),
enable_conditions=True, mode=model_cfg.mode,
bigvgan_vocoder_ckpt=bigvgan_local_dir, need_vae_encoder=False,
).to(device, dtype).eval()
return net, feature_utils, model_cfg, seq_cfg
def _load_hunyuan_model(device, model_size):
"""Load HunyuanFoley model dict + config. Returns (model_dict, cfg)."""
from hunyuanvideo_foley.utils.model_utils import load_model
model_size = model_size.lower()
config_map = {
"xl": "HunyuanVideo-Foley/configs/hunyuanvideo-foley-xl.yaml",
"xxl": "HunyuanVideo-Foley/configs/hunyuanvideo-foley-xxl.yaml",
}
config_path = config_map.get(model_size, config_map["xxl"])
hunyuan_weights_dir = str(HUNYUAN_MODEL_DIR / "HunyuanVideo-Foley")
print(f"[HunyuanFoley] Loading {model_size.upper()} model from {hunyuan_weights_dir}")
return load_model(hunyuan_weights_dir, config_path, device,
enable_offload=True, model_size=model_size)
def mux_video_audio(silent_video: str, audio_path: str, output_path: str,
model: str = None) -> None:
"""Mux a silent video with an audio file into *output_path*.
For HunyuanFoley (*model*="hunyuan") we use its own merge_audio_video which
handles its specific ffmpeg quirks; all other models use stream-copy muxing.
"""
if model == "hunyuan":
_ensure_syspath("HunyuanVideo-Foley")
from hunyuanvideo_foley.utils.media_utils import merge_audio_video
merge_audio_video(audio_path, silent_video, output_path)
else:
v_in = ffmpeg.input(silent_video)
a_in = ffmpeg.input(audio_path)
ffmpeg.output(
v_in["v:0"],
a_in["a:0"],
output_path,
vcodec="libx264", preset="fast", crf=18,
pix_fmt="yuv420p",
acodec="aac", audio_bitrate="128k",
movflags="+faststart",
).run(overwrite_output=True, quiet=True)
# ------------------------------------------------------------------ #
# Shared sliding-window segmentation and crossfade helpers #
# Used by all three models (TARO, MMAudio, HunyuanFoley). #
# ------------------------------------------------------------------ #
def _build_segments(total_dur_s: float, window_s: float, crossfade_s: float) -> list[tuple[float, float]]:
"""Return list of (start, end) pairs covering *total_dur_s*.
Every segment uses the full *window_s* inference window. Segments are
equally spaced so every overlap is identical, guaranteeing the crossfade
setting is honoured at every boundary with no raw bleed.
Algorithm
---------
1. Clamp crossfade_s so the step stays positive.
2. Find the minimum n such that n segments of *window_s* cover
*total_dur_s* with overlap ≥ crossfade_s at every boundary:
n = ceil((total_dur_s - crossfade_s) / (window_s - crossfade_s))
3. Compute equal spacing: step = (total_dur_s - window_s) / (n - 1)
so that every gap is identical and the last segment ends exactly at
total_dur_s.
4. Every segment is exactly *window_s* wide. The trailing audio of each
segment beyond its contact edge is discarded in _stitch_wavs.
"""
crossfade_s = min(crossfade_s, window_s * 0.5)
if total_dur_s <= window_s:
return [(0.0, total_dur_s)]
step_min = window_s - crossfade_s # minimum step to honour crossfade
n = math.ceil((total_dur_s - crossfade_s) / step_min)
n = max(n, 2)
# Equal step so first seg starts at 0 and last seg ends at total_dur_s
step_s = (total_dur_s - window_s) / (n - 1)
return [(i * step_s, i * step_s + window_s) for i in range(n)]
def _cf_join(a: np.ndarray, b: np.ndarray,
crossfade_s: float, db_boost: float, sr: int) -> np.ndarray:
"""Equal-power crossfade join. Works for both mono (T,) and stereo (C, T) arrays.
Stereo arrays are expected in (channels, samples) layout.
db_boost is applied to the overlap region as a whole (after blending), so
it compensates for the -3 dB equal-power dip without doubling amplitude.
Applying gain to each side independently (the common mistake) causes a
+3 dB loudness bump at the seam — this version avoids that."""
stereo = a.ndim == 2
n_a = a.shape[1] if stereo else len(a)
n_b = b.shape[1] if stereo else len(b)
cf = min(int(round(crossfade_s * sr)), n_a, n_b)
if cf <= 0:
return np.concatenate([a, b], axis=1 if stereo else 0)
gain = 10 ** (db_boost / 20.0)
t = np.linspace(0.0, 1.0, cf, dtype=np.float32)
fade_out = np.cos(t * np.pi / 2) # 1 → 0
fade_in = np.sin(t * np.pi / 2) # 0 → 1
if stereo:
# Blend first, then apply boost to the overlap region as a unit
overlap = (a[:, -cf:] * fade_out + b[:, :cf] * fade_in) * gain
return np.concatenate([a[:, :-cf], overlap, b[:, cf:]], axis=1)
else:
overlap = (a[-cf:] * fade_out + b[:cf] * fade_in) * gain
return np.concatenate([a[:-cf], overlap, b[cf:]])
# ================================================================== #
# TARO #
# ================================================================== #
# Constants sourced from TARO/infer.py and TARO/models.py:
# SR=16000, TRUNCATE=131072 → 8.192 s window
# TRUNCATE_FRAME = 4 fps × 131072/16000 = 32 CAVP frames per window
# TRUNCATE_ONSET = 120 onset frames per window
# latent shape: (1, 8, 204, 16) — fixed by MMDiT architecture
# latents_scale: [0.18215]*8 — AudioLDM2 VAE scale factor
# ================================================================== #
# ================================================================== #
# MODEL CONSTANTS & CONFIGURATION REGISTRY #
# ================================================================== #
# All per-model numeric constants live here — MODEL_CONFIGS is the #
# single source of truth consumed by duration estimation, segmentation,#
# and the UI. Standalone names kept only where other code references #
# them by name (TARO geometry, TARGET_SR, GPU_DURATION_CAP). #
# ================================================================== #
# TARO geometry — referenced directly in _taro_infer_segment
TARO_SR = 16000
TARO_TRUNCATE = 131072
TARO_FPS = 4
TARO_TRUNCATE_FRAME = int(TARO_FPS * TARO_TRUNCATE / TARO_SR) # 32
TARO_TRUNCATE_ONSET = 120
TARO_MODEL_DUR = TARO_TRUNCATE / TARO_SR # 8.192 s
GPU_DURATION_CAP = 300 # hard cap per @spaces.GPU call — never reserve more than this
MODEL_CONFIGS = {
"taro": {
"window_s": TARO_MODEL_DUR, # 8.192 s
"sr": TARO_SR, # 16000 (output resampled to TARGET_SR)
"secs_per_step": 0.025, # measured 0.023 s/step on H200
"load_overhead": 35, # CAVP+onset+MMDiT+VAE+vocoder load + feature extraction
# (measured ~18s warm end-to-end on RTX Pro 6000 Blackwell; ~2x margin)
"tab_prefix": "taro",
"label": "TARO",
"regen_fn": None, # set after function definitions (avoids forward-ref)
},
"mmaudio": {
"window_s": 8.0, # MMAudio's fixed generation window
"sr": 48000, # resampled from 44100 in post-processing
"secs_per_step": 0.25, # measured 0.230 s/step on H200
"load_overhead": 35, # 15s warm + 15s model init
# (measured ~23s warm end-to-end on RTX Pro 6000 Blackwell; ~1.5x margin)
"tab_prefix": "mma",
"label": "MMAudio",
"regen_fn": None,
},
"hunyuan": {
"window_s": 15.0, # HunyuanFoley max video duration
"sr": 48000,
"secs_per_step": 0.50, # wall-time ~0.607s/step (incl. VAE decode + feature extraction)
"load_overhead": 55, # SigLIP2 + Synchformer + CLAP + main model + DAC VAE load
# (measured ~39s warm end-to-end on RTX Pro 6000 Blackwell; extra
# margin kept vs. TARO/MMAudio for first-load cold-disk-cache risk
# on the ~10 GB XXL checkpoint)
"tab_prefix": "hf",
"label": "HunyuanFoley",
"regen_fn": None,
},
}
# Convenience aliases used only in the TARO inference path
TARO_SECS_PER_STEP = MODEL_CONFIGS["taro"]["secs_per_step"]
MMAUDIO_WINDOW = MODEL_CONFIGS["mmaudio"]["window_s"]
MMAUDIO_SECS_PER_STEP = MODEL_CONFIGS["mmaudio"]["secs_per_step"]
HUNYUAN_MAX_DUR = MODEL_CONFIGS["hunyuan"]["window_s"]
HUNYUAN_SECS_PER_STEP = MODEL_CONFIGS["hunyuan"]["secs_per_step"]
def _catch_oom(fn):
"""Wrap a @spaces.GPU function so a CUDA OOM can't escape and crash the
ZeroGPU worker. An escaped OOM doesn't just fail that one request — it can
take down the whole Space container, which restarts from scratch (every
checkpoint re-downloaded, ~2 minutes) instead of just showing an error.
Applied as the innermost decorator (below @spaces.GPU) so spaces.GPU's
duration estimator still sees the original call signature via
functools.wraps.
Also re-asserts offline mode here, at the top of the actual function body.
ZeroGPU runs @spaces.GPU functions in a separate worker process, and that
worker does not reliably inherit the main process's post-import os.environ
/ huggingface_hub.constants mutations made at startup (confirmed live: the
same os.environ + constants patch that fixed offline resolution when
tested in-process still didn't stop MMAudio's open_clip CLIP loader from
re-downloading its ~4GB checkpoint from inside an actual GPU call). Setting
it again here runs inside whatever process actually executes the GPU
call, so it can't miss the worker regardless of ZeroGPU's fork/spawn
timing."""
@functools.wraps(fn)
def _wrapped(*args, **kwargs):
os.environ["HF_HUB_OFFLINE"] = "1"
_hf_constants.HF_HUB_OFFLINE = True
try:
return fn(*args, **kwargs)
except torch.cuda.OutOfMemoryError as e:
if torch.cuda.is_available():
torch.cuda.empty_cache()
gc.collect()
raise gr.Error(
f"Out of GPU memory during generation ({fn.__name__}). "
"Try fewer samples, a shorter clip, or fewer steps."
) from e
return _wrapped
def _clamp_duration(secs: float, label: str) -> int:
"""Clamp a raw GPU-seconds estimate to [45, GPU_DURATION_CAP] and log it.
ZeroGPU Pro users get up to 300 s per call; 45 s floor covers per-call
scheduling/queue overhead without over-reserving quota on every warm
call (each model's own load_overhead in MODEL_CONFIGS covers its
model-load cost above this floor)."""
result = min(GPU_DURATION_CAP, max(45, int(secs)))
print(f"[duration] {label}: {secs:.0f}s raw → {result}s reserved")
return result
def _estimate_gpu_duration(model_key: str, num_samples: int, num_steps: int,
total_dur_s: float = None, crossfade_s: float = 0,
video_file: str = None) -> int:
"""Estimate GPU seconds for a full generation call.
Formula: num_samples × n_segs × num_steps × secs_per_step + load_overhead
"""
cfg = MODEL_CONFIGS[model_key]
try:
if total_dur_s is None:
total_dur_s = get_video_duration(video_file)
n_segs = len(_build_segments(total_dur_s, cfg["window_s"], float(crossfade_s)))
except Exception:
n_segs = 1
secs = int(num_samples) * n_segs * int(num_steps) * cfg["secs_per_step"] + cfg["load_overhead"]
print(f"[duration] {cfg['label']}: {int(num_samples)}samp × {n_segs}seg × "
f"{int(num_steps)}steps → {secs:.0f}s → capped ", end="")
return _clamp_duration(secs, cfg["label"])
def _estimate_regen_duration(model_key: str, num_steps: int) -> int:
"""Estimate GPU seconds for a single-segment regen call."""
cfg = MODEL_CONFIGS[model_key]
secs = int(num_steps) * cfg["secs_per_step"] + cfg["load_overhead"]
print(f"[duration] {cfg['label']} regen: 1 seg × {int(num_steps)} steps → ", end="")
return _clamp_duration(secs, f"{cfg['label']} regen")
_TARO_CACHE_MAXLEN = 16 # evict oldest entries beyond this limit
_TARO_INFERENCE_CACHE: dict = {} # keyed by (video_file, seed, cfg, steps, mode, crossfade_s)
_TARO_CACHE_LOCK = threading.Lock()
def _taro_calc_max_samples(total_dur_s: float, num_steps: int, crossfade_s: float) -> int:
n_segs = len(_build_segments(total_dur_s, TARO_MODEL_DUR, crossfade_s))
time_per_seg = num_steps * TARO_SECS_PER_STEP
max_s = int(600.0 / (n_segs * time_per_seg))
return max(1, min(max_s, MAX_SLOTS))
def _taro_duration(video_file, seed_val, cfg_scale, num_steps, mode,
crossfade_s, crossfade_db, num_samples):
"""Pre-GPU callable — must match _taro_gpu_infer's input order exactly."""
return _estimate_gpu_duration("taro", int(num_samples), int(num_steps),
video_file=video_file, crossfade_s=crossfade_s)
def _taro_infer_segment(
model, vae, vocoder,
cavp_feats_full, onset_feats_full,
seg_start_s: float, seg_end_s: float,
device, weight_dtype,
cfg_scale: float, num_steps: int, mode: str,
latents_scale,
euler_sampler, euler_maruyama_sampler,
) -> np.ndarray:
"""Single-segment TARO inference. Returns wav array trimmed to segment length."""
# CAVP features (4 fps)
cavp_start = int(round(seg_start_s * TARO_FPS))
cavp_slice = cavp_feats_full[cavp_start : cavp_start + TARO_TRUNCATE_FRAME]
if cavp_slice.shape[0] < TARO_TRUNCATE_FRAME:
pad = np.zeros(
(TARO_TRUNCATE_FRAME - cavp_slice.shape[0],) + cavp_slice.shape[1:],
dtype=cavp_slice.dtype,
)
cavp_slice = np.concatenate([cavp_slice, pad], axis=0)
video_feats = torch.from_numpy(cavp_slice).unsqueeze(0).to(device, weight_dtype)
# Onset features (onset_fps = TRUNCATE_ONSET / MODEL_DUR ≈ 14.65 fps)
onset_fps = TARO_TRUNCATE_ONSET / TARO_MODEL_DUR
onset_start = int(round(seg_start_s * onset_fps))
onset_slice = onset_feats_full[onset_start : onset_start + TARO_TRUNCATE_ONSET]
if onset_slice.shape[0] < TARO_TRUNCATE_ONSET:
onset_slice = np.pad(
onset_slice,
((0, TARO_TRUNCATE_ONSET - onset_slice.shape[0]),),
mode="constant",
)
onset_feats_t = torch.from_numpy(onset_slice).unsqueeze(0).to(device, weight_dtype)
# Latent noise — shape matches MMDiT architecture (in_channels=8, 204×16 spatial)
z = torch.randn(1, model.in_channels, 204, 16, device=device, dtype=weight_dtype)
sampling_kwargs = dict(
model=model,
latents=z,
y=onset_feats_t,
context=video_feats,
num_steps=int(num_steps),
heun=False,
cfg_scale=float(cfg_scale),
guidance_low=0.0,
guidance_high=0.7,
path_type="linear",
)
with torch.no_grad():
samples = (euler_maruyama_sampler if mode == "sde" else euler_sampler)(**sampling_kwargs)
# samplers return (output_tensor, zs) — index [0] for the audio latent
if isinstance(samples, tuple):
samples = samples[0]
# Decode: AudioLDM2 VAE → mel → vocoder → waveform
samples = vae.decode(samples / latents_scale).sample
wav = vocoder(samples.squeeze().float()).detach().cpu().numpy()
return wav # full window — _stitch_wavs handles contact-edge trimming
# ================================================================== #
# TARO 16 kHz → 48 kHz upsample #
# ================================================================== #
# TARO generates at 16 kHz; all other models output at 44.1/48 kHz.
# We upsample via sinc resampling (torchaudio, CPU-only) so the final
# stitched audio is uniformly at 48 kHz across all three models.
TARGET_SR = 48000 # unified output sample rate for all three models
TARO_SR_OUT = TARGET_SR
def _resample_to_target(wav: np.ndarray, src_sr: int,
dst_sr: int = None) -> np.ndarray:
"""Resample *wav* (mono or stereo numpy float32) from *src_sr* to *dst_sr*.
*dst_sr* defaults to TARGET_SR (48 kHz). No-op if src_sr == dst_sr.
Uses torchaudio Kaiser-windowed sinc resampling — CPU-only, ZeroGPU-safe.
Shape contract: returns same shape as input.
(T,) → (T,) mono
(2, T) → (2, T) stereo
(1, T) → (1, T) 1-channel (caller must apply _to_stereo if needed)
"""
if dst_sr is None:
dst_sr = TARGET_SR
if src_sr == dst_sr:
return wav
# Treat as stereo only when there are genuinely 2 channels
stereo = wav.ndim == 2 and wav.shape[0] == 2
t = torch.from_numpy(np.ascontiguousarray(wav.astype(np.float32)))
if not stereo:
if t.ndim == 1:
t = t.unsqueeze(0) # (T,) → (1, T) for resample
# (1, T) or any other 2D single-channel stays as-is
t = torchaudio.functional.resample(t, src_sr, dst_sr)
if not stereo and wav.ndim == 1:
t = t.squeeze(0) # (1, T) → (T,) to match input shape
return t.numpy()
def _upsample_taro(wav_16k: np.ndarray) -> np.ndarray:
"""Upsample a mono 16 kHz numpy array to 48 kHz via sinc resampling (CPU).
torchaudio.functional.resample uses a Kaiser-windowed sinc filter —
mathematically optimal for bandlimited signals, zero CUDA risk.
Returns a mono float32 numpy array at 48 kHz.
"""
dur_in = len(wav_16k) / TARO_SR
print(f"[TARO upsample] {dur_in:.2f}s @ {TARO_SR}Hz → {TARGET_SR}Hz (sinc, CPU) …")
result = _resample_to_target(wav_16k, TARO_SR)
print(f"[TARO upsample] done — {len(result)/TARGET_SR:.2f}s @ {TARGET_SR}Hz "
f"(expected {dur_in * 3:.2f}s, ratio 3×)")
return result
def _stitch_wavs(wavs: list[np.ndarray], crossfade_s: float, db_boost: float,
total_dur_s: float, sr: int,
segments: list[tuple[float, float]] = None) -> np.ndarray:
"""Crossfade-join a list of wav arrays and trim to *total_dur_s*.
Works for both mono (T,) and stereo (C, T) arrays.
When *segments* is provided (list of (start, end) video-time pairs),
each wav is trimmed to its contact-edge window before joining:
contact_edge[i→i+1] = midpoint of overlap = (seg[i].end + seg[i+1].start) / 2
half_cf = crossfade_s / 2
seg i keep: [contact_edge[i-1→i] - half_cf, contact_edge[i→i+1] + half_cf]
expressed as sample offsets into the generated audio for that segment.
This guarantees every crossfade zone is exactly crossfade_s wide with no
raw bleed regardless of how much the inference windows overlap.
"""
def _trim(wav, start_s, end_s, seg_start_s):
"""Trim wav to [start_s, end_s] expressed in absolute video time,
where the wav starts at seg_start_s in video time."""
s = max(0, int(round((start_s - seg_start_s) * sr)))
e = int(round((end_s - seg_start_s) * sr))
e = min(e, wav.shape[1] if wav.ndim == 2 else len(wav))
return wav[:, s:e] if wav.ndim == 2 else wav[s:e]
if segments is None or len(segments) == 1:
out = wavs[0]
for nw in wavs[1:]:
out = _cf_join(out, nw, crossfade_s, db_boost, sr)
n = int(round(total_dur_s * sr))
return out[:, :n] if out.ndim == 2 else out[:n]
half_cf = crossfade_s / 2.0
# Compute contact edges between consecutive segments
contact_edges = [
(segments[i][1] + segments[i + 1][0]) / 2.0
for i in range(len(segments) - 1)
]
# Trim each segment to its keep window
trimmed = []
for i, (wav, (seg_start, seg_end)) in enumerate(zip(wavs, segments)):
keep_start = (contact_edges[i - 1] - half_cf) if i > 0 else seg_start
keep_end = (contact_edges[i] + half_cf) if i < len(segments) - 1 else total_dur_s
trimmed.append(_trim(wav, keep_start, keep_end, seg_start))
# Crossfade-join the trimmed segments
out = trimmed[0]
for nw in trimmed[1:]:
out = _cf_join(out, nw, crossfade_s, db_boost, sr)
n = int(round(total_dur_s * sr))
return out[:, :n] if out.ndim == 2 else out[:n]
def _save_wav(path: str, wav: np.ndarray, sr: int) -> None:
"""Save a numpy wav array to *path* as stereo (2, T) via soundfile.
Always conforms to stereo so every output WAV is a proper 2-channel file."""
wav = _to_stereo(wav) # (1,T) or (T,) → (2,T); (2,T) is a no-op
sf.write(path, np.ascontiguousarray(wav.T), sr)
def _log_inference_timing(label: str, elapsed: float, n_segs: int,
num_steps: int, constant: float) -> None:
"""Print a standardised inference-timing summary line."""
total_steps = n_segs * num_steps
secs_per_step = elapsed / total_steps if total_steps > 0 else 0
print(f"[{label}] Inference done: {n_segs} seg(s) × {num_steps} steps in "
f"{elapsed:.1f}s wall → {secs_per_step:.3f}s/step "
f"(current constant={constant})")
def _build_seg_meta(*, segments, wav_paths, audio_path, video_path,
silent_video, sr, model, crossfade_s, crossfade_db,
total_dur_s, source_video=None, **extras) -> dict:
"""Build the seg_meta dict shared by all three generate_* functions.
Model-specific keys are passed via **extras.
*source_video* is the original Gradio-managed upload path
(/tmp/gradio/...). It lives for the entire session and is used as a
fallback when *silent_video* (which lives in a managed tmp dir) has been
evicted or the Space has restarted since generation.
"""
meta = {
"segments": segments,
"wav_paths": wav_paths,
"audio_path": audio_path,
"video_path": video_path,
"silent_video": silent_video,
"source_video": source_video or video_path,
"sr": sr,
"model": model,
"crossfade_s": crossfade_s,
"crossfade_db": crossfade_db,
"total_dur_s": total_dur_s,
}
meta.update(extras)
return meta
def _post_process_samples(results: list, *, model: str, tmp_dir: str,
silent_video: str, segments: list,
crossfade_s: float, crossfade_db: float,
total_dur_s: float, sr: int,
source_video: str = None,
extra_meta_fn=None) -> list:
"""Shared CPU post-processing for all three generate_* wrappers.
Each entry in *results* is a tuple whose first element is a list of
per-segment wav arrays. The remaining elements are model-specific
(e.g. TARO returns features, HunyuanFoley returns text_feats).
*extra_meta_fn(sample_idx, result_tuple, tmp_dir) -> dict* is an optional
callback that returns model-specific extra keys to merge into seg_meta
(e.g. cavp_path, onset_path, text_feats_path).
Returns a list of (video_path, audio_path, seg_meta) tuples.
"""
outputs = []
print(f"[_post_process_samples] model={model} num_results={len(results)} tmp_dir={tmp_dir!r}")
for sample_idx, result in enumerate(results):
seg_wavs = result[0]
print(f"[_post_process_samples] sample {sample_idx}: seg_wavs count={len(seg_wavs) if seg_wavs else None}")
full_wav = _stitch_wavs(seg_wavs, crossfade_s, crossfade_db, total_dur_s, sr, segments)
audio_path = os.path.join(tmp_dir, f"{model}_{sample_idx}.wav")
_save_wav(audio_path, full_wav, sr)
print(f"[_post_process_samples] sample {sample_idx}: saved audio={audio_path!r} exists={os.path.exists(audio_path)}")
video_path = os.path.join(tmp_dir, f"{model}_{sample_idx}.mp4")
mux_video_audio(silent_video, audio_path, video_path, model=model)
print(f"[_post_process_samples] sample {sample_idx}: muxed video={video_path!r} exists={os.path.exists(video_path)}")
wav_paths = _save_seg_wavs(seg_wavs, tmp_dir, f"{model}_{sample_idx}")
extras = extra_meta_fn(sample_idx, result, tmp_dir) if extra_meta_fn else {}
seg_meta = _build_seg_meta(
segments=segments, wav_paths=wav_paths, audio_path=audio_path,
video_path=video_path, silent_video=silent_video, sr=sr,
model=model, crossfade_s=crossfade_s, crossfade_db=crossfade_db,
total_dur_s=total_dur_s, source_video=source_video, **extras,
)
outputs.append((video_path, audio_path, seg_meta))
print(f"[_post_process_samples] returning {len(outputs)} output(s)")
return outputs
def _cpu_preprocess(video_file: str, model_dur: float,
crossfade_s: float) -> tuple:
"""Shared CPU pre-processing for all generate_* wrappers.
Returns (tmp_dir, silent_video, total_dur_s, segments)."""
tmp_dir = _register_tmp_dir(tempfile.mkdtemp())
silent_video = os.path.join(tmp_dir, "silent_input.mp4")
strip_audio_from_video(video_file, silent_video)
total_dur_s = get_video_duration(video_file)
segments = _build_segments(total_dur_s, model_dur, crossfade_s)
return tmp_dir, silent_video, total_dur_s, segments
@spaces.GPU(duration=_taro_duration)
@_catch_oom
def _taro_gpu_infer(video_file, seed_val, cfg_scale, num_steps, mode,
crossfade_s, crossfade_db, num_samples):
"""GPU-only TARO inference — model loading + feature extraction + diffusion.
Returns list of (wavs_list, onset_feats) per sample."""
_ensure_offline_in_worker()
seed_val = _resolve_seed(seed_val)
crossfade_s = float(crossfade_s)
num_samples = int(num_samples)
torch.set_grad_enabled(False)
device, weight_dtype = _get_device_and_dtype()
_ensure_syspath("TARO")
from TARO.onset_util import extract_onset
from TARO.samplers import euler_sampler, euler_maruyama_sampler
ctx = _ctx_load("taro_gpu_infer")
print(f"[_taro_gpu_infer] ctx keys={list(ctx.keys())}")
tmp_dir = ctx["tmp_dir"]
silent_video = ctx["silent_video"]
segments = ctx["segments"]
total_dur_s = ctx["total_dur_s"]
print(f"[_taro_gpu_infer] tmp_dir={tmp_dir!r} silent_video={silent_video!r} segments={segments} total_dur_s={total_dur_s}")
print(f"[_taro_gpu_infer] calling _load_taro_feature_extractors")
extract_cavp, onset_model = _load_taro_feature_extractors(device)
print(f"[_taro_gpu_infer] extractors loaded, calling extract_cavp")
cavp_feats = extract_cavp(silent_video, tmp_path=tmp_dir)
print(f"[_taro_gpu_infer] cavp done, calling extract_onset")
# Onset features depend only on the video — extract once for all samples
onset_feats = extract_onset(silent_video, onset_model, tmp_path=tmp_dir, device=device)
print(f"[_taro_gpu_infer] onset done, freeing extractors")
# Free feature extractors before loading the heavier inference models
del extract_cavp, onset_model
if torch.cuda.is_available():
torch.cuda.empty_cache()
print(f"[_taro_gpu_infer] calling _load_taro_models")
model, vae, vocoder, latents_scale = _load_taro_models(device, weight_dtype)
print(f"[_taro_gpu_infer] models loaded")
results = [] # list of (wavs, onset_feats) per sample
for sample_idx in range(num_samples):
sample_seed = seed_val + sample_idx
cache_key = (video_file, sample_seed, float(cfg_scale), int(num_steps), mode, crossfade_s)
with _TARO_CACHE_LOCK:
cached = _TARO_INFERENCE_CACHE.get(cache_key)
if cached is not None:
print(f"[TARO] Sample {sample_idx+1}: cache hit.")
results.append((cached["wavs"], cavp_feats, None))
else:
set_global_seed(sample_seed)
wavs = []
_t_infer_start = time.perf_counter()
for seg_start_s, seg_end_s in segments:
print(f"[TARO] Sample {sample_idx+1} | {seg_start_s:.2f}s – {seg_end_s:.2f}s")
wav = _taro_infer_segment(
model, vae, vocoder,
cavp_feats, onset_feats,
seg_start_s, seg_end_s,
device, weight_dtype,
cfg_scale, num_steps, mode,
latents_scale,
euler_sampler, euler_maruyama_sampler,
)
wavs.append(wav)
_log_inference_timing("TARO", time.perf_counter() - _t_infer_start,
len(segments), int(num_steps), TARO_SECS_PER_STEP)
with _TARO_CACHE_LOCK:
_TARO_INFERENCE_CACHE[cache_key] = {"wavs": wavs}
while len(_TARO_INFERENCE_CACHE) > _TARO_CACHE_MAXLEN:
_TARO_INFERENCE_CACHE.pop(next(iter(_TARO_INFERENCE_CACHE)))
results.append((wavs, cavp_feats, onset_feats))
# Free GPU memory between samples so VRAM fragmentation doesn't
# degrade diffusion quality on samples 2, 3, 4, etc.
if torch.cuda.is_available():
torch.cuda.empty_cache()
print(f"[_taro_gpu_infer] returning {len(results)} results")
return results
def generate_taro(video_file, seed_val, cfg_scale, num_steps, mode,
crossfade_s, crossfade_db, num_samples):
"""TARO: video-conditioned diffusion, 16 kHz, 8.192 s sliding window.
CPU pre/post-processing wraps the GPU-only inference to minimize ZeroGPU cost."""
print(f"[generate_taro] START video_file={video_file!r} num_samples={num_samples}")
crossfade_s = float(crossfade_s)
crossfade_db = float(crossfade_db)
num_samples = int(num_samples)
# ── CPU pre-processing (no GPU needed) ──
tmp_dir, silent_video, total_dur_s, segments = _cpu_preprocess(
video_file, TARO_MODEL_DUR, crossfade_s)
print(f"[generate_taro] preprocess done: total_dur_s={total_dur_s:.2f} segments={segments} tmp_dir={tmp_dir}")
_ctx_store("taro_gpu_infer", {
"tmp_dir": tmp_dir, "silent_video": silent_video,
"segments": segments, "total_dur_s": total_dur_s,
})
# ── GPU inference only ──
results = _taro_gpu_infer(video_file, seed_val, cfg_scale, num_steps, mode,
crossfade_s, crossfade_db, num_samples)
# ── CPU post-processing (no GPU needed) ──
# Upsample 16kHz → 48kHz and normalise result tuples to (seg_wavs, ...)
cavp_path = os.path.join(tmp_dir, "taro_cavp.npy")
onset_path = os.path.join(tmp_dir, "taro_onset.npy")
_feats_saved = False
def _upsample_and_save_feats(result):
nonlocal _feats_saved
wavs, cavp_feats, onset_feats = result
wavs = [_upsample_taro(w) for w in wavs]
if not _feats_saved:
np.save(cavp_path, cavp_feats)
if onset_feats is not None:
np.save(onset_path, onset_feats)
_feats_saved = True
return (wavs, cavp_feats, onset_feats)
print(f"[generate_taro] gpu_infer returned {len(results)} result(s)")
results = [_upsample_and_save_feats(r) for r in results]
print(f"[generate_taro] upsample done, cavp_path={cavp_path} onset_path={onset_path}")
def _taro_extras(sample_idx, result, td):
return {"cavp_path": cavp_path, "onset_path": onset_path}
outputs = _post_process_samples(
results, model="taro", tmp_dir=tmp_dir,
silent_video=silent_video, segments=segments,
crossfade_s=crossfade_s, crossfade_db=crossfade_db,
total_dur_s=total_dur_s, sr=TARO_SR_OUT,
source_video=video_file,
extra_meta_fn=_taro_extras,
)
print(f"[generate_taro] post_process done: {len(outputs)} output(s)")
padded = _pad_outputs(outputs)
print(f"[generate_taro] padded outputs: {[(type(x).__name__, x is not None) for x in padded[:6]]}")
return padded
# ================================================================== #
# MMAudio #
# ================================================================== #
# Constants sourced from MMAudio/mmaudio/model/sequence_config.py:
# CONFIG_44K: duration=8.0 s, sampling_rate=44100
# CLIP encoder: 8 fps, 384×384 px
# Synchformer: 25 fps, 224×224 px
# Default variant: large_44k_v2
# MMAudio uses flow-matching (FlowMatching with euler inference).
# generate() handles all feature extraction + decoding internally.
# ================================================================== #
def _mmaudio_duration(video_file, prompt, negative_prompt, seed_val,
cfg_strength, num_steps, crossfade_s, crossfade_db, num_samples,
silent_video=None, segments_json=None,
clip_start_s=0.0, clip_dur_s=None, **_kwargs):
"""Pre-GPU callable — must match _mmaudio_gpu_infer's input signature exactly.
silent_video, segments_json, clip_start_s, clip_dur_s are extra positional args
that xregen passes; they must appear here so ZeroGPU doesn't raise TypeError
when forwarding all args to this duration fn before the GPU fn runs."""
return _estimate_gpu_duration("mmaudio", int(num_samples), int(num_steps),
video_file=video_file, crossfade_s=crossfade_s)
@spaces.GPU(duration=_mmaudio_duration)
@_catch_oom
def _mmaudio_gpu_infer(video_file, prompt, negative_prompt, seed_val,
cfg_strength, num_steps, crossfade_s, crossfade_db, num_samples,
silent_video, segments_json,
clip_start_s=0.0, clip_dur_s=None):
"""GPU-only MMAudio inference — model loading + flow-matching generation.
Returns list of (seg_audios, sr) per sample.
All video paths and segment data are passed explicitly as positional args
to survive ZeroGPU process isolation (kwargs are silently dropped).
When *clip_dur_s* is set, *silent_video* is the full source and a clip
[clip_start_s, clip_start_s+clip_dur_s] is extracted first inside the
GPU window (ffmpeg is CPU-safe here). This avoids passing pre-extracted
tmp files that don't exist in the GPU worker's process.
"""
print(f"[_mmaudio_gpu_infer] START video={video_file!r} silent={silent_video!r} clip_start={clip_start_s} clip_dur={clip_dur_s} num_samples={num_samples}")
_ensure_offline_in_worker()
_ensure_syspath("MMAudio")
from mmaudio.eval_utils import generate, load_video
from mmaudio.model.flow_matching import FlowMatching
seed_val = _resolve_seed(seed_val)
num_samples = int(num_samples)
num_steps = int(num_steps)
crossfade_s = float(crossfade_s)
device, dtype = _get_device_and_dtype()
net, feature_utils, model_cfg, seq_cfg = _load_mmaudio_models(device, dtype)
# If a clip window is specified, extract it now (inside the GPU fn, so the
# file exists in this worker's /tmp).
tmp_dir = _register_tmp_dir(tempfile.mkdtemp())
if clip_dur_s is not None:
clip_dur_s = float(clip_dur_s)
clip_path = _extract_segment_clip(
silent_video, float(clip_start_s), clip_dur_s,
os.path.join(tmp_dir, "mma_xregen_clip.mp4"),
)
silent_video = clip_path
# Extract per-segment clips from silent_video (now the correct clip source).
segments = json.loads(segments_json)
seg_clip_paths = [
_extract_segment_clip(silent_video, s, e - s,
os.path.join(tmp_dir, f"mma_seg_{i}.mp4"))
for i, (s, e) in enumerate(segments)
]
sr = seq_cfg.sampling_rate # 44100
results = []
for sample_idx in range(num_samples):
rng = torch.Generator(device=device)
rng.manual_seed(seed_val + sample_idx)
seg_audios = []
_t_mma_start = time.perf_counter()
fm = FlowMatching(min_sigma=0, inference_mode="euler", num_steps=num_steps)
for seg_i, (seg_start, seg_end) in enumerate(segments):
seg_dur = seg_end - seg_start
seg_path = seg_clip_paths[seg_i]
video_info = load_video(seg_path, seg_dur)
clip_frames = video_info.clip_frames.unsqueeze(0)
sync_frames = video_info.sync_frames.unsqueeze(0)
actual_dur = video_info.duration_sec
seq_cfg.duration = actual_dur
net.update_seq_lengths(seq_cfg.latent_seq_len, seq_cfg.clip_seq_len, seq_cfg.sync_seq_len)
print(f"[MMAudio] Sample {sample_idx+1} | seg {seg_i+1}/{len(segments)} "
f"{seg_start:.1f}–{seg_end:.1f}s | dur={actual_dur:.2f}s | prompt='{prompt}'")
with torch.no_grad():
audios = generate(
clip_frames,
sync_frames,
[prompt],
negative_text=[negative_prompt] if negative_prompt else None,
feature_utils=feature_utils,
net=net,
fm=fm,
rng=rng,
cfg_strength=float(cfg_strength),
)
wav = audios.float().cpu()[0].numpy() # (C, T) — full window
seg_audios.append(wav)
_log_inference_timing("MMAudio", time.perf_counter() - _t_mma_start,
len(segments), int(num_steps), MMAUDIO_SECS_PER_STEP)
results.append((seg_audios, sr))
# Free GPU memory between samples to prevent VRAM fragmentation
if torch.cuda.is_available():
torch.cuda.empty_cache()
# net + feature_utils (CLIP + Synchformer + VAE) are ~9.5GB combined and,
# unlike TARO/HunyuanFoley, were never explicitly freed here — do it before
# returning so the ZeroGPU worker isn't left holding peak memory longer
# than necessary.
del net, feature_utils
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
return results
def generate_mmaudio(video_file, prompt, negative_prompt, seed_val,
cfg_strength, num_steps, crossfade_s, crossfade_db, num_samples):
"""MMAudio: flow-matching video-to-audio, 44.1 kHz, 8 s sliding window.
CPU pre/post-processing wraps the GPU-only inference to minimize ZeroGPU cost."""
num_samples = int(num_samples)
crossfade_s = float(crossfade_s)
crossfade_db = float(crossfade_db)
# ── CPU pre-processing ──
tmp_dir, silent_video, total_dur_s, segments = _cpu_preprocess(
video_file, MMAUDIO_WINDOW, crossfade_s)
print(f"[MMAudio] Video={total_dur_s:.2f}s | {len(segments)} segment(s) × ≤8 s")
# ── GPU inference only ──
results = _mmaudio_gpu_infer(video_file, prompt, negative_prompt, seed_val,
cfg_strength, num_steps, crossfade_s, crossfade_db,
num_samples,
silent_video, json.dumps(segments))
# ── CPU post-processing ──
# Resample 44100 → 48000 and normalise tuples to (seg_wavs, ...)
resampled = []
for seg_audios, sr in results:
if sr != TARGET_SR:
print(f"[MMAudio upsample] resampling {sr}Hz → {TARGET_SR}Hz (sinc, CPU) …")
seg_audios = [_resample_to_target(w, sr) for w in seg_audios]
print(f"[MMAudio upsample] done — {len(seg_audios)} seg(s) @ {TARGET_SR}Hz")
resampled.append((seg_audios,))
outputs = _post_process_samples(
resampled, model="mmaudio", tmp_dir=tmp_dir,
silent_video=silent_video, segments=segments,
crossfade_s=crossfade_s, crossfade_db=crossfade_db,
total_dur_s=total_dur_s, sr=TARGET_SR,
source_video=video_file,
)
return _pad_outputs(outputs)
# ================================================================== #
# HunyuanVideoFoley #
# ================================================================== #
# Constants sourced from HunyuanVideo-Foley/hunyuanvideo_foley/constants.py
# and configs/hunyuanvideo-foley-xxl.yaml:
# sample_rate = 48000 Hz (from DAC VAE)
# audio_frame_rate = 50 (latent fps, xxl config)
# max video duration = 15 s
# SigLIP2 fps = 8, Synchformer fps = 25
# CLAP text encoder: laion/larger_clap_general (auto-downloaded from HF Hub)
# Default guidance_scale=4.5, num_inference_steps=50
# ================================================================== #
def _hunyuan_duration(video_file, prompt, negative_prompt, seed_val,
guidance_scale, num_steps, model_size, crossfade_s, crossfade_db,
num_samples, silent_video=None, segments_json=None, total_dur_s=None,
clip_start_s=0.0, clip_dur_s=None, **_kwargs):
"""Pre-GPU callable — must match _hunyuan_gpu_infer's input signature exactly.
silent_video, segments_json, total_dur_s, clip_start_s, clip_dur_s are extra
positional args that xregen passes; they must appear here so ZeroGPU doesn't
raise TypeError when forwarding all args to this duration fn."""
return _estimate_gpu_duration("hunyuan", int(num_samples), int(num_steps),
video_file=video_file, crossfade_s=crossfade_s)
@spaces.GPU(duration=_hunyuan_duration)
@_catch_oom
def _hunyuan_gpu_infer(video_file, prompt, negative_prompt, seed_val,
guidance_scale, num_steps, model_size, crossfade_s, crossfade_db,
num_samples, silent_video, segments_json, total_dur_s,
clip_start_s=0.0, clip_dur_s=None):
"""GPU-only HunyuanFoley inference — model loading + feature extraction + denoising.
All segments processed in a single GPU call (Pro ZeroGPU allows up to 300 s).
"""
import traceback as _tb
print(f"[_hunyuan_gpu_infer] START video={video_file!r} model_size={model_size!r} "
f"num_steps={num_steps!r} clip_start={clip_start_s} clip_dur={clip_dur_s}")
_ensure_offline_in_worker()
try:
_ensure_syspath("HunyuanVideo-Foley")
from hunyuanvideo_foley.utils.model_utils import denoise_process
from hunyuanvideo_foley.utils.feature_utils import feature_process, encode_video_features
seed_val = _resolve_seed(seed_val)
num_samples = int(num_samples)
crossfade_s = float(crossfade_s)
total_dur_s = float(total_dur_s)
set_global_seed(seed_val)
device, _ = _get_device_and_dtype()
model_dict, cfg = _load_hunyuan_model(device, model_size.lower())
tmp_dir = _register_tmp_dir(tempfile.mkdtemp())
_sv = silent_video
_total = total_dur_s
if clip_dur_s is not None:
clip_path = _extract_segment_clip(
silent_video, float(clip_start_s), float(clip_dur_s),
os.path.join(tmp_dir, "hny_xregen_clip.mp4"),
)
_sv = clip_path
_total = float(clip_dur_s)
segments = json.loads(segments_json)
dummy_seg_path = _extract_segment_clip(
_sv, 0, min(_total, HUNYUAN_MAX_DUR),
os.path.join(tmp_dir, "_seg_dummy.mp4"),
)
seg_clip_paths = [
_extract_segment_clip(_sv, s, e - s,
os.path.join(tmp_dir, f"hny_seg_{i}.mp4"))
for i, (s, e) in enumerate(segments)
]
# Extract text features once for all segments
_, text_feats, _ = feature_process(
dummy_seg_path,
prompt if prompt else "",
model_dict, cfg,
neg_prompt=negative_prompt if negative_prompt else None,
)
results = []
for sample_idx in range(num_samples):
seg_wavs = []
sr = 48000
_t0 = time.perf_counter()
for seg_i, (seg_start, seg_end) in enumerate(segments):
visual_feats, seg_audio_len = encode_video_features(
seg_clip_paths[seg_i], model_dict)
print(f"[HunyuanFoley] Sample {sample_idx+1} | seg {seg_i+1}/{len(segments)} "
f"{seg_start:.1f}–{seg_end:.1f}s → {seg_audio_len:.2f}s audio")
audio_batch, sr = denoise_process(
visual_feats, text_feats, seg_audio_len, model_dict, cfg,
guidance_scale=float(guidance_scale),
num_inference_steps=int(num_steps),
batch_size=1,
)
seg_wavs.append(audio_batch[0].float().cpu().numpy())
# Free GPU memory between segments — latents/visual_feats from denoise_process
# stay allocated until GC runs; explicit deletion + cache clear prevents OOM
# when processing a second segment (the 15-s latent tensor is ~several GB).
del audio_batch, visual_feats
if torch.cuda.is_available():
torch.cuda.empty_cache()
_log_inference_timing("HunyuanFoley", time.perf_counter() - _t0,
len(segments), int(num_steps), HUNYUAN_SECS_PER_STEP)
# Save text_feats inside the GPU worker — never return CUDA tensors
text_feats_path = os.path.join(tmp_dir, f"hunyuan_{sample_idx}_text_feats.pt")
torch.save(text_feats, text_feats_path)
print(f"[HunyuanFoley] text_feats saved to {text_feats_path}")
results.append((seg_wavs, sr, text_feats_path))
if torch.cuda.is_available():
torch.cuda.empty_cache()
return results
except Exception as _e:
print(f"[_hunyuan_gpu_infer] EXCEPTION: {_e}")
_tb.print_exc()
raise
def generate_hunyuan(video_file, prompt, negative_prompt, seed_val,
guidance_scale, num_steps, model_size, crossfade_s, crossfade_db, num_samples):
"""HunyuanVideoFoley: text-guided foley, 48 kHz, up to 15 s.
CPU pre/post-processing wraps the GPU-only inference to minimize ZeroGPU cost."""
num_samples = int(num_samples)
crossfade_s = float(crossfade_s)
crossfade_db = float(crossfade_db)
# ── CPU pre-processing (no GPU needed) ──
tmp_dir, silent_video, total_dur_s, segments = _cpu_preprocess(
video_file, HUNYUAN_MAX_DUR, crossfade_s)
print(f"[HunyuanFoley] Video={total_dur_s:.2f}s | {len(segments)} segment(s) × ≤15 s")
# ── GPU inference (all segments in one call) ──
results = _hunyuan_gpu_infer(video_file, prompt, negative_prompt, seed_val,
guidance_scale, num_steps, model_size,
crossfade_s, crossfade_db, num_samples,
silent_video, json.dumps(segments), total_dur_s)
# ── CPU post-processing (no GPU needed) ──
def _hunyuan_extras(sample_idx, result, td):
_, _sr, text_feats_path = result
return {"text_feats_path": text_feats_path}
outputs = _post_process_samples(
results, model="hunyuan", tmp_dir=tmp_dir,
silent_video=silent_video, segments=segments,
crossfade_s=crossfade_s, crossfade_db=crossfade_db,
total_dur_s=total_dur_s, sr=48000,
source_video=video_file,
extra_meta_fn=_hunyuan_extras,
)
return _pad_outputs(outputs)
# ================================================================== #
# SEGMENT REGENERATION HELPERS #
# ================================================================== #
# Each regen function:
# 1. Runs inference for ONE segment (random seed, current settings)
# 2. Splices the new wav into the stored wavs list
# 3. Re-stitches the full track, re-saves .wav and re-muxes .mp4
# 4. Returns (new_video_path, new_audio_path, updated_seg_meta, new_waveform_html)
# ================================================================== #
def _splice_and_save(new_wav, seg_idx, meta, slot_id):
"""Replace wavs[seg_idx] with new_wav, re-stitch, re-save, re-mux.
Returns (video_path, audio_path, updated_meta, waveform_html).
All wavs (loaded from disk and new_wav) must be stereo (2, T).
"""
assert new_wav.ndim == 2 and new_wav.shape[0] == 2, \
f"new_wav must be stereo (2, T), got shape {new_wav.shape}"
wavs = _load_seg_wavs(meta["wav_paths"])
wavs[seg_idx]= new_wav
crossfade_s = float(meta["crossfade_s"])
crossfade_db = float(meta["crossfade_db"])
sr = int(meta["sr"])
total_dur_s = float(meta["total_dur_s"])
silent_video = _resolve_silent_video(meta)
segments = meta["segments"]
model = meta["model"]
full_wav = _stitch_wavs(wavs, crossfade_s, crossfade_db, total_dur_s, sr, segments)
# Save new audio — use a new timestamped filename so Gradio / the browser
# treats it as a genuinely different file and reloads the video player.
_ts = int(time.time() * 1000)
tmp_dir = os.path.dirname(meta["audio_path"])
_base = os.path.splitext(os.path.basename(meta["audio_path"]))[0]
# Strip any previous timestamp suffix before adding a new one
_base_clean = _base.rsplit("_regen_", 1)[0]
audio_path = os.path.join(tmp_dir, f"{_base_clean}_regen_{_ts}.wav")
_save_wav(audio_path, full_wav, sr)
# Re-mux into a new video file so the browser is forced to reload it
_vid_base = os.path.splitext(os.path.basename(meta["video_path"]))[0]
_vid_base_clean = _vid_base.rsplit("_regen_", 1)[0]
video_path = os.path.join(tmp_dir, f"{_vid_base_clean}_regen_{_ts}.mp4")
mux_video_audio(silent_video, audio_path, video_path, model=model)
# Save updated segment wavs to .npy files
updated_wav_paths = _save_seg_wavs(wavs, tmp_dir, os.path.splitext(_base_clean)[0])
updated_meta = dict(meta)
updated_meta["wav_paths"] = updated_wav_paths
updated_meta["audio_path"] = audio_path
updated_meta["video_path"] = video_path
state_json_new = json.dumps(updated_meta)
waveform_html = _build_waveform_html(audio_path, segments, slot_id, "",
state_json=state_json_new,
video_path=video_path,
crossfade_s=crossfade_s)
return video_path, audio_path, updated_meta, waveform_html
def _taro_regen_duration(video_file, seg_idx, seg_meta_json,
seed_val, cfg_scale, num_steps, mode,
crossfade_s, crossfade_db, slot_id=None):
# If cached CAVP/onset features exist, skip ~10s feature-extractor overhead
try:
meta = json.loads(seg_meta_json)
cavp_ok = os.path.exists(meta.get("cavp_path", ""))
onset_ok = os.path.exists(meta.get("onset_path", ""))
if cavp_ok and onset_ok:
cfg = MODEL_CONFIGS["taro"]
secs = int(num_steps) * cfg["secs_per_step"] + 5 # 5s model-load only
result = min(GPU_DURATION_CAP, max(30, int(secs)))
print(f"[duration] TARO regen (cache hit): 1 seg × {int(num_steps)} steps → {secs:.0f}s → capped {result}s")
return result
except Exception:
pass
return _estimate_regen_duration("taro", int(num_steps))
@spaces.GPU(duration=_taro_regen_duration)
@_catch_oom
def _regen_taro_gpu(video_file, seg_idx, seg_meta_json,
seed_val, cfg_scale, num_steps, mode,
crossfade_s, crossfade_db, slot_id=None):
"""GPU-only TARO regen — returns new_wav for a single segment."""
_ensure_offline_in_worker()
meta = json.loads(seg_meta_json)
seg_idx = int(seg_idx)
seg_start_s, seg_end_s = meta["segments"][seg_idx]
torch.set_grad_enabled(False)
device, weight_dtype = _get_device_and_dtype()
_ensure_syspath("TARO")
from TARO.samplers import euler_sampler, euler_maruyama_sampler
# Load cached CAVP/onset features from .npy files (CPU I/O, fast, outside GPU budget)
cavp_path = meta.get("cavp_path", "")
onset_path = meta.get("onset_path", "")
if cavp_path and os.path.exists(cavp_path) and onset_path and os.path.exists(onset_path):
print("[TARO regen] Loading cached CAVP + onset features from disk")
cavp_feats = np.load(cavp_path)
onset_feats = np.load(onset_path)
else:
print("[TARO regen] Cache miss — re-extracting CAVP + onset features")
from TARO.onset_util import extract_onset
extract_cavp, onset_model = _load_taro_feature_extractors(device)
silent_video = _resolve_silent_video(meta)
tmp_dir = _register_tmp_dir(tempfile.mkdtemp())
cavp_feats = extract_cavp(silent_video, tmp_path=tmp_dir)
onset_feats = extract_onset(silent_video, onset_model, tmp_path=tmp_dir, device=device)
del extract_cavp, onset_model
if torch.cuda.is_available():
torch.cuda.empty_cache()
model_net, vae, vocoder, latents_scale = _load_taro_models(device, weight_dtype)
set_global_seed(random.randint(0, 2**32 - 1))
return _taro_infer_segment(
model_net, vae, vocoder, cavp_feats, onset_feats,
seg_start_s, seg_end_s, device, weight_dtype,
float(cfg_scale), int(num_steps), mode, latents_scale,
euler_sampler, euler_maruyama_sampler,
)
def regen_taro_segment(video_file, seg_idx, seg_meta_json,
seed_val, cfg_scale, num_steps, mode,
crossfade_s, crossfade_db, slot_id):
"""Regenerate one TARO segment. GPU inference + CPU splice/save."""
meta = json.loads(seg_meta_json)
seg_idx = int(seg_idx)
# GPU: inference — CAVP/onset features loaded from disk paths in seg_meta_json
new_wav = _regen_taro_gpu(video_file, seg_idx, seg_meta_json,
seed_val, cfg_scale, num_steps, mode,
crossfade_s, crossfade_db, slot_id)
# Upsample 16kHz → 48kHz (sinc, CPU)
new_wav = _upsample_taro(new_wav)
# Conform to stereo (2, T) — TARO returns mono (T,)
new_wav = _to_stereo(new_wav)
# CPU: splice, stitch, mux, save
video_path, audio_path, updated_meta, waveform_html = _splice_and_save(
new_wav, seg_idx, meta, slot_id
)
return video_path, audio_path, json.dumps(updated_meta), waveform_html
def _mmaudio_regen_duration(video_file, seg_idx, seg_meta_json,
prompt, negative_prompt, seed_val,
cfg_strength, num_steps, crossfade_s, crossfade_db,
slot_id=None):
return _estimate_regen_duration("mmaudio", int(num_steps))
@spaces.GPU(duration=_mmaudio_regen_duration)
@_catch_oom
def _regen_mmaudio_gpu(video_file, seg_idx, seg_meta_json,
prompt, negative_prompt, seed_val,
cfg_strength, num_steps, crossfade_s, crossfade_db,
slot_id=None):
"""GPU-only MMAudio regen — returns (new_wav, sr) for a single segment."""
_ensure_offline_in_worker()
meta = json.loads(seg_meta_json)
seg_idx = int(seg_idx)
seg_start, seg_end = meta["segments"][seg_idx]
seg_dur = seg_end - seg_start
_ensure_syspath("MMAudio")
from mmaudio.eval_utils import generate, load_video
from mmaudio.model.flow_matching import FlowMatching
device, dtype = _get_device_and_dtype()
net, feature_utils, model_cfg, seq_cfg = _load_mmaudio_models(device, dtype)
sr = seq_cfg.sampling_rate
# Extract segment clip inside the GPU function — ffmpeg is CPU-only and safe here.
# This avoids any cross-process context passing that fails under ZeroGPU isolation.
seg_path = _extract_segment_clip(
_resolve_silent_video(meta), seg_start, seg_dur,
os.path.join(_register_tmp_dir(tempfile.mkdtemp()), "regen_seg.mp4"),
)
rng = torch.Generator(device=device)
rng.manual_seed(random.randint(0, 2**32 - 1))
fm = FlowMatching(min_sigma=0, inference_mode="euler", num_steps=int(num_steps))
video_info = load_video(seg_path, seg_dur)
clip_frames = video_info.clip_frames.unsqueeze(0)
sync_frames = video_info.sync_frames.unsqueeze(0)
actual_dur = video_info.duration_sec
seq_cfg.duration = actual_dur
net.update_seq_lengths(seq_cfg.latent_seq_len, seq_cfg.clip_seq_len, seq_cfg.sync_seq_len)
with torch.no_grad():
audios = generate(
clip_frames, sync_frames, [prompt],
negative_text=[negative_prompt] if negative_prompt else None,
feature_utils=feature_utils, net=net, fm=fm, rng=rng,
cfg_strength=float(cfg_strength),
)
new_wav = audios.float().cpu()[0].numpy() # full window — _stitch_wavs trims
# See _mmaudio_gpu_infer — net + feature_utils (~9.5GB) had no explicit
# cleanup here before returning to the ZeroGPU worker.
del net, feature_utils
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
return new_wav, sr
def regen_mmaudio_segment(video_file, seg_idx, seg_meta_json,
prompt, negative_prompt, seed_val,
cfg_strength, num_steps, crossfade_s, crossfade_db, slot_id):
"""Regenerate one MMAudio segment. GPU inference + CPU splice/save."""
meta = json.loads(seg_meta_json)
seg_idx = int(seg_idx)
# GPU: inference (segment clip extraction happens inside the GPU function)
new_wav, sr = _regen_mmaudio_gpu(video_file, seg_idx, seg_meta_json,
prompt, negative_prompt, seed_val,
cfg_strength, num_steps, crossfade_s, crossfade_db,
slot_id)
# Resample to 48kHz if needed (MMAudio outputs at 44100 Hz)
if sr != TARGET_SR:
print(f"[MMAudio regen upsample] {sr}Hz → {TARGET_SR}Hz (sinc, CPU) …")
new_wav = _resample_to_target(new_wav, sr)
sr = TARGET_SR
meta["sr"] = sr
# Conform to stereo (2, T) — MMAudio returns (2,T) but apply defensively
new_wav = _to_stereo(new_wav)
# CPU: splice, stitch, mux, save
video_path, audio_path, updated_meta, waveform_html = _splice_and_save(
new_wav, seg_idx, meta, slot_id
)
return video_path, audio_path, json.dumps(updated_meta), waveform_html
def _hunyuan_regen_duration(video_file, seg_idx, seg_meta_json,
prompt, negative_prompt, seed_val,
guidance_scale, num_steps, model_size,
crossfade_s, crossfade_db, slot_id=None):
return _estimate_regen_duration("hunyuan", int(num_steps))
@spaces.GPU(duration=_hunyuan_regen_duration)
@_catch_oom
def _regen_hunyuan_gpu(video_file, seg_idx, seg_meta_json,
prompt, negative_prompt, seed_val,
guidance_scale, num_steps, model_size,
crossfade_s, crossfade_db, slot_id=None):
"""GPU-only HunyuanFoley regen — returns (new_wav, sr) for a single segment."""
_ensure_offline_in_worker()
meta = json.loads(seg_meta_json)
seg_idx = int(seg_idx)
seg_start, seg_end = meta["segments"][seg_idx]
seg_dur = seg_end - seg_start
_ensure_syspath("HunyuanVideo-Foley")
from hunyuanvideo_foley.utils.model_utils import denoise_process
from hunyuanvideo_foley.utils.feature_utils import feature_process
device, _ = _get_device_and_dtype()
model_dict, cfg = _load_hunyuan_model(device, model_size)
set_global_seed(random.randint(0, 2**32 - 1))
# Extract segment clip inside the GPU function — ffmpeg is CPU-only and safe here.
seg_path = _extract_segment_clip(
_resolve_silent_video(meta), seg_start, seg_dur,
os.path.join(_register_tmp_dir(tempfile.mkdtemp()), "regen_seg.mp4"),
)
text_feats_path = meta.get("text_feats_path", "")
if text_feats_path and os.path.exists(text_feats_path):
print("[HunyuanFoley regen] Loading cached text features from disk")
from hunyuanvideo_foley.utils.feature_utils import encode_video_features
visual_feats, seg_audio_len = encode_video_features(seg_path, model_dict)
text_feats = torch.load(text_feats_path, map_location=device, weights_only=False)
else:
print("[HunyuanFoley regen] Cache miss — extracting text + visual features")
visual_feats, text_feats, seg_audio_len = feature_process(
seg_path, prompt if prompt else "", model_dict, cfg,
neg_prompt=negative_prompt if negative_prompt else None,
)
audio_batch, sr = denoise_process(
visual_feats, text_feats, seg_audio_len, model_dict, cfg,
guidance_scale=float(guidance_scale),
num_inference_steps=int(num_steps),
batch_size=1,
)
new_wav = audio_batch[0].float().cpu().numpy() # full window — _stitch_wavs trims
return new_wav, sr
def regen_hunyuan_segment(video_file, seg_idx, seg_meta_json,
prompt, negative_prompt, seed_val,
guidance_scale, num_steps, model_size,
crossfade_s, crossfade_db, slot_id):
"""Regenerate one HunyuanFoley segment. GPU inference + CPU splice/save."""
meta = json.loads(seg_meta_json)
seg_idx = int(seg_idx)
# GPU: inference (segment clip extraction happens inside the GPU function)
new_wav, sr = _regen_hunyuan_gpu(video_file, seg_idx, seg_meta_json,
prompt, negative_prompt, seed_val,
guidance_scale, num_steps, model_size,
crossfade_s, crossfade_db, slot_id)
meta["sr"] = sr
# Conform to stereo (2, T) — HunyuanFoley DAC returns (1, T)
new_wav = _to_stereo(new_wav)
# CPU: splice, stitch, mux, save
video_path, audio_path, updated_meta, waveform_html = _splice_and_save(
new_wav, seg_idx, meta, slot_id
)
return video_path, audio_path, json.dumps(updated_meta), waveform_html
# Wire up regen_fn references now that the functions are defined
MODEL_CONFIGS["taro"]["regen_fn"] = regen_taro_segment
MODEL_CONFIGS["mmaudio"]["regen_fn"] = regen_mmaudio_segment
MODEL_CONFIGS["hunyuan"]["regen_fn"] = regen_hunyuan_segment
# ================================================================== #
# CROSS-MODEL REGEN WRAPPERS #
# ================================================================== #
# Three shared endpoints — one per model — that can be called from #
# *any* slot tab. slot_id is passed as plain string data so the #
# result is applied back to the correct slot by the JS listener. #
# The new segment is resampled to match the slot's existing SR before #
# being handed to _splice_and_save, so TARO (16 kHz) / MMAudio #
# (44.1 kHz) / Hunyuan (48 kHz) outputs can all be mixed freely. #
# ================================================================== #
def _resample_to_slot_sr(wav: np.ndarray, src_sr: int, dst_sr: int,
slot_wav_ref: np.ndarray = None) -> np.ndarray:
"""Resample *wav* from src_sr to dst_sr and convert to stereo (2, T).
slot_wav_ref is unused (kept for call-site compatibility) — all wavs
are now always stereo so no per-slot channel matching is needed."""
wav = _resample_to_target(wav, src_sr, dst_sr)
return _to_stereo(wav)
def _resolve_silent_video(meta: dict) -> str:
"""Return a valid silent (audio-stripped) video path for *meta*.
Prefers meta["silent_video"] if it still exists on disk. Falls back to
re-stripping meta["source_video"] (the original Gradio upload path, which
persists for the full session lifetime) into a fresh tmp file.
This prevents xregen failures caused by tmp-dir eviction or Space restarts
between initial generation and the regen call.
"""
sv = meta.get("silent_video", "")
if sv and os.path.exists(sv):
return sv
source = meta.get("source_video") or meta.get("video_path", "")
if not source or not os.path.exists(source):
raise FileNotFoundError(
f"Cannot locate source video for regen — "
f"silent_video={sv!r}, source_video={source!r}"
)
out = os.path.join(_register_tmp_dir(tempfile.mkdtemp()), "silent_input.mp4")
print(f"[regen] silent_video missing, re-stripping from source_video: {source}")
strip_audio_from_video(source, out)
return out
def _xregen_clip_window(meta: dict, seg_idx: int, target_window_s: float) -> tuple:
"""Compute the video clip window for a cross-model regen.
Always returns the full original segment span (clamped to [0, total_dur_s]).
The caller passes clip_dur to _build_segments, which splits it into
model-sized sub-windows exactly as initial generation does — ensuring the
regen wav covers the full segment duration regardless of model window size.
*target_window_s* is unused but kept for call-site compatibility.
"""
total_dur_s = float(meta["total_dur_s"])
seg_start, seg_end = meta["segments"][seg_idx]
clip_start = max(0.0, seg_start)
clip_end = min(total_dur_s, seg_end)
clip_dur = clip_end - clip_start
return clip_start, clip_end, clip_dur
def _xregen_splice(new_wav_raw: np.ndarray, src_sr: int,
meta: dict, seg_idx: int, slot_id: str,
clip_start_s: float = None) -> tuple:
"""Shared epilogue for all xregen_* functions: resample → splice → save.
Returns (video_path, waveform_html).
*clip_start_s* is the absolute video time where new_wav_raw starts.
When the clip was centered on the segment midpoint (not at seg_start),
we need to shift the wav so _stitch_wavs can trim it correctly relative
to the original segment's start. We do this by prepending silence so
the wav's time origin aligns with the original segment's start.
"""
slot_sr = int(meta["sr"])
slot_wavs = _load_seg_wavs(meta["wav_paths"])
new_wav = _resample_to_slot_sr(new_wav_raw, src_sr, slot_sr, slot_wavs[0])
# Align new_wav so sample index 0 corresponds to seg_start in video time.
# _stitch_wavs trims using seg_start as the time origin, so if the clip
# started AFTER seg_start (clip_start_s > seg_start), we prepend silence
# equal to (clip_start_s - seg_start) to shift the audio back to seg_start.
if clip_start_s is not None:
seg_start = meta["segments"][seg_idx][0]
offset_s = seg_start - clip_start_s # negative when clip starts after seg_start
if offset_s < 0:
pad_samples = int(round(abs(offset_s) * slot_sr))
silence = np.zeros((new_wav.shape[0], pad_samples), dtype=new_wav.dtype)
new_wav = np.concatenate([silence, new_wav], axis=1)
video_path, audio_path, updated_meta, waveform_html = _splice_and_save(
new_wav, seg_idx, meta, slot_id
)
return video_path, waveform_html
def _xregen_dispatch(state_json: str, seg_idx: int, slot_id: str, infer_fn):
"""Shared generator skeleton for all xregen_* wrappers.
Yields pending HTML immediately, then calls *infer_fn()* — a zero-argument
callable that runs model-specific CPU prep + GPU inference and returns
(wav_array, src_sr, clip_start_s). For TARO, *infer_fn* should return
the wav already upsampled to 48 kHz; pass TARO_SR_OUT as src_sr.
Yields:
First: (gr.update(), gr.update(value=pending_html)) — shown while GPU runs
Second: (gr.update(value=video_path), gr.update(value=waveform_html))
"""
import traceback as _tb
meta = json.loads(state_json)
pending_html = _build_regen_pending_html(meta["segments"], seg_idx, slot_id, "")
yield gr.update(), gr.update(value=pending_html)
print(f"[_xregen_dispatch] slot={slot_id} seg={seg_idx} calling infer_fn={infer_fn}")
try:
new_wav_raw, src_sr, clip_start_s = infer_fn()
print(f"[_xregen_dispatch] infer_fn returned wav shape={getattr(new_wav_raw,'shape',None)} sr={src_sr} clip_start={clip_start_s}")
except Exception as _e:
print(f"[_xregen_dispatch] EXCEPTION in infer_fn: {_e}")
_tb.print_exc()
raise
video_path, waveform_html = _xregen_splice(new_wav_raw, src_sr, meta, seg_idx, slot_id, clip_start_s)
print(f"[_xregen_dispatch] splice done video_path={video_path!r}")
yield gr.update(value=video_path), gr.update(value=waveform_html)
def xregen_taro(seg_idx, state_json, slot_id,
seed_val, cfg_scale, num_steps, mode,
crossfade_s, crossfade_db,
request: gr.Request = None):
"""Cross-model regen: run TARO on its optimal window, splice into *slot_id*."""
seg_idx = int(seg_idx)
meta = json.loads(state_json)
def _run():
clip_start, clip_end, clip_dur = _xregen_clip_window(meta, seg_idx, TARO_MODEL_DUR)
tmp_dir = _register_tmp_dir(tempfile.mkdtemp())
clip_path = _extract_segment_clip(
_resolve_silent_video(meta), clip_start, clip_dur,
os.path.join(tmp_dir, "xregen_taro_clip.mp4"),
)
# Build a minimal fake-video meta so generate_taro can run on clip_path
sub_segs = _build_segments(clip_dur, TARO_MODEL_DUR, float(crossfade_s))
sub_meta_json = json.dumps({
"segments": sub_segs, "silent_video": clip_path,
"total_dur_s": clip_dur,
})
# Run full TARO generation pipeline on the clip
_ctx_store("taro_gpu_infer", {
"tmp_dir": tmp_dir, "silent_video": clip_path,
"segments": sub_segs, "total_dur_s": clip_dur,
})
results = _taro_gpu_infer(clip_path, seed_val, cfg_scale, num_steps, mode,
crossfade_s, crossfade_db, 1)
wavs, _, _ = results[0]
wavs = [_upsample_taro(w) for w in wavs]
wav = _stitch_wavs(wavs, float(crossfade_s), float(crossfade_db),
clip_dur, TARO_SR_OUT, sub_segs)
return wav, TARO_SR_OUT, clip_start
yield from _xregen_dispatch(state_json, seg_idx, slot_id, _run)
def xregen_mmaudio(seg_idx, state_json, slot_id,
prompt, negative_prompt, seed_val,
cfg_strength, num_steps, crossfade_s, crossfade_db,
request: gr.Request = None):
"""Cross-model regen: run MMAudio on its optimal window, splice into *slot_id*."""
seg_idx = int(seg_idx)
meta = json.loads(state_json)
def _run():
clip_start, clip_end, clip_dur = _xregen_clip_window(meta, seg_idx, MMAUDIO_WINDOW)
source_video = _resolve_silent_video(meta)
sub_segs = _build_segments(clip_dur, MMAUDIO_WINDOW, float(crossfade_s))
print(f"[xregen_mmaudio._run] clip_start={clip_start} clip_dur={clip_dur} source_video={source_video!r} sub_segs={sub_segs}")
# Pass clip_start_s/clip_dur_s so the GPU fn extracts the clip internally —
# pre-extracted tmp files are invisible to the ZeroGPU worker process.
results = _mmaudio_gpu_infer(source_video, prompt, negative_prompt, seed_val,
cfg_strength, num_steps, crossfade_s, crossfade_db, 1,
source_video, json.dumps(sub_segs),
clip_start, clip_dur)
print(f"[xregen_mmaudio._run] gpu_infer returned {len(results)} results")
seg_wavs, sr = results[0]
wav = _stitch_wavs(seg_wavs, float(crossfade_s), float(crossfade_db),
clip_dur, sr, sub_segs)
if sr != TARGET_SR:
wav = _resample_to_target(wav, sr)
sr = TARGET_SR
return wav, sr, clip_start
yield from _xregen_dispatch(state_json, seg_idx, slot_id, _run)
def xregen_hunyuan(seg_idx, state_json, slot_id,
prompt, negative_prompt, seed_val,
guidance_scale, num_steps, model_size,
crossfade_s, crossfade_db,
request: gr.Request = None):
"""Cross-model regen: run HunyuanFoley on its optimal window, splice into *slot_id*."""
seg_idx = int(seg_idx)
meta = json.loads(state_json)
def _run():
clip_start, clip_end, clip_dur = _xregen_clip_window(meta, seg_idx, HUNYUAN_MAX_DUR)
source_video = _resolve_silent_video(meta)
sub_segs = _build_segments(clip_dur, HUNYUAN_MAX_DUR, float(crossfade_s))
results = _hunyuan_gpu_infer(source_video, prompt, negative_prompt, seed_val,
guidance_scale, num_steps, model_size,
crossfade_s, crossfade_db, 1,
source_video, json.dumps(sub_segs), clip_dur,
clip_start, clip_dur)
seg_wavs, sr, _ = results[0]
wav = _stitch_wavs(seg_wavs, float(crossfade_s), float(crossfade_db),
clip_dur, sr, sub_segs)
return wav, sr, clip_start
yield from _xregen_dispatch(state_json, seg_idx, slot_id, _run)
# ================================================================== #
# SHARED UI HELPERS #
# ================================================================== #
def _register_regen_handlers(tab_prefix, model_key, regen_seg_tb, regen_state_tb,
input_components, slot_vids, slot_waves):
"""Register per-slot regen button handlers for a model tab.
This replaces the three nearly-identical for-loops that previously existed
for TARO, MMAudio, and HunyuanFoley tabs.
Args:
tab_prefix: e.g. "taro", "mma", "hf"
model_key: e.g. "taro", "mmaudio", "hunyuan"
regen_seg_tb: gr.Textbox for seg_idx (render=False)
regen_state_tb: gr.Textbox for state_json (render=False)
input_components: list of Gradio input components (video, seed, etc.)
— order must match regen_fn signature after (seg_idx, state_json, video)
slot_vids: list of gr.Video components per slot
slot_waves: list of gr.HTML components per slot
Returns:
list of hidden gr.Buttons (one per slot)
"""
cfg = MODEL_CONFIGS[model_key]
regen_fn = cfg["regen_fn"]
label = cfg["label"]
btns = []
for _i in range(MAX_SLOTS):
_slot_id = f"{tab_prefix}_{_i}"
_btn = gr.Button(render=False, elem_id=f"regen_btn_{_slot_id}")
btns.append(_btn)
print(f"[startup] registering regen handler for slot {_slot_id}")
def _make_regen(_si, _sid, _model_key, _label, _regen_fn):
def _do(seg_idx, state_json, *args):
print(f"[regen {_label}] slot={_sid} seg_idx={seg_idx} "
f"state_json_len={len(state_json) if state_json else 0}")
if not state_json:
print(f"[regen {_label}] early-exit: state_json empty")
yield gr.update(), gr.update()
return
lock = _get_slot_lock(_sid)
with lock:
state = json.loads(state_json)
pending_html = _build_regen_pending_html(
state["segments"], int(seg_idx), _sid, ""
)
yield gr.update(), gr.update(value=pending_html)
print(f"[regen {_label}] slot={_sid} seg_idx={seg_idx} — calling regen")
try:
# args[0] = video, args[1:] = model-specific params
vid, aud, new_meta_json, html = _regen_fn(
args[0], int(seg_idx), state_json, *args[1:], _sid,
)
print(f"[regen {_label}] slot={_sid} seg_idx={seg_idx} — done, vid={vid!r}")
except Exception as _e:
print(f"[regen {_label}] slot={_sid} seg_idx={seg_idx} — ERROR: {_e}")
raise
yield gr.update(value=vid), gr.update(value=html)
return _do
_btn.click(
fn=_make_regen(_i, _slot_id, model_key, label, regen_fn),
inputs=[regen_seg_tb, regen_state_tb] + input_components,
outputs=[slot_vids[_i], slot_waves[_i]],
api_name=f"regen_{tab_prefix}_{_i}",
)
return btns
def _pad_outputs(outputs: list) -> list:
"""Flatten (video, audio, seg_meta) triples and pad to MAX_SLOTS * 3 with None.
Each entry in *outputs* must be a (video_path, audio_path, seg_meta) tuple where
seg_meta = {"segments": [...], "audio_path": str, "video_path": str,
"sr": int, "model": str, "crossfade_s": float,
"crossfade_db": float, "wav_paths": list[str]}
"""
result = []
for i in range(MAX_SLOTS):
if i < len(outputs):
result.extend(outputs[i]) # 3 items: video, audio, meta
else:
result.extend([None, None, None])
return result
# ------------------------------------------------------------------ #
# WaveSurfer waveform + segment marker HTML builder #
# ------------------------------------------------------------------ #
def _build_regen_pending_html(segments: list, regen_seg_idx: int, slot_id: str,
hidden_input_id: str) -> str:
"""Return a waveform placeholder shown while a segment is being regenerated.
Renders a dark bar with the active segment highlighted in amber + a spinner.
"""
segs_json = json.dumps(segments)
seg_colors = [c.format(a="0.25") for c in SEG_COLORS]
active_color = "rgba(255,180,0,0.55)"
duration = segments[-1][1] if segments else 1.0
# Compute contact edges: midpoint of overlap between consecutive segments.
# Matches the contact-edge formula used in _build_waveform_html's canvas JS
# so segment boundaries stay visually identical during pending regen.
contact_edges = []
for i in range(len(segments) - 1):
contact_edges.append((segments[i][1] + segments[i + 1][0]) / 2)
seg_divs = ""
for i, seg in enumerate(segments):
left_t = 0.0 if i == 0 else contact_edges[i - 1]
right_t = duration if i == len(segments) - 1 else contact_edges[i]
left_pct = left_t / duration * 100
width_pct = (right_t - left_t) / duration * 100
color = active_color if i == regen_seg_idx else seg_colors[i % len(seg_colors)]
extra = "border:2px solid #ffb300;animation:wf_pulse 0.8s ease-in-out infinite alternate;" if i == regen_seg_idx else ""
seg_divs += (
f'<div style="position:absolute;top:0;left:{left_pct:.2f}%;'
f'width:{width_pct:.2f}%;height:100%;background:{color};{extra}">'
f'<span style="color:rgba(255,255,255,0.7);font-size:10px;padding:2px 3px;">Seg {i+1}</span>'
f'</div>'
)
spinner = (
'<div style="position:absolute;top:50%;left:50%;transform:translate(-50%,-50%);'
'display:flex;align-items:center;gap:6px;">'
'<div style="width:14px;height:14px;border:2px solid #ffb300;'
'border-top-color:transparent;border-radius:50%;'
'animation:wf_spin 0.7s linear infinite;"></div>'
f'<span style="color:#ffb300;font-size:12px;white-space:nowrap;">'
f'Regenerating Seg {regen_seg_idx+1}…</span>'
'</div>'
)
return f"""
<style>
@keyframes wf_pulse {{from{{opacity:0.5}}to{{opacity:1}}}}
@keyframes wf_spin {{to{{transform:rotate(360deg)}}}}
</style>
<div style="background:#1a1a1a;border-radius:8px;padding:10px;margin-top:6px;">
<div style="position:relative;width:100%;height:80px;background:#1e1e2e;border-radius:4px;overflow:hidden;">
{seg_divs}
{spinner}
</div>
<div style="color:#888;font-size:11px;margin-top:6px;">Regenerating — please wait…</div>
</div>
"""
def _build_waveform_html(audio_path: str, segments: list, slot_id: str,
hidden_input_id: str, state_json: str = "",
fn_index: int = -1, video_path: str = "",
crossfade_s: float = 0.0) -> str:
"""Return a self-contained HTML block with a Canvas waveform (display only),
segment boundary markers, and a download link.
Uses Web Audio API + Canvas — no external libraries.
The waveform is SILENT. The playhead tracks the Gradio <video> element
in the same slot via its timeupdate event.
"""
print(f"[_build_waveform_html] audio_path={audio_path!r} exists={os.path.exists(audio_path) if audio_path else False} slot={slot_id}")
if not audio_path or not os.path.exists(audio_path):
print(f"[_build_waveform_html] returning placeholder — audio missing")
return "<p style='color:#888;font-size:12px'>No audio yet.</p>"
# Serve audio via Gradio's file API instead of base64-encoding the entire
# WAV inline. For a 25s stereo 44.1kHz track this saves ~5 MB per slot.
audio_url = f"/gradio_api/file={audio_path}"
segs_json = json.dumps(segments)
seg_colors = [c.format(a="0.35") for c in SEG_COLORS]
# NOTE: Gradio updates gr.HTML via innerHTML which does NOT execute <script> tags.
# Solution: put the entire waveform (canvas + JS) inside an <iframe srcdoc="...">.
# iframes always execute their scripts. The iframe posts messages to the parent for
# segment-click events; the parent listens and fires the Gradio regen trigger.
# For playhead sync, the iframe polls window.parent for a <video> element.
iframe_inner = f"""<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<style>
* {{ margin:0; padding:0; box-sizing:border-box; }}
body {{ background:#1a1a1a; overflow:hidden; }}
#wrap {{ position:relative; width:100%; height:80px; }}
canvas {{ display:block; }}
#cv {{ position:absolute; top:0; left:0; width:100%; height:100%; }}
#cvp {{ position:absolute; top:0; left:0; width:100%; height:100%; pointer-events:none; }}
</style>
</head>
<body>
<div id="wrap">
<canvas id="cv"></canvas>
<canvas id="cvp"></canvas>
</div>
<script>
(function() {{
const SLOT_ID = '{slot_id}';
const segments = {segs_json};
const segColors = {json.dumps(seg_colors)};
const crossfadeSec = {crossfade_s};
let audioDuration = 0;
// ── Popup via postMessage to parent global listener ─────────────────
// The parent page (Gradio) has a global window.addEventListener('message',...)
// set up via gr.Blocks(js=...) that handles popup show/hide and regen trigger.
function showPopup(idx, mx, my) {{
console.log('[wf showPopup] slot='+SLOT_ID+' idx='+idx+' posting to parent');
// Convert iframe-local coords to parent page coords
try {{
const fr = window.frameElement ? window.frameElement.getBoundingClientRect() : {{left:0,top:0}};
window.parent.postMessage({{
type:'wf_popup', action:'show',
slot_id: SLOT_ID, seg_idx: idx,
t0: segments[idx][0], t1: segments[idx][1],
x: mx + fr.left, y: my + fr.top
}}, '*');
console.log('[wf showPopup] postMessage sent OK');
}} catch(e) {{
console.log('[wf showPopup] postMessage fallback, err='+e.message);
window.parent.postMessage({{
type:'wf_popup', action:'show',
slot_id: SLOT_ID, seg_idx: idx,
t0: segments[idx][0], t1: segments[idx][1],
x: mx, y: my
}}, '*');
}}
}}
function hidePopup() {{
window.parent.postMessage({{type:'wf_popup', action:'hide'}}, '*');
}}
// ── Canvas waveform ────────────────────────────────────────────────
const cv = document.getElementById('cv');
const cvp = document.getElementById('cvp');
const wrap= document.getElementById('wrap');
function drawWaveform(channelData, duration) {{
audioDuration = duration;
const dpr = window.devicePixelRatio || 1;
const W = wrap.getBoundingClientRect().width || window.innerWidth || 600;
const H = 80;
cv.width = W * dpr; cv.height = H * dpr;
const ctx = cv.getContext('2d');
ctx.scale(dpr, dpr);
ctx.fillStyle = '#1e1e2e';
ctx.fillRect(0, 0, W, H);
// Compute contact edges: midpoint of overlap between consecutive segments.
// Each segment is colored from its left contact edge to its right contact edge.
// First segment starts at 0; last segment ends at duration.
const contactEdges = [];
for (let i = 0; i < segments.length - 1; i++) {{
contactEdges.push((segments[i][1] + segments[i+1][0]) / 2);
}}
segments.forEach(function(seg, idx) {{
const x1 = idx === 0 ? 0 : (contactEdges[idx-1] / duration) * W;
const xEnd = idx === segments.length - 1
? W
: (contactEdges[idx] / duration) * W;
ctx.fillStyle = segColors[idx % segColors.length];
ctx.fillRect(x1, 0, xEnd - x1, H);
ctx.fillStyle = 'rgba(255,255,255,0.6)';
ctx.font = '10px sans-serif';
ctx.fillText('Seg '+(idx+1), x1+3, 12);
}});
const samples = channelData.length;
const barW=2, gap=1, step=barW+gap;
const numBars = Math.floor(W / step);
const blockSz = Math.floor(samples / numBars);
ctx.fillStyle = '#4a9eff';
for (let i=0; i<numBars; i++) {{
let max=0;
const s=i*blockSz;
for (let j=0; j<blockSz; j++) {{
const v=Math.abs(channelData[s+j]||0);
if (v>max) max=v;
}}
const barH=Math.max(1, max*H);
ctx.fillRect(i*step, (H-barH)/2, barW, barH);
}}
segments.forEach(function(seg) {{
[seg[0],seg[1]].forEach(function(t) {{
const x=(t/duration)*W;
ctx.strokeStyle='rgba(255,255,255,0.4)';
ctx.lineWidth=1;
ctx.beginPath(); ctx.moveTo(x,0); ctx.lineTo(x,H); ctx.stroke();
}});
}});
// ── Crossfade overlap indicators ──
// The crossfade is applied ±crossfadeSec/2 around each contact edge.
// contactEdges[i] = midpoint of overlap between seg i and seg i+1.
if (crossfadeSec > 0 && segments.length > 1) {{
for (let i = 0; i < segments.length - 1; i++) {{
// Hatch spans crossfadeSec centred on the contact edge
const edgeT = contactEdges[i];
const overlapStart = edgeT - crossfadeSec / 2;
const overlapEnd = edgeT + crossfadeSec / 2;
const xL = (overlapStart / duration) * W;
const xR = (overlapEnd / duration) * W;
// Diagonal hatch pattern over the overlap zone
ctx.save();
ctx.beginPath();
ctx.rect(xL, 0, xR - xL, H);
ctx.clip();
ctx.strokeStyle = 'rgba(255,255,255,0.35)';
ctx.lineWidth = 1;
const spacing = 6;
for (let lx = xL - H; lx < xR + H; lx += spacing) {{
ctx.beginPath();
ctx.moveTo(lx, H);
ctx.lineTo(lx + H, 0);
ctx.stroke();
}}
ctx.restore();
}}
}}
cv.onclick = function(e) {{
const r=cv.getBoundingClientRect();
const xRel=(e.clientX-r.left)/r.width;
const tClick=xRel*duration;
// Pick the segment by contact-edge boundaries — matches the visual coloring.
// Seg 0 owns [0, contactEdges[0]), seg N-1 owns [contactEdges[N-2], duration).
let hit=-1;
segments.forEach(function(seg,idx){{
const visStart = idx === 0 ? 0 : contactEdges[idx-1];
const visEnd = idx === segments.length - 1 ? duration : contactEdges[idx];
if (tClick >= visStart && tClick < visEnd) hit = idx;
}});
console.log('[wf click] tClick='+tClick.toFixed(2)+' hit='+hit+' audioDuration='+audioDuration+' segments='+JSON.stringify(segments));
if (hit>=0) showPopup(hit, e.clientX, e.clientY);
else hidePopup();
}};
}}
function drawPlayhead(progress) {{
const dpr = window.devicePixelRatio || 1;
const W = wrap.getBoundingClientRect().width || window.innerWidth || 600;
const H = 80;
if (cvp.width !== W*dpr) {{ cvp.width=W*dpr; cvp.height=H*dpr; }}
const ctx = cvp.getContext('2d');
ctx.clearRect(0,0,W*dpr,H*dpr);
ctx.save();
ctx.scale(dpr,dpr);
const x=progress*W;
ctx.strokeStyle='#fff';
ctx.lineWidth=2;
ctx.beginPath(); ctx.moveTo(x,0); ctx.lineTo(x,H); ctx.stroke();
ctx.restore();
}}
// Poll parent for video time — find the video in the same wf_container slot
function findSlotVideo() {{
try {{
const par = window.parent.document;
// Walk up from our iframe to find wf_container_{slot_id}, then find its sibling video
const container = par.getElementById('wf_container_{slot_id}');
if (!container) return par.querySelector('video');
// The video is inside the same gr.Group — walk up to find it
let node = container.parentElement;
while (node && node !== par.body) {{
const v = node.querySelector('video');
if (v) return v;
node = node.parentElement;
}}
return null;
}} catch(e) {{ return null; }}
}}
setInterval(function() {{
const vid = findSlotVideo();
if (vid && vid.duration && isFinite(vid.duration) && audioDuration > 0) {{
drawPlayhead(vid.currentTime / vid.duration);
}}
}}, 50);
// ── Fetch + decode audio from Gradio file API ──────────────────────
const audioUrl = '{audio_url}';
fetch(audioUrl)
.then(function(r) {{ return r.arrayBuffer(); }})
.then(function(arrayBuf) {{
const AudioCtx = window.AudioContext || window.webkitAudioContext;
if (!AudioCtx) return;
const tmpCtx = new AudioCtx({{sampleRate:44100}});
tmpCtx.decodeAudioData(arrayBuf,
function(ab) {{
try {{ tmpCtx.close(); }} catch(e) {{}}
function tryDraw() {{
const W = wrap.getBoundingClientRect().width || window.innerWidth;
if (W > 0) {{ drawWaveform(ab.getChannelData(0), ab.duration); }}
else {{ setTimeout(tryDraw, 100); }}
}}
tryDraw();
}},
function(err) {{}}
);
}})
.catch(function(e) {{}});
}})();
</script>
</body>
</html>"""
# Escape for HTML attribute (srcdoc uses HTML entities)
srcdoc = _html.escape(iframe_inner, quote=True)
state_escaped = _html.escape(state_json or "", quote=True)
return f"""
<div id="wf_container_{slot_id}"
data-fn-index="{fn_index}"
data-state="{state_escaped}"
style="background:#1a1a1a;border-radius:8px;padding:10px;margin-top:6px;position:relative;">
<div style="position:relative;width:100%;height:80px;">
<iframe id="wf_iframe_{slot_id}"
srcdoc="{srcdoc}"
sandbox="allow-scripts allow-same-origin"
style="width:100%;height:80px;border:none;border-radius:4px;display:block;"
scrolling="no"></iframe>
</div>
<div style="display:flex;align-items:center;gap:8px;margin-top:6px;">
<span id="wf_statusbar_{slot_id}" style="color:#888;font-size:11px;">Click a segment to regenerate | Playhead syncs to video</span>
<a href="{audio_url}" download="audio_{slot_id}.wav"
style="margin-left:auto;background:#333;color:#eee;border:1px solid #555;
border-radius:4px;padding:3px 10px;font-size:12px;text-decoration:none;">
↓ Audio</a>{f'''
<a href="/gradio_api/file={video_path}" download="video_{slot_id}.mp4"
style="background:#333;color:#eee;border:1px solid #555;
border-radius:4px;padding:3px 10px;font-size:12px;text-decoration:none;">
↓ Video</a>''' if video_path else ''}
</div>
<div id="wf_seglabel_{slot_id}"
style="color:#aaa;font-size:11px;margin-top:4px;min-height:16px;"></div>
</div>
"""
def _make_output_slots(tab_prefix: str) -> tuple:
"""Build MAX_SLOTS output groups for one tab.
Each slot has: video and waveform HTML.
Regen is triggered via direct Gradio queue API calls from JS (no hidden
trigger textboxes needed — DOM event dispatch is unreliable in Gradio 5
Svelte components). State JSON is embedded in the waveform HTML's
data-state attribute and passed directly in the queue API payload.
Returns (grps, vids, waveforms).
"""
grps, vids, waveforms = [], [], []
for i in range(MAX_SLOTS):
slot_id = f"{tab_prefix}_{i}"
with gr.Group(visible=(i == 0), elem_id=f"slot_grp_{slot_id}") as g:
vids.append(gr.Video(label=f"Generation {i+1} — Video",
elem_id=f"slot_vid_{slot_id}",
show_download_button=False))
waveforms.append(gr.HTML(
value="<p style='color:#888;font-size:12px'>Generate audio to see waveform.</p>",
elem_id=f"slot_wave_{slot_id}",
))
grps.append(g)
return grps, vids, waveforms
def _unpack_outputs(flat: list, n: int, tab_prefix: str) -> list:
"""Turn a flat _pad_outputs list into Gradio update lists.
flat has MAX_SLOTS * 3 items: [vid0, aud0, meta0, vid1, aud1, meta1, ...]
Returns updates for vids + waveforms only (NOT grps).
Group visibility is handled separately via .then() to avoid Gradio 5 SSR
'Too many arguments' caused by mixing gr.Group updates with other outputs.
State JSON is embedded in the waveform HTML data-state attribute so JS
can read it when calling the Gradio queue API for regen.
"""
n = int(n)
print(f"[_unpack_outputs] tab={tab_prefix} n={n} flat_len={len(flat)}")
vid_updates = []
wave_updates = []
for i in range(MAX_SLOTS):
vid_path = flat[i * 3]
aud_path = flat[i * 3 + 1]
meta = flat[i * 3 + 2]
print(f"[_unpack_outputs] slot {i}: vid={vid_path is not None} aud={aud_path!r} meta={meta is not None}")
vid_updates.append(gr.update(value=vid_path))
if aud_path and meta:
slot_id = f"{tab_prefix}_{i}"
state_json = json.dumps(meta)
html = _build_waveform_html(aud_path, meta["segments"], slot_id,
"", state_json=state_json,
video_path=meta.get("video_path", ""),
crossfade_s=float(meta.get("crossfade_s", 0)))
wave_updates.append(gr.update(value=html))
else:
wave_updates.append(gr.update(
value="<p style='color:#888;font-size:12px'>Generate audio to see waveform.</p>"
))
return vid_updates + wave_updates
def _on_video_upload_taro(video_file, num_steps, crossfade_s):
if video_file is None:
return gr.update(maximum=MAX_SLOTS, value=1)
try:
D = get_video_duration(video_file)
max_s = _taro_calc_max_samples(D, int(num_steps), float(crossfade_s))
except Exception:
max_s = MAX_SLOTS
return gr.update(maximum=max_s, value=min(1, max_s))
def _btn_generating():
"""Disable a Generate button and show a busy label. Chained as the first
.click() step so it fires immediately, before the (possibly slow) generate
call starts."""
return gr.update(value="Generating…", interactive=False)
def _btn_ready():
"""Re-enable a Generate button. Chained via .then() (not .success()) so it
always runs — even if an earlier step in the chain raised — otherwise a
crashed generation would leave the button stuck disabled."""
return gr.update(value="Generate", interactive=True)
# ================================================================== #
# GRADIO UI #
# ================================================================== #
_SLOT_CSS = """
/* Responsive video: fills column width, height auto from aspect ratio */
.gradio-video video {
width: 100%;
height: auto;
max-height: 60vh;
object-fit: contain;
}
/* Force two-column layout to stay equal-width */
.gradio-container .gradio-row > .gradio-column {
flex: 1 1 0 !important;
min-width: 0 !important;
max-width: 50% !important;
}
/* Hide the built-in download button on output video slots — downloads are
handled by the waveform panel links which always reflect the latest regen. */
[id^="slot_vid_"] .download-icon,
[id^="slot_vid_"] button[aria-label="Download"],
[id^="slot_vid_"] a[download] {
display: none !important;
}
"""
_GLOBAL_JS = """
() => {
// Global postMessage handler for waveform iframe events.
// Runs once on page load (Gradio js= parameter).
// Handles: popup open/close relay, regen trigger via Gradio queue API.
if (window._wf_global_listener) return; // already registered
window._wf_global_listener = true;
// Toggle output-slot visibility by directly flipping display on the slot
// Group's DOM element. This used to be a server round-trip (a .then() that
// returned gr.update(visible=...) for each Group), but Gradio 5.25's client
// rejects Group-visibility outputs with "Too many arguments provided for
// the endpoint" (500 on /queue/join) even when isolated in their own step —
// doing it client-side sidesteps that bug entirely, no server call needed.
window._updateSlotVisibility = function(prefix, n) {
let i = 0;
while (true) {
const el = document.getElementById('slot_grp_' + prefix + '_' + i);
if (!el) break;
el.style.display = (i < n) ? '' : 'none';
i++;
}
};
// Force-reset a stuck "Generating…" button after a generous timeout.
// The normal reset path is a .then() step chained after the (possibly
// crash-prone) generate call — Gradio only fires .then() steps when the
// preceding step's SSE stream reports SOME terminal status. A dropped
// connection (container restart, HTTP/2 protocol error — both observed in
// this app) never delivers one, so the chain stalls forever and the button
// never re-enables. This watchdog is a client-side safety net that doesn't
// depend on the chain completing at all. Token-gated so a stale watchdog
// from an earlier click can't clobber a newer, still-in-flight generation.
window._armGenerateWatchdog = function(btnId, timeoutMs) {
const root = document.getElementById(btnId);
if (!root) return;
const btn = root.tagName === 'BUTTON' ? root : root.querySelector('button');
if (!btn) return;
btn._genToken = (btn._genToken || 0) + 1;
const token = btn._genToken;
setTimeout(function() {
if (btn._genToken === token && btn.disabled) {
btn.disabled = false;
btn.textContent = 'Generate';
console.log('[watchdog] ' + btnId + ' stuck disabled past ' + timeoutMs + 'ms — force re-enabled');
}
}, timeoutMs);
};
// Arm the watchdog via a plain native click listener (event delegation on
// document, so it works regardless of when Gradio/Svelte creates the button
// elements) rather than Gradio's own js= event-chaining mechanism. Combining
// js= with a real server fn= on the SAME .click() event was tried and broke
// the click from ever reaching the server at all (confirmed via server logs
// showing zero activity after the click) — this sidesteps Gradio's event
// wiring entirely so it can't interfere with the real click → fn chain.
document.addEventListener('click', function(ev) {
const hit = ev.target.closest && ev.target.closest('#taro_btn, #mma_btn, #hf_btn');
if (hit) window._armGenerateWatchdog(hit.id, 600000);
});
// ── ZeroGPU quota attribution ──
// HF Spaces run inside an iframe on huggingface.co. Gradio's own JS client
// gets ZeroGPU auth headers (x-zerogpu-token, x-zerogpu-uuid) by sending a
// postMessage("zerogpu-headers") to the parent frame. The parent responds
// with a Map of headers that must be included on queue/join calls.
// We replicate this exact mechanism so our raw regen fetch() calls are
// attributed to the logged-in user's Pro quota.
function _fetchZerogpuHeaders() {
return new Promise(function(resolve) {
// Check if we're in an HF iframe with zerogpu support
if (typeof window === 'undefined' || window.parent === window || !window.supports_zerogpu_headers) {
console.log('[zerogpu] not in HF iframe or no zerogpu support');
resolve({});
return;
}
// Determine origin — same logic as Gradio's client
var hostname = window.location.hostname;
var hfhubdev = 'dev.spaces.huggingface.tech';
var origin = hostname.includes('.dev.')
? 'https://moon-' + hostname.split('.')[1] + '.' + hfhubdev
: 'https://huggingface.co';
// Use MessageChannel just like Gradio's post_message helper
var channel = new MessageChannel();
var done = false;
channel.port1.onmessage = function(ev) {
channel.port1.close();
done = true;
var headers = ev.data;
if (headers && typeof headers === 'object') {
// Convert Map to plain object if needed
var obj = {};
if (typeof headers.forEach === 'function') {
headers.forEach(function(v, k) { obj[k] = v; });
} else {
obj = headers;
}
console.log('[zerogpu] got headers from parent:', Object.keys(obj).join(', '));
resolve(obj);
} else {
resolve({});
}
};
window.parent.postMessage('zerogpu-headers', origin, [channel.port2]);
// Timeout: don't block regen if parent doesn't respond
setTimeout(function() { if (!done) { done = true; channel.port1.close(); resolve({}); } }, 3000);
});
}
// Cache: api_name -> fn_index, built once from gradio_config.dependencies
let _fnIndexCache = null;
function getFnIndex(apiName) {
if (!_fnIndexCache) {
_fnIndexCache = {};
const deps = window.gradio_config && window.gradio_config.dependencies;
if (deps) deps.forEach(function(d, i) {
if (d.api_name) _fnIndexCache[d.api_name] = i;
});
}
return _fnIndexCache[apiName];
}
// Read a component's current DOM value by elem_id.
// For Number/Slider: reads the <input type="number"> or <input type="range">.
// For Textbox/Radio: reads the <textarea> or checked <input type="radio">.
// Returns null if not found.
function readComponentValue(elemId) {
const el = document.getElementById(elemId);
if (!el) return null;
const numInput = el.querySelector('input[type="number"]');
if (numInput) return parseFloat(numInput.value);
const rangeInput = el.querySelector('input[type="range"]');
if (rangeInput) return parseFloat(rangeInput.value);
const radio = el.querySelector('input[type="radio"]:checked');
if (radio) return radio.value;
const ta = el.querySelector('textarea');
if (ta) return ta.value;
const txt = el.querySelector('input[type="text"], input:not([type])');
if (txt) return txt.value;
return null;
}
// Fire regen for a given slot and segment by posting directly to the
// Gradio queue API — bypasses Svelte binding entirely.
// targetModel: 'taro' | 'mma' | 'hf' (which model to use for inference)
// If targetModel matches the slot's own prefix, uses the per-slot regen_* endpoint.
// Otherwise uses the shared xregen_* cross-model endpoint.
function fireRegen(slot_id, seg_idx, targetModel) {
// Block if a regen is already in-flight for this slot
if (_regenInFlight[slot_id]) {
console.log('[fireRegen] blocked — regen already in-flight for', slot_id);
return;
}
_regenInFlight[slot_id] = true;
const prefix = slot_id.split('_')[0]; // owning tab: 'taro'|'mma'|'hf'
const slotNum = parseInt(slot_id.split('_')[1], 10);
// Decide which endpoint to call
const crossModel = (targetModel !== prefix);
let apiName, data;
// Read state_json from the waveform container data-state attribute
const container = document.getElementById('wf_container_' + slot_id);
const stateJson = container ? (container.getAttribute('data-state') || '') : '';
if (!stateJson) {
console.warn('[fireRegen] no state_json for slot', slot_id);
return;
}
if (!crossModel) {
// ── Same-model regen: per-slot endpoint, video passed as null ──
apiName = 'regen_' + prefix + '_' + slotNum;
if (prefix === 'taro') {
data = [seg_idx, stateJson, null,
readComponentValue('taro_seed'), readComponentValue('taro_cfg'),
readComponentValue('taro_steps'), readComponentValue('taro_mode'),
readComponentValue('taro_cf_dur'), readComponentValue('taro_cf_db')];
} else if (prefix === 'mma') {
data = [seg_idx, stateJson, null,
readComponentValue('mma_prompt'), readComponentValue('mma_neg'),
readComponentValue('mma_seed'), readComponentValue('mma_cfg'),
readComponentValue('mma_steps'),
readComponentValue('mma_cf_dur'), readComponentValue('mma_cf_db')];
} else {
data = [seg_idx, stateJson, null,
readComponentValue('hf_prompt'), readComponentValue('hf_neg'),
readComponentValue('hf_seed'), readComponentValue('hf_guidance'),
readComponentValue('hf_steps'), readComponentValue('hf_size'),
readComponentValue('hf_cf_dur'), readComponentValue('hf_cf_db')];
}
} else {
// ── Cross-model regen: shared xregen_* endpoint ──
// slot_id is passed so the server knows which slot's state to splice into.
// UI params are read from the target model's tab inputs.
if (targetModel === 'taro') {
apiName = 'xregen_taro';
data = [seg_idx, stateJson, slot_id,
readComponentValue('taro_seed'), readComponentValue('taro_cfg'),
readComponentValue('taro_steps'), readComponentValue('taro_mode'),
readComponentValue('taro_cf_dur'), readComponentValue('taro_cf_db')];
} else if (targetModel === 'mma') {
apiName = 'xregen_mmaudio';
data = [seg_idx, stateJson, slot_id,
readComponentValue('mma_prompt'), readComponentValue('mma_neg'),
readComponentValue('mma_seed'), readComponentValue('mma_cfg'),
readComponentValue('mma_steps'),
readComponentValue('mma_cf_dur'), readComponentValue('mma_cf_db')];
} else {
apiName = 'xregen_hunyuan';
data = [seg_idx, stateJson, slot_id,
readComponentValue('hf_prompt'), readComponentValue('hf_neg'),
readComponentValue('hf_seed'), readComponentValue('hf_guidance'),
readComponentValue('hf_steps'), readComponentValue('hf_size'),
readComponentValue('hf_cf_dur'), readComponentValue('hf_cf_db')];
}
}
console.log('[fireRegen] calling api', apiName, 'seg', seg_idx);
// Snapshot current waveform HTML + video src before mutating anything,
// so we can restore on error (e.g. quota exceeded).
var _preRegenWaveHtml = null;
var _preRegenVideoSrc = null;
var waveElSnap = document.getElementById('slot_wave_' + slot_id);
if (waveElSnap) _preRegenWaveHtml = waveElSnap.innerHTML;
var vidElSnap = document.getElementById('slot_vid_' + slot_id);
if (vidElSnap) { var vSnap = vidElSnap.querySelector('video'); if (vSnap) _preRegenVideoSrc = vSnap.getAttribute('src'); }
// Show spinner immediately
const lbl = document.getElementById('wf_seglabel_' + slot_id);
if (lbl) lbl.textContent = 'Regenerating Seg ' + (seg_idx + 1) + '...';
const fnIndex = getFnIndex(apiName);
if (fnIndex === undefined) {
console.warn('[fireRegen] fn_index not found for api_name:', apiName);
return;
}
// Get ZeroGPU auth headers from the HF parent frame (same mechanism
// Gradio's own JS client uses), then fire the regen queue/join call.
// Falls back to user-supplied HF token if zerogpu headers aren't available.
_fetchZerogpuHeaders().then(function(zerogpuHeaders) {
var regenHeaders = {'Content-Type': 'application/json'};
var hasZerogpu = zerogpuHeaders && Object.keys(zerogpuHeaders).length > 0;
if (hasZerogpu) {
// Merge zerogpu headers (x-zerogpu-token, x-zerogpu-uuid)
for (var k in zerogpuHeaders) { regenHeaders[k] = zerogpuHeaders[k]; }
console.log('[fireRegen] using zerogpu headers from parent frame');
} else {
console.warn('[fireRegen] no zerogpu headers available — may use anonymous quota');
}
fetch('/gradio_api/queue/join', {
method: 'POST',
credentials: 'include',
headers: regenHeaders,
body: JSON.stringify({
data: data,
fn_index: fnIndex,
api_name: '/' + apiName,
session_hash: window.__gradio_session_hash__,
event_data: null,
trigger_id: null
})
}).then(function(r) { return r.json(); }).then(function(j) {
if (!j.event_id) { console.error('[fireRegen] no event_id:', j); return; }
console.log('[fireRegen] queued, event_id:', j.event_id);
_listenAndApply(j.event_id, slot_id, seg_idx, _preRegenWaveHtml, _preRegenVideoSrc);
}).catch(function(e) {
console.error('[fireRegen] fetch error:', e);
if (lbl) lbl.textContent = 'Error — see console';
var sb = document.getElementById('wf_statusbar_' + slot_id);
if (sb) { sb.style.color = '#e05252'; sb.textContent = '\u26a0 Request failed: ' + e.message; }
});
});
}
// Subscribe to Gradio SSE stream for an event and apply outputs to DOM.
// For regen handlers, output[0] = video update, output[1] = waveform HTML update.
function _applyVideoSrc(slot_id, newSrc) {
var vidEl = document.getElementById('slot_vid_' + slot_id);
if (!vidEl) return false;
var video = vidEl.querySelector('video');
if (!video) return false;
if (video.getAttribute('src') === newSrc) return true; // already correct
video.setAttribute('src', newSrc);
video.src = newSrc;
video.load();
console.log('[_applyVideoSrc] applied src to', 'slot_vid_' + slot_id, 'src:', newSrc.slice(-40));
return true;
}
// Toast notification — styled like ZeroGPU quota warnings.
function _showRegenToast(message, isError) {
var t = document.createElement('div');
t.style.cssText = 'position:fixed;bottom:24px;left:50%;transform:translateX(-50%);' +
'z-index:2147483647;padding:12px 20px;border-radius:8px;font-family:sans-serif;' +
'font-size:13px;max-width:520px;text-align:center;box-shadow:0 4px 20px rgba(0,0,0,.6);' +
'background:' + (isError ? '#7a1c1c' : '#1c4a1c') + ';color:#fff;' +
'border:1px solid ' + (isError ? '#c0392b' : '#27ae60') + ';' +
'pointer-events:none;';
t.textContent = message;
document.body.appendChild(t);
setTimeout(function() {
t.style.transition = 'opacity 0.5s';
t.style.opacity = '0';
setTimeout(function() { t.parentNode && t.parentNode.removeChild(t); }, 600);
}, isError ? 8000 : 3000);
}
function _listenAndApply(eventId, slot_id, seg_idx, preRegenWaveHtml, preRegenVideoSrc) {
var _pendingVideoSrc = null;
const es = new EventSource('/gradio_api/queue/data?session_hash=' + window.__gradio_session_hash__);
es.onmessage = function(e) {
var msg;
try { msg = JSON.parse(e.data); } catch(_) { return; }
if (msg.event_id !== eventId) return;
if (msg.msg === 'process_generating' || msg.msg === 'process_completed') {
var out = msg.output;
if (out && out.data) {
var vidUpdate = out.data[0];
var waveUpdate = out.data[1];
var newSrc = null;
if (vidUpdate) {
if (vidUpdate.value && vidUpdate.value.video && vidUpdate.value.video.url) newSrc = vidUpdate.value.video.url;
else if (vidUpdate.video && vidUpdate.video.url) newSrc = vidUpdate.video.url;
else if (vidUpdate.value && vidUpdate.value.url) newSrc = vidUpdate.value.url;
else if (typeof vidUpdate.value === 'string') newSrc = vidUpdate.value;
else if (vidUpdate.url) newSrc = vidUpdate.url;
}
if (newSrc) _pendingVideoSrc = newSrc;
var waveHtml = null;
if (waveUpdate) {
if (typeof waveUpdate === 'string') waveHtml = waveUpdate;
else if (waveUpdate.value && typeof waveUpdate.value === 'string') waveHtml = waveUpdate.value;
}
if (waveHtml) {
var waveEl = document.getElementById('slot_wave_' + slot_id);
if (waveEl) {
var inner = waveEl.querySelector('.prose') || waveEl.querySelector('div');
if (inner) inner.innerHTML = waveHtml;
else waveEl.innerHTML = waveHtml;
}
}
}
if (msg.msg === 'process_completed') {
es.close();
_regenInFlight[slot_id] = false;
var errMsg = msg.output && msg.output.error;
var hadError = !!errMsg;
console.log('[fireRegen] completed for', slot_id, 'error:', hadError, errMsg || '');
var lbl = document.getElementById('wf_seglabel_' + slot_id);
if (hadError) {
var toastMsg = typeof errMsg === 'string' ? errMsg : JSON.stringify(errMsg);
// Restore previous waveform HTML and video src
if (preRegenWaveHtml !== null) {
var waveEl2 = document.getElementById('slot_wave_' + slot_id);
if (waveEl2) waveEl2.innerHTML = preRegenWaveHtml;
}
if (preRegenVideoSrc !== null) {
var vidElR = document.getElementById('slot_vid_' + slot_id);
if (vidElR) { var vR = vidElR.querySelector('video'); if (vR) { vR.setAttribute('src', preRegenVideoSrc); vR.src = preRegenVideoSrc; vR.load(); } }
}
// Update the statusbar (query after restore so we get the freshly-restored element)
var isAbort = toastMsg.toLowerCase().indexOf('aborted') !== -1;
var isTimeout = toastMsg.toLowerCase().indexOf('timeout') !== -1;
var failMsg = isAbort || isTimeout
? '\u26a0 GPU cold-start — segment unchanged, try again'
: '\u26a0 Regen failed — segment unchanged';
var statusBar = document.getElementById('wf_statusbar_' + slot_id);
if (statusBar) {
statusBar.style.color = '#e05252';
statusBar.textContent = failMsg;
setTimeout(function() { statusBar.style.color = '#888'; statusBar.textContent = 'Click a segment to regenerate \u00a0|\u00a0 Playhead syncs to video'; }, 8000);
}
} else {
if (lbl) lbl.textContent = 'Done';
var src = _pendingVideoSrc;
if (src) {
_applyVideoSrc(slot_id, src);
setTimeout(function() { _applyVideoSrc(slot_id, src); }, 50);
setTimeout(function() { _applyVideoSrc(slot_id, src); }, 300);
setTimeout(function() { _applyVideoSrc(slot_id, src); }, 800);
var vidEl = document.getElementById('slot_vid_' + slot_id);
if (vidEl) {
var obs = new MutationObserver(function() { _applyVideoSrc(slot_id, src); });
obs.observe(vidEl, {subtree: true, attributes: true, attributeFilter: ['src'], childList: true});
setTimeout(function() { obs.disconnect(); }, 2000);
}
}
}
}
}
if (msg.msg === 'close_stream') { es.close(); }
};
es.onerror = function() { es.close(); _regenInFlight[slot_id] = false; };
}
// Track in-flight regen per slot — prevents queuing multiple jobs from rapid clicks
var _regenInFlight = {};
// Shared popup element created once and reused across all slots
let _popup = null;
let _pendingSlot = null, _pendingIdx = null;
function ensurePopup() {
if (_popup) return _popup;
_popup = document.createElement('div');
_popup.style.cssText = 'display:none;position:fixed;z-index:99999;' +
'background:#2a2a2a;border:1px solid #555;border-radius:6px;' +
'padding:8px 12px;box-shadow:0 4px 16px rgba(0,0,0,.5);font-family:sans-serif;';
var btnStyle = 'color:#fff;border:none;border-radius:4px;padding:5px 10px;' +
'font-size:11px;cursor:pointer;flex:1;';
_popup.innerHTML =
'<div id="_wf_popup_lbl" style="color:#ccc;font-size:11px;margin-bottom:6px;white-space:nowrap;"></div>' +
'<div style="display:flex;gap:5px;">' +
'<button id="_wf_popup_taro" style="background:#1d6fa5;' + btnStyle + '">⟳ TARO</button>' +
'<button id="_wf_popup_mma" style="background:#2d7a4a;' + btnStyle + '">⟳ MMAudio</button>' +
'<button id="_wf_popup_hf" style="background:#7a3d8c;' + btnStyle + '">⟳ Hunyuan</button>' +
'</div>';
document.body.appendChild(_popup);
['taro','mma','hf'].forEach(function(model) {
document.getElementById('_wf_popup_' + model).onclick = function(e) {
e.stopPropagation();
var slot = _pendingSlot, idx = _pendingIdx;
hidePopup();
if (slot !== null && idx !== null) fireRegen(slot, idx, model);
};
});
// Use bubble phase (false) so stopPropagation() on the button click prevents this from firing
document.addEventListener('click', function() { hidePopup(); }, false);
return _popup;
}
function hidePopup() {
if (_popup) _popup.style.display = 'none';
_pendingSlot = null; _pendingIdx = null;
}
window.addEventListener('message', function(e) {
const d = e.data;
console.log('[global msg] received type=' + (d && d.type) + ' action=' + (d && d.action));
if (!d || d.type !== 'wf_popup') return;
const p = ensurePopup();
if (d.action === 'hide') { hidePopup(); return; }
// action === 'show'
_pendingSlot = d.slot_id;
_pendingIdx = d.seg_idx;
const lbl = document.getElementById('_wf_popup_lbl');
if (lbl) lbl.textContent = 'Seg ' + (d.seg_idx + 1) +
' (' + d.t0.toFixed(2) + 's \u2013 ' + d.t1.toFixed(2) + 's)';
p.style.display = 'block';
p.style.left = (d.x + 10) + 'px';
p.style.top = (d.y + 10) + 'px';
requestAnimationFrame(function() {
const r = p.getBoundingClientRect();
if (r.right > window.innerWidth - 8) p.style.left = (window.innerWidth - r.width - 8) + 'px';
if (r.bottom > window.innerHeight - 8) p.style.top = (window.innerHeight - r.height - 8) + 'px';
});
});
}
"""
with gr.Blocks(title="Generate Audio for Video", css=_SLOT_CSS, js=_GLOBAL_JS) as demo:
gr.Markdown(
"# Generate Audio for Video\n"
"Choose a model and upload a video to generate synchronized audio.\n\n"
"| Model | Best for | Avoid for |\n"
"|-------|----------|-----------|\n"
"| **TARO** | Natural, physics-driven impacts — footsteps, collisions, water, wind, crackling fire. Excels when the sound is tightly coupled to visible motion without needing a text description. | Dialogue, music, or complex layered soundscapes where semantic context matters. |\n"
"| **MMAudio** | Mixed scenes where you want both visual grounding *and* semantic control via a text prompt — e.g. a busy street scene where you want to emphasize the rain rather than the traffic. Great for ambient textures and nuanced sound design. | Pure impact/foley shots where TARO's motion-coupling would be sharper, or cinematic music beds. |\n"
"| **HunyuanFoley** | Cinematic foley requiring high fidelity and explicit creative direction — dramatic SFX, layered environmental design, or any scene where you have a clear written description of the desired sound palette. | Quick one-shot clips where you don't want to write a prompt, or raw impact sounds where timing precision matters more than richness. |"
)
with gr.Tabs():
# ---------------------------------------------------------- #
# Tab 1 — TARO #
# ---------------------------------------------------------- #
with gr.Tab("TARO"):
with gr.Row():
with gr.Column(scale=1):
taro_video = gr.Video(label="Input Video")
taro_seed = gr.Number(label="Seed (-1 = random)", value=-1, precision=0, elem_id="taro_seed")
taro_cfg = gr.Slider(label="CFG Scale", minimum=1, maximum=15, value=8.0, step=0.5, elem_id="taro_cfg")
taro_steps = gr.Slider(label="Sampling Steps", minimum=10, maximum=50, value=25, step=1, elem_id="taro_steps")
taro_mode = gr.Radio(label="Sampling Mode", choices=["sde", "ode"], value="sde", elem_id="taro_mode")
taro_cf_dur = gr.Slider(label="Crossfade Duration (s)", minimum=0, maximum=4, value=2, step=0.1, elem_id="taro_cf_dur")
taro_cf_db = gr.Textbox(label="Crossfade Boost (dB)", value="3", elem_id="taro_cf_db")
taro_samples = gr.Slider(label="Generations", minimum=1, maximum=MAX_SLOTS, value=1, step=1)
taro_btn = gr.Button("Generate", variant="primary", elem_id="taro_btn")
with gr.Column(scale=1):
(_, taro_slot_vids,
taro_slot_waves) = _make_output_slots("taro")
# Hidden regen plumbing — render=False so no DOM element is created,
# avoiding Gradio's "Too many arguments" Svelte validation error.
# JS passes values directly via queue/join data array at the correct
# positional index (these show up as inputs to the fn but have no DOM).
taro_regen_seg = gr.Textbox(value="0", render=False)
taro_regen_state = gr.Textbox(value="", render=False)
for trigger in [taro_video, taro_steps, taro_cf_dur]:
trigger.change(
fn=_on_video_upload_taro,
inputs=[taro_video, taro_steps, taro_cf_dur],
outputs=[taro_samples],
)
taro_samples.change(
fn=None,
inputs=[taro_samples],
outputs=[],
js="(n) => window._updateSlotVisibility('taro', n)",
)
def _run_taro(video, seed, cfg, steps, mode, cf_dur, cf_db, n):
print(f"[_run_taro] called video={video!r} n={n}")
try:
flat = generate_taro(video, seed, cfg, steps, mode, cf_dur, cf_db, n)
print(f"[_run_taro] generate_taro returned flat_len={len(flat) if flat else None}")
result = _unpack_outputs(flat, n, "taro")
print(f"[_run_taro] _unpack_outputs returned {len(result)} updates")
return result
except Exception as e:
import traceback
print(f"[_run_taro] EXCEPTION: {e}")
traceback.print_exc()
raise
# Split group visibility into a separate .then() to avoid Gradio 5 SSR
# "Too many arguments" caused by including gr.Group in mixed output lists.
(taro_btn.click(
fn=_btn_generating,
inputs=None,
outputs=[taro_btn],
).then(
fn=_run_taro,
inputs=[taro_video, taro_seed, taro_cfg, taro_steps, taro_mode,
taro_cf_dur, taro_cf_db, taro_samples],
outputs=taro_slot_vids + taro_slot_waves,
).then(
fn=None,
inputs=[taro_samples],
outputs=[],
js="(n) => window._updateSlotVisibility('taro', n)",
).then(
fn=_btn_ready,
inputs=None,
outputs=[taro_btn],
))
# Per-slot regen handlers — JS calls /gradio_api/queue/join with
# fn_index (by api_name) + data=[seg_idx, state_json, video, ...params].
taro_regen_btns = _register_regen_handlers(
"taro", "taro", taro_regen_seg, taro_regen_state,
[taro_video, taro_seed, taro_cfg, taro_steps,
taro_mode, taro_cf_dur, taro_cf_db],
taro_slot_vids, taro_slot_waves,
)
# ---------------------------------------------------------- #
# Tab 2 — MMAudio #
# ---------------------------------------------------------- #
with gr.Tab("MMAudio"):
with gr.Row():
with gr.Column(scale=1):
mma_video = gr.Video(label="Input Video")
mma_prompt = gr.Textbox(label="Prompt", placeholder="e.g. footsteps on gravel", elem_id="mma_prompt")
mma_neg = gr.Textbox(label="Negative Prompt", value="music", placeholder="music, speech", elem_id="mma_neg")
mma_seed = gr.Number(label="Seed (-1 = random)", value=-1, precision=0, elem_id="mma_seed")
mma_cfg = gr.Slider(label="CFG Strength", minimum=1, maximum=10, value=4.5, step=0.5, elem_id="mma_cfg")
mma_steps = gr.Slider(label="Steps", minimum=10, maximum=50, value=25, step=1, elem_id="mma_steps")
mma_cf_dur = gr.Slider(label="Crossfade Duration (s)", minimum=0, maximum=4, value=2, step=0.1, elem_id="mma_cf_dur")
mma_cf_db = gr.Textbox(label="Crossfade Boost (dB)", value="3", elem_id="mma_cf_db")
mma_samples = gr.Slider(label="Generations", minimum=1, maximum=MAX_SLOTS, value=1, step=1)
mma_btn = gr.Button("Generate", variant="primary", elem_id="mma_btn")
with gr.Column(scale=1):
(_, mma_slot_vids,
mma_slot_waves) = _make_output_slots("mma")
# Hidden regen plumbing — render=False so no DOM element is created,
# avoiding Gradio's "Too many arguments" Svelte validation error.
mma_regen_seg = gr.Textbox(value="0", render=False)
mma_regen_state = gr.Textbox(value="", render=False)
mma_samples.change(
fn=None,
inputs=[mma_samples],
outputs=[],
js="(n) => window._updateSlotVisibility('mma', n)",
)
def _run_mmaudio(video, prompt, neg, seed, cfg, steps, cf_dur, cf_db, n):
flat = generate_mmaudio(video, prompt, neg, seed, cfg, steps, cf_dur, cf_db, n)
return _unpack_outputs(flat, n, "mma")
(mma_btn.click(
fn=_btn_generating,
inputs=None,
outputs=[mma_btn],
).then(
fn=_run_mmaudio,
inputs=[mma_video, mma_prompt, mma_neg, mma_seed,
mma_cfg, mma_steps, mma_cf_dur, mma_cf_db, mma_samples],
outputs=mma_slot_vids + mma_slot_waves,
).then(
fn=None,
inputs=[mma_samples],
outputs=[],
js="(n) => window._updateSlotVisibility('mma', n)",
).then(
fn=_btn_ready,
inputs=None,
outputs=[mma_btn],
))
mma_regen_btns = _register_regen_handlers(
"mma", "mmaudio", mma_regen_seg, mma_regen_state,
[mma_video, mma_prompt, mma_neg, mma_seed,
mma_cfg, mma_steps, mma_cf_dur, mma_cf_db],
mma_slot_vids, mma_slot_waves,
)
# ---------------------------------------------------------- #
# Tab 3 — HunyuanVideoFoley #
# ---------------------------------------------------------- #
with gr.Tab("HunyuanFoley"):
with gr.Row():
with gr.Column(scale=1):
hf_video = gr.Video(label="Input Video")
hf_prompt = gr.Textbox(label="Prompt", placeholder="e.g. rain hitting a metal roof", elem_id="hf_prompt")
hf_neg = gr.Textbox(label="Negative Prompt", value="noisy, harsh", elem_id="hf_neg")
hf_seed = gr.Number(label="Seed (-1 = random)", value=-1, precision=0, elem_id="hf_seed")
hf_guidance = gr.Slider(label="Guidance Scale", minimum=1, maximum=10, value=4.5, step=0.5, elem_id="hf_guidance")
hf_steps = gr.Slider(label="Steps", minimum=10, maximum=100, value=50, step=5, elem_id="hf_steps")
hf_size = gr.Radio(label="Model Size", choices=["xl", "xxl"], value="xxl", elem_id="hf_size")
hf_cf_dur = gr.Slider(label="Crossfade Duration (s)", minimum=0, maximum=4, value=2, step=0.1, elem_id="hf_cf_dur")
hf_cf_db = gr.Textbox(label="Crossfade Boost (dB)", value="3", elem_id="hf_cf_db")
hf_samples = gr.Slider(label="Generations", minimum=1, maximum=MAX_SLOTS, value=1, step=1)
hf_btn = gr.Button("Generate", variant="primary", elem_id="hf_btn")
with gr.Column(scale=1):
(_, hf_slot_vids,
hf_slot_waves) = _make_output_slots("hf")
# Hidden regen plumbing — render=False so no DOM element is created,
# avoiding Gradio's "Too many arguments" Svelte validation error.
hf_regen_seg = gr.Textbox(value="0", render=False)
hf_regen_state = gr.Textbox(value="", render=False)
hf_samples.change(
fn=None,
inputs=[hf_samples],
outputs=[],
js="(n) => window._updateSlotVisibility('hf', n)",
)
def _run_hunyuan(video, prompt, neg, seed, guidance, steps, size, cf_dur, cf_db, n):
flat = generate_hunyuan(video, prompt, neg, seed, guidance, steps, size, cf_dur, cf_db, n)
return _unpack_outputs(flat, n, "hf")
(hf_btn.click(
fn=_btn_generating,
inputs=None,
outputs=[hf_btn],
).then(
fn=_run_hunyuan,
inputs=[hf_video, hf_prompt, hf_neg, hf_seed,
hf_guidance, hf_steps, hf_size, hf_cf_dur, hf_cf_db, hf_samples],
outputs=hf_slot_vids + hf_slot_waves,
).then(
fn=None,
inputs=[hf_samples],
outputs=[],
js="(n) => window._updateSlotVisibility('hf', n)",
).then(
fn=_btn_ready,
inputs=None,
outputs=[hf_btn],
))
hf_regen_btns = _register_regen_handlers(
"hf", "hunyuan", hf_regen_seg, hf_regen_state,
[hf_video, hf_prompt, hf_neg, hf_seed,
hf_guidance, hf_steps, hf_size, hf_cf_dur, hf_cf_db],
hf_slot_vids, hf_slot_waves,
)
# ---- Browser-safe transcode on upload ----
# Gradio serves the original uploaded file to the browser preview widget,
# so H.265 sources show as blank. We re-encode to H.264 on upload and feed
# the result back so the preview plays. mux_video_audio already re-encodes
# to H.264 during generation, so no double-conversion conflict.
taro_video.upload(fn=_transcode_for_browser, inputs=[taro_video], outputs=[taro_video])
mma_video.upload(fn=_transcode_for_browser, inputs=[mma_video], outputs=[mma_video])
hf_video.upload(fn=_transcode_for_browser, inputs=[hf_video], outputs=[hf_video])
# ---- Cross-tab video sync ----
_sync = lambda v: (gr.update(value=v), gr.update(value=v))
taro_video.change(fn=_sync, inputs=[taro_video], outputs=[mma_video, hf_video])
mma_video.change(fn=_sync, inputs=[mma_video], outputs=[taro_video, hf_video])
hf_video.change(fn=_sync, inputs=[hf_video], outputs=[taro_video, mma_video])
# ---- Cross-model regen endpoints ----
# render=False inputs/outputs: no DOM elements created, no SSR validation impact.
# JS calls these via /gradio_api/queue/join using the api_name and applies
# the returned video+waveform directly to the target slot's DOM elements.
_xr_seg = gr.Textbox(value="0", render=False)
_xr_state = gr.Textbox(value="", render=False)
_xr_slot_id = gr.Textbox(value="", render=False)
# Dummy outputs for xregen events: must be real rendered components so Gradio
# can look them up in session state during postprocess_data. The JS listener
# (_listenAndApply) applies the returned video/HTML directly to the correct
# slot's DOM elements and ignores Gradio's own output routing, so these
# slot-0 components simply act as sinks — their displayed value is overwritten
# by the real JS update immediately after.
_xr_dummy_vid = taro_slot_vids[0]
_xr_dummy_wave = taro_slot_waves[0]
# TARO cross-model regen inputs: seg_idx, state_json, slot_id, seed, cfg, steps, mode, cf_dur, cf_db
_xr_taro_seed = gr.Textbox(value="-1", render=False)
_xr_taro_cfg = gr.Textbox(value="7.5", render=False)
_xr_taro_steps = gr.Textbox(value="25", render=False)
_xr_taro_mode = gr.Textbox(value="sde", render=False)
_xr_taro_cfd = gr.Textbox(value="2", render=False)
_xr_taro_cfdb = gr.Textbox(value="3", render=False)
gr.Button(render=False).click(
fn=xregen_taro,
inputs=[_xr_seg, _xr_state, _xr_slot_id,
_xr_taro_seed, _xr_taro_cfg, _xr_taro_steps,
_xr_taro_mode, _xr_taro_cfd, _xr_taro_cfdb],
outputs=[_xr_dummy_vid, _xr_dummy_wave],
api_name="xregen_taro",
)
# MMAudio cross-model regen inputs: seg_idx, state_json, slot_id, prompt, neg, seed, cfg, steps, cf_dur, cf_db
_xr_mma_prompt = gr.Textbox(value="", render=False)
_xr_mma_neg = gr.Textbox(value="", render=False)
_xr_mma_seed = gr.Textbox(value="-1", render=False)
_xr_mma_cfg = gr.Textbox(value="4.5", render=False)
_xr_mma_steps = gr.Textbox(value="25", render=False)
_xr_mma_cfd = gr.Textbox(value="2", render=False)
_xr_mma_cfdb = gr.Textbox(value="3", render=False)
gr.Button(render=False).click(
fn=xregen_mmaudio,
inputs=[_xr_seg, _xr_state, _xr_slot_id,
_xr_mma_prompt, _xr_mma_neg, _xr_mma_seed,
_xr_mma_cfg, _xr_mma_steps, _xr_mma_cfd, _xr_mma_cfdb],
outputs=[_xr_dummy_vid, _xr_dummy_wave],
api_name="xregen_mmaudio",
)
# HunyuanFoley cross-model regen inputs: seg_idx, state_json, slot_id, prompt, neg, seed, guidance, steps, size, cf_dur, cf_db
_xr_hf_prompt = gr.Textbox(value="", render=False)
_xr_hf_neg = gr.Textbox(value="", render=False)
_xr_hf_seed = gr.Textbox(value="-1", render=False)
_xr_hf_guide = gr.Textbox(value="4.5", render=False)
_xr_hf_steps = gr.Textbox(value="50", render=False)
_xr_hf_size = gr.Textbox(value="xxl", render=False)
_xr_hf_cfd = gr.Textbox(value="2", render=False)
_xr_hf_cfdb = gr.Textbox(value="3", render=False)
gr.Button(render=False).click(
fn=xregen_hunyuan,
inputs=[_xr_seg, _xr_state, _xr_slot_id,
_xr_hf_prompt, _xr_hf_neg, _xr_hf_seed,
_xr_hf_guide, _xr_hf_steps, _xr_hf_size,
_xr_hf_cfd, _xr_hf_cfdb],
outputs=[_xr_dummy_vid, _xr_dummy_wave],
api_name="xregen_hunyuan",
)
# NOTE: ZeroGPU quota attribution is handled via postMessage("zerogpu-headers")
# to the HF parent frame — the same mechanism Gradio's own JS client uses.
# This replaced the old x-ip-token relay approach which was unreliable.
print("[startup] app.py fully loaded — regen handlers registered, SSR disabled")
demo.queue(max_size=10).launch(ssr_mode=False, height=900, allowed_paths=["/tmp"])
"""
Generate Audio for Video — multi-model Gradio app.
Supported models
----------------
TARO – video-conditioned diffusion via CAVP + onset features (16 kHz, 8.192 s window)
MMAudio – multimodal flow-matching with CLIP/Synchformer + text prompt (44 kHz, 8 s window)
HunyuanFoley – text-guided foley via SigLIP2 + Synchformer + CLAP (48 kHz, up to 15 s)
"""
import functools
import gc
import html as _html
import math
import os
import sys
import json
import shutil
import tempfile
import random
import threading
import time
import warnings
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
# HunyuanVideo-Foley's DAC VAE loads its weights via the pre-parametrization
# torch.nn.utils.weight_norm API — harmless, but noisy on every model load.
warnings.filterwarnings(
"ignore",
message="`torch.nn.utils.weight_norm` is deprecated",
category=FutureWarning,
)
import torch
import numpy as np
import torchaudio
import soundfile as sf
import ffmpeg
import spaces
import gradio as gr
from huggingface_hub import HfApi, hf_hub_download, snapshot_download
# ------------------------------------------------------------------ #
# Gradio's own Video component (processing_utils.convert_video_to_ #
# playable_mp4) writes its ffmpeg output to Path(video_path) #
# .with_suffix(".mp4") — a no-op rename whenever the upload is #
# already named *.mp4 (e.g. HEVC-in-mp4). That makes it overwrite the #
# exact file the browser is concurrently streaming for the live #
# preview widget, corrupting the response mid-flight. Since HF Spaces #
# multiplexes a whole browser tab over one HTTP/2 connection, that #
# corrupted stream takes the *entire* connection down with it — even #
# unrelated in-flight requests like the generation job's SSE queue. #
# Patch it to convert a same-named copy instead, so the file the #
# browser already has a handle on is never mutated in place. #
import gradio.processing_utils as _gr_processing_utils
_gr_convert_video_to_playable_mp4 = _gr_processing_utils.convert_video_to_playable_mp4
def _patched_convert_video_to_playable_mp4(video_path: str) -> str:
p = Path(video_path)
if p.suffix.lower() == ".mp4":
copy_path = str(p.with_name(p.stem + "_gr_conv.mp4"))
shutil.copy2(video_path, copy_path)
video_path = copy_path
return _gr_convert_video_to_playable_mp4(video_path)
_gr_processing_utils.convert_video_to_playable_mp4 = _patched_convert_video_to_playable_mp4
# ------------------------------------------------------------------ #
# Gradio 5.x/6.x's Blocks.get_api_info() always returns an empty #
# unnamed_endpoints dict (confirmed still true through 5.50.0 and #
# 6.20.0 — it's simply never populated). But the JS client dispatches #
# every internal (non-api_name) click/change/upload event by numeric #
# fn_index, and looks that fn_index up in exactly that always-empty #
# dict, so endpoint_info resolves to undefined and every such event #
# logs a harmless "Too many arguments provided for the endpoint." #
# console warning (it doesn't drop data or block the request — see #
# gradio's client/src/helpers/api_info.ts map_data_to_params, which #
# just falls back to the raw args array). Since every dependency #
# already gets an auto-generated api_name and a fully-computed entry #
# in named_endpoints regardless, just mirror each of those into #
# unnamed_endpoints keyed by its numeric fn_index so the JS client's #
# lookup actually finds it. #
_orig_get_api_info = gr.Blocks.get_api_info
def _patched_get_api_info(self, all_endpoints: bool = False):
api_info = _orig_get_api_info(self, all_endpoints=all_endpoints)
for fn_index, fn in self.fns.items():
if not fn.fn or fn.api_name is False:
continue
entry = api_info["named_endpoints"].get(f"/{fn.api_name}")
if entry is not None:
api_info["unnamed_endpoints"][str(fn_index)] = entry
return api_info
gr.Blocks.get_api_info = _patched_get_api_info
# ================================================================== #
# CHECKPOINT CONFIGURATION #
# ================================================================== #