diff --git a/streaming_turn_model.py b/streaming_turn_model.py index 6bf6d66..15a2258 100644 --- a/streaming_turn_model.py +++ b/streaming_turn_model.py @@ -149,14 +149,31 @@ def load_qwen_omni( 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 + 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 either " + "transformers.Qwen2_5OmniModel (ROMA fork) or " + "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.",