     1→"""Offline inference demo for Qwen3 TTS via vLLM Omni.
     2→
     3→Provides single and batch sample inputs for CustomVoice, VoiceDesign, and Base
     4→tasks, then runs Omni generation and saves output wav files.
     5→"""
     6→
     7→import asyncio
     8→import logging
     9→import os
    10→import time
    11→from typing import Any, NamedTuple
    12→
    13→import soundfile as sf
    14→import torch
    15→
    16→os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
    17→
    18→from vllm.utils.argparse_utils import FlexibleArgumentParser
    19→
    20→from vllm_omni import AsyncOmni, Omni
    21→
    22→logger = logging.getLogger(__name__)
    23→
    24→
    25→class QueryResult(NamedTuple):
    26→    """Container for a prepared Omni request."""
    27→
    28→    inputs: dict
    29→    model_name: str
    30→
    31→
    32→def _estimate_prompt_len(
    33→    additional_information: dict[str, Any],
    34→    model_name: str,
    35→    _cache: dict[str, Any] = {},
    36→) -> int:
    37→    """Estimate prompt_token_ids placeholder length for the Talker stage.
    38→
    39→    The AR Talker replaces all input embeddings via ``preprocess``, so the
    40→    placeholder values are irrelevant but the **length** must match the
    41→    embeddings that ``preprocess`` will produce.
    42→    """
    43→    try:
    44→        from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
    45→        from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
    46→            Qwen3TTSTalkerForConditionalGeneration,
    47→        )
    48→
    49→        if model_name not in _cache:
    50→            from transformers import AutoTokenizer
    51→
    52→            tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="left")
    53→            cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True)
    54→
    55→            # Load speech tokenizer (codec encoder) for exact ref_code_len.
    56→            speech_tok = None
    57→            try:
    58→                import os
    59→
    60→                from transformers.utils import cached_file
    61→
    62→                from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer
    63→
    64→                st_cfg_path = cached_file(model_name, "speech_tokenizer/config.json")
    65→                if st_cfg_path:
    66→                    speech_tok = Qwen3TTSTokenizer.from_pretrained(
    67→                        os.path.dirname(st_cfg_path), torch_dtype=torch.bfloat16
    68→                    )
    69→                    logger.info("Loaded speech tokenizer for exact ref_code_len estimation")
    70→            except Exception as e:
    71→                logger.debug("Could not load speech tokenizer: %s", e)
    72→
    73→            _cache[model_name] = (tok, getattr(cfg, "talker_config", None), speech_tok)
    74→
    75→        tok, tcfg, speech_tok = _cache[model_name]
    76→        task_type = (additional_information.get("task_type") or ["CustomVoice"])[0]
    77→
    78→        def _estimate_ref_code_len(ref_audio: object) -> int | None:
    79→            """Encode ref_audio with the actual codec to get exact frame count."""
    80→            if not isinstance(ref_audio, (str, list)):
    81→                return None
    82→            audio_path = ref_audio[0] if isinstance(ref_audio, list) else ref_audio
    83→            if not isinstance(audio_path, str) or not audio_path.strip():
    84→                return None
    85→            try:
    86→                from vllm.multimodal.media import MediaConnector
    87→
    88→                connector = MediaConnector(allowed_local_media_path="/")
    89→                audio, sr = connector.fetch_audio(audio_path)
    90→                import numpy as np
    91→
    92→                wav_np = np.asarray(audio, dtype=np.float32)
    93→
    94→                if speech_tok is not None:
    95→                    enc = speech_tok.encode(wav_np, sr=int(sr), return_dict=True)
    96→                    ref_code = getattr(enc, "audio_codes", None)
    97→                    if isinstance(ref_code, list):
    98→                        ref_code = ref_code[0] if ref_code else None
    99→                    if ref_code is not None and hasattr(ref_code, "shape"):
   100→                        shape = ref_code.shape
   101→                        return int(shape[0]) if len(shape) == 2 else int(shape[1]) if len(shape) == 3 else None
   102→
   103→                # Fallback: estimate from duration
   104→                codec_hz = getattr(tcfg, "codec_frame_rate", None) or 12
   105→                return int(len(audio) / sr * codec_hz)
   106→            except Exception:
   107→                return None
   108→
   109→        return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
   110→            additional_information=additional_information,
   111→            task_type=task_type,
   112→            tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
   113→            codec_language_id=getattr(tcfg, "codec_language_id", None),
   114→            spk_is_dialect=getattr(tcfg, "spk_is_dialect", None),
   115→            estimate_ref_code_len=_estimate_ref_code_len,
   116→        )
   117→    except Exception as exc:
   118→        logger.warning("Failed to estimate prompt length, using fallback 2048: %s", exc)
   119→        return 2048
   120→
   121→
   122→def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
   123→    """Build CustomVoice sample inputs.
   124→
   125→    Args:
   126→        use_batch_sample: When True, return a batch of prompts; otherwise a single prompt.
   127→
   128→    Returns:
   129→        QueryResult with Omni inputs and the CustomVoice model path.
   130→    """
   131→    task_type = "CustomVoice"
   132→    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
   133→    if use_batch_sample:
   134→        texts = [
   135→            "其实我真的有发现，我是一个特别善于观察别人情绪的人。",
   136→            "She said she would be here by noon.",
   137→            "I like you very much.",
   138→            "Really, you do?",
   139→            "Yes, absolutely.",
   140→        ]
   141→        instructs = ["", "Very happy.", "Very happy.", "Very happy.", "Very happy."]
   142→        languages = ["Chinese", "English", "English", "English", "English"]
   143→        speakers = ["Vivian", "Ryan", "Ryan", "Ryan", "Ryan"]
   144→        inputs = []
   145→        for text, instruct, language, speaker in zip(texts, instructs, languages, speakers):
   146→            additional_information = {
   147→                "task_type": [task_type],
   148→                "text": [text],
   149→                "instruct": [instruct],
   150→                "language": [language],
   151→                "speaker": [speaker],
   152→                "max_new_tokens": [2048],
   153→            }
   154→            inputs.append(
   155→                {
   156→                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
   157→                    "additional_information": additional_information,
   158→                }
   159→            )
   160→    else:
   161→        text = "其实我真的有发现，我是一个特别善于观察别人情绪的人。"
   162→        language = "Chinese"
   163→        speaker = "Vivian"
   164→        instruct = "用特别愤怒的语气说"
   165→        additional_information = {
   166→            "task_type": [task_type],
   167→            "text": [text],
   168→            "language": [language],
   169→            "speaker": [speaker],
   170→            "instruct": [instruct],
   171→            "max_new_tokens": [2048],
   172→        }
   173→        inputs = {
   174→            "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
   175→            "additional_information": additional_information,
   176→        }
   177→    return QueryResult(
   178→        inputs=inputs,
   179→        model_name=model_name,
   180→    )
   181→
   182→
   183→def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult:
   184→    """Build VoiceDesign sample inputs.
   185→
   186→    Args:
   187→        use_batch_sample: When True, return a batch of prompts; otherwise a single prompt.
   188→
   189→    Returns:
   190→        QueryResult with Omni inputs and the VoiceDesign model path.
   191→    """
   192→    task_type = "VoiceDesign"
   193→    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
   194→    if use_batch_sample:
   195→        texts = [
   196→            "哥哥，你回来啦，人家等了你好久好久了，要抱抱！",
   197→            "It's in the top drawer... wait, it's empty? No way, that's impossible! I'm sure I put it there!",
   198→        ]
   199→        instructs = [
   200→            "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。",
   201→            "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice.",
   202→        ]
   203→        languages = ["Chinese", "English"]
   204→        inputs = []
   205→        for text, instruct, language in zip(texts, instructs, languages):
   206→            additional_information = {
   207→                "task_type": [task_type],
   208→                "text": [text],
   209→                "language": [language],
   210→                "instruct": [instruct],
   211→                "max_new_tokens": [2048],
   212→                "non_streaming_mode": [True],
   213→            }
   214→            inputs.append(
   215→                {
   216→                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
   217→                    "additional_information": additional_information,
   218→                }
   219→            )
   220→    else:
   221→        text = "哥哥，你回来啦，人家等了你好久好久了，要抱抱！"
   222→        instruct = "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。"
   223→        language = "Chinese"
   224→        additional_information = {
   225→            "task_type": [task_type],
   226→            "text": [text],
   227→            "language": [language],
   228→            "instruct": [instruct],
   229→            "max_new_tokens": [2048],
   230→            "non_streaming_mode": [True],
   231→        }
   232→        inputs = {
   233→            "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
   234→            "additional_information": additional_information,
   235→        }
   236→    return QueryResult(
   237→        inputs=inputs,
   238→        model_name=model_name,
   239→    )
   240→
   241→
   242→def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> QueryResult:
   243→    """Build Base (voice clone) sample inputs.
   244→
   245→    Args:
   246→        use_batch_sample: When True, return a batch of prompts (Case 2).
   247→        mode_tag: "icl" or "xvec_only" to control x_vector_only_mode behavior.
   248→
   249→    Returns:
   250→        QueryResult with Omni inputs and the Base model path.
   251→    """
   252→    task_type = "Base"
   253→    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
   254→    ref_audio_path_1 = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
   255→    ref_audio_single = ref_audio_path_1
   256→    ref_text_single = (
   257→        "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
   258→    )
   259→    syn_text_single = "Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye."
   260→    syn_lang_single = "Auto"
   261→    x_vector_only_mode = mode_tag == "xvec_only"
   262→    if use_batch_sample:
   263→        syn_text_batch = [
   264→            "Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye.",
   265→            "其实我真的有发现，我是一个特别善于观察别人情绪的人。",
   266→        ]
   267→        syn_lang_batch = ["Chinese", "English"]
   268→        inputs = []
   269→        for text, language in zip(syn_text_batch, syn_lang_batch):
   270→            additional_information = {
   271→                "task_type": [task_type],
   272→                "ref_audio": [ref_audio_single],
   273→                "ref_text": [ref_text_single],
   274→                "text": [text],
   275→                "language": [language],
   276→                "x_vector_only_mode": [x_vector_only_mode],
   277→                "max_new_tokens": [2048],
   278→            }
   279→            inputs.append(
   280→                {
   281→                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
   282→                    "additional_information": additional_information,
   283→                }
   284→            )
   285→    else:
   286→        additional_information = {
   287→            "task_type": [task_type],
   288→            "ref_audio": [ref_audio_single],
   289→            "ref_text": [ref_text_single],
   290→            "text": [syn_text_single],
   291→            "language": [syn_lang_single],
   292→            "x_vector_only_mode": [x_vector_only_mode],
   293→            "max_new_tokens": [2048],
   294→        }
   295→        inputs = {
   296→            "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
   297→            "additional_information": additional_information,
   298→        }
   299→    return QueryResult(
   300→        inputs=inputs,
   301→        model_name=model_name,
   302→    )
   303→
   304→
   305→query_map = {
   306→    "CustomVoice": get_custom_voice_query,
   307→    "VoiceDesign": get_voice_design_query,
   308→    "Base": get_base_query,
   309→}
   310→
   311→
   312→def _build_inputs(args) -> tuple[str, list]:
   313→    """Resolve model name and inputs list from CLI args."""
   314→    if args.batch_size < 1 or (args.batch_size & (args.batch_size - 1)) != 0:
   315→        raise ValueError(
   316→            f"--batch-size must be a power of two (got {args.batch_size}); "
   317→            "non-power-of-two values do not align with CUDA graph capture sizes "
   318→            "of Code2Wav."
   319→        )
   320→
   321→    query_func = query_map[args.query_type]
   322→    if args.query_type in {"CustomVoice", "VoiceDesign"}:
   323→        query_result = query_func(use_batch_sample=args.use_batch_sample)
   324→    elif args.query_type == "Base":
   325→        query_result = query_func(use_batch_sample=args.use_batch_sample, mode_tag=args.mode_tag)
   326→    else:
   327→        query_result = query_func()
   328→
   329→    model_name = query_result.model_name
   330→
   331→    if args.txt_prompts:
   332→        with open(args.txt_prompts) as f:
   333→            lines = [line.strip() for line in f if line.strip()]
   334→        if not lines:
   335→            raise ValueError(f"No valid prompts found in {args.txt_prompts}")
   336→        template = query_result.inputs if not isinstance(query_result.inputs, list) else query_result.inputs[0]
   337→        template_info = template["additional_information"]
   338→        inputs = [
   339→            {
   340→                "prompt_token_ids": [0] * _estimate_prompt_len({**template_info, "text": [t]}, model_name),
   341→                "additional_information": {**template_info, "text": [t]},
   342→            }
   343→            for t in lines
   344→        ]
   345→    else:
   346→        inputs = query_result.inputs if isinstance(query_result.inputs, list) else [query_result.inputs]
   347→
   348→    return model_name, inputs
   349→
   350→
   351→def _save_wav(output_dir: str, request_id: str, mm: dict) -> None:
   352→    """Concatenate audio chunks and write to a wav file."""
   353→    audio_data = mm["audio"]
   354→    sr_raw = mm["sr"]
   355→    sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
   356→    sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
   357→    audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
   358→    out_wav = os.path.join(output_dir, f"output_{request_id}.wav")
   359→    sf.write(out_wav, audio_tensor.float().cpu().numpy().flatten(), samplerate=sr, format="WAV")
   360→    logger.info(f"Request ID: {request_id}, Saved audio to {out_wav}")
   361→
   362→
   363→def main(args):
   364→    """Run offline inference with Omni."""
   365→    model_name, inputs = _build_inputs(args)
   366→    output_dir = args.output_dir
   367→    os.makedirs(output_dir, exist_ok=True)
   368→
   369→    omni = Omni(
   370→        model=model_name,
   371→        stage_configs_path=args.stage_configs_path,
   372→        log_stats=args.log_stats,
   373→        stage_init_timeout=args.stage_init_timeout,
   374→        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
   375→    )
   376→
   377→    batch_size = args.batch_size
   378→    for batch_start in range(0, len(inputs), batch_size):
   379→        batch = inputs[batch_start : batch_start + batch_size]
   380→        for stage_outputs in omni.generate(batch):
   381→            output = stage_outputs.request_output
   382→            _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
   383→
   384→
   385→async def main_streaming(args):
   386→    """Run offline inference with AsyncOmni, logging each audio chunk as it arrives."""
   387→    model_name, inputs = _build_inputs(args)
   388→    output_dir = args.output_dir
   389→    os.makedirs(output_dir, exist_ok=True)
   390→
   391→    omni = AsyncOmni(
   392→        model=model_name,
   393→        stage_configs_path=args.stage_configs_path,
   394→        log_stats=args.log_stats,
   395→        stage_init_timeout=args.stage_init_timeout,
   396→        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
   397→    )
   398→
   399→    for i, prompt in enumerate(inputs):
   400→        request_id = str(i)
   401→        t_start = time.perf_counter()
   402→        t_prev = t_start
   403→        chunk_idx = 0
   404→        async for stage_output in omni.generate(prompt, request_id=request_id):
   405→            mm = stage_output.request_output.outputs[0].multimodal_output
   406→            if not stage_output.finished:
   407→                t_now = time.perf_counter()
   408→                audio = mm.get("audio")
   409→                n = len(audio) if isinstance(audio, list) else (0 if audio is None else 1)
   410→                dt_ms = (t_now - t_prev) * 1000
   411→                ttfa_ms = (t_now - t_start) * 1000
   412→                if chunk_idx == 0:
   413→                    logger.info(f"Request {request_id}: chunk {chunk_idx} samples={n} TTFA={ttfa_ms:.1f}ms")
   414→                else:
   415→                    logger.info(f"Request {request_id}: chunk {chunk_idx} samples={n} inter_chunk={dt_ms:.1f}ms")
   416→                t_prev = t_now
   417→                chunk_idx += 1
   418→            else:
   419→                t_end = time.perf_counter()
   420→                total_ms = (t_end - t_start) * 1000
   421→                logger.info(f"Request {request_id}: done total={total_ms:.1f}ms chunks={chunk_idx}")
   422→                _save_wav(output_dir, request_id, mm)
   423→
   424→
   425→def parse_args():
   426→    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
   427→    parser.add_argument(
   428→        "--query-type",
   429→        "-q",
   430→        type=str,
   431→        default="CustomVoice",
   432→        choices=query_map.keys(),
   433→        help="Query type.",
   434→    )
   435→    parser.add_argument(
   436→        "--log-stats",
   437→        action="store_true",
   438→        default=False,
   439→        help="Enable writing detailed statistics (default: disabled)",
   440→    )
   441→    parser.add_argument(
   442→        "--stage-init-timeout",
   443→        type=int,
   444→        default=300,
   445→        help="Timeout for initializing a single stage in seconds (default: 300)",
   446→    )
   447→    parser.add_argument(
   448→        "--batch-timeout",
   449→        type=int,
   450→        default=5,
   451→        help="Timeout for batching in seconds (default: 5)",
   452→    )
   453→    parser.add_argument(
   454→        "--init-timeout",
   455→        type=int,
   456→        default=300,
   457→        help="Timeout for initializing stages in seconds (default: 300)",
   458→    )
   459→    parser.add_argument(
   460→        "--shm-threshold-bytes",
   461→        type=int,
   462→        default=65536,
   463→        help="Threshold for using shared memory in bytes (default: 65536)",
   464→    )
   465→    parser.add_argument(
   466→        "--output-dir",
   467→        default="output_audio",
   468→        help="Output directory for generated wav files (default: output_audio).",
   469→    )
   470→    parser.add_argument(
   471→        "--num-prompts",
   472→        type=int,
   473→        default=1,
   474→        help="Number of prompts to generate.",
   475→    )
   476→    parser.add_argument(
   477→        "--txt-prompts",
   478→        type=str,
   479→        default=None,
   480→        help="Path to a .txt file with one prompt per line (preferred).",
   481→    )
   482→    parser.add_argument(
   483→        "--stage-configs-path",
   484→        type=str,
   485→        default=None,
   486→        help="Path to a stage configs file.",
   487→    )
   488→    parser.add_argument(
   489→        "--audio-path",
   490→        "-a",
   491→        type=str,
   492→        default=None,
   493→        help="Path to local audio file. If not provided, uses default audio asset.",
   494→    )
   495→    parser.add_argument(
   496→        "--sampling-rate",
   497→        type=int,
   498→        default=16000,
   499→        help="Sampling rate for audio loading (default: 16000).",
   500→    )
   501→    parser.add_argument(
   502→        "--log-dir",
   503→        type=str,
   504→        default="logs",
   505→        help="Log directory (default: logs).",
   506→    )
   507→    parser.add_argument(
   508→        "--py-generator",
   509→        action="store_true",
   510→        default=False,
   511→        help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
   512→    )
   513→    parser.add_argument(
   514→        "--use-batch-sample",
   515→        action="store_true",
   516→        default=False,
   517→        help="Use batch input sample for CustomVoice/VoiceDesign/Base query.",
   518→    )
   519→    parser.add_argument(
   520→        "--mode-tag",
   521→        type=str,
   522→        default="icl",
   523→        choices=["icl", "xvec_only"],
   524→        help="Mode tag for Base query x_vector_only_mode (default: icl).",
   525→    )
   526→    parser.add_argument(
   527→        "--streaming",
   528→        action="store_true",
   529→        default=False,
   530→        help="Stream audio chunks as they arrive via AsyncOmni (async_chunk mode only).",
   531→    )
   532→    parser.add_argument(
   533→        "--batch-size",
   534→        type=int,
   535→        default=1,
   536→        help="Number of prompts per batch (default: 1, sequential).",
   537→    )
   538→    parser.add_argument(
   539→        "--enable-diffusion-pipeline-profiler",
   540→        action="store_true",
   541→        help="Enable diffusion pipeline profiler to display stage durations.",
   542→    )
   543→
   544→    return parser.parse_args()
   545→
   546→
   547→if __name__ == "__main__":
   548→    args = parse_args()
   549→    if args.streaming:
   550→        asyncio.run(main_streaming(args))
   551→    else:
   552→        main(args)
   553→

"""Offline inference demo for Qwen3 TTS via vLLM Omni.

Provides single and batch sample inputs for CustomVoice, VoiceDesign, and Base
tasks, then runs Omni generation and saves output wav files.
"""

import asyncio
import logging
import os
import time
from typing import Any, NamedTuple

import soundfile as sf
import torch

os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"

from vllm.utils.argparse_utils import FlexibleArgumentParser

from vllm_omni import AsyncOmni, Omni

logger = logging.getLogger(__name__)

class QueryResult(NamedTuple):
    """Container for a prepared Omni request."""

    inputs: dict
    model_name: str

def _estimate_prompt_len(
    additional_information: dict[str, Any],
    model_name: str,
    _cache: dict[str, Any] = {},
) -> int:
    """Estimate prompt_token_ids placeholder length for the Talker stage.

    The AR Talker replaces all input embeddings via ``preprocess``, so the
    placeholder values are irrelevant but the **length** must match the
    embeddings that ``preprocess`` will produce.
    """
    try:
        from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
        from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
            Qwen3TTSTalkerForConditionalGeneration,
        )

        if model_name not in _cache:
            from transformers import AutoTokenizer

            tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="left")
            cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True)

            # Load speech tokenizer (codec encoder) for exact ref_code_len.
            speech_tok = None
            try:
                import os

                from transformers.utils import cached_file

                from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer

                st_cfg_path = cached_file(model_name, "speech_tokenizer/config.json")
                if st_cfg_path:
                    speech_tok = Qwen3TTSTokenizer.from_pretrained(
                        os.path.dirname(st_cfg_path), torch_dtype=torch.bfloat16
                    )
                    logger.info("Loaded speech tokenizer for exact ref_code_len estimation")
            except Exception as e:
                logger.debug("Could not load speech tokenizer: %s", e)

            _cache[model_name] = (tok, getattr(cfg, "talker_config", None), speech_tok)

        tok, tcfg, speech_tok = _cache[model_name]
        task_type = (additional_information.get("task_type") or ["CustomVoice"])[0]

        def _estimate_ref_code_len(ref_audio: object) -> int | None:
            """Encode ref_audio with the actual codec to get exact frame count."""
            if not isinstance(ref_audio, (str, list)):
                return None
            audio_path = ref_audio[0] if isinstance(ref_audio, list) else ref_audio
            if not isinstance(audio_path, str) or not audio_path.strip():
                return None
            try:
                from vllm.multimodal.media import MediaConnector

                connector = MediaConnector(allowed_local_media_path="/")
                audio, sr = connector.fetch_audio(audio_path)
                import numpy as np

                wav_np = np.asarray(audio, dtype=np.float32)

                if speech_tok is not None:
                    enc = speech_tok.encode(wav_np, sr=int(sr), return_dict=True)
                    ref_code = getattr(enc, "audio_codes", None)
                    if isinstance(ref_code, list):
                        ref_code = ref_code[0] if ref_code else None
                    if ref_code is not None and hasattr(ref_code, "shape"):
                        shape = ref_code.shape
                        return int(shape[0]) if len(shape) == 2 else int(shape[1]) if len(shape) == 3 else None

                # Fallback: estimate from duration
                codec_hz = getattr(tcfg, "codec_frame_rate", None) or 12
                return int(len(audio) / sr * codec_hz)
            except Exception:
                return None

        return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
            additional_information=additional_information,
            task_type=task_type,
            tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
            codec_language_id=getattr(tcfg, "codec_language_id", None),
            spk_is_dialect=getattr(tcfg, "spk_is_dialect", None),
            estimate_ref_code_len=_estimate_ref_code_len,
        )
    except Exception as exc:
        logger.warning("Failed to estimate prompt length, using fallback 2048: %s", exc)
        return 2048

def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
    """Build CustomVoice sample inputs.

    Args:
        use_batch_sample: When True, return a batch of prompts; otherwise a single prompt.

    Returns:
        QueryResult with Omni inputs and the CustomVoice model path.
    """
    task_type = "CustomVoice"
    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
    if use_batch_sample:
        texts = [
            "其实我真的有发现，我是一个特别善于观察别人情绪的人。",
            "She said she would be here by noon.",
            "I like you very much.",
            "Really, you do?",
            "Yes, absolutely.",
        ]
        instructs = ["", "Very happy.", "Very happy.", "Very happy.", "Very happy."]
        languages = ["Chinese", "English", "English", "English", "English"]
        speakers = ["Vivian", "Ryan", "Ryan", "Ryan", "Ryan"]
        inputs = []
        for text, instruct, language, speaker in zip(texts, instructs, languages, speakers):
            additional_information = {
                "task_type": [task_type],
                "text": [text],
                "instruct": [instruct],
                "language": [language],
                "speaker": [speaker],
                "max_new_tokens": [2048],
            }
            inputs.append(
                {
                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
                    "additional_information": additional_information,
                }
            )
    else:
        text = "其实我真的有发现，我是一个特别善于观察别人情绪的人。"
        language = "Chinese"
        speaker = "Vivian"
        instruct = "用特别愤怒的语气说"
        additional_information = {
            "task_type": [task_type],
            "text": [text],
            "language": [language],
            "speaker": [speaker],
            "instruct": [instruct],
            "max_new_tokens": [2048],
        }
        inputs = {
            "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
            "additional_information": additional_information,
        }
    return QueryResult(
        inputs=inputs,
        model_name=model_name,
    )

def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult:
    """Build VoiceDesign sample inputs.

    Args:
        use_batch_sample: When True, return a batch of prompts; otherwise a single prompt.

    Returns:
        QueryResult with Omni inputs and the VoiceDesign model path.
    """
    task_type = "VoiceDesign"
    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
    if use_batch_sample:
        texts = [
            "哥哥，你回来啦，人家等了你好久好久了，要抱抱！",
            "It's in the top drawer... wait, it's empty? No way, that's impossible! I'm sure I put it there!",
        ]
        instructs = [
            "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。",
            "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice.",
        ]
        languages = ["Chinese", "English"]
        inputs = []
        for text, instruct, language in zip(texts, instructs, languages):
            additional_information = {
                "task_type": [task_type],
                "text": [text],
                "language": [language],
                "instruct": [instruct],
                "max_new_tokens": [2048],
                "non_streaming_mode": [True],
            }
            inputs.append(
                {
                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
                    "additional_information": additional_information,
                }
            )
    else:
        text = "哥哥，你回来啦，人家等了你好久好久了，要抱抱！"
        instruct = "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。"
        language = "Chinese"
        additional_information = {
            "task_type": [task_type],
            "text": [text],
            "language": [language],
            "instruct": [instruct],
            "max_new_tokens": [2048],
            "non_streaming_mode": [True],
        }
        inputs = {
            "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
            "additional_information": additional_information,
        }
    return QueryResult(
        inputs=inputs,
        model_name=model_name,
    )

def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> QueryResult:
    """Build Base (voice clone) sample inputs.

    Args:
        use_batch_sample: When True, return a batch of prompts (Case 2).
        mode_tag: "icl" or "xvec_only" to control x_vector_only_mode behavior.

    Returns:
        QueryResult with Omni inputs and the Base model path.
    """
    task_type = "Base"
    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
    ref_audio_path_1 = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
    ref_audio_single = ref_audio_path_1
    ref_text_single = (
        "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
    )
    syn_text_single = "Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye."
    syn_lang_single = "Auto"
    x_vector_only_mode = mode_tag == "xvec_only"
    if use_batch_sample:
        syn_text_batch = [
            "Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye.",
            "其实我真的有发现，我是一个特别善于观察别人情绪的人。",
        ]
        syn_lang_batch = ["Chinese", "English"]
        inputs = []
        for text, language in zip(syn_text_batch, syn_lang_batch):
            additional_information = {
                "task_type": [task_type],
                "ref_audio": [ref_audio_single],
                "ref_text": [ref_text_single],
                "text": [text],
                "language": [language],
                "x_vector_only_mode": [x_vector_only_mode],
                "max_new_tokens": [2048],
            }
            inputs.append(
                {
                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
                    "additional_information": additional_information,
                }
            )
    else:
        additional_information = {
            "task_type": [task_type],
            "ref_audio": [ref_audio_single],
            "ref_text": [ref_text_single],
            "text": [syn_text_single],
            "language": [syn_lang_single],
            "x_vector_only_mode": [x_vector_only_mode],
            "max_new_tokens": [2048],
        }
        inputs = {
            "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
            "additional_information": additional_information,
        }
    return QueryResult(
        inputs=inputs,
        model_name=model_name,
    )

query_map = {
    "CustomVoice": get_custom_voice_query,
    "VoiceDesign": get_voice_design_query,
    "Base": get_base_query,
}

def _build_inputs(args) -> tuple[str, list]:
    """Resolve model name and inputs list from CLI args."""
    if args.batch_size < 1 or (args.batch_size & (args.batch_size - 1)) != 0:
        raise ValueError(
            f"--batch-size must be a power of two (got {args.batch_size}); "
            "non-power-of-two values do not align with CUDA graph capture sizes "
            "of Code2Wav."
        )

    query_func = query_map[args.query_type]
    if args.query_type in {"CustomVoice", "VoiceDesign"}:
        query_result = query_func(use_batch_sample=args.use_batch_sample)
    elif args.query_type == "Base":
        query_result = query_func(use_batch_sample=args.use_batch_sample, mode_tag=args.mode_tag)
    else:
        query_result = query_func()

    model_name = query_result.model_name

    if args.txt_prompts:
        with open(args.txt_prompts) as f:
            lines = [line.strip() for line in f if line.strip()]
        if not lines:
            raise ValueError(f"No valid prompts found in {args.txt_prompts}")
        template = query_result.inputs if not isinstance(query_result.inputs, list) else query_result.inputs[0]
        template_info = template["additional_information"]
        inputs = [
            {
                "prompt_token_ids": [0] * _estimate_prompt_len({**template_info, "text": [t]}, model_name),
                "additional_information": {**template_info, "text": [t]},
            }
            for t in lines
        ]
    else:
        inputs = query_result.inputs if isinstance(query_result.inputs, list) else [query_result.inputs]

    return model_name, inputs

def _save_wav(output_dir: str, request_id: str, mm: dict) -> None:
    """Concatenate audio chunks and write to a wav file."""
    audio_data = mm["audio"]
    sr_raw = mm["sr"]
    sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
    sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
    audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
    out_wav = os.path.join(output_dir, f"output_{request_id}.wav")
    sf.write(out_wav, audio_tensor.float().cpu().numpy().flatten(), samplerate=sr, format="WAV")
    logger.info(f"Request ID: {request_id}, Saved audio to {out_wav}")

def main(args):
    """Run offline inference with Omni."""
    model_name, inputs = _build_inputs(args)
    output_dir = args.output_dir
    os.makedirs(output_dir, exist_ok=True)

    omni = Omni(
        model=model_name,
        stage_configs_path=args.stage_configs_path,
        log_stats=args.log_stats,
        stage_init_timeout=args.stage_init_timeout,
        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
    )

    batch_size = args.batch_size
    for batch_start in range(0, len(inputs), batch_size):
        batch = inputs[batch_start : batch_start + batch_size]
        for stage_outputs in omni.generate(batch):
            output = stage_outputs.request_output
            _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)

async def main_streaming(args):
    """Run offline inference with AsyncOmni, logging each audio chunk as it arrives."""
    model_name, inputs = _build_inputs(args)
    output_dir = args.output_dir
    os.makedirs(output_dir, exist_ok=True)

    omni = AsyncOmni(
        model=model_name,
        stage_configs_path=args.stage_configs_path,
        log_stats=args.log_stats,
        stage_init_timeout=args.stage_init_timeout,
        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
    )

    for i, prompt in enumerate(inputs):
        request_id = str(i)
        t_start = time.perf_counter()
        t_prev = t_start
        chunk_idx = 0
        async for stage_output in omni.generate(prompt, request_id=request_id):
            mm = stage_output.request_output.outputs[0].multimodal_output
            if not stage_output.finished:
                t_now = time.perf_counter()
                audio = mm.get("audio")
                n = len(audio) if isinstance(audio, list) else (0 if audio is None else 1)
                dt_ms = (t_now - t_prev) * 1000
                ttfa_ms = (t_now - t_start) * 1000
                if chunk_idx == 0:
                    logger.info(f"Request {request_id}: chunk {chunk_idx} samples={n} TTFA={ttfa_ms:.1f}ms")
                else:
                    logger.info(f"Request {request_id}: chunk {chunk_idx} samples={n} inter_chunk={dt_ms:.1f}ms")
                t_prev = t_now
                chunk_idx += 1
            else:
                t_end = time.perf_counter()
                total_ms = (t_end - t_start) * 1000
                logger.info(f"Request {request_id}: done total={total_ms:.1f}ms chunks={chunk_idx}")
                _save_wav(output_dir, request_id, mm)

def parse_args():
    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
    parser.add_argument(
        "--query-type",
        "-q",
        type=str,
        default="CustomVoice",
        choices=query_map.keys(),
        help="Query type.",
    )
    parser.add_argument(
        "--log-stats",
        action="store_true",
        default=False,
        help="Enable writing detailed statistics (default: disabled)",
    )
    parser.add_argument(
        "--stage-init-timeout",
        type=int,
        default=300,
        help="Timeout for initializing a single stage in seconds (default: 300)",
    )
    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(
        "--output-dir",
        default="output_audio",
        help="Output directory for generated wav files (default: output_audio).",
    )
    parser.add_argument(
        "--num-prompts",
        type=int,
        default=1,
        help="Number of prompts to generate.",
    )
    parser.add_argument(
        "--txt-prompts",
        type=str,
        default=None,
        help="Path to a .txt file with one prompt per line (preferred).",
    )
    parser.add_argument(
        "--stage-configs-path",
        type=str,
        default=None,
        help="Path to a stage configs file.",
    )
    parser.add_argument(
        "--audio-path",
        "-a",
        type=str,
        default=None,
        help="Path to local audio file. If not provided, uses default audio asset.",
    )
    parser.add_argument(
        "--sampling-rate",
        type=int,
        default=16000,
        help="Sampling rate for audio loading (default: 16000).",
    )
    parser.add_argument(
        "--log-dir",
        type=str,
        default="logs",
        help="Log directory (default: logs).",
    )
    parser.add_argument(
        "--py-generator",
        action="store_true",
        default=False,
        help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
    )
    parser.add_argument(
        "--use-batch-sample",
        action="store_true",
        default=False,
        help="Use batch input sample for CustomVoice/VoiceDesign/Base query.",
    )
    parser.add_argument(
        "--mode-tag",
        type=str,
        default="icl",
        choices=["icl", "xvec_only"],
        help="Mode tag for Base query x_vector_only_mode (default: icl).",
    )
    parser.add_argument(
        "--streaming",
        action="store_true",
        default=False,
        help="Stream audio chunks as they arrive via AsyncOmni (async_chunk mode only).",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=1,
        help="Number of prompts per batch (default: 1, sequential).",
    )
    parser.add_argument(
        "--enable-diffusion-pipeline-profiler",
        action="store_true",
        help="Enable diffusion pipeline profiler to display stage durations.",
    )

    return parser.parse_args()

if __name__ == "__main__":
    args = parse_args()
    if args.streaming:
        asyncio.run(main_streaming(args))
    else:
        main(args)


Now I have a clear picture of the vLLM-Omni API. Let me create a benchmark script that uses it for Qwen3-TTS-12Hz-0.6B-Base with voice cloning.


Write

"""
Qwen3-TTS-12Hz-0.6B-Base — vLLM-Omni RTF Benchmark
Compares vLLM-Omni inference engine against baseline results.
"""
import os
import sys
import time
import json
import logging

import numpy as np
import soundfile as sf
import torch

os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

OUTPUT_DIR = "/home/ubuntu/quench_0.6b/results"
os.makedirs(OUTPUT_DIR, exist_ok=True)

MODEL_NAME = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
REF_AUDIO = "/home/ubuntu/quench_0.6b/ref_audio.wav"
REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."

SENTENCES = [
    "The quick brown fox jumps over the lazy dog near the riverbank.",
    "Artificial intelligence is transforming how we approach complex problems in science and engineering.",
    "The Massachusetts Institute of Technology announced groundbreaking research in superconductivity.",
    "She couldn't believe it and then whispered this changes everything we know about quantum mechanics.",
    "Beautiful weather today with clear skies and a gentle breeze coming from the northwest.",
    "The orchestra performed a magnificent symphony that left the audience speechless and deeply moved.",
    "Renewable energy sources such as solar and wind power are becoming increasingly important worldwide.",
    "The philosopher argued that consciousness emerges from the complex interactions of neural networks.",
    "International cooperation is essential for addressing global challenges like climate change and poverty.",
    "The restaurant serves traditional Mediterranean cuisine with fresh ingredients sourced from local farmers.",
    "Machine learning algorithms can identify patterns in large datasets that would be impossible for humans to detect manually.",
    "The astronaut described the view of Earth from space as the most humbling experience of her entire career.",
    "Quantum computing promises to revolutionize cryptography and drug discovery within the next decade.",
    "The ancient library contained thousands of manuscripts dating back to the early medieval period.",
    "Global temperatures have risen significantly over the past century due to increased greenhouse gas emissions.",
]

NUM_WARMUP = 2
NUM_RUNS = 3

def estimate_prompt_len(text, ref_audio, ref_text, model_name):
    """Estimate prompt token length for the Talker stage."""
    try:
        from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
        from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
            Qwen3TTSTalkerForConditionalGeneration,
        )
        from transformers import AutoTokenizer

        tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="left")
        cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True)
        tcfg = getattr(cfg, "talker_config", None)

        additional_information = {
            "task_type": ["Base"],
            "ref_audio": [ref_audio],
            "ref_text": [ref_text],
            "text": [text],
            "language": ["English"],
            "x_vector_only_mode": [False],
            "max_new_tokens": [2048],
        }

        return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
            additional_information=additional_information,
            task_type="Base",
            tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
            codec_language_id=getattr(tcfg, "codec_language_id", None),
            spk_is_dialect=getattr(tcfg, "spk_is_dialect", None),
        )
    except Exception as exc:
        logger.warning("Prompt length estimation failed, using fallback 2048: %s", exc)
        return 2048

def build_input(text, ref_audio, ref_text, model_name):
    """Build a single vLLM-Omni input dict for Base (voice clone) task."""
    additional_information = {
        "task_type": ["Base"],
        "ref_audio": [ref_audio],
        "ref_text": [ref_text],
        "text": [text],
        "language": ["English"],
        "x_vector_only_mode": [False],
        "max_new_tokens": [2048],
    }
    prompt_len = estimate_prompt_len(text, ref_audio, ref_text, model_name)
    return {
        "prompt_token_ids": [0] * prompt_len,
        "additional_information": additional_information,
    }

def extract_audio(output):
    """Extract waveform and sample rate from OmniRequestOutput."""
    mm = output.request_output.outputs[0].multimodal_output
    audio_data = mm["audio"]
    sr_raw = mm["sr"]
    sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
    sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
    audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
    wav = audio_tensor.float().cpu().numpy().flatten()
    return wav, sr

# ═══════════════════════════════════════════════
print("=" * 65)
print("  vLLM-Omni RTF Benchmark — Qwen3-TTS-12Hz-0.6B-Base")
print("  GPU: NVIDIA A100-SXM4-80GB | dtype: bfloat16")
print("=" * 65)

from vllm_omni import Omni

# Build all inputs first
print("\nBuilding inputs...")
inputs = [build_input(text, REF_AUDIO, REF_TEXT, MODEL_NAME) for text in SENTENCES]

# ─────────────────────────────────────────────
# Sequential benchmark (batch_size=1)
# ─────────────────────────────────────────────
print("\n[ENGINE] vLLM-Omni (sequential, batch_size=1)")
print("  Loading model...")

omni = Omni(
    model=MODEL_NAME,
    stage_init_timeout=600,
)

# Warmup
print(f"  Warmup ({NUM_WARMUP} sentences)...")
for i in range(NUM_WARMUP):
    for output in omni.generate([inputs[i % len(inputs)]]):
        _ = extract_audio(output)

# Benchmark each sentence
print(f"  Running {len(SENTENCES)} sentences x {NUM_RUNS} runs...")
results = []
for si, text in enumerate(SENTENCES):
    times = []
    durations = []
    for r in range(NUM_RUNS):
        t0 = time.time()
        for output in omni.generate([inputs[si]]):
            wav, sr = extract_audio(output)
        elapsed = time.time() - t0
        audio_dur = len(wav) / sr
        times.append(elapsed)
        durations.append(audio_dur)

    avg_time = np.mean(times)
    avg_dur = np.mean(durations)
    rtf = avg_time / avg_dur
    results.append({
        "text": text,
        "avg_gen_s": round(avg_time, 3),
        "avg_audio_s": round(avg_dur, 2),
        "rtf": round(rtf, 4),
        "tokens_per_s": round(len(text) / avg_time, 1),
    })
    print(f"    [{si+1:2d}/{len(SENTENCES)}] RTF={rtf:.4f}  gen={avg_time:.2f}s  audio={avg_dur:.1f}s  | {text[:50]}...")

rtfs = [r["rtf"] for r in results]
total_gen = sum(r["avg_gen_s"] for r in results)
total_audio = sum(r["avg_audio_s"] for r in results)
seq_results = {
    "engine": "vllm_omni_sequential",
    "num_sentences": len(SENTENCES),
    "num_runs": NUM_RUNS,
    "avg_rtf": round(np.mean(rtfs), 4),
    "median_rtf": round(np.median(rtfs), 4),
    "min_rtf": round(np.min(rtfs), 4),
    "max_rtf": round(np.max(rtfs), 4),
    "std_rtf": round(np.std(rtfs), 4),
    "total_gen_s": round(total_gen, 2),
    "total_audio_s": round(total_audio, 2),
    "overall_rtf": round(total_gen / total_audio, 4),
    "per_sentence": results,
}

# ─────────────────────────────────────────────
# Streaming TTFB benchmark
# ─────────────────────────────────────────────
import asyncio
from vllm_omni import AsyncOmni

print("\n[ENGINE] vLLM-Omni (streaming TTFB)")

async def measure_ttfb(omni_async, inp, request_id):
    """Measure time to first audio chunk."""
    t_start = time.perf_counter()
    ttfb = None
    total_samples = 0
    async for stage_output in omni_async.generate(inp, request_id=request_id):
        mm = stage_output.request_output.outputs[0].multimodal_output
        if not stage_output.finished:
            if ttfb is None:
                ttfb = (time.perf_counter() - t_start) * 1000  # ms
            audio = mm.get("audio")
            if isinstance(audio, list):
                total_samples += sum(a.numel() for a in audio)
            elif audio is not None:
                total_samples += audio.numel()
        else:
            total_time = (time.perf_counter() - t_start) * 1000
            sr_raw = mm["sr"]
            sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
            sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
            audio_data = mm["audio"]
            audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
            wav = audio_tensor.float().cpu().numpy().flatten()
            audio_dur = len(wav) / sr
            return {
                "ttfb_ms": round(ttfb, 1) if ttfb else None,
                "total_ms": round(total_time, 1),
                "audio_s": round(audio_dur, 2),
                "rtf": round(total_time / 1000 / audio_dur, 4) if audio_dur > 0 else None,
            }

async def run_ttfb_bench():
    omni_async = AsyncOmni(
        model=MODEL_NAME,
        stage_init_timeout=600,
    )

    # Warmup
    print(f"  Warmup ({NUM_WARMUP} requests)...")
    for i in range(NUM_WARMUP):
        async for _ in omni_async.generate(inputs[i % len(inputs)], request_id=f"warmup_{i}"):
            pass

    # Measure TTFB
    print(f"  Measuring TTFB for {min(10, len(SENTENCES))} sentences...")
    ttfb_results = []
    for si in range(min(10, len(SENTENCES))):
        result = await measure_ttfb(omni_async, inputs[si], request_id=f"ttfb_{si}")
        if result:
            ttfb_results.append(result)
            print(f"    [{si+1:2d}] TTFB={result['ttfb_ms']}ms  total={result['total_ms']}ms  audio={result['audio_s']}s")

    return ttfb_results

ttfb_results = asyncio.run(run_ttfb_bench())

ttfb_values = [r["ttfb_ms"] for r in ttfb_results if r["ttfb_ms"]]
streaming_summary = {
    "engine": "vllm_omni_streaming",
    "num_sentences": len(ttfb_results),
    "ttfb_p50_ms": round(np.percentile(ttfb_values, 50), 1) if ttfb_values else None,
    "ttfb_p95_ms": round(np.percentile(ttfb_values, 95), 1) if ttfb_values else None,
    "ttfb_p99_ms": round(np.percentile(ttfb_values, 99), 1) if ttfb_values else None,
    "ttfb_mean_ms": round(np.mean(ttfb_values), 1) if ttfb_values else None,
    "per_request": ttfb_results,
}

# ─────────────────────────────────────────────
# Save and compare
# ─────────────────────────────────────────────
all_vllm = {
    "vllm_omni_sequential": seq_results,
    "vllm_omni_streaming": streaming_summary,
}

results_path = os.path.join(OUTPUT_DIR, "vllm_omni_benchmark.json")
with open(results_path, "w") as f:
    json.dump(all_vllm, f, indent=2)

# Load previous results for comparison
prev_path = os.path.join(OUTPUT_DIR, "rtf_benchmark.json")
prev = {}
if os.path.exists(prev_path):
    with open(prev_path) as f:
        prev = json.load(f)

print("\n" + "=" * 65)
print("  FULL RTF COMPARISON — Qwen3-TTS-12Hz-0.6B-Base")
print("  GPU: NVIDIA A100-SXM4-80GB | English Only")
print("=" * 65)
print()
print(f"  {'Engine':35s}  {'RTF':>8s}  {'Gen/sample':>10s}")
print(f"  {'-'*35}  {'-'*8}  {'-'*10}")

if "qwen_tts_baseline" in prev:
    r1 = prev["qwen_tts_baseline"]
    print(f"  {'qwen_tts (sequential)':35s}  {r1['overall_rtf']:8.4f}  {r1['total_gen_s']/r1['num_sentences']:9.2f}s")

if "qwen_tts_batch" in prev:
    r1b = prev["qwen_tts_batch"]
    print(f"  {'qwen_tts (batch x' + str(r1b.get('batch_size', '?')) + ')':35s}  {r1b['overall_rtf']:8.4f}  {r1b['per_sample_gen_s']:9.3f}s")

if "faster_qwen3_tts" in prev:
    r2 = prev["faster_qwen3_tts"]
    print(f"  {'faster-qwen3-tts (CUDA graphs)':35s}  {r2['overall_rtf']:8.4f}  {r2['total_gen_s']/r2['num_sentences']:9.2f}s")

print(f"  {'vLLM-Omni (sequential)':35s}  {seq_results['overall_rtf']:8.4f}  {seq_results['total_gen_s']/seq_results['num_sentences']:9.2f}s")

if ttfb_values:
    print()
    print(f"  vLLM-Omni Streaming TTFB:")
    print(f"    p50: {streaming_summary['ttfb_p50_ms']}ms  p95: {streaming_summary['ttfb_p95_ms']}ms  p99: {streaming_summary['ttfb_p99_ms']}ms")

print()

# Find best engine
all_rtfs = []
if "qwen_tts_baseline" in prev:
    all_rtfs.append(("qwen_tts_sequential", prev["qwen_tts_baseline"]["overall_rtf"]))
if "qwen_tts_batch" in prev:
    all_rtfs.append(("qwen_tts_batch", prev["qwen_tts_batch"]["overall_rtf"]))
if "faster_qwen3_tts" in prev:
    all_rtfs.append(("faster_qwen3_tts", prev["faster_qwen3_tts"]["overall_rtf"]))
all_rtfs.append(("vllm_omni_sequential", seq_results["overall_rtf"]))

best = min(all_rtfs, key=lambda x: x[1])
print(f"  BEST ENGINE: {best[0]} with RTF = {best[1]:.4f}")
if "qwen_tts_baseline" in prev:
    speedup = prev["qwen_tts_baseline"]["overall_rtf"] / best[1]
    print(f"  Speedup over sequential baseline: {speedup:.2f}x")

print()
print(f"  RTF < 1.0 = faster than real-time")
print(f"  Results saved to: {results_path}")
print("=" * 65)


File created successfully at: /home/ubuntu/quench_0.6b/bench_vllm_omni.py

"""
Qwen3-TTS-12Hz-0.6B-Base — vLLM-Omni RTF Benchmark
Compares vLLM-Omni inference engine against baseline results.
"""
import os
import sys
import time
import json
import logging

import numpy as np
import soundfile as sf
import torch

os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

OUTPUT_DIR = "/home/ubuntu/quench_0.6b/results"
os.makedirs(OUTPUT_DIR, exist_ok=True)

MODEL_NAME = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
REF_AUDIO = "/home/ubuntu/quench_0.6b/ref_audio.wav"
REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."

SENTENCES = [
    "The quick brown fox jumps over the lazy dog near the riverbank.",
    "Artificial intelligence is transforming how we approach complex problems in science and engineering.",
    "The Massachusetts Institute of Technology announced groundbreaking research in superconductivity.",
    "She couldn't believe it and then whispered this changes everything we know about quantum mechanics.",
    "Beautiful weather today with clear skies and a gentle breeze coming from the northwest.",
    "The orchestra performed a magnificent symphony that left the audience speechless and deeply moved.",
    "Renewable energy sources such as solar and wind power are becoming increasingly important worldwide.",
    "The philosopher argued that consciousness emerges from the complex interactions of neural networks.",
    "International cooperation is essential for addressing global challenges like climate change and poverty.",
    "The restaurant serves traditional Mediterranean cuisine with fresh ingredients sourced from local farmers.",
    "Machine learning algorithms can identify patterns in large datasets that would be impossible for humans to detect manually.",
    "The astronaut described the view of Earth from space as the most humbling experience of her entire career.",
    "Quantum computing promises to revolutionize cryptography and drug discovery within the next decade.",
    "The ancient library contained thousands of manuscripts dating back to the early medieval period.",
    "Global temperatures have risen significantly over the past century due to increased greenhouse gas emissions.",
]

NUM_WARMUP = 2
NUM_RUNS = 3

def estimate_prompt_len(text, ref_audio, ref_text, model_name):
    """Estimate prompt token length for the Talker stage."""
    try:
        from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
        from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
            Qwen3TTSTalkerForConditionalGeneration,
        )
        from transformers import AutoTokenizer

        tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="left")
        cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True)
        tcfg = getattr(cfg, "talker_config", None)

        additional_information = {
            "task_type": ["Base"],
            "ref_audio": [ref_audio],
            "ref_text": [ref_text],
            "text": [text],
            "language": ["English"],
            "x_vector_only_mode": [False],
            "max_new_tokens": [2048],
        }

        return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
            additional_information=additional_information,
            task_type="Base",
            tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
            codec_language_id=getattr(tcfg, "codec_language_id", None),
            spk_is_dialect=getattr(tcfg, "spk_is_dialect", None),
        )
    except Exception as exc:
        logger.warning("Prompt length estimation failed, using fallback 2048: %s", exc)
        return 2048

def build_input(text, ref_audio, ref_text, model_name):
    """Build a single vLLM-Omni input dict for Base (voice clone) task."""
    additional_information = {
        "task_type": ["Base"],
        "ref_audio": [ref_audio],
        "ref_text": [ref_text],
        "text": [text],
        "language": ["English"],
        "x_vector_only_mode": [False],
        "max_new_tokens": [2048],
    }
    prompt_len = estimate_prompt_len(text, ref_audio, ref_text, model_name)
    return {
        "prompt_token_ids": [0] * prompt_len,
        "additional_information": additional_information,
    }

def extract_audio(output):
    """Extract waveform and sample rate from OmniRequestOutput."""
    mm = output.request_output.outputs[0].multimodal_output
    audio_data = mm["audio"]
    sr_raw = mm["sr"]
    sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
    sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
    audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
    wav = audio_tensor.float().cpu().numpy().flatten()
    return wav, sr

# ═══════════════════════════════════════════════
print("=" * 65)
print("  vLLM-Omni RTF Benchmark — Qwen3-TTS-12Hz-0.6B-Base")
print("  GPU: NVIDIA A100-SXM4-80GB | dtype: bfloat16")
print("=" * 65)

from vllm_omni import Omni

# Build all inputs first
print("\nBuilding inputs...")
inputs = [build_input(text, REF_AUDIO, REF_TEXT, MODEL_NAME) for text in SENTENCES]

# ─────────────────────────────────────────────
# Sequential benchmark (batch_size=1)
# ─────────────────────────────────────────────
print("\n[ENGINE] vLLM-Omni (sequential, batch_size=1)")
print("  Loading model...")

omni = Omni(
    model=MODEL_NAME,
    stage_init_timeout=600,
)

# Warmup
print(f"  Warmup ({NUM_WARMUP} sentences)...")
for i in range(NUM_WARMUP):
    for output in omni.generate([inputs[i % len(inputs)]]):
        _ = extract_audio(output)

# Benchmark each sentence
print(f"  Running {len(SENTENCES)} sentences x {NUM_RUNS} runs...")
results = []
for si, text in enumerate(SENTENCES):
    times = []
    durations = []
    for r in range(NUM_RUNS):
        t0 = time.time()
        for output in omni.generate([inputs[si]]):
            wav, sr = extract_audio(output)
        elapsed = time.time() - t0
        audio_dur = len(wav) / sr
        times.append(elapsed)
        durations.append(audio_dur)

    avg_time = np.mean(times)
    avg_dur = np.mean(durations)
    rtf = avg_time / avg_dur
    results.append({
        "text": text,
        "avg_gen_s": round(avg_time, 3),
        "avg_audio_s": round(avg_dur, 2),
        "rtf": round(rtf, 4),
        "tokens_per_s": round(len(text) / avg_time, 1),
    })
    print(f"    [{si+1:2d}/{len(SENTENCES)}] RTF={rtf:.4f}  gen={avg_time:.2f}s  audio={avg_dur:.1f}s  | {text[:50]}...")

rtfs = [r["rtf"] for r in results]
total_gen = sum(r["avg_gen_s"] for r in results)
total_audio = sum(r["avg_audio_s"] for r in results)
seq_results = {
    "engine": "vllm_omni_sequential",
    "num_sentences": len(SENTENCES),
    "num_runs": NUM_RUNS,
    "avg_rtf": round(np.mean(rtfs), 4),
    "median_rtf": round(np.median(rtfs), 4),
    "min_rtf": round(np.min(rtfs), 4),
    "max_rtf": round(np.max(rtfs), 4),
    "std_rtf": round(np.std(rtfs), 4),
    "total_gen_s": round(total_gen, 2),
    "total_audio_s": round(total_audio, 2),
    "overall_rtf": round(total_gen / total_audio, 4),
    "per_sentence": results,
}

# ─────────────────────────────────────────────
# Streaming TTFB benchmark
# ─────────────────────────────────────────────
import asyncio
from vllm_omni import AsyncOmni

print("\n[ENGINE] vLLM-Omni (streaming TTFB)")

async def measure_ttfb(omni_async, inp, request_id):
    """Measure time to first audio chunk."""
    t_start = time.perf_counter()
    ttfb = None
    total_samples = 0
    async for stage_output in omni_async.generate(inp, request_id=request_id):
        mm = stage_output.request_output.outputs[0].multimodal_output
        if not stage_output.finished:
            if ttfb is None:
                ttfb = (time.perf_counter() - t_start) * 1000  # ms
            audio = mm.get("audio")
            if isinstance(audio, list):
                total_samples += sum(a.numel() for a in audio)
            elif audio is not None:
                total_samples += audio.numel()
        else:
            total_time = (time.perf_counter() - t_start) * 1000
            sr_raw = mm["sr"]
            sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
            sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
            audio_data = mm["audio"]
            audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
            wav = audio_tensor.float().cpu().numpy().flatten()
            audio_dur = len(wav) / sr
            return {
                "ttfb_ms": round(ttfb, 1) if ttfb else None,
                "total_ms": round(total_time, 1),
                "audio_s": round(audio_dur, 2),
                "rtf": round(total_time / 1000 / audio_dur, 4) if audio_dur > 0 else None,
            }

async def run_ttfb_bench():
    omni_async = AsyncOmni(
        model=MODEL_NAME,
        stage_init_timeout=600,
    )

    # Warmup
    print(f"  Warmup ({NUM_WARMUP} requests)...")
    for i in range(NUM_WARMUP):
        async for _ in omni_async.generate(inputs[i % len(inputs)], request_id=f"warmup_{i}"):
            pass

    # Measure TTFB
    print(f"  Measuring TTFB for {min(10, len(SENTENCES))} sentences...")
    ttfb_results = []
    for si in range(min(10, len(SENTENCES))):
        result = await measure_ttfb(omni_async, inputs[si], request_id=f"ttfb_{si}")
        if result:
            ttfb_results.append(result)
            print(f"    [{si+1:2d}] TTFB={result['ttfb_ms']}ms  total={result['total_ms']}ms  audio={result['audio_s']}s")

    return ttfb_results

ttfb_results = asyncio.run(run_ttfb_bench())

ttfb_values = [r["ttfb_ms"] for r in ttfb_results if r["ttfb_ms"]]
streaming_summary = {
    "engine": "vllm_omni_streaming",
    "num_sentences": len(ttfb_results),
    "ttfb_p50_ms": round(np.percentile(ttfb_values, 50), 1) if ttfb_values else None,
    "ttfb_p95_ms": round(np.percentile(ttfb_values, 95), 1) if ttfb_values else None,
    "ttfb_p99_ms": round(np.percentile(ttfb_values, 99), 1) if ttfb_values else None,
    "ttfb_mean_ms": round(np.mean(ttfb_values), 1) if ttfb_values else None,
    "per_request": ttfb_results,
}

# ─────────────────────────────────────────────
# Save and compare
# ─────────────────────────────────────────────
all_vllm = {
    "vllm_omni_sequential": seq_results,
    "vllm_omni_streaming": streaming_summary,
}

results_path = os.path.join(OUTPUT_DIR, "vllm_omni_benchmark.json")
with open(results_path, "w") as f:
    json.dump(all_vllm, f, indent=2)

# Load previous results for comparison
prev_path = os.path.join(OUTPUT_DIR, "rtf_benchmark.json")
prev = {}
if os.path.exists(prev_path):
    with open(prev_path) as f:
        prev = json.load(f)

print("\n" + "=" * 65)
print("  FULL RTF COMPARISON — Qwen3-TTS-12Hz-0.6B-Base")
print("  GPU: NVIDIA A100-SXM4-80GB | English Only")
print("=" * 65)
print()
print(f"  {'Engine':35s}  {'RTF':>8s}  {'Gen/sample':>10s}")
print(f"  {'-'*35}  {'-'*8}  {'-'*10}")

if "qwen_tts_baseline" in prev:
    r1 = prev["qwen_tts_baseline"]
    print(f"  {'qwen_tts (sequential)':35s}  {r1['overall_rtf']:8.4f}  {r1['total_gen_s']/r1['num_sentences']:9.2f}s")

if "qwen_tts_batch" in prev:
    r1b = prev["qwen_tts_batch"]
    print(f"  {'qwen_tts (batch x' + str(r1b.get('batch_size', '?')) + ')':35s}  {r1b['overall_rtf']:8.4f}  {r1b['per_sample_gen_s']:9.3f}s")

if "faster_qwen3_tts" in prev:
    r2 = prev["faster_qwen3_tts"]
    print(f"  {'faster-qwen3-tts (CUDA graphs)':35s}  {r2['overall_rtf']:8.4f}  {r2['total_gen_s']/r2['num_sentences']:9.2f}s")

print(f"  {'vLLM-Omni (sequential)':35s}  {seq_results['overall_rtf']:8.4f}  {seq_results['total_gen_s']/seq_results['num_sentences']:9.2f}s")

if ttfb_values:
    print()
    print(f"  vLLM-Omni Streaming TTFB:")
    print(f"    p50: {streaming_summary['ttfb_p50_ms']}ms  p95: {streaming_summary['ttfb_p95_ms']}ms  p99: {streaming_summary['ttfb_p99_ms']}ms")

print()

# Find best engine
all_rtfs = []
if "qwen_tts_baseline" in prev:
    all_rtfs.append(("qwen_tts_sequential", prev["qwen_tts_baseline"]["overall_rtf"]))
if "qwen_tts_batch" in prev:
    all_rtfs.append(("qwen_tts_batch", prev["qwen_tts_batch"]["overall_rtf"]))
if "faster_qwen3_tts" in prev:
    all_rtfs.append(("faster_qwen3_tts", prev["faster_qwen3_tts"]["overall_rtf"]))
all_rtfs.append(("vllm_omni_sequential", seq_results["overall_rtf"]))

best = min(all_rtfs, key=lambda x: x[1])
print(f"  BEST ENGINE: {best[0]} with RTF = {best[1]:.4f}")
if "qwen_tts_baseline" in prev:
    speedup = prev["qwen_tts_baseline"]["overall_rtf"] / best[1]
    print(f"  Speedup over sequential baseline: {speedup:.2f}x")

print()
print(f"  RTF < 1.0 = faster than real-time")
print(f"  Results saved to: {results_path}")
print("=" * 65)


Bash

cd /home/ubuntu/quench_0.6b && python3 bench_vllm_omni.py 2>&1

Run vLLM-Omni RTF benchmark for Qwen3-TTS


Output:
  vLLM-Omni RTF Benchmark — Qwen3-TTS-12Hz-0.6B-Base
  GPU: NVIDIA A100-SXM4-80GB | dtype: bfloat16
=================================================================
WARNING 03-21 17:56:29 [envs.py:94] No Flash Attention backend found, using pytorch SDPA implementation


Output:
INFO 03-21 17:56:31 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:31 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:31 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:31,416 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.

=================================================================
  vLLM-Omni RTF Benchmark — Qwen3-TTS-12Hz-0.6B-Base
  GPU: NVIDIA A100-SXM4-80GB | dtype: bfloat16
=================================================================
WARNING 03-21 17:56:29 [envs.py:94] No Flash Attention backend found, using pytorch SDPA implementation
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.

Building inputs...
INFO 03-21 17:56:30 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:30 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:30 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:30 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:30,804 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:31 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:31 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:31 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:31 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:31,416 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.


Output:
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:32,013 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.

=================================================================
  vLLM-Omni RTF Benchmark — Qwen3-TTS-12Hz-0.6B-Base
  GPU: NVIDIA A100-SXM4-80GB | dtype: bfloat16
=================================================================
WARNING 03-21 17:56:29 [envs.py:94] No Flash Attention backend found, using pytorch SDPA implementation
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.

Building inputs...
INFO 03-21 17:56:30 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:30 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:30 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:30 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:30,804 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:31 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:31 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:31 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:31 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:31,416 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:32,013 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.


Output:
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:33,194 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.

O 03-21 17:56:30 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:30 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:30 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:30,804 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:31 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:31 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:31 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:31 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:31,416 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:32,013 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:32,617 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:33,194 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.


Output:
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:34,389 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.

O 03-21 17:56:32 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:32,013 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:32,617 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:33,194 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:33,793 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:34,389 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.


Output:
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:34,980 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.

O 03-21 17:56:32 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:32 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:32,617 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:33,194 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:33,793 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:34,389 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:34,980 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.


Output:
INFO 03-21 17:56:36 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:36 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:36 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:36,149 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.

O 03-21 17:56:33 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:33 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:33,793 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:34,389 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:34 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:34,980 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:35 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:35 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:35 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:35 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:35,564 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.
The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
INFO 03-21 17:56:36 [configuration_qwen3_tts.py:489] talker_config is None. Initializing talker model with default values
INFO 03-21 17:56:36 [configuration_qwen3_tts.py:492] speaker_encoder_config is None. Initializing talker model with default values
INFO 03-21 17:56:36 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
INFO 03-21 17:56:36 [configuration_qwen3_tts.py:441] code_predictor_config is None. Initializing code_predictor model with default values
2026-03-21 17:56:36,149 WARNING Prompt length estimation failed, using fallback 2048: Base in-context voice cloning requires either `voice_clone_prompt.ref_code` or a readable `ref_audio` that can be mapped to a codec frame length.