init

__init__.py

@@ -0,0 +1,18 @@
+"""MossSpeechCodec package
+
+Lightweight, Transformers-style wrapper around the MossSpeech codec used by
+`src/transformers/models/moss_speech/processing_moss_speech.py`.
+
+This module keeps the public API stable for the existing processor while
+organizing the implementation to resemble Hugging Face codec models (e.g.
+`xcodec`, `encodec`). Only the minimal parts required at inference are kept.
+"""
+
+from .configuration_moss_speech_codec import MossSpeechCodecConfig
+from .modeling_moss_speech_codec import MossSpeechCodec
+
+__all__ = [
+    "MossSpeechCodec",
+    "MossSpeechCodecConfig",
+]
+

config.json

@@ -0,0 +1,68 @@
+{
+  "activation_dropout": 0.0,
+  "activation_function": "gelu",
+  "apply_spec_augment": false,
+  "architectures": [
+    "WhisperVQEncoder"
+  ],
+  "auto_map": {
+    "AutoModel": "modeling_moss_speech_codec.MossSpeechCodec",
+    "AutoConfig": "modeling_moss_speech_codec.MossSpeechCodecConfig"
+  },
+  "attention_dropout": 0.0,
+  "begin_suppress_tokens": [
+    220,
+    50257
+  ],
+  "bos_token_id": 50257,
+  "classifier_proj_size": 256,
+  "d_model": 1280,
+  "decoder_attention_heads": 20,
+  "decoder_ffn_dim": 5120,
+  "decoder_layerdrop": 0.0,
+  "decoder_layers": 32,
+  "decoder_start_token_id": 50258,
+  "dropout": 0.0,
+  "encoder_attention_heads": 20,
+  "encoder_causal_attention": true,
+  "encoder_causal_convolution": true,
+  "encoder_ffn_dim": 5120,
+  "encoder_layerdrop": 0.0,
+  "encoder_layers": 32,
+  "eos_token_id": 50257,
+  "init_std": 0.02,
+  "is_encoder_decoder": true,
+  "mask_feature_length": 10,
+  "mask_feature_min_masks": 0,
+  "mask_feature_prob": 0.0,
+  "mask_time_length": 10,
+  "mask_time_min_masks": 2,
+  "mask_time_prob": 0.05,
+  "max_length": 448,
+  "max_source_positions": 1500,
+  "max_target_positions": 448,
+  "median_filter_width": 7,
+  "model_type": "whisper",
+  "num_hidden_layers": 32,
+  "num_mel_bins": 128,
+  "pad_token_id": 50256,
+  "pooling_kernel_size": 4,
+  "pooling_position": 16,
+  "pooling_type": "avg",
+  "quantize_causal_block_size": 200,
+  "quantize_causal_encoder": true,
+  "quantize_commit_coefficient": 0.25,
+  "quantize_ema_decay": 0.99,
+  "quantize_encoder_only": true,
+  "quantize_loss_scale": 10.0,
+  "quantize_position": 16,
+  "quantize_restart_interval": 100,
+  "quantize_vocab_size": 16384,
+  "scale_embedding": false,
+  "skip_language_detection": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.44.1",
+  "use_cache": true,
+  "use_weighted_layer_sum": false,
+  "vocab_size": 51866
+}

configuration_moss_speech_codec.py

@@ -0,0 +1,35 @@
+# coding=utf-8
+# Copyright 2025 OpenMOSS and HuggingFace Inc. teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class MossSpeechCodecConfig(PretrainedConfig):
+    """Lightweight configuration for MossSpeech codec.
+
+    This config is intentionally minimal since the codec assembles a Whisper-VQ
+    encoder and a Flow/HiFT decoder from their own configs and checkpoints.
+    """
+
+    model_type = "moss_speech_codec"
+
+    def __init__(self, sample_rate: int = 16000, return_dict: bool = True, **kwargs):
+        self.sample_rate = int(sample_rate)
+        self.return_dict = bool(return_dict)
+        super().__init__(**kwargs)
+
+
+__all__ = ["MossSpeechCodecConfig"]
+

flow/campplus.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a6ac6a63997761ae2997373e2ee1c47040854b4b759ea41ec48e4e42df0f4d73
+size 28303423

flow/config.yaml

@@ -0,0 +1,111 @@
+# set random seed, so that you may reproduce your result.
+__set_seed1: !apply:random.seed [1986]
+__set_seed2: !apply:numpy.random.seed [1986]
+__set_seed3: !apply:torch.manual_seed [1986]
+__set_seed4: !apply:torch.cuda.manual_seed_all [1986]
+
+# fixed params
+sample_rate: 24000
+llm_input_size: 896
+llm_output_size: 896
+spk_embed_dim: 192
+qwen_pretrain_path: ''
+token_frame_rate: 12.5
+token_mel_ratio: 4
+
+# stream related params
+chunk_size: 5 # streaming inference chunk size, in token
+num_decoding_left_chunks: -1 # streaming inference flow decoder left chunk size, <0 means use all left chunks
+
+
+flow: !new:cosyvoice.flow.flow.CausalMaskedDiffWithXvec
+    input_size: 512
+    output_size: 80
+    spk_embed_dim: !ref <spk_embed_dim>
+    output_type: 'mel'
+    vocab_size: 20480
+    input_frame_rate: !ref <token_frame_rate>
+    only_mask_loss: True
+    token_mel_ratio: !ref <token_mel_ratio>
+    pre_lookahead_len: 3
+    encoder: !new:cosyvoice.transformer.upsample_encoder.UpsampleConformerEncoder
+        output_size: 512
+        attention_heads: 8
+        linear_units: 2048
+        num_blocks: 6
+        dropout_rate: 0.1
+        positional_dropout_rate: 0.1
+        attention_dropout_rate: 0.1
+        normalize_before: True
+        input_layer: 'linear'
+        pos_enc_layer_type: 'rel_pos_espnet'
+        selfattention_layer_type: 'rel_selfattn'
+        input_size: 512
+        upsample_stride: 4
+        use_cnn_module: False
+        macaron_style: False
+        static_chunk_size: !ref <chunk_size>
+    decoder: !new:cosyvoice.flow.flow_matching.CausalConditionalCFM
+        in_channels: 240
+        n_spks: 1
+        spk_emb_dim: 80
+        cfm_params: !new:omegaconf.DictConfig
+            content:
+                sigma_min: 1e-06
+                solver: 'euler'
+                t_scheduler: 'cosine'
+                training_cfg_rate: 0.2
+                inference_cfg_rate: 0.7
+                reg_loss_type: 'l1'
+        estimator: !new:cosyvoice.flow.decoder.CausalConditionalDecoder
+            in_channels: 320
+            out_channels: 80
+            channels: [256]
+            dropout: 0.0
+            attention_head_dim: 64
+            n_blocks: 4
+            num_mid_blocks: 12
+            num_heads: 8
+            act_fn: 'gelu'
+            static_chunk_size: !ref <chunk_size> * <token_mel_ratio>
+            num_decoding_left_chunks: !ref <num_decoding_left_chunks>
+
+hift: !new:cosyvoice.hifigan.generator.HiFTGenerator
+    in_channels: 80
+    base_channels: 512
+    nb_harmonics: 8
+    sampling_rate: !ref <sample_rate>
+    nsf_alpha: 0.1
+    nsf_sigma: 0.003
+    nsf_voiced_threshold: 10
+    upsample_rates: [8, 5, 3]
+    upsample_kernel_sizes: [16, 11, 7]
+    istft_params:
+        n_fft: 16
+        hop_len: 4
+    resblock_kernel_sizes: [3, 7, 11]
+    resblock_dilation_sizes: [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
+    source_resblock_kernel_sizes: [7, 7, 11]
+    source_resblock_dilation_sizes: [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
+    lrelu_slope: 0.1
+    audio_limit: 0.99
+    f0_predictor: !new:cosyvoice.hifigan.f0_predictor.ConvRNNF0Predictor
+        num_class: 1
+        in_channels: 80
+        cond_channels: 512
+
+feat_extractor: !name:matcha.utils.audio.mel_spectrogram
+    n_fft: 1920
+    num_mels: 80
+    sampling_rate: !ref <sample_rate>
+    hop_size: 480
+    win_size: 1920
+    fmin: 0
+    fmax: 8000
+    center: False
+compute_fbank: !name:cosyvoice.dataset.processor.compute_fbank
+    feat_extractor: !ref <feat_extractor>
+    token_mel_ratio: 4
+compute_f0: !name:cosyvoice.dataset.processor.compute_f0
+    sample_rate: !ref <sample_rate>
+    hop_size: 480

flow/flow-chunk-25.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd96a8a6b8b358b97debbf64659a9f7278fcf9d07b9d2fa33088ad991f3a7049
+size 483292901

flow/flow-chunk-5.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0f5640b8ec3bb46bc31320e01d183ce92717d06e62637149423b74fa7e405e68
+size 483292901

flow/flow.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0f5640b8ec3bb46bc31320e01d183ce92717d06e62637149423b74fa7e405e68
+size 483292901

flow/hift.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3386cc880324d4e98e05987b99107f49e40ed925b8ecc87c1f4939432d429879
+size 83390254

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d9a266ec22ad81910a4a2ac4241f8582ded172faba87693e0df32d6570186367
+size 3439132536

modeling_moss_speech_codec.py

@@ -0,0 +1,510 @@
+# coding=utf-8
+# Copyright 2025 OpenMOSS and HuggingFace Inc. teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import os
+import random
+import uuid as uuid_module
+from collections import OrderedDict, defaultdict
+from pathlib import Path
+from typing import List, Optional, Sequence, Tuple, Union
+
+import numpy as np
+import onnxruntime
+from hyperpyyaml import load_hyperpyyaml
+
+import torch
+import torchaudio
+import torchaudio.compliance.kaldi as kaldi
+from safetensors.torch import load_file
+from torch import nn
+from transformers import PreTrainedModel, WhisperFeatureExtractor
+
+from .configuration_moss_speech_codec import MossSpeechCodecConfig
+from .modeling_whisper import WhisperVQEncoder, WhisperVQConfig
+from .utils import extract_speech_token
+
+logger = logging.getLogger(__name__)
+
+def set_seed(seed: int) -> None:
+    if not isinstance(seed, int):
+        raise TypeError("Seed must be an integer.")
+
+    logger.info("Setting random seed to %s", seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+    else:
+        torch.manual_seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    os.environ["TF_CUDNN_DETERMINISTIC"] = "1"
+
+
+def fade_in_out(fade_in_mel, fade_out_mel, window):
+    device = fade_in_mel.device
+    fade_in_mel, fade_out_mel = fade_in_mel.cpu(), fade_out_mel.cpu()
+    mel_overlap_len = int(window.shape[0] / 2)
+    fade_in_mel[..., :mel_overlap_len] = fade_in_mel[..., :mel_overlap_len] * window[:mel_overlap_len] + \
+                                         fade_out_mel[..., -mel_overlap_len:] * window[mel_overlap_len:]
+    return fade_in_mel.to(device)
+
+
+tts_speech_prev = None
+tts_mel_prev = None
+
+
+class AudioDecoder(nn.Module):
+    def __init__(
+        self,
+        config_path: Union[str, os.PathLike],
+        flow_ckpt_path: Union[str, os.PathLike],
+        hift_ckpt_path: Union[str, os.PathLike],
+        campplus_model: Union[str, os.PathLike],
+        device: Union[str, torch.device] = "cuda",
+    ) -> None:
+        super().__init__()
+        self.device = torch.device(device) if isinstance(device, str) else device
+
+        with open(config_path, "r", encoding="utf-8") as config_file:
+            logger.info("Loading decoder configurations from %s", config_path)
+            self.scratch_configs = load_hyperpyyaml(config_file)
+
+        # Load models
+        self.flow = self.scratch_configs["flow"]
+        self.flow.load_state_dict(torch.load(flow_ckpt_path, map_location=self.device), strict=False)
+        self.hift = self.scratch_configs["hift"]
+        self.hift.load_state_dict(torch.load(hift_ckpt_path, map_location=self.device))
+        self.hift = self.hift.eval()
+        self.sample_rate = self.scratch_configs["sample_rate"]
+        self.feat_extractor = self.scratch_configs["feat_extractor"]
+
+        # Move models to the appropriate device
+        self.flow.to(self.device)
+        self.hift.to(self.device)
+        self.mel_overlap_dict = defaultdict(lambda: None)
+        self.hift_cache_dict = defaultdict(lambda: None)
+        self.token_min_hop_len = 2 * self.flow.input_frame_rate
+        self.token_max_hop_len = 4 * self.flow.input_frame_rate
+        self.token_overlap_len = 3.5
+        self.mel_overlap_len = int(self.token_overlap_len / self.flow.input_frame_rate * 24000 / (480 * 2))
+        self.mel_window = np.hamming(2 * self.mel_overlap_len)
+        # hift cache
+        self.mel_cache_len = 1
+        self.source_cache_len = int(self.mel_cache_len * 480)
+        # speech fade in out
+        session_options = onnxruntime.SessionOptions()
+        session_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
+        session_options.intra_op_num_threads = 1
+        self.campplus_session = onnxruntime.InferenceSession(
+            str(campplus_model),
+            sess_options=session_options,
+            providers=["CPUExecutionProvider"],
+        )
+        self.speech_window = np.hamming(2 * self.source_cache_len)
+
+    def token2wav(
+        self,
+        token: torch.Tensor,
+        uuid: str,
+        prompt_token: Optional[torch.Tensor] = None,
+        prompt_feat: Optional[torch.Tensor] = None,
+        embedding: Optional[torch.Tensor] = None,
+        finalize: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        prompt_token = prompt_token if prompt_token is not None else torch.zeros(1, 0, dtype=torch.int32)
+        prompt_feat = prompt_feat if prompt_feat is not None else torch.zeros(1, 0, 80)
+        embedding = embedding if embedding is not None else torch.zeros(1, 192)
+
+        tts_mel = self.flow.inference(
+            token=token.to(self.device),
+            token_len=torch.tensor([token.shape[1]], dtype=torch.int32, device=self.device),
+            prompt_token=prompt_token.to(self.device),
+            prompt_token_len=torch.tensor([prompt_token.shape[1]], dtype=torch.int32, device=self.device),
+            prompt_feat=prompt_feat.to(self.device),
+            prompt_feat_len=torch.tensor([prompt_feat.shape[1]], dtype=torch.int32, device=self.device),
+            embedding=embedding.to(self.device),
+            streaming=False,
+            finalize=finalize,
+        )
+
+        tts_mel = tts_mel[0]
+        if self.mel_overlap_dict[uuid] is not None:
+            tts_mel = fade_in_out(tts_mel, self.mel_overlap_dict[uuid], self.mel_window)
+        # append hift cache
+        if self.hift_cache_dict[uuid] is not None:
+            hift_cache_mel, hift_cache_source = (
+                self.hift_cache_dict[uuid]["mel"],
+                self.hift_cache_dict[uuid]["source"],
+            )
+            tts_mel = torch.cat([hift_cache_mel, tts_mel], dim=2)
+
+        else:
+            hift_cache_source = torch.zeros(1, 1, 0)
+
+        # keep overlap mel and hift cache
+        if not finalize:
+            self.mel_overlap_dict[uuid] = tts_mel[:, :, -self.mel_overlap_len:]
+            tts_mel = tts_mel[:, :, :-self.mel_overlap_len]
+            tts_speech, tts_source = self.hift.inference(speech_feat=tts_mel, cache_source=hift_cache_source)
+
+            self.hift_cache_dict[uuid] = {
+                "mel": tts_mel[:, :, -self.mel_cache_len:],
+                "source": tts_source[:, :, -self.source_cache_len:],
+                "speech": tts_speech[:, -self.source_cache_len:],
+            }
+            tts_speech = tts_speech[:, :-self.source_cache_len]
+
+        else:
+            tts_speech, tts_source = self.hift.inference(speech_feat=tts_mel, cache_source=hift_cache_source)
+            del self.hift_cache_dict[uuid]
+            del self.mel_overlap_dict[uuid]
+        return tts_speech, tts_mel
+
+
+    def offline_inference(self, token: torch.Tensor) -> torch.Tensor:
+        this_uuid = str(uuid_module.uuid1())
+        tts_speech, tts_mel = self.token2wav(token, uuid=this_uuid, finalize=True)
+        return tts_speech.cpu()
+
+    def stream_inference(
+        self,
+        token: torch.Tensor,
+        prompt_token: Optional[torch.Tensor] = None,
+        prompt_feat: Optional[torch.Tensor] = None,
+        embedding: Optional[torch.Tensor] = None,
+        block_size: int = 8,
+    ) -> torch.Tensor:
+        token = token.to(self.device)
+        this_uuid = str(uuid_module.uuid1())
+
+        prompt_tensor = (
+            prompt_token.to(self.device)
+            if prompt_token is not None
+            else torch.zeros(1, 0, dtype=torch.int32, device=self.device)
+        )
+        prompt_speech_feat = (
+            prompt_feat.to(self.device)
+            if prompt_feat is not None
+            else torch.zeros(1, 0, 80, device=self.device)
+        )
+        embedding = embedding.to(self.device) if embedding is not None else torch.zeros(1, 192, device=self.device)
+
+        base_prompt_tensor = prompt_tensor
+        base_prompt_feat = prompt_speech_feat
+
+        tts_speechs: List[torch.Tensor] = []
+        tts_mels: List[torch.Tensor] = []
+        prev_mel: Optional[torch.Tensor] = None
+
+        for idx in range(0, token.size(1), block_size):
+            tts_token = token[:, idx : idx + block_size]
+
+            prompt_tensor_current = base_prompt_tensor
+            prompt_feat_current = base_prompt_feat
+            if prev_mel is not None:
+                prompt_feat_current = torch.cat(
+                    [base_prompt_feat.transpose(1, 2)] + tts_mels,
+                    dim=-1,
+                ).transpose(1, 2)
+                prompt_tensor_current = torch.cat([base_prompt_tensor, token[:, :idx]], dim=-1)
+
+            is_finalize = idx + block_size >= token.size(-1)
+
+            tts_speech, tts_mel = self.token2wav(
+                tts_token,
+                uuid=this_uuid,
+                prompt_token=prompt_tensor_current,
+                prompt_feat=prompt_feat_current,
+                embedding=embedding,
+                finalize=is_finalize,
+            )
+
+            prev_mel = tts_mel
+            tts_speechs.append(tts_speech)
+            tts_mels.append(tts_mel)
+
+        tts_speech = torch.cat(tts_speechs, dim=-1).cpu()
+
+        return tts_speech
+
+    def streaming_inference(
+        self,
+        token: torch.Tensor,
+        prompt_token: Optional[torch.Tensor] = None,
+        prompt_feat: Optional[torch.Tensor] = None,
+        embedding: Optional[torch.Tensor] = None,
+        uuid: Optional[str] = None,
+        prev_mel: Optional[torch.Tensor] = None,
+        prev_token: Optional[torch.Tensor] = None,
+        is_finalize: bool = True,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+        token = token.to(self.device)
+        this_uuid = uuid or str(uuid_module.uuid1())
+
+        prompt_speech_feat = (
+            prompt_feat.to(self.device)
+            if prompt_feat is not None
+            else torch.zeros(1, 0, 80, device=self.device)
+        )
+        flow_prompt_speech_token = (
+            prompt_token.to(self.device)
+            if prompt_token is not None
+            else torch.zeros(1, 0, dtype=torch.int32, device=self.device)
+        )
+        embedding_tensor = (
+            embedding.to(self.device)
+            if embedding is not None
+            else torch.zeros(1, 192, device=self.device)
+        )
+
+        if prev_mel is not None:
+            prompt_speech_feat = prev_mel
+        if prev_token is not None:
+            flow_prompt_speech_token = prev_token
+
+        tts_speech, tts_mel = self.token2wav(
+            token,
+            uuid=this_uuid,
+            prompt_token=flow_prompt_speech_token,
+            prompt_feat=prompt_speech_feat,
+            embedding=embedding_tensor,
+            finalize=is_finalize,
+        )
+
+        if prev_mel is not None:
+            prev_mel = torch.cat([prev_mel, tts_mel], dim=1)
+        else:
+            prev_mel = tts_mel
+        if prev_token is not None:
+            prev_token = torch.cat([prev_token, token], dim=-1)
+        else:
+            prev_token = token
+
+        return tts_speech.cpu(), prev_mel, prev_token
+
+
+class MossSpeechCodec(PreTrainedModel):
+    """MossSpeech codec model (Whisper-VQ encoder + Flow/HiFT decoder).
+
+    Notes
+    - API is designed to be compatible with the existing
+      `MossSpeechProcessor` usages, while adopting a Transformers-style layout
+      similar to HF codec models (`xcodec`, `encodec`).
+    - `encode` accepts raw audio tensors or file paths. It returns a Python
+      list of codec token ids per input sample for backward-compatibility.
+    - `decode` accepts either a 3D LongTensor `(B, 1, T)` or a nested list of
+      token ids, and returns a dict with a list of waveforms under
+      `"syn_wav_list"` (matching current processor expectations).
+    """
+
+    config_class = MossSpeechCodecConfig
+
+    def __init__(
+        self,
+        encoder_weight_path: Union[str, os.PathLike],
+        encoder_config_path: Union[str, os.PathLike],
+        encoder_feature_extractor_path: Union[str, os.PathLike],
+        flow_path: Union[str, os.PathLike],
+    ) -> None:
+        super().__init__(config=MossSpeechCodecConfig())
+
+        # Whisper-VQ encoder
+        self.sample_rate = 16000
+        config = WhisperVQConfig.from_pretrained(str(encoder_config_path))
+        self.whisper_vqmodel = WhisperVQEncoder(config)
+
+        state_dict = load_file(str(encoder_weight_path))
+        new_state_dict: OrderedDict[str, torch.Tensor] = OrderedDict()
+        for k, v in state_dict.items():
+            if k.startswith("encoder."):
+                new_state_dict[k[len("encoder."):]] = v
+        self.whisper_vqmodel.load_state_dict(new_state_dict, strict=False)
+
+        self.feature_extractor = WhisperFeatureExtractor.from_pretrained(
+            str(encoder_feature_extractor_path)
+        )
+
+        # Flow / HiFT decoder stack
+        self.flow_path = str(flow_path)
+        self.audio_decoder = AudioDecoder(
+            config_path=os.path.join(self.flow_path, "config.yaml"),
+            flow_ckpt_path=os.path.join(self.flow_path, "flow.pt"),
+            hift_ckpt_path=os.path.join(self.flow_path, "hift.pt"),
+            campplus_model=os.path.join(self.flow_path, "campplus.onnx"),
+        ).eval()
+
+    @torch.no_grad()
+    def encode(
+        self,
+        inputs: Union[
+            Sequence[Union[str, os.PathLike, Tuple[torch.Tensor, int], torch.Tensor]],
+            torch.Tensor,
+        ],
+        *,
+        sampling_rate: Optional[int] = None,
+        batch_size: int = 128,
+    ) -> List[List[int]]:
+        """Encode audio into codec token ids.
+
+        Accepts one of:
+        - a list of file paths
+        - a list of `(waveform, sr)` tuples
+        - a list of 1D/2D waveforms (sr assumed 16k)
+        - a batched tensor with shape `(B, C, T)` or `(B, T)`
+        """
+        # Normalize to a list the helper can consume
+        if isinstance(inputs, torch.Tensor):
+            if inputs.dim() == 2:
+                inputs = inputs.unsqueeze(1)  # (B, 1, T)
+            if inputs.dim() != 3:
+                raise ValueError("`inputs` must be (B, C, T) when passing a tensor.")
+            sr = sampling_rate or self.sample_rate
+            items: List[Tuple[torch.Tensor, int]] = [
+                (inputs[i].squeeze(0).cpu(), sr) for i in range(inputs.size(0))
+            ]
+        else:
+            items = list(inputs)  # type: ignore[assignment]
+
+        # Use the existing utility (supports file paths, tuples, tensors)
+        audio_tokens: List[List[int]] = extract_speech_token(
+            self.whisper_vqmodel, self.feature_extractor, items, batch_size=batch_size
+        )
+        return audio_tokens
+
+    def _extract_speech_feat(self, speech: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        speech_feat = self.audio_decoder.feat_extractor(speech).squeeze(dim=0).transpose(0, 1)
+        speech_feat = speech_feat.unsqueeze(dim=0)
+        speech_feat_len = torch.tensor([speech_feat.shape[1]], dtype=torch.int32)
+        return speech_feat, speech_feat_len
+
+    def _extract_spk_embedding(self, speech_16k: torch.Tensor) -> torch.Tensor:
+        feat = kaldi.fbank(speech_16k, num_mel_bins=80, dither=0, sample_frequency=16000)
+        feat = feat - feat.mean(dim=0, keepdim=True)
+        embedding = self.audio_decoder.campplus_session.run(
+            None,
+            {self.audio_decoder.campplus_session.get_inputs()[0].name: feat.unsqueeze(0).cpu().numpy()},
+        )[0].flatten().tolist()
+        return torch.tensor([embedding])
+
+    @torch.no_grad()
+    def decode(
+        self,
+        audio_codes: Union[Sequence[Sequence[int]], torch.LongTensor],
+        *,
+        prompt_speech: Optional[Union[str, os.PathLike]] = None,
+        prompt_speech_sample_rate: Optional[int] = None,
+        use_spk_embedding: bool = True,
+        use_prompt_speech: bool = True,
+        finalize: bool = True,
+        device: torch.device = torch.device("cuda"),
+    ) -> dict:
+        """Decode codec token ids back to waveform(s).
+
+        Args
+        - audio_codes: `(B, 1, T)` or Python nested lists per sample.
+        - prompt_speech: path to the enrollment audio used for conditioning.
+        Returns
+        - {"syn_wav_list": List[Tensor(T)]}
+        """
+        if isinstance(audio_codes, torch.Tensor):
+            if audio_codes.dim() == 3 and audio_codes.size(1) == 1:
+                codes_list: List[List[int]] = [
+                    audio_codes[i, 0].detach().cpu().tolist() for i in range(audio_codes.size(0))
+                ]
+            elif audio_codes.dim() == 2:
+                codes_list = [row.detach().cpu().tolist() for row in audio_codes]
+            else:
+                raise ValueError("`audio_codes` must be (B, 1, T) or (B, T) when passing a tensor.")
+        else:
+            codes_list = [list(c) for c in audio_codes]
+
+        if prompt_speech is None or not os.path.exists(str(prompt_speech)):
+            raise ValueError("`prompt_speech` path is required for decoding and must exist.")
+
+        prompt_wav, orig_sr = torchaudio.load(str(prompt_speech))
+        target_sr = self.audio_decoder.sample_rate
+        if orig_sr != target_sr:
+            prompt_wav = torchaudio.transforms.Resample(orig_freq=orig_sr, new_freq=target_sr)(prompt_wav)
+
+        device = device if torch.cuda.is_available() or device.type == "cpu" else torch.device("cpu")
+        speech_token = torch.tensor(self.encode([str(prompt_speech)])[0], device=device).unsqueeze(0)
+        speech_feat, speech_feat_len = self._extract_speech_feat(prompt_wav)
+
+        if target_sr == 24000:
+            token_len = min(int(speech_feat.shape[1] / 4), speech_token.shape[1])
+            speech_feat, speech_feat_len[:] = speech_feat[:, : 4 * token_len], 4 * token_len
+            speech_token, _ = speech_token[:, :token_len], token_len
+
+        prompt_16k = torchaudio.transforms.Resample(orig_freq=target_sr, new_freq=16000)(prompt_wav)
+        embedding = self._extract_spk_embedding(prompt_16k).to(device)
+
+        speech_feat = speech_feat.to(device)
+        speech_feat_len = speech_feat_len.to(device)
+
+        syn_wav_list: List[torch.Tensor] = []
+        for codes in codes_list:
+            codes_t = torch.tensor(codes, device=device).unsqueeze(0)
+            uuid = os.urandom(16).hex()
+
+            kwargs = {"uuid": uuid, "finalize": finalize}
+            if use_prompt_speech:
+                kwargs.update({"prompt_token": speech_token, "prompt_feat": speech_feat})
+            if use_spk_embedding:
+                kwargs.update({"embedding": embedding})
+
+            tts_speech, _ = self.audio_decoder.token2wav(codes_t, **kwargs)
+            syn_wav_list.append(tts_speech.squeeze())
+
+        return {"syn_wav_list": syn_wav_list}
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_dir: Union[str, os.PathLike],
+        *args,
+        **kwargs,
+    ):
+        """Instantiate codec from a directory containing encoder and decoder assets.
+
+        Expected layout:
+        - `model.safetensors` (Whisper VQ encoder weights)
+        - `config.json` (Whisper VQ config)
+        - `preprocessor_config.json` (WhisperFeatureExtractor params)
+        - `flow/{config.yaml, flow.pt, hift.pt, campplus.onnx}`
+        """
+        base = Path(str(model_dir))
+        # Support both layouts:
+        # 1) <base>/{model.safetensors, config.json, preprocessor_config.json, flow/}
+        # 2) <base>/speech_tokenizer/{model.safetensors, ...} and <base>/flow/
+        if (base / "model.safetensors").exists():
+            tokenizer_dir = base
+            flow_dir = base / "flow"
+        else:
+            tokenizer_dir = base / "speech_tokenizer"
+            flow_dir = base / "flow"
+        encoder_weight_path = str(tokenizer_dir / "model.safetensors")
+        encoder_config_path = str(tokenizer_dir / "config.json")
+        encoder_feature_extractor_path = str(tokenizer_dir)
+        flow_path = str(flow_dir)
+
+        return cls(
+            encoder_weight_path=encoder_weight_path,
+            encoder_config_path=encoder_config_path,
+            encoder_feature_extractor_path=encoder_feature_extractor_path,
+            flow_path=flow_path,
+        )

modeling_whisper.py

@@ -0,0 +1,322 @@
+"""Minimal Whisper-VQ encoder for MossSpeech codec.
+
+This file provides only the components used by
+`MossSpeechCodec/modeling_moss_speech_codec.py` during inference:
+ - vector quantization helper
+ - causal conv for streaming
+ - SDPA attention for encoder
+ - WhisperVQEncoderLayer and WhisperVQEncoder
+"""
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import math
+import torch
+from torch import nn
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import EncoderDecoderCache
+from transformers.modeling_outputs import BaseModelOutput
+from transformers.modeling_utils import PreTrainedModel
+
+from .utils import WhisperVQConfig
+
+
+@dataclass
+class QuantizedBaseModelOutput(BaseModelOutput):
+    quantized_token_ids: Optional[torch.LongTensor] = None
+
+
+@dataclass
+class QuantizedBaseModelOutputWithCache(QuantizedBaseModelOutput):
+    past_key_value: Optional[EncoderDecoderCache] = None
+    conv1_cache: Optional[torch.Tensor] = None
+    conv2_cache: Optional[torch.Tensor] = None
+
+
+def vector_quantize(inputs: torch.Tensor, codebook: torch.Tensor):
+    embedding_size = codebook.size(1)
+    inputs_flatten = inputs.reshape(-1, embedding_size)
+    codebook_sqr = torch.sum(codebook ** 2, dim=1)
+    inputs_sqr = torch.sum(inputs_flatten ** 2, dim=1, keepdim=True)
+    distances = torch.addmm(codebook_sqr + inputs_sqr, inputs_flatten, codebook.t(), alpha=-2.0, beta=1.0)
+    _, indices_flatten = torch.min(distances, dim=1)
+    codes_flatten = torch.index_select(codebook, dim=0, index=indices_flatten)
+    return codes_flatten.view_as(inputs), indices_flatten, distances
+
+
+class CausalConv1d(nn.Conv1d):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, **kwargs):
+        super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=0, dilation=dilation, groups=groups, bias=bias, **kwargs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        causal_padding = (self.kernel_size[0] - 1) * self.dilation[0]
+        x = nn.functional.pad(x, (causal_padding, 0))
+        return super().forward(x)
+
+    def forward_causal(self, inp: torch.Tensor, conv_cache: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        k, d = self.kernel_size[0], self.dilation[0]
+        if conv_cache is None:
+            inp_pad = nn.functional.pad(inp, (k - 1, 0))
+        else:
+            inp_pad = torch.cat((conv_cache, inp), dim=-1)
+        out = super().forward(inp_pad)
+        new_cache = inp_pad[:, :, -(k - 1) * d :]
+        return out, new_cache
+
+
+def _prepare_4d_causal_attention_mask_with_cache_position(attention_mask, sequence_length, target_length, cache_position=None, dtype=torch.float32, device=None, min_dtype=None, batch_size=None):
+    if batch_size is None:
+        batch_size = attention_mask.shape[0] if attention_mask is not None else 1
+    if device is None:
+        device = attention_mask.device if attention_mask is not None else None
+    if min_dtype is None:
+        min_dtype = torch.finfo(dtype).min
+    if cache_position is None:
+        target_length = sequence_length
+        sequence_length = target_length
+        if attention_mask is not None:
+            mask_length = attention_mask.shape[-1]
+            target_length = mask_length
+    causal_mask = attention_mask
+    if causal_mask is None:
+        causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)
+        causal_mask = torch.triu(causal_mask, diagonal=1)
+        causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
+        causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
+    else:
+        causal_mask = causal_mask[:, None, None, :].expand(batch_size, 1, sequence_length, target_length).to(dtype)
+        causal_mask = (1.0 - causal_mask) * min_dtype
+    if attention_mask is not None:
+        mask_length = attention_mask.shape[-1]
+        padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
+        padding_mask = padding_mask == 0
+        causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(padding_mask, min_dtype)
+    return causal_mask
+
+
+class WhisperAttention(nn.Module):
+    def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.0, is_causal: bool = False, layer_idx: Optional[int] = None, config: Optional[WhisperVQConfig] = None):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        self.config = config
+        self.is_causal = is_causal
+        self.layer_idx = layer_idx
+        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+
+class WhisperSdpaAttention(WhisperAttention):
+    def forward(self, hidden_states: torch.Tensor, key_value_states: Optional[torch.Tensor] = None, past_key_value: Optional[EncoderDecoderCache] = None, attention_mask: Optional[torch.Tensor] = None, layer_head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.Tensor] = None):
+        bsz, tgt_len, _ = hidden_states.size()
+        query_states = self.q_proj(hidden_states)
+        query_states = query_states.view(bsz, tgt_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+        is_cross_attention = key_value_states is not None
+        current_states = key_value_states if is_cross_attention else hidden_states
+        key_states = self.k_proj(current_states).view(bsz, -1, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+        value_states = self.v_proj(current_states).view(bsz, -1, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+        causal_mask = attention_mask
+        sign = False
+        if self.is_causal and causal_mask is None and tgt_len > 1:
+            if cache_position is not None:
+                dtype, device = query_states.dtype, query_states.device
+                min_dtype = torch.finfo(dtype).min
+                causal_mask = _prepare_4d_causal_attention_mask_with_cache_position(None, query_states.shape[-2], key_states.shape[-2], cache_position=cache_position, dtype=dtype, device=device, min_dtype=min_dtype, batch_size=query_states.shape[0])
+            else:
+                sign = True
+
+        attn_output = torch.nn.functional.scaled_dot_product_attention(query_states, key_states, value_states, attn_mask=causal_mask, dropout_p=self.dropout if self.training else 0.0, is_causal=sign)
+        attn_output = attn_output.transpose(1, 2).reshape(bsz, tgt_len, -1).contiguous()
+        attn_output = self.out_proj(attn_output)
+        return attn_output, None, None
+
+
+WHISPER_ATTENTION_CLASSES = {
+    "sdpa": WhisperSdpaAttention,
+}
+
+
+class WhisperVQEncoderLayer(nn.Module):
+    def __init__(self, config: WhisperVQConfig, is_causal=True, layer_idx=None):
+        super().__init__()
+        self.embed_dim = config.d_model
+        self.kv_cache = True
+        impl = getattr(config, "_attn_implementation", "sdpa")
+        if impl not in WHISPER_ATTENTION_CLASSES:
+            impl = "sdpa"
+        self.self_attn = WHISPER_ATTENTION_CLASSES[impl](embed_dim=self.embed_dim, num_heads=config.encoder_attention_heads, dropout=config.attention_dropout, is_causal=is_causal, layer_idx=layer_idx, config=config)
+        self.is_causal = is_causal
+        if self.is_causal:
+            self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+        self.dropout = config.dropout
+        self.activation_fn = ACT2FN[config.activation_function]
+        self.activation_dropout = config.activation_dropout
+        self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
+        self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
+        self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+    def forward_causal(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, layer_head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False, past_key_value: Optional[EncoderDecoderCache] = None, cache_position: Optional[torch.LongTensor] = None):
+        residual = hidden_states
+        hidden_states = self.self_attn_layer_norm(hidden_states)
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask if not self.is_causal else None, layer_head_mask=layer_head_mask, output_attentions=output_attentions, past_key_value=past_key_value, use_cache=self.kv_cache, cache_position=cache_position)
+        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.final_layer_norm(hidden_states)
+        hidden_states = self.activation_fn(self.fc1(hidden_states))
+        hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+        hidden_states = self.fc2(hidden_states)
+        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+        if output_attentions:
+            outputs += (self_attn_weights,)
+        if self.kv_cache:
+            outputs += (present_key_value,)
+        return outputs, cache_position
+
+
+class WhisperPreTrainedModel(PreTrainedModel):
+    config_class = WhisperVQConfig
+    base_model_prefix = "model"
+    main_input_name = "input_features"
+
+    def _init_weights(self, module):
+        std = self.config.init_std
+        if isinstance(module, (nn.Linear, nn.Conv1d)):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, WhisperVQEncoder):
+            with torch.no_grad():
+                embed_positions = module.embed_positions.weight
+                embed_positions.copy_(sinusoids(*embed_positions.shape))
+
+
+def sinusoids(length: int, channels: int, max_timescale: float = 10000) -> torch.Tensor:
+    if channels % 2 != 0:
+        raise ValueError("channels must be even for sinusoidal positional embeddings")
+    log_timescale_increment = math.log(max_timescale) / (channels // 2 - 1)
+    inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2))
+    scaled_time = torch.arange(length).view(-1, 1) * inv_timescales.view(1, -1)
+    return torch.cat([scaled_time.sin(), scaled_time.cos()], dim=1)
+
+
+class WhisperVQEncoder(WhisperPreTrainedModel):
+    def __init__(self, config: WhisperVQConfig):
+        super().__init__(config)
+        self.config = config
+        self.dropout = config.dropout
+        self.layerdrop = config.encoder_layerdrop
+        embed_dim = config.d_model
+        self.num_mel_bins = config.num_mel_bins
+        self.padding_idx = config.pad_token_id
+        self.max_source_positions = config.max_source_positions
+        if config.encoder_causal_convolution:
+            conv_class = CausalConv1d
+        else:
+            conv_class = nn.Conv1d
+        self.conv1 = conv_class(self.num_mel_bins, embed_dim, kernel_size=3, padding=1)
+        self.conv2 = conv_class(embed_dim, embed_dim, kernel_size=3, stride=2, padding=1)
+        self.embed_positions = nn.Embedding(self.max_source_positions, embed_dim)
+        self.embed_positions.requires_grad_(False)
+        if config.quantize_encoder_only:
+            self.layers = nn.ModuleList([WhisperVQEncoderLayer(config, is_causal=config.encoder_causal_attention or config.quantize_causal_encoder, layer_idx=i) for i in range(config.quantize_position)])
+        else:
+            self.layers = nn.ModuleList([WhisperVQEncoderLayer(config, is_causal=config.encoder_causal_attention or (config.quantize_causal_encoder and layer_id < config.quantize_position), layer_idx=layer_id) for layer_id in range(config.encoder_layers)])
+            self.layer_norm = nn.LayerNorm(config.d_model)
+
+        self.pooling_layer = None
+        if config.pooling_kernel_size is not None:
+            self.pooling_layer = nn.AvgPool1d(kernel_size=config.pooling_kernel_size) if config.pooling_type == "avg" else nn.MaxPool1d(kernel_size=config.pooling_kernel_size)
+
+        self.codebook = None
+        self.embed_positions2 = None
+        if config.quantize_vocab_size is not None:
+            self.codebook = nn.Embedding(config.quantize_vocab_size, config.d_model)
+            pos2_len = self.max_source_positions // max(int(config.pooling_kernel_size or 1), 1)
+            self.embed_positions2 = nn.Embedding(pos2_len, config.d_model)
+            self.embed_positions2.requires_grad_(False)
+
+        self.post_init()
+
+    def forward(self, input_features: torch.FloatTensor, attention_mask: Optional[torch.LongTensor] = None, head_mask: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, past_key_values: Optional[EncoderDecoderCache] = None, cache_position: Optional[torch.LongTensor] = None, quantized_token_ids: Optional[torch.LongTensor] = None, conv1_cache: Optional[torch.Tensor] = None, conv2_cache: Optional[torch.Tensor] = None):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        device = input_features.device
+        if input_features.dim() != 3:
+            raise ValueError("`input_features` should be (batch, feature_size, seq_len)")
+
+        if input_features.shape[-1] % 2 == 1:
+            input_features = nn.functional.pad(input_features, (0, 1))
+        if input_features.shape[1] != self.num_mel_bins:
+            raise ValueError(f"Expected {self.num_mel_bins} mel bins, got {input_features.shape[1]}")
+
+        if isinstance(self.conv1, CausalConv1d):
+            conv1_output, new_conv1_cache = self.conv1.forward_causal(input_features, conv1_cache)
+        else:
+            conv1_output = self.conv1(input_features)
+            new_conv1_cache = None
+        x = nn.functional.gelu(conv1_output)
+        if isinstance(self.conv2, CausalConv1d):
+            conv2_output, new_conv2_cache = self.conv2.forward_causal(x, conv2_cache)
+        else:
+            conv2_output = self.conv2(x)
+            new_conv2_cache = None
+        x = nn.functional.gelu(conv2_output)
+        x = x.permute(0, 2, 1)
+        batch_size, seq_len, _ = x.shape
+        if attention_mask is not None:
+            attention_mask = attention_mask[:, :: self.conv1.stride[0] * self.conv2.stride[0]]
+        if cache_position is None:
+            cache_position = torch.arange(0, seq_len, device=device)
+        embed_pos = self.embed_positions.weight
+        hidden_states = x + embed_pos[cache_position]
+        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+        encoder_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+        if past_key_values is None:
+            past_key_values = EncoderDecoderCache.from_legacy_cache(past_key_values)
+        for idx, layer in enumerate(self.layers):
+            if output_hidden_states:
+                encoder_states = encoder_states + (hidden_states,)
+            layer_outputs, _ = layer.forward_causal(hidden_states, attention_mask=attention_mask, layer_head_mask=(head_mask[idx] if head_mask is not None else None), output_attentions=output_attentions, past_key_value=past_key_values if past_key_values is not None else None, cache_position=cache_position)
+            hidden_states = layer_outputs[0]
+            if output_attentions:
+                all_attentions = all_attentions + (layer_outputs[1],)
+            if idx + 1 == self.config.pooling_position and self.pooling_layer is not None:
+                hs = hidden_states.permute(0, 2, 1)
+                if hs.shape[-1] % self.config.pooling_kernel_size != 0:
+                    hs = nn.functional.pad(hs, (0, self.config.pooling_kernel_size - hs.shape[-1] % self.config.pooling_kernel_size))
+                hidden_states = self.pooling_layer(hs).permute(0, 2, 1)
+            if idx + 1 == self.config.quantize_position and self.codebook is not None:
+                if quantized_token_ids is not None:
+                    hidden_states = self.codebook(quantized_token_ids)
+                else:
+                    hidden_quantized, indices_flat, _ = vector_quantize(hidden_states, self.codebook.weight)
+                    quantized_token_ids = indices_flat.reshape(batch_size, hidden_quantized.shape[1])
+                    hidden_states = hidden_quantized
+                hidden_states = hidden_states + self.embed_positions2.weight[: hidden_states.shape[1]]
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
+        return QuantizedBaseModelOutputWithCache(last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions, quantized_token_ids=quantized_token_ids, past_key_value=past_key_values, conv1_cache=new_conv1_cache, conv2_cache=new_conv2_cache)

preprocessor_config.json

@@ -0,0 +1,14 @@
+{
+  "chunk_length": 30,
+  "feature_extractor_type": "WhisperFeatureExtractor",
+  "feature_size": 128,
+  "hop_length": 160,
+  "n_fft": 400,
+  "n_samples": 480000,
+  "nb_max_frames": 3000,
+  "padding_side": "right",
+  "padding_value": 0.0,
+  "processor_class": "WhisperProcessor",
+  "return_attention_mask": false,
+  "sampling_rate": 16000
+}