Spaces:
Running on Zero
Running on Zero
| """Decode the frozen codec used to tokenize each matched LM's training data. | |
| Decoding goes through a causal multi-resolution iSTFT head declared in the registry entry's ``decoders`` | |
| (``codec_registry.json``, istft_head_a_20261002), keyed by decoder id: ``keep2_plain`` (the default, | |
| ``decoder_default``: the KEEP-2 pruned head, full-unfreezeall-mr-drop30-4x10) and ``full_mr`` (the full-size head, | |
| full-unfreezeall-mr). The heads share the codec, its frame rate and codebooks and differ only in layer counts. The 600k | |
| conv-tail decoder is never used. The 600k generator.pt only supplies the frozen encoder/quantizer and the initial | |
| decoder stages that the head checkpoint then overrides. ``decoder=None`` is the codec's default decoder everywhere.""" | |
| import functools | |
| import json | |
| import os | |
| import time | |
| from pathlib import Path | |
| import torch | |
| from zonos3.utils.hashing import sha256_file | |
| from zonos3.utils.release import hf_join, resolve_file | |
| REGISTRY = Path(__file__).resolve().parents[1] / "registries" / "codec_registry.json" | |
| _CACHE = {} | |
| _INFO = {} | |
| _GRAPHS = {} | |
| def registry(): | |
| return json.loads(REGISTRY.read_text()) | |
| def codec_names(): | |
| return {v["codec"] for v in registry().values()} | |
| def entry(codec): | |
| values = [v for v in registry().values() if v["codec"] == codec] | |
| if len(values) != 1: | |
| raise ValueError(f"Unknown matched codec {codec}") | |
| return values[0] | |
| def decoder_ids(codec): | |
| """The ids of the codec's decoders, its default (``decoder_default``) first, then the registry order.""" | |
| decoders = entry(codec).get("decoders") or {} | |
| default = entry(codec).get("decoder_default") | |
| if not decoders or default not in decoders: | |
| raise RuntimeError( | |
| f"matched codec {codec!r} declares no default iSTFT decoder; the conv-tail decoder is disabled" | |
| ) | |
| return [default, *(name for name in decoders if name != default)] | |
| def decoder_id(codec, decoder=None): | |
| """``decoder`` (None = the codec's default) checked against the codec's decoders (ValueError when unknown).""" | |
| ids = decoder_ids(codec) | |
| if decoder is None: | |
| return ids[0] | |
| if decoder not in ids: | |
| raise ValueError(f"matched codec {codec!r} has no decoder {decoder!r} (decoders: {', '.join(ids)})") | |
| return decoder | |
| def served_decoders(codec, wanted=None): | |
| """The decoder ids a server serves: ``wanted`` (a server flag's ids; None = every decoder of the codec), the | |
| codec's default first when it is among them, else the first wanted; ValueError for an id the codec lacks.""" | |
| ids = decoder_ids(codec) | |
| if wanted is None: | |
| return ids | |
| wanted = list(dict.fromkeys(str(name).strip() for name in wanted if str(name).strip())) | |
| if not wanted: | |
| raise ValueError(f"no decoder given (matched codec {codec!r} has {', '.join(ids)})") | |
| for name in wanted: | |
| decoder_id(codec, name) | |
| return sorted(wanted, key=lambda name: name != ids[0]) | |
| def decoder_spec(codec, decoder=None): | |
| """A registry decoder block; every matched codec decoder must be a causal iSTFT head (no conv-tail fallback).""" | |
| name = decoder_id(codec, decoder) | |
| spec = entry(codec)["decoders"][name] | |
| if spec.get("kind") != "istft_mr": | |
| raise RuntimeError( | |
| f"decoder {name!r} of matched codec {codec!r} is no iSTFT head; the conv-tail decoder is disabled" | |
| ) | |
| return spec | |
| def decoder_headroom_db(codec, decoder=None): | |
| return float(decoder_spec(codec, decoder).get("headroom_db", 0.0)) | |
| def decoder_info(codec, decoder=None): | |
| name = decoder_id(codec, decoder) | |
| spec = decoder_spec(codec, name) | |
| info = { | |
| "id": name, | |
| "kind": spec["kind"], | |
| "label": spec.get("label"), | |
| "checkpoint": hf_join(spec["campaign"], spec["checkpoint"]), | |
| "expect_step": spec.get("expect_step"), | |
| "headroom_db": float(spec.get("headroom_db", 0.0)), | |
| "causal": True, | |
| } | |
| info.update(_INFO.get((codec, name), {})) | |
| return info | |
| def load_model(codec, device, prune=True, decoder=None): | |
| """Build the iSTFT-head decoder exactly as the listening viewer's decode.py does: ISTFTCodec(cfg['model'], cfg['head']), | |
| pretrained generator weights, then the run's trainable state (decoder stages + head) from the pinned checkpoint.""" | |
| e = entry(codec) | |
| name = decoder_id(codec, decoder) | |
| spec = decoder_spec(codec, name) | |
| from zonos3.codec.istft_head_a_20261002.model import ISTFTCodec | |
| campaign = spec["campaign"] | |
| cfg = json.loads(Path(resolve_file(hf_join(campaign, spec["config"]))).read_text()) | |
| assert cfg.get("decoder", {}).get("type") != "original", ( | |
| "the registry points at a conv-tail (original decoder) config" | |
| ) | |
| generator_path = Path(resolve_file(cfg["init"]["generator_checkpoint"])) | |
| assert str(generator_path) == resolve_file(e["checkpoint"]), (str(generator_path), e["checkpoint"]) | |
| generator_state = torch.load(generator_path, map_location="cpu", mmap=True, weights_only=True) | |
| assert generator_state["step"] == e["codec_step"] and generator_state["config"]["model"] == cfg["model"] | |
| model = ISTFTCodec(cfg["model"], cfg["head"]) | |
| init = model.load_pretrained(generator_state["generator"]) | |
| checkpoint = Path(resolve_file(hf_join(campaign, spec["checkpoint"]))) | |
| state = torch.load(checkpoint, map_location="cpu", mmap=True, weights_only=False) | |
| if spec.get("expect_step") is not None: | |
| assert int(state["step"]) == int(spec["expect_step"]), ( | |
| str(checkpoint), | |
| int(state["step"]), | |
| spec["expect_step"], | |
| ) | |
| if "trainable" in state: | |
| model.load_trainable_state_dict(state["trainable"]) | |
| else: | |
| assert not model.trainable_stages | |
| model.head.load_state_dict(state["head"], strict=True) | |
| model.eval().requires_grad_(False) | |
| if prune: | |
| prune_model(model, e) | |
| _INFO[codec, name] = { | |
| "step": int(state["step"]), | |
| "checkpoint_bytes": checkpoint.stat().st_size, | |
| "checkpoint_mtime": time.strftime("%Y-%m-%dT%H:%M:%S", time.localtime(checkpoint.stat().st_mtime)), | |
| "model": type(model).__name__, | |
| "init": init if isinstance(init, dict) else None, | |
| } | |
| del state, generator_state | |
| return model.to(device) | |
| def prune_model(model, e): | |
| # ISTFTCodec never holds the encoder (its __init__ keeps only stages/quantizer/semantic_head/to_external of the base codec); | |
| # semantic_head is used by features() and must stay. | |
| for name in ("encoder", "to_internal", "semantic_encoder"): | |
| if hasattr(model, name): | |
| try: | |
| delattr(model, name) | |
| except Exception: | |
| pass | |
| def decode_model(model, codes, e): | |
| """[B, frames, codebooks] int -> float32 [B, samples]; fp32 with TF32 off, like the viewers' decode service.""" | |
| device = next(model.parameters()).device | |
| packed = codes.permute(0, 2, 1).to(device=device, dtype=torch.int64) | |
| expected = packed.shape[-1] * e["hop"] | |
| prev = (torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32) | |
| torch.backends.cuda.matmul.allow_tf32 = False | |
| torch.backends.cudnn.allow_tf32 = False | |
| try: | |
| with torch.inference_mode(): | |
| audio = model(packed) | |
| finally: | |
| torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32 = prev | |
| if audio.shape != (packed.shape[0], 1, expected): | |
| raise RuntimeError(f"Codec output shape {tuple(audio.shape)}; expected {expected} samples") | |
| if not torch.isfinite(audio).all(): | |
| raise RuntimeError("Nonfinite codec audio") | |
| return audio.float().squeeze(1).cpu() | |
| def _key(codec, decoder): | |
| return codec, decoder_id(codec, decoder), os.environ.get("ZONOS3_AUDIO_CODEC_DEVICE", "cuda:0") | |
| def _model(codec, decoder=None): | |
| """A decoder of the codec on ZONOS3_AUDIO_CODEC_DEVICE (default cuda:0), loaded once after the checkpoint's sha256.""" | |
| e = entry(codec) | |
| key = _key(codec, decoder) | |
| if key not in _CACHE: | |
| actual = sha256_file(resolve_file(e["checkpoint"])) | |
| assert actual == e["checkpoint_sha256"], ("generator checkpoint sha256", actual, e["checkpoint_sha256"]) | |
| _CACHE[key] = load_model(codec, key[2], decoder=key[1]) | |
| return _CACHE[key] | |
| def decode(codes, codec, decoder=None): | |
| e = entry(codec) | |
| if codes.ndim != 3 or codes.shape[-1] != e["codebooks"] or bool(((codes < 0) | (codes >= 4096)).any()): | |
| raise ValueError("Invalid matched codec tokens") | |
| return decode_model(_model(codec, decoder), codes, e) | |
| def stream(codec, decoder=None): | |
| """A fresh stateful decoder (ISTFTStream) on the cached model: decode(codes [n, codebooks]) -> final samples, flush(). | |
| It replays the CUDA graphs captured by ``capture_stream_graphs`` for that decoder and matching chunk shapes.""" | |
| from zonos3.codec.istft_head_a_20261002.stream import ISTFTStream | |
| return ISTFTStream(_model(codec, decoder), graphs=_GRAPHS.get(_key(codec, decoder))) | |
| def capture_stream_graphs(codec, chunks, decoder=None) -> None: | |
| """CUDA-graph the chunk shapes of a stream decoding ``chunks`` frames per call from the start (bit-identical to | |
| the eager decode; a no-op unless the codec runs on CUDA). Each decoder keeps its own graphs.""" | |
| from zonos3.codec.istft_head_a_20261002.stream import ChunkGraphs | |
| model = _model(codec, decoder) | |
| if next(model.parameters()).device.type != "cuda": | |
| return | |
| _GRAPHS.setdefault(_key(codec, decoder), ChunkGraphs()).capture(model, chunks) | |