diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index 5213d2cec..581e48253 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -226,6 +226,14 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: guidance_scale = getattr(args, "guidance_scale", None) if guidance_scale is not None and args.runtime != "onnx-genai": raise SystemExit("Error: --guidance-scale can only be used with --runtime onnx-genai.") + input_sampling_rate = getattr(args, "input_sampling_rate", None) + bwe_sampling_rate = getattr(args, "bwe_sampling_rate", None) + for option, value in ( + ("--input-sample-rate", input_sampling_rate), + ("--bwe-sample-rate", bwe_sampling_rate), + ): + if value is not None and value <= 0: + raise SystemExit(f"Error: {option} must be a positive integer.") # Validate static-cache + --task compatibility. if args.static_cache and args.task is not None: @@ -304,6 +312,13 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: task = CausalLMTask(paged_cache=True) trust_remote_code = args.trust_remote_code revision = args.revision + if args.model == "nvidia/RE-USE" and revision is None: + # Pin every Hub probe, including the early Diffusers detector. A + # mutable model_index.json on Hub main must not reroute this checkpoint + # before the bespoke builder applies its immutable default. + from mobius.models.reuse import REUSE_REVISION + + revision = REUSE_REVISION output_dir = args.output_dir os.makedirs(output_dir, exist_ok=True) dtype_override = resolve_dtype(args.dtype) @@ -318,6 +333,11 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: if args.model and not args.config and not args.text_only: pipeline_index = _load_diffusers_pipeline_index(args.model, revision=revision) if pipeline_index is not None: + if input_sampling_rate is not None or bwe_sampling_rate is not None: + raise SystemExit( + "Error: --input-sample-rate and --bwe-sample-rate are only " + "supported for RE-USE speech-enhancement checkpoints." + ) print( f"Detected diffusers pipeline: {pipeline_index.get('_class_name', 'Unknown')}" ) @@ -363,6 +383,31 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: import transformers config_path = args.config + from mobius.models.reuse import _is_reuse_checkpoint, build_reuse + + if _is_reuse_checkpoint(config_path): + if task not in (None, "speech-enhancement"): + from mobius.tasks import SpeechEnhancementTask + + if not isinstance(task, SpeechEnhancementTask): + raise SystemExit( + "Error: RE-USE checkpoints only support --task speech-enhancement." + ) + pkg = build_reuse( + config_path, + dtype=dtype_override, + execution_provider=execution_provider, + load_weights=load_weights, + input_sampling_rate=input_sampling_rate, + bwe_sampling_rate=bwe_sampling_rate, + ) + _save_package(pkg, output_dir, args, optimize, component_filter) + return + if input_sampling_rate is not None or bwe_sampling_rate is not None: + raise SystemExit( + "Error: --input-sample-rate and --bwe-sample-rate are only " + "supported for RE-USE speech-enhancement checkpoints." + ) try: hf_config = transformers.AutoConfig.from_pretrained( config_path, trust_remote_code=trust_remote_code @@ -501,6 +546,8 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: glm_full_attention=args.glm_full_attention, export_paged_attention=export_paged_attention, keep_quantized=keep_quantized, + input_sampling_rate=input_sampling_rate, + bwe_sampling_rate=bwe_sampling_rate, ) _save_package(pkg, output_dir, args, optimize, component_filter) @@ -1372,6 +1419,29 @@ def build_parser() -> argparse.ArgumentParser: "for unguided generation." ), ) + sample_rate_group = build_parser.add_mutually_exclusive_group() + sample_rate_group.add_argument( + "--input-sample-rate", + dest="input_sampling_rate", + type=int, + default=None, + metavar="HZ", + help=( + "Build RE-USE for a known native input rate with static FFT geometry. " + "Omit to preserve native-rate dynamic geometry." + ), + ) + sample_rate_group.add_argument( + "--bwe-sample-rate", + dest="bwe_sampling_rate", + type=int, + default=None, + metavar="HZ", + help=( + "Build RE-USE with NVIDIA BWE semantics: resample input audio to this " + "target rate and use consistently scaled FFT geometry." + ), + ) build_parser.add_argument( "--kv-cache-scale-file", dest="kv_cache_scale_file", diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index b20f2a66a..51e0e5d4e 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -150,7 +150,9 @@ Qwen35VL3ModelCausalLMModel, Qwen35VLTextModel, QwenCausalLMModel, + ReUseConfig, RND1Model, + SEMambaSpeechEnhancementModel, SmallThinkerGGUFCausalLMModel, SmolLM3CausalLMModel, SortformerDiarizationModel, @@ -1066,6 +1068,16 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: ), "fastconformer_rnnt": ModelRegistration(EncDecRNNTModel, task="fastconformer-rnnt"), "sortformer": ModelRegistration(SortformerDiarizationModel, task="diarization"), + "reuse": ModelRegistration( + SEMambaSpeechEnhancementModel, + task="speech-enhancement", + config_class=ReUseConfig, + ), + "semamba": ModelRegistration( + SEMambaSpeechEnhancementModel, + task="speech-enhancement", + config_class=ReUseConfig, + ), } diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index af7edbe6e..65389842f 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -116,6 +116,8 @@ "RmsNorm2d", "ScaleFreeRMSNorm", "SelectiveScan", + "SequenceMambaBlock", + "SequenceSelectiveScan", "SiLU", "Siglip2NaFlexVisionEmbeddings", "Siglip2NaFlexVisionModel", @@ -280,6 +282,7 @@ from mobius.components._lora import LoRALinear from mobius.components._mamba_block import Mamba2Block as Mamba2Block from mobius.components._mamba_block import MambaBlock as MambaBlock +from mobius.components._mamba_block import SequenceMambaBlock from mobius.components._mimo_minimax_vision import ( DualTemporalPatchEmbedding as MiMoDualTemporalPatchEmbedding, ) @@ -467,6 +470,7 @@ ) from mobius.components._ssm import ( SelectiveScan, + SequenceSelectiveScan, ) from mobius.components._vision import ( PatchEmbedding, diff --git a/src/mobius/components/_conv.py b/src/mobius/components/_conv.py index 0f11187e1..e8cabed97 100644 --- a/src/mobius/components/_conv.py +++ b/src/mobius/components/_conv.py @@ -77,6 +77,7 @@ def __init__( stride: int | tuple[int, int] = 1, padding: int | tuple[int, int, int, int] = 0, groups: int = 1, + dilation: int | tuple[int, int] = 1, ): super().__init__() kernel_h, kernel_w = _pair(kernel_size, "kernel_size") @@ -87,6 +88,7 @@ def __init__( self._strides = _pair(stride, "stride") self._pads = _resolve_pads(padding) self._groups = groups + self._dilations = _pair(dilation, "dilation") def forward(self, op: OpBuilder, x: ir.Value): return op.Conv( @@ -96,6 +98,7 @@ def forward(self, op: OpBuilder, x: ir.Value): kernel_shape=list(self._kernel_size), strides=list(self._strides), pads=self._pads, + dilations=list(self._dilations), group=self._groups, ) diff --git a/src/mobius/components/_conv_test.py b/src/mobius/components/_conv_test.py index 65e209646..cf5b340d1 100644 --- a/src/mobius/components/_conv_test.py +++ b/src/mobius/components/_conv_test.py @@ -113,6 +113,33 @@ def test_asymmetric_kernel_stride_and_padding(self): assert conv._strides == (1, 4) assert conv._pads == [0, 2, 0, 2] + def test_dilation_defaults_to_one(self): + conv = Conv2d(3, 16, kernel_size=3, padding=1) + assert conv._dilations == (1, 1) + + def test_dilation_accepts_int_and_pair(self): + assert Conv2d(3, 16, kernel_size=3, dilation=2)._dilations == (2, 2) + assert Conv2d(3, 16, kernel_size=3, dilation=(4, 1))._dilations == (4, 1) + + def test_positional_call_is_unchanged_by_dilation(self): + """``dilation`` was appended, so existing positional calls still work.""" + conv = Conv2d(3, 16, 3, 2, 1, 1) + assert list(conv.weight.shape) == [16, 3, 3, 3] + assert conv._strides == (2, 2) + assert conv._pads == [1, 1, 1, 1] + assert conv._groups == 1 + assert conv._dilations == (1, 1) + + def test_dilation_reaches_the_conv_node(self): + conv = Conv2d(4, 4, kernel_size=(3, 3), padding=(2, 1, 2, 1), dilation=(2, 1)) + test_builder, op, graph = create_test_builder() + x = create_test_input(test_builder, "x", [1, 4, 8, 8]) + + conv(op, x) + + node = next(n for n in graph if n.op_type == "Conv") + assert list(node.attributes["dilations"].value) == [2, 1] + class TestConv2dNoBias: """Tests for 2D convolution without bias.""" diff --git a/src/mobius/components/_mamba_block.py b/src/mobius/components/_mamba_block.py index 74c365c94..32c9c0417 100644 --- a/src/mobius/components/_mamba_block.py +++ b/src/mobius/components/_mamba_block.py @@ -26,7 +26,7 @@ from mobius.components._common import INT64_MAX, Linear from mobius.components._rms_norm import GatedRMSNorm, PostGatedRMSNorm -from mobius.components._ssm import SelectiveScan +from mobius.components._ssm import SelectiveScan, SequenceSelectiveScan class _DepthwiseConv1d(nn.Module): @@ -173,6 +173,84 @@ def forward( return output, new_conv_state, new_ssm_state +class SequenceMambaBlock(nn.Module): + """Mamba1 layer applied to a whole sequence, with no carried state. + + Same parameters and math as :class:`MambaBlock` (input projection → + causal Conv1D → selective scan → SiLU gate → output projection), but the + entire ``(batch, seq_len, d_model)`` sequence is processed in one call + and both the conv and SSM states start at zero. This mirrors the + reference ``mamba_ssm.Mamba`` module's offline forward pass and is what + non-autoregressive backbones (e.g. speech enhancement) need. + + Args: + d_model: Model hidden dimension. + d_inner: Expanded inner dimension (typically ``expand * d_model``). + d_state: SSM state dimension (typically 16). + dt_rank: Rank of the SSM time-step projection. + Defaults to ``ceil(d_model / 16)`` (Mamba convention). + conv_kernel: Causal Conv1D kernel size (typically 4). + """ + + def __init__( + self, + d_model: int, + d_inner: int, + d_state: int = 16, + dt_rank: int | None = None, + conv_kernel: int = 4, + ): + super().__init__() + self.d_model = d_model + self.d_inner = d_inner + self.d_state = d_state + self.conv_kernel = conv_kernel + self.dt_rank = dt_rank if dt_rank is not None else -(-d_model // 16) + + self.in_proj = Linear(d_model, 2 * d_inner, bias=False) + self.conv1d = _DepthwiseConv1d(d_inner, conv_kernel, bias=True) + self.ssm = SequenceSelectiveScan(d_inner, d_state, self.dt_rank) + self.out_proj = Linear(d_inner, d_model, bias=False) + + def forward(self, op: OpBuilder, hidden_states: ir.Value): + """Run the Mamba layer over a full sequence. + + Args: + op: ONNX op builder. + hidden_states: (batch, seq_len, d_model). + + Returns: + (batch, seq_len, d_model) output. + """ + # --- Step 1: Input projection, split into SSM branch and gate --- + projected = self.in_proj(op, hidden_states) # (batch, seq_len, 2*d_inner) + x_branch, z_gate = op.Split( + projected, + [self.d_inner, self.d_inner], + axis=-1, + _outputs=2, + ) + + # --- Step 2: Causal depthwise Conv1D over the sequence --- + # (batch, seq_len, d_inner) → (batch, d_inner, seq_len) for Conv. + x_t = op.Transpose(x_branch, perm=[0, 2, 1]) + # Left-pad by (conv_kernel - 1) so output position t only sees inputs + # up to t; the conv itself uses pads=[0, 0] and keeps seq_len. + x_padded = op.Pad( + x_t, + op.Constant(value_ints=[0, 0, self.conv_kernel - 1, 0, 0, 0]), + ) + conv_out = op.Swish(self.conv1d(op, x_padded)) # (batch, d_inner, seq_len) + + # --- Step 3: Selective scan over the sequence --- + x_ssm = op.Transpose(conv_out, perm=[0, 2, 1]) # (batch, seq_len, d_inner) + y = self.ssm(op, x_ssm) # (batch, seq_len, d_inner) + + # --- Step 4: Gating and output projection --- + gated = op.Mul(y, op.Swish(z_gate)) + return self.out_proj(op, gated) + + # ===================================================================== # Mamba2 block using LinearAttention # ===================================================================== diff --git a/src/mobius/components/_mamba_block_test.py b/src/mobius/components/_mamba_block_test.py index 55f8b0124..3c5d948f8 100644 --- a/src/mobius/components/_mamba_block_test.py +++ b/src/mobius/components/_mamba_block_test.py @@ -5,12 +5,18 @@ from __future__ import annotations +import numpy as np +import onnx_ir as ir +import pytest +from onnxscript import GraphBuilder + +from mobius._constants import OPSET_VERSION from mobius._testing import ( count_op_type, create_test_builder, create_test_input, ) -from mobius.components._mamba_block import MambaBlock +from mobius.components._mamba_block import MambaBlock, SequenceMambaBlock class TestMambaBlock: @@ -131,3 +137,116 @@ def test_custom_conv_kernel(self): block = MambaBlock(d_model=32, d_inner=64, conv_kernel=8) assert list(block.conv1d.weight.shape) == [64, 1, 8] assert block.conv_kernel == 8 + + +class TestSequenceMambaBlock: + """Tests for the full-sequence (stateless) Mamba1 block.""" + + def test_parameters_match_decode_block(self): + """Parameter names and shapes are identical to MambaBlock's.""" + + def spec(module): + return {name: list(p.shape) for name, p in module.named_parameters()} + + assert spec(SequenceMambaBlock(d_model=32, d_inner=64, d_state=8)) == spec( + MambaBlock(d_model=32, d_inner=64, d_state=8) + ) + + def test_forward_needs_no_state(self): + """The whole sequence is consumed in one call, with no carried state.""" + block = SequenceMambaBlock(d_model=32, d_inner=64, d_state=8) + test_builder, op, graph = create_test_builder() + hidden = create_test_input(test_builder, "hidden_states", [2, 7, 32]) + + out = block(op, hidden) + + assert out is not None + assert count_op_type(graph, "Scan") == 1 + # in_proj, x_proj, dt_proj, out_proj + assert count_op_type(graph, "MatMul") >= 4 + + def test_conv_is_left_padded_for_causality(self): + """Conv1D is left-padded by kernel_size - 1 so it stays causal.""" + block = SequenceMambaBlock(d_model=16, d_inner=32, d_state=4, conv_kernel=4) + test_builder, op, graph = create_test_builder() + hidden = create_test_input(test_builder, "hidden_states", [1, 5, 16]) + + block(op, hidden) + + pad = next(node for node in graph if node.op_type == "Pad") + pads = pad.inputs[1].const_value.numpy().tolist() + # (batch, d_inner, seq): 3 begin values then 3 end values. + assert pads == [0, 0, 3, 0, 0, 0] + + def test_matches_reference_recurrence(self): + """ONNX output matches an explicit PyTorch selective-scan reference.""" + pytest.importorskip("onnxruntime") + import onnxruntime as ort + import torch + import torch.nn.functional as torch_f + + d_model, d_state, conv_kernel, expand = 8, 4, 4, 2 + d_inner = d_model * expand + dt_rank = -(-d_model // 16) + batch, seq_len = 2, 6 + + torch.manual_seed(0) + weights = { + "in_proj.weight": torch.randn(2 * d_inner, d_model) * 0.2, + "conv1d.weight": torch.randn(d_inner, 1, conv_kernel) * 0.2, + "conv1d.bias": torch.randn(d_inner) * 0.2, + "ssm.x_proj.weight": torch.randn(dt_rank + 2 * d_state, d_inner) * 0.2, + "ssm.dt_proj.weight": torch.randn(d_inner, dt_rank) * 0.2, + "ssm.dt_proj.bias": torch.randn(d_inner) * 0.2, + "ssm.A_log": torch.randn(d_inner, d_state) * 0.2, + "ssm.D": torch.randn(d_inner) * 0.2, + "out_proj.weight": torch.randn(d_model, d_inner) * 0.2, + } + + def reference(u): + """``mamba_ssm.Mamba`` forward, written out with plain torch ops.""" + x, z = (u @ weights["in_proj.weight"].t()).chunk(2, dim=-1) + x = torch_f.conv1d( + torch_f.pad(x.transpose(1, 2), (conv_kernel - 1, 0)), + weights["conv1d.weight"], + weights["conv1d.bias"], + groups=d_inner, + ) + x = torch_f.silu(x).transpose(1, 2) + + x_dbl = x @ weights["ssm.x_proj.weight"].t() + dt_raw, b_mat, c_mat = torch.split(x_dbl, [dt_rank, d_state, d_state], dim=-1) + dt = torch_f.softplus( + dt_raw @ weights["ssm.dt_proj.weight"].t() + weights["ssm.dt_proj.bias"] + ) + + a_neg = -torch.exp(weights["ssm.A_log"]) + state = torch.zeros(u.shape[0], d_inner, d_state) + outputs = [] + for t in range(u.shape[1]): + dt_col = dt[:, t].unsqueeze(-1) + decay = torch.exp(dt_col * a_neg.unsqueeze(0)) + update = dt_col * x[:, t].unsqueeze(-1) * b_mat[:, t].unsqueeze(1) + state = decay * state + update + outputs.append((state * c_mat[:, t].unsqueeze(1)).sum(-1)) + y = torch.stack(outputs, dim=1) + weights["ssm.D"] * x + return (y * torch_f.silu(z)) @ weights["out_proj.weight"].t() + + graph = ir.Graph([], [], nodes=[], name="g", opset_imports={"": OPSET_VERSION}) + builder = GraphBuilder(graph) + hidden = builder.input( + "hidden_states", dtype=ir.DataType.FLOAT, shape=["batch", "seq", d_model] + ) + block = SequenceMambaBlock(d_model, d_inner, d_state, dt_rank, conv_kernel) + for name, param in block.named_parameters(): + param.const_value = ir.tensor(weights[name].numpy()) + builder.add_output(block(builder.op, hidden), "y") + + session = ort.InferenceSession( + ir.to_proto(ir.Model(graph, ir_version=11)).SerializeToString(), + providers=["CPUExecutionProvider"], + ) + hidden_states = torch.randn(batch, seq_len, d_model) + (got,) = session.run(None, {"hidden_states": hidden_states.numpy()}) + + np.testing.assert_allclose(got, reference(hidden_states).detach().numpy(), atol=1e-5) diff --git a/src/mobius/components/_ssm.py b/src/mobius/components/_ssm.py index 3039ff6c8..f2a207c5c 100644 --- a/src/mobius/components/_ssm.py +++ b/src/mobius/components/_ssm.py @@ -32,10 +32,11 @@ from onnxscript import OpBuilder, nn from mobius.components._common import Linear +from mobius.components._scan_utils import create_body_graph, rename_subgraph_values -class SelectiveScan(nn.Module): - """Core selective scan (S6) operation. +class _SelectiveScanBase(nn.Module): + """Parameters and projections shared by step-wise and sequence scans. Args: d_inner: Expanded hidden dimension (``expand * d_model``). @@ -80,6 +81,10 @@ def _project_ssm_params(self, op: OpBuilder, x_db): ) return dt_raw, b_mat, c_mat + +class SelectiveScan(_SelectiveScanBase): + """Core single-step selective scan (S6) operation.""" + def forward( self, op: OpBuilder, @@ -243,6 +248,165 @@ def forward( ) +class SequenceSelectiveScan(_SelectiveScanBase): + """Selective scan (S6) over a whole sequence, with no carried state. + + Unlike :class:`SelectiveScan`, which advances one token per call and + threads ``ssm_state`` through the graph for autoregressive decode, this + variant consumes the full ``(batch, seq_len, d_inner)`` activation in a + single ONNX ``Scan`` and starts from a zero state. It is the right + building block for non-autoregressive encoders (speech enhancement, + audio/vision Mamba backbones) that see the entire sequence at once. + + The recurrence body computes ``dA`` and ``dBx`` per step rather than + materialising them for all timesteps up front. Precomputing them would + require two ``(batch, seq_len, d_inner, d_state)`` tensors, which for a + typical audio backbone is orders of magnitude larger than the + ``(batch, seq_len, d_inner)`` activations themselves. + + Parameters are identical to :class:`SelectiveScan` (``x_proj``, + ``dt_proj``, ``A_log``, ``D``), so the two share checkpoints. + + Precision: the recurrence runs in float32 regardless of model dtype, + matching the reference CUDA ``selective_scan_fn`` which accumulates the + state in float32. + """ + + def forward( + self, + op: OpBuilder, + x: ir.Value, + ): + """Run the selective scan over an entire sequence. + + Args: + op: ONNX op builder. + x: (batch, seq_len, d_inner) — input after conv1d + activation. + + Returns: + y: (batch, seq_len, d_inner) — scan output including the ``D`` + skip connection. + """ + # --- Project x to get dt, B, C for every timestep --- + # x_db: (batch, seq_len, dt_rank + 2*d_state) + x_db = self.x_proj(op, x) + dt_raw, b_mat, c_mat = self._project_ssm_params(op, x_db) + + # dt: (batch, seq_len, d_inner). Upcast before softplus so the whole + # recurrence (softplus/exp/state accumulation) stays in float32. + dt = op.Softplus(op.Cast(self.dt_proj(op, dt_raw), to=ir.DataType.FLOAT)) + b_f32 = op.Cast(b_mat, to=ir.DataType.FLOAT) # (batch, seq_len, d_state) + c_f32 = op.Cast(c_mat, to=ir.DataType.FLOAT) # (batch, seq_len, d_state) + x_f32 = op.Cast(x, to=ir.DataType.FLOAT) # (batch, seq_len, d_inner) + + # a_neg = -exp(A_log): (d_inner, d_state), the continuous-time decay. + a_neg = op.Neg(op.Exp(op.Cast(self.A_log, to=ir.DataType.FLOAT))) + + # Zero initial state: (batch, d_inner, d_state). The batch dim is + # dynamic, so build the shape from the activation at run time. + batch_dim = op.Shape(x, start=0, end=1) + state_shape = op.Concat( + batch_dim, + op.Constant(value_ints=[self.d_inner, self.d_state]), + axis=0, + ) + # ONNX's ConstantOfShape defaults to a float32 zero. + initial_state = op.ConstantOfShape(state_shape) + + body = _build_sequence_scan_body() + + # Transpose to time-major for the Scan. `scan_input_axes` could name + # axis 1 directly and skip these, but several execution providers + # (e.g. the MLX plugin EP) only claim Scan nodes that iterate axis 0, + # and falling back to CPU for the recurrence costs far more than the + # transposes do. + dt_t = op.Transpose(dt, perm=[1, 0, 2]) # (seq_len, batch, d_inner) + b_t = op.Transpose(b_f32, perm=[1, 0, 2]) # (seq_len, batch, d_state) + c_t = op.Transpose(c_f32, perm=[1, 0, 2]) # (seq_len, batch, d_state) + x_t = op.Transpose(x_f32, perm=[1, 0, 2]) # (seq_len, batch, d_inner) + + # `a_neg` rides along as a pass-through carry so the body never has to + # reach into the enclosing graph's scope for an initializer. + _final_state, _a_out, y = op.Scan( + initial_state, + a_neg, + dt_t, + b_t, + c_t, + x_t, + body=body, + num_scan_inputs=4, + _outputs=3, + ) + # y: (seq_len, batch, d_inner) — restore batch-major. + y = op.Transpose(y, perm=[1, 0, 2]) # (batch, seq_len, d_inner) + + # --- Skip connection: y += D * x --- + y = op.Add(y, op.Mul(op.Cast(self.D, to=ir.DataType.FLOAT), x_f32)) + + return op.CastLike(y, x) + + +def _build_sequence_scan_body() -> ir.Graph: + """Build the ``Scan`` body for one timestep of the sequence selective scan. + + Body inputs (in order): + 1. ``state``: (batch, d_inner, d_state) — carry + 2. ``a_neg``: (d_inner, d_state) — carry, passed through unchanged + 3. ``dt_t``: (batch, d_inner) — scan input + 4. ``b_t``: (batch, d_state) — scan input + 5. ``c_t``: (batch, d_state) — scan input + 6. ``x_t``: (batch, d_inner) — scan input + + Body outputs: + 1. ``new_state``: (batch, d_inner, d_state) — carry + 2. ``a_neg_out``: (d_inner, d_state) — carry + 3. ``y_t``: (batch, d_inner) — scan output + """ + f32 = ir.TensorType(ir.DataType.FLOAT) + state_in = ir.Value(name="ssm_state", type=f32) + a_in = ir.Value(name="a_neg", type=f32) + dt_t = ir.Value(name="dt_t", type=f32) + b_t = ir.Value(name="b_t", type=f32) + c_t = ir.Value(name="c_t", type=f32) + x_t = ir.Value(name="x_t", type=f32) + + body_graph, body_builder = create_body_graph( + state_inputs=[state_in, a_in], + scan_inputs=[dt_t, b_t, c_t, x_t], + name="selective_scan_step", + ) + bop = body_builder.op + + # dt_col: (batch, d_inner, 1) — broadcasts over the d_state axis. + dt_col = bop.Unsqueeze(dt_t, [-1]) + + # Discretised decay dA = exp(dt * A): (batch, d_inner, d_state). + da = bop.Exp(bop.Mul(dt_col, bop.Unsqueeze(a_in, [0]))) + + # Input contribution dBx = (dt * x) ⊗ B: (batch, d_inner, d_state). + dbx = bop.Mul( + bop.Mul(dt_col, bop.Unsqueeze(x_t, [-1])), + bop.Unsqueeze(b_t, [1]), + ) + + # State update: h = dA * h_prev + dBx. + new_state = bop.Add(bop.Mul(da, state_in), dbx) + + # Readout: y = C · h, summing over d_state → (batch, d_inner). + y_t = bop.ReduceSum(bop.Mul(new_state, bop.Unsqueeze(c_t, [1])), [-1], keepdims=False) + + # The pass-through carry must be a distinct value from the graph input. + a_out = bop.Identity(a_in) + + new_state.name = "ssm_state_out" + a_out.name = "a_neg_out" + y_t.name = "y_t" + body_graph.outputs.extend([new_state, a_out, y_t]) + rename_subgraph_values(body_graph, "seq_scan_") + return body_graph + + class Mamba2Scan(nn.Module): """Multi-head selective scan for Mamba2/SSD architecture. diff --git a/src/mobius/components/_ssm_test.py b/src/mobius/components/_ssm_test.py index 59919502e..e71e96385 100644 --- a/src/mobius/components/_ssm_test.py +++ b/src/mobius/components/_ssm_test.py @@ -12,7 +12,7 @@ create_test_builder, create_test_input, ) -from mobius.components._ssm import SelectiveScan +from mobius.components._ssm import SelectiveScan, SequenceSelectiveScan class TestSelectiveScan: @@ -99,3 +99,87 @@ def test_state_input_dtype(self): y, _new_state = ssm(op, x, state) assert y is not None + + +class TestSequenceSelectiveScan: + """Tests for the full-sequence selective scan.""" + + def test_shares_parameters_with_decode_variant(self): + """Parameter names and shapes match the single-token SelectiveScan.""" + decode = SelectiveScan(d_inner=64, d_state=16, dt_rank=4) + sequence = SequenceSelectiveScan(d_inner=64, d_state=16, dt_rank=4) + + def spec(module): + return {name: list(p.shape) for name, p in module.named_parameters()} + + assert spec(sequence) == spec(decode) + + def test_forward_builds_graph(self): + """Forward pass constructs a graph with no carried state.""" + ssm = SequenceSelectiveScan(d_inner=32, d_state=8, dt_rank=2) + test_builder, op, graph = create_test_builder() + x = create_test_input(test_builder, "x", [2, 5, 32]) + + y = ssm(op, x) + + assert y is not None + # The recurrence is expressed as a single Scan over the sequence. + assert count_op_type(graph, "Scan") == 1 + + def test_scan_body_carries_state_and_decay(self): + """The Scan body takes 2 carries + 4 per-step inputs and emits 3 values.""" + ssm = SequenceSelectiveScan(d_inner=16, d_state=4, dt_rank=2) + test_builder, op, graph = create_test_builder() + x = create_test_input(test_builder, "x", [1, 3, 16]) + + ssm(op, x) + + scan = next(node for node in graph if node.op_type == "Scan") + assert scan.attributes["num_scan_inputs"].value == 4 + body = scan.attributes["body"].value + # ssm_state + a_neg carries, then dt/B/C/x scan inputs. + assert len(body.inputs) == 6 + # ssm_state_out + a_neg_out carries, then the per-step output. + assert len(body.outputs) == 3 + + def test_scan_iterates_over_axis_zero(self): + """Inputs are made time-major so the Scan iterates axis 0. + + Several execution providers (e.g. the MLX plugin EP) only claim + Scan nodes with the default axes; naming axis 1 instead would send + the whole recurrence back to CPU. + """ + ssm = SequenceSelectiveScan(d_inner=16, d_state=4, dt_rank=2) + test_builder, op, graph = create_test_builder() + x = create_test_input(test_builder, "x", [1, 3, 16]) + + ssm(op, x) + + scan = next(node for node in graph if node.op_type == "Scan") + assert "scan_input_axes" not in scan.attributes + assert "scan_output_axes" not in scan.attributes + # Each scan input arrives through a batch<->time transpose. + for scan_input in scan.inputs[2:]: + producer = scan_input.producer() + assert producer.op_type == "Transpose" + assert list(producer.attributes["perm"].value) == [1, 0, 2] + + def test_recurrence_runs_in_float32(self): + """The scan body stays in float32 even for a float16 activation.""" + ssm = SequenceSelectiveScan(d_inner=16, d_state=4, dt_rank=2) + test_builder, op, graph = create_test_builder() + x = create_test_input(test_builder, "x", [1, 3, 16], ir.DataType.FLOAT16) + + ssm(op, x) + + scan = next(node for node in graph if node.op_type == "Scan") + for inp in scan.inputs: + producer = inp.producer() + if producer is not None and producer.op_type == "Cast": + assert producer.attributes["to"].value == ir.DataType.FLOAT + else: + # ConstantOfShape / Neg / Softplus already operate on float32. + assert inp.dtype in (ir.DataType.FLOAT, None) + + # The scan result is handed back in the activation's dtype. + assert count_op_type(graph, "CastLike") >= 1 diff --git a/src/mobius/integrations/onnx_genai/__init__.py b/src/mobius/integrations/onnx_genai/__init__.py index fc82c6e59..acecf4a76 100644 --- a/src/mobius/integrations/onnx_genai/__init__.py +++ b/src/mobius/integrations/onnx_genai/__init__.py @@ -92,6 +92,7 @@ build_image_edit_workflow_metadata, build_language_diffusion_pipeline_metadata, build_speculative_workflow_metadata, + build_speech_enhancement_workflow_metadata, build_tts_workflow_metadata, build_video_diffusion_workflow_metadata, build_vlm_workflow_metadata, @@ -103,6 +104,7 @@ write_image_edit_workflow_metadata, write_language_diffusion_workflow_metadata, write_speculative_workflow_metadata, + write_speech_enhancement_workflow_metadata, write_tts_workflow_metadata, write_video_diffusion_workflow_metadata, write_vlm_workflow_metadata, @@ -138,6 +140,7 @@ "build_speculative_workflow_metadata", "build_speech_to_text_pipeline_metadata", "build_shared_state_pixel_flow_workflow_metadata", + "build_speech_enhancement_workflow_metadata", "build_tokenizer_facts", "build_tts_workflow_metadata", "build_video_diffusion_workflow_metadata", @@ -165,6 +168,7 @@ "write_speculative_workflow_metadata", "write_speech_to_text_pipeline_metadata", "write_shared_state_pixel_flow_workflow_metadata", + "write_speech_enhancement_workflow_metadata", "write_tts_workflow_metadata", "write_video_diffusion_workflow_metadata", "write_vlm_workflow_metadata", diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 9d3b2d09b..56335a446 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -33,6 +33,7 @@ ) from mobius.integrations.onnx_genai.workflow_metadata import ( HierarchicalAudioWorkflowConfig, + _validate_reuse_rate_selection, write_audio_codec_workflow_metadata, write_ctc_asr_workflow_metadata, write_decoder_workflow_metadata, @@ -42,11 +43,13 @@ write_image_edit_workflow_metadata, write_language_diffusion_workflow_metadata, write_speculative_workflow_metadata, + write_speech_enhancement_workflow_metadata, write_speech_to_text_workflow_metadata, write_tts_workflow_metadata, write_video_diffusion_workflow_metadata, write_vlm_workflow_metadata, ) +from mobius.models.reuse import ReUseConfig _LOGGER = logging.getLogger(__name__) @@ -687,6 +690,31 @@ def _looks_like_encoder_embedding(pkg: Any) -> bool: return not any(str(name).startswith("past_key_values") for name in inputs) +def _looks_like_speech_enhancement(pkg: Any) -> bool: + """Detect a spectral speech-enhancement package. + + The signal is structural: a single ``model`` component that consumes a + noisy magnitude and phase spectrogram and emits the enhanced pair. There + is no ``logits`` port and no KV cache, so nothing about it is generative + -- it is a single pure spectrum-to-spectrum call. + """ + try: + names = set(pkg.keys()) + except AttributeError: + return False + if names != {"model"}: + return False + try: + model = pkg["model"] + inputs = {str(value.name) for value in model.graph.inputs} + outputs = {str(value.name) for value in model.graph.outputs} + except (AttributeError, KeyError): + return False + if not {"noisy_mag", "noisy_pha"} <= inputs: + return False + return {"denoised_mag", "denoised_pha"} <= outputs + + def _looks_like_audio_codec(pkg: Any) -> bool: """Detect an audio-to-audio neural codec package. @@ -903,6 +931,9 @@ def write_onnx_genai_config( ``scheduler`` / ``guidance_scale`` set the loop. """ package_config = getattr(pkg, "config", None) + resolved_config = config if config is not None else package_config + if isinstance(resolved_config, ReUseConfig): + _validate_reuse_rate_selection(resolved_config) config_types = { getattr(candidate, "model_type", None) for candidate in (package_config, config) @@ -931,7 +962,6 @@ def write_onnx_genai_config( return _write_advisory_component_contract(pkg, output_dir, warning=warning) os.makedirs(output_dir, exist_ok=True) if is_shared_state_pixel_flow_package(pkg): - resolved_config = config if config is not None else getattr(pkg, "config", None) if resolved_config is None: raise ValueError("shared-state pixel-flow metadata requires a model config") path = write_shared_state_pixel_flow_workflow_metadata( @@ -1150,6 +1180,15 @@ def write_onnx_genai_config( ) return artifacts + if _looks_like_speech_enhancement(pkg): + # A spectral enhancement model has no logits and no cache: one pure + # spectrum-to-spectrum call. Emit before the config requirement below + # because nothing here needs a decoder config. + path = write_speech_enhancement_workflow_metadata( + pkg, output_dir, config if config is not None else getattr(pkg, "config", None) + ) + return {"inference_metadata": path} + if _looks_like_encoder_embedding(pkg): # A bidirectional encoder has no logits and no cache: it runs once and # returns one hidden vector per position. Emit before the config diff --git a/src/mobius/integrations/onnx_genai/speech_enhancement_metadata_test.py b/src/mobius/integrations/onnx_genai/speech_enhancement_metadata_test.py new file mode 100644 index 000000000..6c532603a --- /dev/null +++ b/src/mobius/integrations/onnx_genai/speech_enhancement_metadata_test.py @@ -0,0 +1,618 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for spectral speech-enhancement onnx-genai metadata.""" + +from __future__ import annotations + +import json +import os + +import numpy as np +import onnx_ir as ir +import pytest + +from mobius import build_from_module +from mobius.integrations.onnx_genai import ( + build_speech_enhancement_workflow_metadata, + write_onnx_genai_config, + write_speech_enhancement_workflow_metadata, +) +from mobius.integrations.onnx_genai.auto_export import _looks_like_speech_enhancement +from mobius.integrations.onnx_genai.inference_metadata_test import ( + _onnx_genai_schema_path, +) +from mobius.models.reuse import ReUseConfig, SEMambaSpeechEnhancementModel + +_TINY = { + "model_cfg": { + "hid_feature": 8, + "num_tfmamba": 1, + "d_state": 4, + "expand": 2, + "compress_factor": "relu_log1p", + }, + "stft_cfg": {"n_fft": 32, "hop_size": 4, "win_size": 32, "sampling_rate": 8000}, +} + + +def _package(**config_overrides): + config = ReUseConfig.from_json(_TINY) + for key, value in config_overrides.items(): + setattr(config, key, value) + module = SEMambaSpeechEnhancementModel(config) + return build_from_module(module, config, task="speech-enhancement"), config + + +def _runtime_package(**config_overrides): + config = ReUseConfig( + hid_feature=4, + num_tfmamba=1, + d_state=2, + d_conv=2, + expand=2, + n_fft=320, + hop_size=40, + win_size=320, + sampling_rate=8000, + **config_overrides, + ) + module = SEMambaSpeechEnhancementModel(config) + rng = np.random.default_rng(0) + for _name, parameter in module.named_parameters(): + shape = [ + dimension if isinstance(dimension, int) else 1 for dimension in parameter.shape + ] + parameter.const_value = ir.tensor( + (rng.standard_normal(shape) * 0.02).astype(np.float32) + ) + package = build_from_module(module, config, task="speech-enhancement") + ort = pytest.importorskip("onnxruntime") + session = ort.InferenceSession( + ir.to_proto(package["model"]).SerializeToString(), + providers=["CPUExecutionProvider"], + ) + return package, config, session + + +def _even_scaled(value: int, sample_rate: int, reference_rate: int) -> int: + scaled = value * sample_rate // reference_rate + return scaled if scaled % 2 == 0 else scaled + 1 + + +def _run_declared_workflow(package, config, session, waveform, source_rate): + """Execute the declared native/BWE transform, graph, and inverse contract.""" + torch = pytest.importorskip("torch") + metadata = build_speech_enhancement_workflow_metadata(package, config) + transforms = metadata["preprocessing"]["audio"]["transforms"] + samples = np.asarray(waveform, dtype=np.float32) + sample_rate = source_rate + + for transform in transforms: + op = transform["op"] + if op in {"decode", "downmix", "log1p"}: + continue + if op == "require_sample_rate": + assert sample_rate == transform["sample_rate"] + continue + if op == "resample": + target_rate = transform["sample_rate"] + target_length = round(samples.shape[-1] * target_rate / sample_rate) + source_positions = np.arange(samples.shape[-1], dtype=np.float64) + target_positions = np.linspace( + 0.0, samples.shape[-1] - 1, target_length, dtype=np.float64 + ) + samples = np.interp(target_positions, source_positions, samples).astype(np.float32) + sample_rate = target_rate + continue + if op in {"spectrogram", "scaled_spectrogram"}: + if op == "scaled_spectrogram": + reference_rate = transform["sample_rate"] + geometry = tuple( + _even_scaled(transform[name], sample_rate, reference_rate) + for name in ("n_fft", "hop_length", "win_length") + ) + else: + geometry = ( + transform["n_fft"], + transform["hop_length"], + transform["win_length"], + ) + break + else: + raise AssertionError("workflow declares no spectrogram transform") + + n_fft, hop_length, win_length = geometry + tensor = torch.from_numpy(samples[None]) + spectrum = torch.stft( + tensor, + n_fft, + hop_length=hop_length, + win_length=win_length, + window=torch.hann_window(win_length), + center=True, + pad_mode="reflect", + normalized=False, + return_complex=True, + ) + noisy_mag = torch.log1p(torch.abs(spectrum)).numpy() + noisy_pha = torch.angle(spectrum).numpy() + denoised_mag, denoised_pha, _denoised_com = session.run( + None, + {"noisy_mag": noisy_mag, "noisy_pha": noisy_pha}, + ) + + # NVIDIA suppresses frames whose decompressed magnitude is zero in more + # than half the bins, then performs ISTFT and pads/trims to the input. + decompressed = np.expm1(np.maximum(denoised_mag, 0.0)) + bad_frames = np.mean(np.equal(decompressed, 0.0), axis=1) > 0.5 + denoised_mag = denoised_mag.copy() + denoised_mag[:, :, bad_frames[0]] = 0.0 + magnitude = torch.from_numpy(np.expm1(np.maximum(denoised_mag, 0.0))) + phase = torch.from_numpy(denoised_pha) + enhanced = torch.istft( + torch.complex(magnitude * torch.cos(phase), magnitude * torch.sin(phase)), + n_fft, + hop_length=hop_length, + win_length=win_length, + window=torch.hann_window(win_length), + center=True, + ).numpy()[0] + target_length = samples.shape[-1] + if enhanced.shape[-1] < target_length: + enhanced = np.pad( + enhanced, + (0, target_length - enhanced.shape[-1]), + constant_values=1e-8, + ) + else: + enhanced = enhanced[:target_length] + return { + "audio": enhanced, + "sample_rate": sample_rate, + "sample_length": target_length, + "geometry": geometry, + "input_shape": noisy_mag.shape, + } + + +class TestDetection: + """The dispatcher must recognise the package structurally.""" + + def test_detects_enhancement_package(self): + pkg, _config = _package() + assert _looks_like_speech_enhancement(pkg) is True + + def test_rejects_a_decoder_package(self): + from mobius._testing import make_config + from mobius.models import CausalLMModel + + config = make_config() + pkg = build_from_module(CausalLMModel(config), config) + + assert _looks_like_speech_enhancement(pkg) is False + + +class TestMetadata: + """The emitted document must describe the graph truthfully.""" + + def test_declares_a_single_pure_invocation(self): + pkg, config = _package() + + metadata = build_speech_enhancement_workflow_metadata(pkg, config) + + workflow = metadata["pipeline"]["workflow"] + # No generation loop: nothing to sample, nothing carried between calls. + invokes = [s for s in workflow["steps"] if s["kind"] == "invoke"] + assert [s["component"] for s in invokes] == [ + "audio_preprocess", + "enhancer", + "audio_postprocess", + ] + emits = [s for s in workflow["steps"] if s["kind"] == "emit"] + assert len(emits) == 6 + assert not any(s["kind"] == "loop" for s in workflow["steps"]) + for effect in workflow["effects"].values(): + assert effect["retry"] == "pure" + + def test_publishes_every_graph_output(self): + pkg, config = _package() + + metadata = build_speech_enhancement_workflow_metadata(pkg, config) + + outputs = metadata["pipeline"]["workflow"]["outputs"] + assert set(outputs) == { + "denoised_mag", + "denoised_pha", + "denoised_com", + "audio", + "sample_rate", + "sample_lengths", + } + # The complex spectrogram carries a trailing (real, imag) pair. + assert outputs["denoised_com"]["contract"]["rank"] == 4 + assert outputs["denoised_com"]["contract"]["shape"][-1] == 2 + assert outputs["denoised_mag"]["contract"]["rank"] == 3 + profile = metadata["profiles"]["speech_enhancement"] + assert set(profile["outputs"]) == set(outputs) + + def test_declares_the_stft_front_end(self): + pkg, config = _package() + + metadata = build_speech_enhancement_workflow_metadata(pkg, config) + + transforms = metadata["preprocessing"]["audio"]["transforms"] + by_op = {t["op"]: t for t in transforms} + assert "resample" not in by_op + spectrogram = by_op["scaled_spectrogram"] + assert spectrogram["n_fft"] == config.n_fft + assert spectrogram["hop_length"] == config.hop_size + assert spectrogram["win_length"] == config.win_size + assert spectrogram["sample_rate"] == config.sampling_rate + assert spectrogram["mode"] == ( + "native_scaled_floor_then_even_center_reflect_unnormalized" + ) + # The model is trained on log1p-compressed magnitudes. + assert "log1p" in by_op + assert by_op["log1p"]["inputs"] == ["magnitude"] + + def test_audio_outputs_bind_to_the_graph_inputs(self): + pkg, config = _package() + + metadata = build_speech_enhancement_workflow_metadata(pkg, config) + + bindings = {b["name"]: b for b in metadata["preprocessing"]["audio"]["outputs"]} + assert set(bindings) == { + "noisy_mag", + "noisy_pha", + "reference_audio", + "sample_rate", + "sample_lengths", + } + assert bindings["noisy_mag"]["source"] == "magnitude" + assert bindings["noisy_pha"]["source"] == "phase" + assert bindings["reference_audio"]["source"] == "samples" + assert bindings["sample_rate"]["source"] == "sample_rate" + assert bindings["sample_lengths"]["source"] == "sample_lengths" + assert bindings["noisy_mag"]["contract"]["rank"] == 3 + assert bindings["reference_audio"]["contract"]["rank"] == 2 + + def test_declares_the_adapter_abi_when_the_front_end_ships(self): + """A runtime must be able to version-check the STFT adapter it has to supply. + + The component already names the ABI, but a consumer reads the manifest to + decide whether it can run the package at all — the other audio workflows + declare it there, and omitting it hides the requirement until binding time. + """ + pkg, config = _package() + + metadata = build_speech_enhancement_workflow_metadata(pkg, config) + + manifest = metadata["pipeline"]["workflow"]["manifest"] + components = metadata["pipeline"]["workflow"]["components"] + for name in ("audio_preprocess", "audio_postprocess"): + component = components[name] + abi = component["implementation"]["abi"] + assert manifest["adapter_abis"][abi] == component["implementation"]["version"] + + def test_omits_the_adapter_abi_when_no_front_end_ships(self): + """No adapter in the package means nothing for a runtime to version-check. + + Declaring the ABI unconditionally would advertise a requirement the caller + does not have to satisfy, since it supplies the spectra itself. + """ + pkg, config = _package() + config.sampling_rate = None # type: ignore[assignment] + + metadata = build_speech_enhancement_workflow_metadata(pkg, config) + + manifest = metadata["pipeline"]["workflow"]["manifest"] + assert "adapter_abis" not in manifest + assert "audio_preprocess" not in metadata["pipeline"]["workflow"]["components"] + assert "audio_postprocess" not in metadata["pipeline"]["workflow"]["components"] + + def test_omits_the_program_when_geometry_is_unknown(self): + """Never invent a transform program the config cannot justify.""" + pkg, config = _package() + config.sampling_rate = None # type: ignore[assignment] + + metadata = build_speech_enhancement_workflow_metadata(pkg, config) + + assert "preprocessing" not in metadata + # The caller then supplies the spectra directly. + inputs = metadata["pipeline"]["workflow"]["inputs"] + assert set(inputs) == {"request.noisy_mag", "request.noisy_pha"} + + def test_fixed_native_rate_requires_that_rate_without_resampling(self): + pkg, config = _package(input_sampling_rate=16_000) + + metadata = build_speech_enhancement_workflow_metadata(pkg, config) + transforms = metadata["preprocessing"]["audio"]["transforms"] + by_op = {transform["op"]: transform for transform in transforms} + + assert by_op["require_sample_rate"]["sample_rate"] == 16_000 + assert "resample" not in by_op + assert by_op["spectrogram"] == { + "op": "spectrogram", + "n_fft": 64, + "hop_length": 8, + "win_length": 64, + "window": "hann", + "mode": "center_reflect_unnormalized", + "inputs": ["samples"], + "outputs": ["magnitude", "phase"], + } + + def test_bwe_rate_resamples_and_scales_static_geometry(self): + pkg, config = _package(bwe_sampling_rate=48_000) + + metadata = build_speech_enhancement_workflow_metadata(pkg, config) + transforms = metadata["preprocessing"]["audio"]["transforms"] + by_op = {transform["op"]: transform for transform in transforms} + + assert by_op["resample"]["sample_rate"] == 48_000 + assert "require_sample_rate" not in by_op + assert by_op["spectrogram"]["n_fft"] == 192 + assert by_op["spectrogram"]["hop_length"] == 24 + assert by_op["spectrogram"]["win_length"] == 192 + + @pytest.mark.parametrize("bwe_sampling_rate", [16_000, 48_000]) + def test_rejects_native_and_bwe_rates_together(self, bwe_sampling_rate): + pkg, config = _package() + config.input_sampling_rate = 16_000 + config.bwe_sampling_rate = bwe_sampling_rate + + with pytest.raises(ValueError, match="mutually exclusive"): + build_speech_enhancement_workflow_metadata(pkg, config) + + @pytest.mark.parametrize( + "missing_field", ["sampling_rate", "n_fft", "hop_size", "win_size"] + ) + @pytest.mark.parametrize("bwe_sampling_rate", [16_000, 48_000]) + def test_rejects_dual_rates_before_incomplete_geometry( + self, missing_field, bwe_sampling_rate + ): + pkg, config = _package() + setattr(config, missing_field, None) + config.input_sampling_rate = 16_000 + config.bwe_sampling_rate = bwe_sampling_rate + + with pytest.raises(ValueError, match="mutually exclusive"): + build_speech_enhancement_workflow_metadata(pkg, config) + + @pytest.mark.parametrize("rate_name", ["input_sampling_rate", "bwe_sampling_rate"]) + @pytest.mark.parametrize("value", ["16000", 16_000.0, True, False, 0, -1]) + def test_rejects_invalid_rate_before_incomplete_geometry(self, rate_name, value): + pkg, config = _package() + config.sampling_rate = None # type: ignore[assignment] + setattr(config, rate_name, value) + + with pytest.raises(ValueError, match=f"{rate_name} must be a positive integer"): + build_speech_enhancement_workflow_metadata(pkg, config) + + @pytest.mark.parametrize( + ("rate_name", "rate_value"), + [(None, None), ("input_sampling_rate", 16_000), ("bwe_sampling_rate", 48_000)], + ) + @pytest.mark.parametrize( + "missing_field", ["sampling_rate", "n_fft", "hop_size", "win_size"] + ) + def test_valid_rates_with_incomplete_geometry_use_spectrum_inputs( + self, missing_field, rate_name, rate_value + ): + pkg, config = _package() + setattr(config, missing_field, None) + if rate_name is not None: + setattr(config, rate_name, rate_value) + + metadata = build_speech_enhancement_workflow_metadata(pkg, config) + + assert "preprocessing" not in metadata + assert set(metadata["pipeline"]["workflow"]["inputs"]) == { + "request.noisy_mag", + "request.noisy_pha", + } + + @pytest.mark.parametrize("bwe_sampling_rate", [16_000, 48_000]) + def test_builder_validates_rates_before_graph_inspection(self, bwe_sampling_rate): + _, config = _package() + config.sampling_rate = None # type: ignore[assignment] + config.input_sampling_rate = 16_000 + config.bwe_sampling_rate = bwe_sampling_rate + + with pytest.raises(ValueError, match="mutually exclusive"): + build_speech_enhancement_workflow_metadata({}, config) + + def test_postprocess_declares_inverse_and_reference_length_contract(self): + pkg, config = _package() + + metadata = build_speech_enhancement_workflow_metadata(pkg, config) + component = metadata["pipeline"]["workflow"]["components"]["audio_postprocess"] + parameters = component["contract"]["parameters"] + + assert parameters["geometry_mode"] == "native_scaled" + assert parameters["rounding"] == "floor_then_even" + assert parameters["magnitude_decompression"] == "relu_log1p" + assert parameters["zero_frame_fraction_threshold"] == pytest.approx(0.5) + assert parameters["length_alignment"] == "pad_or_trim_to_reference" + assert parameters["pad_value"] == pytest.approx(1e-8) + + def test_rejects_a_graph_without_the_expected_ports(self): + pkg, config = _package() + pkg["model"].graph.inputs[0].name = "something_else" + + with pytest.raises(ValueError, match="noisy_mag"): + build_speech_enhancement_workflow_metadata(pkg, config) + + +class TestNativeRateRuntime: + """Execute the pinned source's rate scaling and length alignment contract.""" + + @pytest.fixture(scope="class") + def native_runtime(self): + return _runtime_package() + + @pytest.mark.parametrize( + ("sample_rate", "geometry"), + [ + (8_000, (320, 40, 320)), + (16_000, (640, 80, 640)), + (48_000, (1920, 240, 1920)), + ], + ) + @pytest.mark.parametrize("length_delta", [0, 1]) + def test_native_rate_preserves_rate_and_odd_even_length( + self, + native_runtime, + sample_rate, + geometry, + length_delta, + ): + package, config, session = native_runtime + length = sample_rate // 20 + length_delta + time = np.arange(length, dtype=np.float32) / sample_rate + waveform = 0.1 * np.sin(2 * np.pi * 440 * time) + + result = _run_declared_workflow( + package, + config, + session, + waveform, + sample_rate, + ) + + assert result["sample_rate"] == sample_rate + assert result["sample_length"] == length + assert result["audio"].shape == (length,) + assert result["geometry"] == geometry + assert result["input_shape"][1] == geometry[0] // 2 + 1 + assert np.isfinite(result["audio"]).all() + + def test_explicit_bwe_resamples_and_returns_target_rate(self): + package, config, session = _runtime_package(bwe_sampling_rate=16_000) + source_rate = 8_000 + source_length = 401 + time = np.arange(source_length, dtype=np.float32) / source_rate + waveform = 0.1 * np.sin(2 * np.pi * 440 * time) + + result = _run_declared_workflow( + package, + config, + session, + waveform, + source_rate, + ) + + assert result["sample_rate"] == 16_000 + assert result["sample_length"] == 802 + assert result["audio"].shape == (802,) + assert result["geometry"] == (640, 80, 640) + assert result["input_shape"][1] == 321 + + +class TestAutoExport: + """`write_onnx_genai_config` must route the package here.""" + + @pytest.mark.parametrize("bwe_sampling_rate", [16_000, 48_000]) + def test_rejects_dual_rates_before_package_inspection_or_output( + self, tmp_path, bwe_sampling_rate + ): + _, config = _package() + config.sampling_rate = None # type: ignore[assignment] + config.input_sampling_rate = 16_000 + config.bwe_sampling_rate = bwe_sampling_rate + output_dir = tmp_path / "metadata" + + with pytest.raises(ValueError, match="mutually exclusive"): + write_onnx_genai_config({}, str(output_dir), config=config) + + assert not output_dir.exists() + + @pytest.mark.parametrize("rate_name", ["input_sampling_rate", "bwe_sampling_rate"]) + @pytest.mark.parametrize("value", ["16000", 16_000.0, True, False, 0, -1]) + def test_rejects_malformed_rate_before_package_inspection_or_output( + self, tmp_path, rate_name, value + ): + _, config = _package() + config.n_fft = None # type: ignore[assignment] + setattr(config, rate_name, value) + output_dir = tmp_path / "metadata" + + with pytest.raises(ValueError, match=f"{rate_name} must be a positive integer"): + write_onnx_genai_config({}, str(output_dir), config=config) + + assert not output_dir.exists() + + def test_validates_package_config_before_structural_dispatch(self, tmp_path): + pkg, config = _package() + pkg.clear() + config.hop_size = None # type: ignore[assignment] + config.input_sampling_rate = 16_000 + config.bwe_sampling_rate = 48_000 + output_dir = tmp_path / "metadata" + + with pytest.raises(ValueError, match="mutually exclusive"): + write_onnx_genai_config(pkg, str(output_dir)) + + assert not output_dir.exists() + + def test_dispatches_to_the_enhancement_writer(self, tmp_path): + pkg, config = _package() + + artifacts = write_onnx_genai_config(pkg, str(tmp_path), config=config) + + assert os.path.isfile(artifacts["inference_metadata"]) + + def test_written_document_round_trips(self, tmp_path): + pkg, config = _package() + yaml = pytest.importorskip("yaml") + + path = write_speech_enhancement_workflow_metadata(pkg, str(tmp_path), config) + + with open(path, encoding="utf-8") as handle: + document = yaml.safe_load(handle) + assert document["schema_version"] == "v1" + assert "speech_enhancement" in document["profiles"] + + def test_writer_validates_rates_before_output_or_graph_inspection(self, tmp_path): + _, config = _package() + config.n_fft = None # type: ignore[assignment] + config.input_sampling_rate = 16_000 + config.bwe_sampling_rate = 48_000 + output_dir = tmp_path / "metadata" + + with pytest.raises(ValueError, match="mutually exclusive"): + write_speech_enhancement_workflow_metadata({}, str(output_dir), config) + + assert not output_dir.exists() + + +class TestSchema: + """The document must validate against onnx-genai's committed schema.""" + + def test_matches_onnx_genai_json_schema(self): + jsonschema = pytest.importorskip("jsonschema") + schema_path = _onnx_genai_schema_path() + if schema_path is None: + pytest.skip("onnx-genai schema not found (set ONNX_GENAI_SCHEMA)") + pkg, config = _package() + + metadata = build_speech_enhancement_workflow_metadata(pkg, config) + + with open(schema_path, encoding="utf-8") as handle: + schema = json.load(handle) + jsonschema.validate(metadata, schema) + + def test_matches_schema_without_preprocessing(self): + jsonschema = pytest.importorskip("jsonschema") + schema_path = _onnx_genai_schema_path() + if schema_path is None: + pytest.skip("onnx-genai schema not found (set ONNX_GENAI_SCHEMA)") + pkg, config = _package() + config.n_fft = None # type: ignore[assignment] + + metadata = build_speech_enhancement_workflow_metadata(pkg, config) + + with open(schema_path, encoding="utf-8") as handle: + schema = json.load(handle) + jsonschema.validate(metadata, schema) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 9ff6a70a2..989cfd698 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -9901,6 +9901,8 @@ def write_language_diffusion_workflow_metadata( _AUDIO_PREPROCESS_ABI = "onnx-genai.audio-preprocess" _AUDIO_PREPROCESS_ABI_VERSION = "1" +_AUDIO_POSTPROCESS_ABI = "onnx-genai.audio-postprocess" +_AUDIO_POSTPROCESS_ABI_VERSION = "1" def _audio_preprocess_component( @@ -10363,6 +10365,572 @@ def build_encoder_embedding_workflow_metadata( } +#: Graph inputs of a spectral speech-enhancement model, in binding order. +_SPEECH_ENHANCEMENT_INPUTS: tuple[str, ...] = ("noisy_mag", "noisy_pha") + +#: Graph outputs, in the order the enhancement task publishes them. +_SPEECH_ENHANCEMENT_OUTPUTS: tuple[str, ...] = ( + "denoised_mag", + "denoised_pha", + "denoised_com", +) + + +def _stft_preprocess_component( + mag_contract: dict[str, Any], + pha_contract: dict[str, Any], + waveform_contract: dict[str, Any], + sample_rate_contract: dict[str, Any], + lengths_contract: dict[str, Any], +) -> dict[str, Any]: + """Declare the audio-preprocessing adapter that produces the noisy STFT. + + The enhancement graph consumes a spectrum, not a waveform, so the adapter + turns request-supplied encoded audio bytes into the magnitude and phase + the model was trained on. Declaring its ports lets a runtime type-check + the binding without knowing which model family produced the package. + """ + return { + "implementation": { + "kind": "adapter", + "abi": _AUDIO_PREPROCESS_ABI, + "version": _AUDIO_PREPROCESS_ABI_VERSION, + }, + "ports": { + "inputs": { + "encoded": {"dtype": "uint8", "rank": 1, "shape": ["bytes"]}, + }, + "outputs": { + "noisy_mag": mag_contract, + "noisy_pha": pha_contract, + "reference_audio": waveform_contract, + "sample_rate": sample_rate_contract, + "sample_lengths": lengths_contract, + }, + }, + "contract": { + "id": _AUDIO_PREPROCESS_ABI, + "version": _AUDIO_PREPROCESS_ABI_VERSION, + "bindings": { + "encoded": "encoded", + "noisy_mag": "noisy_mag", + "noisy_pha": "noisy_pha", + "reference_audio": "reference_audio", + "sample_rate": "sample_rate", + "sample_lengths": "sample_lengths", + }, + }, + "effects": ["audio_preprocess"], + } + + +def _validate_reuse_rate_selection(config: Any) -> tuple[int | None, int | None]: + """Validate the mutually exclusive native-rate and BWE selections.""" + input_rate = getattr(config, "input_sampling_rate", None) + bwe_rate = getattr(config, "bwe_sampling_rate", None) + for name, value in (("input_sampling_rate", input_rate), ("bwe_sampling_rate", bwe_rate)): + if value is not None and ( + not isinstance(value, int) or isinstance(value, bool) or value <= 0 + ): + raise ValueError(f"{name} must be a positive integer when provided") + if input_rate is not None and bwe_rate is not None: + raise ValueError("input_sampling_rate and bwe_sampling_rate are mutually exclusive") + return input_rate, bwe_rate + + +def _reuse_stft_geometry(config: Any) -> dict[str, Any] | None: + """Resolve native-dynamic, fixed-native, or explicit-BWE STFT geometry.""" + input_rate, bwe_rate = _validate_reuse_rate_selection(config) + raw_reference = { + "sample_rate": getattr(config, "sampling_rate", None), + "n_fft": getattr(config, "n_fft", None), + "hop_length": getattr(config, "hop_size", None), + "win_length": getattr(config, "win_size", None), + } + if any(not isinstance(value, int) or value <= 0 for value in raw_reference.values()): + return None + reference = {name: int(value) for name, value in raw_reference.items()} + + selected_rate = bwe_rate or input_rate + if selected_rate is None: + return { + "mode": "native_scaled", + "reference_sample_rate": reference["sample_rate"], + "n_fft": reference["n_fft"], + "hop_length": reference["hop_length"], + "win_length": reference["win_length"], + "rounding": "floor_then_even", + } + + def _scaled(value: int) -> int: + scaled = value * selected_rate // reference["sample_rate"] + return scaled if scaled % 2 == 0 else scaled + 1 + + scaled_geometry = { + "n_fft": _scaled(reference["n_fft"]), + "hop_length": _scaled(reference["hop_length"]), + "win_length": _scaled(reference["win_length"]), + } + if any(value <= 0 for value in scaled_geometry.values()): + raise ValueError( + f"selected sample rate {selected_rate} is too small for RE-USE STFT geometry" + ) + return { + "mode": "bwe" if bwe_rate is not None else "fixed_native", + "sample_rate": selected_rate, + **scaled_geometry, + } + + +def _stft_transforms(config: Any) -> list[dict[str, Any]] | None: + """Describe the STFT front-end the enhancement model was trained with. + + Returns ``None`` when the config does not carry the STFT geometry, so a + package is never given an invented transform program. + """ + geometry = _reuse_stft_geometry(config) + if geometry is None: + return None + + transforms: list[dict[str, Any]] = [ + { + "op": "decode", + "outputs": ["samples", "sample_rate", "sample_lengths"], + }, + {"op": "downmix", "channels": 1}, + ] + if geometry["mode"] == "bwe": + transforms.append( + { + "op": "resample", + "sample_rate": geometry["sample_rate"], + "inputs": ["samples"], + "outputs": ["samples", "sample_rate", "sample_lengths"], + } + ) + elif geometry["mode"] == "fixed_native": + # A static native-rate graph must reject a differently sampled input, + # not silently analyze it with mismatched geometry. + transforms.append( + { + "op": "require_sample_rate", + "sample_rate": geometry["sample_rate"], + "inputs": ["sample_rate"], + } + ) + + spectrogram = { + "op": "scaled_spectrogram" if geometry["mode"] == "native_scaled" else "spectrogram", + "n_fft": geometry["n_fft"], + "hop_length": geometry["hop_length"], + "win_length": geometry["win_length"], + "window": "hann", + # NVIDIA uses torch.stft(center=True, pad_mode="reflect", + # normalized=False). Keep those coupled semantics explicit rather + # than letting an adapter select different defaults. + "mode": "center_reflect_unnormalized", + "inputs": ["samples"], + "outputs": ["magnitude", "phase"], + } + if geometry["mode"] == "native_scaled": + spectrogram.update( + { + "sample_rate": geometry["reference_sample_rate"], + "mode": "native_scaled_floor_then_even_center_reflect_unnormalized", + } + ) + transforms.append(spectrogram) + + # RE-USE trains on log1p-compressed magnitudes; a caller that skips the + # compression feeds the model a different distribution. The transform + # vocabulary is open for extension values, so the compression is declared + # rather than left as an undocumented assumption. + compression = getattr(config, "compress_factor", None) + if isinstance(compression, str) and compression.endswith("log1p"): + transforms.append({"op": "log1p", "inputs": ["magnitude"], "outputs": ["magnitude"]}) + return transforms + + +def _stft_postprocess_component( + mag_contract: dict[str, Any], + pha_contract: dict[str, Any], + waveform_contract: dict[str, Any], + sample_rate_contract: dict[str, Any], + lengths_contract: dict[str, Any], + config: Any, +) -> dict[str, Any]: + """Declare NVIDIA's inverse STFT, artifact suppression, and length alignment.""" + geometry = _reuse_stft_geometry(config) + assert geometry is not None + parameters: dict[str, str | int | float | bool | None] = { + "geometry_mode": geometry["mode"], + "sample_rate": geometry.get("sample_rate") or geometry.get("reference_sample_rate"), + "n_fft": geometry["n_fft"], + "hop_length": geometry["hop_length"], + "win_length": geometry["win_length"], + "window": "hann", + "stft_mode": "center_reflect_unnormalized", + "rounding": geometry.get("rounding"), + "magnitude_decompression": getattr(config, "compress_factor", None), + "zero_frame_fraction_threshold": 0.5, + "length_alignment": "pad_or_trim_to_reference", + "pad_value": 1e-8, + } + return { + "implementation": { + "kind": "adapter", + "abi": _AUDIO_POSTPROCESS_ABI, + "version": _AUDIO_POSTPROCESS_ABI_VERSION, + }, + "ports": { + "inputs": { + "denoised_mag": mag_contract, + "denoised_pha": pha_contract, + "reference_audio": waveform_contract, + "sample_rate": sample_rate_contract, + "sample_lengths": lengths_contract, + }, + "outputs": { + "audio": waveform_contract, + "sample_rate": sample_rate_contract, + "sample_lengths": lengths_contract, + }, + }, + "contract": { + "id": _AUDIO_POSTPROCESS_ABI, + "version": _AUDIO_POSTPROCESS_ABI_VERSION, + "bindings": { + "denoised_mag": "denoised_mag", + "denoised_pha": "denoised_pha", + "reference_audio": "reference_audio", + "sample_rate": "sample_rate", + "sample_lengths": "sample_lengths", + "audio": "audio", + }, + "parameters": parameters, + }, + "effects": ["audio_postprocess"], + } + + +def build_speech_enhancement_workflow_metadata( + pkg: Any, + config: Any = None, + *, + artifact: str = "model.onnx", +) -> dict[str, Any]: + """Build one-file metadata for a spectral speech-enhancement model. + + An enhancement model such as RE-USE / SEMamba maps a noisy STFT to a clean + one. It is not generative: it reads the whole spectrogram at once, carries + no state between calls and has no ``logits`` to sample. The workflow is a + pure preprocess → enhance → postprocess sequence; describing it with decoder + metadata would publish a generation loop the artifact cannot execute. + + The STFT lives outside the graph, so when the config carries the STFT + geometry it is published as an audio preprocessing program and the + workflow accepts encoded audio. A paired postprocessing adapter declares + NVIDIA's magnitude decompression, zero-frame suppression, inverse STFT, + and exact input-length alignment. + + Args: + pkg: The built :class:`ModelPackage`; must hold a single ``model``. + config: The resolved architecture config. When it carries STFT + geometry (``sampling_rate``, ``n_fft``, ``hop_size``, + ``win_size``) the workflow takes encoded audio and declares the + transform program; otherwise the spectra are request-supplied. + artifact: Model artifact path relative to the package root. + + Returns: + A metadata document with a ``speech_enhancement`` profile and a pure, + single-request ``pipeline.workflow``. + """ + _validate_reuse_rate_selection(config) + if "model" not in pkg: + raise ValueError("speech enhancement workflow requires a 'model' component") + model = pkg["model"] + + graph_inputs = {str(value.name): value for value in model.graph.inputs} + graph_outputs = {str(value.name): value for value in model.graph.outputs} + missing = [n for n in _SPEECH_ENHANCEMENT_INPUTS if n not in graph_inputs] + if missing: + raise ValueError(f"speech enhancement graph must declare inputs {missing}") + emitted = [n for n in _SPEECH_ENHANCEMENT_OUTPUTS if n in graph_outputs] + if not emitted: + raise ValueError( + "speech enhancement graph must declare at least one of " + f"{list(_SPEECH_ENHANCEMENT_OUTPUTS)}" + ) + if config is not None and any( + name not in graph_outputs for name in ("denoised_mag", "denoised_pha") + ): + raise ValueError( + "speech enhancement audio postprocessing requires denoised_mag and denoised_pha" + ) + + mag_contract = _contract(graph_inputs["noisy_mag"]) + pha_contract = _contract(graph_inputs["noisy_pha"]) + transforms = _stft_transforms(config) + waveform_contract = { + "dtype": "float32", + "rank": 2, + "shape": ["batch", "audio_samples"], + } + sample_rate_contract = {"dtype": "int64", "rank": 0, "shape": []} + lengths_contract = {"dtype": "int64", "rank": 1, "shape": ["batch"]} + geometry = _reuse_stft_geometry(config) + if geometry is not None and geometry["mode"] != "native_scaled": + expected_bins = geometry["n_fft"] // 2 + 1 + for name in _SPEECH_ENHANCEMENT_INPUTS: + shape = list(graph_inputs[name].shape or []) + if len(shape) > 1 and isinstance(shape[1], int) and shape[1] != expected_bins: + raise ValueError( + f"{name} frequency extent {shape[1]} does not match selected " + f"STFT geometry ({expected_bins} bins)" + ) + + workflow_outputs: dict[str, Any] = {} + emit_nodes: list[dict[str, Any]] = [] + profile_outputs: dict[str, str] = {} + for name in emitted: + workflow_outputs[name] = { + "contract": _contract(graph_outputs[name]), + "role": "tensor", + "stage": "post_adapter", + } + emit_nodes.append( + { + "kind": "emit", + "value": f"enhancer.{name}", + "output": name, + "mode": "replace", + } + ) + profile_outputs[name] = name + invoke_outputs = {name: f"enhancer.{name}" for name in emitted} + + effects: dict[str, Any] = { + # One pure call: the model observes nothing outside its inputs, so a + # retry replays it exactly and a speculative clone is safe. + "enhance": {"retry": "pure", "speculation_safety": {"kind": "clonable"}}, + } + components: dict[str, Any] = { + "enhancer": _component(model, artifact, effects=("enhance",)) + } + initial_effects: dict[str, str] = {"enhance": "enhance.0"} + nodes: list[dict[str, Any]] = [] + + if transforms is not None: + effects["audio_preprocess"] = { + "retry": "pure", + "speculation_safety": {"kind": "clonable"}, + } + effects["audio_postprocess"] = { + "retry": "pure", + "speculation_safety": {"kind": "clonable"}, + } + components["audio_preprocess"] = _stft_preprocess_component( + mag_contract, + pha_contract, + waveform_contract, + sample_rate_contract, + lengths_contract, + ) + components["audio_postprocess"] = _stft_postprocess_component( + _contract(graph_outputs["denoised_mag"]), + _contract(graph_outputs["denoised_pha"]), + waveform_contract, + sample_rate_contract, + lengths_contract, + config, + ) + initial_effects["audio_preprocess"] = "audio_preprocess.0" + initial_effects["audio_postprocess"] = "audio_postprocess.0" + workflow_inputs = { + "request.audio": { + "contract": {"dtype": "uint8", "rank": 1, "shape": ["bytes"]}, + "role": {"kind": "runtime", "version": "1.0", "role": "media"}, + "source": {"kind": "request", "field": "media"}, + "required": True, + } + } + nodes.append( + _invoke( + "audio_preprocess", + {"encoded": "request.audio"}, + { + "noisy_mag": "audio.noisy_mag", + "noisy_pha": "audio.noisy_pha", + "reference_audio": "audio.reference", + "sample_rate": "audio.sample_rate", + "sample_lengths": "audio.sample_lengths", + }, + ) + ) + invoke_inputs = { + "noisy_mag": "audio.noisy_mag", + "noisy_pha": "audio.noisy_pha", + } + else: + # Without the STFT geometry we cannot state how a waveform becomes a + # spectrum, so the caller supplies the spectra directly. The portable + # role vocabulary has no term for a magnitude or phase spectrogram, + # so these stay opaque rather than being mislabelled as audio. + workflow_inputs = { + f"request.{name}": { + "contract": _contract(graph_inputs[name]), + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": f"request.{name}"}, + "required": True, + } + for name in _SPEECH_ENHANCEMENT_INPUTS + } + invoke_inputs = {name: f"request.{name}" for name in _SPEECH_ENHANCEMENT_INPUTS} + + nodes.append(_invoke("enhancer", invoke_inputs, invoke_outputs)) + if transforms is not None: + nodes.append( + _invoke( + "audio_postprocess", + { + "denoised_mag": "enhancer.denoised_mag", + "denoised_pha": "enhancer.denoised_pha", + "reference_audio": "audio.reference", + "sample_rate": "audio.sample_rate", + "sample_lengths": "audio.sample_lengths", + }, + { + "audio": "enhanced.audio", + "sample_rate": "enhanced.sample_rate", + "sample_lengths": "enhanced.sample_lengths", + }, + ) + ) + for name, contract, role in ( + ("audio", waveform_contract, "audio"), + ("sample_rate", sample_rate_contract, "tensor"), + ("sample_lengths", lengths_contract, "tensor"), + ): + workflow_outputs[name] = { + "contract": contract, + "role": role, + "stage": "post_adapter", + } + profile_outputs[name] = name + emit_nodes.append( + { + "kind": "emit", + "value": f"enhanced.{name}", + "output": name, + "mode": "replace", + } + ) + nodes.extend(emit_nodes) + + workflow = { + "manifest": { + # Declared only when the STFT adapter is actually shipped. Without + # `transforms` the caller supplies the spectra directly, so there is no + # adapter for a runtime to version-check. + **( + { + "adapter_abis": { + _AUDIO_PREPROCESS_ABI: _AUDIO_PREPROCESS_ABI_VERSION, + _AUDIO_POSTPROCESS_ABI: _AUDIO_POSTPROCESS_ABI_VERSION, + } + } + if transforms is not None + else {} + ), + "capabilities": ["workflow_ssa", "linear_effects", "typed_emit"], + }, + "effects": effects, + "inputs": workflow_inputs, + "outputs": workflow_outputs, + "components": components, + "initial_effects": initial_effects, + "graph": {"kind": "sequence", "nodes": nodes}, + } + + profile: dict[str, Any] = { + "kind": "speech_enhancement", + "version": "1.0", + "requirement": "required", + "outputs": profile_outputs, + } + + metadata: dict[str, Any] = {"schema_version": "v1"} + if transforms is not None: + metadata["preprocessing"] = { + "audio": { + "transforms": transforms, + # The full contract is published alongside dtype/rank because + # this package declares a `pipeline.workflow`, whose binding + # the runtime type-checks against these ports. + "outputs": [ + { + "name": "noisy_mag", + "source": "magnitude", + "content": "features", + "dtype": mag_contract["dtype"], + "rank": mag_contract["rank"], + "contract": mag_contract, + }, + { + "name": "noisy_pha", + "source": "phase", + "content": "features", + "dtype": pha_contract["dtype"], + "rank": pha_contract["rank"], + "contract": pha_contract, + }, + { + "name": "reference_audio", + "source": "samples", + "content": "waveform", + "dtype": waveform_contract["dtype"], + "rank": waveform_contract["rank"], + "contract": waveform_contract, + }, + { + "name": "sample_rate", + "source": "sample_rate", + "content": "sample_rate", + "dtype": sample_rate_contract["dtype"], + "rank": sample_rate_contract["rank"], + "contract": sample_rate_contract, + }, + { + "name": "sample_lengths", + "source": "sample_lengths", + "content": "sample_lengths", + "dtype": lengths_contract["dtype"], + "rank": lengths_contract["rank"], + "contract": lengths_contract, + }, + ], + } + } + metadata["profiles"] = {"speech_enhancement": profile} + metadata["pipeline"] = {"workflow": _publish_workflow_v1(workflow)} + return metadata + + +def write_speech_enhancement_workflow_metadata( + pkg: Any, + output_dir: str, + config: Any = None, +) -> str: + """Write one-file speech-enhancement metadata into *output_dir*.""" + _validate_reuse_rate_selection(config) + os.makedirs(output_dir, exist_ok=True) + metadata = build_speech_enhancement_workflow_metadata(pkg, config) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + _dump_yaml(metadata, handle) + return path + + def write_encoder_embedding_workflow_metadata( pkg: Any, output_dir: str, diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index 459c60c0d..72099a822 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -183,6 +183,8 @@ def build_transformers_model( prune_prefill_prefix: bool = False, glm_full_attention: bool = False, export_paged_attention: bool = False, + input_sampling_rate: int | None = None, + bwe_sampling_rate: int | None = None, ) -> ModelPackage: """Build a model package from a Transformers checkpoint. @@ -199,18 +201,68 @@ def build_transformers_model( their native block-weight representation. Set it to ``False`` only to request explicit dense reconstruction. """ + if input_sampling_rate is not None and bwe_sampling_rate is not None: + raise ValueError("input_sampling_rate and bwe_sampling_rate are mutually exclusive") + from mobius.integrations.diffusers import build_diffusers_pipeline from mobius.integrations.transformers._config_resolver import ( _config_from_hf, _default_task_for_model, ) + detection_revision = revision + if model_id == "nvidia/RE-USE" and detection_revision is None: + # Pin the very first AutoConfig/raw-JSON probe, not only the later + # bespoke loader. Otherwise mutable Hub main could change dispatch + # before RE-USE's pinned default ever takes effect. + from mobius.models.reuse import REUSE_REVISION + + detection_revision = REUSE_REVISION + hf_config, loaded_from_raw_json = _load_transformers_config( model_id, - revision=revision, + revision=detection_revision, trust_remote_code=trust_remote_code, ) if hf_config is None or (loaded_from_raw_json and hf_config.model_type not in registry): + from mobius.models.reuse import _is_reuse_checkpoint, build_reuse + + if module_class is None and _is_reuse_checkpoint(model_id, detection_revision): + from mobius.tasks import SpeechEnhancementTask + + if task not in (None, "speech-enhancement") and not isinstance( + task, SpeechEnhancementTask + ): + raise ValueError("RE-USE checkpoints only support task='speech-enhancement'.") + unsupported = { + "output_layer_indices": output_layer_indices is not None, + "text_only": text_only, + "fp8_kv_cache": fp8_kv_cache, + "kv_cache_scales": kv_cache_scales is not None, + "prune_prefill_prefix": prune_prefill_prefix, + "glm_full_attention": glm_full_attention, + "export_paged_attention": export_paged_attention, + } + selected = sorted(name for name, enabled in unsupported.items() if enabled) + if selected: + raise ValueError( + "RE-USE checkpoints do not support these decoder-only options: " + + ", ".join(selected) + ) + return build_reuse( + model_id, + revision=detection_revision, + dtype=dtype, + execution_provider=execution_provider, + load_weights=load_weights, + input_sampling_rate=input_sampling_rate, + bwe_sampling_rate=bwe_sampling_rate, + ) + if input_sampling_rate is not None or bwe_sampling_rate is not None: + raise ValueError( + "input_sampling_rate and bwe_sampling_rate are only supported " + "for RE-USE speech-enhancement checkpoints" + ) if text_only: raise ValueError( f"text_only=True is not supported for '{model_id}': it does not " @@ -231,6 +283,12 @@ def build_transformers_model( execution_provider=execution_provider, ) + if input_sampling_rate is not None or bwe_sampling_rate is not None: + raise ValueError( + "input_sampling_rate and bwe_sampling_rate are only supported " + "for RE-USE speech-enhancement checkpoints" + ) + hf_config, parent_config, model_type = _select_primary_config(hf_config) compressed_tensors_config = CompressedTensorsConfig.from_hf_config(parent_config) diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 6769da63c..05d6e0d2b 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -176,6 +176,9 @@ "Qwen4ExpCausalLMModel", "SenseNovaU1Model", "SenseVoiceSmallModel", + "ReUseConfig", + "SEMambaSpeechEnhancementModel", + "build_reuse", "SortformerConfig", "SortformerDiarizationModel", "Qwen3TTSCodePredictorModel", @@ -433,6 +436,11 @@ Qwen25VLTextModel, Qwen25VLVisionEncoderModel, ) +from mobius.models.reuse import ( + ReUseConfig, + SEMambaSpeechEnhancementModel, + build_reuse, +) from mobius.models.sensenova_u1 import SenseNovaU1Model from mobius.models.sensevoice_small import SenseVoiceSmallModel from mobius.models.smallthinker import SmallThinkerGGUFCausalLMModel diff --git a/src/mobius/models/reuse.py b/src/mobius/models/reuse.py new file mode 100644 index 000000000..a429d75ca --- /dev/null +++ b/src/mobius/models/reuse.py @@ -0,0 +1,900 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""RE-USE / SEMamba universal speech enhancement (NVIDIA). + +Replicates the forward pass of NVIDIA's ``SEMamba`` generator as published +in `nvidia/RE-USE `_. The exported +ONNX graph consumes the magnitude and phase of a noisy STFT and produces +the enhanced magnitude, phase, and complex spectrogram. The STFT/ISTFT +themselves stay outside the graph, matching how every other audio model in +mobius is exported. + +Pipeline (matching ``SEMamba.forward``):: + + noisy_mag [B, F, T], noisy_pha [B, F, T] + -> stack as 2 channels, zero-pad time and freq by 2 [B, 2, T+2, F+2] + -> DenseEncoder (1x1 conv, dilated dense block, strided conv) + [B, C, T', F'] + -> num_tfmamba x TFMambaBlock (bidirectional Mamba over time, + then over frequency) [B, C, T', F'] + -> MagDecoder -> denoised_mag [B, F, T] + -> PhaseDecoder -> denoised_pha [B, F, T] + -> denoised_com = stack(mag*cos(pha), mag*sin(pha)) [B, F, T, 2] + +``T' = floor((T + 1) / 4) + 1`` and ``F' = floor((F - 1) / 2) + 1`` follow +from the encoder's ``stride=(4, 2)`` convolution; the decoders undo both +strides and the result is cropped back to the input ``(F, T)``. + +The SSM layers are the original (Mamba1) selective scan run over the whole +sequence in both directions, so they use +:class:`~mobius.components.SequenceMambaBlock` rather than the decode-time +:class:`~mobius.components.MambaBlock`. + +Reference implementation: ``models/generator_SEMamba_time_d4.py``, +``models/codec_module_time_d4.py``, and ``models/mamba_block2_SEMamba.py`` +in the ``nvidia/RE-USE`` repository. +""" + +from __future__ import annotations + +import dataclasses +import json +import math +import os + +import onnx_ir as ir +import torch +from onnxscript import OpBuilder, nn + +from mobius._configs import BaseModelConfig +from mobius._model_package import ModelPackage +from mobius.components import Conv2d, LayerNorm, Linear, SequenceMambaBlock + +# Slice "start" sentinel for a reverse (negative-step) slice. +_INT64_MIN = -9223372036854775808 + +# Frequency-axis geometry of the encoder, kept here because two places depend on +# it: ``SEMambaSpeechEnhancementModel.forward`` emits the tail pad, and +# ``DenseEncoder`` builds the strided convolution, while +# ``ReUseConfig.encoder_freq_bins`` predicts the resulting extent at build time. +# ``TestEncoderFreqBins`` pins the prediction against the graph the model +# actually builds, so these cannot drift apart silently. +_ENCODER_FREQ_TAIL_PAD = 2 +_ENCODER_FREQ_KERNEL = 3 +_ENCODER_FREQ_STRIDE = 2 + +# Immutable NVIDIA checkpoint revision used unless callers explicitly select +# another revision. Config, source, and weights are therefore resolved together. +REUSE_REVISION = "761905064ea1ea882e015e20a64e2e9d28458890" + + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class ReUseConfig(BaseModelConfig): + """Configuration for the RE-USE / SEMamba speech-enhancement model. + + Fields mirror the ``model_cfg`` and ``stft_cfg`` sections of the + ``nvidia/RE-USE`` ``config.json``. + """ + + # --- model_cfg --- + #: Input channels of the encoder (magnitude + phase). + input_channel: int = 2 + #: Output channels of each decoder head. + output_channel: int = 1 + #: Encoder/decoder feature width (Mamba ``d_model``). + hid_feature: int = 64 + #: Number of stacked time-frequency Mamba blocks. + num_tfmamba: int = 30 + #: SSM state dimension. + d_state: int = 16 + #: Causal Conv1D kernel size inside each Mamba block. + d_conv: int = 4 + #: Inner expansion factor (``d_inner = expand * hid_feature``). + expand: int = 4 + #: Depth of each dense block (number of dilated convolutions). + dense_depth: int = 4 + #: Epsilon for the instance / layer normalizations. + norm_epsilon: float = 1e-5 + + # --- stft_cfg --- + #: FFT size; the frequency axis has ``n_fft // 2 + 1`` bins. + n_fft: int = 320 + #: STFT hop size, in samples. + hop_size: int = 40 + #: STFT window size, in samples. + win_size: int = 320 + #: Audio sample rate the model was trained for. + sampling_rate: int = 8000 + #: Known native input rate for a static-frequency export. ``None`` keeps + #: frequency dynamic and follows the decoded audio's native rate. + input_sampling_rate: int | None = None + #: Explicit NVIDIA ``BWE`` target. The workflow resamples to this rate + #: before analysis; ``None`` preserves the native input rate. + bwe_sampling_rate: int | None = None + #: Magnitude compression applied before the model. Informational only: it is + #: applied by the STFT front-end, outside this graph. Note it is declared + #: under ``model_cfg`` in the checkpoint's config.json, not ``stft_cfg``, + #: despite describing a front-end step. + compress_factor: str = "relu_log1p" + + model_type: str | None = "reuse" + + @property + def num_freq_bins(self) -> int: + """Number of bins at the checkpoint's reference sample rate.""" + return self.n_fft // 2 + 1 + + @staticmethod + def _make_even(value: int) -> int: + """Match NVIDIA's floor-then-round-up-to-even geometry rule.""" + return value if value % 2 == 0 else value + 1 + + def stft_geometry(self, sample_rate: int) -> tuple[int, int, int]: + """Scale reference FFT, hop, and window sizes to *sample_rate*.""" + if sample_rate <= 0: + raise ValueError("sample_rate must be positive") + geometry = [ + self._make_even(value * sample_rate // self.sampling_rate) + for value in (self.n_fft, self.hop_size, self.win_size) + ] + if any(value <= 0 for value in geometry): + raise ValueError( + f"sample_rate {sample_rate} is too small for reference STFT geometry" + ) + return geometry[0], geometry[1], geometry[2] + + @property + def analysis_sampling_rate(self) -> int | None: + """Static analysis rate, or ``None`` for native-rate dynamic export.""" + return self.bwe_sampling_rate or self.input_sampling_rate + + @property + def analysis_stft_geometry(self) -> tuple[int, int, int] | None: + """Static analysis geometry, or ``None`` when native rate is runtime data.""" + rate = self.analysis_sampling_rate + return self.stft_geometry(rate) if rate is not None else None + + @property + def static_num_freq_bins(self) -> int | None: + """Static graph frequency extent for an explicitly selected rate.""" + geometry = self.analysis_stft_geometry + return geometry[0] // 2 + 1 if geometry is not None else None + + @property + def encoder_freq_bins(self) -> int: + """Reference-rate frequency extent of the encoder output.""" + return self._encoder_freq_bins(self.num_freq_bins) + + @property + def static_encoder_freq_bins(self) -> int | None: + """Static encoded extent for an explicitly selected analysis rate.""" + bins = self.static_num_freq_bins + return self._encoder_freq_bins(bins) if bins is not None else None + + @staticmethod + def _encoder_freq_bins(num_freq_bins: int) -> int: + """Frequency extent of the encoder output, i.e. what the TF blocks see. + + This is statically derivable when an export selects a native or BWE + sample rate. Native-rate exports leave it dynamic because the decoded + sample rate determines the FFT geometry at invocation time. + + The derivation mirrors what the graph actually does, in order: + + 1. ``forward`` zero-pads the tail of the frequency axis by + ``_ENCODER_FREQ_TAIL_PAD``, so the strided convolution never drops a + partial window. + 2. ``DenseEncoder.dense_conv_2`` is a ``kernel=(1, 3)``, ``stride=(4, 2)`` + convolution with no padding, giving the usual + ``floor((in - kernel) / stride) + 1``. + + ``n_fft`` is even, so ``num_freq_bins`` is odd and the floor is exact. + The decoder's ``up_conv1`` doubles this back to ``2 * encoder_freq_bins``, + which is ``>= num_freq_bins``; ``forward`` crops off the overshoot. + """ + padded = num_freq_bins + _ENCODER_FREQ_TAIL_PAD + return (padded - _ENCODER_FREQ_KERNEL) // _ENCODER_FREQ_STRIDE + 1 + + @property + def d_inner(self) -> int: + """Expanded Mamba inner dimension.""" + return self.expand * self.hid_feature + + @property + def dt_rank(self) -> int: + """Rank of the SSM time-step projection (``mamba_ssm`` "auto").""" + return math.ceil(self.hid_feature / 16) + + def validate(self) -> None: + if self.hid_feature <= 0: + raise ValueError("hid_feature must be positive") + if self.num_tfmamba <= 0: + raise ValueError("num_tfmamba must be positive") + if self.n_fft <= 0 or self.n_fft % 2 != 0: + raise ValueError("n_fft must be a positive even number") + if self.sampling_rate <= 0: + raise ValueError("sampling_rate must be positive") + for name in ("input_sampling_rate", "bwe_sampling_rate"): + value = getattr(self, name) + if value is not None and value <= 0: + raise ValueError(f"{name} must be positive when provided") + if self.input_sampling_rate is not None and self.bwe_sampling_rate is not None: + raise ValueError( + "input_sampling_rate and bwe_sampling_rate are mutually exclusive" + ) + + @classmethod + def from_json(cls, cfg: dict) -> ReUseConfig: + """Build a config from the parsed ``nvidia/RE-USE`` ``config.json``.""" + model_cfg = cfg.get("model_cfg", {}) + stft_cfg = cfg.get("stft_cfg", {}) + return cls( + input_channel=int(model_cfg.get("input_channel", 2)), + output_channel=int(model_cfg.get("output_channel", 1)), + hid_feature=int(model_cfg.get("hid_feature", 64)), + num_tfmamba=int(model_cfg.get("num_tfmamba", 30)), + d_state=int(model_cfg.get("d_state", 16)), + d_conv=int(model_cfg.get("d_conv", 4)), + expand=int(model_cfg.get("expand", 4)), + norm_epsilon=float(model_cfg.get("norm_epsilon", 1e-5)), + n_fft=int(stft_cfg.get("n_fft", 320)), + hop_size=int(stft_cfg.get("hop_size", 40)), + win_size=int(stft_cfg.get("win_size", 320)), + sampling_rate=int(stft_cfg.get("sampling_rate", 8000)), + compress_factor=str(model_cfg.get("compress_factor", "relu_log1p")), + model_type="reuse", + ) + + @classmethod + def from_pretrained( + cls, + model_id: str = "nvidia/RE-USE", + *, + revision: str | None = None, + ) -> ReUseConfig: + """Read ``config.json`` from a local directory or the HuggingFace Hub. + + RE-USE ships a bespoke ``config.json`` with no ``model_type`` or + ``architectures`` field, so ``transformers.AutoConfig`` cannot read + it and the generic :func:`mobius.build` entry point does not apply. + """ + return cls.from_json(_load_config_dict(model_id, revision)) + + +# --------------------------------------------------------------------------- +# Low-level ONNX helpers +# --------------------------------------------------------------------------- + + +class _InstanceNorm2d(nn.Module): + """``torch.nn.InstanceNorm2d(affine=True, track_running_stats=False)``. + + ONNX's ``InstanceNormalization`` normalizes each (batch, channel) plane + over its spatial extent using the batch's own statistics, which is + exactly PyTorch's behaviour when running statistics are disabled. + """ + + def __init__(self, num_features: int, eps: float = 1e-5): + super().__init__() + self.weight = nn.Parameter([num_features]) + self.bias = nn.Parameter([num_features]) + self._eps = eps + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + return op.InstanceNormalization(x, self.weight, self.bias, epsilon=self._eps) + + +class _PReLU2d(nn.Module): + """``torch.nn.PReLU(num_parameters=channels)`` for NCHW activations.""" + + def __init__(self, num_parameters: int): + super().__init__() + self.weight = nn.Parameter([num_parameters]) + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + # slope: (C,) → (C, 1, 1) so it broadcasts across H and W but not N. + return op.PRelu(x, op.Unsqueeze(self.weight, [-1, -2])) + + +def _norm_act_stage(conv: nn.Module, channels: int, eps: float) -> nn.ModuleList: + """``Sequential(conv, InstanceNorm2d, PReLU)`` shared by encoder/decoder. + + A plain :class:`~onnxscript.nn.ModuleList` (rather than a wrapper module) + keeps the checkpoint's ``nn.Sequential`` indices — ``.0``/``.1``/``.2`` — + as the parameter-name segments, so no weight renaming is needed. + """ + return nn.ModuleList([conv, _InstanceNorm2d(channels, eps), _PReLU2d(channels)]) + + +def _apply_stage(op: OpBuilder, stage: nn.ModuleList, x: ir.Value) -> ir.Value: + """Run a :func:`_norm_act_stage` list in order.""" + for layer in stage: + x = layer(op, x) + return x + + +def _atan2(op: OpBuilder, y: ir.Value, x: ir.Value) -> ir.Value: + """Two-argument arctangent, matching ``torch.atan2``. + + ONNX has no ``Atan2``, so the quadrant correction is spelled out: + ``atan(y/x)`` is only valid for ``x > 0``, needs a ``±pi`` shift for + ``x < 0``, and degenerates to ``±pi/2`` on the ``x == 0`` axis. + + Signed zeros are not distinguished: ``atan2(-0.0, x<0)`` returns ``+pi`` + where IEEE specifies ``-pi``. Both name the same angle, so the phase + decoder's ``cos``/``sin`` consumers cannot tell them apart. + """ + zero = op.CastLike(op.Constant(value_float=0.0), x) + one = op.CastLike(op.Constant(value_float=1.0), x) + pi = op.CastLike(op.Constant(value_float=math.pi), x) + half_pi = op.CastLike(op.Constant(value_float=math.pi / 2.0), x) + + x_is_zero = op.Equal(x, zero) + # Substitute 1 for x where it is zero so the division never produces a + # NaN/Inf that a later Where could not mask out. + safe_x = op.Where(x_is_zero, one, x) + base = op.Atan(op.Div(y, safe_x)) + + # On the x == 0 axis the angle is sign(y) * pi/2 (and 0 at the origin, + # which Sign(0) == 0 gives for free). + base = op.Where(x_is_zero, op.Mul(op.Sign(y), half_pi), base) + + # For x < 0, atan(y/x) lands in the wrong half-plane: shift by +pi when + # y >= 0 and by -pi otherwise. + shift = op.Where( + op.Less(x, zero), + op.Where(op.GreaterOrEqual(y, zero), pi, op.Neg(pi)), + zero, + ) + return op.Add(base, shift) + + +class _SPConvTranspose2d(nn.Module): + """Sub-pixel "transposed" convolution used by both decoders. + + Produces ``out_channels * r`` maps with an ordinary convolution and then + interleaves the ``r`` copies along the last axis, expanding it by ``r``. + The input is first padded by one column on each side of the last axis so + the ``(1, 3)`` kernel preserves that axis before expansion. + """ + + def __init__(self, in_channels: int, out_channels: int, kernel_size, r: int): + super().__init__() + self.conv = Conv2d(in_channels, out_channels * r, kernel_size) + self._out_channels = out_channels + self._r = r + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + # Pad last axis by 1 on each side: pads = [begin..., end...] over 4 dims. + x = op.Pad(x, op.Constant(value_ints=[0, 0, 0, 1, 0, 0, 0, 1])) + out = self.conv(op, x) # (B, out_channels * r, H, W) + + batch = op.Shape(out, start=0, end=1) + height = op.Shape(out, start=2, end=3) + width = op.Shape(out, start=3, end=4) + + # (B, r, C, H, W) -> (B, C, H, W, r) -> (B, C, H, W * r). + out = op.Reshape( + out, + op.Concat( + batch, + op.Constant(value_ints=[self._r, self._out_channels]), + height, + width, + axis=0, + ), + ) + out = op.Transpose(out, perm=[0, 2, 3, 4, 1]) + return op.Reshape( + out, + op.Concat( + batch, + op.Constant(value_ints=[self._out_channels]), + height, + op.Constant(value_ints=[-1]), + axis=0, + ), + ) + + +# --------------------------------------------------------------------------- +# Dense encoder / decoders +# --------------------------------------------------------------------------- + + +class DenseBlock(nn.Module): + """Densely connected stack of dilated 2D convolutions. + + Each step convolves the concatenation of all previous outputs, so the + ``i``-th convolution reads ``hid_feature * (i + 1)`` channels. Dilation + doubles per step along the time axis, widening the receptive field + without changing the spatial dimensions. + """ + + def __init__(self, config: ReUseConfig, depth: int = 4): + super().__init__() + self._depth = depth + hid = config.hid_feature + blocks = [] + for i in range(depth): + dilation = 2**i + conv = Conv2d( + hid * (i + 1), + hid, + kernel_size=(3, 3), + # ONNX pad order is [top, left, bottom, right]; the reference + # uses PyTorch padding=(dilation, 1) on a (3, 3) kernel. + padding=(dilation, 1, dilation, 1), + dilation=(dilation, 1), + ) + blocks.append(_norm_act_stage(conv, hid, config.norm_epsilon)) + self.dense_block = nn.ModuleList(blocks) + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + skip = x + for block in self.dense_block: + # x: (B, hid, H, W); skip grows by hid channels each step. + x = _apply_stage(op, block, skip) + skip = op.Concat(x, skip, axis=1) + return x + + +class DenseEncoder(nn.Module): + """Channel lift → dense block → strided downsample. + + ``dense_conv_2`` strides the time axis by 4 and the frequency axis by 2, + which is what makes the Mamba stack affordable. + """ + + def __init__(self, config: ReUseConfig): + super().__init__() + hid = config.hid_feature + eps = config.norm_epsilon + self.dense_conv_1 = _norm_act_stage( + Conv2d(config.input_channel, hid, kernel_size=(1, 1)), hid, eps + ) + self.dense_block = DenseBlock(config, depth=config.dense_depth) + self.dense_conv_2 = _norm_act_stage( + Conv2d( + hid, + hid, + kernel_size=(1, _ENCODER_FREQ_KERNEL), + stride=(4, _ENCODER_FREQ_STRIDE), + ), + hid, + eps, + ) + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + x = _apply_stage(op, self.dense_conv_1, x) # (B, hid, T, F) + x = self.dense_block(op, x) # (B, hid, T, F) + return _apply_stage(op, self.dense_conv_2, x) # (B, hid, T', F') + + +class _Decoder(nn.Module): + """Shared decoder trunk: dense block → freq upsample → time upsample. + + ``up_conv1`` expands the frequency axis by 2; ``up_conv2`` is applied to + a time/frequency-transposed view so that it expands the time axis by 4, + undoing the encoder's ``stride=(4, 2)``. + """ + + def __init__(self, config: ReUseConfig): + super().__init__() + hid = config.hid_feature + eps = config.norm_epsilon + self.dense_block = DenseBlock(config, depth=config.dense_depth) + self.up_conv1 = _norm_act_stage(_SPConvTranspose2d(hid, hid, (1, 3), 2), hid, eps) + self.up_conv2 = _norm_act_stage(_SPConvTranspose2d(hid, hid, (1, 3), 4), hid, eps) + + def _trunk(self, op: OpBuilder, x: ir.Value) -> ir.Value: + x = self.dense_block(op, x) + x = _apply_stage(op, self.up_conv1, x) # (B, hid, T', F' * 2) + # Swap time and frequency so up_conv2 expands the time axis, then + # swap back. + x = op.Transpose(x, perm=[0, 1, 3, 2]) + x = _apply_stage(op, self.up_conv2, x) # (B, hid, F' * 2, T' * 4) + return op.Transpose(x, perm=[0, 1, 3, 2]) + + +class MagDecoder(_Decoder): + """Decoder head producing the enhanced magnitude spectrogram.""" + + def __init__(self, config: ReUseConfig): + super().__init__(config) + self.final_conv = Conv2d(config.hid_feature, config.output_channel, kernel_size=(1, 1)) + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + return self.final_conv(op, self._trunk(op, x)) # (B, 1, T'*4, F'*2) + + +class PhaseDecoder(_Decoder): + """Decoder head producing the enhanced phase via ``atan2(imag, real)``.""" + + def __init__(self, config: ReUseConfig): + super().__init__(config) + self.phase_conv_r = Conv2d( + config.hid_feature, config.output_channel, kernel_size=(1, 1) + ) + self.phase_conv_i = Conv2d( + config.hid_feature, config.output_channel, kernel_size=(1, 1) + ) + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + x = self._trunk(op, x) + # Predict a unit-vector-like (real, imag) pair and read off its angle, + # which keeps the output inherently wrapped to (-pi, pi]. + return _atan2(op, self.phase_conv_i(op, x), self.phase_conv_r(op, x)) + + +# --------------------------------------------------------------------------- +# Bidirectional Mamba +# --------------------------------------------------------------------------- + + +class BiMambaBlock(nn.Module): + """Bidirectional Mamba over one axis, with a residual inside each branch. + + Runs the forward branch on the sequence and the backward branch on its + reverse, re-reverses the latter, concatenates both, projects back to + ``d_model``, and layer-normalizes. + """ + + def __init__(self, config: ReUseConfig): + super().__init__() + d_model = config.hid_feature + self.forward_blocks = SequenceMambaBlock( + d_model, + config.d_inner, + config.d_state, + config.dt_rank, + config.d_conv, + ) + self.backward_blocks = SequenceMambaBlock( + d_model, + config.d_inner, + config.d_state, + config.dt_rank, + config.d_conv, + ) + self.output_proj = Linear(2 * d_model, d_model, bias=True) + self.norm = LayerNorm(d_model, eps=config.norm_epsilon) + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + # x: (batch, seq_len, d_model) + out_fw = op.Add(self.forward_blocks(op, x), x) + + # Reverse along the sequence axis. ONNX has no Flip, and a + # negative-step Slice is the standard spelling for it. + # + # ReverseSequence would also work and needs no negative-step support, + # but it costs an extra Expand to build the per-row lengths and says + # something the model does not mean: nothing here is padded, so every + # row is reversed in full. Slice states the intent directly. + x_rev = op.Slice(x, [-1], [_INT64_MIN], [1], [-1]) + out_bw = op.Add(self.backward_blocks(op, x_rev), x_rev) + out_bw = op.Slice(out_bw, [-1], [_INT64_MIN], [1], [-1]) + + out = op.Concat(out_fw, out_bw, axis=-1) # (batch, seq_len, 2*d_model) + return self.norm(op, self.output_proj(op, out)) + + +class TFMambaBlock(nn.Module): + """Bidirectional Mamba along time, then along frequency. + + The feature map ``(B, C, T, F)`` is folded so that each axis in turn + becomes the sequence dimension of a :class:`BiMambaBlock`, with the other + axis absorbed into the batch. + """ + + def __init__(self, config: ReUseConfig): + super().__init__() + self.time_mamba = BiMambaBlock(config) + self.freq_mamba = BiMambaBlock(config) + self._channels = config.hid_feature + self._freq = config.static_encoder_freq_bins + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + channels = op.Constant(value_ints=[self._channels]) + batch = op.Shape(x, start=0, end=1) + time = op.Shape(x, start=2, end=3) + # An explicitly selected native/BWE rate gives an EP-claimable static + # frequency Scan. The faithful default follows the decoded native rate, + # so frequency is runtime data just as it is in NVIDIA's PyTorch model. + freq = ( + op.Constant(value_ints=[self._freq]) + if self._freq is not None + else op.Shape(x, start=3, end=4) + ) + minus_one = op.Constant(value_ints=[-1]) + + # --- Time branch: (B, C, T, F) -> (B*F, T, C) --- + h = op.Transpose(x, perm=[0, 3, 2, 1]) # (B, F, T, C) + h = op.Reshape(h, op.Concat(minus_one, time, channels, axis=0)) + h = op.Add(self.time_mamba(op, h), h) + + # --- Frequency branch: (B*F, T, C) -> (B*T, F, C) --- + h = op.Reshape(h, op.Concat(batch, freq, time, channels, axis=0)) + h = op.Transpose(h, perm=[0, 2, 1, 3]) # (B, T, F, C) + h = op.Reshape(h, op.Concat(minus_one, freq, channels, axis=0)) + h = op.Add(self.freq_mamba(op, h), h) + + # --- Back to (B, C, T, F) --- + h = op.Reshape(h, op.Concat(batch, time, freq, channels, axis=0)) + return op.Transpose(h, perm=[0, 3, 1, 2]) + + +# --------------------------------------------------------------------------- +# Top-level model +# --------------------------------------------------------------------------- + + +class SEMambaSpeechEnhancementModel(nn.Module): + """NVIDIA RE-USE (SEMamba) universal speech enhancement generator. + + Consumes the magnitude and phase of a noisy STFT and predicts the + enhanced magnitude, phase, and complex spectrogram. The magnitude is + expected to already be compressed by the STFT front-end (RE-USE uses + ``log1p``); decompression and the ISTFT happen after this graph. + """ + + default_task: str = "speech-enhancement" + category: str = "Audio" + + def __init__(self, config: ReUseConfig): + super().__init__() + self.config = config + self.dense_encoder = DenseEncoder(config) + self.TSMamba = nn.ModuleList([TFMambaBlock(config) for _ in range(config.num_tfmamba)]) + self.mask_decoder = MagDecoder(config) + self.phase_decoder = PhaseDecoder(config) + + def forward( + self, + op: OpBuilder, + noisy_mag: ir.Value, + noisy_pha: ir.Value, + ) -> tuple[ir.Value, ir.Value, ir.Value]: + """Enhance a noisy spectrogram. + + Args: + op: ONNX op builder. + noisy_mag: (batch, freq, time) — compressed noisy magnitude. + noisy_pha: (batch, freq, time) — noisy phase, in radians. + + Returns: + ``(denoised_mag, denoised_pha, denoised_com)`` with shapes + ``(batch, freq, time)``, ``(batch, freq, time)`` and + ``(batch, freq, time, 2)``. + """ + static_num_freq = self.config.static_num_freq_bins + # Remember the input time extent; the decoders overshoot it and the + # result is cropped back at the end. + time = op.Shape(noisy_mag, start=2, end=3) + + # (B, F, T) -> (B, 1, T, F) for each of magnitude and phase, then + # stack them as the encoder's two input channels. + mag = op.Unsqueeze(op.Transpose(noisy_mag, perm=[0, 2, 1]), [1]) + pha = op.Unsqueeze(op.Transpose(noisy_pha, perm=[0, 2, 1]), [1]) + x = op.Concat(mag, pha, axis=1) # (B, 2, T, F) + + # Zero-pad the tail of both the time and frequency axes. The reference + # does this so the strided encoder convolution never has to drop a + # partial window. The frequency pad is part of how + # ``ReUseConfig.encoder_freq_bins`` predicts the encoder's output extent. + x = op.Pad( + x, + op.Constant( + value_ints=[0, 0, 0, 0, 0, 0, _ENCODER_FREQ_TAIL_PAD, _ENCODER_FREQ_TAIL_PAD] + ), + ) + + x = self.dense_encoder(op, x) # (B, C, T', F') + for block in self.TSMamba: + x = block(op, x) + + # (B, 1, T'*4, F'*2) -> (B, F'*2, T'*4) + denoised_mag = op.Squeeze( + op.Transpose(self.mask_decoder(op, x), perm=[0, 3, 2, 1]), [-1] + ) + denoised_pha = op.Squeeze( + op.Transpose(self.phase_decoder(op, x), perm=[0, 3, 2, 1]), [-1] + ) + + # Crop the upsampled output back to the input (freq, time) extent. + starts = op.Constant(value_ints=[0, 0]) + num_freq = ( + op.Constant(value_ints=[static_num_freq]) + if static_num_freq is not None + else op.Shape(noisy_mag, start=1, end=2) + ) + ends = op.Concat(num_freq, time, axis=0) + axes = op.Constant(value_ints=[1, 2]) + denoised_mag = op.Slice(denoised_mag, starts, ends, axes) + denoised_pha = op.Slice(denoised_pha, starts, ends, axes) + + # Complex spectrogram as a trailing (real, imag) pair. + denoised_com = op.Concat( + op.Unsqueeze(op.Mul(denoised_mag, op.Cos(denoised_pha)), [-1]), + op.Unsqueeze(op.Mul(denoised_mag, op.Sin(denoised_pha)), [-1]), + axis=-1, + ) + + return denoised_mag, denoised_pha, denoised_com + + def preprocess_weights( + self, + state_dict: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + """Map ``nvidia/RE-USE`` checkpoint names onto this module tree. + + The checkpoint stores the selective-scan parameters flat on each + Mamba module (``forward_blocks.A_log``), while + :class:`~mobius.components.SequenceMambaBlock` nests them under an + ``ssm`` submodule (``forward_blocks.ssm.A_log``) — the same offset + the Mamba causal-LM models correct for. Everything else, including + the encoder/decoder ``nn.Sequential`` indices, already lines up. + + The rename is idempotent: a state dict that already uses the nested + names is returned unchanged, so re-running this (or loading an + already-converted checkpoint) cannot produce ``ssm.ssm.A_log``. + """ + renames = {} + for key in list(state_dict): + for param in _SSM_PARAMS: + suffix = f".{param}" + if not key.endswith(suffix): + continue + prefix = key[: -len(suffix)] + # Already nested — leave it alone. + if prefix.endswith(".ssm") or prefix == "ssm": + break + renames[key] = f"{prefix}.ssm{suffix}" + break + for old_key, new_key in renames.items(): + state_dict[new_key] = state_dict.pop(old_key) + return state_dict + + +#: Selective-scan parameters the checkpoint keeps flat on the Mamba module +#: but :class:`~mobius.components.SequenceMambaBlock` nests under ``ssm``. +_SSM_PARAMS: tuple[str, ...] = ( + "A_log", + "D", + "x_proj.weight", + "dt_proj.weight", + "dt_proj.bias", +) + + +def build_reuse( + model_id: str = "nvidia/RE-USE", + *, + revision: str | None = None, + dtype: str | None = None, + execution_provider: str = "default", + load_weights: bool = True, + input_sampling_rate: int | None = None, + bwe_sampling_rate: int | None = None, +) -> ModelPackage: + """Build an ONNX :class:`ModelPackage` for a RE-USE / SEMamba checkpoint. + + RE-USE is published with a bespoke ``config.json`` (no ``model_type``, + no ``architectures``) and a ``PyTorchModelHubMixin`` checkpoint, so + :func:`mobius.build` — which goes through ``transformers.AutoConfig`` — + cannot discover it. This function is the equivalent entry point. + + Args: + model_id: Local directory or HuggingFace Hub repo holding + ``config.json`` and ``model.safetensors``. + revision: Hub revision (branch, tag, or commit SHA) to pin downloads. + Defaults to the immutable revision in :data:`REUSE_REVISION` and + is ignored for local directories. + dtype: Override model dtype (e.g. ``"f16"``). Defaults to float32. + execution_provider: Target execution provider for EP-aware + optimizations. + load_weights: When false, build the graph structure only. + input_sampling_rate: Known native input rate for a static-frequency + export. Omit it to preserve NVIDIA's native-rate behavior with a + dynamic frequency axis. + bwe_sampling_rate: Explicit NVIDIA ``BWE`` target rate. The workflow + resamples audio to this rate before analysis. + + Returns: + A :class:`ModelPackage` whose ``"model"`` entry is the enhancement + network. + """ + if input_sampling_rate is not None and bwe_sampling_rate is not None: + raise ValueError("input_sampling_rate and bwe_sampling_rate are mutually exclusive") + + from mobius._builder import build_from_module, resolve_dtype + from mobius.integrations._weight_loading import apply_weights + + config = ReUseConfig.from_pretrained(model_id, revision=revision) + config.input_sampling_rate = input_sampling_rate + config.bwe_sampling_rate = bwe_sampling_rate + config.validate() + resolved = resolve_dtype(dtype) + if resolved is not None: + config.dtype = resolved + + module = SEMambaSpeechEnhancementModel(config) + package = build_from_module( + module, + config, + task="speech-enhancement", + execution_provider=execution_provider, + ) + if load_weights: + apply_weights( + package["model"], + module.preprocess_weights(_load_state_dict(model_id, revision)), + ) + for model in package.values(): + model.metadata_props["mobius.source_revision"] = _effective_revision( + model_id, revision + ) + return package + + +def _effective_revision(model_id: str, revision: str | None) -> str: + """Return the checkpoint revision recorded on exported artifacts.""" + if os.path.isdir(model_id): + return "local" + return revision or REUSE_REVISION + + +def _load_config_dict(model_id: str, revision: str | None) -> dict: + """Read the bespoke config from a local directory or pinned Hub revision.""" + local = os.path.join(model_id, "config.json") + if not os.path.isfile(local): + from huggingface_hub import hf_hub_download + + local = hf_hub_download( + model_id, + "config.json", + revision=_effective_revision(model_id, revision), + ) + with open(local, encoding="utf-8") as handle: + return json.load(handle) + + +def _is_reuse_checkpoint(model_id: str, revision: str | None = None) -> bool: + """Return whether a non-Transformers checkpoint has the RE-USE contract.""" + try: + config = _load_config_dict(model_id, revision) + except (OSError, ValueError, json.JSONDecodeError): + return False + model_cfg = config.get("model_cfg") + stft_cfg = config.get("stft_cfg") + return ( + isinstance(model_cfg, dict) + and isinstance(stft_cfg, dict) + and all( + name in model_cfg + for name in ("hid_feature", "num_tfmamba", "d_state", "d_conv", "expand") + ) + and all(name in stft_cfg for name in ("n_fft", "hop_size", "win_size")) + ) + + +def _load_state_dict(model_id: str, revision: str | None) -> dict[str, torch.Tensor]: + """Read ``model.safetensors`` from a local directory or the Hub.""" + from safetensors.torch import load_file + + local = os.path.join(model_id, "model.safetensors") + if not os.path.isfile(local): + from huggingface_hub import hf_hub_download + + local = hf_hub_download( + model_id, + "model.safetensors", + revision=_effective_revision(model_id, revision), + ) + return load_file(local) diff --git a/src/mobius/models/reuse_test.py b/src/mobius/models/reuse_test.py new file mode 100644 index 000000000..6273b3031 --- /dev/null +++ b/src/mobius/models/reuse_test.py @@ -0,0 +1,777 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for the RE-USE / SEMamba speech-enhancement model.""" + +from __future__ import annotations + +import math + +import numpy as np +import onnx_ir as ir +import pytest + +from mobius import build, build_from_module +from mobius.models.reuse import ( + REUSE_REVISION, + ReUseConfig, + SEMambaSpeechEnhancementModel, + _atan2, + _effective_revision, + build_reuse, +) +from mobius.tasks import SpeechEnhancementTask, get_task + +# The real nvidia/RE-USE config, scaled down: same structure, tiny widths. +_TINY_CONFIG = { + "model_cfg": { + "hid_feature": 8, + "num_tfmamba": 2, + "d_state": 4, + "d_conv": 4, + "expand": 2, + "input_channel": 2, + "output_channel": 1, + "norm_epsilon": 1e-5, + "compress_factor": "relu_log1p", + }, + "stft_cfg": {"n_fft": 32, "hop_size": 4, "win_size": 32, "sampling_rate": 8000}, +} + + +def _tiny_config() -> ReUseConfig: + return ReUseConfig.from_json(_TINY_CONFIG) + + +def _build(): + config = _tiny_config() + module = SEMambaSpeechEnhancementModel(config) + return config, build_from_module(module, config, task=SpeechEnhancementTask()) + + +class TestReUseConfig: + """Config extraction from the nvidia/RE-USE config.json layout.""" + + def test_from_json_reads_both_sections(self): + config = _tiny_config() + assert config.hid_feature == 8 + assert config.num_tfmamba == 2 + assert config.d_state == 4 + assert config.expand == 2 + assert config.n_fft == 32 + assert config.hop_size == 4 + assert config.sampling_rate == 8000 + assert config.model_type == "reuse" + + def test_from_json_defaults_match_the_dataclass(self): + """An absent field must fall back to the same value the dataclass declares. + + ``from_json`` previously defaulted ``num_tfmamba`` to 4 while the dataclass + said 30, so a config.json missing that key would silently build a model an + eighth of the real depth — and still load, because every layer is + independently named. + """ + from_empty = ReUseConfig.from_json({}) + declared = ReUseConfig() + for field in ( + "input_channel", + "output_channel", + "hid_feature", + "num_tfmamba", + "d_state", + "d_conv", + "expand", + "norm_epsilon", + "n_fft", + "hop_size", + "win_size", + "sampling_rate", + "input_sampling_rate", + "bwe_sampling_rate", + "compress_factor", + ): + assert getattr(from_empty, field) == getattr(declared, field), field + + def test_compress_factor_is_read_from_model_cfg(self): + """It describes a front-end step but is declared under ``model_cfg`` upstream. + + Checked against the published nvidia/RE-USE config.json, where + ``compress_factor`` sits in ``model_cfg`` alongside ``hid_feature``, not in + ``stft_cfg`` with ``n_fft``. Reading it from ``stft_cfg`` would silently + fall back to the default for the real checkpoint. + """ + config = ReUseConfig.from_json( + { + "model_cfg": {"compress_factor": "custom"}, + "stft_cfg": {"compress_factor": "wrong"}, + } + ) + assert config.compress_factor == "custom" + + def test_real_checkpoint_shape_parameters(self): + """The published config yields the checkpoint's real dimensions.""" + config = ReUseConfig.from_json( + { + "model_cfg": { + "hid_feature": 64, + "num_tfmamba": 30, + "d_state": 16, + "d_conv": 4, + "expand": 4, + }, + "stft_cfg": {"n_fft": 320, "hop_size": 40, "win_size": 320}, + } + ) + # in_proj is (2 * d_inner, hid_feature) = (512, 64) in the checkpoint. + assert config.d_inner == 256 + # x_proj is (dt_rank + 2 * d_state, d_inner) = (36, 256). + assert config.dt_rank == 4 + assert config.dt_rank + 2 * config.d_state == 36 + assert config.num_freq_bins == 161 + + def test_dt_rank_follows_mamba_convention(self): + for hid in (8, 64, 100): + config = ReUseConfig(hid_feature=hid) + assert config.dt_rank == math.ceil(hid / 16) + + def test_validate_rejects_odd_n_fft(self): + with pytest.raises(ValueError, match="n_fft"): + ReUseConfig(n_fft=33).validate() + + def test_scaled_geometry_rejects_rate_too_small_for_hop(self): + with pytest.raises(ValueError, match="too small"): + ReUseConfig().stft_geometry(1) + + @pytest.mark.parametrize( + ("input_rate", "bwe_rate"), + [(8_000, 16_000), (16_000, 16_000)], + ) + def test_validate_rejects_ambiguous_native_and_bwe_rates(self, input_rate, bwe_rate): + with pytest.raises(ValueError, match="mutually exclusive"): + ReUseConfig( + input_sampling_rate=input_rate, + bwe_sampling_rate=bwe_rate, + ).validate() + + def test_remote_default_is_pinned(self): + assert _effective_revision("nvidia/RE-USE", None) == REUSE_REVISION + + +class TestBuildGraphReUse: + """Graph construction for the SEMamba generator.""" + + def test_task_registry_lookup(self): + assert isinstance(get_task("speech-enhancement"), SpeechEnhancementTask) + + def test_default_task(self): + assert SEMambaSpeechEnhancementModel.default_task == "speech-enhancement" + + def test_package_has_single_model(self): + _config, pkg = _build() + assert set(pkg) == {"model"} + + def test_model_io(self): + _config, pkg = _build() + graph = pkg["model"].graph + + assert [inp.name for inp in graph.inputs] == ["noisy_mag", "noisy_pha"] + assert [out.name for out in graph.outputs] == [ + "denoised_mag", + "denoised_pha", + "denoised_com", + ] + # Native-rate export follows the decoded sample rate; all three axes + # are dynamic while the complex pair remains fixed. + for value in (*graph.inputs, *graph.outputs[:2]): + assert str(value.shape[1]) == "freq" + assert graph.outputs[2].shape[3] == 2 + + def test_explicit_input_rate_makes_frequency_static(self): + config = _tiny_config() + config.input_sampling_rate = 16_000 + pkg = build_from_module( + SEMambaSpeechEnhancementModel(config), + config, + task=SpeechEnhancementTask(), + ) + + assert config.analysis_stft_geometry == (64, 8, 64) + for value in (*pkg["model"].graph.inputs, *pkg["model"].graph.outputs[:2]): + assert value.shape[1] == 33 + + def test_initializer_names_match_checkpoint_layout(self): + """Parameter names line up with the nvidia/RE-USE checkpoint.""" + config = _tiny_config() + module = SEMambaSpeechEnhancementModel(config) + names = {name for name, _ in module.named_parameters()} + + # nn.Sequential stages keep their integer indices. + assert "dense_encoder.dense_conv_1.0.weight" in names + assert "dense_encoder.dense_conv_1.1.bias" in names + assert "dense_encoder.dense_conv_1.2.weight" in names + assert "dense_encoder.dense_block.dense_block.3.0.weight" in names + assert "mask_decoder.up_conv1.0.conv.weight" in names + assert "mask_decoder.final_conv.weight" in names + assert "phase_decoder.phase_conv_r.weight" in names + assert "phase_decoder.phase_conv_i.weight" in names + # Mamba parameters keep the checkpoint's block names. + assert "TSMamba.0.time_mamba.forward_blocks.in_proj.weight" in names + assert "TSMamba.0.freq_mamba.backward_blocks.conv1d.bias" in names + assert "TSMamba.1.time_mamba.output_proj.bias" in names + assert "TSMamba.1.freq_mamba.norm.weight" in names + + def test_dense_block_dilation_grows_by_powers_of_two(self): + """The i-th dense conv reads i+1 stacked feature maps at dilation 2**i.""" + config = _tiny_config() + module = SEMambaSpeechEnhancementModel(config) + block = module.dense_encoder.dense_block + for i in range(config.dense_depth): + conv = block.dense_block[i][0] + assert list(conv.weight.shape) == [ + config.hid_feature, + config.hid_feature * (i + 1), + 3, + 3, + ] + assert conv._dilations == (2**i, 1) + assert conv._pads == [2**i, 1, 2**i, 1] + + def test_preprocess_weights_nests_ssm_parameters(self): + """Flat checkpoint SSM parameters move under the ``ssm`` submodule.""" + config = _tiny_config() + module = SEMambaSpeechEnhancementModel(config) + prefix = "TSMamba.0.time_mamba.forward_blocks" + state_dict = { + f"{prefix}.A_log": "a", + f"{prefix}.D": "d", + f"{prefix}.x_proj.weight": "x", + f"{prefix}.dt_proj.weight": "w", + f"{prefix}.dt_proj.bias": "b", + f"{prefix}.in_proj.weight": "i", + f"{prefix}.conv1d.bias": "c", + } + + result = module.preprocess_weights(state_dict) + + assert result[f"{prefix}.ssm.A_log"] == "a" + assert result[f"{prefix}.ssm.dt_proj.bias"] == "b" + # Non-SSM parameters are untouched. + assert result[f"{prefix}.in_proj.weight"] == "i" + assert result[f"{prefix}.conv1d.bias"] == "c" + assert f"{prefix}.A_log" not in result + + def test_preprocess_weights_is_idempotent(self): + """Already-nested names must survive a second pass unchanged. + + Renaming on suffix alone would turn ``ssm.A_log`` into + ``ssm.ssm.A_log``, silently dropping every SSM parameter when a + converted state dict is preprocessed again. + """ + config = _tiny_config() + module = SEMambaSpeechEnhancementModel(config) + onnx_names = {name for name, _ in module.named_parameters()} + checkpoint_names = {name.replace(".ssm.", ".") for name in onnx_names} + + once = module.preprocess_weights(dict.fromkeys(checkpoint_names, 0)) + twice = module.preprocess_weights(dict(once)) + + assert set(once) == onnx_names + assert set(twice) == onnx_names + assert not any(".ssm.ssm." in name for name in twice) + + def test_preprocess_weights_covers_every_parameter(self): + """Renaming a checkpoint-shaped state dict yields exactly our names.""" + config = _tiny_config() + module = SEMambaSpeechEnhancementModel(config) + onnx_names = {name for name, _ in module.named_parameters()} + # Reconstruct the checkpoint's flat naming from ours. + checkpoint_names = {name.replace(".ssm.", ".") for name in onnx_names} + + renamed = module.preprocess_weights(dict.fromkeys(checkpoint_names, 0)) + + assert set(renamed) == onnx_names + + def test_scan_count_matches_mamba_module_count(self): + """One Scan per Mamba module: blocks x 2 axes x 2 directions.""" + config = _tiny_config() + _config, pkg = _build() + scans = sum(1 for node in pkg["model"].graph if node.op_type == "Scan") + assert scans == config.num_tfmamba * 4 + + +class TestReUseRuntime: + """End-to-end ONNX Runtime execution with random weights.""" + + def _session(self): + ort = pytest.importorskip("onnxruntime") + config = _tiny_config() + module = SEMambaSpeechEnhancementModel(config) + + rng = np.random.default_rng(0) + for _name, param in module.named_parameters(): + shape = [d if isinstance(d, int) else 1 for d in param.shape] + param.const_value = ir.tensor( + (rng.standard_normal(shape) * 0.05).astype(np.float32) + ) + + pkg = build_from_module(module, config, task=SpeechEnhancementTask()) + session = ort.InferenceSession( + ir.to_proto(pkg["model"]).SerializeToString(), + providers=["CPUExecutionProvider"], + ) + return config, session + + @pytest.mark.parametrize("time_steps", [1, 5, 17]) + def test_output_shapes_track_input_length(self, time_steps): + """The model is spectrally shape preserving for any input length.""" + config, session = self._session() + rng = np.random.default_rng(1) + shape = (2, config.num_freq_bins, time_steps) + mag = np.abs(rng.standard_normal(shape)).astype(np.float32) + pha = (rng.standard_normal(shape) * np.pi).astype(np.float32) + + denoised_mag, denoised_pha, denoised_com = session.run( + None, {"noisy_mag": mag, "noisy_pha": pha} + ) + + assert denoised_mag.shape == shape + assert denoised_pha.shape == shape + assert denoised_com.shape == (*shape, 2) + assert np.isfinite(denoised_mag).all() + assert np.isfinite(denoised_pha).all() + + @pytest.mark.parametrize("freq_bins", [17, 33, 97]) + def test_one_native_rate_graph_accepts_scaled_frequency_extents(self, freq_bins): + """The default graph follows 8/16/48 kHz FFT geometry at runtime.""" + _config, session = self._session() + rng = np.random.default_rng(freq_bins) + shape = (1, freq_bins, 5) + outputs = session.run( + None, + { + "noisy_mag": np.abs(rng.standard_normal(shape)).astype(np.float32), + "noisy_pha": rng.standard_normal(shape).astype(np.float32), + }, + ) + assert [output.shape for output in outputs] == [shape, shape, (*shape, 2)] + + def test_phase_is_wrapped_and_consistent_with_complex_output(self): + """Phase stays in (-pi, pi] and denoised_com is its polar form.""" + config, session = self._session() + rng = np.random.default_rng(2) + shape = (1, config.num_freq_bins, 7) + mag = np.abs(rng.standard_normal(shape)).astype(np.float32) + pha = (rng.standard_normal(shape) * np.pi).astype(np.float32) + + denoised_mag, denoised_pha, denoised_com = session.run( + None, {"noisy_mag": mag, "noisy_pha": pha} + ) + + assert np.abs(denoised_pha).max() <= np.pi + 1e-5 + np.testing.assert_allclose( + denoised_com[..., 0], denoised_mag * np.cos(denoised_pha), atol=1e-5 + ) + np.testing.assert_allclose( + denoised_com[..., 1], denoised_mag * np.sin(denoised_pha), atol=1e-5 + ) + + +class TestAtan2: + """The hand-rolled atan2 (ONNX has no Atan2 op).""" + + def test_matches_numpy_over_all_quadrants(self): + ort = pytest.importorskip("onnxruntime") + session = self._session() + + # Cover all four quadrants plus both axes and the origin. + grid = np.array([-2.0, -1.0, 0.0, 0.5, 2.0], dtype=np.float32) + ys, xs = (a.ravel() for a in np.meshgrid(grid, grid)) + (got,) = session.run(None, {"y": ys, "x": xs}) + + np.testing.assert_allclose(got, np.arctan2(ys, xs), atol=1e-6) + assert ort is not None + + def test_negative_zero_numerator_picks_the_positive_branch(self): + """Signed zeros collapse: -0.0 is treated as +0.0. + + IEEE distinguishes ``atan2(-0.0, -1) == -pi`` from + ``atan2(+0.0, -1) == +pi``. This implementation returns ``+pi`` for + both, which is the same angle and therefore indistinguishable to the + ``cos``/``sin`` consumers downstream. + """ + pytest.importorskip("onnxruntime") + session = self._session() + + ys = np.array([-0.0, 0.0], dtype=np.float32) + xs = np.array([-1.0, -1.0], dtype=np.float32) + (got,) = session.run(None, {"y": ys, "x": xs}) + + np.testing.assert_allclose(got, [np.pi, np.pi], atol=1e-6) + # Same point on the unit circle either way. + np.testing.assert_allclose(np.cos(got), np.cos(np.arctan2(ys, xs)), atol=1e-6) + np.testing.assert_allclose(np.sin(got), np.sin(np.arctan2(ys, xs)), atol=1e-6) + + def _session(self): + import onnxruntime as ort + from onnxscript import GraphBuilder + + from mobius._constants import OPSET_VERSION + + graph = ir.Graph([], [], nodes=[], name="g", opset_imports={"": OPSET_VERSION}) + builder = GraphBuilder(graph) + y = builder.input("y", dtype=ir.DataType.FLOAT, shape=["n"]) + x = builder.input("x", dtype=ir.DataType.FLOAT, shape=["n"]) + builder.add_output(_atan2(builder.op, y, x), "out") + + return ort.InferenceSession( + ir.to_proto(ir.Model(graph, ir_version=11)).SerializeToString(), + providers=["CPUExecutionProvider"], + ) + + +class TestBuildReUse: + """The bespoke Hub/directory loader (RE-USE has no transformers config).""" + + def _checkpoint_dir(self, tmp_path): + """Write a tiny RE-USE-shaped repo: config.json + model.safetensors.""" + import json + + import torch + from safetensors.torch import save_file + + (tmp_path / "config.json").write_text(json.dumps(_TINY_CONFIG)) + + # The checkpoint keeps SSM parameters flat, so undo our ``ssm`` nesting. + module = SEMambaSpeechEnhancementModel(_tiny_config()) + state = { + name.replace(".ssm.", "."): torch.zeros( + [d if isinstance(d, int) else 1 for d in param.shape] + ) + for name, param in module.named_parameters() + } + save_file(state, str(tmp_path / "model.safetensors")) + return tmp_path + + def test_reads_config_from_a_directory(self, tmp_path): + pytest.importorskip("safetensors") + config = ReUseConfig.from_pretrained(str(self._checkpoint_dir(tmp_path))) + + assert config.hid_feature == 8 + assert config.num_tfmamba == 2 + assert config.n_fft == 32 + + def test_builds_and_fills_every_initializer(self, tmp_path): + """A full local build leaves no initializer unpopulated.""" + pytest.importorskip("safetensors") + pkg = build_reuse(str(self._checkpoint_dir(tmp_path))) + + graph = pkg["model"].graph + unfilled = [ + name for name, value in graph.initializers.items() if value.const_value is None + ] + assert unfilled == [] + assert [inp.name for inp in graph.inputs] == ["noisy_mag", "noisy_pha"] + + def test_structure_only_build_skips_weights(self, tmp_path): + pytest.importorskip("safetensors") + pkg = build_reuse(str(self._checkpoint_dir(tmp_path)), load_weights=False) + + assert "model" in pkg + + @pytest.mark.parametrize( + ("kwargs", "expected_bins"), + [ + ({"input_sampling_rate": 16_000}, 33), + ({"bwe_sampling_rate": 48_000}, 97), + ], + ) + def test_build_exposes_static_native_and_bwe_rates(self, tmp_path, kwargs, expected_bins): + pytest.importorskip("safetensors") + pkg = build_reuse( + str(self._checkpoint_dir(tmp_path)), + load_weights=False, + **kwargs, + ) + + for value in pkg["model"].graph.inputs: + assert value.shape[1] == expected_bins + + @pytest.mark.parametrize( + ("input_rate", "bwe_rate"), + [(8_000, 16_000), (16_000, 16_000)], + ) + def test_build_reuse_rejects_both_rate_modes( + self, + input_rate, + bwe_rate, + ): + with pytest.raises(ValueError, match="mutually exclusive"): + build_reuse( + "config-must-not-be-read", + load_weights=False, + input_sampling_rate=input_rate, + bwe_sampling_rate=bwe_rate, + ) + + def test_public_build_detects_bespoke_checkpoint(self, tmp_path): + """The normal API/CLI path detects RE-USE without Transformers metadata.""" + pytest.importorskip("safetensors") + checkpoint = self._checkpoint_dir(tmp_path) + + pkg = build(str(checkpoint), load_weights=False) + + assert set(pkg) == {"model"} + assert pkg["model"].metadata_props["mobius.source_revision"] == "local" + + def test_public_build_forwards_native_rate_selection(self, tmp_path): + pytest.importorskip("safetensors") + checkpoint = self._checkpoint_dir(tmp_path) + + pkg = build( + str(checkpoint), + load_weights=False, + input_sampling_rate=16_000, + ) + + assert pkg.config.input_sampling_rate == 16_000 + assert pkg["model"].graph.inputs[0].shape[1] == 33 + + @pytest.mark.parametrize( + ("input_rate", "bwe_rate"), + [(8_000, 16_000), (16_000, 16_000)], + ) + def test_public_build_rejects_both_rate_modes( + self, + monkeypatch, + input_rate, + bwe_rate, + ): + from mobius.integrations.transformers import _builder + + def _unexpected_probe(*_args, **_kwargs): + raise AssertionError("config probe must not run for ambiguous rate modes") + + monkeypatch.setattr(_builder, "_load_transformers_config", _unexpected_probe) + + with pytest.raises(ValueError, match="mutually exclusive"): + build( + "config-must-not-be-read", + load_weights=False, + input_sampling_rate=input_rate, + bwe_sampling_rate=bwe_rate, + ) + + def test_canonical_detection_is_pinned_before_auto_config(self, monkeypatch): + from mobius.integrations.transformers import _builder + + observed = {} + + def _capture_probe(_model_id, *, revision, trust_remote_code): + observed["revision"] = revision + raise RuntimeError("stop after first probe") + + monkeypatch.setattr(_builder, "_load_transformers_config", _capture_probe) + + with pytest.raises(RuntimeError, match="stop after first probe"): + build("nvidia/RE-USE", load_weights=False) + assert observed["revision"] == REUSE_REVISION + + @pytest.mark.parametrize( + "kwargs", + [{"input_sampling_rate": 16_000}, {"bwe_sampling_rate": 48_000}], + ) + def test_rate_selection_is_rejected_for_transformers_models(self, monkeypatch, kwargs): + import types + + from mobius.integrations.transformers import _builder + + monkeypatch.setattr( + _builder, + "_load_transformers_config", + lambda *_args, **_kwargs: (types.SimpleNamespace(model_type="llama"), False), + ) + + with pytest.raises(ValueError, match="only supported for RE-USE"): + _builder.build_transformers_model("example/llama", load_weights=False, **kwargs) + + def test_public_build_accepts_task_object(self, tmp_path): + pytest.importorskip("safetensors") + checkpoint = self._checkpoint_dir(tmp_path) + + pkg = build( + str(checkpoint), + task=SpeechEnhancementTask(), + load_weights=False, + ) + + assert set(pkg) == {"model"} + + def test_local_revision_metadata_never_claims_a_remote_revision(self, tmp_path): + checkpoint = self._checkpoint_dir(tmp_path) + assert _effective_revision(str(checkpoint), "unrelated-remote-sha") == "local" + + def test_norm_epsilon_reaches_every_normalization(self): + """``norm_epsilon`` must govern the LayerNorms too, not just the InstanceNorms. + + The BiMamba LayerNorm hardcoded 1e-5, which happens to equal both PyTorch's + default and the published config's value, so nothing diverged in practice — + but a config setting a different epsilon would have been half-applied. + """ + config = _tiny_config() + config.norm_epsilon = 3e-3 + module = SEMambaSpeechEnhancementModel(config) + pkg = build_from_module(module, config, task=SpeechEnhancementTask()) + + epsilons = { + node.attributes["epsilon"].as_float() + for node in pkg["model"].graph + if node.op_type in ("LayerNormalization", "InstanceNormalization") + and "epsilon" in node.attributes + } + assert epsilons, "no normalization nodes found" + assert all(eps == pytest.approx(3e-3) for eps in epsilons), ( + f"some normalizations ignore config.norm_epsilon: {sorted(epsilons)}" + ) + + +class TestEncoderFreqBins: + """Frequency is dynamic by default and static for an explicit rate. + + NVIDIA scales FFT geometry from the decoded native sample rate. The + default graph therefore reads frequency at runtime. Explicit native/BWE + exports retain a constant extent for provider partitioning. + """ + + @pytest.mark.parametrize("n_fft", [320, 400, 512, 322]) + def test_matches_the_conv_geometry(self, n_fft): + """Derivation agrees with the pad-then-strided-conv arithmetic. + + ``322`` is included on purpose: it is the ``n_fft % 4 == 2`` case, the + only one where the floor division actually truncates. + """ + config = ReUseConfig(n_fft=n_fft) + padded = config.num_freq_bins + 2 # tail pad on the frequency axis + expected = (padded - 3) // 2 + 1 # kernel 3, stride 2, no padding + assert config.encoder_freq_bins == expected + + @pytest.mark.parametrize("n_fft", [320, 400, 512, 322]) + def test_decoder_can_cover_the_input_extent(self, n_fft): + """``up_conv1`` doubles the extent, which must reach the input width. + + ``forward`` crops the overshoot, so the requirement is ``>=``. If this + ever became ``<`` the model would silently return a truncated spectrum. + """ + config = ReUseConfig(n_fft=n_fft) + assert 2 * config.encoder_freq_bins >= config.num_freq_bins + + @pytest.mark.parametrize( + ("sample_rate", "expected"), + [(8_000, (320, 40, 320)), (16_000, (640, 80, 640)), (48_000, (1920, 240, 1920))], + ) + def test_scaled_geometry_matches_pinned_inference(self, sample_rate, expected): + assert ReUseConfig().stft_geometry(sample_rate) == expected + + def test_native_graph_reads_frequency_at_runtime(self): + _config, pkg = _build() + scans = [ + node + for node in pkg["model"].graph + if node.op_type == "Scan" and "freq_mamba" in (node.name or "") + ] + assert scans, "no frequency-axis Scan found" + for scan in scans: + swept = scan.inputs[2].shape + assert swept is not None + assert not isinstance(swept[0], int) + + def test_explicit_rate_emits_constant_frequency_extent(self): + config = _tiny_config() + config.input_sampling_rate = 16_000 + pkg = build_from_module( + SEMambaSpeechEnhancementModel(config), + config, + task=SpeechEnhancementTask(), + ) + graph = pkg["model"].graph + producers = {v.name: node for node in graph for v in node.outputs} + for node in graph: + if node.op_type != "Scan" or "freq_mamba" not in (node.name or ""): + continue + # Walk back from the swept input; the extent must originate in a + # Constant, never a Shape read of the encoder output. + seen: set[str] = set() + frontier = [node.inputs[2]] + saw_shape_on_freq_axis = False + while frontier: + value = frontier.pop() + if value is None or value.name in seen: + continue + seen.add(value.name) + producer = producers.get(value.name) + if producer is None or producer.op_type == "Scan": + continue + if ( + producer.op_type == "Shape" + and producer.attributes.get("start") is not None + ): + if producer.attributes["start"].as_int() == 3: + saw_shape_on_freq_axis = True + frontier.extend(producer.inputs) + assert not saw_shape_on_freq_axis, ( + "an explicitly selected rate should make frequency static" + ) + + +class TestExecutionProviderPartitioning: + """Lock in graph shapes that plugin EPs need in order to claim the SSM. + + Both assertions look like arbitrary spelling choices. They are not: each was + measured on the MLX plugin EP (Apple M1 Max, 2 s of audio at 8 kHz, steady + state), and getting either wrong costs far more than the transposes and + slices they buy. + + Quote steady-state numbers only when citing this model's throughput. The + first run is roughly 4x slower — one-time graph translation and kernel + compilation — and the gap widens with the frequency extent, so a + median-of-three silently reports a warm-up figure. + """ + + def test_scan_iterates_axis_zero(self): + """Scan must iterate axis 0, not name a `scan_input_axes` of 1. + + The MLX EP rejects any Scan whose scan axis is not 0 + ("only scan_input_axes=0 is supported"). Since an unclaimed node is a + partition boundary, that sends all 120 recurrences back to CPU and + gives up the whole 7.0x the EP is worth on this model. + """ + _config, pkg = _build() + + scans = [node for node in pkg["model"].graph if node.op_type == "Scan"] + assert scans + for scan in scans: + assert "scan_input_axes" not in scan.attributes + assert "scan_output_axes" not in scan.attributes + + def test_sequence_reverse_uses_negative_step_slice(self): + """The backward branch reverses with a negative-step Slice. + + ReverseSequence would also reverse the sequence, but it needs an + extra Expand to build per-row lengths and implies padding semantics + this model does not have — nothing here is padded, so every row is + reversed in full. + """ + _config, pkg = _build() + + reverse_slices = [ + node + for node in pkg["model"].graph + if node.op_type == "Slice" + and len(node.inputs) == 5 + and node.inputs[4] is not None + and node.inputs[4].const_value is not None + and node.inputs[4].const_value.numpy().tolist() == [-1] + ] + # Two reversals (in and out) per bidirectional block, and each + # TFMambaBlock has a time-axis and a frequency-axis block. + assert len(reverse_slices) == _tiny_config().num_tfmamba * 2 * 2 + assert not any(node.op_type == "ReverseSequence" for node in pkg["model"].graph) diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index ef8326581..8c73cff49 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -99,6 +99,7 @@ "SSM2CausalLMTask", "SSMCausalLMTask", "Seq2SeqTask", + "SpeechEnhancementTask", "SpeechLanguageTask", "SpeechToTextTask", "TASK_REGISTRY", @@ -196,6 +197,7 @@ from mobius.tasks._rnnt import RNNTTask from mobius.tasks._sensenova_u1 import SenseNovaU1Task from mobius.tasks._seq2seq import Seq2SeqTask +from mobius.tasks._speech_enhancement import SpeechEnhancementTask from mobius.tasks._speech_language import SpeechLanguageTask from mobius.tasks._speech_to_text import SpeechToTextTask from mobius.tasks._ssm_causal_lm import SSM2CausalLMTask, SSMCausalLMTask @@ -301,6 +303,7 @@ "fun-asr-speech-language": FunASRSpeechLanguageTask, "glmasr-speech-language": GlmAsrSpeechLanguageTask, "fastconformer-rnnt": RNNTTask, + "speech-enhancement": SpeechEnhancementTask, "speech-language": SpeechLanguageTask, "speech-to-text": SpeechToTextTask, "ssm-text-generation": SSMCausalLMTask, diff --git a/src/mobius/tasks/_speech_enhancement.py b/src/mobius/tasks/_speech_enhancement.py new file mode 100644 index 000000000..fe7bee905 --- /dev/null +++ b/src/mobius/tasks/_speech_enhancement.py @@ -0,0 +1,72 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Speech-enhancement task. + +Builds a single ONNX graph for spectral speech-enhancement models (e.g. +RE-USE / SEMamba) that map a noisy STFT magnitude and phase to an enhanced +magnitude, phase, and complex spectrogram. The STFT and ISTFT stay outside +the graph, as they do for every other audio model in mobius. +""" + +from __future__ import annotations + +from typing import ClassVar + +import onnx_ir as ir + +from mobius._configs import BaseModelConfig +from mobius._model_package import ModelPackage +from mobius.tasks._base import ModelTask, _make_graph, _make_model + + +class SpeechEnhancementTask(ModelTask): + """Build an ONNX graph for spectral speech enhancement (encoder-only). + + Inputs: + ``noisy_mag`` — ``[batch, freq, time]`` noisy STFT magnitude. + ``noisy_pha`` — ``[batch, freq, time]`` noisy STFT phase, in radians. + + Outputs: + ``denoised_mag`` — ``[batch, freq, time]`` + ``denoised_pha`` — ``[batch, freq, time]`` + ``denoised_com`` — ``[batch, freq, time, 2]`` real/imaginary pair. + """ + + model_roles: ClassVar[dict[str, str]] = {"model": "encoder"} + + def build( + self, + module, + config: BaseModelConfig, + ) -> ModelPackage: + graph, builder = _make_graph() + + # Native-rate RE-USE scales its FFT geometry from the decoded sample + # rate, so frequency is dynamic by default. An explicit native/BWE + # export selects a static geometry for provider partitioning. + num_freq = getattr(config, "static_num_freq_bins", None) or "freq" + shape = ["batch", num_freq, "time"] + + noisy_mag = builder.input("noisy_mag", dtype=config.dtype, shape=shape) + noisy_pha = builder.input("noisy_pha", dtype=config.dtype, shape=shape) + + denoised_mag, denoised_pha, denoised_com = module( + builder.op, + noisy_mag=noisy_mag, + noisy_pha=noisy_pha, + ) + + # The model is spectrally shape preserving, but symbolic shape + # inference cannot see through the Scan-based SSM recurrence, so + # republish the input's named dimensions rather than leaving the + # outputs fully anonymous. + denoised_mag.shape = ir.Shape(shape) + denoised_pha.shape = ir.Shape(shape) + denoised_com.shape = ir.Shape([*shape, 2]) + + builder.add_output(denoised_mag, "denoised_mag") + builder.add_output(denoised_pha, "denoised_pha") + builder.add_output(denoised_com, "denoised_com") + + return ModelPackage({"model": _make_model(graph)}, config=config) diff --git a/tests/_test_configs.py b/tests/_test_configs.py index 7d8e3d07b..276c8ed84 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -63,7 +63,7 @@ YolosConfig, Zamba2Config, ) -from mobius.models import EsmConfig +from mobius.models import EsmConfig, ReUseConfig # --------------------------------------------------------------------------- # Tiny model dimensions shared by all configs @@ -3738,6 +3738,25 @@ def vl_overrides(model_type: str) -> dict: }, True, ), + # --- RE-USE / SEMamba (spectral speech enhancement) --- + # Bidirectional Mamba over time and frequency; consumes a noisy STFT + # magnitude/phase pair rather than audio features. + ( + "reuse", + { + "_config_cls": ReUseConfig, + "hid_feature": 8, + "num_tfmamba": 1, + "d_state": 4, + "d_conv": 4, + "expand": 2, + "n_fft": 32, + "hop_size": 4, + "win_size": 32, + "sampling_rate": 8000, + }, + True, + ), ] ALL_CONFIGS: list[tuple[str, dict, bool]] = ( CAUSAL_LM_CONFIGS diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index ee4b312f2..20853703e 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -6167,6 +6167,9 @@ def test_jamba_preprocesses_fused_experts_in_numeric_order(self): "falcon_mamba", "mamba", "mamba2", + # Speech enhancement: "reuse" is driven by SPEECH_CONFIGS; "semamba" is + # a bare alias of the same class, so it has no config of its own. + "semamba", # Hybrid SSM+Attention dedicated tests "bamba", "jamba", @@ -6879,6 +6882,7 @@ def test_outputs_have_shapes_and_dtypes(self, model_type: str, config_overrides: "codec": {"decoder", "encoder"}, "audio-feature-extraction": {"model"}, "feature-ctc-asr": {"model"}, + "speech-enhancement": {"model"}, } diff --git a/tests/cli_test.py b/tests/cli_test.py index 6206ee07f..4e8aee556 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -100,6 +100,35 @@ def test_max_workers_defaults_to_eight(self): assert save_package.call_args.args[2].max_workers == 8 + def test_reuse_revision_is_pinned_before_diffusers_probe(self): + from mobius.models.reuse import REUSE_REVISION + + with ( + tempfile.TemporaryDirectory() as tmpdir, + mock.patch( + "mobius.integrations.diffusers._builder._load_diffusers_pipeline_index", + return_value=None, + ) as pipeline_probe, + mock.patch("mobius.__main__.build", return_value=mock.MagicMock()) as build_model, + mock.patch("mobius.__main__._save_package"), + ): + main(["build", "--model", "nvidia/RE-USE", tmpdir, "--no-weights"]) + + assert pipeline_probe.call_args.kwargs["revision"] == REUSE_REVISION + assert build_model.call_args.kwargs["revision"] == REUSE_REVISION + + @pytest.mark.parametrize("option", ["--input-sample-rate", "--bwe-sample-rate"]) + def test_reuse_rate_options_are_rejected_for_diffusers(self, option): + with ( + tempfile.TemporaryDirectory() as tmpdir, + mock.patch( + "mobius.integrations.diffusers._builder._load_diffusers_pipeline_index", + return_value={"_class_name": "ExamplePipeline"}, + ), + pytest.raises(SystemExit, match="only supported for RE-USE"), + ): + main(["build", "--model", "example/diffusers", tmpdir, option, "16000"]) + @pytest.mark.parametrize( ("extra_args", "expected"), [([], True), (["--dequantize"], False)], diff --git a/tests/model_coverage_test.py b/tests/model_coverage_test.py index 47f2b303e..6925890ab 100644 --- a/tests/model_coverage_test.py +++ b/tests/model_coverage_test.py @@ -175,6 +175,11 @@ def _all_registered_with_test_id() -> dict[str, str]: "hy_v3": "L1 graph construction, native Transformers L2 config loading, and " "payload-free L3 trunk parity are covered; the immutable official checkpoint is " "597,578,239,288 bytes, so L4/L5 real-weight goldens exceed the 16 GiB policy.", + "reuse": "RE-USE / SEMamba speech enhancement — L1/L3 run from " + "SPEECH_CONFIGS and src/mobius/models/reuse_test.py. No L2: the published " + "nvidia/RE-USE config.json is a bespoke model_cfg/stft_cfg document with no " + "model_type field, which arch_validation_test requires, so the generic " + "download-and-build path cannot drive it.", # --- Internal / duplicate aliases --- "code_llama": "Alias for llama — covered by llama", "command_r": "Alias for cohere — covered by cohere", @@ -245,6 +250,7 @@ def _all_registered_with_test_id() -> dict[str, str]: "generic L2/L4/L5; pinned HF-to-GGUF config semantics and synthetic ORT parity " "are covered by _specialized_encoders_test.py.", "seed_oss": "Internal model — no public HF checkpoint", + "semamba": "Alias for reuse — covered by reuse", "shieldgemma2": "Alias for gemma2 — covered by gemma2", "yi": "Alias for llama — covered by llama", # --- VL models / text-decoder submodels (L1 graph-build only) ---