File size: 3,690 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | """VoxCPM2 AudioVAE V2 decoder: latent frames ``.npy`` → 48 kHz float WAV (listening samples).
The input of each request is a latent array ``[num_frames, latent_dim]`` (float32, *VAE*
scale, i.e. already de-normalised) whose path is passed as ``audio_path`` (the worker
protocol's input file). Output: ``params["out_path"]``, 32-bit float mono WAV at the
decoder's output rate, written atomically; with ``params["out_path_16k"]`` also a 16 kHz
16-bit copy (for ASR checks). Same pinned weights as :mod:`vae_encode`.
"""
from __future__ import annotations
import importlib
from collections.abc import Mapping, Sequence
from pathlib import Path
from nida.bench.protocol import WorkRequest
from nida.bench.workers.runner import main_for
from nida.bench.workers.vae_voxcpm import VoxCPMVAEAdapter
__all__ = ["ASR_RATE_HZ", "VAEDecodeAdapter", "wav_out_path"]
#: Sample rate of the optional ASR copy.
ASR_RATE_HZ = 16000
def wav_out_path(request: WorkRequest) -> Path:
"""The ``.wav`` output path named by ``request.params["out_path"]``.
Raises:
ValueError: if the parameter is missing or does not end in ``.wav``.
"""
raw = request.params.get("out_path")
if not raw:
raise ValueError(f"{request.clip_id}: missing params.out_path")
path = Path(raw)
if path.suffix != ".wav":
raise ValueError(f"{request.clip_id}: params.out_path must end in .wav, got {raw!r}")
return path
class VAEDecodeAdapter(VoxCPMVAEAdapter):
name = "vae_decode"
def _decode(self, request: WorkRequest) -> Mapping[str, object]:
out_path = wav_out_path(request)
torch = self._torch
numpy = importlib.import_module("numpy")
soundfile = importlib.import_module("soundfile")
latents = numpy.load(request.audio_path, allow_pickle=False)
if latents.ndim != 2 or latents.shape[0] < 1:
raise ValueError(f"{request.clip_id}: expected [frames, dim], got {latents.shape}")
if not bool(numpy.isfinite(latents).all()):
raise ValueError(f"{request.clip_id}: non-finite latents")
z = torch.from_numpy(latents.astype("float32")).transpose(0, 1).unsqueeze(0)
with torch.inference_mode():
decoded = self._vae.decode(z.to(self._device)) # (1, 1, samples)
out = decoded.reshape(-1).float().cpu().numpy()
rate_hz = int(self._vae.out_sample_rate)
target = Path(out_path)
target.parent.mkdir(parents=True, exist_ok=True)
tmp = target.with_name(f".{target.name}.tmp.wav") # never leave a partial output
soundfile.write(str(tmp), out, rate_hz, subtype="FLOAT")
tmp.replace(target)
payload: dict[str, object] = {}
if request.params.get("out_path_16k"):
functional = importlib.import_module("torchaudio.functional")
low = functional.resample(torch.from_numpy(out), rate_hz, ASR_RATE_HZ).numpy()
target_16k = Path(request.params["out_path_16k"])
tmp_16k = target_16k.with_name(f".{target_16k.name}.tmp.wav")
soundfile.write(str(tmp_16k), numpy.clip(low, -1.0, 1.0), ASR_RATE_HZ, "PCM_16")
tmp_16k.replace(target_16k)
payload["out_path_16k"] = str(target_16k)
return {
**payload,
"out_path": str(out_path),
"out_sample_rate_hz": rate_hz,
"num_samples": int(out.shape[0]),
"num_latent_frames": int(latents.shape[0]),
}
def run_batch(self, batch: Sequence[WorkRequest]) -> list[Mapping[str, object]]:
return [self._decode(r) for r in batch]
if __name__ == "__main__":
raise SystemExit(main_for(VAEDecodeAdapter()))
|