import os
import modal


app = modal.App("vllm-audio-language")

vllm_image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install(
        "vllm[audio]==0.7.3",
        "huggingface_hub[hf_transfer]==0.26.2",
        "flashinfer-python==0.2.0.post2",  # pinning, very unstable
        extra_index_url="https://flashinfer.ai/whl/cu124/torch2.5",
    )
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})  # faster model transfers
)

hf_cache_vol = modal.Volume.from_name("huggingface-cache", create_if_missing=True)
vllm_cache_vol = modal.Volume.from_name("vllm-cache", create_if_missing=True)

# PP need close v1
vllm_image = vllm_image.env(
    {
        "VLLM_USE_V1": os.getenv("VLLM_USE_V1", "1"),
    }
)


@app.function(
    image=vllm_image,
    gpu=os.getenv("IMAGE_GPU", "L4"),
    volumes={
        "/root/.cache/huggingface": hf_cache_vol,
        "/root/.cache/vllm": vllm_cache_vol,
    },
    retries=0,
    # how long should we stay up with no requests?
    scaledown_window=15 * 60,
)
def run_audio(model_type: str, num_prompts: int, num_audios: int):
    from transformers import AutoTokenizer

    from vllm import LLM, SamplingParams
    from vllm.assets.audio import AudioAsset
    from vllm.utils import FlexibleArgumentParser

    audio_assets = [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
    question_per_audio_count = {
        0: "What is 1+1?",
        1: "What is recited in the audio?",
        2: "What sport and what nursery rhyme are referenced?",
    }

    # NOTE: The default `max_num_seqs` and `max_model_len` may result in OOM on
    # lower-end GPUs.
    # Unless specified, these settings have been tested to work on a single L4.

    # Ultravox 0.5-1B
    def run_ultravox(question: str, audio_count: int):
        model_name = "fixie-ai/ultravox-v0_5-llama-3_2-1b"

        tokenizer = AutoTokenizer.from_pretrained(model_name)
        messages = [{"role": "user", "content": "<|audio|>\n" * audio_count + question}]
        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

        llm = LLM(
            model=model_name,
            max_model_len=4096,
            max_num_seqs=5,
            trust_remote_code=True,
            limit_mm_per_prompt={"audio": audio_count},
        )
        stop_token_ids = None
        return llm, prompt, stop_token_ids

    # Qwen2-Audio
    def run_qwen2_audio(question: str, audio_count: int):
        model_name = "Qwen/Qwen2-Audio-7B-Instruct"

        llm = LLM(
            model=model_name,
            max_model_len=4096,
            max_num_seqs=5,
            limit_mm_per_prompt={"audio": audio_count},
        )

        audio_in_prompt = "".join(
            [
                f"Audio {idx+1}: " f"<|audio_bos|><|AUDIO|><|audio_eos|>\n"
                for idx in range(audio_count)
            ]
        )

        prompt = (
            "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
            "<|im_start|>user\n"
            f"{audio_in_prompt}{question}<|im_end|>\n"
            "<|im_start|>assistant\n"
        )
        stop_token_ids = None
        return llm, prompt, stop_token_ids

    def run_minicpmo(question: str, audio_count: int):
        model_name = "openbmb/MiniCPM-o-2_6"
        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
        llm = LLM(
            model=model_name,
            trust_remote_code=True,
            max_model_len=4096,
            max_num_seqs=5,
            limit_mm_per_prompt={"audio": audio_count},
        )

        stop_tokens = ["<|im_end|>", "<|endoftext|>"]
        stop_token_ids = [tokenizer.convert_tokens_to_ids(i) for i in stop_tokens]

        audio_placeholder = "(<audio>./</audio>)" * audio_count
        audio_chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n<|spk_bos|><|spk|><|spk_eos|><|tts_bos|>' }}{% endif %}"  # noqa: E501
        messages = [{"role": "user", "content": f"{audio_placeholder}\n{question}"}]
        prompt = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True, chat_template=audio_chat_template
        )
        return llm, prompt, stop_token_ids

    model_example_map = {
        "ultravox": run_ultravox,
        "qwen2_audio": run_qwen2_audio,
        "minicpmo": run_minicpmo,
    }

    model = model_type
    if model not in model_example_map:
        raise ValueError(f"Model type {model} is not supported.")

    audio_count = num_audios
    llm, prompt, stop_token_ids = model_example_map[model](
        question_per_audio_count[audio_count], audio_count
    )

    # We set temperature to 0.2 so that outputs can be different
    # even when all prompts are identical when running batch inference.
    sampling_params = SamplingParams(temperature=0.2, max_tokens=64, stop_token_ids=stop_token_ids)

    mm_data = {}
    if audio_count > 0:
        mm_data = {"audio": [asset.audio_and_sample_rate for asset in audio_assets[:audio_count]]}

    assert num_prompts > 0
    inputs = {"prompt": prompt, "multi_modal_data": mm_data}
    if num_prompts > 1:
        # Batch inference
        inputs = [inputs] * num_prompts

    outputs = llm.generate(inputs, sampling_params=sampling_params)

    for o in outputs:
        generated_text = o.outputs[0].text
        print(generated_text)


@app.local_entrypoint()
def main(model_type: str = "qwen2_audio", num_prompts: int = 1, num_audios: int = 1):
    run_audio.remote(model_type, num_prompts, num_audios)
