Add ROMA frozen backbone extraction support

extract_ego4d_batch.py

@@ -87,6 +87,11 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
         "--model-id",
         default="/data/wangx/Models/Qwen2.5-Omni-3B",
     )
+    parser.add_argument(
+        "--backbone-loader",
+        choices=("qwen", "roma"),
+        default="qwen",
+    )
     parser.add_argument(
         "--start-index",
         type=int,
@@ -283,6 +288,8 @@ def clip_argv(
     values = [
         "--model-id",
         str(args.model_id),
+        "--backbone-loader",
+        args.backbone_loader,
         "--video",
         str(args.video_root / f"{record.uid}.mp4"),
         "--target",
@@ -373,6 +380,7 @@ def main(argv: Sequence[str] | None = None) -> 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)
     print(

run_extract_easycom_direct.sh

@@ -8,6 +8,7 @@ set -euo pipefail
 
 PROJECT_ROOT="${PROJECT_ROOT:-/mnt/58462F4C462F29E8/Code/icassp2027-experiment}"
 MODEL="${MODEL:-/mnt/58462F4C462F29E8/Models/Qwen2.5-Omni-3B}"
+BACKBONE_LOADER="${BACKBONE_LOADER:-qwen}"
 VIDEO_ROOT="${VIDEO_ROOT:-/mnt/58462F4C462F29E8/Datasets/easycom/Main/Video_Compressed}"
 TARGET_ROOT="${TARGET_ROOT:-/mnt/58462F4C462F29E8/Code/EgoSpeak/dataset/EasyCom/target_perframe}"
 FEATURE_ROOT="${FEATURE_ROOT:-${PROJECT_ROOT}/features/easycom_direct}"
@@ -23,6 +24,9 @@ DTYPE="${DTYPE:-bfloat16}"
 DECODER="${DECODER:-decord}"
 DECODER_GPU_ID="${DECODER_GPU_ID:-0}"
 FRAME_CACHE_SIZE="${FRAME_CACHE_SIZE:-64}"
+# Delay worker startup by GPU index to avoid eight processes reading a large
+# full ROMA checkpoint at the same instant. Zero preserves the old behavior.
+MODEL_LOAD_STAGGER_SECONDS="${MODEL_LOAD_STAGGER_SECONDS:-0}"
 # Space-separated host GPU ids, e.g. "0" or "0 1 2 3". One session per GPU at a time.
 GPUS="${GPUS:-0}"
 # Optional: only run selected session ids, e.g. SESSIONS="1 2 3"
@@ -70,6 +74,7 @@ run_one() {
   {
     echo "$(date): GPU ${gpu} starting ${session}"
     echo "model=${MODEL}"
+    echo "backbone_loader=${BACKBONE_LOADER}"
     echo "session_dir=${session_dir}"
     echo "target=${target}"
     echo "output=${output}"
@@ -81,6 +86,7 @@ run_one() {
     if ! CUDA_VISIBLE_DEVICES="${gpu}" PYTHONUNBUFFERED=1 \
       python streaming_turn_model_directly.py \
         --model-id "${MODEL}" \
+        --backbone-loader "${BACKBONE_LOADER}" \
         --easycom-session-dir "${session_dir}" \
         --target "${target}" \
         --use-audio-in-video \
@@ -124,15 +130,30 @@ done
 
 pids=()
 status=0
+worker_index=0
 for gpu in "${GPU_ARR[@]}"; do
   sids="${WORKERS[${gpu}]}"
   [[ -z "${sids// }" ]] && continue
+  delay_seconds="$(python - "${MODEL_LOAD_STAGGER_SECONDS}" "${worker_index}" <<'PY'
+import sys
+print(float(sys.argv[1]) * int(sys.argv[2]))
+PY
+)"
   (
+    if python - "${delay_seconds}" <<'PY'
+import sys
+raise SystemExit(0 if float(sys.argv[1]) > 0 else 1)
+PY
+    then
+      echo "GPU ${gpu} worker waiting ${delay_seconds}s before model load"
+      sleep "${delay_seconds}"
+    fi
     for sid in ${sids}; do
       run_one "${gpu}" "${sid}" || exit 1
     done
   ) &
   pids+=("$!")
+  worker_index=$((worker_index + 1))
 done
 
 for pid in "${pids[@]}"; do

smoke_test.py

@@ -22,6 +22,12 @@ from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcess
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
     parser.add_argument("--model", required=True, help="Local model directory or HF model ID")
+    parser.add_argument(
+        "--backbone-loader",
+        choices=("qwen", "roma"),
+        default="qwen",
+        help="Use the normal full Qwen smoke path or the ROMA frozen-Thinker loader",
+    )
     parser.add_argument("--video", help="Optional local .mp4 file for an audio-video test")
     parser.add_argument("--device", default="cuda:0", help="One CUDA device, e.g. cuda:0")
     parser.add_argument("--max-new-tokens", type=int, default=64)
@@ -73,13 +79,25 @@ def main() -> None:
     if args.flash_attn2:
         model_load_kwargs["attn_implementation"] = "flash_attention_2"
 
-    model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
+    if args.backbone_loader == "roma":
-        args.model,
+        from streaming_turn_model import backbone_metadata, load_qwen_omni
-        **model_load_kwargs,
+
-    )
+        model, processor = load_qwen_omni(
-    model.disable_talker()  # We need Thinker text states only; saves VRAM.
+            args.model,
-    model.eval()
+            dtype="bfloat16",
-    processor = Qwen2_5OmniProcessor.from_pretrained(args.model)
+            device_map={"": args.device},
+            use_flash_attention_2=args.flash_attn2,
+            backbone_loader="roma",
+        )
+        print(f"Backbone metadata: {backbone_metadata(model)}")
+    else:
+        model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
+            args.model,
+            **model_load_kwargs,
+        )
+        model.disable_talker()  # We need Thinker text states only; saves VRAM.
+        model.eval()
+        processor = Qwen2_5OmniProcessor.from_pretrained(args.model)
     print(f"Model loaded in {time.perf_counter() - started:.1f}s")
 
     if args.video:
@@ -127,12 +145,13 @@ def main() -> None:
 
     started = time.perf_counter()
     with torch.inference_mode():
-        output_ids = model.generate(
+        generation_kwargs = {
-            **inputs,
+            "max_new_tokens": args.max_new_tokens,
-            max_new_tokens=args.max_new_tokens,
+            "do_sample": False,
-            do_sample=False,
+        }
-            return_audio=False,
+        if args.backbone_loader == "qwen":
-        )
+            generation_kwargs["return_audio"] = False
+        output_ids = model.generate(**inputs, **generation_kwargs)
 
     # Remove the prompt tokens before decoding the generated completion.
     generated_ids = output_ids[:, inputs["input_ids"].shape[1] :]

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"

streaming_turn_model_directly.py

@@ -1030,6 +1030,15 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
         default="open-stream-replay",
     )
     parser.add_argument("--model-id", default=core.MODEL_ID)
+    parser.add_argument(
+        "--backbone-loader",
+        choices=("qwen", "roma"),
+        default="qwen",
+        help=(
+            "qwen loads a Thinker checkpoint directly; roma loads the full "
+            "ROMA checkpoint and keeps only its frozen Thinker"
+        ),
+    )
     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)
@@ -1360,8 +1369,15 @@ def main(
             dtype=args.dtype,
             device_map={"": args.device},
             use_flash_attention_2=args.flash_attention_2,
+            backbone_loader=args.backbone_loader,
         )
     else:
+        loaded_backbone = core.backbone_metadata(model).get("loader")
+        if loaded_backbone not in {"unknown", args.backbone_loader}:
+            raise RuntimeError(
+                "Preloaded backbone does not match CLI provenance: "
+                f"model={loaded_backbone}, requested={args.backbone_loader}"
+            )
         model.eval()
     model.requires_grad_(False)
     media_path_for_api = reader.clips[0].path
@@ -1447,6 +1463,8 @@ def main(
     )
     payload: dict[str, Any] = {
         "model_id": args.model_id,
+        "backbone_loader": args.backbone_loader,
+        "backbone_metadata": core.backbone_metadata(model),
         "mode": args.mode,
         "stream_representation": (
             "single_causal_prefix"