Port changes from feat/multi-request-stream (vs main)

Signed-off-by: 齐保元 <qibaoyuan@xiaomi.com>

examples/offline_inference/qwen_2_5_omni/extract_prompts.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+import argparse
+from typing import Optional
+import os
+import torch
+
+
+def extract_prompt(line: str) -> Optional[str]:
+    # 提取第一个 '|' 与第二个 '|' 之间的内容
+    i = line.find('|')
+    if i == -1:
+        return None
+    j = line.find('|', i + 1)
+    if j == -1:
+        return None
+    return line[i + 1:j].strip()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input", "-i", required=True, help="输入 .lst 文件路径")
+    parser.add_argument("--output", "-o", required=True, help="输出文件路径")
+    parser.add_argument("--topk", "-k", type=int, default=100, help="提取前 K 个 prompt（默认 100）")
+    parser.add_argument("--pt-output", "-p", type=str, default=None, help="可选：输出 .pt 文件路径（默认与 --output 同名 .pt）")
+    args = parser.parse_args()
+
+    prompts = []
+    with open(args.input, "r", encoding="utf-8", errors="ignore") as f:
+        for line in f:
+            if len(prompts) >= args.topk:
+                break
+            p = extract_prompt(line.rstrip("\n"))
+            if p:
+                prompts.append(p)
+
+    with open(args.output, "w", encoding="utf-8") as f:
+        for p in prompts:
+            f.write(p + "\n")
+
+    # 另存为 PyTorch 序列化的 .pt，内容为字符串列表
+    pt_path = args.pt_output if args.pt_output else os.path.splitext(args.output)[0] + ".pt"
+    torch.save(prompts, pt_path)
+
+
+if __name__ == "__main__":
+    main()
+
+

examples/offline_inference/qwen_2_5_omni/hf_qwen25_omni_batch_no_template.py

@@ -0,0 +1,279 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+Batch text generation with Qwen 2.5 Omni using Transformers without chat templates.
+
+- Processes 10 fixed prompts in batches of 3
+- Does NOT use tokenizer.apply_chat_template (raw prompts only)
+- Prints only generated continuations for each prompt
+"""
+
+import os
+import argparse
+from typing import Iterable, List
+import time
+import torch
+
+import random
+import numpy as np
+import torch
+from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor
+from tqdm import tqdm
+import json
+import soundfile as sf
+
+PROMPTS: List[str] = [
+    "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words.",
+    "Generate a friendly greeting message suitable for a voice assistant. Answer in 15 words.",
+    "Summarize the benefits of on-device inference for edge deployments. Answer in 15 words.",
+    "Offer troubleshooting steps for degraded audio quality in a TTS pipeline. Answer in 15 words.",
+    "List best practices for optimizing latency in real-time voice synthesis. Answer in 15 words.",
+    "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words.",
+    "Generate a friendly greeting message suitable for a voice assistant. Answer in 15 words.",
+    "Summarize the benefits of on-device inference for edge deployments. Answer in 15 words.",
+    "Offer troubleshooting steps for degraded audio quality in a TTS pipeline. Answer in 15 words.",
+    "List best practices for optimizing latency in real-time voice synthesis. Answer in 15 words.",
+]
+def make_text_prompt(prompt):
+    return [{
+        "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": [
+            {"type": "text", "text": prompt},
+        ],
+    }]
+
+def chunked(items: List[str], size: int) -> Iterable[List[str]]:
+    for i in range(0, len(items), size):
+        yield items[i : i + size]
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Qwen 2.5 Omni batch generation without chat templates",
+    )
+    parser.add_argument(
+        "--model-id",
+        type=str,
+        default=os.getenv("QWEN_OMNI_MODEL_ID", "Qwen/Qwen2.5-Omni"),
+        help="HF model repo id (e.g., Qwen/Qwen2.5-Omni)",
+    )
+    parser.add_argument(
+        "--max-new-tokens",
+        type=int,
+        default=2048,
+        help="Maximum number of new tokens to generate",
+    )
+    parser.add_argument(
+        "--temperature", type=float, default=0.0, help="Sampling temperature"
+    )
+    parser.add_argument("--top-p", type=float, default=1.0, help="Top-p sampling")
+    parser.add_argument(
+        "--repetition-penalty",
+        type=float,
+        default=1.1,
+        help="Repetition penalty (matches pipelined_end2end defaults)",
+    )
+    parser.add_argument(
+        "--stop-token-ids",
+        type=str,
+        default="8294",
+        help="Comma-separated stop token ids (added to eos_token_id)",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Random seed for reproducibility (sampling)",
+    )
+    parser.add_argument(
+        "--pt-prompts",
+        type=str,
+        default=None,
+        help="Path to .pt file containing List[str] prompts (overrides built-in PROMPTS)",
+    )
+    parser.add_argument(
+        "--tokenize",
+        action="store_true",
+        help="If set, utils.make_text_prompt will return token ids; keep unset for raw text",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="outputs",
+        help="Output directory to save aligned comparison files (same folder as pipelined_end2end)",
+    )
+    args = parser.parse_args()
+
+    if torch.cuda.is_available():
+        torch_dtype = torch.bfloat16
+    else:
+        torch_dtype = torch.float32
+
+    # Match pipelined_end2end.py deterministic settings
+    if args.seed is not None:
+        os.environ["PYTHONHASHSEED"] = str(args.seed)
+        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+        try:
+            random.seed(args.seed)
+            np.random.seed(args.seed)
+            torch.manual_seed(args.seed)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed(args.seed)
+                torch.cuda.manual_seed_all(args.seed)
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
+        except Exception:
+            pass
+
+    print(f"[Info] Loading model: {args.model_id}")
+    processor = Qwen2_5OmniProcessor.from_pretrained(args.model_id)
+    model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
+        args.model_id,
+        trust_remote_code=True,
+        torch_dtype=torch_dtype,
+        device_map="auto",
+        attn_implementation="flash_attention_2",
+    )
+
+    # Load prompts: prefer .pt if provided
+    if args.pt_prompts is not None:
+        loaded = torch.load(args.pt_prompts)
+        if not isinstance(loaded, list) or (len(loaded) > 0 and not isinstance(loaded[0], str)):
+            raise ValueError("--pt-prompts must point to a .pt containing List[str]")
+        prompts: List[str] = loaded
+        print(f"[Info] Loaded {len(prompts)} prompts from {args.pt_prompts}")
+    else:
+        prompts = PROMPTS
+
+    # Ensure utils.make_text_prompt can find expected fields
+    # utils expects args.model and args.tokenize
+    setattr(args, "model", args.model_id)
+    if not hasattr(args, "tokenize"):
+        setattr(args, "tokenize", False)
+
+    print(f"[Info] Running single-sample generation (batch size fixed to 1) ...\n")
+    global_index = 0
+    total_requests = len(prompts)
+    sum_request_time_ms = 0.0
+    wall_t0 = time.time()
+    stats_path = os.path.join(args.output_dir, "hf_transformers.stats.jsonl")
+
+    for rid, prompt in tqdm(list(enumerate(prompts)), total=len(prompts), desc="Generating"):
+        _t0 = time.time()
+        # Build single-sample inputs
+        tp = make_text_prompt(prompt)
+        tp = processor.apply_chat_template(tp, add_generation_prompt=True, tokenize=False)
+        inputs = processor(text=tp, return_tensors="pt", padding=True, use_audio_in_video=True)
+        inputs = inputs.to(model.device).to(model.dtype)
+
+        with torch.no_grad():
+            text_ids, audio = model.generate(
+                **inputs,
+                use_audio_in_video=True,
+            )
+
+        # Decode
+        try:
+            texts = processor.batch_decode(text_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)  # type: ignore[arg-type]
+        except Exception:
+            try:
+                texts = processor.tokenizer.batch_decode(text_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)  # type: ignore[attr-defined]
+            except Exception:
+                texts = [str(x) for x in text_ids]
+
+        text_str = texts[0] if len(texts) > 0 else ""
+        os.makedirs(args.output_dir, exist_ok=True)
+        out_txt = os.path.join(args.output_dir, f"{rid:05d}.txt")
+
+        content = None
+        if os.path.exists(out_txt):
+            try:
+                with open(out_txt, "r", encoding="utf-8") as f:
+                    content = f.read()
+            except Exception:
+                content = None
+        if content and "transformers_text_output:" in content:
+            try:
+                parts = content.split("transformers_text_output:\n", 1)
+                new_content = parts[0] + "transformers_text_output:\n" + text_str.strip() + "\n"
+                with open(out_txt, "w", encoding="utf-8") as f:
+                    f.write(new_content)
+            except Exception:
+                pass
+        else:
+            lines = [
+                "Prompt:\n",
+                str(prompt) + "\n",
+                "vllm_text_output:\n\n",
+                "transformers_text_output:\n",
+                text_str.strip() + "\n",
+            ]
+            try:
+                with open(out_txt, "w", encoding="utf-8") as f:
+                    f.writelines(lines)
+            except Exception:
+                pass
+
+        # Save wav per request
+        try:
+            if audio is not None:
+                if isinstance(audio, (list, tuple)):
+                    wav = audio[0]
+                else:
+                    # Tensor or array: take first sample if batched, else as-is
+                    if hasattr(audio, "ndim") and hasattr(audio, "shape") and len(audio.shape) > 1:
+                        wav = audio[0]
+                    else:
+                        wav = audio
+                wav_path = os.path.join(args.output_dir, f"output_{rid}.wav")
+                sf.write(wav_path, torch.as_tensor(wav).reshape(-1).detach().cpu().numpy(), samplerate=24000)
+        except Exception:
+            pass
+
+        # Stats per request
+        req_ms = (time.time() - _t0) * 1000.0
+        sum_request_time_ms += req_ms
+        try:
+            rec = {
+                "type": "transformers_request_stats",
+                "request_id": rid,
+                "batch_size": 1,
+                "per_request_time_ms": req_ms,
+                "num_tokens_out": int(text_ids.shape[-1]) if hasattr(text_ids, "shape") else None,
+            }
+            with open(stats_path, "a", encoding="utf-8") as f:
+                f.write(json.dumps(rec, ensure_ascii=False) + "\n")
+        except Exception:
+            pass
+
+        global_index += 1
+
+    print("[Done] Generated", global_index, "responses.")
+    wall_ms = (time.time() - wall_t0) * 1000.0
+    avg_e2e_ms = (sum_request_time_ms / total_requests) if total_requests > 0 else 0.0
+    print(f"[E2E] total_requests={total_requests}, wall_time_ms={wall_ms:.2f}, avg_e2e_per_request_ms={avg_e2e_ms:.2f}")
+    try:
+        summary = {
+            "type": "transformers_summary",
+            "total_requests": total_requests,
+            "wall_time_ms": wall_ms,
+            "sum_request_time_ms": sum_request_time_ms,
+            "avg_e2e_per_request_ms": avg_e2e_ms,
+        }
+        with open(stats_path, "a", encoding="utf-8") as f:
+            f.write(json.dumps(summary, ensure_ascii=False) + "\n")
+    except Exception:
+        pass
+
+
+if __name__ == "__main__":
+    main()
+
+

examples/offline_inference/qwen_2_5_omni/pipelined_end2end.py

@@ -0,0 +1,264 @@
+import argparse
+import os
+import os as _os_env_toggle
+import random
+
+import numpy as np
+import soundfile as sf
+import torch
+from utils import make_omni_prompt
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.entrypoints.pipelined_omni_llm import PipelinedOmniLLM
+
+_os_env_toggle.environ["VLLM_USE_V1"] = "1"
+
+SEED = 42
+# Set all random seeds
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+torch.cuda.manual_seed(SEED)
+torch.cuda.manual_seed_all(SEED)
+
+# Make PyTorch deterministic
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+
+# Set environment variables for deterministic behavior
+os.environ["PYTHONHASHSEED"] = str(SEED)
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model",
+        required=True,
+        help="Path to merged model directory (will be created if downloading).",
+    )
+    parser.add_argument("--thinker-model", type=str, default=None)
+    parser.add_argument("--talker-model", type=str, default=None)
+    parser.add_argument("--code2wav-model", type=str, default=None)
+    parser.add_argument(
+        "--hf-hub-id",
+        default="Qwen/Qwen2.5-Omni-7B",
+        help="Hugging Face repo id to download if needed.",
+    )
+    parser.add_argument(
+        "--hf-revision", default=None, help="Optional HF revision (branch/tag/commit)."
+    )
+    parser.add_argument(
+        "--prompts", nargs="+", default=None, help="Input text prompts."
+    )
+    parser.add_argument(
+        "--pt-prompts", type=str, default=None, help="Path to .pt file containing List[str] prompts (overrides --prompts when prompt_type=text)."
+    )
+    parser.add_argument(
+        "--voice-type", default="default", help="Voice type, e.g., m02, f030, default."
+    )
+    parser.add_argument(
+        "--code2wav-dir",
+        default=None,
+        help="Path to code2wav folder (contains spk_dict.pt).",
+    )
+    parser.add_argument(
+        "--dit-ckpt", default=None, help="Path to DiT checkpoint file (e.g., dit.pt)."
+    )
+    parser.add_argument(
+        "--bigvgan-ckpt", default=None, help="Path to BigVGAN checkpoint file."
+    )
+    parser.add_argument(
+        "--dtype", default="bfloat16", choices=["float16", "bfloat16", "float32"]
+    )
+    parser.add_argument("--max-model-len", type=int, default=32768)
+    parser.add_argument(
+        "--init-sleep-seconds",
+        type=int,
+        default=20,
+        help="Sleep seconds after starting each stage process to allow initialization (default: 20)",
+    )
+
+    parser.add_argument("--thinker-only", action="store_true")
+    parser.add_argument("--text-only", action="store_true")
+    parser.add_argument("--do-wave", action="store_true")
+    parser.add_argument(
+        "--prompt_type",
+        choices=[
+            "text",
+            "audio",
+            "audio-long",
+            "audio-long-chunks",
+            "audio-long-expand-chunks",
+            "image",
+            "video",
+            "video-frames",
+            "audio-in-video",
+            "audio-in-video-v2",
+            "audio-multi-round",
+            "badcase-vl",
+            "badcase-text",
+            "badcase-image-early-stop",
+            "badcase-two-audios",
+            "badcase-two-videos",
+            "badcase-multi-round",
+            "badcase-voice-type",
+            "badcase-voice-type-v2",
+            "badcase-audio-tower-1",
+            "badcase-audio-only",
+        ],
+        default="text",
+    )
+    parser.add_argument("--use-torchvision", action="store_true")
+    parser.add_argument("--tokenize", action="store_true")
+    parser.add_argument(
+        "--output-wav", default="output.wav", help="[Deprecated] Output wav directory (use --output-dir)."
+    )
+    parser.add_argument(
+        "--output-dir", default="outputs", help="Output directory to save text and wav files together."
+    )
+    parser.add_argument(
+        "--thinker-hidden-states-dir",
+        default="thinker_hidden_states",
+        help="Path to thinker hidden states directory.",
+    )
+    parser.add_argument(
+        "--batch-timeout",
+        type=int,
+        default=5,
+        help="Timeout for batching in seconds (default: 5)",
+    )
+    parser.add_argument(
+        "--init-timeout",
+        type=int,
+        default=300,
+        help="Timeout for initializing stages in seconds (default: 300)",
+    )
+    parser.add_argument(
+        "--shm-threshold-bytes",
+        type=int,
+        default=65536,
+        help="Threshold for using shared memory in bytes (default: 65536)",
+    )
+    parser.add_argument(
+        "--enable-stats",
+        action="store_true",
+        default=True,
+        help="Enable writing detailed statistics (default: enabled)",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+    model_name = args.model
+    # If pt-prompts provided and prompt_type is text, load and override args.prompts
+    try:
+        if getattr(args, "pt_prompts", None) and args.prompt_type == "text":
+            loaded = torch.load(args.pt_prompts)
+            if not isinstance(loaded, list) or (len(loaded) > 0 and not isinstance(loaded[0], str)):
+                raise ValueError("--pt-prompts must point to a .pt containing List[str]")
+            args.prompts = loaded
+            print(f"[Info] Loaded {len(args.prompts)} prompts from {args.pt_prompts}")
+    except Exception as e:
+        raise
+
+    if args.prompts is None:
+        raise ValueError("No prompts provided. Use --prompts ... or --pt-prompts <file.pt> (with --prompt_type text)")
+    omni_llm = PipelinedOmniLLM(
+        model=model_name,
+        log_stats=args.enable_stats,
+        log_file="omni_llm_pipeline.log",
+        init_sleep_seconds=args.init_sleep_seconds,
+        batch_timeout=args.batch_timeout,
+        init_timeout=args.init_timeout,
+        shm_threshold_bytes=args.shm_threshold_bytes,
+    )
+    thinker_sampling_params = SamplingParams(
+        temperature=0.0,  # Deterministic - no randomness
+        top_p=1.0,  # Disable nucleus sampling
+        top_k=-1,  # Disable top-k sampling
+        max_tokens=2048,
+        seed=SEED,  # Fixed seed for sampling
+        detokenize=True,
+        repetition_penalty=1.1,
+    )
+    talker_sampling_params = SamplingParams(
+        temperature=0.9,
+        top_p=0.8,
+        top_k=40,
+        max_tokens=2048,
+        seed=SEED,  # Fixed seed for sampling
+        detokenize=True,
+        repetition_penalty=1.05,
+        stop_token_ids=[8294],
+    )
+    code2wav_sampling_params = SamplingParams(
+        temperature=0.0,  # Deterministic - no randomness
+        top_p=1.0,  # Disable nucleus sampling
+        top_k=-1,  # Disable top-k sampling
+        max_tokens=2048,
+        seed=SEED,  # Fixed seed for sampling
+        detokenize=True,
+        repetition_penalty=1.1,
+    )
+
+    sampling_params_list = [
+        thinker_sampling_params,
+        talker_sampling_params,
+        code2wav_sampling_params,
+    ]
+
+    prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts]
+    omni_outputs = omni_llm.generate(prompt, sampling_params_list)
+
+    # Determine output directory: prefer --output-dir; fallback to --output-wav
+    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
+    os.makedirs(output_dir, exist_ok=True)
+    for stage_outputs in omni_outputs:
+        if stage_outputs.final_output_type == "text":
+            for output in stage_outputs.request_output:
+                request_id = int(output.request_id)
+                text_output = output.outputs[0].text
+                # Save aligned text file per request
+                prompt_text = args.prompts[request_id]
+                out_txt = os.path.join(output_dir, f"{request_id:05d}.txt")
+                lines = []
+                lines.append("Prompt:\n")
+                lines.append(str(prompt_text) + "\n")
+                lines.append("vllm_text_output:\n")
+                lines.append(str(text_output).strip() + "\n")
+                lines.append("transformers_text_output:\n\n")
+                try:
+                    with open(out_txt, "w", encoding="utf-8") as f:
+                        f.writelines(lines)
+                except Exception:
+                    pass
+                print(f"Request ID: {request_id}, Text saved to {out_txt}")
+        elif stage_outputs.final_output_type == "audio":
+            for output in stage_outputs.request_output:
+                request_id = int(output.request_id)
+                audio_tensor = output.multimodal_output["audio"]
+                output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav")
+                sf.write(
+                    output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000
+                )
+                # Save aligned text file with empty text (no text output in this stage)
+                prompt_text = args.prompts[request_id]
+                out_txt = os.path.join(output_dir, f"{request_id:05d}.txt")
+                lines = []
+                lines.append("Prompt:\n")
+                lines.append(str(prompt_text) + "\n")
+                lines.append("vllm_text_output:\n\n")
+                lines.append("transformers_text_output:\n\n")
+                try:
+                    with open(out_txt, "w", encoding="utf-8") as f:
+                        f.writelines(lines)
+                except Exception:
+                    pass
+                print(f"Request ID: {request_id}, Saved audio to {output_wav} and text to {out_txt}")
+
+
+if __name__ == "__main__":
+    main()