Support upstream Omni class for ROMA loading

streaming_turn_model.py

@@ -149,14 +149,31 @@ def load_qwen_omni(
     elif backbone_loader == "roma":
         try:
             from transformers import Qwen2_5OmniModel
+        except ImportError:
+            # The ROMA transformers fork exposes the full wrapper as
+            # Qwen2_5OmniModel.  Upstream Transformers releases instead use
+            # Qwen2_5OmniForConditionalGeneration for the same full
+            # Thinker+Talker checkpoint layout.  The latter is already
+            # available in the Qwen/MS-SWIFT image used by this project.
+            try:
+                from transformers import Qwen2_5OmniForConditionalGeneration
             except ImportError as error:
                 raise RuntimeError(
-                "The ROMA loader requires transformers.Qwen2_5OmniModel. "
+                    "The ROMA loader requires either "
-                "Use the ROMA/Qwen2.5-Omni environment rather than an older "
+                    "transformers.Qwen2_5OmniModel (ROMA fork) or "
-                "Transformers installation."
+                    "Qwen2_5OmniForConditionalGeneration (upstream "
+                    "Transformers)."
                 ) from error
 
-        loaded = Qwen2_5OmniModel.from_pretrained(
+            full_model_class = Qwen2_5OmniForConditionalGeneration
+            full_model_class_name = (
+                "Qwen2_5OmniForConditionalGeneration"
+            )
+        else:
+            full_model_class = Qwen2_5OmniModel
+            full_model_class_name = "Qwen2_5OmniModel"
+
+        loaded = full_model_class.from_pretrained(
             model_id,
             output_loading_info=True,
             **model_kwargs,
@@ -203,6 +220,7 @@ def load_qwen_omni(
 
         metadata = {
             "loader": "roma",
+            "full_model_class": full_model_class_name,
             "checkpoint_layout": "full_qwen2_5_omni",
             "ignored_checkpoint_keys": unexpected,
             "ignored_checkpoint_key_policy": "thinker.gate_*_only",
@@ -218,6 +236,7 @@ def load_qwen_omni(
         }
         print(
             "ROMA backbone loaded cleanly: "
+            f"full_model_class={full_model_class_name}, "
             f"gate_modules_present={gate_modules_present}, "
             f"ignored_gate_tensors={len(unexpected)}; "
             "Talker/Token2Wav released.",