| """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"] |
|
|
| |
| 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)) |
| 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") |
| 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())) |
|
|
