Add ROMA frozen backbone extraction support

streaming_turn_model.py

@@ -31,6 +31,7 @@ from __future__ import annotations
 
 import argparse
 import copy
+import gc
 import json
 import logging
 import math
@@ -64,6 +65,7 @@ Mode = Literal[
     "open-stream-replay",
     "overlap-working-memory",
 ]
+BackboneLoader = Literal["qwen", "roma"]
 MODEL_ID = "Qwen/Qwen2.5-Omni-3B"
 _KNOWN_QWEN_WARNING_TEXT = (
     "System prompt modified, audio output may not work",
@@ -116,8 +118,17 @@ def load_qwen_omni(
     dtype: Literal["auto", "bfloat16", "float16", "float32"] = "auto",
     device_map: Any = "auto",
     use_flash_attention_2: bool = False,
+    backbone_loader: BackboneLoader = "qwen",
 ) -> tuple[Qwen2_5OmniThinkerForConditionalGeneration, Qwen2_5OmniProcessor]:
-    """Load only Qwen's multimodal Thinker and its processor."""
+    """Load a multimodal Thinker and its processor.
+
+    The normal Qwen checkpoints can be loaded directly as a Thinker.  ROMA is
+    released as a complete Qwen2.5-Omni model whose safetensors also contain
+    custom ``thinker.gate_*`` modules.  For a frozen-backbone experiment we
+    deliberately ignore those binary proactive-speaking heads, load the full
+    checkpoint so the ``thinker.*`` prefix is handled correctly, detach the
+    Thinker, and release the speech-generation modules.
+    """
 
     model_kwargs: dict[str, Any] = {
         "torch_dtype": _choose_dtype(dtype),
@@ -125,14 +136,128 @@ def load_qwen_omni(
     }
     if use_flash_attention_2:
         model_kwargs["attn_implementation"] = "flash_attention_2"
-    model = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(
-        model_id, **model_kwargs
-    )
+
+    if backbone_loader == "qwen":
+        model = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(
+            model_id, **model_kwargs
+        )
+        metadata: dict[str, Any] = {
+            "loader": "qwen",
+            "checkpoint_layout": "thinker",
+            "original_speak_head": "not_applicable",
+        }
+    elif backbone_loader == "roma":
+        try:
+            from transformers import Qwen2_5OmniModel
+        except ImportError as error:
+            raise RuntimeError(
+                "The ROMA loader requires transformers.Qwen2_5OmniModel. "
+                "Use the ROMA/Qwen2.5-Omni environment rather than an older "
+                "Transformers installation."
+            ) from error
+
+        loaded = Qwen2_5OmniModel.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", [])]
+        non_gate_unexpected = [
+            key for key in unexpected if not key.startswith("thinker.gate_")
+        ]
+        if missing or mismatched or errors or non_gate_unexpected:
+            raise RuntimeError(
+                "ROMA checkpoint did not load cleanly as a Qwen2.5-Omni "
+                "backbone: "
+                f"missing={missing[:8]}, mismatched={mismatched[:8]}, "
+                f"errors={errors[:4]}, unexpected_non_gate="
+                f"{non_gate_unexpected[:8]}"
+            )
+        if not hasattr(full_model, "thinker"):
+            raise RuntimeError("ROMA full checkpoint has no .thinker module")
+
+        model = full_model.thinker
+        gate_modules_present = bool(
+            hasattr(model, "gate_mixer")
+            and (
+                hasattr(model, "gate_head")
+                or hasattr(model, "gate_head_pro_fc1")
+            )
+        )
+        full_model.thinker = None
+        del full_model
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        metadata = {
+            "loader": "roma",
+            "checkpoint_layout": "full_qwen2_5_omni",
+            "ignored_checkpoint_keys": unexpected,
+            "ignored_checkpoint_key_policy": "thinker.gate_*_only",
+            "original_speak_head": (
+                "loaded_but_not_called; EasyCom head is trained anew"
+                if gate_modules_present
+                else (
+                    "checkpoint_weights_ignored_by_standard_transformers; "
+                    "EasyCom head is trained anew"
+                )
+            ),
+            "original_speak_head_modules_present": gate_modules_present,
+        }
+        print(
+            "ROMA backbone loaded cleanly: "
+            f"gate_modules_present={gate_modules_present}, "
+            f"ignored_gate_tensors={len(unexpected)}; "
+            "Talker/Token2Wav released.",
+            flush=True,
+        )
+    else:
+        raise ValueError(f"Unknown backbone_loader: {backbone_loader}")
+
     processor = Qwen2_5OmniProcessor.from_pretrained(model_id)
     model.eval()
+    text_config = getattr(model.config, "text_config", None)
+    if text_config is None and hasattr(model.config, "get_text_config"):
+        text_config = model.config.get_text_config()
+    metadata.update(
+        {
+            "model_id": str(model_id),
+            "hidden_size": int(getattr(text_config, "hidden_size")),
+            "num_hidden_layers": int(
+                getattr(text_config, "num_hidden_layers")
+            ),
+            "seconds_per_chunk": float(
+                getattr(model.config, "seconds_per_chunk")
+            ),
+        }
+    )
+    setattr(model, "_turn_backbone_metadata", metadata)
     return model, processor
 
 
+def backbone_metadata(
+    model: Qwen2_5OmniThinkerForConditionalGeneration,
+) -> dict[str, Any]:
+    """Return JSON/pickle-safe provenance attached by ``load_qwen_omni``."""
+
+    metadata = getattr(model, "_turn_backbone_metadata", None)
+    if not isinstance(metadata, dict):
+        return {"loader": "unknown"}
+    return copy.deepcopy(metadata)
+
+
 @dataclass(frozen=True)
 class StreamConfig:
     """Temporal and visual settings for one extraction run."""
@@ -2189,6 +2314,16 @@ def parse_args() -> argparse.Namespace:
         default="open-stream-replay",
     )
     parser.add_argument("--model-id", default=MODEL_ID)
+    parser.add_argument(
+        "--backbone-loader",
+        choices=("qwen", "roma"),
+        default="qwen",
+        help=(
+            "qwen loads a Thinker checkpoint directly; roma loads the full "
+            "ROMA Qwen2.5-Omni checkpoint, keeps its Thinker, and ignores the "
+            "original binary Speak Head"
+        ),
+    )
     parser.add_argument("--start", type=float, default=0.0)
     parser.add_argument("--end", type=float)
     parser.add_argument("--step-seconds", type=float, default=0.2)
@@ -2318,6 +2453,7 @@ def main() -> None:
         dtype=args.dtype,
         device_map={"": args.device},
         use_flash_attention_2=args.flash_attention_2,
+        backbone_loader=args.backbone_loader,
     )
     model.requires_grad_(False)
 
@@ -2350,6 +2486,8 @@ def main() -> None:
 
     payload: dict[str, Any] = {
         "model_id": args.model_id,
+        "backbone_loader": args.backbone_loader,
+        "backbone_metadata": backbone_metadata(model),
         "mode": args.mode,
         "stream_representation": (
             "single_causal_window"