add support for optional num_quantizers argument

modeling_moss_audio_tokenizer.py

@@ -1,4 +1,3 @@
-# coding=utf-8
 # Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -26,9 +25,8 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from transformers.modeling_utils import PreTrainedAudioTokenizerBase
-from transformers.utils import ModelOutput, auto_docstring, logging
-
+from ...modeling_utils import PreTrainedAudioTokenizerBase
+from ...utils import ModelOutput, auto_docstring, logging
 from .configuration_moss_audio_tokenizer import MossAudioTokenizerConfig
 
 
@@ -372,7 +370,7 @@ def create_sin_embedding(
     dim: int,
     max_period: float = 10000,
     dtype: torch.dtype = torch.float32,
-):
+) -> torch.Tensor:
     """Create sinusoidal positional embedding with shape [B, T, C]."""
     if dim % 2 != 0:
         raise ValueError(f"Sinusoidal embedding requires even dim, got dim={dim}")
@@ -406,7 +404,7 @@ class KVCacheResult:
         return iter((self.keys, self.values, self.positions))
 
     @staticmethod
-    def from_kv(keys: torch.Tensor, values: torch.Tensor) -> "KVCacheResult":
+    def from_kv(keys: torch.Tensor, values: torch.Tensor) -> KVCacheResult:
         B, H, T, D = keys.shape
         positions = torch.arange(T, device=keys.device, dtype=torch.long)
         return KVCacheResult(keys, values, positions.expand(B, -1))
@@ -506,7 +504,7 @@ def apply_weights_per_step(
     schedule: list[int] | None,
     x: torch.Tensor,
     offset: int | None,
-):
+) -> torch.Tensor:
     """Apply different weights for each time step."""
     if len(modules) == 1:
         return modules[0](x)
@@ -1088,7 +1086,6 @@ class MossAudioTokenizerLFQ(nn.Module):
             - 2 * encodings @ codebook.t()
             + codebook.pow(2).sum(1, keepdim=True).t()
         )
-
         indices = (-dist).max(1)[1]
         indices = indices.reshape(latents.size(0), -1)
         z_q = self.decode_code_wo_out_proj(indices).float()
@@ -1306,10 +1303,6 @@ class MossAudioTokenizerPreTrainedModel(PreTrainedAudioTokenizerBase):
         "MossAudioTokenizerResidualLFQ",
     ]
 
-    def _init_weights(self, module: nn.Module) -> None:
-        if isinstance(module, MossAudioTokenizerLayerScale):
-            nn.init.constant_(module.scale, 1e-4)
-
 
 @auto_docstring(
     custom_intro="""
@@ -1348,7 +1341,7 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
                         context=int(current_frame_rate * self.causal_transformer_context_duration),
                     )
                 )
-            current_frame_rate /= cast(MossAudioTokenizerPatchedPretransform, self.encoder[-1]).downsample_ratio
+            current_frame_rate /= self.encoder[-1].downsample_ratio
 
         # Build quantizer
         quantizer_kwargs = dict(config.quantizer_kwargs)
@@ -1375,7 +1368,7 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
                         context=int(current_frame_rate * self.causal_transformer_context_duration),
                     )
                 )
-            current_frame_rate *= cast(MossAudioTokenizerPatchedPretransform, self.decoder[-1]).downsample_ratio
+            current_frame_rate *= self.decoder[-1].downsample_ratio
 
         self.post_init()
 
@@ -1407,11 +1400,14 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
             self._stop_streaming()
 
     @torch.no_grad()
-    def batch_encode(self, wav_list: list[torch.Tensor]) -> MossAudioTokenizerEncoderOutput:
+    def batch_encode(
+        self, wav_list: list[torch.Tensor], num_quantizers: int | None = None
+    ) -> MossAudioTokenizerEncoderOutput:
         """Batch encode a list of audio waveforms.
 
         Args:
             wav_list: List of audio tensors, each of shape `(num_samples,)`.
+            num_quantizers: Number of quantizers to use. By default, all quantizers are used.
 
         Returns:
             [`MossAudioTokenizerEncoderOutput`] with `audio_codes` and `audio_codes_lengths`.
@@ -1430,14 +1426,18 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
             input_values[i, 0, : wav.shape[-1]] = wav
             input_lengths[i] = wav.shape[-1]
 
-        return self._encode_frame(input_values, input_lengths)
+        return self._encode_frame(input_values, input_lengths, n_quantizers=num_quantizers)
 
     @torch.no_grad()
-    def batch_decode(self, codes_list: list[torch.Tensor]) -> MossAudioTokenizerDecoderOutput:
+    def batch_decode(
+        self, codes_list: list[torch.Tensor], num_quantizers: int | None = None
+    ) -> MossAudioTokenizerDecoderOutput:
         """Batch decode a list of audio codes.
 
         Args:
             codes_list: List of audio code tensors, each of shape `(num_quantizers, codes_length)`.
+            num_quantizers: If provided, decode only the first `num_quantizers` quantizers from each element in
+                `codes_list`. If omitted, all elements in `codes_list` must have the same number of quantizers.
 
         Returns:
             [`MossAudioTokenizerDecoderOutput`] with `audio` and `audio_lengths`.
@@ -1447,13 +1447,28 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
 
         batch_size = len(codes_list)
         device = codes_list[0].device
-        num_quantizers = codes_list[0].shape[0]
+        nqs = [codes.shape[0] for codes in codes_list]
+        if num_quantizers is None:
+            num_quantizers = nqs[0]
+            if any(nq != num_quantizers for nq in nqs):
+                raise ValueError(
+                    "All elements in `codes_list` must have the same number of quantizers when `num_quantizers` is None. "
+                    "Pass `num_quantizers=...` to decode a common prefix."
+                )
+        else:
+            min_nq = min(nqs)
+            if min_nq < num_quantizers:
+                raise ValueError(
+                    "`num_quantizers` must be <= the number of quantizers for every element in `codes_list`. "
+                    f"Got num_quantizers={num_quantizers}, min(codes.shape[0])={min_nq}."
+                )
         max_length = max(codes.shape[-1] for codes in codes_list)
 
         audio_codes = torch.zeros(num_quantizers, batch_size, max_length, device=device, dtype=torch.long)
         audio_codes_lengths = torch.zeros(batch_size, device=device, dtype=torch.long)
 
         for i, codes in enumerate(codes_list):
+            codes = codes[:num_quantizers]
             audio_codes[:, i, : codes.shape[-1]] = codes
             audio_codes_lengths[i] = codes.shape[-1]
 
@@ -1638,6 +1653,7 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
         padding_mask: torch.Tensor | None = None,
         return_dict: bool | None = None,
         chunk_duration: float | None = None,
+        num_quantizers: int | None = None,
     ):
         """
         Decodes the given codes into an output audio waveform.
@@ -1653,6 +1669,9 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
                 If provided, decode the input codes in successive chunks of `chunk_duration` seconds while keeping a
                 streaming KV cache for the causal transformers.
 
+            num_quantizers (`int`, *optional*):
+                Number of quantizers to use. By default, all quantizers in `audio_codes` are used.
+
                 `chunk_duration` must be <= `config.causal_transformer_context_duration`, and
                 `chunk_duration * config.sampling_rate` must be divisible by `config.downsample_rate`.
 
@@ -1664,6 +1683,13 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
         if audio_codes.dim() == 2:
             audio_codes = audio_codes.unsqueeze(1)  # nq, T -> nq, B=1, T
 
+        if num_quantizers is not None:
+            if num_quantizers > audio_codes.shape[0]:
+                raise ValueError(
+                    f"`num_quantizers` ({num_quantizers}) must be <= audio_codes.shape[0] ({audio_codes.shape[0]})."
+                )
+            audio_codes = audio_codes[:num_quantizers]
+
         _, B, T = audio_codes.shape
         device = audio_codes.device
 
@@ -1793,7 +1819,12 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
             if decoded_from_encoded_codes and output_audio_codes_lengths is not None:
                 decoder_output = self._decode_frame(audio_codes, output_audio_codes_lengths)
             else:
-                decoder_output = self.decode(audio_codes, padding_mask=padding_mask, return_dict=True)
+                decoder_output = self.decode(
+                    audio_codes,
+                    padding_mask=padding_mask,
+                    return_dict=True,
+                    num_quantizers=num_quantizers,
+                )
                 decoder_output = cast(MossAudioTokenizerDecoderOutput, decoder_output)
             output_audio = decoder_output.audio
             output_audio_lengths = decoder_output.audio_lengths