File size: 2,016 Bytes
a2b7cfd 81ff557 f2527c6 a2b7cfd 81ff557 8641305 81ff557 a2b7cfd 81ff557 8641305 81ff557 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
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:
# Guarda el audio en un archivo temporal y retorna la ruta
sf.write("output.wav", audio.reshape(-1).detach().cpu().numpy(), samplerate=24000)
result["audio_path"] = "output.wav"
return [result] |