[Feature] Add Gradio Demo for Qwen2.5Omni (#60)

Signed-off-by: SamitHuang <285365963@qq.com>
Signed-off-by: samithuang <285365963@qq.com>

examples/online_serving/README.md

@@ -38,3 +38,40 @@ If you encounter error about backend of librosa, try to install ffmpeg with comm
 sudo apt update
 sudo apt install ffmpeg
 ```
+
+## Run Local Web UI Demo
+
+This Web UI demo allows users to interact with the model through a web browser.
+
+### Running Gradio Demo
+
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by
+
+```bash
+python gradio_demo.py  --model Qwen/Qwen2.5-Omni-7B --port 7861
+```
+
+Then open `http://localhost:7861/` on your local browser to interact with the web UI.
+
+
+### Options
+
+The gradio demo also supports running with an existing API server and can be customized with the following arguments.
+
+
+```bash
+python gradio_demo.py \
+    --model Qwen/Qwen2.5-Omni-7B \
+    --use-api-server \
+    --api-base http://localhost:8091/v1 \
+    --ip 127.0.0.1 \
+    --port 7861
+```
+
+- `--model`: Model name
+- `--use-api-server`: If set, connect to an existing vLLM HTTP API server instead of running AsyncOmniLLM locally.
+- `--api-base`: Base URL for vllm serve (only used when `use-api-server` is set, default: http://localhost:8091/v1)
+- `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
+- `--port`: Port for Gradio server (default: 7861)
+- `--stage-configs-path`: Path to custom stage configs YAML file (optional)
+- `--share`: Share the Gradio demo publicly (creates a public link)

examples/online_serving/gradio_demo.py

@@ -0,0 +1,535 @@
+import argparse
+import os
+import random
+import signal
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any, Optional
+
+import gradio as gr
+import numpy as np
+import soundfile as sf
+import torch
+from PIL import Image
+from vllm.assets.video import video_get_metadata, video_to_ndarrays
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
+
+# Import utils from offline inference example
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../offline_inference/qwen2_5_omni"))
+try:
+    from utils import make_omni_prompt
+except ImportError:
+    # Fallback if utils doesn't exist - we'll build prompts directly
+    make_omni_prompt = None
+
+SEED = 42
+ASYNC_INIT_TIMEOUT = 600
+
+SUPPORTED_MODELS: dict[str, dict[str, Any]] = {
+    "Qwen/Qwen2.5-Omni-7B": {
+        "sampling_params": {
+            "thinker": {
+                "temperature": 0.0,
+                "top_p": 1.0,
+                "top_k": -1,
+                "max_tokens": 2048,
+                "detokenize": True,
+                "repetition_penalty": 1.1,
+            },
+            "talker": {
+                "temperature": 0.0,
+                "top_p": 1.0,
+                "top_k": -1,
+                "max_tokens": 2048,
+                "detokenize": True,
+                "repetition_penalty": 1.1,
+                "stop_token_ids": [8294],
+            },
+            "code2wav": {
+                "temperature": 0.0,
+                "top_p": 1.0,
+                "top_k": -1,
+                "max_tokens": 2048,
+                "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 Qwen2.5-Omni online inference.")
+    parser.add_argument(
+        "--model",
+        default="Qwen/Qwen2.5-Omni-7B",
+        help="Path to local model checkpoint.",
+    )
+    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.")
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=None,
+        help="Path to custom stage configs YAML file (optional).",
+    )
+    return parser.parse_args()
+
+
+def build_sampling_params(seed: int, model_key: str) -> list[SamplingParams]:
+    """Build SamplingParams objects by reusing the dict definitions."""
+    return [SamplingParams(**params_dict) for params_dict in build_sampling_params_dict(seed, model_key)]
+
+
+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 create_prompt_args(base_args: argparse.Namespace) -> SimpleNamespace:
+    # The prompt builder expects a minimal namespace with these attributes.
+    return SimpleNamespace(
+        model=base_args.model,
+        prompt_type="text",
+        tokenize=True,
+        use_torchvision=True,
+        legacy_omni_video=False,
+    )
+
+
+def process_audio_file(
+    audio_file: Optional[Any],
+) -> Optional[tuple[np.ndarray, int]]:
+    """Normalize Gradio audio input to (np.ndarray, sample_rate)."""
+    if audio_file is None:
+        return None
+
+    sample_rate: Optional[int] = None
+    audio_np: Optional[np.ndarray] = None
+
+    def _load_from_path(path_str: str) -> Optional[tuple[np.ndarray, int]]:
+        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: Optional[Image.Image]) -> Optional[Image.Image]:
+    """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 process_video_file(
+    video_file: Optional[str],
+    enable_audio_in_video: bool = False,
+    max_frames: int = 32,
+) -> Optional[tuple[np.ndarray, dict[str, Any], Optional[tuple[np.ndarray, int]]]]:
+    """Process video file and optionally extract audio track."""
+    if video_file is None:
+        return None
+
+    video_path = Path(video_file)
+    if not video_path.exists():
+        print(f"Video file not found: {video_path}")
+        return None
+
+    try:
+        frames = video_to_ndarrays(str(video_path), num_frames=max_frames)
+        metadata = video_get_metadata(str(video_path), num_frames=frames.shape[0])
+    except Exception as exc:  # pylint: disable=broad-except
+        print(f"Failed to decode video {video_path}: {exc}")
+        return None
+
+    audio_tuple: Optional[tuple[np.ndarray, int]] = None
+    if enable_audio_in_video:
+        try:
+            import librosa  # type: ignore import
+
+            audio_signal, sampling_rate = librosa.load(str(video_path), sr=16000)
+            audio_tuple = (audio_signal.astype(np.float32), sampling_rate)
+        except Exception as exc:  # pylint: disable=broad-except
+            print(f"Failed to extract audio from video {video_path}: {exc}")
+
+    return frames, metadata, audio_tuple
+
+
+async def run_inference_async_omni_llm(
+    omni_llm: AsyncOmniLLM,
+    sampling_params: list[SamplingParams],
+    prompt_args_template: SimpleNamespace,
+    user_prompt: str,
+    audio_file: Optional[tuple[str, tuple[int, np.ndarray]]] = None,
+    image_file: Optional[Image.Image] = None,
+    video_file: Optional[str] = None,
+    use_audio_in_video: bool = False,
+):
+    """Run inference using AsyncOmniLLM directly with multimodal support."""
+    if not user_prompt.strip() and not audio_file and not image_file and not video_file:
+        return "Please provide at least a text prompt or multimodal input.", None
+
+    try:
+        # Build prompt with multimodal data
+        prompt_args = SimpleNamespace(**prompt_args_template.__dict__)
+
+        # Process multimodal inputs
+        multi_modal_data = {}
+        mm_processor_kwargs = {}
+
+        # Process audio
+        audio_data = process_audio_file(audio_file)
+        if audio_data is not None:
+            multi_modal_data["audio"] = audio_data
+
+        # Process image
+        image_data = process_image_file(image_file)
+        if image_data is not None:
+            multi_modal_data["image"] = image_data
+
+        # Process video
+        video_payload = process_video_file(video_file, enable_audio_in_video=use_audio_in_video)
+        if video_payload is not None:
+            video_frames, video_metadata, extracted_audio = video_payload
+            video_entry: Any
+            if video_metadata:
+                video_entry = (video_frames, video_metadata)
+            else:
+                video_entry = video_frames
+            multi_modal_data["video"] = video_entry
+            if use_audio_in_video and extracted_audio is not None and "audio" not in multi_modal_data:
+                multi_modal_data["audio"] = extracted_audio
+                mm_processor_kwargs["use_audio_in_video"] = True
+
+        # Build the prompt input
+        if make_omni_prompt is not None:
+            omni_prompt = make_omni_prompt(prompt_args, user_prompt)
+            # Add multimodal data if present
+            if multi_modal_data:
+                if isinstance(omni_prompt, dict):
+                    omni_prompt["multi_modal_data"] = multi_modal_data
+                    if mm_processor_kwargs:
+                        omni_prompt["mm_processor_kwargs"] = mm_processor_kwargs
+                else:
+                    # If make_omni_prompt returns a string, we need to wrap it
+                    omni_prompt = {
+                        "prompt": omni_prompt,
+                        "multi_modal_data": multi_modal_data,
+                    }
+                    if mm_processor_kwargs:
+                        omni_prompt["mm_processor_kwargs"] = mm_processor_kwargs
+        else:
+            # Fallback: build prompt directly
+            default_system = (
+                "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."
+            )
+            prompt = f"<|im_start|>system\n{default_system}<|im_end|>\n<|im_start|>user\n"
+            if audio_data:
+                prompt += "<|audio_bos|><|AUDIO|><|audio_eos|>"
+            if image_data:
+                prompt += "<|vision_bos|><|IMAGE|><|vision_eos|>"
+            if video_payload is not None:
+                prompt += "<|vision_bos|><|VIDEO|><|vision_eos|>"
+            if user_prompt.strip():
+                prompt += f"{user_prompt}"
+            prompt += "<|im_end|>\n<|im_start|>assistant\n"
+
+            omni_prompt = {
+                "prompt": prompt,
+                "multi_modal_data": multi_modal_data,
+            }
+            if mm_processor_kwargs:
+                omni_prompt["mm_processor_kwargs"] = mm_processor_kwargs
+
+        request_id = "0"
+        text_outputs: list[str] = []
+        audio_output = None
+
+        async for stage_outputs in omni_llm.generate(
+            prompt=omni_prompt,
+            request_id=request_id,
+            sampling_params_list=sampling_params,
+        ):
+            # stage_outputs.request_output is a RequestOutput object, not a list
+            request_output = stage_outputs.request_output
+            if stage_outputs.final_output_type == "text":
+                if request_output.outputs:
+                    for output in request_output.outputs:
+                        if output.text:
+                            text_outputs.append(output.text)
+            elif stage_outputs.final_output_type == "audio":
+                # multimodal_output is on the RequestOutput object
+                # See vllm_omni/entrypoints/openai/serving_chat.py:680 for reference
+                if hasattr(request_output, "multimodal_output") and request_output.multimodal_output:
+                    if "audio" in request_output.multimodal_output:
+                        audio_tensor = request_output.multimodal_output["audio"]
+                        # Ensure audio is 1D (flatten if needed)
+                        if hasattr(audio_tensor, "ndim") and audio_tensor.ndim > 1:
+                            audio_tensor = audio_tensor.flatten()
+                        audio_np = audio_tensor.detach().cpu().numpy()
+                        audio_output = (
+                            24000,  # sampling rate in Hz
+                            audio_np,
+                        )
+
+        text_response = "\n\n".join(text_outputs) if text_outputs else "No text output."
+        return text_response, audio_output
+    except Exception as exc:  # pylint: disable=broad-except
+        return f"Inference failed: {exc}", None
+
+
+def build_interface(
+    omni_llm: AsyncOmniLLM,
+    sampling_params: list[SamplingParams],
+    prompt_args_template: SimpleNamespace,
+    model: str,
+):
+    """Build Gradio interface for AsyncOmniLLM mode."""
+
+    async def run_inference(
+        user_prompt: str,
+        audio_file: Optional[tuple[str, tuple[int, np.ndarray]]],
+        image_file: Optional[Image.Image],
+        video_file: Optional[str],
+        use_audio_in_video: bool,
+    ):
+        return await run_inference_async_omni_llm(
+            omni_llm,
+            sampling_params,
+            prompt_args_template,
+            user_prompt,
+            audio_file,
+            image_file,
+            video_file,
+            use_audio_in_video,
+        )
+
+    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():
+            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],
+            outputs=[text_output, audio_output],
+        )
+        demo.queue()
+    return demo
+
+
+def main():
+    args = parse_args()
+    omni_llm = None
+
+    model_name = "/".join(args.model.split("/")[-2:])
+    assert model_name in SUPPORTED_MODELS, (
+        f"Unsupported model '{model_name}'. Supported models: {SUPPORTED_MODELS.keys()}"
+    )
+
+    # Register signal handlers for graceful shutdown
+    def signal_handler(sig, frame):
+        print("\nReceived interrupt signal, shutting down...")
+        if omni_llm is not None:
+            try:
+                omni_llm.shutdown()
+            except Exception as e:
+                print(f"Error during shutdown: {e}")
+        sys.exit(0)
+
+    signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
+
+    print(f"Initializing AsyncOmniLLM with model: {args.model}")
+    if args.stage_configs_path:
+        print(f"Using custom stage configs: {args.stage_configs_path}")
+
+    sampling_params = build_sampling_params(SEED, model_name)
+    omni_llm = AsyncOmniLLM(
+        model=args.model,
+        stage_configs_path=args.stage_configs_path,
+        init_timeout=ASYNC_INIT_TIMEOUT,
+    )
+    print("✓ AsyncOmniLLM initialized successfully")
+    prompt_args_template = create_prompt_args(args)
+
+    demo = build_interface(
+        omni_llm,
+        sampling_params,
+        prompt_args_template,
+        args.model,
+    )
+    try:
+        demo.launch(
+            server_name=args.ip,
+            server_port=args.port,
+            share=args.share,
+        )
+    except KeyboardInterrupt:
+        print("\nShutting down...")
+    finally:
+        # Cleanup
+        if omni_llm is not None:
+            try:
+                omni_llm.shutdown()
+            except Exception as e:
+                print(f"Error during cleanup: {e}")
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -46,6 +46,8 @@ dev = [
     "isort>=5.12.0",
     "mypy>=1.0.0",
     "pre-commit>=3.0.0",
+    "gradio>=5.49.1",
+    "soundfile>=0.13.1"
 ]
 
 docs = [