fix(loader): cast audio_tokenizer inputs to param dtype via forward hook

After casting the audio_tokenizer's weights to fp16/bf16 (saves ~3 GB
VRAM), incoming fp32 wav tensors hit Linear/Conv layers with fp16
weights and crash with "expected scalar type Float but found Half" deep
inside encode_audios_from_wav.

Register a forward-pre-hook on every leaf-ish module of the
audio_tokenizer that casts floating-point input tensors to the module's
parameter dtype on the fly. Covers both the encode (voice reference)
and decode paths without needing to enumerate the model's submodule
names.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

lib/model_state.py

@@ -241,6 +241,42 @@ def _fix_moss_model_config_token_ids(processor) -> None:
         cls._moss_tts_lang_attr_fallback_patched = True
 
 
+def _cast_inputs_to_param_dtype_hook(module, args):
+    """Forward pre-hook: cast floating tensor inputs to module's param dtype.
+
+    AutoProcessor loads the audio_tokenizer in fp32. After we cast its weights
+    to fp16/bf16 to save VRAM, raw fp32 wav inputs flow through Linear/Conv
+    layers whose weights are now half-precision -> RuntimeError "expected
+    scalar type Float but found Half". The hook fires before each leaf
+    module's forward and aligns input dtype to the module's own parameters.
+    """
+    if not args:
+        return args
+    params = list(module.parameters(recurse=False))
+    if not params:
+        return args
+    target_dtype = params[0].dtype
+    new_args = tuple(
+        a.to(target_dtype)
+        if isinstance(a, torch.Tensor)
+        and a.is_floating_point()
+        and a.dtype != target_dtype
+        else a
+        for a in args
+    )
+    return new_args
+
+
+def _patch_audio_tokenizer_dtype(audio_tokenizer) -> None:
+    """Register the dtype-cast pre-hook on every leaf-ish module of the audio
+    tokenizer. Covers both encode (voice ref) and decode paths without
+    needing to know the model's submodule names.
+    """
+    for module in audio_tokenizer.modules():
+        if list(module.parameters(recurse=False)):
+            module.register_forward_pre_hook(_cast_inputs_to_param_dtype_hook)
+
+
 def resolve_attn_impl(requested: str, device: str) -> str:
     if requested == "auto":
         return "sdpa" if device == "cuda" else "eager"
@@ -314,6 +350,7 @@ def get_or_load(
             processor.audio_tokenizer = processor.audio_tokenizer.to(
                 resolved_at_device, dtype=resolved_dtype
             )
+            _patch_audio_tokenizer_dtype(processor.audio_tokenizer)
         else:
             processor.audio_tokenizer = processor.audio_tokenizer.to(resolved_at_device)
 

tests/test_model_state.py

@@ -133,6 +133,30 @@ def test_audio_tokenizer_cast_to_dtype_on_cuda():
         fake_at.to.assert_called_with("cuda", dtype=torch.float16)
 
 
+def test_dtype_cast_hook_aligns_fp32_input_with_fp16_layer():
+    """End-to-end: a Linear with fp16 weights normally errors on fp32 input,
+    but with the hook installed it casts the input and produces fp16 output.
+    """
+    import torch
+    layer = torch.nn.Linear(4, 4).to(torch.float16)
+    inp_fp32 = torch.randn(2, 4, dtype=torch.float32)
+    with pytest.raises(RuntimeError, match=r"expected scalar type"):
+        layer(inp_fp32)
+    model_state._patch_audio_tokenizer_dtype(layer)
+    out = layer(inp_fp32)
+    assert out.dtype == torch.float16
+
+
+def test_dtype_cast_hook_passes_through_matching_dtype():
+    """Hook is a no-op when input dtype already matches param dtype."""
+    import torch
+    layer = torch.nn.Linear(4, 4).to(torch.float16)
+    model_state._patch_audio_tokenizer_dtype(layer)
+    inp_fp16 = torch.randn(2, 4, dtype=torch.float16)
+    out = layer(inp_fp16)
+    assert out.dtype == torch.float16
+
+
 def test_audio_tokenizer_no_dtype_cast_on_cpu():
     """When the audio_tokenizer is on CPU, do not force fp16 — fp16 ops are
     slow on most CPUs and the VRAM concern doesn't apply off-GPU.