diff --git a/baselines/main.py b/baselines/main.py index 56d73ff..2302c73 100644 --- a/baselines/main.py +++ b/baselines/main.py @@ -31,6 +31,8 @@ from models.mllm.ola import load_model as load_ola from models.mllm.qwen_omni import generate as generate_qwen_omni from models.mllm.qwen_omni import load_model as load_qwen_omni +from models.mllm.qwen3_omni import generate as generate_qwen3_omni +from models.mllm.qwen3_omni import load_model as load_qwen3_omni # import speech models from models.speech.desta import generate as generate_desta @@ -141,6 +143,9 @@ def load_model(model_name): elif model_name == "qwen_omni": model = load_qwen_omni() generate_func = generate_qwen_omni + elif model_name == "qwen3_omni": + model = load_qwen3_omni() + generate_func = generate_qwen3_omni elif model_name == "ming_lite_omni": model = load_ming_lite_omni() generate_func = generate_ming_lite_omni @@ -228,6 +233,7 @@ def main(in_data_folder, out_folder, model, lang, track, modality, prompt): "ming_lite_omni", "ola", "qwen_omni", + "qwen3_omni", "gpt_oss", "gemini", ] diff --git a/baselines/models/mllm/qwen3_omni.py b/baselines/models/mllm/qwen3_omni.py new file mode 100644 index 0000000..a5e0ec1 --- /dev/null +++ b/baselines/models/mllm/qwen3_omni.py @@ -0,0 +1,118 @@ +# Copyright 2026 FBK, KIT + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License + +from utils import read_txt_file + + +def load_model(): + from transformers import Qwen3OmniMoeForConditionalGeneration, Qwen3OmniMoeProcessor + + model = Qwen3OmniMoeForConditionalGeneration.from_pretrained( + "Qwen/Qwen3-Omni-30B-A3B-Instruct", + torch_dtype="auto", + device_map="auto", + attn_implementation="flash_attention_2", + ) + processor = Qwen3OmniMoeProcessor.from_pretrained("Qwen/Qwen3-Omni-30B-A3B-Instruct") + + return model, processor + + +def generate(model_processor, prompt, example_path, modality): + from qwen_omni_utils import process_mm_info + + model, processor = model_processor + + if modality == "mllm": + USE_AUDIO_IN_VIDEO = True + user_conv_content = [ + {"type": "video", "video": example_path}, + {"type": "text", "text": prompt}, + ] + + elif modality == "video": + USE_AUDIO_IN_VIDEO = False + user_conv_content = [ + {"type": "video", "video": example_path}, + {"type": "text", "text": prompt}, + ] + + elif modality == "audio": + USE_AUDIO_IN_VIDEO = False + user_conv_content = [ + {"type": "audio", "audio": example_path}, + {"type": "text", "text": prompt}, + ] + + elif modality == "text": + USE_AUDIO_IN_VIDEO = False + example = read_txt_file(example_path) + user_conv_content = [ + {"type": "text", "text": f"{example}\n{prompt}\n"}, + ] + + system_conv = { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, " + "capable of perceiving auditory and visual inputs, as well as generating " + "text and speech. Only return the answer requested. Do not include any " + "explanation or introductions.", + } + ], + } + + user_conv = { + "role": "user", + "content": user_conv_content, + } + + conversation = [system_conv, user_conv] + + # Preparation for inference + text = processor.apply_chat_template( + conversation, add_generation_prompt=True, tokenize=False + ) + audios, images, videos = process_mm_info( + conversation, use_audio_in_video=USE_AUDIO_IN_VIDEO + ) + inputs = processor( + text=text, + audio=audios, + images=images, + videos=videos, + return_tensors="pt", + padding=True, + use_audio_in_video=USE_AUDIO_IN_VIDEO, + ) + inputs = inputs.to(model.device).to(model.dtype) + + # Qwen3-Omni always returns a (text_ids, audio) tuple; with return_audio=False the + # second element is None. thinker_max_new_tokens/thinker_do_sample are Qwen2.5-Omni + # specific parameters and are not supported by Qwen3-Omni. + text_ids, _ = model.generate( + **inputs, + use_audio_in_video=USE_AUDIO_IN_VIDEO, + return_audio=False, + max_new_tokens=4096, + ) + text = processor.batch_decode( + text_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False + ) + + # postprocess + response = text[-1].split("\nassistant")[-1].strip() + return response