Add ROMA frozen backbone extraction support
This commit is contained in:
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user