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
    )