| from typing import Dict, List, Any |
| import soundfile as sf |
| from transformers import Qwen3OmniMoeForConditionalGeneration, Qwen3OmniMoeProcessor |
| from qwen_omni_utils import process_mm_info |
|
|
| class EndpointHandler(): |
| def __init__(self, path="./"): |
| self.model = Qwen3OmniMoeForConditionalGeneration.from_pretrained( |
| path, |
| dtype="auto", |
| device_map="auto", |
| ) |
| self.processor = Qwen3OmniMoeProcessor.from_pretrained(path) |
|
|
| def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
| messages = data.get("messages", []) |
| use_audio_in_video = data.get("use_audio_in_video", True) |
| speaker = data.get("speaker", "Ethan") |
|
|
| text = self.processor.apply_chat_template( |
| messages, |
| tokenize=False, |
| add_generation_prompt=True, |
| ) |
| audios, images, videos = process_mm_info(messages, use_audio_in_video=use_audio_in_video) |
| inputs = self.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(self.model.device).to(self.model.dtype) |
|
|
| text_ids, audio = self.model.generate( |
| **inputs, |
| speaker=speaker, |
| thinker_return_dict_in_generate=True, |
| use_audio_in_video=use_audio_in_video |
| ) |
| text_output = self.processor.batch_decode( |
| text_ids.sequences[:, inputs["input_ids"].shape[1]:], |
| skip_special_tokens=True, |
| clean_up_tokenization_spaces=False |
| ) |
| result = {"generated_text": text_output} |
| if audio is not None: |
| |
| sf.write("output.wav", audio.reshape(-1).detach().cpu().numpy(), samplerate=24000) |
| result["audio_path"] = "output.wav" |
| return [result] |
