Add ROMA frozen backbone extraction support

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