wx44wx/icassp2027-experiment
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add ROMA frozen backbone extraction support

Browse Source
This commit is contained in:
Xin Wang
2026-08-27 13:35:59 +08:00
parent 1925ff71a7
commit 655d27a4ca
5 changed files with 221 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
extract_ego4d_batch.py
View File

@@ -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(

21
run_extract_easycom_direct.sh
View File

@@ -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

45
smoke_test.py
View File

@@ -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] :]

146
streaming_turn_model.py
View File

@@ -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"

18
streaming_turn_model_directly.py
View File

@@ -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"