{"repo": "celeste-jq/vllm-omni-hunyuanimage3", "files": {"vllm_omni/diffusion/data.py": {"82cce633a3058b98248863b716ed7e64017bd0ef": {"t": 1785768250.0273838, "n": 57879, "h": "d2fcd73f84f92dff", "c": "# adapted from sglang and fastvideo\n# SPDX-License-Identifier: Apache-2.0\n# SPDX-FileCopyrightText: Copyright contributors to the vLLM project\nimport copy\nimport os\nimport random\nfrom collections.abc import Callable, Mapping\nfrom dataclasses import dataclass, field, fields\nfrom typing import TYPE_CHECKING, Any\n\nimport diffusers\nimport torch\nfrom PIL import Image\nfrom pydantic import Field, model_validator\nfrom typing_extensions import Self\nfrom vllm.config.utils import config\nfrom vllm.logger import init_logger\nfrom vllm.model_executor.layers.quantization.base_config import (\n    QuantizationConfig,\n)\n\nfrom vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata\nfrom vllm_omni.diffusion.utils.network_utils import is_port_available\nfrom vllm_omni.quantization import build_quant_config\n\nif TYPE_CHECKING:\n    from vllm.config import ProfilerConfig\n\n# Import after TYPE_CHECKING to avoid circular imports at runtime\n# The actual import is deferred to __post_init__ to avoid import order issues\n\nlogger = init_logger(__name__)\n\n\ndef parse_kv_cache_skip_selector(\n    selector: str | list[int] | tuple[int, ...] | set[int] | None,\n) -> set[int] | None:\n    \"\"\"Parse a non-negative index selector such as \"0-9,20,25-30\".\"\"\"\n    if selector is None:\n        return None\n    if isinstance(selector, set):\n        values = selector\n    elif isinstance(selector, (list, tuple)):\n        values = set(selector)\n    elif isinstance(selector, str):\n        text = selector.strip()\n        if not text:\n            return None\n        values: set[int] = set()\n        for chunk in text.split(\",\"):\n            token = chunk.strip()\n            if not token:\n                continue\n            if \"-\" in token:\n                start_str, end_str = token.split(\"-\", 1)\n                try:\n                    start = int(start_str.strip())\n                    end = int(end_str.strip())\n                except ValueError as exc:\n                    raise ValueError(f\"Invalid range token '{token}' in selector '{selector}'.\") from exc\n                if start < 0 or end < 0 or start > end:\n                    raise ValueError(f\"Invalid range token '{token}' in selector '{selector}'.\")\n                values.update(range(start, end + 1))\n            else:\n                try:\n                    index = int(token)\n                except ValueError as exc:\n                    raise ValueError(f\"Invalid index token '{token}' in selector '{selector}'.\") from exc\n                if index < 0:\n                    raise ValueError(f\"Negative index '{index}' is not allowed in selector '{selector}'.\")\n                values.add(index)\n    else:\n        raise TypeError(f\"Unsupported selector type: {type(selector)!r}\")\n\n    for idx in values:\n        if not isinstance(idx, int):\n            raise TypeError(f\"Selector index must be int, got {type(idx)!r}\")\n        if idx < 0:\n            raise ValueError(\"Selector indices must be non-negative.\")\n    return values\n\n\n@config\n@dataclass\nclass DiffusionParallelConfig:\n    \"\"\"Configuration for diffusion model distributed execution.\"\"\"\n\n    pipeline_parallel_size: int = 1\n    \"\"\"Number of pipeline parallel stages.\"\"\"\n\n    data_parallel_size: int = 1\n    \"\"\"Number of data parallel groups.\"\"\"\n\n    tensor_parallel_size: int = 1\n    \"\"\"Number of tensor parallel groups.\"\"\"\n\n    enable_expert_parallel: bool = False\n    \"\"\"Enable expert parallelism for MoE layers (TP is still used for non-MoE layers).\"\"\"\n\n    sequence_parallel_size: int | None = None\n    \"\"\"Number of sequence parallel groups. sequence_parallel_size = ring_degree * ulysses_degree\"\"\"\n\n    ulysses_degree: int = 1\n    \"\"\"Number of GPUs used for ulysses sequence parallelism.\"\"\"\n\n    ring_degree: int = 1\n    \"\"\"Number of GPUs used for ring sequence parallelism.\"\"\"\n\n    ulysses_mode: str = \"strict\"\n    \"\"\"Ulysses sequence-parallel mode.\n\n    - \"strict\": Require divisibility constraints (fastest, default).\n    - \"advanced_uaa\": Enable UAA (\"Ulysses Anything Attention\") to support\n      uneven sequence lengths and non-divisible head counts.\n\n    Note:\n    - Ring attention does not support `attention_mask`, so models that rely on\n      mask-based auto-padding are still incompatible with Ring.\n    - When used in hybrid Ulysses+Ring, Ring requires consistent per-rank\n      sequence shapes across the ring group.\n    \"\"\"\n\n    cfg_parallel_size: int = 1\n    \"\"\"Number of Classifier Free Guidance (CFG) parallel groups.\"\"\"\n\n    vae_patch_parallel_size: int = 1\n    \"\"\"Number of ranks used for VAE patch/tile parallelism (decode/encode).\"\"\"\n\n    use_hsdp: bool = False\n    \"\"\"Enable Hybrid Sharded Data Parallel (HSDP) for model weight sharding.\"\"\"\n\n    hsdp_shard_size: int = -1\n    \"\"\"Number of GPUs to shard weights across within each replica group. -1 means auto-calculate.\"\"\"\n\n    hsdp_replicate_size: int = 1\n    \"\"\"Number of replica groups for HSDP. Each replica holds a full sharded copy.\"\"\"\n\n    @model_validator(mode=\"after\")\n    def _validate_parallel_config(self) -> Self:\n        \"\"\"Validates the config relationships among the parallel strategies.\"\"\"\n        assert self.pipeline_parallel_size > 0, \"Pipeline parallel size must be > 0\"\n        assert self.data_parallel_size > 0, \"Data parallel size must be > 0\"\n        assert self.tensor_parallel_size > 0, \"Tensor parallel size must be > 0\"\n        assert self.sequence_parallel_size > 0, \"Sequence parallel size must be > 0\"\n        assert self.ulysses_degree > 0, \"Ulysses degree must be > 0\"\n        assert self.ring_degree > 0, \"Ring degree must be > 0\"\n        assert self.cfg_parallel_size > 0, \"CFG parallel size must be > 0\"\n        assert self.cfg_parallel_size in [1, 2, 3], (\n            f\"CFG parallel size must be 1, 2, or 3, but got {self.cfg_parallel_size}\"\n        )\n        assert self.vae_patch_parallel_size > 0, \"VAE patch parallel size must be > 0\"\n        assert self.sequence_parallel_size == self.ulysses_degree * self.ring_degree, (\n            \"Sequence parallel size must be equal to the product of ulysses degree and ring degree,\"\n            f\" but got {self.sequence_parallel_size} != {self.ulysses_degree} * {self.ring_degree}\"\n        )\n        assert self.ulysses_mode in {\"strict\", \"advanced_uaa\"}, (\n            f\"ulysses_mode must be one of {{'strict','advanced_uaa'}}, but got {self.ulysses_mode!r}.\"\n        )\n\n        # Validate HSDP configuration\n        if self.use_hsdp:\n            assert self.hsdp_replicate_size > 0, \"HSDP replicate size must be > 0\"\n            assert self.hsdp_shard_size > 0, \"HSDP shard size must be > 0 (should be set in __post_init__)\"\n        return self\n\n    def __post_init__(self) -> None:\n        if self.sequence_parallel_size is None:\n            self.sequence_parallel_size = self.ulysses_degree * self.ring_degree\n\n        # Calculate world_size from other parallelism dimensions\n        other_parallel_world_size = (\n            self.pipeline_parallel_size\n            * self.data_parallel_size\n            * self.tensor_parallel_size\n            * self.ulysses_degree\n            * self.ring_degree\n            * self.cfg_parallel_size\n        )\n\n        # Handle HSDP configuration\n        # HSDP can work in two modes:\n        # 1. Standalone: when other parallelism is all 1, HSDP determines world_size\n        # 2. Combined: HSDP overlays on top of other parallelism\n        if self.use_hsdp:\n            if self.tensor_parallel_size > 1 or self.data_parallel_size > 1:\n                raise ValueError(\n                    \"HSDP (use_hsdp=True) cannot be used with TP or DP \"\n                    f\"(tensor_parallel_size={self.tensor_parallel_size}, \"\n                    f\"data_parallel_size={self.data_parallel_size}). \"\n                    \"Set tensor_parallel_size=1 and data_parallel_size=1 when using HSDP.\"\n                )\n            if self.hsdp_shard_size == -1:\n                # Auto-calculate: use other_parallel_world_size as shard_size\n                if self.hsdp_replicate_size <= 0:\n                    raise ValueError(\"hsdp_replicate_size must be > 0\")\n                if other_parallel_world_size == 1:\n                    raise ValueError(\n                        \"Cannot auto-calculate hsdp_shard_size when other parallelism is all 1. \"\n                        \"Please specify hsdp_shard_size explicitly for standalone HSDP.\"\n                    )\n                if other_parallel_world_size % self.hsdp_replicate_size != 0:\n                    raise ValueError(\n                        f\"Invalid HSDP configuration: replicate_size ({self.hsdp_replicate_size}) \"\n                        f\"must evenly divide world_size ({other_parallel_world_size}) when shard_size is -1.\"\n                    )\n                self.hsdp_shard_size = other_parallel_world_size // self.hsdp_replicate_size\n                self.world_size = other_parallel_world_size\n            else:\n                # Explicit shard_size: HSDP can work standalone or combined\n                hsdp_world_size = self.hsdp_replicate_size * self.hsdp_shard_size\n                if other_parallel_world_size == 1:\n                    # Standalone HSDP: world_size is determined by HSDP\n                    self.world_size = hsdp_world_size\n                else:\n                    # Combined: HSDP must match other parallelism world_size\n                    if hsdp_world_size != other_parallel_world_size:\n                        raise ValueError(\n                            f\"HSDP dimensions \"\n                            f\"({self.hsdp_replicate_size} × {self.hsdp_shard_size} = {hsdp_world_size}) \"\n                            f\"must equal world_size from other parallelism ({other_parallel_world_size})\"\n                        )\n                    self.world_size = other_parallel_world_size\n        else:\n            self.world_size = other_parallel_world_size\n\n    @classmethod\n    def from_dict(cls, data: dict[str, Any]) -> \"DiffusionParallelConfig\":\n        \"\"\"\n        Create DiffusionParallelConfig from a dictionary.\n\n        Args:\n            data: Dictionary containing parallel configuration parameters\n\n        Returns:\n            DiffusionParallelConfig instance with parameters set from dict\n        \"\"\"\n        if not isinstance(data, dict):\n            raise TypeError(f\"Expected parallel config dict, got {type(data)!r}\")\n        return cls(**data)\n\n\n@dataclass\nclass TransformerConfig:\n    \"\"\"Container for raw transformer configuration dictionaries.\"\"\"\n\n    params: dict[str, Any] = field(default_factory=dict)\n    quant_method: str | None = None\n    quant_config: \"QuantizationConfig | None\" = None\n\n    @classmethod\n    def from_dict(cls, data: dict[str, Any]) -> \"TransformerConfig\":\n        if not isinstance(data, dict):\n            raise TypeError(f\"Expected transformer config dict, got {type(data)!r}\")\n        params = dict(data)  # copy to avoid mutating caller's dict\n\n        quant_method: str | None = None\n        quant_config: QuantizationConfig | None = None\n        disk_qc = params.get(\"quantization_config\")\n        if isinstance(disk_qc, dict):\n            raw_quant_method = disk_qc.get(\"quant_method\", disk_qc.get(\"method\"))\n            quant_config = build_quant_config(disk_qc)\n            if quant_config is not None:\n                quant_method = raw_quant_method if raw_quant_method is not None else quant_config.get_name()\n\n        return cls(params=params, quant_method=quant_method, quant_config=quant_config)\n\n    def to_dict(self) -> dict[str, Any]:\n        return dict(self.params)\n\n    def get(self, key: str, default: Any | None = None) -> Any:\n        return self.params.get(key, default)\n\n    def __getattr__(self, item: str) -> Any:\n        params = object.__getattribute__(self, \"params\")\n        try:\n            return params[item]\n        except KeyError as exc:\n            raise AttributeError(item) from exc\n\n\n@dataclass\nclass DiffusionCacheConfig:\n    \"\"\"\n    Configuration for cache adapters (TeaCache, cache-dit, MagCache, etc.).\n\n    This dataclass provides a unified interface for cache configuration parameters.\n    It can be initialized from a dictionary and accessed via attributes.\n\n    Common parameters:\n        - TeaCache: rel_l1_thresh, coefficients (optional)\n        - cache-dit: Fn_compute_blocks, Bn_compute_blocks, max_warmup_steps,\n                    residual_diff_threshold, enable_taylorseer, taylorseer_order,\n                    scm_steps_mask_policy, scm_steps_policy\n        - MagCache: mag_threshold, mag_max_skip_steps, mag_retention_ratio,\n                    mag_ratios, mag_calibrate\n\n    Example:\n        >>> # From dict (user-facing API) - partial config uses defaults for missing keys\n        >>> config = DiffusionCacheConfig.from_dict({\"rel_l1_thresh\": 0.3})\n        >>> # Access via attribute\n        >>> print(config.rel_l1_thresh)  # 0.3 (from dict)\n        >>> print(config.Fn_compute_blocks)  # 8 (default)\n        >>> # Empty dict uses all defaults\n        >>> default_config = DiffusionCacheConfig.from_dict({})\n        >>> print(config.rel_l1_thresh)  # 0.2 (default)\n    \"\"\"\n\n    # TeaCache parameters [tea_cache only]\n    # Default: 0.2 provides ~1.5x speedup with minimal quality loss (optimal balance)\n    rel_l1_thresh: float = 0.2\n    coefficients: list[float] | None = None  # Uses model-specific defaults if None\n\n    # MagCache parameters [mag_cache only]\n    # Default: 0.24 threshold for accumulated magnitude error\n    mag_threshold: float = 0.24\n    # Default: 5 maximum consecutive skip steps (K)\n    mag_max_skip_steps: int = 5\n    # Default: 0.1 fraction of initial steps where skipping is disabled (stability)\n    mag_retention_ratio: float = 0.1\n    # Default: None magnitude ratios (model-specific, required for inference)\n    mag_ratios: list[float] | None = None\n    # Default: False calibration mode (computes mag_ratios on first run)\n    mag_calibrate: bool = False\n\n    # cache-dit parameters [cache-dit only]\n    # Default: 1 forward compute block (optimized for single-transformer models)\n    # Use 1 as default instead of cache-dit's 8, optimized for single-transformer models\n    # This provides better performance while maintaining quality for most use cases\n    Fn_compute_blocks: int = 1\n    # Default: 0 backward compute blocks (no fusion by default)\n    Bn_compute_blocks: int = 0\n    # Default: 4 warmup steps (optimized for few-step distilled models like Z-Image with 8 steps)\n    # Use 4 as default warmup steps instead of 8 in cache-dit, making DBCache work\n    # for few-step distilled models (e.g., Z-Image with 8 steps)\n    max_warmup_steps: int = 4\n    # Default: -1 (unlimited cached steps) - DBCache disables caching when previous cached steps exceed this value\n    # to prevent precision degradation. Set to -1 for unlimited caching (cache-dit default).\n    max_cached_steps: int = -1\n    # Default: 0.24 residual difference threshold (higher for more aggressive caching)\n    # Use a relatively higher residual diff threshold (0.24) as default to allow more\n    # aggressive caching. This is safe because we have max_continuous_cached_steps limit.\n    # Without this limit, a lower threshold like 0.12 would be needed.\n    residual_diff_threshold: float = 0.24\n    # Default: Limit consecutive cached steps to 3 to prevent precision degradation\n    # This allows us to use a higher residual_diff_threshold for more aggressive caching\n    max_continuous_cached_steps: int = 3\n    # Default: Disable TaylorSeer (not suitable for few-step distilled models)\n    # TaylorSeer is not suitable for few-step distilled models, so we disable it by default.\n    # References:\n    # - From Reusing to Forecasting: Accelerating Diffusion Models with TaylorSeers\n    # - Forecast then Calibrate: Feature Caching as ODE for Efficient Diffusion Transformers\n    enable_taylorseer: bool = False\n    # Default: 1st order TaylorSeer polynomial\n    taylorseer_order: int = 1\n    # Default: None SCM mask policy (disabled by default)\n    scm_steps_mask_policy: str | None = None\n    # Default: \"dynamic\" steps policy for adaptive caching\n    scm_steps_policy: str = \"dynamic\"\n    # Used by cache-dit for scm mask generation. If this value changes during inference,\n    # we will re-generate the scm mask and refresh the cache context.\n    num_inference_steps: int | None = None\n    # Force refresh the cache at a specific step index hint, useful for models like\n    # GLM-Image (image preprocessing step in editing mode).\n    force_refresh_step_hint: int | None = None\n    # Policy for force refresh: \"once\" refreshes only at the hint step,\n    # \"repeat\" refreshes every force_refresh_step_hint steps.\n    force_refresh_step_policy: str = \"once\"\n\n    # Additional parameters that may be passed but not explicitly defined\n    _extra_params: dict[str, Any] = field(default_factory=dict, repr=False)\n\n    @classmethod\n    def from_dict(cls, data: dict[str, Any]) -> \"DiffusionCacheConfig\":\n        \"\"\"\n        Create DiffusionCacheConfig from a dictionary.\n\n        Args:\n            data: Dictionary containing cache configuration parameters\n\n        Returns:\n            DiffusionCacheConfig instance with parameters set from dict\n        \"\"\"\n        if not isinstance(data, dict):\n            raise TypeError(f\"Expected cache config dict, got {type(data)!r}\")\n\n        # Get all dataclass field names automatically\n        field_names = {f.name for f in fields(cls)}\n\n        # Extract parameters that match dataclass fields (excluding private fields)\n        known_params = {k: v for k, v in data.items() if k in field_names and not k.startswith(\"_\")}\n\n        # Store extra parameters\n        extra_params = {k: v for k, v in data.items() if k not in field_names}\n\n        # Create instance with known params (missing ones will use defaults)\n        # Then update _extra_params after creation since it's a private field\n        instance = cls(**known_params, _extra_params=extra_params)\n        return instance\n\n    def __getattr__(self, item: str) -> Any:\n        \"\"\"\n        Allow access to extra parameters via attribute access.\n\n        This enables accessing parameters that weren't explicitly defined\n        in the dataclass fields but were passed in the dict.\n        \"\"\"\n        if item == \"_extra_params\" or item.startswith(\"_\"):\n            return object.__getattribute__(self, item)\n\n        extra = object.__getattribute__(self, \"_extra_params\")\n        if item in extra:\n            return extra[item]\n\n        raise AttributeError(f\"'{self.__class__.__name__}' object has no attribute '{item}'\")\n\n\n@dataclass\nclass OmniDiffusionConfig:\n    # Model and path configuration (for convenience)\n    stage_id: int = 0\n\n    model: str | None = None\n\n    model_class_name: str | None = None\n\n    dtype: torch.dtype = torch.bfloat16\n\n    model_config: dict[str, Any] = field(default_factory=dict)\n    tf_model_config: TransformerConfig = field(default_factory=TransformerConfig)\n\n    # Attention\n    diffusion_attention_config: \"AttentionConfig\" = field(default_factory=lambda: AttentionConfig())\n\n    # Running mode\n    # mode: ExecutionMode = ExecutionMode.INFERENCE\n\n    # Workload type\n    # workload_type: WorkloadType = WorkloadType.T2V\n\n    # Cache strategy (legacy)\n    cache_strategy: str = \"none\"\n    parallel_config: DiffusionParallelConfig = field(default_factory=DiffusionParallelConfig)\n\n    # Cache backend configuration (NEW)\n    cache_backend: str = \"none\"  # \"tea_cache\", \"deep_cache\", etc.\n    cache_config: DiffusionCacheConfig | dict[str, Any] = field(default_factory=dict)\n    enable_cache_dit_summary: bool = False\n\n    # Prompt-embedding cache. When enabled, ``DiffusionModelRunner`` wraps the\n    # pipeline's ``encode_prompt`` so repeated calls with identical prompt\n    # arguments (e.g. GRPO rollouts that sample the same prompt many times\n    # with different seeds) reuse the text-encoder output instead of re-running\n    # it. Safe against inputs that cannot be hashed (tensors, PIL images):\n    # those calls transparently bypass the cache.\n    enable_prompt_embed_cache: bool = False\n    prompt_embed_cache_size: int = 32\n\n    # Distributed executor backend\n    distributed_executor_backend: str = \"mp\"\n    nccl_port: int | None = None\n\n    # HuggingFace specific parameters\n    trust_remote_code: bool = False\n    revision: str | None = None\n\n    num_gpus: int | None = None\n\n    dist_timeout: int | None = None  # timeout for torch.distributed\n\n    # pipeline_config: PipelineConfig = field(default_factory=PipelineConfig, repr=False)\n\n    # LoRA parameters\n    lora_path: str | None = None\n    lora_scale: float = 1.0\n    max_cpu_loras: int | None = None\n\n    output_type: str = \"pil\"\n\n    # CPU offload parameters\n    # When enabled, DiT and encoders swap GPU access (mutual exclusion):\n    # - Text encoders run on GPU while DiT is on CPU\n    # - DiT runs on GPU while encoders are on CPU\n    enable_cpu_offload: bool = False\n    # Layer-wise offloading (block-level offloading) parameters\n    enable_layerwise_offload: bool = False\n\n    pin_cpu_memory: bool = True  # Use pinned memory for faster transfers when offloading\n\n    # VAE memory optimization parameters\n    vae_use_slicing: bool = False\n    vae_use_tiling: bool = False\n\n    # STA (Sliding Tile Attention) parameters\n    mask_strategy_file_path: str | None = None\n    # STA_mode: STA_Mode = STA_Mode.STA_INFERENCE\n    skip_time_steps: int = 15\n\n    # Compilation\n    enforce_eager: bool = False\n\n    # Parallel weight loading (for faster diffusion model startup)\n    enable_multithread_weight_load: bool = True\n    num_weight_load_threads: int = 4\n\n    # Enable sleep mode\n    enable_sleep_mode: bool = False\n\n    disable_autocast: bool = False\n\n    # VSA parameters\n    VSA_sparsity: float = 0.0  # inference/validation sparsity\n\n    # V-MoBA parameters\n    moba_config_path: str | None = None\n    # moba_config: dict[str, Any] = field(default_factory=dict)\n\n    # Master port for distributed inference\n    # TODO: do not hard code\n    master_port: int | None = None\n\n    # Worker extension class for custom functionality\n    worker_extension_cls: str | None = None\n\n    # Custom pipeline arguments for custom pipelines\n    custom_pipeline_args: dict[str, Any] | None = None\n\n    # Diffusion model loading format\n    # \"default\", \"custom_pipeline\", \"dummy\", \"diffusers\" (HF diffusers adapter)\n    diffusion_load_format: str = \"default\"\n\n    # Diffusers adapter kwargs\n    # kwargs forwarded to DiffusionPipeline.from_pretrained()\n    diffusers_load_kwargs: dict[str, Any] = field(default_factory=dict)\n    # kwargs forwarded to pipeline.__call__()\n    diffusers_call_kwargs: dict[str, Any] = field(default_factory=dict)\n    # Actual diffusers pipeline object (to determine inputs of the dummy run)\n    diffusers_pipeline_cls: type[diffusers.DiffusionPipeline] | None = None  # pyright: ignore[reportPrivateImportUsage]\n\n    # http server endpoint config, would be ignored in local mode\n    host: str | None = None\n    port: int | None = None\n\n    scheduler_port: int = 5555\n\n    # Stage verification\n    enable_stage_verification: bool = True\n\n    # Prompt text file for batch processing\n    prompt_file_path: str | None = None\n\n    # model paths for correct deallocation\n    model_paths: dict[str, str] = field(default_factory=dict)\n    model_loaded: dict[str, bool] = field(\n        default_factory=lambda: {\n            \"transformer\": True,\n            \"vae\": True,\n        }\n    )\n    override_transformer_cls_name: str | None = None\n\n    # # DMD parameters\n    # dmd_denoising_steps: List[int] | None = field(default=None)\n\n    # MoE parameters used by Wan2.2\n    boundary_ratio: float | None = None\n    # Scheduler flow_shift for Wan2.2 (12.0 for 480p, 5.0 for 720p)\n    flow_shift: float | None = None\n\n    # Support multi-image inputs and expose any model-specific request limit\n    # through a generic config field so serving code stays model-agnostic.\n    supports_multimodal_inputs: bool = False\n    max_multimodal_image_inputs: int | None = None\n\n    log_level: str = \"info\"\n\n    # Omni configuration (injected from stage config)\n    omni_kv_config: dict[str, Any] = field(default_factory=dict)\n    additional_config: dict[str, Any] = field(default_factory=dict)\n\n    profiler_config: \"ProfilerConfig | dict[str, Any] | None\" = None\n\n    # Model-specific function for collecting CFG KV caches (set at runtime)\n    cfg_kv_collect_func: Any | None = None\n\n    # Quantization: str method name, dict config, QuantizationConfig, or None.\n    # str is resolved to {\"method\": <str>} internally.\n    # Per-component: {\"transformer\": {\"method\": \"fp8\"}, \"vae\": None}\n    quantization_config: str | QuantizationConfig | dict[str, Any] | None = None\n    # Explicit runtime override for ModelOpt FP8 diffusion checkpoints. This\n    # does not enable FP8 by itself; it only selects CUTLASS once the checkpoint\n    # has already resolved to vLLM's ModelOpt FP8 linear method.\n    force_cutlass_fp8: bool = False\n\n    # Diffusion attention KV cache dtype (not vLLM's --kv-cache-dtype for AR models).\n    # None = native dtype (no quantization).\n    # \"fp8\" = dynamic FP8 (float8_e4m3fn) quantization per forward pass.\n    # On Hopper+FA3: native FP8 attention (memory + compute savings).\n    # On other backends: no benefit, backends skip quantization.\n    diffusion_kv_cache_dtype: str | None = None\n    # Optional skip selectors for KV-cache quantization. Format: \"0-9,20,25-30\".\n    # Listed steps/layers skip quantization; others keep quantized execution.\n    diffusion_kv_cache_skip_steps: str | None = None\n    diffusion_kv_cache_skip_layers: str | None = None\n    diffusion_kv_cache_skip_step_indices: set[int] | None = None\n    diffusion_kv_cache_skip_layer_indices: set[int] | None = None\n\n    # Diffusion pipeline Profiling config\n    enable_diffusion_pipeline_profiler: bool = False\n\n    # Step mode settings\n    step_execution: bool = False\n\n    # sleep mode\n    enable_sleep_mode: bool = False\n\n    # Maximum number of sequences to generate in a batch\n    max_num_seqs: int = 1\n\n    # Supplementary model specific parameters\n    extras: dict[str, Any] = Field(default_factory=dict)\n\n    @property\n    def is_moe(self) -> bool:\n        num_experts = self.tf_model_config.get(\"num_experts\", None)\n        if not isinstance(num_experts, (list, tuple, int)):\n            return False\n        if isinstance(num_experts, int):\n            return num_experts > 0\n\n        if isinstance(num_experts, (list, tuple)):\n            return any(isinstance(n, int) and n > 0 for n in num_experts)\n\n        return False\n\n    def _resolve_master_port(self) -> int:\n        \"\"\"Resolve torch.distributed master port without unnecessary random jitter.\n\n        Precedence:\n        1. ``MASTER_PORT`` environment variable (set by orchestrators for multi-replica launch).\n        2. Explicit ``master_port`` passed at construction time.\n        3. An OS-assigned ephemeral port when neither is provided.\n        \"\"\"\n        from vllm.utils.network_utils import get_open_port\n\n        from vllm_omni.diffusion import envs\n\n        env_port = envs.MASTER_PORT\n        if env_port is not None:\n            return self.settle_port(env_port, port_inc=37)\n        if self.master_port is not None:\n            return self.settle_port(self.master_port, port_inc=37)\n        return self.settle_port(get_open_port(), port_inc=37)\n\n    def settle_port(self, port: int, port_inc: int = 42, max_attempts: int = 100) -> int:\n        \"\"\"\n        Find an available port with retry logic.\n\n        Args:\n            port: Initial port to check\n            port_inc: Port increment for each attempt\n            max_attempts: Maximum number of attempts to find an available port\n\n        Returns:\n            An available port number\n\n        Raises:\n            RuntimeError: If no available port is found after max_attempts\n        \"\"\"\n        attempts = 0\n        original_port = port\n\n        while attempts < max_attempts:\n            if is_port_available(port):\n                if attempts > 0:\n                    logger.info(f\"Port {original_port} was unavailable, using port {port} instead\")\n                return port\n\n            attempts += 1\n            if port < 60000:\n                port += port_inc\n            else:\n                # Wrap around with randomization to avoid collision\n                port = 5000 + random.randint(0, 1000)\n\n        raise RuntimeError(\n            f\"Failed to find available port after {max_attempts} attempts (started from port {original_port})\"\n        )\n\n    def __post_init__(self):\n        self.master_port = self._resolve_master_port()\n\n        if isinstance(self.profiler_config, dict):\n            from vllm.config import ProfilerConfig\n\n            self.profiler_config = ProfilerConfig(**self.profiler_config)\n\n        if self.additional_config is None:\n            self.additional_config = {}\n        elif isinstance(self.additional_config, Mapping):\n            self.additional_config = dict(self.additional_config)\n        else:\n            raise TypeError(f\"additional_config must be a mapping or None, got {type(self.additional_config)!r}\")\n\n        # Convert parallel_config dict/DictConfig to DiffusionParallelConfig\n        # Use Mapping to handle both plain dicts and OmegaConf DictConfig\n        if isinstance(self.parallel_config, Mapping):\n            self.parallel_config = DiffusionParallelConfig.from_dict(dict(self.parallel_config))\n        elif not isinstance(self.parallel_config, DiffusionParallelConfig):\n            self.parallel_config = DiffusionParallelConfig()\n\n        if self.num_gpus is None:\n            if self.parallel_config is not None:\n                self.num_gpus = self.parallel_config.world_size\n            else:\n                self.num_gpus = 1\n\n        if self.num_gpus < self.parallel_config.world_size:\n            raise ValueError(\n                f\"num_gpus ({self.num_gpus}) < parallel_config.world_size ({self.parallel_config.world_size})\"\n            )\n\n        # Convert string dtype to torch.dtype if needed\n        if isinstance(self.dtype, str):\n            dtype_map = {\n                \"auto\": torch.bfloat16,\n                \"bfloat16\": torch.bfloat16,\n                \"bf16\": torch.bfloat16,\n                \"float16\": torch.float16,\n                \"fp16\": torch.float16,\n                \"half\": torch.float16,\n                \"float32\": torch.float32,\n                \"fp32\": torch.float32,\n                \"float\": torch.float32,\n            }\n            dtype_lower = self.dtype.lower()\n            if dtype_lower in dtype_map:\n                self.dtype = dtype_map[dtype_lower]\n            else:\n                logger.warning(f\"Unknown dtype string '{self.dtype}', defaulting to bfloat16\")\n                self.dtype = torch.bfloat16\n\n        # Convert cache_config dict to DiffusionCacheConfig if needed\n        if isinstance(self.cache_config, dict):\n            self.cache_config = DiffusionCacheConfig.from_dict(self.cache_config)\n        elif not isinstance(self.cache_config, DiffusionCacheConfig):\n            # If it's neither dict nor DiffusionCacheConfig, convert to empty config\n            self.cache_config = DiffusionCacheConfig()\n\n        # Auto-detect quantization from TransformerConfig if not explicitly set.\n        # This covers the case where tf_model_config is passed at construction\n        # time. For late (post-construction) assignment, callers should use\n        # set_tf_model_config() which propagates quant_config automatically.\n        self._propagate_quantization_from_tf_config(self.tf_model_config)\n\n        # Resolve quantization_config: str/dict -> QuantizationConfig via build_quant_config.\n        if self.quantization_config is not None:\n            if isinstance(self.quantization_config, QuantizationConfig):\n                pass  # Already built\n            elif isinstance(self.quantization_config, str):\n                self.quantization_config = build_quant_config(self.quantization_config)\n            elif isinstance(self.quantization_config, Mapping):\n                self.quantization_config = build_quant_config(dict(self.quantization_config))\n            else:\n                raise TypeError(\n                    f\"quantization_config must be str, dict, QuantizationConfig, or None, \"\n                    f\"got {type(self.quantization_config)!r}\"\n                )\n\n        # Match vLLM's config flow: parse entrypoint shorthands before the\n        # config object is built, and keep a single runtime truth source.\n        self.diffusion_attention_config = build_attention_config(self.diffusion_attention_config)\n        self.diffusion_kv_cache_skip_step_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_steps)\n        self.diffusion_kv_cache_skip_layer_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_layers)\n\n        if self.max_cpu_loras is None:\n            self.max_cpu_loras = 1\n        elif self.max_cpu_loras < 1:\n            raise ValueError(\"max_cpu_loras must be >= 1 for diffusion LoRA\")\n\n        if self.diffusion_load_format != \"diffusers\" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs):\n            raise ValueError(\n                \"diffusers_load_kwargs and diffusers_call_kwargs are only \"\n                \"valid together with diffusion_load_format=diffusers\"\n            )\n\n    def _propagate_quantization_from_tf_config(self, tf_config: \"TransformerConfig\") -> None:\n        if tf_config.quant_config is None:\n            return\n\n        is_checkpoint_fp8 = bool(getattr(tf_config.quant_config, \"is_checkpoint_fp8_serialized\", False))\n        is_checkpoint_nvfp4 = bool(getattr(tf_config.quant_config, \"is_checkpoint_nvfp4_serialized\", False))\n        should_use_checkpoint_config = (\n            self.quantization_config is None\n            or (is_checkpoint_fp8 and self._is_generic_fp8_quant_config(self.quantization_config))\n            or (is_checkpoint_nvfp4 and self._is_generic_nvfp4_quant_config(self.quantization_config))\n        )\n        if should_use_checkpoint_config:\n            self.quantization_config = tf_config.quant_config\n            logger.info(\n                \"Auto-detected quantization '%s' from model config\",\n                tf_config.quant_method,\n            )\n\n    @staticmethod\n    def _is_generic_fp8_quant_config(quant_config: object) -> bool:\n        if isinstance(quant_config, str):\n            return quant_config.lower() == \"fp8\"\n        if isinstance(quant_config, Mapping):\n            method = quant_config.get(\"method\", quant_config.get(\"quant_method\"))\n            return isinstance(method, str) and method.lower() == \"fp8\"\n        if hasattr(quant_config, \"get_name\"):\n            return quant_config.get_name() == \"fp8\"\n        return False\n\n    @staticmethod\n    def _is_generic_nvfp4_quant_config(quant_config: object) -> bool:\n        if isinstance(quant_config, str):\n            return quant_config.lower() in {\"fp4\", \"nvfp4\", \"modelopt_fp4\"}\n        if isinstance(quant_config, Mapping):\n            method = quant_config.get(\"method\", quant_config.get(\"quant_method\"))\n            return isinstance(method, str) and method.lower() in {\"fp4\", \"nvfp4\", \"modelopt_fp4\"}\n        if hasattr(quant_config, \"get_name\"):\n            return quant_config.get_name() == \"modelopt_fp4\"\n        return False\n\n    def set_tf_model_config(self, tf_config: \"TransformerConfig\") -> None:\n        \"\"\"Assign `tf_model_config` and propagate quantization if detected.\n\n        In the normal startup flow `OmniDiffusionConfig` is created\n        *before* the transformer `config.json` is loaded from disk, so\n        `__post_init__` sees an empty `TransformerConfig`.  Callers\n        that load the config later should use this method instead of bare\n        assignment so that an embedded `quant_config` is propagated to\n        `self.quantization_config` automatically.\n\n        Args:\n            tf_config: Transformer configuration, typically built via\n                `TransformerConfig.from_dict`.\n        \"\"\"\n        self.tf_model_config = tf_config\n        self._propagate_quantization_from_tf_config(tf_config)\n\n    def update_multimodal_support(self) -> None:\n        # Resolve serving-visible multimodal behavior from shared metadata\n        # instead of importing concrete pipeline modules into the config layer.\n        metadata = get_diffusion_model_metadata(self.model_class_name)\n        self.supports_multimodal_inputs = metadata.supports_multimodal_inputs\n        self.max_multimodal_image_inputs = metadata.max_multimodal_image_inputs\n\n    @staticmethod\n    def _looks_like_lance_subfolder(model: str | None) -> bool:\n        \"\"\"Return True when ``--model`` points at a Lance per-component subfolder.\n\n        Lance's HF repo bundles ``Lance_3B/``, ``Lance_3B_Video/`` and\n        ``Qwen2.5-VL-ViT/`` under a single top-level ``config.json``; users may\n        reasonably hand the AR-style sub-checkpoint path directly.  The\n        ``LancePipeline`` constructor knows to walk up to the repo root from\n        either subfolder name.\n        \"\"\"\n        if not model:\n            return False\n        base = os.path.basename(str(model).rstrip(\"/\"))\n        return base in {\"Lance_3B\", \"Lance_3B_Video\"}\n\n    def enrich_config(self) -> None:\n        \"\"\"Load model metadata from HuggingFace and populate config fields.\n\n        Diffusers-style models expose ``model_index.json`` with ``_class_name``.\n        Non-diffusers models (e.g. Bagel, NextStep) only have ``config.json``,\n        so we fall back to reading that and mapping model_type manually.\n        \"\"\"\n        from vllm.transformers_utils.config import get_hf_file_to_dict\n\n        # Default model_class_name for diffusers adapter\n        if self.model_class_name is None and self.diffusion_load_format == \"diffusers\":\n            self.model_class_name = \"DiffusersAdapterPipeline\"\n\n        try:\n            config_dict = get_hf_file_to_dict(\"model_index.json\", self.model)\n            if config_dict is not None:\n                if self.model_class_name is None:\n                    self.model_class_name = config_dict.get(\"_class_name\", None)\n                self.update_multimodal_support()\n\n                # Skip transformer config loading for diffusers adapter\n                # (non-DiT models don't have a separate transformer folder/config)\n                if self.diffusion_load_format == \"diffusers\":\n                    self.set_tf_model_config(TransformerConfig())\n                    try:\n                        diffusers_pipeline_cls_name = config_dict[\"_class_name\"]\n                        self.diffusers_pipeline_cls = getattr(diffusers, diffusers_pipeline_cls_name)\n                    except (KeyError, AttributeError) as exc:\n                        logger.warning(\n                            \"Could not find valid _class_name for diffusers pipeline in model_index.json: %s. \"\n                            \"Without the underlying pipeline class the dummy run may omit required inputs.\",\n                            exc,\n                        )\n                else:\n                    tf_config_dict = get_hf_file_to_dict(\"transformer/config.json\", self.model)\n                    self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict))\n            else:\n                raise FileNotFoundError(\"model_index.json not found\")\n        except (AttributeError, OSError, ValueError, FileNotFoundError):\n            # Skip transformer config loading for diffusers adapter\n            # (non-DiT models don't have a separate transformer folder/config)\n            if self.diffusion_load_format == \"diffusers\":\n                self.set_tf_model_config(TransformerConfig())\n                logger.warning(\n                    \"Could not find valid model_index.json per diffusers format. \"\n                    \"This model is likely unsupported by the diffusers backend. \"\n                    \"Also, without knowing the underlying diffusers pipeline class from model_index.json, \"\n                    \"the dummy run will input only text prompt, which may cause errors for pipelines \"\n                    \"that require additional inputs.\"\n                )\n            else:\n                cfg = get_hf_file_to_dict(\"config.json\", self.model)\n                if cfg is None:\n                    # Lance ships its top-level config.json one directory above\n                    # the per-checkpoint subfolders (``Lance_3B/`` or\n                    # ``Lance_3B_Video/``).  Try to recover that case before\n                    # raising.\n                    if self._looks_like_lance_subfolder(self.model):\n                        self.model_class_name = \"LancePipeline\"\n                        self.set_tf_model_config(TransformerConfig())\n                        self.update_multimodal_support()\n                        return\n                    raise ValueError(f\"Could not find config.json or model_index.json for model {self.model}\")\n\n                self.set_tf_model_config(TransformerConfig.from_dict(cfg))\n                model_type = cfg.get(\"model_type\")\n                architectures = cfg.get(\"architectures\") or []\n\n                if model_type == \"bagel\" or \"BagelForConditionalGeneration\" in architectures:\n                    self.model_class_name = \"BagelPipeline\"\n                    self.set_tf_model_config(TransformerConfig())\n                    self.update_multimodal_support()\n                elif (\n                    model_type == \"lance\"\n                    or \"LancePipeline\" in architectures\n                    or cfg.get(\"model_name\") == \"Lance\"\n                    or self._looks_like_lance_subfolder(self.model)\n                ):\n                    # Lance ships a non-HF top-level config.json (model_name only)\n                    # plus per-component subfolders; resolve to the Lance pipeline.\n                    # Also accept --model pointing directly at the ``Lance_3B`` or\n                    # ``Lance_3B_Video`` subfolder by walking up to the repo root.\n                    self.model_class_name = \"LancePipeline\"\n                    self.set_tf_model_config(TransformerConfig())\n                    self.update_multimodal_support()\n                elif model_type == \"neo_chat\":\n                    self.model_class_name = \"SenseNovaU1Pipeline\"\n                    self.tf_model_config = TransformerConfig()\n                    self.update_multimodal_support()\n                elif \"BailingMM2NativeForConditionalGeneration\" in architectures or model_type in (\n                    \"bailingmm_moe_v2_lite\",\n                    \"ming_flash_omni\",\n                    \"ming_flash_omni_thinker\",\n                ):\n                    # Ming-flash-omni-2.0 — imagegen stage uses the custom\n                    # ``MingImagePipeline`` (ZImage DiT + Qwen2 connector). See\n                    # vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py.\n                    self.model_class_name = \"MingImagePipeline\"\n                    self.tf_model_config = TransformerConfig()\n                    self.update_multimodal_support()\n                elif model_type == \"nextstep\":\n                    if self.model_class_name is None:\n                        self.model_class_name = \"NextStep11Pipeline\"\n                    self.set_tf_model_config(TransformerConfig())\n                    self.update_multimodal_support()\n                elif model_type == \"s2v\":\n                    if self.model_class_name is None:\n                        self.model_class_name = \"WanS2VPipeline\"\n                    self.tf_model_config = TransformerConfig()\n                    self.update_multimodal_support()\n                elif model_type == \"vla\":\n                    from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero\n\n                    if _looks_like_dreamzero(self.model):\n                        self.model_class_name = \"DreamZeroPipeline\"\n                        self.set_tf_model_config(TransformerConfig())\n                        self.update_multimodal_support()\n                    else:\n                        raise\n                elif architectures and len(architectures) == 1:\n                    architecture = architectures[0]\n                    from vllm_omni.diffusion.registry import DiffusionModelRegistry\n\n                    if (\n                        self.model_class_name is None\n                        or DiffusionModelRegistry._try_load_model_cls(architecture) is not None\n                    ):\n                        self.model_class_name = architecture\n                else:\n                    raise\n\n    @classmethod\n    def from_kwargs(cls, **kwargs: Any) -> \"OmniDiffusionConfig\":\n        # Backwards-compatibility: older callers may use a diffusion-specific\n        # \"static_lora_scale\" kwarg. Normalize it to the canonical \"lora_scale\"\n        # before constructing the dataclass to avoid TypeError on unknown fields.\n        if \"static_lora_scale\" in kwargs:\n            if \"lora_scale\" not in kwargs:\n                kwargs[\"lora_scale\"] = kwargs[\"static_lora_scale\"]\n            kwargs.pop(\"static_lora_scale\", None)\n\n        # Backwards-compatibility: map \"quantization\" to \"quantization_config\"\n        # so callers using the old field name still work.\n        if \"quantization\" in kwargs and kwargs.get(\"quantization_config\", None) is None:\n            kwargs[\"quantization_config\"] = kwargs.pop(\"quantization\")\n        else:\n            kwargs.pop(\"quantization\", None)\n\n        # Renamed from kv_cache_* to avoid clashing with vLLM's --kv-cache-dtype.\n        if kwargs.get(\"diffusion_kv_cache_dtype\") is None and \"kv_cache_dtype\" in kwargs:\n            kwargs[\"diffusion_kv_cache_dtype\"] = kwargs.pop(\"kv_cache_dtype\")\n        else:\n            kwargs.pop(\"kv_cache_dtype\", None)\n        if kwargs.get(\"diffusion_kv_cache_skip_steps\") is None and \"kv_cache_skip_steps\" in kwargs:\n            kwargs[\"diffusion_kv_cache_skip_steps\"] = kwargs.pop(\"kv_cache_skip_steps\")\n        else:\n            kwargs.pop(\"kv_cache_skip_steps\", None)\n        if kwargs.get(\"diffusion_kv_cache_skip_layers\") is None and \"kv_cache_skip_layers\" in kwargs:\n            kwargs[\"diffusion_kv_cache_skip_layers\"] = kwargs.pop(\"kv_cache_skip_layers\")\n        else:\n            kwargs.pop(\"kv_cache_skip_layers\", None)\n\n        # Handle \"diffusion_attention_backend\" shorthand: merge into\n        # diffusion_attention_config before field filtering.\n        diffusion_attn_backend = kwargs.pop(\"diffusion_attention_backend\", None)\n        if diffusion_attn_backend is not None:\n            existing = kwargs.get(\"diffusion_attention_config\")\n            kwargs[\"diffusion_attention_config\"] = parse_attention_config(\n                existing,\n                attention_backend=diffusion_attn_backend,\n            )\n\n        # Check environment variable as fallback for cache_backend\n        # Support both old DIFFUSION_CACHE_ADAPTER and new DIFFUSION_CACHE_BACKEND for backwards compatibility\n        if \"cache_backend\" not in kwargs:\n            cache_backend = os.environ.get(\"DIFFUSION_CACHE_BACKEND\") or os.environ.get(\"DIFFUSION_CACHE_ADAPTER\")\n            kwargs[\"cache_backend\"] = cache_backend.lower() if cache_backend else \"none\"\n\n        # Convert optional YAML null values to empty containers.\n        for key in (\"diffusers_load_kwargs\", \"diffusers_call_kwargs\"):\n            if key in kwargs and kwargs[key] is None:\n                kwargs[key] = {}\n\n        # Filter kwargs to only include valid fields\n        valid_fields = {f.name for f in fields(cls)}\n        filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields}\n\n        return cls(**filtered_kwargs)\n\n\n@dataclass\nclass DiffusionOutput:\n    \"\"\"\n    Final output (after pipeline completion)\n    \"\"\"\n\n    # Fields may be replaced with SHM handle dicts by ipc.pack_diffusion_output_shm\n    output: torch.Tensor | tuple[Any, ...] | dict[str, Any] | None = None\n    trajectory_timesteps: torch.Tensor | dict[str, Any] | None = None\n    trajectory_latents: torch.Tensor | dict[str, Any] | None = None\n    trajectory_log_probs: torch.Tensor | dict[str, Any] | None = None\n    trajectory_decoded: list[Image.Image] | None = None\n    error: str | None = None\n    aborted: bool = False\n    abort_message: str | None = None\n\n    post_process_func: Callable[..., Any] | None = None\n\n    # Extra custom output data (e.g. latent trajectories, prompt embeds)\n    # passed through to OmniRequestOutput.custom_output\n    custom_output: dict[str, Any] = field(default_factory=dict)\n\n    # logged timings info, directly from Req.timings\n    # timings: Optional[\"RequestTimings\"] = None\n\n    # logged duration of stages\n    stage_durations: dict[str, float] = field(default_factory=dict)\n\n    # memory usage info\n    peak_memory_mb: float = 0.0\n\n    # When True, move all tensor fields (including tensors inside\n    # ``custom_output``) to CPU at construction time. Useful when the output\n    # is shipped across process boundaries (e.g. step-execution mode) and the\n    # receiving side must not initialise a stray CUDA context.\n    to_cpu: bool = False\n\n    def __post_init__(self) -> None:\n        if not self.to_cpu:\n            return\n\n        def _maybe_to_cpu(value: Any) -> Any:\n            if isinstance(value, torch.Tensor):\n                return value.detach().cpu()\n            return value\n\n        self.output = _maybe_to_cpu(self.output)\n        self.trajectory_timesteps = _maybe_to_cpu(self.trajectory_timesteps)\n        self.trajectory_latents = _maybe_to_cpu(self.trajectory_latents)\n        self.trajectory_log_probs = _maybe_to_cpu(self.trajectory_log_probs)\n        if self.custom_output:\n            self.custom_output = {k: _maybe_to_cpu(v) for k, v in self.custom_output.items()}\n\n\nclass DiffusionRequestAbortedError(RuntimeError):\n    \"\"\"Raised when a diffusion request ends via user-visible abort.\"\"\"\n\n\n@dataclass\nclass AttentionSpec:\n    \"\"\"Specifies a backend and its backend-specific parameters for one attention role.\"\"\"\n\n    backend: str  # registry name, e.g. \"FLASH_ATTN\"\n    extra: dict[str, Any] = field(default_factory=dict)  # backend-specific kwargs\n\n    def __post_init__(self) -> None:\n        if not isinstance(self.backend, str):\n            raise TypeError(f\"Expected str for AttentionSpec.backend, got {type(self.backend)!r}\")\n\n        if self.extra is None:\n            self.extra = {}\n        elif isinstance(self.extra, Mapping):\n            self.extra = dict(self.extra)\n        else:\n            raise TypeError(f\"Expected dict for AttentionSpec.extra, got {type(self.extra)!r}\")\n\n\n@dataclass\nclass AttentionConfig:\n    \"\"\"Per-role attention backend configuration.\n\n    Lookup precedence for a given (role, role_category):\n      1. per_role[role]         — exact match\n      2. per_role[role_category] — category fallback (e.g. \"ltx2.audio_to_video\" → \"cross\")\n      3. default                — global default\n      4. platform default       — unchanged platform logic\n    \"\"\"\n\n    default: AttentionSpec | None = None\n    per_role: dict[str, AttentionSpec] = field(default_factory=dict)\n\n    def __post_init__(self) -> None:\n        if self.default is not None:\n            self.default = self._coerce_spec_or_none(self.default, \"default\")\n\n        normalized_per_role: dict[str, AttentionSpec] = {}\n        for role_key, spec_data in self._normalize_per_role_mapping(self.per_role).items():\n            spec = self._coerce_spec_or_none(spec_data, f\"per_role[{role_key!r}]\")\n            if spec is not None:\n                normalized_per_role[role_key] = spec\n        self.per_role = normalized_per_role\n\n    @staticmethod\n    def _coerce_spec(spec_data: Any, field_name: str) -> AttentionSpec:\n        if isinstance(spec_data, AttentionSpec):\n            return spec_data\n        if isinstance(spec_data, str):\n            return AttentionSpec(backend=spec_data)\n        if isinstance(spec_data, Mapping):\n            return AttentionSpec(**dict(spec_data))\n        raise TypeError(f\"Expected str, dict, or AttentionSpec for {field_name}, got {type(spec_data)!r}\")\n\n    @classmethod\n    def _coerce_spec_or_none(cls, spec_data: Any, field_name: str) -> AttentionSpec | None:\n        spec = cls._coerce_spec(spec_data, field_name)\n        if spec.backend.lower() == \"auto\":\n            return None\n        return spec\n\n    @classmethod\n    def _normalize_per_role_mapping(cls, raw_per_role: Any) -> dict[str, Any]:\n        if raw_per_role is None:\n            return {}\n        if not isinstance(raw_per_role, Mapping):\n            raise TypeError(f\"Expected dict for AttentionConfig.per_role, got {type(raw_per_role)!r}\")\n\n        normalized: dict[str, Any] = {}\n        for role_key, spec_data in raw_per_role.items():\n            cls._flatten_per_role_entry([role_key], spec_data, normalized)\n        return normalized\n\n    @classmethod\n    def _flatten_per_role_entry(\n        cls,\n        path: list[str],\n        node: Any,\n        normalized: dict[str, Any],\n    ) -> None:\n        role = \".\".join(path)\n        if not isinstance(node, Mapping):\n            normalized[role] = node\n            return\n\n        spec_keys = {\"backend\", \"extra\"}\n        node_dict = dict(node)\n        node_keys = set(node_dict)\n        if node_keys & spec_keys:\n            if not node_keys <= spec_keys:\n                raise ValueError(\n                    f\"Invalid per_role entry for role {role!r}: cannot mix backend/extra with nested role keys.\"\n                )\n            normalized[role] = node_dict\n            return\n\n        if not node_dict:\n            raise ValueError(f\"Empty per_role entry for role {role!r}\")\n\n        for child_key, child_value in node_dict.items():\n            cls._flatten_per_role_entry([*path, child_key], child_value, normalized)\n\n    def resolve_with_source(\n        self,\n        role: str = \"self\",\n        role_category: str | None = None,\n    ) -> tuple[AttentionSpec | None, str | None]:\n        \"\"\"Resolve the AttentionSpec and report which config entry matched.\"\"\"\n        spec = self.per_role.get(role)\n        if spec is not None:\n            return spec, f\"attention_config.per_role[{role!r}]\"\n        if role_category is not None:\n            spec = self.per_role.get(role_category)\n            if spec is not None:\n                return spec, f\"attention_config.per_role[{role_category!r}] (role_category fallback)\"\n        if self.default is not None:\n            return self.default, \"attention_config.default\"\n        return None, None\n\n\ndef parse_attention_config(\n    attention_config: AttentionConfig | Mapping[str, Any] | None = None,\n    *,\n    attention_backend: str | None = None,\n) -> AttentionConfig:\n    \"\"\"Pure type-conversion: coerce *attention_config* to an AttentionConfig.\n\n    Optionally merges an ``attention_backend`` shorthand into the config's\n    ``default`` field.  This does **not** read environment variables —\n    use :func:`build_attention_config` for the full normalisation that\n    should happen exactly once in ``OmniDiffusionConfig.__post_init__``.\n    \"\"\"\n    if attention_config is None:\n        normalized = AttentionConfig()\n    elif isinstance(attention_config, AttentionConfig):\n        normalized = copy.deepcopy(attention_config)\n    elif isinstance(attention_config, Mapping):\n        normalized = AttentionConfig(**dict(attention_config))\n    else:\n        raise TypeError(\n            f\"attention_config must be an AttentionConfig, mapping, or None; got {type(attention_config)!r}\"\n        )\n\n    if attention_backend is not None:\n        if normalized.default is not None:\n            raise ValueError(\n                \"--diffusion-attention-backend is mutually exclusive with --diffusion-attention-config.default.backend.\"\n            )\n        if attention_backend.lower() != \"auto\":\n            normalized.default = AttentionSpec(backend=attention_backend)\n\n    return normalized\n\n\ndef build_attention_config(\n    attention_config: AttentionConfig | Mapping[str, Any] | None = None,\n) -> AttentionConfig:\n    \"\"\"Normalize diffusion attention config — the single authoritative entry point.\n\n    Called exactly once in ``OmniDiffusionConfig.__post_init__``.\n    Handles type-conversion **and** env-var fallback\n    (``DIFFUSION_ATTENTION_BACKEND``).\n    \"\"\"\n    normalized = parse_attention_config(attention_config)\n\n    if normalized.default is not None:\n        return normalized\n\n    env_attention_backend = os.environ.get(\"DIFFUSION_ATTENTION_BACKEND\")\n    if env_attention_backend is None:\n        return normalized\n\n    if env_attention_backend.lower() == \"auto\":\n        return normalized\n\n    normalized.default = AttentionSpec(backend=env_attention_backend)\n    logger.info(\n        \"Parsed attention config from DIFFUSION_ATTENTION_BACKEND '%s': default=%s, per_role=%s\",\n        env_attention_backend,\n        normalized.default,\n        {k: v.backend for k, v in normalized.per_role.items()},\n    )\n    return normalized\n\n\n@dataclass\nclass OmniACK:\n    \"\"\"\n    Handshake payload from Workers to Orchestrator.\n    \"\"\"\n\n    task_id: str\n    status: str\n    stage_id: int | None = None\n    rank: int | None = None\n    freed_bytes: int = 0\n    metadata: dict[str, Any] = field(default_factory=dict)\n    \"\"\"\n    Additional telemetry such as:\n    - max_contiguous_block: for fragmentation analysis.\n    - cuda_graph_recalled: boolean if graphs were successfully destroyed/rebuilt.\n    - latency_ms: time taken for the D2H/H2D transfer.\n    \"\"\"\n    error_msg: str | None = None\n\n\n@dataclass\nclass OmniSleepTask:\n    \"\"\"Structured sleep instruction.\"\"\"\n\n    task_id: str\n    level: int = 2\n    metadata: dict[str, Any] = field(default_factory=dict)\n\n\n@dataclass\nclass OmniWakeTask:\n    \"\"\"Structured wake-up instruction.\"\"\"\n\n    task_id: str\n    tags: list[str] | None = None\n\n\n# Special message broadcast via scheduler queues to signal worker shutdown.\nSHUTDOWN_MESSAGE = {\"type\": \"shutdown\"}\n"}}}}