128 lines
4.7 KiB
Python
128 lines
4.7 KiB
Python
"""In-memory shim for lightx2v.utils.audio_io.load_audio_file.
|
|
|
|
Routes ComfyUI AUDIO tensors into the runner without round-tripping through
|
|
a temp WAV file. Sentinel paths start with SENTINEL_PREFIX; the patched
|
|
loader returns a stashed tensor instead of touching disk.
|
|
|
|
Scope: only the single-AUDIO ComfyUI input path uses this shim. V3 multi-talker
|
|
padding still writes real WAV files (its inputs are external file paths / URLs,
|
|
not ComfyUI tensors), and those reads fall through to the original loader.
|
|
|
|
Patching strategy: lightx2v consumers do `from lightx2v.utils.audio_io import
|
|
load_audio_file`, which captures the original function at import time. So
|
|
patching only `lightx2v.utils.audio_io.load_audio_file` would miss them — we
|
|
rebind every known import site too. New consumers must be added to _PATCH_SITES.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import threading
|
|
import uuid
|
|
from typing import Dict, Tuple
|
|
|
|
import torch
|
|
|
|
SENTINEL_PREFIX = "<lightx2v-mem-audio:"
|
|
SENTINEL_SUFFIX = ">"
|
|
|
|
_REGISTRY: Dict[str, Tuple[torch.Tensor, int]] = {}
|
|
_LOCK = threading.Lock()
|
|
_PATCHED = False
|
|
|
|
# (module dotted-path, attribute name). Each entry rebinds that module's
|
|
# `load_audio_file` attribute to the shim. Add new early-binders here as
|
|
# upstream changes.
|
|
_PATCH_SITES = (
|
|
("lightx2v.utils.audio_io", "load_audio_file"),
|
|
("lightx2v.models.runners.wan.wan_audio_runner", "load_audio_file"),
|
|
("lightx2v.shot_runner.rs2v_infer", "load_audio_file"),
|
|
("lightx2v.shot_runner.stream_infer", "load_audio_file"),
|
|
)
|
|
|
|
|
|
def register(waveform: torch.Tensor, sample_rate: int) -> str:
|
|
"""Register a [C, T] float32 waveform; return a sentinel "path" for it."""
|
|
if waveform.dim() != 2:
|
|
raise ValueError(f"waveform must be [C, T]; got shape {tuple(waveform.shape)}")
|
|
token = uuid.uuid4().hex
|
|
sentinel = f"{SENTINEL_PREFIX}{token}{SENTINEL_SUFFIX}"
|
|
stashed = waveform.detach().to(torch.float32).cpu().contiguous()
|
|
with _LOCK:
|
|
_REGISTRY[sentinel] = (stashed, int(sample_rate))
|
|
return sentinel
|
|
|
|
|
|
def release(sentinel: str) -> None:
|
|
with _LOCK:
|
|
_REGISTRY.pop(sentinel, None)
|
|
|
|
|
|
def is_sentinel(path) -> bool:
|
|
return isinstance(path, str) and path.startswith(SENTINEL_PREFIX)
|
|
|
|
|
|
def comfyui_audio_to_loader_pair(audio_dict) -> Tuple[torch.Tensor, int]:
|
|
"""ComfyUI AUDIO {"waveform": [B, C, T], "sample_rate": int} -> ([C, T], sr).
|
|
|
|
ComfyUI gives waveform as [B, C, T] float32 in [-1, 1] (B usually 1).
|
|
lightx2v's load_audio_file returns [C, T] (channels_first=True). We
|
|
squeeze the batch dim here; mono-down and resample happen downstream in
|
|
AudioProcessor / ShotRS2VPipeline so the shim stays format-agnostic.
|
|
"""
|
|
waveform = audio_dict["waveform"]
|
|
if waveform.dim() == 3:
|
|
waveform = waveform[0]
|
|
if waveform.dim() != 2:
|
|
raise ValueError(f"Unexpected ComfyUI AUDIO waveform shape {tuple(waveform.shape)}")
|
|
return waveform, int(audio_dict["sample_rate"])
|
|
|
|
|
|
def install() -> None:
|
|
"""Idempotent monkey-patch. Safe to call multiple times."""
|
|
global _PATCHED
|
|
if _PATCHED:
|
|
return
|
|
import importlib
|
|
|
|
original = None
|
|
targets = []
|
|
for module_path, attr in _PATCH_SITES:
|
|
try:
|
|
module = importlib.import_module(module_path)
|
|
except ImportError:
|
|
logging.warning(f"_audio_shim: {module_path} not importable; skipping")
|
|
continue
|
|
fn = getattr(module, attr, None)
|
|
if fn is None:
|
|
logging.warning(f"_audio_shim: {module_path}.{attr} missing; skipping")
|
|
continue
|
|
if original is None:
|
|
original = fn
|
|
targets.append((module, attr))
|
|
|
|
if original is None:
|
|
raise RuntimeError("_audio_shim.install: no patch sites resolved; lightx2v not installed?")
|
|
|
|
def _patched(uri, frame_offset: int = 0, num_frames: int = -1, channels_first: bool = True):
|
|
if is_sentinel(uri):
|
|
with _LOCK:
|
|
entry = _REGISTRY.get(uri)
|
|
if entry is None:
|
|
raise FileNotFoundError(f"Stale lightx2v in-memory audio sentinel: {uri}")
|
|
tensor, sr = entry
|
|
# Slice semantics mirror torchaudio.load(frame_offset, num_frames).
|
|
if frame_offset > 0 or num_frames > 0:
|
|
end = tensor.shape[-1] if num_frames < 0 else frame_offset + num_frames
|
|
tensor = tensor[..., frame_offset:end]
|
|
if not channels_first:
|
|
tensor = tensor.transpose(0, 1).contiguous()
|
|
return tensor, sr
|
|
return original(uri, frame_offset=frame_offset, num_frames=num_frames, channels_first=channels_first)
|
|
|
|
for module, attr in targets:
|
|
setattr(module, attr, _patched)
|
|
|
|
_PATCHED = True
|
|
logging.info(f"_audio_shim installed at {len(targets)} site(s)")
|