Add ROMA frozen backbone extraction support

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

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"