File size: 2,777 Bytes
c0bce59 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | """S9 — VoxCPM2 AudioVAE V2 encoder: audio → latent frames (25 Hz × 64) as ``.npy``.
Reuses the pinned, hash-verified model loading of the B2 round-trip worker
(:class:`nida.bench.workers.vae_voxcpm.VoxCPMVAEAdapter`): same commit, same weights, same
16 kHz input path. Output per request: float32 ``[num_frames, latent_dim]`` (the posterior
mean, deterministic) at ``params["out_path"]``, with the frame/duration facts of
:func:`nida.data.workers.finish_frame_output`.
"""
from __future__ import annotations
import importlib
from collections.abc import Mapping, Sequence
from nida.bench.protocol import WorkRequest
from nida.bench.workers.runner import main_for
from nida.bench.workers.vae_voxcpm import VoxCPMVAEAdapter
from nida.data.workers import finish_frame_output, out_path_param
__all__ = ["FRAME_HOP_MS", "FRAME_TOLERANCE_FRAMES", "VAEEncodeAdapter"]
#: 25.00 Hz latent rate (R2 / B2 design note: verified from the model source); every
#: output is checked against the audio duration, so a wrong value fails loudly.
FRAME_HOP_MS = 40.0
#: The encoder pads the input to whole frames; allow two frames of slack vs the duration.
FRAME_TOLERANCE_FRAMES = 2
class VAEEncodeAdapter(VoxCPMVAEAdapter):
name = "vae_encode"
def _encode(self, request: WorkRequest) -> Mapping[str, object]:
out_path = out_path_param(request)
torch = self._torch
soundfile = importlib.import_module("soundfile")
functional = importlib.import_module("torchaudio.functional")
audio, source_rate_hz = soundfile.read(request.audio_path, dtype="float32")
if audio.ndim != 1:
raise ValueError(f"{request.clip_id}: expected mono audio, got shape {audio.shape}")
in_rate_hz = int(self._vae.sample_rate)
wave = torch.from_numpy(audio).unsqueeze(0)
if source_rate_hz != in_rate_hz:
wave = functional.resample(wave, source_rate_hz, in_rate_hz)
duration_ms = 1000.0 * int(wave.shape[-1]) / in_rate_hz
with torch.inference_mode():
latents = self._vae.encode(wave.to(self._device), in_rate_hz) # (1, D, frames)
frames = latents[0].transpose(0, 1).float().cpu().numpy() # (frames, D)
payload = finish_frame_output(
out_path=out_path,
array=frames,
frame_hop_ms=FRAME_HOP_MS,
duration_ms=duration_ms,
tolerance_ms=FRAME_TOLERANCE_FRAMES * FRAME_HOP_MS,
what=request.clip_id,
)
payload["latent_dim"] = int(frames.shape[1])
return payload
def run_batch(self, batch: Sequence[WorkRequest]) -> list[Mapping[str, object]]:
return [self._encode(r) for r in batch]
if __name__ == "__main__":
raise SystemExit(main_for(VAEEncodeAdapter()))
|