Upload infer_qwen2.5omni.py

infer_qwen2.5omni.py

@@ -0,0 +1,68 @@
+from qwen_omni_utils import process_mm_info
+import torch
+from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor
+import librosa
+import os 
+from io import BytesIO
+from urllib.request import urlopen
+import argparse
+# @title inference function
+def inference(audio_path,model,processor,prompt, sys_prompt):
+    messages = [
+        {"role": "system", "content": [{"type": "text", "text": sys_prompt}]},
+        {"role": "user", "content": [
+                {"type": "audio", "audio": audio_path},
+                {"type": "text", "text": prompt},
+            ]
+        },
+    ]
+    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    audios, images, videos = process_mm_info(messages, use_audio_in_video=True)
+    inputs = processor(text=text, audio=audios, images=images, videos=videos, return_tensors="pt", padding=True, use_audio_in_video=True)
+    inputs = inputs.to(model.device).to(model.dtype)
+
+    output = model.generate(**inputs, use_audio_in_video=True, return_audio=False, thinker_max_new_tokens=256, thinker_do_sample=False)
+
+    text = processor.batch_decode(output, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+    return text
+
+def transcribe(wavs_path, out_path, gpu_id, model):
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
+    model_path = model
+    model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
+                model_path,
+                torch_dtype=torch.bfloat16,
+                device_map="auto",
+        )
+    prompt = "请将这段中文语音转换为纯文本，去掉标点符号。"
+    processor = Qwen2_5OmniProcessor.from_pretrained(model_path)
+    with open(wavs_path, "r") as f_in, open(out_path, "w") as f_out:
+        for line in f_in:
+            utt, path = line.strip().split(" ", maxsplit=1)
+            try:
+                response=inference(path,model,processor, prompt=prompt, sys_prompt="You are a speech recognition model.")
+            except Exception as e:
+                print(f"Inference failed: {str(e)}")  
+                response="None"
+            text = response[0].strip()
+            lines = text.strip().splitlines()
+            text = lines[-1]
+            print(f"[{utt}] >>> {text}")
+            f_out.write(f"{utt} {text}\n")
+
+
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--wavs_path", type=str)
+    parser.add_argument("--out_path", type=str)
+    parser.add_argument("--gpu", type=int, default=0)
+    parser.add_argument("--model", type=str)
+    args = parser.parse_args()
+    transcribe(
+        wavs_path=args.wavs_path,
+        out_path=args.out_path,
+        gpu_id=args.gpu,
+        model=args.model
+    )