Qwen3-Omni: Online serving¶

import argparse
import base64
import io
import os
import random
from pathlib import Path
from typing import Any

try:
    import gradio as gr
except ImportError:
    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
import numpy as np
import soundfile as sf
import torch
from openai import OpenAI
from PIL import Image

SEED = 42

SUPPORTED_MODELS: dict[str, dict[str, Any]] = {
    "Qwen/Qwen3-Omni-30B-A3B-Instruct": {
        "sampling_params": {
            "thinker": {
                "temperature": 0.4,
                "top_p": 0.9,
                "top_k": 1,
                "max_tokens": 16384,
                "detokenize": True,
                "repetition_penalty": 1.05,
                "stop_token_ids": [151645],
                "seed": SEED,
            },
            "talker": {
                "temperature": 0.9,
                "top_k": 50,
                "max_tokens": 4096,
                "seed": SEED,
                "detokenize": False,
                "repetition_penalty": 1.05,
                "stop_token_ids": [2150],
            },
            "code2wav": {
                "temperature": 0.0,
                "top_p": 1.0,
                "top_k": -1,
                "max_tokens": 4096 * 16,
                "seed": SEED,
                "detokenize": True,
                "repetition_penalty": 1.1,
            },
        },
    },
}
# Ensure deterministic behavior across runs.
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
os.environ["PYTHONHASHSEED"] = str(SEED)
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"


def parse_args():
    parser = argparse.ArgumentParser(description="Gradio demo for Qwen3-Omni online inference.")
    parser.add_argument(
        "--model",
        default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
        help="Model name/path (should match the server model).",
    )
    parser.add_argument(
        "--api-base",
        default="http://localhost:8091/v1",
        help="Base URL for the vLLM API server.",
    )
    parser.add_argument(
        "--ip",
        default="127.0.0.1",
        help="Host/IP for gradio `launch`.",
    )
    parser.add_argument("--port", type=int, default=7861, help="Port for gradio `launch`.")
    parser.add_argument("--share", action="store_true", help="Share the Gradio demo publicly.")
    return parser.parse_args()


def build_sampling_params_dict(seed: int, model_key: str) -> list[dict]:
    """Build sampling params as dict for HTTP API mode."""
    model_conf = SUPPORTED_MODELS.get(model_key)
    if model_conf is None:
        raise ValueError(f"Unsupported model '{model_key}'")

    sampling_templates: dict[str, dict[str, Any]] = model_conf["sampling_params"]
    sampling_params: list[dict] = []
    for stage_name, template in sampling_templates.items():
        params = dict(template)
        params["seed"] = seed
        sampling_params.append(params)
    return sampling_params


def image_to_base64_data_url(image: Image.Image) -> str:
    """Convert PIL Image to base64 data URL."""
    buffered = io.BytesIO()
    # Convert to RGB if needed
    if image.mode != "RGB":
        image = image.convert("RGB")
    image.save(buffered, format="JPEG")
    img_bytes = buffered.getvalue()
    img_b64 = base64.b64encode(img_bytes).decode("utf-8")
    return f"data:image/jpeg;base64,{img_b64}"


def audio_to_base64_data_url(audio_data: tuple[np.ndarray, int]) -> str:
    """Convert audio (numpy array, sample_rate) to base64 data URL."""
    audio_np, sample_rate = audio_data
    # Convert to int16 format for WAV
    if audio_np.dtype != np.int16:
        # Normalize to [-1, 1] range if needed
        if audio_np.dtype == np.float32 or audio_np.dtype == np.float64:
            audio_np = np.clip(audio_np, -1.0, 1.0)
            audio_np = (audio_np * 32767).astype(np.int16)
        else:
            audio_np = audio_np.astype(np.int16)

    # Write to WAV bytes
    buffered = io.BytesIO()
    sf.write(buffered, audio_np, sample_rate, format="WAV")
    wav_bytes = buffered.getvalue()
    wav_b64 = base64.b64encode(wav_bytes).decode("utf-8")
    return f"data:audio/wav;base64,{wav_b64}"


def video_to_base64_data_url(video_file: str) -> str:
    """Convert video file to base64 data URL."""
    video_path = Path(video_file)
    if not video_path.exists():
        raise FileNotFoundError(f"Video file not found: {video_file}")

    # Detect MIME type from extension
    video_path_lower = str(video_path).lower()
    if video_path_lower.endswith(".mp4"):
        mime_type = "video/mp4"
    elif video_path_lower.endswith(".webm"):
        mime_type = "video/webm"
    elif video_path_lower.endswith(".mov"):
        mime_type = "video/quicktime"
    elif video_path_lower.endswith(".avi"):
        mime_type = "video/x-msvideo"
    elif video_path_lower.endswith(".mkv"):
        mime_type = "video/x-matroska"
    else:
        mime_type = "video/mp4"

    with open(video_path, "rb") as f:
        video_bytes = f.read()
    video_b64 = base64.b64encode(video_bytes).decode("utf-8")
    return f"data:{mime_type};base64,{video_b64}"


def process_audio_file(
    audio_file: Any | None,
) -> tuple[np.ndarray, int] | None:
    """Normalize Gradio audio input to (np.ndarray, sample_rate)."""
    if audio_file is None:
        return None

    sample_rate: int | None = None
    audio_np: np.ndarray | None = None

    def _load_from_path(path_str: str) -> tuple[np.ndarray, int] | None:
        if not path_str:
            return None
        path = Path(path_str)
        if not path.exists():
            return None
        data, sr = sf.read(path)
        if data.ndim > 1:
            data = data[:, 0]
        return data.astype(np.float32), int(sr)

    if isinstance(audio_file, tuple):
        if len(audio_file) == 2:
            first, second = audio_file
            # Case 1: (sample_rate, np.ndarray)
            if isinstance(first, (int, float)) and isinstance(second, np.ndarray):
                sample_rate = int(first)
                audio_np = second
            # Case 2: (filepath, (sample_rate, np.ndarray or list))
            elif isinstance(first, str):
                if isinstance(second, tuple) and len(second) == 2:
                    sr_candidate, data_candidate = second
                    if isinstance(sr_candidate, (int, float)) and isinstance(data_candidate, np.ndarray):
                        sample_rate = int(sr_candidate)
                        audio_np = data_candidate
                if audio_np is None:
                    loaded = _load_from_path(first)
                    if loaded is not None:
                        audio_np, sample_rate = loaded
            # Case 3: (None, (sample_rate, np.ndarray))
            elif first is None and isinstance(second, tuple) and len(second) == 2:
                sr_candidate, data_candidate = second
                if isinstance(sr_candidate, (int, float)) and isinstance(data_candidate, np.ndarray):
                    sample_rate = int(sr_candidate)
                    audio_np = data_candidate
        elif len(audio_file) == 1 and isinstance(audio_file[0], str):
            loaded = _load_from_path(audio_file[0])
            if loaded is not None:
                audio_np, sample_rate = loaded
    elif isinstance(audio_file, str):
        loaded = _load_from_path(audio_file)
        if loaded is not None:
            audio_np, sample_rate = loaded

    if audio_np is None or sample_rate is None:
        return None

    if audio_np.ndim > 1:
        audio_np = audio_np[:, 0]

    return audio_np.astype(np.float32), sample_rate


def process_image_file(image_file: Image.Image | None) -> Image.Image | None:
    """Process image file from Gradio input.

    Returns:
        PIL Image in RGB mode or None if no image provided.
    """
    if image_file is None:
        return None
    # Convert to RGB if needed
    if image_file.mode != "RGB":
        image_file = image_file.convert("RGB")
    return image_file


def run_inference_api(
    client: OpenAI,
    model: str,
    sampling_params_dict: list[dict],
    user_prompt: str,
    audio_file: tuple[str, tuple[int, np.ndarray]] | None = None,
    image_file: Image.Image | None = None,
    video_file: str | None = None,
    use_audio_in_video: bool = False,
    output_modalities: str | None = None,
    stream: bool = False,
):
    """Run inference using OpenAI API client with multimodal support."""
    if not user_prompt.strip() and not audio_file and not image_file and not video_file:
        yield "Please provide at least a text prompt or multimodal input.", None

    try:
        # Build message content list
        content_list = []

        # Process audio
        audio_data = process_audio_file(audio_file)
        if audio_data is not None:
            audio_url = audio_to_base64_data_url(audio_data)
            content_list.append(
                {
                    "type": "audio_url",
                    "audio_url": {"url": audio_url},
                }
            )

        # Process image
        if image_file is not None:
            image_data = process_image_file(image_file)
            if image_data is not None:
                image_url = image_to_base64_data_url(image_data)
                content_list.append(
                    {
                        "type": "image_url",
                        "image_url": {"url": image_url},
                    }
                )

        # Process video
        mm_processor_kwargs = {}
        if video_file is not None:
            video_url = video_to_base64_data_url(video_file)
            video_content = {
                "type": "video_url",
                "video_url": {"url": video_url},
            }
            if use_audio_in_video:
                video_content["video_url"]["num_frames"] = 32  # Default max frames
                mm_processor_kwargs["use_audio_in_video"] = True
            content_list.append(video_content)

        # Add text prompt
        if user_prompt.strip():
            content_list.append(
                {
                    "type": "text",
                    "text": user_prompt,
                }
            )

        # Build messages
        messages = [
            {
                "role": "system",
                "content": [
                    {
                        "type": "text",
                        "text": (
                            "You are Qwen, a virtual human developed by the Qwen Team, "
                            "Alibaba Group, capable of perceiving auditory and visual inputs, "
                            "as well as generating text and speech."
                        ),
                    }
                ],
            },
            {
                "role": "user",
                "content": content_list,
            },
        ]

        # Build extra_body
        extra_body = {
            "sampling_params_list": sampling_params_dict,
        }
        if mm_processor_kwargs:
            extra_body["mm_processor_kwargs"] = mm_processor_kwargs

        # Parse output modalities
        if output_modalities and output_modalities.strip():
            output_modalities_list = [m.strip() for m in output_modalities.split(",")]
        else:
            output_modalities_list = None

        # Call API
        chat_completion = client.chat.completions.create(
            messages=messages,
            model=model,
            modalities=output_modalities_list,
            extra_body=extra_body,
            stream=stream,
        )

        if not stream:
            # Non-streaming mode: extract outputs and yield once
            text_outputs: list[str] = []
            audio_output = None

            for choice in chat_completion.choices:
                if choice.message.content:
                    text_outputs.append(choice.message.content)
                if choice.message.audio:
                    # Decode base64 audio
                    audio_data = base64.b64decode(choice.message.audio.data)
                    # Load audio from bytes
                    audio_np, sample_rate = sf.read(io.BytesIO(audio_data))
                    # Convert to mono if needed
                    if audio_np.ndim > 1:
                        audio_np = audio_np[:, 0]
                    audio_output = (int(sample_rate), audio_np.astype(np.float32))

            text_response = "\n\n".join(text_outputs) if text_outputs else "No text output."
            yield text_response, audio_output
        else:
            # Streaming mode: yield incremental updates
            text_content = ""
            audio_output = None

            for chunk in chat_completion:
                for choice in chunk.choices:
                    if hasattr(choice, "delta"):
                        content = getattr(choice.delta, "content", None)
                    else:
                        content = None

                    # Handle audio modality
                    if getattr(chunk, "modality", None) == "audio" and content:
                        try:
                            # Decode base64 audio
                            audio_data = base64.b64decode(content)
                            # Load audio from bytes
                            audio_np, sample_rate = sf.read(io.BytesIO(audio_data))
                            # Convert to mono if needed
                            if audio_np.ndim > 1:
                                audio_np = audio_np[:, 0]
                            audio_output = (int(sample_rate), audio_np.astype(np.float32))
                            # Yield current text and audio
                            yield text_content if text_content else "", audio_output
                        except Exception:  # pylint: disable=broad-except
                            # If audio processing fails, just yield text
                            yield text_content if text_content else "", None

                    # Handle text modality
                    elif getattr(chunk, "modality", None) == "text":
                        if content:
                            text_content += content
                            # Yield updated text content (keep existing audio if any)
                            yield text_content, audio_output

            # Final yield with accumulated text and last audio (if any)
            yield text_content if text_content else "No text output.", audio_output

    except Exception as exc:  # pylint: disable=broad-except
        error_msg = f"Inference failed: {exc}"
        yield error_msg, None


def build_interface(
    client: OpenAI,
    model: str,
    sampling_params_dict: list[dict],
):
    """Build Gradio interface for API server mode."""

    def run_inference(
        user_prompt: str,
        audio_file: tuple[str, tuple[int, np.ndarray]] | None,
        image_file: Image.Image | None,
        video_file: str | None,
        use_audio_in_video: bool,
        output_modalities: str | None = None,
        stream: bool = False,
    ):
        # Always yield from the API function to maintain consistent generator behavior
        yield from run_inference_api(
            client,
            model,
            sampling_params_dict,
            user_prompt,
            audio_file,
            image_file,
            video_file,
            use_audio_in_video,
            output_modalities,
            stream,
        )

    css = """
    .media-input-container {
        display: flex;
        gap: 10px;
    }
    .media-input-container > div {
        flex: 1;
    }
    .media-input-container .image-input,
    .media-input-container .audio-input {
        height: 300px;
    }
    .media-input-container .video-column {
        height: 300px;
        display: flex;
        flex-direction: column;
    }
    .media-input-container .video-input {
        flex: 1;
        min-height: 0;
    }
    #generate-btn button {
        width: 100%;
    }
    """

    with gr.Blocks(css=css) as demo:
        gr.Markdown("# vLLM-Omni Online Serving Demo")
        gr.Markdown(f"**Model:** {model} \n\n")

        with gr.Column():
            with gr.Row():
                input_box = gr.Textbox(
                    label="Text Prompt",
                    placeholder="For example: Describe what happens in the media inputs.",
                    lines=4,
                    scale=1,
                )
            with gr.Row(elem_classes="media-input-container"):
                image_input = gr.Image(
                    label="Image Input (optional)",
                    type="pil",
                    sources=["upload"],
                    scale=1,
                    elem_classes="image-input",
                )
                with gr.Column(scale=1, elem_classes="video-column"):
                    video_input = gr.Video(
                        label="Video Input (optional)",
                        sources=["upload"],
                        elem_classes="video-input",
                    )
                    use_audio_in_video_checkbox = gr.Checkbox(
                        label="Use audio from video",
                        value=False,
                        info="Extract the video's audio track when provided.",
                    )
                audio_input = gr.Audio(
                    label="Audio Input (optional)",
                    type="numpy",
                    sources=["upload", "microphone"],
                    scale=1,
                    elem_classes="audio-input",
                )

        with gr.Row():
            output_modalities = gr.Textbox(
                label="Output Modalities",
                value=None,
                placeholder="For example: text, image, video. Use comma to separate multiple modalities.",
                lines=1,
                scale=2,
            )
            stream_checkbox = gr.Checkbox(
                label="Stream output",
                value=False,
                info="Enable streaming to see output as it's generated.",
                scale=1,
            )

        with gr.Row():
            generate_btn = gr.Button(
                "Generate",
                variant="primary",
                size="lg",
                elem_id="generate-btn",
            )

        with gr.Row():
            text_output = gr.Textbox(label="Text Output", lines=10, scale=2)
            audio_output = gr.Audio(label="Audio Output", interactive=False, scale=1)

        generate_btn.click(
            fn=run_inference,
            inputs=[
                input_box,
                audio_input,
                image_input,
                video_input,
                use_audio_in_video_checkbox,
                output_modalities,
                stream_checkbox,
            ],
            outputs=[text_output, audio_output],
        )
        demo.queue()
    return demo


def main():
    args = parse_args()

    model_name = "/".join(args.model.split("/")[-2:])
    assert model_name in SUPPORTED_MODELS, (
        f"Unsupported model '{model_name}'. Supported models: {SUPPORTED_MODELS.keys()}"
    )

    # Initialize OpenAI client
    print(f"Connecting to API server at: {args.api_base}")
    client = OpenAI(
        api_key="EMPTY",
        base_url=args.api_base,
    )
    print("✓ Connected to API server")

    # Build sampling params
    sampling_params_dict = build_sampling_params_dict(SEED, model_name)

    demo = build_interface(
        client,
        args.model,
        sampling_params_dict,
    )
    try:
        demo.launch(
            server_name=args.ip,
            server_port=args.port,
            share=args.share,
        )
    except KeyboardInterrupt:
        print("\nShutting down...")


if __name__ == "__main__":
    main()