vllm-omni/examples/offline_inference/text_to_image/text_to_image.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import argparse
import functools
import json
import time
from pathlib import Path
from typing import Any

import numpy as np
import torch
from diffusers.utils import numpy_to_pil

from vllm_omni.diffusion.data import logger
from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.lora.request import LoRARequest
from vllm_omni.lora.utils import stable_lora_int_id
from vllm_omni.model_extras import (
    build_text_to_image_prompt as build_model_text_to_image_prompt,
)
from vllm_omni.model_extras import (
    get_ar_input_builder,
    get_ar_tokenizer_validator,
    get_extra_body_params,
    get_model_class_name,
    should_init_extra_args_for_non_diffusion_stages,
)
from vllm_omni.platforms import current_omni_platform


def is_nextstep_model(model_name: str) -> bool:
    """Check if the model is a NextStep model by reading its config."""
    from vllm.transformers_utils.config import get_hf_file_to_dict

    try:
        cfg = get_hf_file_to_dict("config.json", model_name)
        if cfg and cfg.get("model_type") == "nextstep":
            return True
    except Exception:
        pass
    return False


def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]:
    """Parse a CLI value as a JSON object, attributing errors to ``flag_name``."""
    try:
        config = json.loads(value)
    except json.JSONDecodeError as e:
        raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e
    if not isinstance(config, dict):
        raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object")
    return config


parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")


def build_text_to_image_prompt(prompt: str, negative_prompt: str | None) -> dict[str, Any]:
    """Build the canonical request envelope for the shared T2I example."""
    result: dict[str, Any] = {
        "prompt": prompt,
        "modalities": ["image"],
    }
    if negative_prompt is not None:
        result["negative_prompt"] = negative_prompt
    return result


def _normalize_images_for_save(images: list[Any]) -> list[Any]:
    """Convert NumPy diffusion outputs to PIL images before saving."""
    normalized = []
    for image in images:
        if isinstance(image, np.ndarray):
            normalized.extend(numpy_to_pil(image))
        else:
            normalized.append(image)
    return normalized


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.")
    parser.add_argument(
        "--model",
        default="Qwen/Qwen-Image",
        help="Diffusion model name or local path. Supported models: "
        "Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, stepfun-ai/NextStep-1.1, "
        "black-forest-labs/FLUX.1-dev, black-forest-labs/FLUX.2-klein-9B, "
        "black-forest-labs/FLUX.2-dev, tencent/HunyuanImage-3.0-Instruct, "
        "HiDream-ai/HiDream-O1-Image, meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, "
        "stabilityai/stable-diffusion-3.5-medium, Tongyi-MAI/Z-Image-Turbo and etc.",
    )
    parser.add_argument(
        "--deploy-config",
        type=str,
        default=None,
        help=(
            "Path to a deploy YAML (new pipeline/stages format). Required for multi-stage "
            "text-to-image pipelines whose deploy config is not auto-loaded."
        ),
    )
    parser.add_argument("--prompt", default="a cup of coffee on the table", help="Text prompt for image generation.")
    parser.add_argument(
        "--negative-prompt",
        default=None,
        help="negative prompt for classifier-free conditional guidance.",
    )
    parser.add_argument("--seed", type=int, default=142, help="Random seed for deterministic results.")
    parser.add_argument(
        "--cfg-scale",
        type=float,
        default=4.0,
        help="True classifier-free guidance scale specific to Qwen-Image.",
    )
    parser.add_argument(
        "--guidance-scale",
        type=float,
        default=4.0,
        help="Classifier-free guidance scale. HunyuanImage3 recommends 4.0-5.0.",
    )
    parser.add_argument("--height", type=int, default=1024, help="Height of generated image.")
    parser.add_argument("--width", type=int, default=1024, help="Width of generated image.")
    parser.add_argument(
        "--output",
        type=str,
        default="qwen_image_output.png",
        help="Path to save the generated image (PNG).",
    )
    parser.add_argument(
        "--num-images-per-prompt",
        type=int,
        default=1,
        help="Number of images to generate for the given prompt.",
    )
    parser.add_argument(
        "--num-inference-steps",
        type=int,
        default=50,
        help="Number of denoising steps for the diffusion sampler.",
    )
    parser.add_argument(
        "--cache-backend",
        type=str,
        default=None,
        choices=["cache_dit", "tea_cache"],
        help=(
            "Cache backend to use for acceleration. "
            "Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). "
            "Default: None (no cache acceleration)."
        ),
    )
    parser.add_argument(
        "--enable-cache-dit-summary",
        action="store_true",
        help="Enable cache-dit summary logging after diffusion forward passes.",
    )
    parser.add_argument(
        "--ulysses-degree",
        type=int,
        default=1,
        help="Number of GPUs used for ulysses sequence parallelism.",
    )
    parser.add_argument(
        "--ulysses-mode",
        type=str,
        default="strict",
        choices=["strict", "advanced_uaa"],
        help="Ulysses sequence-parallel mode: 'strict' (divisibility required) or 'advanced_uaa' (UAA).",
    )
    parser.add_argument(
        "--ring-degree",
        type=int,
        default=1,
        help="Number of GPUs used for ring sequence parallelism.",
    )
    parser.add_argument(
        "--cfg-parallel-size",
        type=int,
        default=1,
        choices=[1, 2],
        help="Number of GPUs used for classifier free guidance parallel size.",
    )
    parser.add_argument(
        "--enforce-eager",
        action="store_true",
        default=None,
        help=(
            "Disable torch.compile and force eager execution. Left unset (None) "
            "so it is only forwarded when explicitly given; "
            "otherwise the per-stage deploy YAML value wins."
        ),
    )
    parser.add_argument(
        "--enable-cpu-offload",
        action="store_true",
        help="Enable CPU offloading for diffusion models.",
    )
    parser.add_argument(
        "--enable-layerwise-offload",
        action="store_true",
        help="Enable layerwise (blockwise) offloading on DiT modules.",
    )
    parser.add_argument(
        "--use-hsdp",
        action="store_true",
        help="Enable HSDP (Hybrid Sharded Data Parallel) for diffusion models.",
    )
    parser.add_argument(
        "--hsdp-shard-size",
        type=int,
        default=1,
        help="Number of GPUs to shard weights across for HSDP.",
    )
    parser.add_argument(
        "--hsdp-replicate-size",
        type=int,
        default=1,
        help="Number of HSDP replica groups.",
    )
    parser.add_argument(
        "--quantization",
        type=str,
        default=None,
        choices=["fp8", "int8", "bitsandbytes"],
        help="Quantization method for the transformer. "
        "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), "
        "'int8' (Int8 W8A8), 'bitsandbytes' (NF4 4-bit weight-only, CUDA SM 75+). "
        "Default: None (no quantization, uses BF16).",
    )
    parser.add_argument(
        "--ignored-layers",
        type=str,
        default=None,
        help="Comma-separated list of layer name patterns to skip quantization. "
        "Only used when --quantization is set. "
        "Available layers: to_qkv, to_out, add_kv_proj, to_add_out, img_mlp, txt_mlp, proj_out. "
        "Example: --ignored-layers 'add_kv_proj,to_add_out'",
    )
    parser.add_argument(
        "--vae-use-slicing",
        action="store_true",
        help="Enable VAE slicing for memory optimization.",
    )
    parser.add_argument(
        "--vae-use-tiling",
        action="store_true",
        help="Enable VAE tiling for memory optimization.",
    )
    parser.add_argument(
        "--tensor-parallel-size",
        type=int,
        default=None,
        help=(
            "Number of GPUs used for tensor parallelism (TP) inside the DiT. "
            "Left unset so it is only forwarded when explicitly given; "
            "otherwise the per-stage deploy YAML (or engine default of 1) wins. "
            "Passing it always overrides the deploy YAML."
        ),
    )
    parser.add_argument(
        "--enable-expert-parallel",
        action="store_true",
        help="Enable expert parallelism for MoE layers.",
    )
    parser.add_argument(
        "--lora-path",
        type=str,
        nargs="+",
        default=None,
        help="Path to LoRA adapter folder (PEFT format). Loaded at initialization and used for generation.",
    )
    parser.add_argument(
        "--lora-scale",
        type=float,
        default=1.0,
        help="Scale factor for LoRA weights (default: 1.0).",
    )
    parser.add_argument(
        "--lora-backend",
        type=str,
        default="peft",
        choices=["peft", "distill"],
        help="LoRA backend for loading LoRA adapters. Default: peft"
        "'peft' loads a PEFT-format adapter folder, used e.g. for RL"
        "'distill' fuses one or more concrete LoRA checkpoint files, used e.g. for distilled few-step LoRAs",
    )
    parser.add_argument(
        "--vae-patch-parallel-size",
        type=int,
        default=1,
        help="Number of ranks used for VAE patch/tile parallelism (decode/encode).",
    )
    # NextStep-1.1 specific arguments
    parser.add_argument(
        "--guidance-scale-2",
        type=float,
        default=1.0,
        help="Secondary guidance scale (e.g. image-level CFG for NextStep-1.1).",
    )
    parser.add_argument(
        "--timesteps-shift",
        type=float,
        default=1.0,
        help="[NextStep-1.1 only] Timesteps shift parameter for sampling.",
    )
    parser.add_argument(
        "--cfg-schedule",
        type=str,
        default="constant",
        choices=["constant", "linear"],
        help="[NextStep-1.1 only] CFG schedule type.",
    )
    parser.add_argument(
        "--use-norm",
        action="store_true",
        help="[NextStep-1.1 only] Apply layer normalization to sampled tokens.",
    )
    parser.add_argument(
        "--extra-body",
        type=functools.partial(parse_json_object, flag_name="--extra-body"),
        default=None,
        help=(
            "Model-specific generation params as a JSON object, e.g. "
            '\'{"timestep_shift": 3.0, "cfg_text_scale": 4.0, "cfg_interval": [0.4, 1.0]}\'. '
            "Each key is filtered against the model's declared extra_body_params "
            "(see vllm_omni/model_extras), so unknown keys for the chosen model are "
            "silently dropped. Values here take precedence over the equivalent "
            "model-specific flags above."
        ),
    )
    parser.add_argument(
        "--enable-diffusion-pipeline-profiler",
        action="store_true",
        help="Enable diffusion pipeline profiler to display stage durations.",
    )
    parser.add_argument(
        "--profiler-config",
        type=parse_profiler_config,
        default=None,
        help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
    )
    parser.add_argument(
        "--log-stats",
        action="store_true",
        help="Enable logging of diffusion pipeline stats.",
    )
    parser.add_argument(
        "--init-timeout",
        type=int,
        default=600,
        help="Timeout for initializing a single stage in seconds (default: 600s)",
    )
    parser.add_argument(
        "--stage-init-timeout",
        type=int,
        default=600,
        help="Timeout for initializing a single stage in seconds (default: 600s)",
    )
    parser.add_argument(
        "--use-system-prompt",
        type=str,
        default=None,
        choices=["None", "dynamic", "en_vanilla", "en_recaption", "en_think_recaption", "en_unified", "custom"],
        help="System prompt preset for generation. Recommended: en_unified.",
    )
    parser.add_argument(
        "--system-prompt",
        type=str,
        default=None,
        help=("Custom system prompt. Used when --use-system-prompt is custom. "),
    )
    parser.add_argument(
        "--auxiliary-text-encoder",
        type=str,
        default=None,
        help="Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).",
    )
    parser.add_argument(
        "--trust-remote-code",
        action="store_true",
        help="Trust and execute custom modeling code from the model repo (required by e.g. HunyuanImage-3.0).",
    )
    parser.add_argument(
        "--allow-tokenizer-fallback",
        action="store_true",
        help=(
            "For models with a declared ar_input_builder (e.g. HunyuanImage-3.0): if loading "
            "the AR tokenizer fails, degrade to the string-prompt form (no BPE parity) instead "
            "of failing the run. Off by default -- a tokenizer load failure usually means a "
            "missing --trust-remote-code or a network/cache issue that's worth surfacing, not "
            "silently masking with a possibly-wrong prompt."
        ),
    )
    current_omni_platform.pre_register_and_update(parser)
    return parser.parse_args()


def _apply_ar_stage_inputs(
    ar_input_builder: Any,
    *,
    model: str,
    prompt_text: str,
    extra_body: dict[str, Any],
    num_images: int,
    height: int | None,
    width: int | None,
    prompt_dict: dict[str, Any],
    sampling_params_list: list[Any],
    text_output: bool = False,
    trust_remote_code: bool = False,
    validate_tokenizer: Any | None = None,
    allow_tokenizer_fallback: bool = False,
) -> None:
    """Apply a model's declared AR-stage inputs to the request in place.

    Loads the model tokenizer (for byte-for-byte HF-parity segment
    tokenization) and asks the model's ``ar_input_builder`` for the AR
    prefill + stop tokens, then writes them onto ``prompt_dict`` and the
    non-diffusion (AR) stage sampling params. Kept model-agnostic: the
    example only knows the declared-hook contract.

    ``trust_remote_code`` must be threaded in from the caller's own resolved
    ``--trust-remote-code`` flag -- this is a separate tokenizer load outside
    the main engine, so it needs the same explicit user opt-in rather than
    defaulting to trusting whatever ``--model`` was passed.

    ``validate_tokenizer``, when the model declares one via
    ``get_ar_tokenizer_validator``, is called on a successfully-loaded real
    tokenizer -- outside the load's try/except, so a validation failure (a
    model/tokenizer revision drifting from hardcoded special-token ids)
    raises instead of being swallowed into the string-prompt fallback.

    ``allow_tokenizer_fallback`` (default ``False``, i.e. fail fast): a
    tokenizer load failure normally raises, since silently degrading to the
    string-prompt form can produce a request that completes with subtly
    wrong stop tokens instead of surfacing the real problem (missing
    ``--trust-remote-code``, network/cache issue, etc.). Set ``True`` only
    for explicit unit tests or a deliberate offline-compat run where the
    caller has already decided a degraded prompt is acceptable.
    """
    try:
        from transformers import AutoTokenizer

        ar_tokenizer: Any | None = AutoTokenizer.from_pretrained(model, trust_remote_code=trust_remote_code)
    except Exception as exc:  # noqa: BLE001 - re-raised unless the caller opts into the fallback
        if not allow_tokenizer_fallback:
            raise
        logger.warning(f"AR tokenizer load failed ({exc}); falling back to string prompt (no BPE parity).")
        ar_tokenizer = None

    if ar_tokenizer is not None and validate_tokenizer is not None:
        validate_tokenizer(ar_tokenizer)

    ar_inputs = ar_input_builder(
        prompt=prompt_text,
        tokenizer=ar_tokenizer,
        extra_body=extra_body,
        num_images=num_images,
        height=height,
        width=width,
        text_output=text_output,
    )

    if ar_inputs.prompt_token_ids is not None:
        prompt_dict["prompt_token_ids"] = ar_inputs.prompt_token_ids
    elif ar_inputs.prompt is not None:
        prompt_dict["prompt"] = ar_inputs.prompt
    if ar_inputs.use_system_prompt:
        prompt_dict["use_system_prompt"] = ar_inputs.use_system_prompt
    prompt_dict["modalities"] = ar_inputs.modalities

    # Apply stop_token_ids to exactly the stage(s) the builder names -- not
    # "whichever stage isn't a diffusion stage," which would misfire on any
    # future topology with more than one non-diffusion stage.
    for stage_index in ar_inputs.stage_indices:
        if stage_index >= len(sampling_params_list):
            continue
        stage_params = sampling_params_list[stage_index]
        if not isinstance(stage_params, OmniDiffusionSamplingParams) and hasattr(stage_params, "stop_token_ids"):
            stage_params.stop_token_ids = ar_inputs.stop_token_ids


def main():
    args = parse_args()
    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
    use_nextstep = is_nextstep_model(args.model)

    cache_config = None
    cache_backend = args.cache_backend

    if cache_backend == "cache_dit":
        # cache-dit configuration: Hybrid DBCache + SCM + TaylorSeer
        # All parameters marked with [cache-dit only] in DiffusionCacheConfig
        cache_config = {
            # DBCache parameters [cache-dit only]
            "Fn_compute_blocks": 1,  # Optimized for single-transformer models
            "Bn_compute_blocks": 0,  # Number of backward compute blocks
            "max_warmup_steps": 4,  # Maximum warmup steps (works for few-step models)
            "residual_diff_threshold": 0.24,  # Higher threshold for more aggressive caching
            "max_continuous_cached_steps": 3,  # Limit to prevent precision degradation
            # TaylorSeer parameters [cache-dit only]
            "enable_taylorseer": False,  # Disabled by default (not suitable for few-step models)
            "taylorseer_order": 1,  # TaylorSeer polynomial order
            # SCM (Step Computation Masking) parameters [cache-dit only]
            "scm_steps_mask_policy": None,  # SCM mask policy: None (disabled), "slow", "medium", "fast", "ultra"
            "scm_steps_policy": "dynamic",  # SCM steps policy: "dynamic" or "static"
        }
    elif cache_backend == "tea_cache":
        # TeaCache configuration
        # All parameters marked with [tea_cache only] in DiffusionCacheConfig
        cache_config = {
            # TeaCache parameters [tea_cache only]
            "rel_l1_thresh": 0.2,  # Threshold for accumulated relative L1 distance
            # Note: coefficients will use model-specific defaults based on model_type
            #       (e.g., QwenImagePipeline or FluxPipeline)
        }

    profiler_enabled = args.profiler_config is not None

    # Prepare LoRA kwargs for Omni initialization
    lora_args: dict[str, Any] = {}
    if args.lora_path:
        lora_path = args.lora_path
        if len(lora_path) == 1:
            lora_path = lora_path[0]
        lora_args["lora_path"] = lora_path
        lora_args["lora_backend"] = args.lora_backend
        print(f"Using LoRA from: {lora_path} with backend: {args.lora_backend}")

    # Build quantization kwargs: use quantization_config dict when
    # ignored_layers is specified so the list flows through OmniDiffusionConfig
    quant_kwargs: dict[str, Any] = {}
    ignored_layers = [s.strip() for s in args.ignored_layers.split(",") if s.strip()] if args.ignored_layers else None
    if args.quantization and ignored_layers:
        quant_kwargs["quantization_config"] = {
            "method": args.quantization,
            "ignored_layers": ignored_layers,
        }
    elif args.quantization:
        quant_kwargs["quantization"] = args.quantization

    omni_kwargs = {
        "model": args.model,
        "enable_layerwise_offload": args.enable_layerwise_offload,
        "vae_use_slicing": args.vae_use_slicing,
        "vae_use_tiling": args.vae_use_tiling,
        "cache_backend": args.cache_backend,
        "cache_config": cache_config,
        "enable_cache_dit_summary": args.enable_cache_dit_summary,
        "ulysses_degree": args.ulysses_degree,
        "ring_degree": args.ring_degree,
        "ulysses_mode": args.ulysses_mode,
        "cfg_parallel_size": args.cfg_parallel_size,
        "vae_patch_parallel_size": args.vae_patch_parallel_size,
        "enable_expert_parallel": args.enable_expert_parallel,
        "enable_cpu_offload": args.enable_cpu_offload,
        "mode": "text-to-image",
        "log_stats": args.log_stats,
        "enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler,
        "profiler_config": args.profiler_config,
        "init_timeout": args.init_timeout,
        "stage_init_timeout": args.stage_init_timeout,
        "auxiliary_text_encoder": args.auxiliary_text_encoder,
        **lora_args,
        **quant_kwargs,
    }
    if args.tensor_parallel_size is not None:
        omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size
    if args.enforce_eager is not None:
        omni_kwargs["enforce_eager"] = args.enforce_eager
    if args.trust_remote_code:
        omni_kwargs["trust_remote_code"] = True
    if args.deploy_config:
        omni_kwargs["deploy_config"] = args.deploy_config
    if use_nextstep:
        # NextStep-1.1 requires explicit pipeline class
        omni_kwargs["model_class_name"] = "NextStep11Pipeline"
    # Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
    # gate is an engine-level config (offline analog of the server's --no-guardrails).
    if args.extra_body and "guardrails" in args.extra_body:
        omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])}
    omni = Omni(**omni_kwargs)
    model_class_name = get_model_class_name(omni)
    declared_extra_body_params = get_extra_body_params(model_class_name)

    if profiler_enabled:
        print("[Profiler] Starting profiling...")
        omni.start_profile()

    # Time profiling for generation
    print(f"\n{'=' * 60}")
    print("Generation Configuration:")
    print(f"  Model: {args.model}")
    print(f"  Inference steps: {args.num_inference_steps}")
    print(f"  Cache backend: {cache_backend if cache_backend else 'None (no acceleration)'}")
    print(f"  Quantization: {args.quantization if args.quantization else 'None (BF16)'}")
    if ignored_layers:
        print(f"  Ignored layers: {ignored_layers}")
    tp_display = args.tensor_parallel_size if args.tensor_parallel_size is not None else "deploy/default"
    print(
        f"  Parallel configuration: tensor_parallel_size={tp_display}, "
        f"ulysses_degree={args.ulysses_degree}, ulysses_mode={args.ulysses_mode}, "
        f"ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, "
        f"vae_patch_parallel_size={args.vae_patch_parallel_size}, "
        f"enable_expert_parallel={args.enable_expert_parallel}."
    )
    print(f"  CPU offload: {args.enable_cpu_offload}; CPU Layerwise Offload: {args.enable_layerwise_offload}")
    print(f"  Image size: {args.width}x{args.height}")
    if args.lora_path:
        print(f"  LoRA: scale={args.lora_scale}")
    if args.deploy_config:
        print(f"  deploy-config: {args.deploy_config}")
    print(f"{'=' * 60}\n")

    # Build LoRA request when --lora-path is set
    lora_request = None
    if args.lora_path and args.lora_backend == "peft":
        if len(args.lora_path) != 1:
            raise ValueError("Only one LoRA path is expected for PEFT backend.")

        lora_path = args.lora_path[0]
        lora_request_id = stable_lora_int_id(lora_path)
        lora_request = LoRARequest(
            lora_name=Path(lora_path).stem,
            lora_int_id=lora_request_id,
            lora_path=lora_path,
        )

    generation_start = time.perf_counter()

    prompt_dict = build_text_to_image_prompt(
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
    )
    prompt_dict = build_model_text_to_image_prompt(
        model_class_name=model_class_name,
        prompt=prompt_dict,
        height=args.height,
        width=args.width,
    )

    diffusion_params = OmniDiffusionSamplingParams(
        height=args.height,
        width=args.width,
        seed=args.seed,
        generator=generator,
        true_cfg_scale=args.cfg_scale,
        guidance_scale=args.guidance_scale,
        guidance_scale_2=args.guidance_scale_2,
        num_inference_steps=args.num_inference_steps,
        num_outputs_per_prompt=args.num_images_per_prompt,
    )

    # Base layer: backward-compatible model-specific flags. New model params should
    # instead be declared in vllm_omni/model_extras and passed via --extra-body, so
    # this dict does not need to grow per-model.
    user_extra = {
        "cfg_scale": args.cfg_scale,
        "cfg_text_scale": args.cfg_scale,
        "negative_prompt": args.negative_prompt,
        "timestep_shift": args.timesteps_shift,
        "timesteps_shift": args.timesteps_shift,
        "cfg_schedule": args.cfg_schedule,
        "use_norm": args.use_norm,
        "use_system_prompt": args.use_system_prompt,
        "system_prompt": args.system_prompt,
    }
    # Override layer: generic JSON passthrough wins over the flags above. Keys are
    # still filtered against the model's declared extra_body_params downstream.
    if args.extra_body:
        user_extra.update(args.extra_body)
    if declared_extra_body_params:
        apply_declared_extra_args(diffusion_params, declared_extra_body_params, user_extra)
    else:
        diffusion_params.extra_args.update({k: v for k, v in user_extra.items() if v is not None})

    if lora_request:
        diffusion_params.lora_request = lora_request
        diffusion_params.lora_scale = args.lora_scale

    # Build per-stage sampling params for multi-stage models (e.g. BAGEL),
    # or wrap single diffusion params for single-stage models.
    init_non_diffusion = should_init_extra_args_for_non_diffusion_stages(
        model_class_name,
    )
    defaults = list(omni.default_sampling_params_list or [])
    sampling_params_list = [clone_sampling_params(p) for p in defaults]
    if not sampling_params_list:
        sampling_params_list = [diffusion_params]

    diffusion_replaced = False
    for idx, params in enumerate(sampling_params_list):
        if isinstance(params, OmniDiffusionSamplingParams):
            sampling_params_list[idx] = diffusion_params
            diffusion_replaced = True
        elif init_non_diffusion and hasattr(params, "extra_args"):
            if params.extra_args is None:
                params.extra_args = {}
            params.extra_args.update(diffusion_params.extra_args or {})
            if args.seed is not None and hasattr(params, "seed"):
                params.seed = args.seed

            # MammothModa2's AR stage emits one visual token per grid cell,
            # one EOL token per row, and one final look-ahead token whose hidden
            # state is unavailable. Size the first stage from the prompt metadata
            # instead of SamplingParams' default of 16 tokens.
            prompt_info = prompt_dict.get("additional_information", {})
            if idx == 0 and prompt_info.get("omni_task") == ["t2i"]:
                ar_width = int(prompt_info.get("ar_width", [0])[0])
                ar_height = int(prompt_info.get("ar_height", [0])[0])
                if ar_width > 0 and ar_height > 0:
                    params.max_tokens = ar_height * (ar_width + 1) + 1

    if not diffusion_replaced and len(sampling_params_list) == 1:
        sampling_params_list = [diffusion_params]

    # Models with an AR text stage (e.g. HunyuanImage3) declare an
    # ar_input_builder. When present, build the AR prefill token-ids and AR
    # stop-token-ids declaratively from the plain prompt + extra_body, so this
    # example stays model-agnostic. Models without one are untouched.
    ar_input_builder = get_ar_input_builder(model_class_name)
    if ar_input_builder is not None:
        _apply_ar_stage_inputs(
            ar_input_builder,
            model=args.model,
            prompt_text=args.prompt,
            extra_body=user_extra,
            num_images=0,
            height=args.height,
            width=args.width,
            prompt_dict=prompt_dict,
            sampling_params_list=sampling_params_list,
            trust_remote_code=args.trust_remote_code,
            validate_tokenizer=get_ar_tokenizer_validator(model_class_name),
            allow_tokenizer_fallback=args.allow_tokenizer_fallback,
        )

    outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list)

    generation_end = time.perf_counter()
    generation_time = generation_end - generation_start

    # Print profiling results
    print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")

    if profiler_enabled:
        print("\n[Profiler] Stopping profiler and collecting results...")
        profile_results = omni.stop_profile()
        if profile_results and isinstance(profile_results, dict):
            traces = profile_results.get("traces", [])
            print("\n" + "=" * 60)
            print("PROFILING RESULTS:")
            for rank, trace in enumerate(traces):
                print(f"\nRank {rank}:")
                if trace:
                    print(f"  • Trace: {trace}")
            if not traces:
                print("  No traces collected.")
            print("=" * 60)
        else:
            print("[Profiler] No valid profiling data returned.")

    # omni.generate() returns list[OmniRequestOutput]
    if not outputs or len(outputs) == 0:
        raise ValueError("No output generated from omni.generate()")
    logger.info(f"Outputs: {outputs}")

    images = None
    for output in outputs:
        images = getattr(output, "images", None)
        if images:
            break
        req_out = output
        images = getattr(req_out, "images", None) if req_out is not None else None
        if images:
            break

    if not images:
        images = extract_images_from_outputs(outputs)

    if not images:
        raise ValueError("No images found in request_output")
    images = _normalize_images_for_save(images)

    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    suffix = output_path.suffix or ".png"
    stem = output_path.stem or "qwen_image_output"
    if len(images) <= 1:
        images[0].save(output_path)
        print(f"Saved generated image to {output_path}")
    else:
        for idx, img in enumerate(images):
            save_path = output_path.parent / f"{stem}_{idx}{suffix}"
            img.save(save_path)
            print(f"Saved generated image to {save_path}")


if __name__ == "__main__":
    main()