From d34725778f7f3dcb9ca527726bc570a646365351 Mon Sep 17 00:00:00 2001
From: Xin Wang <your-email@example.com>
Date: Thu, 27 Aug 2026 13:50:02 +0800
Subject: [PATCH] Avoid upstream Omni loading-info speaker crash

---
 streaming_turn_model.py | 92 ++++++++++++++++++++++++++++++++++-------
 1 file changed, 76 insertions(+), 16 deletions(-)

diff --git a/streaming_turn_model.py b/streaming_turn_model.py
index 15a2258..68ce124 100644
--- a/streaming_turn_model.py
+++ b/streaming_turn_model.py
@@ -112,6 +112,29 @@ def _choose_dtype(
     return getattr(torch, dtype)
 
 
+def _local_sharded_checkpoint_keys(model_id: str) -> set[str] | None:
+    """Read weight names without loading shard tensors when an index exists."""
+
+    root = Path(model_id)
+    if not root.is_dir():
+        return None
+    for name in (
+        "model.safetensors.index.json",
+        "pytorch_model.bin.index.json",
+    ):
+        index_path = root / name
+        if not index_path.is_file():
+            continue
+        payload = json.loads(index_path.read_text(encoding="utf-8"))
+        weight_map = payload.get("weight_map")
+        if not isinstance(weight_map, dict):
+            raise RuntimeError(
+                f"Checkpoint index has no weight_map object: {index_path}"
+            )
+        return {str(key) for key in weight_map}
+    return None
+
+
 def load_qwen_omni(
     model_id: str = MODEL_ID,
     *,
@@ -169,27 +192,63 @@ def load_qwen_omni(
             full_model_class_name = (
                 "Qwen2_5OmniForConditionalGeneration"
             )
+            supports_loading_info = False
         else:
             full_model_class = Qwen2_5OmniModel
             full_model_class_name = "Qwen2_5OmniModel"
+            supports_loading_info = True
 
-        loaded = full_model_class.from_pretrained(
-            model_id,
-            output_loading_info=True,
-            **model_kwargs,
-        )
-        if not isinstance(loaded, tuple) or len(loaded) != 2:
-            raise RuntimeError("ROMA loading did not return checkpoint audit data")
-        full_model, loading_info = loaded
+        if supports_loading_info:
+            loaded = full_model_class.from_pretrained(
+                model_id,
+                output_loading_info=True,
+                **model_kwargs,
+            )
+            if not isinstance(loaded, tuple) or len(loaded) != 2:
+                raise RuntimeError(
+                    "ROMA loading did not return checkpoint audit data"
+                )
+            full_model, loading_info = loaded
+            missing = [
+                str(key) for key in loading_info.get("missing_keys", [])
+            ]
+            unexpected = [
+                str(key) for key in loading_info.get("unexpected_keys", [])
+            ]
+            mismatched = [
+                str(item)
+                for item in loading_info.get("mismatched_keys", [])
+            ]
+            errors = [
+                str(item) for item in loading_info.get("error_msgs", [])
+            ]
+            checkpoint_audit = "transformers_output_loading_info"
+        else:
+            # Upstream Qwen2.5-Omni overrides from_pretrained() and calls
+            # load_speakers() on its return value.  output_loading_info=True
+            # changes that return value into (model, info), causing the
+            # upstream override to crash before returning.  Load normally and
+            # recover the unexpected-key audit from the local shard index.
+            full_model = full_model_class.from_pretrained(
+                model_id,
+                **model_kwargs,
+            )
+            missing = []
+            mismatched = []
+            errors = []
+            checkpoint_keys = _local_sharded_checkpoint_keys(model_id)
+            if checkpoint_keys is None:
+                unexpected = []
+                checkpoint_audit = (
+                    "upstream_loader; local_shard_index_unavailable"
+                )
+            else:
+                model_keys = set(full_model.state_dict())
+                unexpected = sorted(checkpoint_keys - model_keys)
+                checkpoint_audit = (
+                    "upstream_loader_plus_local_shard_key_audit"
+                )
 
-        missing = [str(key) for key in loading_info.get("missing_keys", [])]
-        unexpected = [
-            str(key) for key in loading_info.get("unexpected_keys", [])
-        ]
-        mismatched = [
-            str(item) for item in loading_info.get("mismatched_keys", [])
-        ]
-        errors = [str(item) for item in loading_info.get("error_msgs", [])]
         non_gate_unexpected = [
             key for key in unexpected if not key.startswith("thinker.gate_")
         ]
@@ -222,6 +281,7 @@ def load_qwen_omni(
             "loader": "roma",
             "full_model_class": full_model_class_name,
             "checkpoint_layout": "full_qwen2_5_omni",
+            "checkpoint_audit": checkpoint_audit,
             "ignored_checkpoint_keys": unexpected,
             "ignored_checkpoint_key_policy": "thinker.gate_*_only",
             "original_speak_head": (