diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 867ba2af8..a1f4b2148 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -1125,6 +1125,14 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: max_position_embeddings=getattr(ec, "max_position_embeddings", 8000), num_quantizers=getattr(ec, "num_quantizers", 32), num_semantic_quantizers=getattr(ec, "num_semantic_quantizers", 1), + audio_channels=getattr(ec, "audio_channels", 1), + num_filters=getattr(ec, "num_filters", 64), + num_residual_layers=getattr(ec, "num_residual_layers", 1), + kernel_size=getattr(ec, "kernel_size", 7), + last_kernel_size=getattr(ec, "last_kernel_size", 3), + residual_kernel_size=getattr(ec, "residual_kernel_size", 3), + compress=getattr(ec, "compress", 2), + upsampling_ratios=list(getattr(ec, "upsampling_ratios", [8, 6, 5, 4])), ) # Model dtype diff --git a/src/mobius/_configs/_base_test.py b/src/mobius/_configs/_base_test.py index 810f6a555..51f5ec1c0 100644 --- a/src/mobius/_configs/_base_test.py +++ b/src/mobius/_configs/_base_test.py @@ -79,3 +79,108 @@ def test_list_intermediate_size_collapses_to_first_element(): out = ArchitectureConfig.from_transformers(cfg) assert out.intermediate_size == 8192 assert isinstance(out.intermediate_size, int) + + +def _codec_hf_config(**encoder_overrides): + """A Qwen3-TTS-Tokenizer-style HF config with nested encoder/decoder.""" + encoder = { + "codebook_dim": 256, + "codebook_size": 2048, + "hidden_size": 512, + "intermediate_size": 2048, + "num_hidden_layers": 8, + "num_attention_heads": 8, + "num_key_value_heads": 8, + "head_dim": 64, + "num_quantizers": 32, + "num_semantic_quantizers": 1, + "audio_channels": 1, + "num_filters": 64, + "num_residual_layers": 1, + "kernel_size": 7, + "last_kernel_size": 3, + "residual_kernel_size": 3, + "compress": 2, + "upsampling_ratios": [8, 6, 5, 4], + } + encoder.update(encoder_overrides) + return _FakeHFConfig( + model_type="qwen3_tts_tokenizer_12hz", + hidden_size=512, + num_attention_heads=8, + num_hidden_layers=8, + vocab_size=2048, + decoder_config={"hidden_size": 512, "codebook_dim": 512}, + encoder_config=encoder, + ) + + +def test_codec_encoder_conv_fields_extracted_from_nested_config(): + """Nested ``encoder_config`` values drive the derived conv stack. + + These fields are read with ``getattr`` off a nested config, so a wrong + key or default would silently fall back to the checkpoint defaults and + build the wrong architecture. + """ + out = ArchitectureConfig.from_transformers(_codec_hf_config()) + + enc = out.codec_encoder + assert enc is not None + assert enc.audio_channels == 1 + assert enc.num_filters == 64 + assert enc.num_residual_layers == 1 + assert enc.kernel_size == 7 + assert enc.last_kernel_size == 3 + assert enc.residual_kernel_size == 3 + assert enc.compress == 2 + assert enc.upsampling_ratios == [8, 6, 5, 4] + + +def test_codec_encoder_conv_fields_honor_non_default_values(): + """Non-default nested values must survive extraction unchanged.""" + out = ArchitectureConfig.from_transformers( + _codec_hf_config( + hidden_size=64, + audio_channels=2, + num_filters=8, + num_residual_layers=3, + kernel_size=5, + last_kernel_size=1, + residual_kernel_size=7, + compress=4, + upsampling_ratios=[4, 2], + ) + ) + + enc = out.codec_encoder + assert enc is not None + assert enc.hidden_size == 64 + assert enc.audio_channels == 2 + assert enc.num_filters == 8 + assert enc.num_residual_layers == 3 + assert enc.kernel_size == 5 + assert enc.last_kernel_size == 1 + assert enc.residual_kernel_size == 7 + assert enc.compress == 4 + assert enc.upsampling_ratios == [4, 2] + + +def test_codec_encoder_conv_fields_fall_back_to_checkpoint_defaults(): + """A config omitting the conv fields still yields the real architecture.""" + out = ArchitectureConfig.from_transformers( + _FakeHFConfig( + model_type="qwen3_tts_tokenizer_12hz", + hidden_size=512, + num_attention_heads=8, + num_hidden_layers=8, + vocab_size=2048, + decoder_config={"hidden_size": 512}, + encoder_config={"hidden_size": 512}, + ) + ) + + enc = out.codec_encoder + assert enc is not None + assert enc.num_filters == 64 + assert enc.upsampling_ratios == [8, 6, 5, 4] + assert enc.num_residual_layers == 1 diff --git a/src/mobius/_configs/_sub_configs.py b/src/mobius/_configs/_sub_configs.py index 76c375636..a4bca1539 100644 --- a/src/mobius/_configs/_sub_configs.py +++ b/src/mobius/_configs/_sub_configs.py @@ -148,6 +148,18 @@ class CodecEncoderConfig: max_position_embeddings: int = 8000 num_quantizers: int = 32 num_semantic_quantizers: int = 1 + # Mimi convolutional encoder shape (HF ``MimiConfig`` defaults). The conv + # stack is derived from these: 1 leading conv, then per upsampling ratio + # ``num_residual_layers`` residual blocks + ELU + a strided conv that + # doubles the channel count, then a trailing ELU + conv to ``hidden_size``. + audio_channels: int = 1 + num_filters: int = 64 + num_residual_layers: int = 1 + kernel_size: int = 7 + last_kernel_size: int = 3 + residual_kernel_size: int = 3 + compress: int = 2 + upsampling_ratios: list[int] = dataclasses.field(default_factory=lambda: [8, 6, 5, 4]) @dataclasses.dataclass diff --git a/src/mobius/models/qwen3_tts_tokenizer.py b/src/mobius/models/qwen3_tts_tokenizer.py index 9a41f90d3..4d1f27f85 100644 --- a/src/mobius/models/qwen3_tts_tokenizer.py +++ b/src/mobius/models/qwen3_tts_tokenizer.py @@ -40,6 +40,8 @@ ) if TYPE_CHECKING: + from collections.abc import Sequence + import onnx_ir as ir import torch @@ -223,20 +225,22 @@ def __init__(self, config: ArchitectureConfig): head_dim = enc.head_dim if enc else 64 intermediate = enc.intermediate_size if enc else 2048 - # Conv encoder: series of Conv1d + optional residual blocks - # Based on MimiModel encoder with num_filters=64, ratios=[8,6,5,4] - # Layer structure from actual weights: - # 0: Conv1d(1->64, k=7) - # 1: ResBlock(64) [conv1d k=3 + conv1d k=1] - # 3: Conv1d(64->128, k=8, stride=4) - # 4: ResBlock(128) - # 6: Conv1d(128->256, k=10, stride=5) - # 7: ResBlock(256) - # 9: Conv1d(256->512, k=12, stride=6) - # 10: ResBlock(512) - # 12: Conv1d(512->1024, k=16, stride=8) - # 14: Conv1d(1024->512, k=3) - self.encoder = _MimiConvEncoder() + # Conv encoder: series of Conv1d + residual blocks, derived from + # config exactly like HF ``MimiEncoder.__init__``. With the + # checkpoint's defaults (num_filters=64, upsampling_ratios=[8,6,5,4], + # num_residual_layers=1, kernel_size=7, last_kernel_size=3) this + # yields 1->64->128->256->512->1024->hidden_size. + self.encoder = _MimiConvEncoder( + hidden_size=hidden_size, + audio_channels=enc.audio_channels if enc else 1, + num_filters=enc.num_filters if enc else 64, + num_residual_layers=enc.num_residual_layers if enc else 1, + kernel_size=enc.kernel_size if enc else 7, + last_kernel_size=enc.last_kernel_size if enc else 3, + residual_kernel_size=enc.residual_kernel_size if enc else 3, + compress=enc.compress if enc else 2, + upsampling_ratios=tuple(enc.upsampling_ratios) if enc else (8, 6, 5, 4), + ) # Transformer self.encoder_transformer = CodecEncoderTransformerModel( @@ -248,7 +252,7 @@ def __init__(self, config: ArchitectureConfig): head_dim=head_dim, ) - # Downsample: Conv1d(512->512, k=4, stride=2) + # Downsample: Conv1d(hidden_size->hidden_size, k=4, stride=2) self.downsample = _DownsampleConv(hidden_size, hidden_size, 4, 2) # Quantizer (encoder-side): uses argmin for encoding @@ -264,11 +268,11 @@ def forward(self, op: OpBuilder, waveform: ir.Value): Returns: codes: (B, 16, T) int64 audio codes. """ - # 1. Conv encoder: (B, 1, samples) -> (B, 512, T') + # 1. Conv encoder: (B, 1, samples) -> (B, hidden_size, T') hidden = self.encoder(op, waveform) # 2. Transformer: channels-last - # (B, 512, T') -> (B, T', 512) + # (B, hidden_size, T') -> (B, T', hidden_size) hidden = op.Transpose(hidden, perm=[0, 2, 1]) seq_len = op.Shape(hidden, start=1, end=2) position_ids = op.Unsqueeze( @@ -280,25 +284,33 @@ def forward(self, op: OpBuilder, waveform: ir.Value): [0], ) hidden = self.encoder_transformer(op, hidden, position_ids) - # (B, T', 512) -> (B, 512, T') + # (B, T', hidden_size) -> (B, hidden_size, T') hidden = op.Transpose(hidden, perm=[0, 2, 1]) - # 3. Downsample: (B, 512, T') -> (B, 512, T'/2) + # 3. Downsample: (B, hidden_size, T') -> (B, hidden_size, T'/2) hidden = self.downsample(op, hidden) - # 4. Quantize: (B, 512, T) -> (B, 16, T) + # 4. Quantize: (B, hidden_size, T) -> (B, 16, T) codes = self.quantizer(op, hidden) return codes class _MimiConvEncoder(nn.Module): - """Mimi-style convolutional encoder. + """Mimi-style convolutional encoder, derived from config. + + Mirrors HF ``MimiEncoder.__init__``: one leading conv, then for each + upsampling ratio (in reverse) ``num_residual_layers`` residual blocks, + an ELU and a strided conv that doubles the channel count, then a + trailing ELU and a conv down to ``hidden_size``. - Progressively downsamples and increases channels: - 1->64->128->256->512->1024->512 + The parameterless ELU modules occupy list slots so that the resulting + parameter names line up with the checkpoint's + ``encoder.encoder.layers.*`` numbering. With the checkpoint defaults + (``num_filters=64``, ``upsampling_ratios=[8, 6, 5, 4]``, + ``num_residual_layers=1``, ``kernel_size=7``, ``last_kernel_size=3``) + the stack is 1->64->128->256->512->1024->512 and numbers as: - Weight structure from HF (encoder.encoder.layers.*): 0: Conv1d(1->64, k=7) 1: ResBlock(64): block.1 Conv1d(64->32,k=3), block.3 Conv1d(32->64,k=1) 3: Conv1d(64->128, k=8, stride=4) @@ -309,37 +321,69 @@ class _MimiConvEncoder(nn.Module): 10: ResBlock(512) 12: Conv1d(512->1024, k=16, stride=8) 14: Conv1d(1024->512, k=3) + + Parameters: + hidden_size: Output channels of the final conv. + audio_channels: Input channels of the waveform (1 for mono). + num_filters: Channel count after the leading conv. + num_residual_layers: Residual blocks per downsampling stage. + kernel_size: Kernel of the leading conv. + last_kernel_size: Kernel of the final conv. + residual_kernel_size: Kernel of the first conv inside a residual block. + compress: Channel-reduction factor inside a residual block. + upsampling_ratios: Downsampling strides, applied in reverse order. """ - def __init__(self): + def __init__( + self, + hidden_size: int = 512, + audio_channels: int = 1, + num_filters: int = 64, + num_residual_layers: int = 1, + kernel_size: int = 7, + last_kernel_size: int = 3, + residual_kernel_size: int = 3, + compress: int = 2, + upsampling_ratios: Sequence[int] = (8, 6, 5, 4), + ): super().__init__() - # Non-sequential layer indices to match HF weight names - self.layers = nn.Sequential( - _EncoderConvLayer(1, 64, 7, 1), # 0 - _EncoderResBlock(64), # 1 - _ELUModule(), # 2: ELU before downsample - _EncoderConvLayer(64, 128, 8, 4), # 3 - _EncoderResBlock(128), # 4 - _ELUModule(), # 5 - _EncoderConvLayer(128, 256, 10, 5), # 6 - _EncoderResBlock(256), # 7 - _ELUModule(), # 8 - _EncoderConvLayer(256, 512, 12, 6), # 9 - _EncoderResBlock(512), # 10 - _ELUModule(), # 11 - _EncoderConvLayer(512, 1024, 16, 8), # 12 - _ELUModule(), # 13 - _EncoderConvLayer(1024, 512, 3, 1), # 14 + layers: list[nn.Module] = [ + _EncoderConvLayer(audio_channels, num_filters, kernel_size, 1), + ] + scaling = 1 + # Encoder downsamples in reverse of the decoder's upsampling order: + # ratios [8, 6, 5, 4] -> strides 4, 5, 6, 8. + for ratio in reversed(upsampling_ratios): + current_scale = scaling * num_filters + for _ in range(num_residual_layers): + layers.append( + _EncoderResBlock( + current_scale, + kernel=residual_kernel_size, + compress=compress, + ) + ) + layers.append(_ELUModule()) + # Strided conv doubles channels; kernel is 2x the stride. + layers.append( + _EncoderConvLayer(current_scale, current_scale * 2, ratio * 2, ratio) + ) + scaling *= 2 + layers.append(_ELUModule()) + layers.append( + _EncoderConvLayer(scaling * num_filters, hidden_size, last_kernel_size, 1) ) + self.layers = nn.Sequential(*layers) + self.output_channels = hidden_size def forward(self, op: OpBuilder, x: ir.Value): """Encode waveform through conv layers. Args: - x: (B, 1, audio_samples). + x: (B, audio_channels, audio_samples). Returns: - (B, 512, T). + (B, hidden_size, T). """ return self.layers(op, x) @@ -397,14 +441,19 @@ class _EncoderResBlock(nn.Module): HF weight structure: block.1.conv (dilated), block.3.conv (pointwise). block.0 and block.2 are ELU activations. + + Parameters: + dim: Input/output channels. + kernel: Kernel of the first (channel-reducing) conv. + compress: Channel-reduction factor for the block's inner width. """ - def __init__(self, dim: int): + def __init__(self, dim: int, kernel: int = 3, compress: int = 2): super().__init__() - half = dim // 2 + half = dim // compress self.block = nn.Sequential( _ELUModule(), # 0: ELU - _EncoderConvLayer(dim, half, 3, 1), # 1: dilated conv + _EncoderConvLayer(dim, half, kernel, 1), # 1: dilated conv _ELUModule(), # 2: ELU _EncoderConvLayer(half, dim, 1, 1), # 3: pointwise conv ) @@ -456,28 +505,32 @@ class _EncoderSplitRVQ(nn.Module): def __init__(self, config: ArchitectureConfig): super().__init__() enc = config.codec_encoder - codebook_dim = enc.codebook_dim if enc else 512 + codebook_dim = enc.codebook_dim if enc else 256 codebook_size = enc.codebook_size if enc else 2048 num_quantizers = enc.num_quantizers if enc else 32 num_semantic = enc.num_semantic_quantizers if enc else 1 + hidden_size = enc.hidden_size if enc else 512 - dim = codebook_dim // 2 + # The RVQ consumes the encoder's hidden features and projects them + # down to the codebook dimension, matching HF ``MimiResidualVector + # Quantizer``: input_proj is Conv1d(hidden_size -> codebook_dim, k=1). + dim = codebook_dim # Semantic quantizer self.semantic_residual_vector_quantizer = _EncoderRVQ( num_quantizers=num_semantic, codebook_size=codebook_size, dim=dim, - input_dim=codebook_dim, - output_dim=codebook_dim, + input_dim=hidden_size, + output_dim=hidden_size, ) # Acoustic quantizer self.acoustic_residual_vector_quantizer = _EncoderRVQ( num_quantizers=num_quantizers - num_semantic, codebook_size=codebook_size, dim=dim, - input_dim=codebook_dim, - output_dim=codebook_dim, + input_dim=hidden_size, + output_dim=hidden_size, ) self._num_semantic = num_semantic self._num_valid = config.num_quantizers if hasattr(config, "num_quantizers") else 16 @@ -486,7 +539,7 @@ def forward(self, op: OpBuilder, hidden: ir.Value): """Encode features to discrete codes. Args: - hidden: (B, codebook_dim, T) continuous features. + hidden: (B, hidden_size, T) continuous features. Returns: codes: (B, num_valid_quantizers, T) int64. diff --git a/src/mobius/tasks/_codec.py b/src/mobius/tasks/_codec.py index ace4f3afb..9be29ca3b 100644 --- a/src/mobius/tasks/_codec.py +++ b/src/mobius/tasks/_codec.py @@ -81,16 +81,21 @@ def _build_encoder( """Build encoder: waveform → codes. Inputs: - waveform: (B, 1, audio_samples) float32 + waveform: (B, audio_channels, audio_samples) float32 Outputs: codes: (B, num_quantizers, T) int64 """ batch = ir.SymbolicDim("batch") audio_len = ir.SymbolicDim("audio_length") + # The first conv is sized from codec_encoder.audio_channels, so the + # graph input must declare the same channel count. + audio_channels = config.codec_encoder.audio_channels if config.codec_encoder else 1 graph, builder = _make_graph(name="encoder") waveform = builder.input( - "waveform", dtype=ir.DataType.FLOAT, shape=[batch, 1, audio_len] + "waveform", + dtype=ir.DataType.FLOAT, + shape=[batch, audio_channels, audio_len], ) codes = encoder(builder.op, waveform) diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index cc25e140e..a73c29cee 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -4811,6 +4811,10 @@ def _codec_config(): max_position_embeddings=128, num_quantizers=8, num_semantic_quantizers=1, + # Narrow conv stack (1->4->8->16->32->64->32) but the same + # depth as the real checkpoint, so `layers.*` numbering and + # weight names match what the real model relies on. + num_filters=4, ), ) @@ -4864,6 +4868,164 @@ def test_encoder_io(self): output_names = {out.name for out in encoder.graph.outputs} assert "codes" in output_names + @staticmethod + def _conv_encoder_weights(encoder_config): + """Return {param name: shape} for the conv stack of an encoder config.""" + from mobius.models.qwen3_tts_tokenizer import Qwen3TTSCodecEncoderModel + + config = ArchitectureConfig( + hidden_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=8, + intermediate_size=64, + vocab_size=256, + max_position_embeddings=128, + codec_encoder=encoder_config, + ) + module = Qwen3TTSCodecEncoderModel(config) + return { + name: tuple(param.shape) + for name, param in module.named_parameters() + if name.startswith("encoder.layers") and name.endswith("conv.weight") + } + + def test_default_config_conv_stack_matches_checkpoint(self): + """Default config must reproduce the real checkpoint's conv stack. + + Both names and shapes are asserted: the ``layers.*`` indices are + derived from ``upsampling_ratios``/``num_residual_layers``, so a + drift in indexing would only surface at weight-load time. + """ + from mobius._configs import CodecEncoderConfig + + weights = self._conv_encoder_weights(CodecEncoderConfig()) + + # Matches Qwen/Qwen3-TTS-Tokenizer-12Hz encoder.encoder.layers.* + assert weights == { + "encoder.layers.0.conv.weight": (64, 1, 7), + "encoder.layers.1.block.1.conv.weight": (32, 64, 3), + "encoder.layers.1.block.3.conv.weight": (64, 32, 1), + "encoder.layers.3.conv.weight": (128, 64, 8), + "encoder.layers.4.block.1.conv.weight": (64, 128, 3), + "encoder.layers.4.block.3.conv.weight": (128, 64, 1), + "encoder.layers.6.conv.weight": (256, 128, 10), + "encoder.layers.7.block.1.conv.weight": (128, 256, 3), + "encoder.layers.7.block.3.conv.weight": (256, 128, 1), + "encoder.layers.9.conv.weight": (512, 256, 12), + "encoder.layers.10.block.1.conv.weight": (256, 512, 3), + "encoder.layers.10.block.3.conv.weight": (512, 256, 1), + "encoder.layers.12.conv.weight": (1024, 512, 16), + "encoder.layers.14.conv.weight": (512, 1024, 3), + } + + def test_tiny_config_keeps_checkpoint_layer_names(self): + """The tiny test config must keep the checkpoint's layer numbering.""" + from mobius._configs import CodecEncoderConfig + + tiny = self._conv_encoder_weights(self._codec_config().codec_encoder) + default = self._conv_encoder_weights(CodecEncoderConfig()) + + assert set(tiny) == set(default) + # Only the widths shrink: 1 -> 4 -> 8 -> 16 -> 32 -> 64 -> 32. + assert tiny["encoder.layers.0.conv.weight"] == (4, 1, 7) + assert tiny["encoder.layers.12.conv.weight"] == (64, 32, 16) + assert tiny["encoder.layers.14.conv.weight"] == (32, 64, 3) + + def test_conv_stack_indices_follow_config(self): + """Layer indices are derived from ratios and residual-layer count.""" + from mobius._configs import CodecEncoderConfig + + # 2 ratios -> final conv lands at layers.8 + two_ratios = self._conv_encoder_weights( + CodecEncoderConfig(hidden_size=8, num_filters=2, upsampling_ratios=[4, 2]) + ) + assert "encoder.layers.8.conv.weight" in two_ratios + assert two_ratios["encoder.layers.8.conv.weight"] == (8, 8, 3) + + # 4 ratios with 2 residual layers each -> final conv at layers.18 + deep = self._conv_encoder_weights( + CodecEncoderConfig(hidden_size=8, num_filters=2, num_residual_layers=2) + ) + assert "encoder.layers.18.conv.weight" in deep + + def test_hidden_size_drives_conv_output_width(self): + """The final conv width follows ``hidden_size`` (no hardcoded 512). + + Regression guard: a hardcoded final width silently produced a + malformed graph (LayerNormalization over a mismatched width). + """ + from mobius._configs import CodecEncoderConfig + + weights = self._conv_encoder_weights(CodecEncoderConfig(hidden_size=32, num_filters=4)) + assert weights["encoder.layers.14.conv.weight"][0] == 32 + + def test_default_config_rvq_projections_match_checkpoint(self): + """Encoder RVQ projections must match the real checkpoint's shapes. + + The RVQ consumes ``hidden_size``-wide features and projects to + ``codebook_dim``, mirroring HF ``MimiResidualVectorQuantizer``. + Deriving these from ``codebook_dim`` alone produced projections + that could not load the checkpoint's weights at all. + """ + from mobius._configs import CodecEncoderConfig + from mobius.models.qwen3_tts_tokenizer import Qwen3TTSCodecEncoderModel + + config = ArchitectureConfig( + hidden_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=8, + intermediate_size=64, + vocab_size=256, + max_position_embeddings=128, + codec_encoder=CodecEncoderConfig(), + ) + module = Qwen3TTSCodecEncoderModel(config) + shapes = { + name: tuple(param.shape) + for name, param in module.named_parameters() + if "_proj.weight" in name and "quantizer" in name + } + + # Matches Qwen/Qwen3-TTS-Tokenizer-12Hz encoder.quantizer.* weights. + for prefix in ( + "quantizer.semantic_residual_vector_quantizer", + "quantizer.acoustic_residual_vector_quantizer", + ): + assert shapes[f"{prefix}.input_proj.weight"] == (256, 512, 1) + assert shapes[f"{prefix}.output_proj.weight"] == (512, 256, 1) + + # Codebooks are codebook_dim-wide, matching codebook.embed_sum. + codebooks = { + tuple(param.shape) + for name, param in module.named_parameters() + if name.endswith("codebook.embedding") + } + assert codebooks == {(2048, 256)} + + def test_encoder_input_declares_configured_audio_channels(self): + """The graph input channel count must match the first conv. + + ``audio_channels`` sizes ``encoder.layers.0.conv``, so a task that + always declared a mono input would feed a 1-channel tensor into a + conv expecting more. + """ + import dataclasses + + from mobius.models.qwen3_tts_tokenizer import Qwen3TTSTokenizerV2Model + from mobius.tasks import CodecTask + + config = self._codec_config() + config.codec_encoder = dataclasses.replace(config.codec_encoder, audio_channels=2) + module = Qwen3TTSTokenizerV2Model(config) + pkg = build_from_module(module, config, task=CodecTask()) + + waveform = next(inp for inp in pkg["encoder"].graph.inputs if inp.name == "waveform") + assert waveform.shape[1] == 2 + def test_registry_lookup(self): """Verify qwen3_tts_tokenizer_12hz is registered with codec task.""" model_cls = registry.get("qwen3_tts_tokenizer_12hz")