Add mimo-audio model

bonsai/models/mimo_audio/README.md

@@ -0,0 +1,94 @@
+# MiMo-Audio in JAX
+
+This directory contains a pure JAX implementation of the [MiMo-Audio multimodal language model](https://github.com/XiaomiMiMo/MiMo), using the [Flax NNX](https://flax.readthedocs.io/en/v0.8.3/experimental/nnx/index.html) API.
+
+MiMo-Audio is a unified speech-text model that supports:
+- **Text-to-Speech (TTS)**: Generate natural speech from text
+- **Speech-to-Text (ASR)**: Transcribe speech to text
+- **Speech-to-Speech**: Direct speech translation and conversion
+
+## Model Configuration Support Status
+
+| Model Name | Config Support Status |
+| :--- | :--- |
+| **Main Models** | |
+| [MiMo-Audio-7B-Base](https://huggingface.co/XiaomiMiMo/MiMo-Audio-7B-Base) | **✅ Supported** |
+| [MiMo-Audio-7B-Instruct](https://huggingface.co/XiaomiMiMo/MiMo-Audio-7B-Instruct) | **✅ Supported** |
+| **Audio Tokenizer** | |
+| [MiMo-Audio-Tokenizer](https://huggingface.co/XiaomiMiMo/MiMo-Audio-Tokenizer) | **✅ Supported** |
+
+
+## Running this model
+
+Run MiMo-Audio inference with a minimal example:
+
+```sh
+python3 -m bonsai.models.mimo_audio.test.run_model
+```
+## Model Architecture
+
+MiMo-Audio consists of three main components:
+
+### 1. Main Transformer (Qwen2-based)
+- **Layers**: 36
+- **Hidden size**: 4096
+- **Attention heads**: 32 (8 KV heads)
+- **Intermediate size**: 11008
+- **Max position embeddings**: 8192
+
+### 2. Local Transformer (for audio generation)
+- **Layers**: 16
+- **Hidden size**: 1024
+- **Attention heads**: 64
+- **FFN dimension**: 4096
+
+### 3. Input Local Transformer (for audio encoding)
+- **Layers**: 6
+- **Hidden size**: 1024
+- **Attention heads**: 64
+- **Bidirectional attention**: Yes
+
+### 4. Audio Tokenizer
+- **Encoder layers**: 32
+- **Decoder layers**: 32
+- **Quantizers**: 20
+- **Sampling rate**: 24000 Hz
+- **Audio channels**: 8 (used for multi-codebook representation)
+
+## Special Tokens
+
+MiMo-Audio uses special tokens for controlling speech generation:
+
+- `<|sostm|>` (151648): Start of speech/stream
+- `<|eostm|>` (151649): End of speech/stream
+- `<|sosp|>` (151646): Start of speech
+- `<|eosp|>` (151647): End of speech
+- `<|empty|>` (151645): Empty token (indicates audio generation)
+- `<|eot|>` (151643): End of turn
+
+## Input Format
+
+MiMo-Audio uses an interleaved format where each group contains:
+- 1 text token (repeated `group_size` times)
+- 8 audio channel tokens (one per channel, repeated `group_size` times)
+
+Shape: `[batch, audio_channels + 1, num_groups * group_size]`
+
+For text-only input (TTS), audio channels are filled with channel-specific empty IDs.
+
+## Audio Processing
+
+### Encoding (Speech → Tokens)
+1. Waveform → Mel spectrogram
+2. Mel spectrogram → Encoder
+3. Encoder → Quantizer → Audio tokens (8 channels)
+
+### Decoding (Tokens → Speech)
+1. Audio tokens → Decoder
+2. Decoder → Vocoder
+3. Vocoder → Waveform (24kHz)
+
+## References
+
+- [MiMo-Audio Paper](https://github.com/XiaomiMiMo/MiMo-Audio/blob/main/MiMo-Audio-Technical-Report.pdf)
+- [Official Implementation](https://github.com/XiaomiMiMo/MiMo-Audio)

bonsai/models/mimo_audio/__init__.py

@@ -0,0 +1,19 @@
+from bonsai.models.mimo_audio.modeling import (
+    MiMoAudioConfig,
+    MiMoAudioArguments,
+    FlaxMiMoAudioForCausalLM,
+)
+from bonsai.models.mimo_audio.mimo_audio_tokenizer import (
+    FlaxMiMoAudioTokenizer,
+    MiMoAudioTokenizerConfig,
+    MelSpectrogram,
+)
+
+__all__ = [
+    "MiMoAudioConfig",
+    "MiMoAudioArguments",
+    "FlaxMiMoAudioForCausalLM",
+    "FlaxMiMoAudioTokenizer",
+    "MiMoAudioTokenizerConfig",
+    "MelSpectrogram",
+]

bonsai/models/mimo_audio/mimo_audio_configuration.py

@@ -0,0 +1,110 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+from bonsai.models.qwen3.modeling import ShardingCfg
+
+if TYPE_CHECKING:
+    from bonsai.models.qwen2.modeling import ModelConfig as Qwen2Config
+
+
+@dataclass
+class MiMoAudioConfig:
+    vocab_size: int = 151680
+    hidden_size: int = 4096
+    num_hidden_layers: int = 36
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8
+    intermediate_size: int = 11008
+    max_position_embeddings: int = 8192
+    rope_theta: int = 640000
+    head_dim: int = 128
+
+    group_size: int = 4
+    audio_channels: int = 8
+
+    local_dim: int = 1024
+    local_layers: int = 16
+    local_attn_heads: int = 64
+    local_ffn_dim: int = 4096
+    local_attn_dropout: float = 0.1
+
+    input_local_layers: int = 6
+    input_local_dim: int = 1024
+    input_full_attention: bool = True
+
+    shd_cfg: ShardingCfg = ShardingCfg.no_sharding()
+
+    @classmethod
+    def with_sharding(cls, **kwargs):
+        kwargs["shd_cfg"] = ShardingCfg.default()
+        return cls(**kwargs)
+
+    def create_qwen2_config(self) -> "Qwen2Config":
+        from bonsai.models.qwen2.modeling import ModelConfig as Qwen2Config
+
+        return Qwen2Config(
+            num_layers=36,
+            vocab_size=151680,
+            emb_dim=4096,
+            mlp_dim=11008,
+            num_heads=32,
+            head_dim=128,
+            num_kv_heads=8,
+            rope_theta=640000,
+            norm_eps=1e-6,
+            tie_word_embeddings=False,
+            shd_cfg=self.shd_cfg,
+        )
+
+    def create_local_qwen2_config(self) -> "Qwen2Config":
+        from bonsai.models.qwen2.modeling import ModelConfig as Qwen2Config
+
+        return Qwen2Config(
+            num_layers=16,
+            vocab_size=151680,
+            emb_dim=1024,
+            mlp_dim=4096,
+            num_heads=64,
+            head_dim=16,  # 1024 // 64 = 16
+            num_kv_heads=64,
+            rope_theta=640000,
+            norm_eps=1e-6,
+            tie_word_embeddings=False,
+            shd_cfg=self.shd_cfg,
+        )
+
+    def create_input_local_qwen2_config(self) -> "Qwen2Config":
+        from bonsai.models.qwen2.modeling import ModelConfig as Qwen2Config
+
+        return Qwen2Config(
+            num_layers=6,
+            vocab_size=151680,
+            emb_dim=1024,
+            mlp_dim=4096,
+            num_heads=64,
+            head_dim=16,
+            num_kv_heads=64,
+            rope_theta=640000,
+            norm_eps=1e-6,
+            tie_word_embeddings=False,
+            use_causal_mask=False,
+            shd_cfg=self.shd_cfg,
+        )
+
+
+@dataclass
+class MiMoAudioArguments:
+    model_name_or_path: str
+    sosp_idx: int
+    eosp_idx: int
+    sostm_idx: int
+    eostm_idx: int
+    eot_idx: int
+    empty_idx: int
+
+
+@dataclass
+class MiMoSamplerConfig:
+    do_sample: bool = True
+    temperature: float = 1.0
+    top_k: int = 50
+    top_p: float = 0.95