| """Unified checkpoint loading and audio-to-semantic-token inference.""" |
|
|
| import json |
| import math |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from safetensors import safe_open |
| from torch import nn |
|
|
| from .encoder import MERT2Encoder, SAMPLE_RATE as ENCODER_RATE |
|
|
| CODEC_SIZE = 32768 |
| HEAD_WINDOW = 512 |
|
|
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, width: int, hidden: int): |
| super().__init__() |
| self.norm = nn.LayerNorm(width) |
| self.gate = nn.Linear(width, hidden, bias=False) |
| self.up = nn.Linear(width, hidden, bias=False) |
| self.down = nn.Linear(hidden, width, bias=False) |
|
|
| def forward(self, x): |
| hidden = self.norm(x) |
| return x + self.down(F.silu(self.gate(hidden)) * self.up(hidden)) |
|
|
|
|
| class CVQHead(nn.Module): |
| def __init__(self, width: int, dim: int, hidden: int, tau: float): |
| super().__init__() |
| self.adapter = SwiGLU(width, hidden) if hidden else None |
| self.proj = nn.Linear(width, dim) |
| self.codebook = nn.Parameter(torch.empty(CODEC_SIZE, dim)) |
| self.tau = tau |
|
|
| def forward(self, x): |
| if self.adapter is not None: |
| x = self.adapter(x) |
| return F.normalize(self.proj(x), dim=-1) @ F.normalize(self.codebook, dim=-1).T / self.tau |
|
|
|
|
| def to_token_clock(per_frame: np.ndarray, offset: int) -> np.ndarray: |
| """Delay predictions so feature frame t supplies token t + offset.""" |
| indices = np.clip(np.arange(len(per_frame)) - offset, 0, len(per_frame) - 1) |
| return per_frame[indices] |
|
|
|
|
| class WannabeTokenizer(nn.Module): |
| """Load a unified WannabeTokenizer safetensors file without remote code.""" |
|
|
| def __init__(self, checkpoint: str | Path, device: str | torch.device = "cuda"): |
| super().__init__() |
| self.checkpoint = Path(checkpoint) |
| if not self.checkpoint.is_file(): |
| raise FileNotFoundError(f"Tokenizer checkpoint not found: {self.checkpoint}") |
| self.device = torch.device(device) |
| with safe_open(self.checkpoint, framework="pt") as f: |
| metadata = dict(f.metadata() or {}) |
| self._validate_format(metadata) |
| config = self._json_object(metadata, "encoder.config") |
| source = self._json_object(metadata, "head.source_spec") |
| self.layer = self._integer(source, "layer", minimum=0) |
| causal = source.get("causal") |
| if type(causal) is not bool: |
| raise ValueError("head.source_spec.causal must be a boolean") |
| self.chunk = float(source.get("chunk", 0)) |
| if not math.isfinite(self.chunk) or self.chunk < 0 or 0 < self.chunk < 1: |
| raise ValueError("head.source_spec.chunk must be 0 or at least 1 second") |
| self.offset = self._integer(metadata, "head.offset", minimum=0) |
| self.instnorm = self._boolean(metadata, "head.instnorm") |
| tau = float(metadata.get("head.tau", 0.07)) |
| if not math.isfinite(tau) or tau <= 0: |
| raise ValueError("head.tau must be a finite positive number") |
|
|
| encoder_keys = [key for key in f.keys() if key.startswith("encoder.")] |
| head_keys = [key for key in f.keys() if key.startswith("head.")] |
| if not encoder_keys or not head_keys: |
| raise ValueError("Unified checkpoint must contain encoder.* and head.* tensors") |
| encoder_state = {key[8:]: f.get_tensor(key) for key in encoder_keys} |
| head_state = {key[5:]: f.get_tensor(key).float() for key in head_keys} |
|
|
| self.encoder = MERT2Encoder(config, causal=causal, layers=self.layer + 1) |
| kept_encoder_state = { |
| key: value |
| for key, value in encoder_state.items() |
| if not (key.startswith("layers.") and int(key.split(".")[1]) > self.layer) |
| } |
| self.encoder.load_state_dict(kept_encoder_state, strict=True) |
| try: |
| dim = head_state["codebook"].shape[1] |
| width = head_state["proj.weight"].shape[1] |
| hidden = head_state["adapter.up.weight"].shape[0] if "adapter.up.weight" in head_state else 0 |
| except (KeyError, IndexError) as exc: |
| raise ValueError("Checkpoint is missing valid CVQ head tensors") from exc |
| if head_state["codebook"].shape[0] != CODEC_SIZE: |
| raise ValueError(f"Expected {CODEC_SIZE} codebook entries") |
| self.head = CVQHead(width, dim, hidden, tau) |
| self.head.load_state_dict(head_state, strict=True) |
| self.to(self.device).eval().requires_grad_(False) |
|
|
| @staticmethod |
| def _validate_format(metadata: dict) -> None: |
| if metadata.get("format") != "yue2sem-tokenizer-1": |
| raise ValueError("Expected a unified yue2sem-tokenizer-1 checkpoint") |
| if metadata.get("head.kind") != "cvq": |
| raise ValueError("This inference package supports CVQ tokenizer checkpoints") |
|
|
| @staticmethod |
| def _json_object(metadata: dict, key: str) -> dict: |
| try: |
| value = json.loads(metadata[key]) |
| except (KeyError, json.JSONDecodeError, TypeError) as exc: |
| raise ValueError(f"Checkpoint has invalid or missing {key}") from exc |
| if not isinstance(value, dict): |
| raise ValueError(f"{key} must contain a JSON object") |
| return value |
|
|
| @staticmethod |
| def _integer(values: dict, key: str, minimum: int) -> int: |
| try: |
| value = int(values[key]) |
| except (KeyError, TypeError, ValueError) as exc: |
| raise ValueError(f"Checkpoint has invalid or missing {key}") from exc |
| if isinstance(values[key], float) and not values[key].is_integer(): |
| raise ValueError(f"{key} must be an integer") |
| if value < minimum: |
| raise ValueError(f"{key} must be at least {minimum}") |
| return value |
|
|
| @staticmethod |
| def _boolean(metadata: dict, key: str) -> bool: |
| value = metadata.get(key) |
| if value not in ("0", "1"): |
| raise ValueError(f"{key} must be '0' or '1'") |
| return value == "1" |
|
|
| @torch.inference_mode() |
| def tokenize(self, waveform: torch.Tensor, batch_windows: int = 8) -> np.ndarray: |
| """Convert fp32 [channels, samples] 48 kHz audio to int32 tokens at 25 Hz.""" |
| import torchaudio |
|
|
| if waveform.ndim != 2 or waveform.shape[0] < 1: |
| raise ValueError("Expected audio shaped [channels, samples]") |
| if batch_windows < 1: |
| raise ValueError("batch_windows must be positive") |
| mono = torchaudio.functional.resample( |
| waveform.float().mean(0), 48000, ENCODER_RATE |
| ).to(self.device) |
| features = self.encoder.features(mono, self.layer, self.chunk) |
| |
| features = features.half().cpu().numpy() |
| return self._predict(features, batch_windows).astype(np.int32) |
|
|
| def _predict(self, features: np.ndarray, batch_windows: int) -> np.ndarray: |
| x = features.astype(np.float32) |
| if self.instnorm: |
| x = (x - x.mean(0)) / (x.std(0) + 1e-5) |
| x = torch.from_numpy(x) |
| total = len(x) |
| starts = list(range(0, max(1, total - HEAD_WINDOW + 1), HEAD_WINDOW // 2)) |
| if starts[-1] + HEAD_WINDOW < total: |
| starts.append(max(0, total - HEAD_WINDOW)) |
| output = torch.zeros(total, dtype=torch.long) |
| for first in range(0, len(starts), batch_windows): |
| group = starts[first : first + batch_windows] |
| windows = torch.stack( |
| [ |
| F.pad( |
| x[start : start + HEAD_WINDOW], |
| (0, 0, 0, HEAD_WINDOW - min(HEAD_WINDOW, total - start)), |
| ) |
| for start in group |
| ] |
| ).to(self.device) |
| with torch.autocast( |
| self.device.type, |
| dtype=torch.bfloat16, |
| enabled=self.device.type == "cuda", |
| ): |
| logits = self.head(windows).float() |
| for start, window_logits in zip(group, logits): |
| count = min(HEAD_WINDOW, total - start) |
| low = start + (0 if start == 0 else HEAD_WINDOW // 4) |
| high = start + count - (0 if start + count >= total else HEAD_WINDOW // 4) |
| prediction = window_logits[:count].argmax(-1) |
| output[low:high] = prediction[low - start : high - start].cpu() |
| return to_token_clock(output.numpy(), self.offset) |
|
|
