nida-tts / src /nida /data /workers /vae_encode.py

Nida v0.5 (research preview): weights, 16 voices, code, model card

c0bce59 verified 23 days ago

2.78 kB

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

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