You are researching how to maximize throughput of the neucodec audio codec encoder (from `neuphonic/neucodec` on HuggingFace).

Context from benchmarking already done:
- neucodec encodes audio at 50 tokens/sec, 16kHz input
- The encode_code() pipeline has these stages:
  1. Feature extraction (CPU, 15.8% time) - uses SeamlessM4T fbank, runs on CPU via numpy
  2. Acoustic encoder CodecEnc (GPU, 35% time) - CNN-based
  3. Semantic model wav2vec2-bert (GPU, 39.5% time) - transformer
  4. Quantize FSQ (GPU, 5.3% time)
- Model is 3.1GB FP32, uses ~4.5GB VRAM per process
- Batching with padding HURT performance (padding waste for variable-length audio)
- torch.compile FAILS (numpy ops in feature extractor)
- FP16/BF16 are SLOWER (LayerNorm needs FP32)
- CUDA MPS with 3 processes gave only 5% gain
- RTF is ~48-50x on A100

The encode_code source code:
```python
def encode_code(self, audio_or_path):
    y = self._prepare_audio(audio_or_path)
    all_semantic_features = []
    for i in range(y.size(0)):
        semantic_features = (
            self.feature_extractor(
                y[i, :].cpu(), sampling_rate=16_000, return_tensors="pt",
            ).input_features.to(self.device)
        )
        all_semantic_features.append(semantic_features)
    semantic_features = torch.vstack(all_semantic_features)
    acoustic_emb = self.CodecEnc(y.to(self.device))
    acoustic_emb = acoustic_emb.transpose(1, 2)
    semantic_output = (
        self.semantic_model(semantic_features[:]).hidden_states[16].transpose(1, 2)
    )
    semantic_encoded = self.SemanticEncoder_module(semantic_output)
    if acoustic_emb.shape[-1] != semantic_encoded.shape[-1]:
        min_len = min(acoustic_emb.shape[-1], semantic_encoded.shape[-1])
        acoustic_emb = acoustic_emb[:, :, :min_len]
        semantic_encoded = semantic_encoded[:, :, :min_len]
    concat_emb = torch.cat([semantic_encoded, acoustic_emb], dim=1)
    concat_emb = self.fc_prior(concat_emb.transpose(1, 2)).transpose(1, 2)
    _, fsq_codes, _ = self.generator(concat_emb, vq=True)
    return fsq_codes
```

RESEARCH TASKS:
1. Can we write a custom batched encode that pads ONLY the fbank features (fixed-size after feature extraction) rather than raw audio? The fbank features might be more uniform in size. Check if SeamlessM4T feature extractor outputs fixed-size features per chunk.
2. Can we bucket audio by similar duration and batch within buckets to minimize padding waste?
3. Can we use ONNX export or TensorRT for the GPU-bound parts (CodecEnc, semantic_model)?
4. Can we replace the numpy fbank with a GPU-based fbank (like torchaudio)?
5. Any way to use torch.compile on just the GPU parts (CodecEnc, semantic_model, generator) separately, skipping the problematic feature_extractor?
6. Can we use torch.jit.trace on the individual submodules?
7. Look at the neucodec package source code at /home/ubuntu/.local/lib/python3.10/site-packages/neucodec/ - read the model.py, and any other relevant files to understand the architecture. Look for optimization opportunities.

For each approach, assess: implementation difficulty, expected speedup, and risk of quality degradation. Be thorough - read the actual source code files. Report findings concisely.


I'll research all these optimization opportunities by examining the actual source code. Let me start by reading the relevant files.


activations.py
alias_free_torch
bs_roformer5.py
codec_decoder_vocos.py
codec_encoder_distill.py
codec_encoder.py
distill_layers.py
__init__.py
model.py
module.py
__pycache__
tconv


Bash

pip show neucodec 2>/dev/null | head -20


python3 -c "from transformers import SeamlessM4TFeatureExtractor; help(SeamlessM4TFeatureExtractor.__call__)" 2>/dev/null | head -80


Name: neucodec
Version: 0.0.5
Summary: A package for NeuCodec, based on xcodec2.
Home-page:
Author: Neuphonic
Author-email: support@neuphonic.com
License: MIT
Location: /home/ubuntu/.local/lib/python3.10/site-packages
Requires: local_attention, numpy, torch, torchao, torchaudio, torchtune, transformers, vector-quantize-pytorch
Required-by:


Help on function __call__ in module transformers.models.seamless_m4t.feature_extraction_seamless_m4t:

__call__(self, raw_speech: Union[numpy.ndarray, list[float], list[numpy.ndarray], list[list[float]]], padding: Union[bool, str, transformers.utils.generic.PaddingStrategy] = True, pad_to_multiple_of: Optional[int] = 2, max_length: Optional[int] = None, truncation: bool = False, return_tensors: Union[str, transformers.utils.generic.TensorType, NoneType] = None, sampling_rate: Optional[int] = None, return_attention_mask: Optional[bool] = None, do_normalize_per_mel_bins: Optional[bool] = True, **kwargs) -> transformers.feature_extraction_utils.BatchFeature
    Main method to featurize and prepare for the model one or several sequence(s).

    Args:
        raw_speech (`np.ndarray`, `torch.Tensor`, `list[float]`, `list[np.ndarray]`, `list[torch.Tensor]`,
        `list[list[float]]`, `list[list[list[float]]]`):
            The sequence or batch of sequences to be padded. Each sequence can be a numpy array,
            a torch tensor, a list of float values, a list of numpy arrays, a list of torch tensors,
            a list of list of float values or a list of a list of list of float values.
            If `raw_speech` is a one-dimensional `np.ndarray`, `torch.Tensor` or a `list[float]`, `raw_speech` is
            considered a single-channel, single-sample sound. In all other cases, the first dimension of
            `raw_speech`, whether from an `np.ndarray`, a `torch.Tensor` or a `list[...]`,
            corresponds to the number of samples in the batch, and the number of channels
            (i.e. mono or stereo character) is derived from the other dimensions
            (1D -> single-channel waveform batches; 2D-> stereo-channel waveform batches).
        padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):
            Select a strategy to pad the returned sequences (according to the model's padding side and padding
            index) among:

            - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
              sequence if provided).
            - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
              acceptable input length for the model if that argument is not provided.
            - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
              lengths).
        pad_to_multiple_of (`int`, *optional*, defaults to 2):
            If set will pad the sequence to a multiple of the provided value.

            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability
            `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128.
        max_length (`int`, *optional*):
            Maximum length of the returned list and optionally padding length (see above).
        truncation (`bool`):
            Activates truncation to cut input sequences longer than *max_length* to *max_length*.
        return_attention_mask (`bool`, *optional*):
            Whether to return the attention mask. If left to the default, will return the attention mask according
            to the specific feature_extractor's default.

            [What are attention masks?](../glossary#attention-mask)

            <Tip>

            For SeamlessM4T models, `attention_mask` should always be passed for batched inference, to avoid subtle
            bugs.

            </Tip>

        return_tensors (`str` or [`~utils.TensorType`], *optional*):
            If set, will return tensors instead of list of python integers. Acceptable values are:

            - `'tf'`: Return TensorFlow `tf.constant` objects.
            - `'pt'`: Return PyTorch `torch.Tensor` objects.
            - `'np'`: Return Numpy `np.ndarray` objects.
        sampling_rate (`int`, *optional*):
            The sampling rate at which the `raw_speech` input was sampled. It is strongly recommended to pass
            `sampling_rate` at the forward call to prevent silent errors.
        do_normalize_per_mel_bins (`bool`, *optional*, defaults to `True`):
            Whether or not to zero-mean unit-variance normalize the input per mel-channel.
        kwargs (*optional*):
            Remaining dictionary of keyword arguments that will be passed to the tokenizer or the feature
            extractor.


1	from typing import Optional, Dict
2	from pathlib import Path
3	import numpy as np
4	import torch
5	import torch.nn as nn
6	import torch.nn.functional as F
7	import torchaudio
8	from torchaudio import transforms as T
9	from huggingface_hub import PyTorchModelHubMixin, ModelHubMixin, hf_hub_download
10	from transformers import AutoFeatureExtractor, HubertModel, Wav2Vec2BertModel
11
12	from .codec_encoder import CodecEncoder
13	from .codec_encoder_distill import DistillCodecEncoder
14	from .codec_decoder_vocos import CodecDecoderVocos
15	from .module import SemanticEncoder
16
17
18	class NeuCodec(
19	    nn.Module,
20	    PyTorchModelHubMixin,
21	    repo_url="https://github.com/neuphonic/neucodec",
22	    license="apache-2.0",
23	):
24
25	    def __init__(self, sample_rate: int, hop_length: int):
26	        super().__init__()
27	        self.sample_rate = sample_rate
28	        self.hop_length = hop_length
29	        self.semantic_model = Wav2Vec2BertModel.from_pretrained(
30	            "facebook/w2v-bert-2.0", output_hidden_states=True
31	        )
32
33	        # (harryjulian): remove layers after #16 as we don't need the features after that
34	        self.semantic_model.encoder.layers[:16]
35	        self.feature_extractor = AutoFeatureExtractor.from_pretrained(
36	            "facebook/w2v-bert-2.0"
37	        )
38	        self.SemanticEncoder_module = SemanticEncoder(1024, 1024, 1024)
39	        self.CodecEnc = CodecEncoder()
40	        self.generator = CodecDecoderVocos(hop_length=hop_length)
41	        self.fc_prior = nn.Linear(2048, 2048)
42	        self.fc_post_a = nn.Linear(2048, 1024)
43
44	    @property
45	    def device(self):
46	        return next(self.parameters()).device
47
48	    @classmethod
49	    def _from_pretrained(
50	        cls,
51	        *,
52	        model_id: str,
53	        revision: Optional[str] = None,
54	        cache_dir: Optional[str] = None,
55	        force_download: bool = False,
56	        proxies: Optional[Dict] = None,
57	        resume_download: bool = False,
58	        local_files_only: bool = False,
59	        token: Optional[str] = None,
60	        map_location: str = "cpu",
61	        strict: bool = True,
62	        **model_kwargs,
63	    ):
64
65	        assert model_id in ["neuphonic/neucodec", "neuphonic/distill-neucodec"]
66	        if model_id == "neuphonic/neucodec":
67	            ignore_keys = ["fc_post_s", "SemanticDecoder"]
68	        elif model_id == "neuphonic/distill-neucodec":
69	            ignore_keys = []
70
71	        # download the model weights file
72	        ckpt_path = hf_hub_download(
73	            repo_id=model_id,
74	            filename="pytorch_model.bin",
75	            revision=revision,
76	            cache_dir=cache_dir,
77	            force_download=force_download,
78	            proxies=proxies,
79	            resume_download=resume_download,
80	            local_files_only=local_files_only,
81	            token=token,
82	        )
83
84	        # download meta.yaml to track number of downloads
85	        _ = hf_hub_download(
86	            repo_id=model_id,
87	            filename="meta.yaml",
88	            revision=revision,
89	            cache_dir=cache_dir,
90	            force_download=force_download,
91	            proxies=proxies,
92	            resume_download=resume_download,
93	            local_files_only=local_files_only,
94	            token=token,
95	        )
96
97	        # initialize model
98	        model = cls(24_000, 480)
99
100	        # load weights
101	        state_dict = torch.load(ckpt_path, map_location)
102	        contains_list = lambda s, l: any(i in s for i in l)
103	        state_dict = {
104	            k:v for k, v in state_dict.items()
105	            if not contains_list(k, ignore_keys)
106	        }
107
108	        # TODO: we can move to strict loading once we clean up the checkpoints
109	        model.load_state_dict(state_dict, strict=False)
110
111	        return model
112
113	    def _prepare_audio(self, audio_or_path: torch.Tensor | Path | str):
114
115	        # load from file
116	        if isinstance(audio_or_path, (Path, str)):
117	            y, sr = torchaudio.load(audio_or_path)
118	            if sr != 16_000:
119	                y, sr = (T.Resample(sr, 16_000)(y), 16_000)
120	                y = y[None, :]  # [1, T] -> [B, 1, T]
121
122	        # ensure input tensor is of correct shape
123	        elif isinstance(audio_or_path, torch.Tensor):
124	            y = audio_or_path
125	            if len(y.shape) == 3:
126	                y = audio_or_path
127	            else:
128	                raise ValueError(
129	                    f"NeuCodec expects tensor audio input to be of shape [B, 1, T] -- received shape: {y.shape}"
130	                )
131
132	        # pad audio
133	        pad_for_wav = 320 - (y.shape[-1] % 320)
134	        y = torch.nn.functional.pad(y, (0, pad_for_wav))
135
136	        return y
137
138	    def encode_code(self, audio_or_path: torch.Tensor | Path | str) -> torch.Tensor:
139	        """
140	        Args:
141	            audio_or_path: torch.Tensor [B, 1, T] | Path | str, input audio
142
143	        Returns:
144	            fsq_codes: torch.Tensor [B, 1, F], 50hz FSQ codes
145	        """
146
147	        # prepare inputs
148	        y = self._prepare_audio(audio_or_path)
149	        all_semantic_features = []
150	        for i in range(y.size(0)):
151	            semantic_features = (
152	                self.feature_extractor(
153	                    y[i, :].cpu(),
154	                    sampling_rate=16_000,
155	                    return_tensors="pt",
156	                )
157	                .input_features.to(self.device)
158	            )
159	            all_semantic_features.append(semantic_features)
160	        semantic_features = torch.vstack(all_semantic_features)
161
162	        # acoustic encoding
163	        acoustic_emb = self.CodecEnc(y.to(self.device))
164	        acoustic_emb = acoustic_emb.transpose(1, 2)
165
166	        # semantic encoding
167	        semantic_output = (
168	            self.semantic_model(semantic_features[:]).hidden_states[16].transpose(1, 2)
169	        )
170	        semantic_encoded = self.SemanticEncoder_module(semantic_output)
171
172	        # concatenate embeddings
173	        if acoustic_emb.shape[-1] != semantic_encoded.shape[-1]:
174	            min_len = min(acoustic_emb.shape[-1], semantic_encoded.shape[-1])
175	            acoustic_emb = acoustic_emb[:, :, :min_len]
176	            semantic_encoded = semantic_encoded[:, :, :min_len]
177	        concat_emb = torch.cat([semantic_encoded, acoustic_emb], dim=1)
178	        concat_emb = self.fc_prior(concat_emb.transpose(1, 2)).transpose(1, 2)
179
180	        # quantize
181	        _, fsq_codes, _ = self.generator(concat_emb, vq=True)
182	        return fsq_codes
183
184	    def decode_code(self, fsq_codes: torch.Tensor) -> torch.Tensor:
185	        """
186	        Args:
187	            fsq_codes: torch.Tensor [B, 1, F], 50hz FSQ codes
188
189	        Returns:
190	            recon: torch.Tensor [B, 1, T], reconstructed 24kHz audio
191	        """
192
193	        fsq_post_emb = self.generator.quantizer.get_output_from_indices(fsq_codes.transpose(1, 2))
194	        fsq_post_emb = fsq_post_emb.transpose(1, 2)
195	        fsq_post_emb = self.fc_post_a(fsq_post_emb.transpose(1, 2)).transpose(1, 2)
196	        recon = self.generator(fsq_post_emb.transpose(1, 2), vq=False)[0]
197	        return recon
198
199
200	class DistillNeuCodec(NeuCodec):
201	    def __init__(self, sample_rate: int, hop_length: int):
202	        nn.Module.__init__(self)
203	        self.sample_rate = sample_rate
204	        self.hop_length = hop_length
205	        self.semantic_model = HubertModel.from_pretrained(
206	            "ntu-spml/distilhubert", output_hidden_states=True
207	        )
208	        self.feature_extractor = AutoFeatureExtractor.from_pretrained(
209	            "ntu-spml/distilhubert"
210	        )
211	        self.SemanticEncoder_module = SemanticEncoder(768, 768, 1024)
212	        self.codec_encoder = DistillCodecEncoder()
213	        self.generator = CodecDecoderVocos(hop_length=hop_length)
214	        self.fc_prior = nn.Linear(
215	            768  # acoustic model
216	            + 768,  # semantic model
217	            2048,
218	        )
219	        self.fc_sq_prior = nn.Linear(512, 768)
220	        self.fc_post_a = nn.Linear(2048, 1024)
221
222	    def encode_code(self, audio_or_path:  torch.Tensor | Path | str) -> torch.Tensor:
223	        """
224	        Args:
225	            audio_or_path: torch.Tensor [B, 1, T] | Path | str, input audio
226
227	        Returns:
228	            fsq_codes: torch.Tensor [B, 1, F], 50hz FSQ codes
229	        """
230
231	        # prepare inputs
232	        y = self._prepare_audio(audio_or_path)
233
234	        all_semantic_features = []
235	        for i in range(y.size(0)):
236	            semantic_features = (
237	                self.feature_extractor(
238	                    F.pad(y[i, :].cpu(), (160, 160)),
239	                    sampling_rate=16_000,
240	                    return_tensors="pt",
241	                )
242	                .input_values.to(self.device)
243	                .squeeze(0)
244	            )
245	            all_semantic_features.append(semantic_features)
246	        semantic_features = torch.vstack(all_semantic_features)
247
248	        # acoustic encoding
249	        fsq_emb = self.fc_sq_prior(self.codec_encoder(y.to(self.device)))
250	        fsq_emb = fsq_emb.transpose(1, 2)
251
252	        # semantic encoding
253	        semantic_target = self.semantic_model(
254	            semantic_features
255	        ).last_hidden_state.transpose(1, 2)
256	        semantic_target = self.SemanticEncoder_module(semantic_target)
257
258	        if fsq_emb.shape[-1] != semantic_target.shape[-1]:
259	            min_len = min(fsq_emb.shape[-1], semantic_target.shape[-1])
260	            fsq_emb = fsq_emb[:, :, :min_len]
261	            semantic_target = semantic_target[:, :, :min_len]
262
263	        concat_emb = torch.cat([semantic_target, fsq_emb], dim=1)
264	        concat_emb = self.fc_prior(concat_emb.transpose(1, 2)).transpose(1, 2)
265	        _, fsq_codes, _ = self.generator(concat_emb, vq=True)
266	        return fsq_codes
267
268
269	class NeuCodecOnnxDecoder(
270	    ModelHubMixin,
271	    repo_url="https://github.com/neuphonic/neucodec",
272	    license="apache-2.0",
273	):
274
275	    def __init__(self, onnx_path):
276
277	        # onnx import
278	        try:
279	            import onnxruntime
280	        except ImportError as e:
281	            raise ImportError("Failed to import `onnxruntime`. Install with the following command: pip install onnxruntime") from e
282
283	        # load model
284	        so = onnxruntime.SessionOptions()
285	        so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
286	        self.session = onnxruntime.InferenceSession(
287	            onnx_path,
288	            sess_options=so
289	        )
290	        self.sample_rate = 24_000
291
292	    @classmethod
293	    def _from_pretrained(
294	        cls,
295	        *,
296	        model_id: str,
297	        revision: Optional[str] = None,
298	        cache_dir: Optional[str] = None,
299	        force_download: bool = False,
300	        proxies: Optional[Dict] = None,
301	        resume_download: bool = False,
302	        local_files_only: bool = False,
303	        token: Optional[str] = None,
304	        map_location: str = "cpu",
305	        strict: bool = True,
306	        **model_kwargs,
307	    ):
308
309	        # download the model weights file
310	        onnx_path = hf_hub_download(
311	            repo_id=model_id,
312	            filename="model.onnx",
313	            revision=revision,
314	            cache_dir=cache_dir,
315	            force_download=force_download,
316	            proxies=proxies,
317	            resume_download=resume_download,
318	            local_files_only=local_files_only,
319	            token=token,
320	        )
321
322	        # download meta.yaml to track number of downloads
323	        _ = hf_hub_download(
324	            repo_id=model_id,
325	            filename="meta.yaml",
326	            revision=revision,
327	            cache_dir=cache_dir,
328	            force_download=force_download,
329	            proxies=proxies,
330	            resume_download=resume_download,
331	            local_files_only=local_files_only,
332	            token=token,
333	        )
334
335	        # initialize model
336	        model = cls(onnx_path) #cls(onnx_path)
337
338	        # only support CPU
339	        if map_location != "cpu":
340	            raise ValueError("The onnx decoder currently only supports CPU runtimes.")
341
342	        return model
343
344	    def encode_code(self, *args, **kwargs):
345	        raise NotImplementedError(
346	            "The onnx decoder has no functionality to encode codes, as it only contains the compiled decoder graph."
347	        )
348
349	    def decode_code(self, codes: np.ndarray) -> np.ndarray:
350	        """
351	        Args:
352	            fsq_codes: np.array [B, 1, F], 50hz FSQ codes
353
354	        Returns:
355	            recon: np.array [B, 1, T], reconstructed 24kHz audio
356	        """
357
358	        # validate inputs
359	        if not isinstance(codes, np.ndarray):
360	            raise ValueError("`Codes` should be an np.array.")
361	        if not len(codes.shape) == 3 or codes.shape[1] != 1:
362	            raise ValueError("`Codes` should be of shape [B, 1, F].")
363
364	        # run decoder
365	        recon = self.session.run(
366	            None, {"codes": codes}
367	        )[0].astype(np.float32)
368
369	        return recon
370


1	import torch
2	import numpy as np
3
4	from torch import nn
5
6	from .module import WNConv1d, EncoderBlock
7	from .alias_free_torch import Activation1d
8	from . import activations
9
10
11	def init_weights(m):
12	    if isinstance(m, nn.Conv1d):
13	        nn.init.trunc_normal_(m.weight, std=0.02)
14	        nn.init.constant_(m.bias, 0)
15
16
17	class CodecEncoder(nn.Module):
18	    def __init__(
19	        self,
20	        ngf=48,
21	        up_ratios=[2, 2, 4, 4, 5],
22	        dilations=(1, 3, 9),
23	        hidden_dim=1024,
24	        depth=12,
25	        heads=12,
26	        pos_meb_dim=64,
27	    ):
28	        super().__init__()
29	        self.hop_length = np.prod(up_ratios)
30	        self.ngf = ngf
31	        self.up_ratios = up_ratios
32
33	        d_model = ngf
34	        self.conv_blocks = [WNConv1d(1, d_model, kernel_size=7, padding=3)]
35
36	        for i, stride in enumerate(up_ratios):
37	            d_model *= 2
38	            self.conv_blocks += [
39	                EncoderBlock(d_model, stride=stride, dilations=dilations)
40	            ]
41
42	        self.conv_blocks = nn.Sequential(*self.conv_blocks)
43
44	        self.conv_final_block = [
45	            Activation1d(
46	                activation=activations.SnakeBeta(d_model, alpha_logscale=True)
47	            ),
48	            WNConv1d(d_model, hidden_dim, kernel_size=3, padding=1),
49	        ]
50	        self.conv_final_block = nn.Sequential(*self.conv_final_block)
51
52	        self.reset_parameters()
53
54	    def forward(self, x):
55	        x = self.conv_blocks(x)
56	        x = self.conv_final_block(x)
57	        x = x.permute(0, 2, 1)
58	        return x
59
60	    def inference(self, x):
61	        return self.block(x)
62
63	    def remove_weight_norm(self):
64	        """Remove weight normalization module from all of the layers."""
65
66	        def _remove_weight_norm(m):
67	            try:
68	                torch.nn.utils.remove_weight_norm(m)
69	            except ValueError:  # this module didn't have weight norm
70	                return
71
72	        self.apply(_remove_weight_norm)
73
74	    def apply_weight_norm(self):
75	        """Apply weight normalization module from all of the layers."""
76
77	        def _apply_weight_norm(m):
78	            if isinstance(m, nn.Conv1d):
79	                torch.nn.utils.weight_norm(m)
80
81	        self.apply(_apply_weight_norm)
82
83	    def reset_parameters(self):
84	        self.apply(init_weights)
85


1	from .model import NeuCodec, DistillNeuCodec, NeuCodecOnnxDecoder
2


1	import torch.nn as nn
2
3	from torch.nn.utils import weight_norm
4
5	from .activations import SnakeBeta
6	from .alias_free_torch import Activation1d
7
8
9	def WNConv1d(*args, **kwargs):
10	    return weight_norm(nn.Conv1d(*args, **kwargs))
11
12
13	class ResidualUnit(nn.Module):
14	    def __init__(self, dim: int = 16, dilation: int = 1):
15	        super().__init__()
16	        pad = ((7 - 1) * dilation) // 2
17	        self.block = nn.Sequential(
18	            Activation1d(activation=SnakeBeta(dim, alpha_logscale=True)),
19	            WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
20	            Activation1d(activation=SnakeBeta(dim, alpha_logscale=True)),
21	            WNConv1d(dim, dim, kernel_size=1),
22	        )
23
24	    def forward(self, x):
25	        return x + self.block(x)
26
27
28	class EncoderBlock(nn.Module):
29	    def __init__(self, dim: int = 16, stride: int = 1, dilations=(1, 3, 9)):
30	        super().__init__()
31	        runits = [ResidualUnit(dim // 2, dilation=d) for d in dilations]
32	        self.block = nn.Sequential(
33	            *runits,
34	            Activation1d(activation=SnakeBeta(dim // 2, alpha_logscale=True)),
35	            WNConv1d(
36	                dim // 2,
37	                dim,
38	                kernel_size=2 * stride,
39	                stride=stride,
40	                padding=stride // 2 + stride % 2,
41	            ),
42	        )
43
44	    def forward(self, x):
45	        return self.block(x)
46
47
48	class SemanticEncoder(nn.Module):
49	    def __init__(
50	        self,
51	        input_channels: int,
52	        code_dim: int,
53	        encode_channels: int,
54	        kernel_size: int = 3,
55	        bias: bool = True,
56	    ):
57	        super(SemanticEncoder, self).__init__()
58
59	        self.initial_conv = nn.Conv1d(
60	            in_channels=input_channels,
61	            out_channels=encode_channels,
62	            kernel_size=kernel_size,
63	            stride=1,
64	            padding=(kernel_size - 1) // 2,
65	            bias=False,
66	        )
67
68	        self.residual_blocks = nn.Sequential(
69	            nn.ReLU(inplace=True),
70	            nn.Conv1d(
71	                encode_channels,
72	                encode_channels,
73	                kernel_size=kernel_size,
74	                stride=1,
75	                padding=(kernel_size - 1) // 2,
76	                bias=bias,
77	            ),
78	            nn.ReLU(inplace=True),
79	            nn.Conv1d(
80	                encode_channels,
81	                encode_channels,
82	                kernel_size=kernel_size,
83	                stride=1,
84	                padding=(kernel_size - 1) // 2,
85	                bias=bias,
86	            ),
87	        )
88
89	        self.final_conv = nn.Conv1d(
90	            in_channels=encode_channels,
91	            out_channels=code_dim,
92	            kernel_size=kernel_size,
93	            stride=1,
94	            padding=(kernel_size - 1) // 2,
95	            bias=False,
96	        )
97
98	    def forward(self, x):
99	        x = self.initial_conv(x)  # (Batch, Encode_channels, Length)
100	        x = self.residual_blocks(x) + x  # 残差连接
101	        x = self.final_conv(x)  # (Batch, Code_dim, Length)
102	        return x
103


1	import torch
2	import torch.nn as nn
3
4	from typing import List
5	from torchtune.modules import RotaryPositionalEmbeddings
6	from vector_quantize_pytorch import ResidualFSQ
7
8	from .bs_roformer5 import TransformerBlock
9
10
11	class ISTFT(nn.Module):
12	    """
13	    Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with
14	    windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges.
15	    See issue: https://github.com/pytorch/pytorch/issues/62323
16	    Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs.
17	    The NOLA constraint is met as we trim padded samples anyway.
18
19	    Args:
20	        n_fft (int): Size of Fourier transform.
21	        hop_length (int): The distance between neighboring sliding window frames.
22	        win_length (int): The size of window frame and STFT filter.
23	        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
24	    """
25
26	    def __init__(
27	        self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"
28	    ):
29	        super().__init__()
30	        if padding not in ["center", "same"]:
31	            raise ValueError("Padding must be 'center' or 'same'.")
32	        self.padding = padding
33	        self.n_fft = n_fft
34	        self.hop_length = hop_length
35	        self.win_length = win_length
36	        window = torch.hann_window(win_length)
37	        self.register_buffer("window", window)
38
39	    def forward(self, spec: torch.Tensor) -> torch.Tensor:
40	        """
41	        Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram.
42
43	        Args:
44	            spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size,
45	                            N is the number of frequency bins, and T is the number of time frames.
46
47	        Returns:
48	            Tensor: Reconstructed time-domain signal of shape (B, L), where L is the length of the output signal.
49	        """
50	        if self.padding == "center":
51	            # Fallback to pytorch native implementation
52	            return torch.istft(
53	                spec,
54	                self.n_fft,
55	                self.hop_length,
56	                self.win_length,
57	                self.window,
58	                center=True,
59	            )
60	        elif self.padding == "same":
61	            pad = (self.win_length - self.hop_length) // 2
62	        else:
63	            raise ValueError("Padding must be 'center' or 'same'.")
64
65	        assert spec.dim() == 3, "Expected a 3D tensor as input"
66	        B, N, T = spec.shape
67
68	        # Inverse FFT
69	        ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
70	        ifft = ifft * self.window[None, :, None]
71
72	        # Overlap and Add
73	        output_size = (T - 1) * self.hop_length + self.win_length
74	        y = torch.nn.functional.fold(
75	            ifft,
76	            output_size=(1, output_size),
77	            kernel_size=(1, self.win_length),
78	            stride=(1, self.hop_length),
79	        )[:, 0, 0, pad:-pad]
80
81	        # Window envelope
82	        window_sq = self.window.square().expand(1, T, -1).transpose(1, 2)
83	        window_envelope = torch.nn.functional.fold(
84	            window_sq,
85	            output_size=(1, output_size),
86	            kernel_size=(1, self.win_length),
87	            stride=(1, self.hop_length),
88	        ).squeeze()[pad:-pad]
89
90	        # Normalize
91	        assert (window_envelope > 1e-11).all()
92	        y = y / window_envelope
93
94	        return y
95
96
97	class FourierHead(nn.Module):
98	    """Base class for inverse fourier modules."""
99
100	    def forward(self, x: torch.Tensor) -> torch.Tensor:
101	        """
102	        Args:
103	            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
104	                        L is the sequence length, and H denotes the model dimension.
105
106	        Returns:
107	            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
108	        """
109	        raise NotImplementedError("Subclasses must implement the forward method.")
110
111
112	class ISTFTHead(FourierHead):
113	    """
114	    ISTFT Head module for predicting STFT complex coefficients.
115
116	    Args:
117	        dim (int): Hidden dimension of the model.
118	        n_fft (int): Size of Fourier transform.
119	        hop_length (int): The distance between neighboring sliding window frames, which should align with
120	                          the resolution of the input features.
121	        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
122	    """
123
124	    def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"):
125	        super().__init__()
126	        out_dim = n_fft + 2
127	        self.out = torch.nn.Linear(dim, out_dim)
128	        self.istft = ISTFT(
129	            n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding
130	        )
131
132	    def forward(self, x: torch.Tensor) -> torch.Tensor:
133	        """
134	        Forward pass of the ISTFTHead module.
135
136	        Args:
137	            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
138	                        L is the sequence length, and H denotes the model dimension.
139
140	        Returns:
141	            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
142	        """
143	        x_pred = self.out(x)
144	        # x_pred = x
145	        x_pred = x_pred.transpose(1, 2)
146	        mag, p = x_pred.chunk(2, dim=1)
147	        mag = torch.exp(mag)
148	        mag = torch.clip(
149	            mag, max=1e2
150	        )  # safeguard to prevent excessively large magnitudes
151	        # wrapping happens here. These two lines produce real and imaginary value
152	        x = torch.cos(p)
153	        y = torch.sin(p)
154	        # recalculating phase here does not produce anything new
155	        # only costs time
156	        # phase = torch.atan2(y, x)
157	        # S = mag * torch.exp(phase * 1j)
158	        # better directly produce the complex value
159	        S = mag * (x + 1j * y)
160	        audio = self.istft(S)
161	        return audio.unsqueeze(1), x_pred
162
163
164	def nonlinearity(x):
165	    # swish
166	    return x * torch.sigmoid(x)
167
168
169	def Normalize(in_channels, num_groups=32):
170	    return torch.nn.GroupNorm(
171	        num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True
172	    )
173
174
175	class ResnetBlock(nn.Module):
176	    def __init__(
177	        self,
178	        *,
179	        in_channels,
180	        out_channels=None,
181	        conv_shortcut=False,
182	        dropout,
183	        temb_channels=512,
184	    ):
185	        super().__init__()
186	        self.in_channels = in_channels
187	        out_channels = in_channels if out_channels is None else out_channels
188	        self.out_channels = out_channels
189	        self.use_conv_shortcut = conv_shortcut
190
191	        self.norm1 = Normalize(in_channels)
192	        self.conv1 = torch.nn.Conv1d(
193	            in_channels, out_channels, kernel_size=3, stride=1, padding=1
194	        )
195	        if temb_channels > 0:
196	            self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
197	        self.norm2 = Normalize(out_channels)
198	        self.dropout = torch.nn.Dropout(dropout)
199	        self.conv2 = torch.nn.Conv1d(
200	            out_channels, out_channels, kernel_size=3, stride=1, padding=1
201	        )
202	        if self.in_channels != self.out_channels:
203	            if self.use_conv_shortcut:
204	                self.conv_shortcut = torch.nn.Conv1d(
205	                    in_channels, out_channels, kernel_size=3, stride=1, padding=1
206	                )
207	            else:
208	                self.nin_shortcut = torch.nn.Conv1d(
209	                    in_channels, out_channels, kernel_size=1, stride=1, padding=0
210	                )
211
212	    def forward(self, x, temb=None):
213	        h = x
214	        h = self.norm1(h)
215	        h = nonlinearity(h)
216	        h = self.conv1(h)
217
218	        if temb is not None:
219	            h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
220
221	        h = self.norm2(h)
222	        h = nonlinearity(h)
223	        h = self.dropout(h)
224	        h = self.conv2(h)
225
226	        if self.in_channels != self.out_channels:
227	            if self.use_conv_shortcut:
228	                x = self.conv_shortcut(x)
229	            else:
230	                x = self.nin_shortcut(x)
231
232	        return x + h
233
234
235	class Backbone(nn.Module):
236	    """Base class for the generator's backbone. It preserves the same temporal resolution across all layers."""
237
238	    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
239	        """
240	        Args:
241	            x (Tensor): Input tensor of shape (B, C, L), where B is the batch size,
242	                        C denotes output features, and L is the sequence length.
243
244	        Returns:
245	            Tensor: Output of shape (B, L, H), where B is the batch size, L is the sequence length,
246	                    and H denotes the model dimension.
247	        """
248	        raise NotImplementedError("Subclasses must implement the forward method.")
249
250
251	class VocosBackbone(Backbone):
252	    """
253	    Vocos backbone module built with ConvNeXt blocks. Supports additional conditioning with Adaptive Layer Normalization
254
255	    Args:
256	        input_channels (int): Number of input features channels.
257	        dim (int): Hidden dimension of the model.
258	        intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock.
259	        num_layers (int): Number of ConvNeXtBlock layers.
260	        layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to `1 / num_layers`.
261	        adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm.
262	                                                None means non-conditional model. Defaults to None.
263	    """
264
265	    def __init__(self, hidden_dim=1024, depth=12, heads=16, pos_meb_dim=64):
266	        super().__init__()
267
268	        self.embed = nn.Conv1d(hidden_dim, hidden_dim, kernel_size=7, padding=3)
269
270	        self.temb_ch = 0
271	        block_in = hidden_dim
272	        dropout = 0.1
273
274	        prior_net: List[nn.Module] = [
275	            ResnetBlock(
276	                in_channels=block_in,
277	                out_channels=block_in,
278	                temb_channels=self.temb_ch,
279	                dropout=dropout,
280	            ),
281	            ResnetBlock(
282	                in_channels=block_in,
283	                out_channels=block_in,
284	                temb_channels=self.temb_ch,
285	                dropout=dropout,
286	            ),
287	        ]
288	        self.prior_net = nn.Sequential(*prior_net)
289
290	        depth = depth
291	        time_rotary_embed = RotaryPositionalEmbeddings(dim=pos_meb_dim)
292
293	        transformer_blocks = [
294	            TransformerBlock(
295	                dim=hidden_dim, n_heads=heads, rotary_embed=time_rotary_embed
296	            )
297	            for _ in range(depth)
298	        ]
299
300	        self.transformers = nn.Sequential(*transformer_blocks)
301	        self.final_layer_norm = nn.LayerNorm(hidden_dim, eps=1e-6)
302	        post_net: List[nn.Module] = [
303	            ResnetBlock(
304	                in_channels=block_in,
305	                out_channels=block_in,
306	                temb_channels=self.temb_ch,
307	                dropout=dropout,
308	            ),
309	            ResnetBlock(
310	                in_channels=block_in,
311	                out_channels=block_in,
312	                temb_channels=self.temb_ch,
313	                dropout=dropout,
314	            ),
315	        ]
316	        self.post_net = nn.Sequential(*post_net)
317
318	    def forward(self, x: torch.Tensor) -> torch.Tensor:
319	        x = x.transpose(1, 2)
320	        x = self.embed(x)
321	        x = self.prior_net(x)
322	        x = x.transpose(1, 2)
323	        x = self.transformers(x)
324	        x = x.transpose(1, 2)
325	        x = self.post_net(x)
326	        x = x.transpose(1, 2)
327	        x = self.final_layer_norm(x)
328	        return x
329
330
331	def init_weights(m):
332	    if isinstance(m, nn.Conv1d):
333	        nn.init.trunc_normal_(m.weight, std=0.02)
334	        nn.init.constant_(m.bias, 0)
335
336
337	class CodecDecoderVocos(nn.Module):
338	    def __init__(
339	        self,
340	        hidden_dim=1024,
341	        depth=12,
342	        heads=16,
343	        pos_meb_dim=64,
344	        hop_length=320,
345	        vq_num_quantizers=1,
346	        vq_dim=2048,  # 1024 2048
347	        vq_commit_weight=0.25,
348	        vq_weight_init=False,
349	        vq_full_commit_loss=False,
350	        codebook_size=16384,
351	        codebook_dim=16,
352	    ):
353	        super().__init__()
354	        self.hop_length = hop_length
355
356	        self.quantizer = ResidualFSQ(
357	            dim=vq_dim, levels=[4, 4, 4, 4, 4, 4, 4, 4], num_quantizers=1
358	        )
359
360	        self.backbone = VocosBackbone(
361	            hidden_dim=hidden_dim, depth=depth, heads=heads, pos_meb_dim=pos_meb_dim
362	        )
363
364	        self.head = ISTFTHead(
365	            dim=hidden_dim,
366	            n_fft=self.hop_length * 4,
367	            hop_length=self.hop_length,
368	            padding="same",
369	        )
370
371	        self.reset_parameters()
372
373	    def forward(self, x, vq=True):
374	        if vq is True:
375	            # x, q, commit_loss = self.quantizer(x)
376	            x = x.permute(0, 2, 1)
377	            x, q = self.quantizer(x)
378	            x = x.permute(0, 2, 1)
379	            q = q.permute(0, 2, 1)
380	            return x, q, None
381	        x = self.backbone(x)
382	        x, _ = self.head(x)
383
384	        return x, _
385
386	    def vq2emb(self, vq):
387	        self.quantizer = self.quantizer.eval()
388	        x = self.quantizer.vq2emb(vq)
389	        return x
390
391	    def get_emb(self):
392	        self.quantizer = self.quantizer.eval()
393	        embs = self.quantizer.get_emb()
394	        return embs
395
396	    def inference_vq(self, vq):
397	        x = vq[None, :, :]
398	        x = self.model(x)
399	        return x
400
401	    def inference_0(self, x):
402	        x, q, loss, perp = self.quantizer(x)
403	        x = self.model(x)
404	        return x, None
405
406	    def inference(self, x):
407	        x = self.model(x)
408	        return x, None
409
410	    def remove_weight_norm(self):
411	        """Remove weight normalization module from all of the layers."""
412
413	        def _remove_weight_norm(m):
414	            try:
415	                torch.nn.utils.remove_weight_norm(m)
416	            except ValueError:  # this module didn't have weight norm
417	                return
418
419	        self.apply(_remove_weight_norm)
420
421	    def apply_weight_norm(self):
422	        """Apply weight normalization module from all of the layers."""
423
424	        def _apply_weight_norm(m):
425	            if isinstance(m, nn.Conv1d) or isinstance(m, nn.ConvTranspose1d):
426	                torch.nn.utils.weight_norm(m)
427
428	        self.apply(_apply_weight_norm)
429
430	    def reset_parameters(self):
431	        self.apply(init_weights)
432


Now let me check the SeamlessM4T feature extractor internals and the semantic model structure.


python3 -c "
from transformers import AutoFeatureExtractor
fe = AutoFeatureExtractor.from_pretrained('facebook/w2v-bert-2.0')
print('Type:', type(fe))
print('Config:', fe.to_dict())
" 2>/dev/null


python3 -c "
import torch, numpy as np
from transformers import AutoFeatureExtractor
fe = AutoFeatureExtractor.from_pretrained('facebook/w2v-bert-2.0')
# Test different audio lengths to see output shapes
for dur_sec in [0.5, 1.0, 2.0, 5.0, 10.0]:
    samples = int(dur_sec * 16000)
    audio = np.random.randn(1, samples).astype(np.float32)
    out = fe(audio, sampling_rate=16000, return_tensors='pt')
    print(f'Input: {dur_sec}s ({samples} samples) -> fbank shape: {out.input_features.shape}')
" 2>/dev/null


python3 -c "
from transformers import Wav2Vec2BertModel
m = Wav2Vec2BertModel.from_pretrained('facebook/w2v-bert-2.0', output_hidden_states=True)
print('Number of encoder layers:', len(m.encoder.layers))
print('Model config adapter_kernel_size:', m.config.adapter_kernel_size if hasattr(m.config, 'adapter_kernel_size') else 'N/A')
print('Model config adapter_stride:', m.config.adapter_stride if hasattr(m.config, 'adapter_stride') else 'N/A')
# Check if model has attention_mask support
import inspect
sig = inspect.signature(m.forward)
print('Forward params:', list(sig.parameters.keys()))
" 2>/dev/null


Type: <class 'transformers.models.seamless_m4t.feature_extraction_seamless_m4t.SeamlessM4TFeatureExtractor'>
Config: {'num_mel_bins': 80, 'return_attention_mask': True, 'stride': 2, 'feature_size': 80, 'sampling_rate': 16000, 'padding_value': 1, 'padding_side': 'right', '_processor_class': 'Wav2Vec2BertProcessor', 'feature_extractor_type': 'SeamlessM4TFeatureExtractor'}


Input: 0.5s (8000 samples) -> fbank shape: torch.Size([1, 24, 160])
Input: 1.0s (16000 samples) -> fbank shape: torch.Size([1, 49, 160])
Input: 2.0s (32000 samples) -> fbank shape: torch.Size([1, 99, 160])
Input: 5.0s (80000 samples) -> fbank shape: torch.Size([1, 249, 160])
Input: 10.0s (160000 samples) -> fbank shape: torch.Size([1, 499, 160])


Number of encoder layers: 24
Model config adapter_kernel_size: 3
Model config adapter_stride: 2
Forward params: ['input_features', 'attention_mask', 'mask_time_indices', 'output_attentions', 'output_hidden_states', 'return_dict']


python3 -c "
# Check the fbank feature extractor source - is it pure numpy?
import inspect
from transformers.models.seamless_m4t.feature_extraction_seamless_m4t import SeamlessM4TFeatureExtractor
src = inspect.getsource(SeamlessM4TFeatureExtractor._extract_fbank_features)
print(src[:3000])
" 2>/dev/null


python3 -c "
# Check if torchaudio has a compatible fbank
import torchaudio
print(dir(torchaudio.compliance))
print(dir(torchaudio.transforms))
# Check kaldi fbank
import inspect
src = inspect.getsource(torchaudio.compliance.kaldi.fbank)
print(src[:2000])
" 2>/dev/null


    def _extract_fbank_features(
        self,
        waveform: np.ndarray,
    ) -> np.ndarray:
        """
        Get mel-filter bank features using TorchAudio. Note that TorchAudio requires 16-bit signed integers as inputs
        and hence the waveform should not be normalized before feature extraction.
        """
        # by default, it extracts the left channel if stereo
        if len(waveform.shape) == 2:
            waveform = waveform[0]

        waveform = np.squeeze(waveform) * (2**15)  # Kaldi compliance: 16-bit signed integers
        features = spectrogram(
            waveform,
            self.window,
            frame_length=400,
            hop_length=160,
            fft_length=512,
            power=2.0,
            center=False,
            preemphasis=0.97,
            mel_filters=self.mel_filters,
            log_mel="log",
            mel_floor=1.192092955078125e-07,
            remove_dc_offset=True,
        ).T
        return features


['__all__', '__builtins__', '__cached__', '__doc__', '__file__', '__loader__', '__name__', '__package__', '__path__', '__spec__', 'kaldi']
['AddNoise', 'AmplitudeToDB', 'ComputeDeltas', 'Convolve', 'Deemphasis', 'FFTConvolve', 'Fade', 'FrequencyMasking', 'GriffinLim', 'InverseMelScale', 'InverseSpectrogram', 'LFCC', 'Loudness', 'MFCC', 'MVDR', 'MelScale', 'MelSpectrogram', 'MuLawDecoding', 'MuLawEncoding', 'PSD', 'PitchShift', 'Preemphasis', 'RNNTLoss', 'RTFMVDR', 'Resample', 'SlidingWindowCmn', 'SoudenMVDR', 'SpecAugment', 'SpectralCentroid', 'Spectrogram', 'Speed', 'SpeedPerturbation', 'TimeMasking', 'TimeStretch', 'Vad', 'Vol', '_RNNTLoss', '__all__', '__builtins__', '__cached__', '__doc__', '__file__', '__loader__', '__name__', '__package__', '__path__', '__spec__', '_multi_channel', '_transforms', 'dropping_class_support']
def fbank(
    waveform: Tensor,
    blackman_coeff: float = 0.42,
    channel: int = -1,
    dither: float = 0.0,
    energy_floor: float = 1.0,
    frame_length: float = 25.0,
    frame_shift: float = 10.0,
    high_freq: float = 0.0,
    htk_compat: bool = False,
    low_freq: float = 20.0,
    min_duration: float = 0.0,
    num_mel_bins: int = 23,
    preemphasis_coefficient: float = 0.97,
    raw_energy: bool = True,
    remove_dc_offset: bool = True,
    round_to_power_of_two: bool = True,
    sample_frequency: float = 16000.0,
    snip_edges: bool = True,
    subtract_mean: bool = False,
    use_energy: bool = False,
    use_log_fbank: bool = True,
    use_power: bool = True,
    vtln_high: float = -500.0,
    vtln_low: float = 100.0,
    vtln_warp: float = 1.0,
    window_type: str = POVEY,
) -> Tensor:
    r"""Create a fbank from a raw audio signal. This matches the input/output of Kaldi's
    compute-fbank-feats.

    Args:
        waveform (Tensor): Tensor of audio of size (c, n) where c is in the range [0,2)
        blackman_coeff (float, optional): Constant coefficient for generalized Blackman window. (Default: ``0.42``)
        channel (int, optional): Channel to extract (-1 -> expect mono, 0 -> left, 1 -> right) (Default: ``-1``)
        dither (float, optional): Dithering constant (0.0 means no dither). If you turn this off, you should set
            the energy_floor option, e.g. to 1.0 or 0.1 (Default: ``0.0``)
        energy_floor (float, optional): Floor on energy (absolute, not relative) in Spectrogram computation.  Caution:
            this floor is applied to the zeroth component, representing the total signal energy.  The floor on the
            individual spectrogram elements is fixed at std::numeric_limits<float>::epsilon(). (Default: ``1.0``)
        frame_length (float, optional): Frame length in milliseconds (Default: ``25.0``)
        frame_shift (float, optional): Frame shift in milliseconds (Default: ``10.0``)
        high_freq (floa


1	import torch
2	import torch.nn as nn
3	from einops import rearrange
4	from torchtune.modules import RotaryPositionalEmbeddings
5
6
7	class RMSNorm(torch.nn.Module):
8	    def __init__(self, dim: int, eps: float = 1e-6):
9	        r"""https://github.com/meta-llama/llama/blob/main/llama/model.py"""
10	        super().__init__()
11	        self.eps = eps
12	        self.weight = nn.Parameter(torch.ones(dim))
13
14	    def forward(self, x):
15	        norm_x = torch.mean(x**2, dim=-1, keepdim=True)
16	        output = x * torch.rsqrt(norm_x + self.eps) * self.weight
17	        return output
18
19
20	class MLP(nn.Module):
21	    def __init__(self, dim: int) -> None:
22	        super().__init__()
23
24	        self.fc1 = nn.Linear(dim, 4 * dim, bias=False)
25	        self.silu = nn.SiLU()
26	        self.fc2 = nn.Linear(4 * dim, dim, bias=False)
27
28	    def forward(self, x):
29	        x = self.fc1(x)
30	        x = self.silu(x)
31	        x = self.fc2(x)
32	        return x
33
34
35	class Attention(nn.Module):
36	    def __init__(
37	        self, dim: int, n_heads: int, rotary_embed: RotaryPositionalEmbeddings
38	    ):
39	        super().__init__()
40
41	        assert dim % n_heads == 0
42
43	        self.n_heads = n_heads
44	        self.dim = dim
45	        self.rotary_embed = rotary_embed
46
47	        self.flash = hasattr(torch.nn.functional, "scaled_dot_product_attention")
48	        assert self.flash, "Must have flash attention."
49
50	        self.c_attn = nn.Linear(dim, 3 * dim, bias=False)
51	        self.c_proj = nn.Linear(dim, dim, bias=False)
52
53	    def forward(self, x):
54	        r"""
55	        Args:
56	            x: (b, t, h*d)
57
58	        Constants:
59	            b: batch_size
60	            t: time steps
61	            r: 3
62	            h: heads_num
63	            d: heads_dim
64	        """
65	        B, T, C = x.size()
66
67	        q, k, v = rearrange(
68	            self.c_attn(x), "b t (r h d) -> r b h t d", r=3, h=self.n_heads
69	        )
70	        # q, k, v: (b, h, t, d)
71
72	        q = self.rotary_embed(q)
73	        k = self.rotary_embed(k)
74
75	        if self.flash:
76	            y = torch.nn.functional.scaled_dot_product_attention(
77	                q, k, v, attn_mask=None, dropout_p=0, is_causal=False
78	            )
79
80	        y = rearrange(y, "b h t d -> b t (h d)")
81
82	        y = self.c_proj(y)
83	        # shape: (b, t, h*d)
84
85	        return y
86
87
88	class TransformerBlock(nn.Module):
89	    def __init__(
90	        self, dim: int, n_heads: int, rotary_embed: RotaryPositionalEmbeddings
91	    ):
92	        super().__init__()
93	        self.dim = dim
94	        self.n_heads = n_heads
95
96	        self.att_norm = RMSNorm(dim)
97	        self.ffn_norm = RMSNorm(dim)
98	        self.att = Attention(dim=dim, n_heads=n_heads, rotary_embed=rotary_embed)
99	        self.mlp = MLP(dim=dim)
100
101	    def forward(
102	        self,
103	        x: torch.Tensor,
104	    ):
105	        x = x + self.att(self.att_norm(x))
106	        x = x + self.mlp(self.ffn_norm(x))
107	        return x
108
109
110	if __name__ == "__main__":
111	    rotary_embed_128 = RotaryPositionalEmbeddings(dim=128)
112	    transformer_block = TransformerBlock(
113	        dim=1024, n_heads=8, rotary_embed=rotary_embed_128
114	    )
115	    x = torch.randn(2, 128, 1024)
116	    y = transformer_block(x)
117	    print(y.shape)
118	    c = 1
119


python3 -c "
# Check what 'spectrogram' function the feature extractor uses - it's from transformers audio_utils
from transformers.audio_utils import spectrogram
import inspect
src = inspect.getsource(spectrogram)
print(src[:2000])
" 2>/dev/null


def spectrogram(
    waveform: np.ndarray,
    window: np.ndarray,
    frame_length: int,
    hop_length: int,
    fft_length: Optional[int] = None,
    power: Optional[float] = 1.0,
    center: bool = True,
    pad_mode: str = "reflect",
    onesided: bool = True,
    dither: float = 0.0,
    preemphasis: Optional[float] = None,
    mel_filters: Optional[np.ndarray] = None,
    mel_floor: float = 1e-10,
    log_mel: Optional[str] = None,
    reference: float = 1.0,
    min_value: float = 1e-10,
    db_range: Optional[float] = None,
    remove_dc_offset: bool = False,
    dtype: np.dtype = np.float32,
) -> np.ndarray:
    """
    Calculates a spectrogram over one waveform using the Short-Time Fourier Transform.

    This function can create the following kinds of spectrograms:

      - amplitude spectrogram (`power = 1.0`)
      - power spectrogram (`power = 2.0`)
      - complex-valued spectrogram (`power = None`)
      - log spectrogram (use `log_mel` argument)
      - mel spectrogram (provide `mel_filters`)
      - log-mel spectrogram (provide `mel_filters` and `log_mel`)

    How this works:

      1. The input waveform is split into frames of size `frame_length` that are partially overlapping by `frame_length
         - hop_length` samples.
      2. Each frame is multiplied by the window and placed into a buffer of size `fft_length`.
      3. The DFT is taken of each windowed frame.
      4. The results are stacked into a spectrogram.

    We make a distinction between the following "blocks" of sample data, each of which may have a different lengths:

      - The analysis frame. This is the size of the time slices that the input waveform is split into.
      - The window. Each analysis frame is multiplied by the window to avoid spectral leakage.
      - The FFT input buffer. The length of this determines how many frequency bins are in the spectrogram.

    In this implementation, the window is assumed to be zero-padded to have the same size as the analysis fram


python3 -c "
# Check if torch.compile works on individual submodules
import torch
print('torch version:', torch.__version__)
print('has compile:', hasattr(torch, 'compile'))

# Check if we can truncate wav2vec2bert encoder layers
from transformers import Wav2Vec2BertModel
m = Wav2Vec2BertModel.from_pretrained('facebook/w2v-bert-2.0', output_hidden_states=True)
print('Original layers:', len(m.encoder.layers))
# Line 34 in model.py just evaluates self.semantic_model.encoder.layers[:16] without assignment!
# This is a BUG - layers are NOT actually truncated
print('This is a no-op in the source code!')
" 2>/dev/null


python3 -c "
# Check the CodecEncoder hop_length / downsampling factor
import numpy as np
up_ratios = [2, 2, 4, 4, 5]
total_stride = np.prod(up_ratios)
print(f'CodecEnc total stride: {total_stride}')
# For 16kHz input, output rate = 16000/320 = 50 tokens/sec
print(f'Output rate at 16kHz: {16000/total_stride} tokens/sec')

# For wav2vec2-bert with adapter_stride=2, adapter_kernel_size=3
# Input fbank is at 100Hz (160 hop / 16000 sr = 10ms), then adapter downsamples by 2 -> 50Hz
print('wav2vec2-bert fbank rate: 100Hz -> adapter -> 50Hz')
" 2>/dev/null


torch version: 2.9.1+cu128
has compile: True
Original layers: 24
This is a no-op in the source code!


CodecEnc total stride: 320
Output rate at 16kHz: 50.0 tokens/sec
wav2vec2-bert fbank rate: 100Hz -> adapter -> 50Hz


python3 -c "
# Test torch.compile on CodecEncoder alone
import torch
from neucodec.codec_encoder import CodecEncoder

enc = CodecEncoder().cuda().eval()
x = torch.randn(1, 1, 16000).cuda()

# Test basic forward
with torch.no_grad():
    out = enc(x)
    print('CodecEnc output shape:', out.shape)

# Test torch.compile
compiled_enc = torch.compile(enc, mode='reduce-overhead')
with torch.no_grad():
    out2 = compiled_enc(x)
    print('Compiled CodecEnc output shape:', out2.shape)
    print('torch.compile works on CodecEnc!')
" 2>/dev/null


python3 -c "
# Test torch.compile on Wav2Vec2BertModel
import torch
from transformers import Wav2Vec2BertModel

model = Wav2Vec2BertModel.from_pretrained('facebook/w2v-bert-2.0', output_hidden_states=True).cuda().eval()
x = torch.randn(1, 49, 160).cuda()  # fbank features for 1s audio

with torch.no_grad():
    out = model(x)
    print('Semantic model output hidden state 16 shape:', out.hidden_states[16].shape)

# Try torch.compile
compiled = torch.compile(model, mode='reduce-overhead')
with torch.no_grad():
    out2 = compiled(x)
    print('Compiled semantic model works! Shape:', out2.hidden_states[16].shape)
" 2>/dev/null


python3 -c "
# Test torch.jit.trace on CodecEncoder
import torch
from neucodec.codec_encoder import CodecEncoder

enc = CodecEncoder().cuda().eval()
x = torch.randn(1, 1, 16000).cuda()

with torch.no_grad():
    traced = torch.jit.trace(enc, x)
    out = traced(x)
    print('JIT trace on CodecEnc works! Shape:', out.shape)
" 2>/dev/null


CodecEnc output shape: torch.Size([1, 50, 1024])
Compiled CodecEnc output shape: torch.Size([1, 50, 1024])
torch.compile works on CodecEnc!


Semantic model output hidden state 16 shape: torch.Size([1, 49, 1024])
Compiled semantic model works! Shape: torch.Size([1, 49, 1024])


JIT trace on CodecEnc works! Shape: torch.Size([1, 50, 1024])


python3 -c "
# Test torch.jit.trace on Wav2Vec2BertModel
import torch
from transformers import Wav2Vec2BertModel

model = Wav2Vec2BertModel.from_pretrained('facebook/w2v-bert-2.0', output_hidden_states=True).cuda().eval()
x = torch.randn(1, 49, 160).cuda()

with torch.no_grad():
    try:
        traced = torch.jit.trace(model, x)
        out = traced(x)
        print('JIT trace on semantic model works!')
    except Exception as e:
        print(f'JIT trace FAILED: {e}')
" 2>/dev/null


python3 -c "
# Check if we can replace numpy fbank with a torch-based one on GPU
# The SeamlessM4T extractor uses: frame_length=400, hop_length=160, fft_length=512,
# power=2.0, center=False, preemphasis=0.97, mel_filters, log_mel='log', remove_dc_offset=True
# 80 mel bins

# Can we use torchaudio.transforms.MelSpectrogram on GPU?
import torch
import torchaudio

# Create a GPU-based mel spectrogram
mel_spec = torchaudio.transforms.MelSpectrogram(
    sample_rate=16000,
    n_fft=512,
    win_length=400,
    hop_length=160,
    center=False,
    n_mels=80,
    power=2.0,
).cuda()

x = torch.randn(1, 16000).cuda()
out = mel_spec(x)
print('MelSpectrogram output shape:', out.shape)
print('Expected: ~(1, 80, 99) for 1s audio')
# Note: still need preemphasis, remove_dc_offset, log, and per-mel normalization
" 2>/dev/null


JIT trace FAILED: Tracer cannot infer type of Wav2Vec2BaseModelOutput(last_hidden_state=tensor([[[ 0.0338,  0.0299, -0.0152,  ...,  0.0235,  0.0326, -0.0012],
         [ 0.0168, -0.0285,  0.0463,  ..., -0.0723,  0.0209,  0.0522],
         [-0.0068, -0.0059,  0.0156,  ..., -0.0461,  0.0006,  0.0513],
         ...,
         [-0.0375, -0.0818,  0.0037,  ..., -0.0650, -0.0081,  0.0365],
         [ 0.0174,  0.0181, -0.0053,  ...,  0.0060,  0.0394, -0.0123],
         [ 0.0356, -0.0620, -0.0073,  ..., -0.0807,  0.0313,  0.0304]]],
       device='cuda:0'), extract_features=tensor([[[-9.2007, -0.2769, -3.9063,  ..., -0.8255,  1.7240, -2.4054],
         [ 2.1918, -3.9839,  5.0094,  ..., -4.3784, -0.4852, -3.2457],
         [ 7.8914, -0.8944,  3.0996,  ..., -0.6127, -0.4051, -1.9401],
         ...,
         [ 6.0368, -0.9419, -0.1237,  ..., -6.1062,  0.9141, -5.4367],
         [ 3.3271,  0.2633, -2.0592,  ..., -2.0963,  1.7059,  2.1499],
         [ 2.5975,  0.8916,  0.8181,  ..., -2.0145,  0.9051, -4.8502]]],
       device='cuda:0'), hidden_states=(tensor([[[-3.2356,  3.9294, -6.5169,  ..., -7.4877,  0.4945, -7.1125],
         [-2.3062, -0.1526, -2.3617,  ...,  1.9337, -3.5641,  2.1413],
         [10.4035,  1.9082,  7.1749,  ..., -0.3213,  2.3395,  3.9023],
         ...,
         [-6.2755,  3.0096,  6.0477,  ...,  2.7962,  1.5157,  5.5045],
         [ 1.9296,  1.8087, -1.7949,  ...,  1.7674,  4.1199, -0.0910],
         [-0.7492, -2.6836,  1.9642,  ...,  0.7835,  3.9289, -2.7955]]],
       device='cuda:0'), tensor([[[-1.1124,  0.3362, -0.5096,  ..., -0.2007, -0.2038,  0.0470],
         [-0.4777, -0.4853, -0.1870,  ...,  0.3801, -0.6546,  1.1555],
         [ 0.1407,  0.0685, -0.0423,  ...,  0.2229,  0.1136,  0.6811],
         ...,
         [-0.9373,  0.0972,  0.2020,  ...,  0.3665, -0.1152,  0.8886],
         [-0.0124,  0.0514, -0.3225,  ..., -0.5155,  0.1708,  0.5894],
         [-0.3817, -0.1746,  0.1110,  ..., -0.0015,  0.2694,  0.7083]]],
       device='cuda:0'), tensor([[[-0.3218,  0.3154, -0.5130,  ...,  0.3759,  0.1231,  0.3793],
         [-0.0794,  0.3788, -0.2263,  ...,  0.8228,  0.2176,  0.4715],
         [-0.1938,  0.4898,  0.0226,  ...,  0.2429,  0.0425,  0.2892],
         ...,
         [-0.0018,  0.4851,  0.1951,  ...,  0.2953,  0.4097,  0.3401],
         [-0.1416,  0.6747, -0.0144,  ..., -0.2081,  0.2870,  0.3004],
         [ 0.0355,  0.7331,  0.4466,  ..., -0.1977,  0.2736,  0.3495]]],
       device='cuda:0'), tensor([[[ 0.2297, -0.4026,  0.3519,  ...,  0.4844,  0.2443,  0.5040],
         [-0.5617, -0.2115,  0.0155,  ...,  0.7365,  0.2984,  0.5134],
         [-0.1501, -0.3028,  0.2203,  ...,  0.7040, -0.0853,  0.3357],
         ...,
         [-0.0548, -0.5085,  0.2535,  ...,  0.1183,  0.3713,  0.1900],
         [-0.1920, -0.3242, -0.2995,  ..., -0.2048,  0.2009, -0.2241],
         [ 0.3046, -0.1471,  0.1825,  ...,  0.0484, -0.0211,  0.0727]]],
       device='cuda:0'), tensor([[[ 0.2640,  0.7802, -0.2302,  ...,  0.0178,  0.0278, -0.2736],
         [-0.0561,  0.4083, -0.1288,  ...,  0.5370, -0.0377, -0.0034],
         [ 0.2269,  0.4742, -0.0116,  ...,  0.4923, -0.0616, -0.1667],
         ...,
         [-0.0845,  0.0243,  0.3352,  ..., -0.1264,  0.1711, -0.8234],
         [-0.0770, -0.0596,  0.1442,  ..., -0.0436, -0.2130, -0.3855],
         [ 0.2597,  0.4799,  0.2887,  ..., -0.0827, -0.2379, -0.4332]]],
       device='cuda:0'), tensor([[[-0.3020,  0.1635, -0.2993,  ...,  0.1666, -0.0205,  0.2402],
         [-0.1842,  0.0414,  0.1032,  ...,  0.3212,  0.4652,  0.6508],
         [-0.3631,  0.1203,  0.2554,  ...,  0.3111,  0.5734,  0.2393],
         ...,
         [-0.3090, -0.1827,  0.5709,  ...,  0.2884,  0.4172, -0.6424],
         [-0.6153, -0.2413,  0.9199,  ..., -0.3123, -0.0733, -0.2641],
         [-0.5262, -0.0680,  0.4121,  ..., -0.0661,  0.0841, -0.2996]]],
       device='cuda:0'), tensor([[[-0.4084, -0.3193, -0.3898,  ...,  0.2345, -0.0318, -0.0201],
         [-0.6283, -0.2987, -0.3013,  ...,  0.5254,  0.1141,  0.4576],
         [-0.5322,  0.0925, -0.3446,  ...,  0.7586,  0.1802,  0.3648],
         ...,
         [-0.6571, -0.3912, -0.3548,  ...,  0.3940, -0.1805, -0.2007],
         [-0.5800, -0.0666, -0.4303,  ...,  0.1647,  0.0126, -0.4044],
         [-0.6339, -0.0903, -0.5236,  ...,  0.3975, -0.2372, -0.0325]]],
       device='cuda:0'), tensor([[[ 0.0794, -0.0421, -0.2001,  ..., -0.2809,  0.4390, -0.3924],
         [ 0.0340,  0.1246, -0.1858,  ...,  0.2465,  0.4117,  0.2475],
         [-0.1012,  0.0921,  0.0083,  ...,  0.3694,  0.5167,  0.3589],
         ...,
         [-0.0500, -0.0082, -0.1003,  ...,  0.1937,  0.2552,  0.2263],
         [ 0.3356, -0.0301,  0.0400,  ..., -0.0374,  0.4068, -0.3374],
         [-0.0673, -0.0506, -0.1732,  ...,  0.2463,  0.2822,  0.0482]]],
       device='cuda:0'), tensor([[[ 0.0944, -0.1995, -0.0445,  ...,  0.0240, -0.2173, -0.2578],
         [ 0.0753, -0.2340, -0.0964,  ...,  0.1247, -0.1097, -0.0632],
         [-0.0801, -0.1624, -0.0337,  ...,  0.1595, -0.2053, -0.1212],
         ...,
         [-0.0284, -0.3011, -0.0695,  ...,  0.2039, -0.0857,  0.0177],
         [ 0.1289, -0.2051, -0.1052,  ...,  0.1244, -0.1617, -0.1792],
         [-0.0779, -0.2201, -0.1073,  ...,  0.0686, -0.1877, -0.1409]]],
       device='cuda:0'), tensor([[[ 0.1232,  0.0916,  0.1072,  ..., -0.1072,  0.2189, -0.1932],
         [ 0.5058,  0.3602, -0.5817,  ..., -0.0699,  0.9494,  0.1224],
         [ 0.2086,  0.2802, -0.4191,  ...,  0.2202,  0.9467,  0.0395],
         ...,
         [ 0.1342,  0.0632, -0.8452,  ..., -0.2314,  0.9535,  0.0989],
         [ 0.0281,  0.1114, -0.2264,  ..., -0.2917,  0.1745,  0.0070],
         [ 0.2722,  0.1387, -0.4840,  ..., -0.0091,  0.6823, -0.0886]]],
       device='cuda:0'), tensor([[[ 0.5359,  0.2592,  0.2531,  ...,  0.4151,  0.2098, -0.1725],
         [ 0.3078, -0.3775, -0.8966,  ..., -0.0977,  0.7835,  0.1195],
         [ 0.0089, -0.3181, -0.9150,  ...,  0.2482,  0.5157, -0.0126],
         ...,
         [ 0.3635, -0.6296, -1.0328,  ..., -0.1099,  0.6666,  0.1908],
         [ 0.2266, -0.2342, -0.2392,  ..., -0.3046,  0.0398, -0.1527],
         [ 0.5094, -0.4256, -0.8317,  ..., -0.1081,  0.5159, -0.0781]]],
       device='cuda:0'), tensor([[[ 0.2308,  0.6091, -0.2331,  ...,  0.7738,  0.1225, -0.2650],
         [ 0.6199, -0.1058, -1.1384,  ..., -0.4595,  0.5721,  0.1937],
         [ 0.6247, -0.1368, -1.1933,  ..., -0.1544,  0.3810,  0.2789],
         ...,
         [ 0.3719, -0.3218, -0.7789,  ..., -0.5569,  0.2554,  0.3081],
         [-0.0771,  0.0992, -0.4011,  ..., -0.5757, -0.0110,  0.2254],
         [ 0.3572, -0.2645, -0.6512,  ..., -0.9068,  0.1115,  0.1401]]],
       device='cuda:0'), tensor([[[ 0.2601,  0.0229, -0.0851,  ...,  0.2828, -0.0268, -0.1584],
         [ 0.1406, -0.0570, -0.9548,  ..., -0.1791,  0.1741,  0.1205],
         [ 0.2045, -0.1129, -0.9564,  ..., -0.1714,  0.1655,  0.0530],
         ...,
         [ 0.0298, -0.2945, -0.8160,  ..., -0.5713,  0.0104,  0.1446],
         [-0.4000, -0.2452, -0.0960,  ..., -0.4683, -0.0462,  0.3516],
         [-0.1398, -0.2164, -0.6027,  ..., -0.9295, -0.0194,  0.0191]]],
       device='cuda:0'), tensor([[[ 0.1847,  0.1137, -0.1059,  ...,  0.0551,  0.0855, -0.2080],
         [ 0.2163,  0.1739, -0.4497,  ...,  0.0288,  0.1000, -0.3129],
         [ 0.1073,  0.0506, -0.4962,  ...,  0.0627,  0.0886, -0.4518],
         ...,
         [ 0.0334, -0.1721, -0.5852,  ..., -0.2339,  0.2250, -0.2779],
         [-0.5107,  0.0480, -0.0734,  ..., -0.2845,  0.0146,  0.0010],
         [-0.0344, -0.0742, -0.3728,  ..., -0.3214,  0.2304, -0.2916]]],
       device='cuda:0'), tensor([[[ 0.2347, -0.1506,  0.4854,  ..., -0.4073, -0.1390, -0.3774],
         [ 0.2732, -0.3458, -0.0262,  ..., -0.1931,  0.0116, -0.2534],
         [ 0.2396, -0.4045, -0.1124,  ..., -0.1686,  0.0652, -0.2638],
         ...,
         [-0.1267, -0.5686, -0.0879,  ..., -0.2155,  0.1941, -0.0760],
         [-0.1960, -0.1293,  0.3018,  ..., -0.2407, -0.0978, -0.0139],
         [-0.0744, -0.4919,  0.0396,  ..., -0.2896,  0.2305, -0.0428]]],
       device='cuda:0'), tensor([[[ 0.4011,  0.1458,  0.1308,  ..., -0.3605,  0.4148, -0.4748],
         [ 0.4030, -0.0951, -0.2213,  ..., -0.4513,  0.1464, -0.2151],
         [ 0.2223, -0.0662, -0.2571,  ..., -0.5812,  0.3183, -0.1709],
         ...,
         [-0.0100,  0.0526, -0.0937,  ..., -0.4223,  0.4574, -0.2727],
         [ 0.0612,  0.1694, -0.0052,  ...,  0.0295,  0.4626, -0.1101],
         [-0.0055,  0.0817, -0.0589,  ..., -0.3790,  0.4439, -0.2235]]],
       device='cuda:0'), tensor([[[ 0.5697, -0.3376, -0.0199,  ...,  0.5842,  0.1612, -0.4484],
         [ 0.1004, -0.8812, -0.3376,  ..., -0.3023,  0.0684,  0.1830],
         [ 0.1814, -0.7755, -0.4685,  ..., -0.6108,  0.1038,  0.0221],
         ...,
         [ 0.3034, -0.6664, -0.1870,  ..., -0.5318,  0.2769,  0.0269],
         [-0.0449, -0.5270, -0.2551,  ...,  0.2133,  0.6460,  0.0820],
         [ 0.1326, -0.6496, -0.4286,  ..., -0.5081,  0.3573,  0.2194]]],
       device='cuda:0'), tensor([[[-0.0662, -0.0969,  0.4244,  ...,  0.2656,  0.0067, -0.1476],
         [-0.1419, -0.3081, -0.3009,  ..., -0.4577,  0.4326, -0.1521],
         [-0.2688, -0.3234, -0.4971,  ..., -0.4097,  0.6101, -0.2382],
         ...,
         [-0.1102, -0.5205, -0.0705,  ..., -0.4572,  0.6062, -0.1152],
         [-0.6211, -0.5700, -0.0110,  ..., -0.1581,  1.0141,  0.0716],
         [-0.2604, -0.2587, -0.1561,  ..., -0.1416,  0.6386,  0.0076]]],
       device='cuda:0'), tensor([[[ 0.1795, -1.0566,  0.1002,  ..., -0.1394,  0.1057,  0.5246],
         [ 0.1719, -1.1530, -0.1117,  ..., -0.1021, -0.0420,  0.1809],
         [-0.0367, -0.8510, -0.0286,  ...,  0.0982,  0.0356,  0.1516],
         ...,
         [ 0.2770, -1.0588,  0.4345,  ..., -0.2482,  0.0980,  0.1352],
         [-0.2375, -0.8931,  0.2813,  ...,  0.2215,  0.5758,  0.3155],
         [ 0.2263, -0.7043,  0.1551,  ...,  0.1339,  0.1218,  0.3671]]],
       device='cuda:0'), tensor([[[ 0.1116, -0.3757,  0.3904,  ..., -0.0410,  0.6068,  0.1904],
         [ 0.3873, -0.8491, -0.1043,  ...,  0.2039, -0.2020,  0.3331],
         [ 0.2491, -0.6019, -0.2343,  ...,  0.2190, -0.1746,  0.1472],
         ...,
         [ 0.2772, -0.7071, -0.0305,  ...,  0.0830, -0.1620, -0.0335],
         [-0.0074, -0.7141, -0.0014,  ...,  0.1611,  0.4871,  0.1673],
         [ 0.3784, -0.5136, -0.3873,  ...,  0.2674, -0.0334,  0.2901]]],
       device='cuda:0'), tensor([[[ 0.1942, -0.0840, -0.0230,  ..., -0.2757,  0.7651, -0.0417],
         [ 0.1584, -0.7506, -0.4724,  ..., -0.0377,  0.3797, -0.2194],
         [-0.0998, -0.5148, -0.3983,  ..., -0.3114,  0.1952, -0.1920],
         ...,
         [ 0.3077, -0.4228,  0.0416,  ..., -0.4202,  0.2987, -0.1988],
         [-0.0225, -0.4284, -0.3918,  ...,  0.0095,  0.8021, -0.2512],
         [ 0.3871, -0.4587, -0.2416,  ..., -0.2258,  0.5465,  0.0590]]],
       device='cuda:0'), tensor([[[ 0.3400,  0.0249,  0.4201,  ..., -0.2709,  0.4625,  0.2509],
         [ 0.3704, -0.3258, -0.1867,  ..., -0.3521,  0.5888, -0.2620],
         [ 0.3981, -0.5729,  0.0528,  ..., -0.8177,  0.4877, -0.1373],
         ...,
         [ 0.4972, -0.7975,  0.4959,  ..., -0.7930,  0.4658, -0.2721],
         [-0.0296,  0.0363, -0.0035,  ...,  0.2228,  0.4452,  0.1546],
         [ 0.3928, -0.5975,  0.1264,  ..., -0.7736,  0.5738,  0.1957]]],
       device='cuda:0'), tensor([[[ 0.0212,  0.2778, -0.1559,  ..., -0.1167,  0.2267,  0.1283],
         [ 0.0051,  0.2455, -0.1393,  ..., -0.4033,  0.2743, -0.1595],
         [ 0.1814,  0.2721, -0.1808,  ..., -0.4883,  0.3803, -0.0453],
         ...,
         [ 0.2063, -0.2924,  0.0103,  ..., -0.4266,  0.2488, -0.1423],
         [ 0.2942, -0.0622, -0.4165,  ...,  0.0289,  0.1789,  0.0489],
         [ 0.0470,  0.0214, -0.0934,  ..., -0.5871,  0.1671,  0.2788]]],
       device='cuda:0'), tensor([[[ 0.2339,  0.1900,  0.0613,  ..., -0.0188,  0.0510,  0.1409],
         [-0.1574, -0.0505,  0.1995,  ..., -0.4683, -0.0541,  0.2886],
         [-0.0392, -0.0411,  0.0412,  ..., -0.3830,  0.0454,  0.2353],
         ...,
         [ 0.1512, -0.5771,  0.1143,  ..., -0.4071, -0.1587,  0.2289],
         [ 0.3090, -0.2457, -0.0778,  ..., -0.0597,  0.2074,  0.1692],
         [ 0.3077, -0.1574,  0.1786,  ..., -0.6329, -0.2438,  0.4567]]],
       device='cuda:0'), tensor([[[ 0.0338,  0.0299, -0.0152,  ...,  0.0235,  0.0326, -0.0012],
         [ 0.0168, -0.0285,  0.0463,  ..., -0.0723,  0.0209,  0.0522],
         [-0.0068, -0.0059,  0.0156,  ..., -0.0461,  0.0006,  0.0513],
         ...,
         [-0.0375, -0.0818,  0.0037,  ..., -0.0650, -0.0081,  0.0365],
         [ 0.0174,  0.0181, -0.0053,  ...,  0.0060,  0.0394, -0.0123],
         [ 0.0356, -0.0620, -0.0073,  ..., -0.0807,  0.0313,  0.0304]]],
       device='cuda:0')), attentions=None)
:Dictionary inputs to traced functions must have consistent type. Found Tensor and Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]


MelSpectrogram output shape: torch.Size([1, 80, 97])
Expected: ~(1, 80, 99) for 1s audio


python3 -c "
# Check if ONNX export works for CodecEncoder
import torch
from neucodec.codec_encoder import CodecEncoder

enc = CodecEncoder().eval()
x = torch.randn(1, 1, 16000)

try:
    torch.onnx.export(enc, x, '/tmp/codec_enc.onnx',
                      input_names=['audio'], output_names=['embedding'],
                      dynamic_axes={'audio': {2: 'time'}, 'embedding': {1: 'frames'}},
                      opset_version=17)
    print('ONNX export of CodecEncoder succeeded!')
    import os
    print(f'File size: {os.path.getsize(\"/tmp/codec_enc.onnx\") / 1024 / 1024:.1f} MB')
except Exception as e:
    print(f'ONNX export FAILED: {e}')
" 2>/dev/null


python3 -c "
# Check the bug on line 34 more carefully and test truncating layers
from transformers import Wav2Vec2BertModel
m = Wav2Vec2BertModel.from_pretrained('facebook/w2v-bert-2.0', output_hidden_states=True)
print('Before truncation:', len(m.encoder.layers))

# Actually truncate
import torch.nn as nn
m.encoder.layers = nn.ModuleList(list(m.encoder.layers)[:17])  # Keep 0-16 (17 layers)
print('After truncation:', len(m.encoder.layers))

# Test it works
import torch
x = torch.randn(1, 49, 160)
with torch.no_grad():
    out = m(x)
    print('Hidden states count:', len(out.hidden_states))
    print('Layer 16 shape:', out.hidden_states[16].shape)
" 2>/dev/null


python3 -c "
# Check batched fbank feature extraction - can we batch the feature_extractor call?
import numpy as np
from transformers import AutoFeatureExtractor

fe = AutoFeatureExtractor.from_pretrained('facebook/w2v-bert-2.0')

# Test batched call with different length audio
audio1 = np.random.randn(8000).astype(np.float32)  # 0.5s
audio2 = np.random.randn(32000).astype(np.float32)  # 2s
audio3 = np.random.randn(16000).astype(np.float32)  # 1s

out = fe([audio1, audio2, audio3], sampling_rate=16000, return_tensors='pt',
         padding=True, return_attention_mask=True)
print('Batched fbank shape:', out.input_features.shape)
print('Attention mask shape:', out.attention_mask.shape)
print('Attention mask[0] sum:', out.attention_mask[0].sum().item())
print('Attention mask[1] sum:', out.attention_mask[1].sum().item())
print('Attention mask[2] sum:', out.attention_mask[2].sum().item())
" 2>/dev/null


[torch.onnx] Obtain model graph for `CodecEncoder([...]` with `torch.export.export(..., strict=False)`...
[torch.onnx] Obtain model graph for `CodecEncoder([...]` with `torch.export.export(..., strict=False)`... ✅
[torch.onnx] Run decomposition...
[torch.onnx] Run decomposition... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
Applied 73 of general pattern rewrite rules.
ONNX export of CodecEncoder succeeded!
File size: 1.3 MB


Before truncation: 24
After truncation: 17
Hidden states count: 18
Layer 16 shape: torch.Size([1, 49, 1024])


Batched fbank shape: torch.Size([3, 99, 160])
Attention mask shape: torch.Size([3, 99])
Attention mask[0] sum: 24
Attention mask[1] sum: 99
Attention mask[2] sum: 49


python3 -c "
# Test ONNX export for the semantic_model (Wav2Vec2BertModel)
import torch
from transformers import Wav2Vec2BertModel
import torch.nn as nn

model = Wav2Vec2BertModel.from_pretrained('facebook/w2v-bert-2.0', output_hidden_states=True).eval()
# Truncate to 17 layers
model.encoder.layers = nn.ModuleList(list(model.encoder.layers)[:17])

x = torch.randn(1, 49, 160)

try:
    torch.onnx.export(model, x, '/tmp/semantic.onnx',
                      input_names=['features'], output_names=['output'],
                      dynamic_axes={'features': {0: 'batch', 1: 'seq'}},
                      opset_version=17)
    print('ONNX export succeeded!')
    import os
    print(f'File size: {os.path.getsize(\"/tmp/semantic.onnx\") / 1024 / 1024:.1f} MB')
except Exception as e:
    print(f'ONNX export FAILED: {type(e).__name__}: {str(e)[:500]}')
" 2>/dev/null


python3 -c "
# Test whether Wav2Vec2BertModel supports attention_mask for batched inference with padding
import torch
from transformers import Wav2Vec2BertModel

model = Wav2Vec2BertModel.from_pretrained('facebook/w2v-bert-2.0', output_hidden_states=True).eval()

# Simulate batched input with attention mask
features = torch.randn(2, 99, 160)  # batch of 2, padded to 99 frames
attention_mask = torch.ones(2, 99)
attention_mask[0, 49:] = 0  # first sample is shorter

with torch.no_grad():
    # With attention mask
    out_masked = model(features, attention_mask=attention_mask)
    print('With attention_mask - hidden state 16 shape:', out_masked.hidden_states[16].shape)

    # Without attention mask (single sample, no padding)
    out_single = model(features[0:1, :49, :])
    print('Single sample shape:', out_single.hidden_states[16].shape)

    # Check if values roughly match for valid frames
    h16_masked = out_masked.hidden_states[16][0]
    h16_single = out_single.hidden_states[16][0]
    print('Shapes - masked first:', h16_masked.shape, 'single:', h16_single.shape)
    # The adapter in wav2vec2-bert downsamples by 2, so 99 frames -> ~50, 49 frames -> ~25
" 2>/dev/null


[torch.onnx] Obtain model graph for `Wav2Vec2BertModel([...]` with `torch.export.export(..., strict=False)`...
[torch.onnx] Obtain model graph for `Wav2Vec2BertModel([...]` with `torch.export.export(..., strict=False)`... ✅
[torch.onnx] Run decomposition...
[torch.onnx] Run decomposition... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
Applied 89 of general pattern rewrite rules.
ONNX export succeeded!
File size: 2.4 MB


With attention_mask - hidden state 16 shape: torch.Size([2, 99, 1024])
Single sample shape: torch.Size([1, 49, 1024])
Shapes - masked first: torch.Size([99, 1024]) single: torch.Size([49, 1024])


python3 -c "
# Check the adapter layer in wav2vec2-bert - does it downsample?
from transformers import Wav2Vec2BertModel
model = Wav2Vec2BertModel.from_pretrained('facebook/w2v-bert-2.0', output_hidden_states=True)
print('Has adapter:', hasattr(model, 'adapter'))
print('Config adapter_kernel_size:', model.config.adapter_kernel_size)
print('Config adapter_stride:', model.config.adapter_stride)
# So the adapter downsamples by stride 2, meaning the fbank @ 100Hz -> 50Hz after adapter
# But the hidden_states[16] is BEFORE or AFTER adapter?

import torch
x = torch.randn(1, 99, 160)
with torch.no_grad():
    out = model(x)
    print('Input fbank frames:', 99)
    print('hidden_states[0] shape:', out.hidden_states[0].shape)  # after feature projection
    print('hidden_states[1] shape:', out.hidden_states[1].shape)  # after first encoder layer
    print('hidden_states[16] shape:', out.hidden_states[16].shape)  # after layer 16
    print('last_hidden_state shape:', out.last_hidden_state.shape)
" 2>/dev/null


python3 -c "
# Test running CodecEnc and semantic model in parallel with CUDA streams
import torch
import time
from neucodec.model import NeuCodec

model = NeuCodec.from_pretrained('neuphonic/neucodec').cuda().eval()

# Create test audio
audio = torch.randn(1, 1, 80000).cuda()  # 5 seconds

# Warm up
with torch.no_grad():
    _ = model.encode_code(audio)
torch.cuda.synchronize()

# Benchmark sequential
times = []
for _ in range(5):
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    with torch.no_grad():
        _ = model.encode_code(audio)
    torch.cuda.synchronize()
    times.append(time.perf_counter() - t0)
print(f'Sequential avg: {sum(times)/len(times)*1000:.1f} ms')
" 2>/dev/null


Has adapter: True
Config adapter_kernel_size: 3
Config adapter_stride: 2
Input fbank frames: 99
hidden_states[0] shape: torch.Size([1, 99, 1024])
hidden_states[1] shape: torch.Size([1, 99, 1024])
hidden_states[16] shape: torch.Size([1, 99, 1024])
last_hidden_state shape: torch.Size([1, 99, 1024])


Sequential avg: 164.0 ms


python3 -c "
# Test CUDA streams - run CodecEnc and feature extraction concurrently
import torch
import time
from neucodec.model import NeuCodec

model = NeuCodec.from_pretrained('neuphonic/neucodec').cuda().eval()
audio = torch.randn(1, 1, 80000)  # 5 seconds

# Prepare audio
y = model._prepare_audio(audio)

# Warm up
with torch.no_grad():
    _ = model.encode_code(audio)
torch.cuda.synchronize()

# Test with CUDA streams - overlap fbank (CPU) with CodecEnc (GPU)
stream1 = torch.cuda.Stream()
stream2 = torch.cuda.Stream()

times = []
for _ in range(5):
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    with torch.no_grad():
        # Stream 1: Acoustic encoding on GPU
        with torch.cuda.stream(stream1):
            acoustic_emb = model.CodecEnc(y.to(model.device))
            acoustic_emb = acoustic_emb.transpose(1, 2)

        # Meanwhile, feature extraction on CPU (runs concurrently with GPU)
        all_semantic_features = []
        for i in range(y.size(0)):
            sf = model.feature_extractor(
                y[i, :].cpu(), sampling_rate=16_000, return_tensors='pt'
            ).input_features.to(model.device)
            all_semantic_features.append(sf)
        semantic_features = torch.vstack(all_semantic_features)

        # Wait for stream1 to finish
        stream1.synchronize()

        # Stream 2: Semantic model
        with torch.cuda.stream(stream2):
            semantic_output = model.semantic_model(semantic_features).hidden_states[16].transpose(1, 2)
            semantic_encoded = model.SemanticEncoder_module(semantic_output)
        stream2.synchronize()

        # Combine
        if acoustic_emb.shape[-1] != semantic_encoded.shape[-1]:
            min_len = min(acoustic_emb.shape[-1], semantic_encoded.shape[-1])
            acoustic_emb = acoustic_emb[:, :, :min_len]
            semantic_encoded = semantic_encoded[:, :, :min_len]
        concat_emb = torch.cat([semantic_encoded, acoustic_emb], dim=1)
        concat_emb = model.fc_prior(concat_emb.transpose(1, 2)).transpose(1, 2)
        _, fsq_codes, _ = model.generator(concat_emb, vq=True)

    torch.cuda.synchronize()
    times.append(time.perf_counter() - t0)

print(f'With CUDA streams avg: {sum(times)/len(times)*1000:.1f} ms')
" 2>/dev/null


python3 -c "
# Test torch.compile on individual GPU submodules
import torch
import time
from neucodec.model import NeuCodec

model = NeuCodec.from_pretrained('neuphonic/neucodec').cuda().eval()

# Compile individual submodules
model.CodecEnc = torch.compile(model.CodecEnc, mode='reduce-overhead')
model.semantic_model = torch.compile(model.semantic_model, mode='reduce-overhead')

audio = torch.randn(1, 1, 80000)

# Warm up (triggers compilation)
print('Warming up (compiling)...')
with torch.no_grad():
    _ = model.encode_code(audio)
    _ = model.encode_code(audio)
    _ = model.encode_code(audio)
torch.cuda.synchronize()
print('Warm up done')

# Benchmark
times = []
for _ in range(10):
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    with torch.no_grad():
        _ = model.encode_code(audio)
    torch.cuda.synchronize()
    times.append(time.perf_counter() - t0)

print(f'With torch.compile on submodules avg: {sum(times)/len(times)*1000:.1f} ms')
print(f'Min: {min(times)*1000:.1f} ms, Max: {max(times)*1000:.1f} ms')
" 2>/dev/null


With CUDA streams avg: 168.4 ms


Exit code 1
Warming up (compiling)...

Error: Exit code 1
Warming up (compiling)...


python3 -c "
# torch.compile with mode='default' instead of 'reduce-overhead'
import torch
import time
from neucodec.model import NeuCodec

model = NeuCodec.from_pretrained('neuphonic/neucodec').cuda().eval()

# Compile individual submodules with default mode
model.CodecEnc = torch.compile(model.CodecEnc, mode='default')
model.semantic_model = torch.compile(model.semantic_model, mode='default')

audio = torch.randn(1, 1, 80000)

print('Warming up (compiling)...')
with torch.no_grad():
    for i in range(3):
        _ = model.encode_code(audio)
        print(f'  warmup {i+1} done')
torch.cuda.synchronize()

times = []
for _ in range(10):
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    with torch.no_grad():
        _ = model.encode_code(audio)
    torch.cuda.synchronize()
    times.append(time.perf_counter() - t0)

print(f'With torch.compile (default) avg: {sum(times)/len(times)*1000:.1f} ms')
print(f'Min: {min(times)*1000:.1f} ms')
" 2>/dev/null


Warming up (compiling)...
  warmup 1 done
  warmup 2 done
  warmup 3 done
With torch.compile (default) avg: 169.0 ms
Min: 157.1 ms


python3 -c "
# Test the HUGE optimization: truncating wav2vec2-bert encoder layers
# This should save ~30% of semantic model time (8 fewer layers out of 24)
import torch
import torch.nn as nn
import time
from neucodec.model import NeuCodec

model = NeuCodec.from_pretrained('neuphonic/neucodec').cuda().eval()

# Baseline
audio = torch.randn(1, 1, 80000)
with torch.no_grad():
    _ = model.encode_code(audio)
torch.cuda.synchronize()

times = []
for _ in range(10):
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    with torch.no_grad():
        _ = model.encode_code(audio)
    torch.cuda.synchronize()
    times.append(time.perf_counter() - t0)
baseline = sum(times)/len(times)
print(f'Baseline (24 layers): {baseline*1000:.1f} ms')

# Truncate to 17 layers (layers 0-16, since we only need hidden_states[16])
model.semantic_model.encoder.layers = nn.ModuleList(list(model.semantic_model.encoder.layers)[:17])
print(f'Encoder layers after truncation: {len(model.semantic_model.encoder.layers)}')

# Verify correctness
with torch.no_grad():
    result = model.encode_code(audio)
    print(f'Output shape: {result.shape}')

# Benchmark truncated
times2 = []
for _ in range(10):
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    with torch.no_grad():
        _ = model.encode_code(audio)
    torch.cuda.synchronize()
    times2.append(time.perf_counter() - t0)
truncated = sum(times2)/len(times2)
print(f'Truncated (17 layers): {truncated*1000:.1f} ms')
print(f'Speedup: {baseline/truncated:.2f}x ({(1-truncated/baseline)*100:.1f}% faster)')
" 2>/dev/null


Baseline (24 layers): 225.9 ms
Encoder layers after truncation: 17
Output shape: torch.Size([1, 1, 250])
Truncated (17 layers): 185.7 ms
Speedup: 1.22x (17.8% faster)


python3 -c "
# Verify output is IDENTICAL with truncation (not just similar)
import torch
import torch.nn as nn
from neucodec.model import NeuCodec

model1 = NeuCodec.from_pretrained('neuphonic/neucodec').cuda().eval()
model2 = NeuCodec.from_pretrained('neuphonic/neucodec').cuda().eval()

# Truncate model2
model2.semantic_model.encoder.layers = nn.ModuleList(list(model2.semantic_model.encoder.layers)[:17])

audio = torch.randn(1, 1, 80000)

with torch.no_grad():
    codes1 = model1.encode_code(audio)
    codes2 = model2.encode_code(audio)
    print('Codes match exactly:', (codes1 == codes2).all().item())
    print('codes1 sample:', codes1[0, 0, :10])
    print('codes2 sample:', codes2[0, 0, :10])
" 2>/dev/null


Codes match exactly: True
codes1 sample: tensor([12303, 52931, 16176, 12540,   199, 65475, 13104, 12350,  2060,   960],
       device='cuda:0', dtype=torch.int32)
codes2 sample: tensor([12303, 52931, 16176, 12540,   199, 65475, 13104, 12350,  2060,   960],
       device='cuda:0', dtype=torch.int32)


This is a major finding. Let me now test bucketed batching and the GPU fbank replacement.


python3 -c "
# Test bucketed batching: batch audio of similar duration together
# and use the batched feature_extractor with attention_mask
import torch
import torch.nn as nn
import time
import numpy as np
from neucodec.model import NeuCodec

model = NeuCodec.from_pretrained('neuphonic/neucodec').cuda().eval()
model.semantic_model.encoder.layers = nn.ModuleList(list(model.semantic_model.encoder.layers)[:17])

# Simulate a batch of 8 similar-length audio clips (4.5-5.5 seconds)
np.random.seed(42)
durations = np.random.uniform(4.5, 5.5, size=8)
audios = [torch.randn(1, 1, int(d * 16000)) for d in durations]

# Sequential processing baseline
with torch.no_grad():
    for a in audios:
        _ = model.encode_code(a)
torch.cuda.synchronize()

times = []
for _ in range(5):
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    with torch.no_grad():
        for a in audios:
            _ = model.encode_code(a)
    torch.cuda.synchronize()
    times.append(time.perf_counter() - t0)
print(f'Sequential (8 clips): {sum(times)/len(times)*1000:.1f} ms')

# Batched: pad to max length and process together
max_len = max(a.shape[-1] for a in audios)
max_len = max_len + (320 - max_len % 320)  # pad to multiple of 320
batch = torch.zeros(8, 1, max_len)
for i, a in enumerate(audios):
    batch[i, :, :a.shape[-1]] = a

# Warm up
with torch.no_grad():
    _ = model.encode_code(batch)
torch.cuda.synchronize()

times2 = []
for _ in range(5):
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    with torch.no_grad():
        _ = model.encode_code(batch)
    torch.cuda.synchronize()
    times2.append(time.perf_counter() - t0)
print(f'Batched (8 similar-len clips): {sum(times2)/len(times2)*1000:.1f} ms')
print(f'Speedup: {sum(times)/len(times) / (sum(times2)/len(times2)):.2f}x')
" 2>/dev/null


Sequential (8 clips): 1346.1 ms
Batched (8 similar-len clips): 1163.2 ms
Speedup: 1.16x


python3 -c "
# Test with a tighter bucket (very similar durations) and smaller batch
import torch
import torch.nn as nn
import time
from neucodec.model import NeuCodec

model = NeuCodec.from_pretrained('neuphonic/neucodec').cuda().eval()
model.semantic_model.encoder.layers = nn.ModuleList(list(model.semantic_model.encoder.layers)[:17])

# 4 clips, all ~5s (tight bucket)
audios = [torch.randn(1, 1, int(d * 16000)) for d in [4.9, 5.0, 5.0, 5.1]]

# Sequential
with torch.no_grad():
    for a in audios:
        _ = model.encode_code(a)
torch.cuda.synchronize()

times = []
for _ in range(5):
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    with torch.no_grad():
        for a in audios:
            _ = model.encode_code(a)
    torch.cuda.synchronize()
    times.append(time.perf_counter() - t0)
seq_avg = sum(times)/len(times)

# Batched
max_len = max(a.shape[-1] for a in audios)
max_len = max_len + (320 - max_len % 320)
batch = torch.zeros(4, 1, max_len)
for i, a in enumerate(audios):
    batch[i, :, :a.shape[-1]] = a

with torch.no_grad():
    _ = model.encode_code(batch)
torch.cuda.synchronize()

times2 = []
for _ in range(5):
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    with torch.no_grad():
        _ = model.encode_code(batch)
    torch.cuda.synchronize()
    times2.append(time.perf_counter() - t0)
bat_avg = sum(times2)/len(times2)

print(f'Sequential (4 clips ~5s): {seq_avg*1000:.1f} ms')
print(f'Batched (4 clips tight bucket): {bat_avg*1000:.1f} ms')
print(f'Speedup: {seq_avg/bat_avg:.2f}x')
print(f'Per-clip: sequential={seq_avg/4*1000:.1f}ms, batched={bat_avg/4*1000:.1f}ms')
" 2>/dev/null


Sequential (4 clips ~5s): 594.2 ms
Batched (4 clips tight bucket): 535.8 ms
Speedup: 1.11x
Per-clip: sequential=148.6ms, batched=133.9ms


python3 -c "
# Test the GPU fbank replacement approach
# The CPU fbank takes ~15.8% of time. Let's see if we can replicate it on GPU.
import torch
import numpy as np
import time
from transformers import AutoFeatureExtractor

fe = AutoFeatureExtractor.from_pretrained('facebook/w2v-bert-2.0')

# Check the mel_filters shape and window
print('mel_filters shape:', fe.mel_filters.shape)
print('window shape:', fe.window.shape)

# Time the CPU fbank on 5s audio
audio_np = np.random.randn(1, 80000).astype(np.float32)
audio_pt = torch.from_numpy(audio_np)

# CPU timing
times = []
for _ in range(20):
    t0 = time.perf_counter()
    out = fe(audio_np, sampling_rate=16000, return_tensors='pt')
    times.append(time.perf_counter() - t0)
print(f'CPU fbank avg: {sum(times)/len(times)*1000:.1f} ms')
print(f'Output shape: {out.input_features.shape}')
" 2>/dev/null


mel_filters shape: (257, 80)
window shape: (400,)
CPU fbank avg: 24.8 ms
Output shape: torch.Size([1, 249, 160])


python3 -c "
# Build a GPU-based fbank that matches SeamlessM4T exactly
# Parameters: frame_length=400, hop_length=160, fft_length=512, power=2.0,
# center=False, preemphasis=0.97, mel_filters from the extractor, log_mel='log',
# mel_floor=1.192092955078125e-07, remove_dc_offset=True
# Then stride=2 subsampling and per-mel normalization

import torch
import torch.nn.functional as F
import numpy as np
from transformers import AutoFeatureExtractor

fe = AutoFeatureExtractor.from_pretrained('facebook/w2v-bert-2.0')

class GPUFbank(torch.nn.Module):
    def __init__(self, feature_extractor):
        super().__init__()
        # Store mel filters as buffer
        self.register_buffer('mel_filters', torch.from_numpy(feature_extractor.mel_filters.T).float())  # (80, 257)
        self.register_buffer('window', torch.from_numpy(feature_extractor.window).float())
        self.mel_floor = 1.192092955078125e-07

        # Store mean/std for normalization (from the __call__ method)
        # do_normalize_per_mel_bins requires computing mean/std from the features themselves

    def forward(self, waveform):
        # waveform: (B, 1, T) or (B, T)
        if waveform.dim() == 3:
            waveform = waveform.squeeze(1)

        # Preemphasis: y[n] = x[n] - 0.97 * x[n-1]
        preemph = torch.cat([waveform[:, :1], waveform[:, 1:] - 0.97 * waveform[:, :-1]], dim=1)

        # Remove DC offset per frame (will do via STFT)
        # Scale to 16-bit
        preemph = preemph * (2**15)

        # STFT
        stft = torch.stft(preemph, n_fft=512, hop_length=160, win_length=400,
                          window=self.window, center=False, return_complex=True)
        # Power spectrogram
        power_spec = stft.abs().pow(2)  # (B, 257, T)

        # Apply mel filters
        mel_spec = torch.matmul(self.mel_filters, power_spec)  # (B, 80, T)

        # Log mel
        mel_spec = torch.clamp(mel_spec, min=self.mel_floor)
        log_mel = torch.log(mel_spec)  # (B, 80, T)

        # Transpose to (B, T, 80)
        log_mel = log_mel.transpose(1, 2)

        # Stride 2 subsampling (as SeamlessM4T does)
        # Actually the stride=2 is in the feature extractor's config but it's applied differently
        # Let me check...
        # pad_to_multiple_of=2 means pad time to even, then stride 2 subsampling
        if log_mel.shape[1] % 2 == 1:
            log_mel = F.pad(log_mel, (0, 0, 0, 1), value=1.0)

        # Reshape with stride 2: (B, T, 80) -> (B, T//2, 160)
        B, T, F = log_mel.shape
        log_mel = log_mel.reshape(B, T // 2, F * 2)

        # Per-mel-bin normalization (zero-mean, unit-variance)
        mean = log_mel.mean(dim=1, keepdim=True)
        std = log_mel.std(dim=1, keepdim=True) + 1e-7
        log_mel = (log_mel - mean) / std

        return log_mel

# Test
gpu_fbank = GPUFbank(fe).cuda()
audio = torch.randn(1, 1, 80000)

# CPU reference
cpu_out = fe(audio[0].numpy(), sampling_rate=16000, return_tensors='pt').input_features
print('CPU fbank shape:', cpu_out.shape)

# GPU
with torch.no_grad():
    gpu_out = gpu_fbank(audio.cuda())
print('GPU fbank shape:', gpu_out.shape)

# Compare values
cpu_flat = cpu_out.flatten()
gpu_flat = gpu_out.cpu().flatten()
# Check correlation
corr = torch.corrcoef(torch.stack([cpu_flat, gpu_flat]))[0, 1]
print(f'Correlation: {corr:.6f}')
print(f'Max abs diff: {(cpu_flat - gpu_flat).abs().max():.4f}')
print(f'Mean abs diff: {(cpu_flat - gpu_flat).abs().mean():.4f}')
" 2>/dev/null


Error: Exit code 1
CPU fbank shape: torch.Size([1, 249, 160])


python3 -c "
import torch
import torch.nn.functional as F
import numpy as np
from transformers import AutoFeatureExtractor

fe = AutoFeatureExtractor.from_pretrained('facebook/w2v-bert-2.0')

# Check the exact _extract_fbank_features implementation details more carefully
# The spectrogram function applies: remove_dc_offset, preemphasis, windowing, fft, mel, log
# Then __call__ does: stride 2 reshape, per-mel normalization

audio = np.random.randn(80000).astype(np.float32)

# Step 1: get raw fbank features (before stride/norm)
raw_features = fe._extract_fbank_features(audio[np.newaxis, :])
print('Raw features shape:', raw_features.shape)

# Now check what __call__ does
out = fe(audio[np.newaxis, :], sampling_rate=16000, return_tensors='pt', do_normalize_per_mel_bins=False, padding=False)
print('Without normalization shape:', out.input_features.shape)
" 2>/dev/null


Raw features shape: (498, 80)
Without normalization shape: torch.Size([1, 249, 160])