#!/usr/bin/env python3
"""Minimal single-GPU smoke test for a local Qwen2.5-Omni checkpoint.

Examples:
  python smoke_test.py --model /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B
  python smoke_test.py --model /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \
      --video /data1/example.mp4

Qwen2.5-Omni generation currently is more reliable on a single GPU for this
use case.  In particular, do not replace the explicit device map below with
``device_map=\"auto\"``: it can place cache tensors on different GPUs.
"""

import argparse
import time

import torch
from qwen_omni_utils import process_mm_info
from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", required=True, help="Local model directory or HF model ID")
    parser.add_argument(
        "--backbone-loader",
        choices=("qwen", "roma"),
        default="qwen",
        help="Use the normal full Qwen smoke path or the ROMA frozen-Thinker loader",
    )
    parser.add_argument("--video", help="Optional local .mp4 file for an audio-video test")
    parser.add_argument("--device", default="cuda:0", help="One CUDA device, e.g. cuda:0")
    parser.add_argument("--max-new-tokens", type=int, default=64)
    parser.add_argument(
        "--video-fps",
        type=float,
        default=1.0,
        help="Frames per second sampled by the Omni video processor (default: 1.0)",
    )
    parser.add_argument(
        "--video-max-pixels",
        type=int,
        default=256 * 28 * 28,
        help="Maximum pixels per sampled video frame (default: 200704)",
    )
    parser.add_argument(
        "--flash-attn2",
        action="store_true",
        help="Use FlashAttention-2; requires the flash_attn package.",
    )
    return parser.parse_args()


def move_tensors_to_device(batch: dict, device: str) -> dict:
    """Move processor tensors without accidentally converting integer IDs."""
    return {
        name: value.to(device) if torch.is_tensor(value) else value
        for name, value in batch.items()
    }


def main() -> None:
    args = parse_args()
    if not torch.cuda.is_available():
        raise RuntimeError("A CUDA GPU is required for this smoke test.")

    print(f"GPU: {torch.cuda.get_device_name(args.device)}")
    print(f"CUDA: {torch.version.cuda}")
    print(f"PyTorch: {torch.__version__}")
    print(f"Loading: {args.model}")

    started = time.perf_counter()
    # A 3B model fits comfortably on one L20 (48 GB).  Keep all weights and
    # generation cache on the same device; device_map='auto' may split them.
    model_load_kwargs = dict(
        torch_dtype=torch.bfloat16,
        device_map={"": args.device},
    )
    if args.flash_attn2:
        model_load_kwargs["attn_implementation"] = "flash_attention_2"

    if args.backbone_loader == "roma":
        from streaming_turn_model import backbone_metadata, load_qwen_omni

        model, processor = load_qwen_omni(
            args.model,
            dtype="bfloat16",
            device_map={"": args.device},
            use_flash_attention_2=args.flash_attn2,
            backbone_loader="roma",
        )
        print(f"Backbone metadata: {backbone_metadata(model)}")
    else:
        model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
            args.model,
            **model_load_kwargs,
        )
        model.disable_talker()  # We need Thinker text states only; saves VRAM.
        model.eval()
        processor = Qwen2_5OmniProcessor.from_pretrained(args.model)
    print(f"Model loaded in {time.perf_counter() - started:.1f}s")

    if args.video:
        content = [
            {
                "type": "video",
                "video": args.video,
                # These are consumed by qwen_omni_utils while sampling the
                # video.  They prevent quadratic vision attention from
                # allocating tens of GB for a short, high-resolution clip.
                "fps": args.video_fps,
                "max_pixels": args.video_max_pixels,
            },
            {
                "type": "text",
                "text": "Briefly describe the people, speech, and interaction in this video.",
            },
        ]
        print(
            f"Video: {args.video} "
            f"(sample_fps={args.video_fps}, max_pixels={args.video_max_pixels})"
        )
    else:
        content = [
            {"type": "text", "text": "Reply with exactly: Qwen Omni smoke test passed."},
        ]

    messages = [{"role": "user", "content": content}]
    prompt = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    audios, images, videos = process_mm_info(
        messages, use_audio_in_video=bool(args.video)
    )
    inputs = processor(
        text=prompt,
        audio=audios,
        images=images,
        videos=videos,
        return_tensors="pt",
        padding=True,
        use_audio_in_video=bool(args.video),
    )
    inputs = move_tensors_to_device(inputs, args.device)

    started = time.perf_counter()
    with torch.inference_mode():
        generation_kwargs = {
            "max_new_tokens": args.max_new_tokens,
            "do_sample": False,
        }
        if args.backbone_loader == "qwen":
            generation_kwargs["return_audio"] = False
        output_ids = model.generate(**inputs, **generation_kwargs)

    # Remove the prompt tokens before decoding the generated completion.
    generated_ids = output_ids[:, inputs["input_ids"].shape[1] :]
    answer = processor.batch_decode(
        generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )[0]
    print(f"Generation completed in {time.perf_counter() - started:.1f}s")
    print(f"\nModel response:\n{answer.strip()}")


if __name__ == "__main__":
    main()