# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni 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") parse_custom_pipeline_args = functools.partial(parse_json_object, flag_name="--custom-pipeline-args") 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( "--model-class-name", type=str, default=None, help="Override the diffusion pipeline class name (e.g. AnimaPipeline).", ) parser.add_argument( "--custom-pipeline-args", type=parse_custom_pipeline_args, default=None, help='JSON object passed to native/custom pipelines (e.g. \'{"components_path": "/path"}\').', ) 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, "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 args.model_class_name: omni_kwargs["model_class_name"] = args.model_class_name elif use_nextstep: # NextStep-1.1 requires explicit pipeline class omni_kwargs["model_class_name"] = "NextStep11Pipeline" if args.custom_pipeline_args is not None: omni_kwargs["custom_pipeline_args"] = args.custom_pipeline_args # 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}") if args.model_class_name: print(f" Model class name: {args.model_class_name}") if args.custom_pipeline_args is not None: print(f" Custom pipeline args: {args.custom_pipeline_args}") 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) # A model can also be deployed with only its diffusion stage. Keep those # requests on the string-prompt path, as in the single-stage images API. has_ar_stage = any(not isinstance(params, OmniDiffusionSamplingParams) for params in sampling_params_list) if ar_input_builder is not None and has_ar_stage: _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()