diff --git a/extract_ego4d_batch.py b/extract_ego4d_batch.py index 9309129..6de4eb5 100644 --- a/extract_ego4d_batch.py +++ b/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( diff --git a/run_extract_easycom_direct.sh b/run_extract_easycom_direct.sh index 697436f..930b862 100644 --- a/run_extract_easycom_direct.sh +++ b/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 diff --git a/smoke_test.py b/smoke_test.py index b564a71..a4f42c8 100644 --- a/smoke_test.py +++ b/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( - 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) + if args.backbone_loader == "roma": + from streaming_turn_model import backbone_metadata, load_qwen_omni + + model, processor = load_qwen_omni( + args.model, + dtype="bfloat16", + 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( - **inputs, - max_new_tokens=args.max_new_tokens, - do_sample=False, - return_audio=False, - ) + generation_kwargs = { + "max_new_tokens": args.max_new_tokens, + "do_sample": 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] :] diff --git a/streaming_turn_model.py b/streaming_turn_model.py index 8c7963a..6bf6d66 100644 --- a/streaming_turn_model.py +++ b/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" diff --git a/streaming_turn_model_directly.py b/streaming_turn_model_directly.py index 91fcdac..f2416c7 100644 --- a/streaming_turn_model_directly.py +++ b/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"