1 line
59 KiB
JSON
1 line
59 KiB
JSON
{"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"}}}} |