Spaces:
Running on Zero
Running on Zero
| """Frozen dual codec: native SSL teacher -> semantic Q1, acoustic residual RVQ. | |
| Each independent lane retains both external-rate teacher audio and full-file | |
| resampled codec audio. Packing external hop samples into the second channel is | |
| lossless and only aligns the two sample grids for the existing batch scheduler. | |
| The teacher sees 20-second lanes, including the same 10-second real left context | |
| as the acoustic encoder. It retains native bidirectional attention within a lane. | |
| """ | |
| import json | |
| import time | |
| from types import SimpleNamespace | |
| import numpy as np | |
| import torch | |
| from torch.nn import functional as F | |
| from zonos3.codec.data_processing.independent_batcher import IndependentBatcher | |
| from zonos3.codec.data_processing.torch_utils import compile_fixed | |
| from zonos3.codec.hybrid10_dual_encoding.inference import extract_ssl | |
| from zonos3.codec.hybrid10_dual_encoding.model import MossCodec | |
| from zonos3.codec.resample_exact import resample_exact | |
| from zonos3.codec.s1dac.losses import SemanticTeacher | |
| class MossAutoencoder: | |
| def __init__( | |
| self, | |
| args, | |
| *, | |
| generator_path: str, | |
| teacher_dir: str, | |
| step: int, | |
| codebook_sizes: list[int], | |
| name: str, | |
| checkpoint_sha256: str, | |
| device: str, | |
| ): | |
| torch.set_num_threads(1) | |
| torch.backends.cuda.matmul.allow_tf32 = False | |
| torch.backends.cudnn.allow_tf32 = False | |
| self.device = torch.device(device) | |
| if self.device.type != "cuda": | |
| raise RuntimeError("Production encoding requires an allocated GPU") | |
| torch.cuda.set_device(self.device) | |
| self.receipt = { | |
| "step": step, | |
| "codebook_sizes": list(codebook_sizes), | |
| "output_name": name, | |
| "sha256": checkpoint_sha256, | |
| "checkpoint": generator_path, | |
| } | |
| self.teacher_dir = teacher_dir | |
| state = torch.load(generator_path, map_location="cpu", weights_only=True, mmap=True) | |
| cfg = state["config"]["model"] | |
| self.cfg = cfg | |
| assert state["step"] == self.receipt["step"] | |
| self.codebook_sizes = self.receipt["codebook_sizes"] | |
| assert self.codebook_sizes == [cfg["semantic_codebook_size"]] + [cfg["residual_codebook_size"]] * 9 | |
| assert cfg["residual_codebooks"] == 9 and cfg["frame_rate"] == 15 and cfg["dual_encoding"] | |
| assert ( | |
| cfg["sample_rate"], | |
| cfg["internal_sample_rate"], | |
| cfg["internal_hop_length"], | |
| cfg["external_samples_per_frame"], | |
| ) == (44100, 44160, 2944, 2940) | |
| model = MossCodec(cfg) | |
| model.load_state_dict(state["generator"], strict=True, assign=True) | |
| model.eval().requires_grad_(False) | |
| self.resample = model.to_internal.cpu() | |
| # Both semantic modules are required for token encoding. | |
| del model.decoder, model.to_external, model.to_internal | |
| self.model = model.to(self.device) | |
| self.teacher = ( | |
| SemanticTeacher(teacher_dir, layer=state["config"]["teacher"]["layer"], dtype=torch.bfloat16) | |
| .to(self.device) | |
| .eval() | |
| .requires_grad_(False) | |
| ) | |
| del state | |
| self.sample_rate = 44100 | |
| self.feature_dim = 10 | |
| self.name = self.receipt["output_name"] | |
| self.batch_size = args.hybrid_batch_size | |
| self.chunk_seconds = args.hybrid_chunk_seconds | |
| self.context_seconds = args.hybrid_context_seconds | |
| self.execution = ( | |
| "cuda_graph_acoustic_encoder_native_ssl_and_dual_rvq" | |
| if args.hybrid_cuda_graph | |
| else "eager_native_ssl_and_dual_rvq" | |
| ) | |
| self.packer = IndependentBatcher( | |
| torch=torch, | |
| device=self.device, | |
| sample_rate=44160, | |
| hop_length=2944, | |
| feature_dim=10, | |
| prepare=self._prepare, | |
| run=self._run, | |
| batch_size=self.batch_size, | |
| chunk_seconds=self.chunk_seconds, | |
| continuation_context_seconds=self.context_seconds, | |
| window=32, | |
| output_dtype=np.int16, | |
| label=self.name, | |
| input_dtype=torch.float32, | |
| input_channels=2, | |
| ) | |
| self.optimized = None | |
| print(json.dumps(dict(event="model_loaded", **self.metadata())), flush=True) | |
| with torch.inference_mode(): | |
| g = torch.Generator(device=self.device).manual_seed(728) | |
| warmup = ( | |
| torch.randn((self.batch_size, 1, self.packer.chunk_samples), generator=g, device=self.device) * 0.05 | |
| ) | |
| expected = self._encoder(warmup).clone() | |
| self.optimized = compile_fixed( | |
| torch, | |
| self._encoder, | |
| warmup, | |
| enabled=args.hybrid_cuda_graph, | |
| cache_name=self.name, | |
| strategy="cuda_graph", | |
| ) | |
| assert torch.equal(self.optimized(warmup) if self.optimized else self._encoder(warmup), expected), ( | |
| "Graph altered acoustic latent" | |
| ) | |
| for scale in [0.05, 0.1]: | |
| packed = ( | |
| torch.randn((self.batch_size, 2, self.packer.chunk_samples), generator=g, device=self.device) | |
| * scale | |
| ) | |
| assert torch.equal(self._run(packed), self._codes(packed)), "Graph/native dual tokens disagree" | |
| report = { | |
| "event": "inference_validation_passed", | |
| "step": self.receipt["step"], | |
| "codebooks": 10, | |
| "codebook_sizes": self.codebook_sizes, | |
| "graph_token_agreement": 1.0, | |
| "native_quantize_dual": True, | |
| "teacher_layer": 16, | |
| "batch_size": self.batch_size, | |
| "chunk_seconds": self.chunk_seconds, | |
| "peak_gib": torch.cuda.max_memory_allocated() / 1024**3, | |
| } | |
| print(json.dumps(report), flush=True) | |
| def _external(self, packed): | |
| b = packed.shape[0] | |
| return packed[:, 1:2].reshape(b, 1, -1, 2944)[..., :2940].reshape(b, 1, -1) | |
| def _encoder(self, audio): | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| return self.model.encoder(audio) | |
| def _dual(self, packed, optimized): | |
| # Uses the exact CPU feature preprocessor and frozen teacher from training. | |
| audio = packed[:, 0:1].contiguous() | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| ssl = extract_ssl(self._external(packed), self.teacher) | |
| latent = self.optimized(audio) if optimized and self.optimized is not None else self.model.encoder(audio) | |
| result = self.model.quantize_dual(latent, ssl) | |
| return result["codes"].transpose(1, 2) | |
| def _codes(self, packed): | |
| return self._dual(packed, False) | |
| def _run(self, packed): | |
| return self._dual(packed, True) | |
| def _prepare(self, audio, sample_rate): | |
| if sample_rate != 44100: | |
| raise ValueError("Expected 44100 Hz") | |
| audio = np.asarray(audio, dtype=np.float32) | |
| if audio.ndim != 2 or audio.shape[0] != 1 or not audio.shape[-1]: | |
| raise ValueError("Expected nonempty mono audio") | |
| if not np.isfinite(audio).all(): | |
| raise ValueError("Nonfinite source audio") | |
| with torch.inference_mode(): | |
| wave = torch.from_numpy(np.ascontiguousarray(audio)).unsqueeze(0) | |
| frames = (wave.shape[-1] + 2939) // 2940 | |
| internal = resample_exact(wave, self.resample) | |
| internal = F.pad(internal, (0, frames * 2944 - internal.shape[-1])) | |
| external = F.pad(wave, (0, frames * 2940 - wave.shape[-1])).reshape(1, 1, frames, 2940) | |
| external = F.pad(external, (0, 4)).reshape(1, 1, frames * 2944) | |
| return torch.cat([internal, external], dim=1) | |
| def metadata(self): | |
| return { | |
| "name": self.name, | |
| "model": "s1_s2_dac/hybrid_dual_encoding", | |
| "checkpoint_step": self.receipt["step"], | |
| "checkpoint_sha256": self.receipt["sha256"], | |
| "checkpoint": self.receipt["checkpoint"], | |
| "sample_rate": 44100, | |
| "internal_sample_rate": 44160, | |
| "hop_length": 2940, | |
| "frame_rate": 15, | |
| "n_codebooks": 10, | |
| "codebook_sizes": self.codebook_sizes, | |
| "semantic_codebook_index": 0, | |
| "representation": "discrete_codes", | |
| "layout": ["time", "codebook"], | |
| "output_dtype": np.dtype(np.int16).str, | |
| "batch_size": self.batch_size, | |
| "chunk_seconds": self.chunk_seconds, | |
| "continuation_context_seconds": self.context_seconds, | |
| "resample_policy": "full_file_before_hop_aligned_chunking", | |
| "long_file_policy": "overlap_left_context_then_discard_context_tokens", | |
| "precision": "fp32_weights_bf16_encoder_teacher_semantic_modules_fp32_rvq", | |
| "execution": self.execution, | |
| "ssl_teacher_layer": 16, | |
| "ssl_teacher": self.teacher_dir, | |
| "ssl_preprocessing": "native_training_44100_to_16000_fbank_stride2", | |
| "ssl_context_policy": "native_bidirectional_teacher_on_independent_padded_lanes", | |
| } | |
| def reference_metadata(self, audio, sample_rate, output): | |
| expected = (audio.shape[-1] + 2939) // 2940 | |
| if output.shape != (expected, 10): | |
| raise ValueError(f"Incorrect token shape {output.shape}, expected {(expected, 10)}") | |
| if output.min() < 0 or np.any(output >= np.asarray(self.codebook_sizes)[None, :]): | |
| raise ValueError("Invalid codebook index") | |
| return { | |
| "audio_samples": int(audio.shape[-1]), | |
| "audio_sample_rate": sample_rate, | |
| "checkpoint_step": self.receipt["step"], | |
| } | |
| def encode_stream(self, samples): | |
| count = 0 | |
| seconds = 0.0 | |
| start = time.monotonic() | |
| for sample, codes in self.packer.encode_stream(samples): | |
| if sample.error is None: | |
| self.reference_metadata(sample.audio, sample.sample_rate, codes) | |
| seconds += sample.audio.shape[-1] / sample.sample_rate | |
| count += 1 | |
| if count % 100 == 0: | |
| print( | |
| json.dumps( | |
| { | |
| "event": "encoding_progress", | |
| "files": count, | |
| "audio_hours": seconds / 3600, | |
| "audio_seconds_per_wall_second": seconds / (time.monotonic() - start), | |
| } | |
| ), | |
| flush=True, | |
| ) | |
| yield sample, codes | |
| def encode(self, audio, sample_rate): | |
| return next(iter(self.encode_stream([SimpleNamespace(audio=audio, sample_rate=sample_rate, error=None)])))[1] | |