| """IndiCodec inference graph. No weight download. |
| |
| Architecture classes still come from installed `neucodec` (Apache 2.0, commit ed3e6cd). |
| The release model.safetensors holds every tensor, including the 16 w2v-BERT 2.0 (MIT) layers we use, built from the |
| shipped w2v_bert/config.json. The internal .pt has no w2v-BERT, so that path still fetches it (pinned revision). |
| """ |
| from __future__ import annotations |
|
|
| import math |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| SEM_LAYER = 16 |
| HOP16, HOP24 = 320, 480 |
| LEVELS_S = [3] * 8 |
| W2V = "facebook/w2v-bert-2.0" |
| W2V_REV = "da985ba0987f70aaeb84a80f2851cfac8c697a7b" |
| W2V_DIR = Path(__file__).resolve().parent / "w2v_bert" |
| NEED_GEN = ( |
| "fsq_s.", "fsq_a.", "backbone.", "head.", |
| "CodecEnc.", "SemanticEncoder_module.", |
| ) |
| |
| TRAIN_ONLY = ("SemanticDecoder_module.",) |
|
|
|
|
| class IndiNet(nn.Module): |
| """IndiCodec encode / decode only. Refuses phase2d, rope_fix, continuous, joint FSQ.""" |
|
|
| def __init__(self, levels_a: list[int], w2v_bundled: bool = True): |
| super().__init__() |
| from transformers import Wav2Vec2BertConfig, Wav2Vec2BertModel |
| from vector_quantize_pytorch import FSQ |
| from neucodec.codec_encoder import CodecEncoder |
| from neucodec.codec_decoder_vocos import ISTFTHead, VocosBackbone |
| from neucodec.module import SemanticEncoder |
|
|
| self.w2v_bundled = bool(w2v_bundled) |
| if self.w2v_bundled: |
| sm = Wav2Vec2BertModel(Wav2Vec2BertConfig.from_json_file(str(W2V_DIR / "config.json"))) |
| else: |
| sm = Wav2Vec2BertModel.from_pretrained(W2V, revision=W2V_REV, output_hidden_states=True) |
| sm.encoder.layers = sm.encoder.layers[:SEM_LAYER] |
| if len(sm.encoder.layers) != SEM_LAYER: |
| raise SystemExit(f"w2v-BERT has {len(sm.encoder.layers)} layers, IndiCodec needs {SEM_LAYER}") |
| self.semantic_model = sm.eval().requires_grad_(False) |
| self.SemanticEncoder_module = SemanticEncoder(1024, 1024, 1024) |
| self.CodecEnc = CodecEncoder() |
| self.backbone = VocosBackbone(hidden_dim=1024, depth=12, heads=16, pos_meb_dim=64) |
| self.head = ISTFTHead(dim=1024, n_fft=HOP24 * 4, hop_length=HOP24, padding="same") |
| self.fsq_s = FSQ(levels=LEVELS_S, dim=1024) |
| self.levels_a = list(levels_a) |
| self.fsq_a = FSQ(levels=self.levels_a, dim=1024) |
| self.tail_frames = 0 |
|
|
| def load_gen(self, gen: dict) -> None: |
| gen = {k: v for k, v in gen.items() if not k.startswith(TRAIN_ONLY)} |
| missing_pref = [p for p in NEED_GEN if not any(k.startswith(p) for k in gen)] |
| if missing_pref: |
| raise SystemExit(f"ckpt gen missing {missing_pref}; this infer needs the IndiCodec trained towers") |
| r = self.load_state_dict(gen, strict=False) |
| bad_u = [k for k in r.unexpected_keys if k.startswith("head.refine.")] |
| if bad_u: |
| raise RuntimeError(f"checkpoint has PhaseRefiner2D ({bad_u[0]}); IndiCodec infer refuses it") |
| bad_m = [m for m in r.missing_keys if self.w2v_bundled or not m.startswith("semantic_model.")] |
| if bad_m: |
| raise SystemExit(f"ckpt missing {len(bad_m)} keys, e.g. {bad_m[:5]}") |
|
|
| @torch.no_grad() |
| def semantic_target(self, feats: torch.Tensor) -> torch.Tensor: |
| return self.semantic_model(feats).hidden_states[SEM_LAYER].transpose(1, 2) |
|
|
| def towers(self, wav16: torch.Tensor, feats: torch.Tensor): |
| return self.semantic_target(feats), self.CodecEnc(wav16).transpose(1, 2) |
|
|
| def encode(self, wav16: torch.Tensor, feats: torch.Tensor): |
| sem, ac = self.towers(wav16, feats) |
| s_enc = self.SemanticEncoder_module(sem) |
| t = min(sem.shape[-1], ac.shape[-1], s_enc.shape[-1]) |
| t -= t % 2 |
| sem, ac, s_enc = sem[..., :t], ac[..., :t], s_enc[..., :t] |
| s25 = F.avg_pool1d(s_enc, 2, 2) |
| zq_s, idx_s = self.fsq_s(s25.transpose(1, 2)) |
| c50 = zq_s.repeat_interleave(2, dim=1) |
| res = ac.transpose(1, 2) - c50.detach() |
| zq_a, idx_a = self.fsq_a(res) |
| return {"idx_s": idx_s, "idx_a": idx_a, "t": t} |
|
|
| def codes_to_latent(self, idx_s: torch.Tensor, idx_a: torch.Tensor | None) -> torch.Tensor: |
| c50 = self.fsq_s.indices_to_codes(idx_s).repeat_interleave(2, dim=1) |
| if idx_a is None: |
| return c50 |
| za = self.fsq_a.indices_to_codes(idx_a) |
| t = min(c50.shape[1], za.shape[1]) |
| return c50[:, :t] + za[:, :t] |
|
|
| def decode_latent(self, h: torch.Tensor) -> torch.Tensor: |
| return self.head(self.backbone(h))[0] |
|
|
| @torch.no_grad() |
| def decode_codes(self, idx_s, idx_a, mode: str = "full") -> torch.Tensor: |
| if mode == "zero_c": |
| return self.decode_latent(self.fsq_a.indices_to_codes(idx_a)) |
| if mode == "zero_ac": |
| return self.decode_latent(self.codes_to_latent(idx_s, None)) |
| return self.decode_latent(self.codes_to_latent(idx_s, idx_a)) |
|
|
|
|
| _FE = None |
|
|
|
|
| def features(wav16_np_list): |
| global _FE |
| if _FE is None: |
| from transformers import AutoFeatureExtractor |
| _FE = AutoFeatureExtractor.from_pretrained(str(W2V_DIR)) |
| return _FE(wav16_np_list, sampling_rate=16_000, return_tensors="pt").input_features |
|
|
|
|
| def to_rate(w, src_sr, tgt): |
| from scipy.signal import resample_poly |
| if int(src_sr) == int(tgt): |
| return np.ascontiguousarray(w, dtype=np.float32) |
| g = math.gcd(int(src_sr), int(tgt)) |
| return np.ascontiguousarray(resample_poly(np.asarray(w, np.float64), tgt // g, src_sr // g), dtype=np.float32) |
|
|
|
|
| def pad_w16(w16: np.ndarray, tail_frames: int) -> np.ndarray: |
| return np.pad(w16, (0, (-len(w16)) % HOP16 + int(tail_frames) * HOP16)) |
|
|
|
|
| @torch.no_grad() |
| def encode_wav(model: IndiNet, wav, sr, device, amp: bool = True): |
| w16 = pad_w16(to_rate(wav, sr, 16000), getattr(model, "tail_frames", 0)) |
| feats = features([w16]).to(device) |
| x = torch.from_numpy(w16)[None, None].to(device) |
| with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp and device != "cpu"): |
| e = model.encode(x, feats) |
| return e["idx_s"], e["idx_a"], e["t"] |
|
|
