# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Inspired by SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/server_args.py """The arguments of sglang-diffusion Inference.""" import argparse import dataclasses import json import math import os import random import sys import tempfile from dataclasses import field from enum import Enum from typing import Any, List, Literal, Optional import addict import yaml from sglang.multimodal_gen import envs from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( LTX2PipelineConfig, is_ltx23_native_variant, ) from sglang.multimodal_gen.configs.quantization.nunchaku import NunchakuSVDQuantArgs from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import ( NunchakuConfig, ) from sglang.multimodal_gen.runtime.loader.utils import BYTES_PER_GB from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import ( LAYERWISE_OFFLOAD_ALL_COMPONENTS, LAYERWISE_OFFLOAD_DIT_GROUP, cpu_offload_flags_for_layerwise_components, layerwise_component_matches_any_selection, normalize_layerwise_offload_components, ) from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) from sglang.multimodal_gen.runtime.server_args.auto_tune import ( PERFORMANCE_MODES, ServerArgsAutoTuner, ) from sglang.multimodal_gen.runtime.server_args.disagg import DisaggServerArgsMixin from sglang.multimodal_gen.runtime.utils.common import ( is_port_available, is_valid_ipv6_address, normalize_gpu_ids, ) from sglang.multimodal_gen.runtime.utils.logging_utils import ( _sanitize_for_logging, configure_logger, init_logger, ) from sglang.multimodal_gen.utils import ( FlexibleArgumentParser, StoreBoolean, expand_path_fields, ) logger = init_logger(__name__) LTX2_TWO_STAGE_DEVICE_MODES = ("original", "resident") LTX2_TWO_STAGE_DEVICE_MODE_CHOICES = (*LTX2_TWO_STAGE_DEVICE_MODES, "snapshot") LTX2_TWO_STAGE_PIPELINE_NAMES = ("LTX2TwoStagePipeline", "LTX2TwoStageHQPipeline") # H200-class GPUs (>=130 GiB total) can usually keep both LTX2 DiTs resident. LTX2_RESIDENT_AUTO_ENABLE_MEM_GB = 130 LORA_MERGE_MODES = ("auto", "merge", "dynamic") def _normalize_ltx2_two_stage_device_mode(mode: str | None) -> str | None: if mode is None: return None mode = mode.lower() if mode == "snapshot": logger.warning( "ltx2_two_stage_device_mode=snapshot is deprecated and is treated " "as original. Please use ltx2_two_stage_device_mode=original or " "resident instead. This alias may be removed after two release cycles." ) return "original" return mode def is_ltx2_two_stage_pipeline_name(pipeline_class_name: str | None) -> bool: return pipeline_class_name in LTX2_TWO_STAGE_PIPELINE_NAMES class Backend(str, Enum): """ Enumeration for different model backends. - AUTO: Automatically select backend (prefer sglang native, fallback to diffusers) - SGLANG: Use sglang's native optimized implementation - DIFFUSERS: Use vanilla diffusers pipeline (supports all diffusers models) """ AUTO = "auto" SGLANG = "sglang" DIFFUSERS = "diffusers" @classmethod def from_string(cls, value: str) -> "Backend": """Convert string to Backend enum.""" try: return cls(value.lower()) except ValueError: raise ValueError( f"Invalid backend: {value}. Must be one of: {', '.join([m.value for m in cls])}" ) from None @classmethod def choices(cls) -> list[str]: """Get all available choices as strings for argparse.""" return [backend.value for backend in cls] WARMUP_MODES = ("off", "request", "server") # Default prompt sequence-length buckets for breakable CUDA graph (BCG) padding. # Prompt-conditioning is padded up to the smallest bucket that fits so prompts # of different lengths share one captured graph. DEFAULT_BCG_TEXT_BUCKETS = (64, 128, 256, 512, 1024) BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS = frozenset( { "comfy-org/ideogram-4", "glm-image", "ideogram-4", "ideogram-4-fp8", "ideogram-4-nf4", "ideogram-ai/ideogram-4-fp8", "ideogram-ai/ideogram-4-nf4", "qwen/qwen-image", "qwen/qwen-image-2512", "qwen-image", "qwen-image-2512", "tongyi-mai/z-image", "tongyi-mai/z-image-turbo", "zai-org/glm-image", "z-image", "z-image-turbo", } ) BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS = frozenset( { "GlmImagePipelineConfig", "Ideogram4PipelineConfig", "QwenImagePipelineConfig", "ZImagePipelineConfig", } ) def _normalized_bcg_model_refs(model_ref: str | None) -> set[str]: if not model_ref: return set() normalized = str(model_ref).strip().rstrip("/").lower() refs = {normalized, os.path.basename(normalized)} if "models--" in normalized: hf_cache_name = normalized.split("models--", 1)[1].split("/", 1)[0] refs.add(hf_cache_name.replace("--", "/")) return refs @dataclasses.dataclass class ServerArgs(DisaggServerArgsMixin): # Model and path configuration (for convenience) model_path: str # explicit model ID override (e.g. "Qwen-Image") model_id: str | None = None # Model backend (sglang native or diffusers) backend: Backend = Backend.AUTO # Attention attention_backend: str = None attention_backend_config: addict.Dict | None = None component_attention_backends: dict[str, str] | str | None = field( default_factory=dict ) cache_dit_config: str | dict[str, Any] | None = ( None # cache-dit config for diffusers ) # Distributed executor backend nccl_port: Optional[int] = None # HuggingFace specific parameters trust_remote_code: bool = False revision: str | None = None # Parallelism num_gpus: int = 1 performance_mode: str = "auto" base_gpu_id: int = 0 gpu_ids: list[int] | None = None tp_size: Optional[int] = None sp_degree: Optional[int] = None # sequence parallelism ulysses_degree: Optional[int] = None ring_degree: Optional[int] = None # data parallelism # number of data parallelism groups dp_size: int = 1 # number of gpu in a dp group dp_degree: int = 1 # cfg parallel (None = auto-decide based on num_gpus) enable_cfg_parallel: Optional[bool] = None # number of GPUs in each CFG parallel group (None = auto, 1 = disabled, N > 1 = enabled) cfg_parallel_degree: Optional[int] = None hsdp_replicate_dim: int = 1 hsdp_shard_dim: Optional[int] = None dist_timeout: int | None = 3600 # 1 hour pipeline_config: PipelineConfig = field(default_factory=PipelineConfig, repr=False) # Pipeline override pipeline_class_name: str | None = ( None # Override pipeline class from model_index.json ) # LoRA parameters # (Wenxuan) prefer to keep it here instead of in pipeline config to not make it complicated. lora_path: str | None = None lora_nickname: str = "default" # for swapping adapters in the pipeline lora_scale: float = 1.0 # LoRA scale for merging (e.g., 0.125 for Hyper-SD) lora_merge_mode: str = "auto" lora_weight_name: str | None = None # Component path overrides (key = model_index.json component name, value = path) component_paths: dict[str, str] = field(default_factory=dict) # path to pre-quantized transformer weights (single .safetensors or directory). transformer_weights_path: str | None = None # Per-component transformer weight overrides (key = model_index.json component name). # Pipelines use this when a checkpoint ships separate quantized weights for # secondary DiT components; the generic loader consumes it without model-specific # filename logic. component_transformer_weights_paths: dict[str, str] = field(default_factory=dict) # Quantization method for online quantization quantization: str | None = None # Layer name patterns to skip during online quantization quantization_ignored_layers: list[str] | None = None # can restrict layers to adapt, e.g. ["q_proj"] # Will adapt only q, k, v, o by default. lora_target_modules: list[str] | None = None # CPU offload parameters dit_cpu_offload: bool | None = None # if true, select the DiT layerwise group dit_layerwise_offload: bool | None = None layerwise_offload_components: list[str] | None = None dit_offload_prefetch_size: float = 0.0 offload_during_compile: bool = True text_encoder_cpu_offload: bool | None = None image_encoder_cpu_offload: bool | None = None vae_cpu_offload: bool | None = False use_fsdp_inference: bool | None = None pin_cpu_memory: bool = True ltx2_two_stage_device_mode: str | None = None _explicit_arg_names: set[str] = field(default_factory=set, repr=False) # ComfyUI integration comfyui_mode: bool = False # Compilation enable_torch_compile: bool = False # Breakable CUDA graph (BCG): capture the DiT forward as CUDA-graph # segments split at attention modules (SP all-to-all / dynamic attention # stay eager). Mutually exclusive with --enable-torch-compile and # Cache-DiT; BCG takes priority when more than one is requested. # # BCG graphs are resolution-specific, so --warmup-resolutions is required # when BCG is enabled: every requested resolution is captured at warmup so # serving never triggers a fresh capture. enable_breakable_cuda_graph: bool = False # Text/prompt sequence-length padding budget for BCG. Prompt-conditioning # inputs are padded up to the smallest bucket that fits, so prompts of # different lengths reuse one captured graph. Warmup captures one graph per # bucket; a prompt longer than the largest bucket falls back to eager. # ``None`` resolves to DEFAULT_BCG_TEXT_BUCKETS. bcg_text_buckets: list[int] = None # NVTX profiling enable_layerwise_nvtx_marker: bool = False # warmup # `warmup_mode` is the canonical knob: one of WARMUP_MODES # - "off": no warmup. # - "server": server-based warmup — a synthetic request right after the # server is ready, before real traffic # - "request": request-based warmup — warm on the first real request(s). # This is a BENCHMARK aid # existing consumers keep working) and as deprecated CLI aliases. None means # "derive the mode from the legacy booleans"; _adjust_warmup resolves it. warmup_mode: str | None = None # deprecated: warmup and server_warmup warmup: bool = False server_warmup: bool = False warmup_resolutions: list[str] = None warmup_steps: int = 1 disable_autocast: bool | None = None # Explicit quantization method override (e.g. "mxfp8", "fp8", "modelslim"). # When set, the transformer loader will use this instead of auto-detection. quantization: str | None = None # Quantization / Nunchaku SVDQuant configuration nunchaku_config: NunchakuSVDQuantArgs | NunchakuConfig | None = field( default_factory=NunchakuSVDQuantArgs, repr=False ) # Master port for distributed inference master_port: int = 30005 # http server endpoint config host: str | None = "127.0.0.1" port: int | None = 30000 # TODO: webui and their endpoint, check if webui_port is available. webui: bool = False webui_port: int | None = 12312 scheduler_port: int = 5555 batching_mode: str = "dynamic" batching_max_size: int = 1 batching_delay_ms: float = 0.0 batching_config: str | None = None enable_batching_metrics: bool = False # Strict port mode: fail if requested port is unavailable instead of auto-selecting strict_ports: bool = False output_path: str | None = "outputs/" input_save_path: str | None = "inputs/uploads" # Prompt text file for batch processing prompt_file_path: str | None = None # model paths for correct deallocation model_paths: dict[str, str] = field(default_factory=dict) model_loaded: dict[str, bool] = field( default_factory=lambda: { "transformer": True, "vae": True, "video_vae": True, "audio_vae": True, "video_dit": True, "audio_dit": True, "dual_tower_bridge": True, } ) # # DMD parameters # dmd_denoising_steps: List[int] | None = field(default=None) # MoE parameters used by Wan2.2 boundary_ratio: float | None = None # Disaggregation (pool mode only — launched via launch_pool_disagg_server()) disagg_role: RoleType = RoleType.MONOLITHIC disagg_timeout: int = 3600 disagg_downstream_wait_timeout: int = 1800 disagg_dispatch_policy: str = "round_robin" disagg_mode: bool = False disagg_instance_id: int = 0 disagg_max_slots_per_instance: int = 8 disagg_transfer_redundancy: float = 1.25 disagg_role_device: Literal["auto", "cpu", "cuda"] = "auto" disagg_transfer_backend: Literal["auto", "mock", "mooncake"] = "auto" disagg_transfer_pool_size: int = 256 * 1024 * 1024 disagg_transfer_pin_memory: Literal["auto", "off", "required"] = "auto" disagg_p2p_hostname: str = "127.0.0.1" disagg_ib_device: str | None = None disagg_server_addr: str | None = None encoder_urls: str | None = None denoiser_urls: str | None = None decoder_urls: str | None = None encoder_tp: int | None = None denoiser_tp: int | None = None denoiser_sp: int | None = None denoiser_ulysses: int | None = None denoiser_ring: int | None = None decoder_sp: int | None = None decoder_tp: int | None = None pool_work_endpoint: str | None = None pool_result_endpoint: str | None = None pool_control_endpoint: str | None = None pool_control_advertised_endpoint: str | None = None # Logging log_level: str = "info" log_requests: bool = False log_requests_level: int = 2 log_requests_format: str = "text" log_requests_target: Optional[List[str]] = None uvicorn_access_log_exclude_prefixes: list[str] = field(default_factory=list) # Tracing enable_trace: bool = False otlp_traces_endpoint: str = "localhost:4317" # SGLang backend for encoder stage srt_encoder_url: str | None = None srt_encoder_connect_timeout: int = 3.05 srt_encoder_timeout: int = 100 @property def broker_port(self) -> int: return self.port + 1 @property def is_local_mode(self) -> bool: """ If no server is running when a generation task begins, 'local_mode' will be enabled: a dedicated server will be launched """ return self.host is None or self.port is None def _adjust_path(self): expand_path_fields(self) self._adjust_save_paths() def _adjust_parameters(self): """set defaults and normalize values.""" auto_tuner = ServerArgsAutoTuner(self) auto_tuner.adjust_based_on_performance_mode() self._adjust_disagg_parallelism_aliases() if auto_tuner.could_override_server_args(): self._adjust_offload() auto_tuner.maybe_adjust_auto_default_layerwise_offload() self._adjust_ltx2_two_stage_device_mode() if auto_tuner.could_override_server_args(): auto_tuner.maybe_adjust_auto_component_residency_after_offload() auto_tuner.maybe_adjust_auto_fsdp_with_offload_enabled() auto_tuner.maybe_replace_cpu_offloaded_components_with_layerwise() self._adjust_path() self._adjust_quant_config() self._adjust_breakable_cuda_graph_support() self._adjust_warmup() self._adjust_network_ports() # adjust parallelism before attention backend self._adjust_parallelism() self._adjust_attention_backend() self._adjust_platform_specific() self._adjust_layerwise_offload_components() self._adjust_autocast() auto_tuner.finalize_auto_flags() self.adjust_pipeline_config() def _adjust_disagg_parallelism_aliases(self): if self.decoder_tp is None: return if self.decoder_sp is not None and self.decoder_sp != self.decoder_tp: raise ValueError( "decoder_tp is deprecated in favor of decoder_sp; " "please set only one of them or keep the same value." ) if self.decoder_sp is None: logger.warning( "decoder_tp is deprecated and is treated as decoder_sp for " "decoder/VAE parallel decode. Please use decoder_sp instead." ) self.decoder_sp = self.decoder_tp def _validate_parameters(self): """check consistency and raise errors for invalid configs""" self._validate_pipeline() self._validate_offload() if not current_platform.is_cpu(): self._validate_parallelism() self._validate_cfg_parallel() self._validate_batching() self._validate_breakable_cuda_graph() def resolved_bcg_text_buckets(self) -> tuple[int, ...]: """Sorted, de-duplicated, positive BCG text buckets. Falls back to :data:`DEFAULT_BCG_TEXT_BUCKETS` when ``--bcg-text-buckets`` is unset, so both prompt padding and warmup capture share one source of truth instead of the legacy ``SGLANG_BCG_TEXT_BUCKETS`` env var. """ raw = self.bcg_text_buckets if not raw: return DEFAULT_BCG_TEXT_BUCKETS buckets = sorted({int(b) for b in raw if int(b) > 0}) return tuple(buckets) or DEFAULT_BCG_TEXT_BUCKETS def _validate_breakable_cuda_graph(self): if not self.enable_breakable_cuda_graph: return # BCG graphs are captured per resolution and only replay for that exact # latent shape, so the user must declare the resolutions up front. We # capture every one of them at warmup; serving then never re-captures. if not self.warmup_resolutions: raise ValueError( "--enable-breakable-cuda-graph requires --warmup-resolutions: " "diffusion CUDA graphs only replay for a fixed resolution, so " "every served resolution must be declared and captured at " "warmup, e.g. --warmup-resolutions 1024x1024 1328x1328." ) if self.bcg_text_buckets is not None and not any( int(b) > 0 for b in self.bcg_text_buckets ): raise ValueError( "--bcg-text-buckets must contain at least one positive integer." ) def _adjust_breakable_cuda_graph_support(self): if not self.enable_breakable_cuda_graph: return pipeline_config = getattr(self, "pipeline_config", None) pipeline_config_name = type(pipeline_config).__name__ if ( pipeline_config_name in BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS and self._is_breakable_cuda_graph_supported_model() ): return logger.warning( "[Diffusion BCG] disabled for %s: only Ideogram-4, Qwen/Qwen-Image, " "Qwen/Qwen-Image-2512, Tongyi-MAI/Z-Image/Z-Image-Turbo, " "and zai-org/GLM-Image are currently supported.", pipeline_config_name, ) self.enable_breakable_cuda_graph = False def _is_breakable_cuda_graph_supported_model(self) -> bool: refs = _normalized_bcg_model_refs(self.model_id) refs.update(_normalized_bcg_model_refs(self.model_path)) return bool(refs & BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS) def _adjust_save_paths(self): """Normalize empty-string save paths to None (disabled).""" if self.output_path is not None and self.output_path.strip() == "": self.output_path = None if self.input_save_path is not None and self.input_save_path.strip() == "": self.input_save_path = None def _adjust_quant_config(self): """ resolve, validate and adjust quantization config handles only nunchaku for now """ ncfg = self.nunchaku_config if ncfg is None or isinstance(ncfg, NunchakuConfig): return resolution = ncfg.resolve_runtime_config() if resolution.transformer_weights_path: self.transformer_weights_path = resolution.transformer_weights_path self.nunchaku_config = resolution.nunchaku_config def adjust_pipeline_config(self): # 1. adjust for encoder parallel folding tp_size = self.tp_size or 1 dp_size = self.dp_size or 1 sp_degree = self.sp_degree or 1 # one replica = all its GPUs replica_size = (self.num_gpus or tp_size) // dp_size fold_world = dp_size == 1 and not self.disagg_mode and replica_size > tp_size if fold_world: mode = "world" elif tp_size == 1 and sp_degree > 1: # Preserve prior behavior for dp>1 / disaggregated SP runs. mode = "sp" else: return # Propose the fold group from the parallelism for every encoder. The # loader keeps it only for encoders wide enough to benefit at their real # (post-load) size and whose dims divide the group -- see # finalize_encoder_folding. Deciding on real size (not architecture) # handles the same encoder family at different parameter counts. encoder_configs = list(self.pipeline_config.text_encoder_configs) + list( getattr(self.pipeline_config, "image_encoder_configs", ()) or () ) for encoder_config in encoder_configs: encoder_config.parallel_folding_mode = mode logger.info( "Proposed encoder parallel folding (mode=%s) for %s " "(tp=%s sp=%s cfg=%s replica=%s); the loader keeps it for encoders " "wide enough to benefit.", mode, self.__class__.__name__, tp_size, sp_degree, self.cfg_parallel_degree or 1, replica_size, ) def _adjust_offload(self): if current_platform.is_cpu(): # CPU platform does not need offload return if self.pipeline_config.task_type.is_action_gen(): if self.dit_cpu_offload is None: self.dit_cpu_offload = False if self.text_encoder_cpu_offload is None: self.text_encoder_cpu_offload = False if self.image_encoder_cpu_offload is None: self.image_encoder_cpu_offload = False if self.vae_cpu_offload is None: self.vae_cpu_offload = False return # TODO: to be handled by each platform if current_platform.get_device_total_memory() / BYTES_PER_GB < 30: logger.info( "Enabling large component offloading for GPU with low device memory" ) if self.dit_cpu_offload is None: self.dit_cpu_offload = True if self.text_encoder_cpu_offload is None: self.text_encoder_cpu_offload = True if self.image_encoder_cpu_offload is None: self.image_encoder_cpu_offload = True elif self.pipeline_config.task_type.is_image_gen(): if self.dit_cpu_offload is None: self.dit_cpu_offload = True if self.text_encoder_cpu_offload is None: self.text_encoder_cpu_offload = True if self.image_encoder_cpu_offload is None: self.image_encoder_cpu_offload = False else: if self.dit_cpu_offload is None: self.dit_cpu_offload = True if self.text_encoder_cpu_offload is None: self.text_encoder_cpu_offload = True if self.image_encoder_cpu_offload is None: self.image_encoder_cpu_offload = True def _adjust_ltx2_two_stage_device_mode(self): if not self._is_ltx23_two_stage_pipeline(): return mode = self.ltx2_two_stage_device_mode if mode is None: env_mode = os.getenv("SGLANG_LTX2_TWO_STAGE_DEVICE_MODE") mode = ( _normalize_ltx2_two_stage_device_mode(env_mode) if env_mode else self._resolve_default_ltx2_two_stage_device_mode() ) else: mode = _normalize_ltx2_two_stage_device_mode(mode) if mode not in LTX2_TWO_STAGE_DEVICE_MODES: raise ValueError( f"Invalid ltx2_two_stage_device_mode={mode!r}. " f"Expected one of {LTX2_TWO_STAGE_DEVICE_MODE_CHOICES}." ) self.ltx2_two_stage_device_mode = mode def _resolve_default_ltx2_two_stage_device_mode(self) -> str: if not current_platform.is_cuda(): logger.info( "Automatically set ltx2_two_stage_device_mode=original on non-CUDA platform" ) return "original" device_name = str(current_platform.get_device_name(0)).upper() device_total_memory_gb = ( current_platform.get_device_total_memory() / BYTES_PER_GB ) if ( "H200" in device_name or device_total_memory_gb >= LTX2_RESIDENT_AUTO_ENABLE_MEM_GB ): logger.info( "Automatically set ltx2_two_stage_device_mode=resident for high-memory CUDA GPU (%s, %.2f GiB total)", device_name, device_total_memory_gb, ) return "resident" logger.info( "Automatically set ltx2_two_stage_device_mode=original for CUDA GPU (%s, %.2f GiB total)", device_name, device_total_memory_gb, ) return "original" def _is_ltx23_two_stage_pipeline(self) -> bool: return is_ltx2_two_stage_pipeline_name(self.pipeline_class_name) and ( self._is_ltx23_model_path(self.model_path) or is_ltx23_native_variant(self.pipeline_config.vae_config.arch_config) ) def _uses_ltx23_high_memory_resident_two_stage_mode(self) -> bool: if ( self.ltx2_two_stage_device_mode != "resident" or not self._is_ltx23_two_stage_pipeline() or not current_platform.is_cuda() ): return False return ( current_platform.get_device_total_memory() / BYTES_PER_GB >= LTX2_RESIDENT_AUTO_ENABLE_MEM_GB ) def _adjust_attention_backend(self): if self.attention_backend in ["fa3", "fa4"]: self.attention_backend = "fa" self.component_attention_backends = ( self._normalize_component_attention_backends( self.component_attention_backends ) ) # attention_backend_config if self.attention_backend_config is None: self.attention_backend_config = addict.Dict() elif isinstance(self.attention_backend_config, str): self.attention_backend_config = addict.Dict( self._parse_attention_backend_config(self.attention_backend_config) ) if self.backend != Backend.DIFFUSERS and isinstance( self.pipeline_config, LTX2PipelineConfig ): text_backend = self.component_attention_backends.get("text_encoder") if text_backend != "torch_sdpa": if text_backend is None: logger.info( "Automatically set torch_sdpa backend for component text_encoder to preserve LTX2 official attention semantics" ) else: logger.warning( "Overriding %s backend with torch_sdpa for component text_encoder to preserve LTX2 official attention semantics", text_backend, ) self.component_attention_backends["text_encoder"] = "torch_sdpa" if self.ring_degree > 1: if self.attention_backend is not None and self.attention_backend not in ( "fa", "sage_attn", ): raise ValueError( "Ring Attention is only supported for flash attention or sage attention backend for now" ) if self.attention_backend is None: self.attention_backend = "fa" logger.info( "Ring Attention is currently only supported for flash attention or sage attention; " "attention_backend has been automatically set to flash attention" ) if self.attention_backend is None and self.backend != Backend.DIFFUSERS: if ( current_platform.is_cuda() and self.pipeline_class_name is None and self.num_gpus == 1 and self.tp_size == 1 and self.sp_degree == 1 and self.ulysses_degree == 1 and self.ring_degree == 1 and self._is_ltx23_model_path(self.model_path) ): self.attention_backend = "fa" logger.info( "Automatically set attention_backend=fa for LTX-2.3 one-stage on 1 GPU to preserve precision" ) return self._set_default_attention_backend() @staticmethod def _normalize_attention_backend_name(backend: str) -> str: if not isinstance(backend, str): raise ValueError("Attention backend name must be a string") normalized = backend.strip().lower() if normalized in ("fa3", "fa4"): normalized = "fa" try: return AttentionBackendEnum[normalized.upper()].name.lower() except KeyError: raise ValueError( f"Invalid attention backend '{backend}'. " f"Available options are: {[e.name.lower() for e in AttentionBackendEnum]}" ) from None @staticmethod def _parse_component_attention_backend_map( value: dict[str, str] | str | None, ) -> dict[str, str]: if value is None or value == "": return {} if isinstance(value, dict): return dict(value) if not isinstance(value, str): raise ValueError( "component_attention_backends must be a dict or a comma-separated component=backend string" ) try: parsed = json.loads(value) if not isinstance(parsed, dict): raise ValueError return parsed except (json.JSONDecodeError, ValueError): pass result: dict[str, str] = {} for pair in value.split(","): pair = pair.strip() if not pair: continue if "=" not in pair: raise ValueError( "component_attention_backends must use component=backend entries" ) component, backend = pair.split("=", 1) result[component.strip()] = backend.strip() return result @classmethod def _normalize_component_attention_backends( cls, value: dict[str, str] | str | None ) -> dict[str, str]: raw = cls._parse_component_attention_backend_map(value) normalized: dict[str, str] = {} for component, backend in raw.items(): if not isinstance(component, str): raise ValueError("Component attention backend key must be a string") component_name = component.strip().replace("-", "_") if not component_name: raise ValueError("Component attention backend key must not be empty") normalized[component_name] = cls._normalize_attention_backend_name(backend) return normalized def resolve_component_attention_backend( self, *component_names: str | None ) -> tuple[AttentionBackendEnum | None, str | None]: for component_name in component_names: if component_name is None: continue key = component_name.replace("-", "_") fallback_keys = [key] if key.endswith("_2"): # Secondary two-stage components inherit the base component # backend unless explicitly overridden. fallback_keys.append(key[:-2]) for backend_key in fallback_keys: backend = self.component_attention_backends.get(backend_key) if backend is not None: return AttentionBackendEnum[backend.upper()], backend_key return None, None def _adjust_warmup(self): # --warmup-mode > --warmup/--server-warmup mode_explicit = self.is_arg_explicitly_set("warmup_mode") legacy_explicit = self.is_arg_explicitly_set( "warmup" ) or self.is_arg_explicitly_set("server_warmup") if self.warmup_mode is not None: if self.warmup_mode not in WARMUP_MODES: raise ValueError( f"Invalid --warmup-mode {self.warmup_mode!r}; " f"expected one of {WARMUP_MODES}." ) if mode_explicit and legacy_explicit: logger.warning( "Both --warmup-mode and the deprecated --warmup/--server-warmup " "were set; --warmup-mode=%s takes precedence.", self.warmup_mode, ) if mode_explicit or not legacy_explicit: self.warmup = self.warmup_mode != "off" self.server_warmup = self.warmup_mode == "server" elif self.warmup: self.server_warmup = self.server_warmup or self.warmup_mode == "server" # Explicit resolutions imply warmup is on (request-based). if self.warmup_resolutions is not None: self.warmup = True if ( self.enable_torch_compile and self.warmup_mode is None and not mode_explicit and not legacy_explicit ): self.warmup = True self.server_warmup = True logger.info( "Automatically enabled server warmup for torch.compile so first " "real requests do not pay compile latency. Set --warmup-mode off " "to disable this behavior." ) # BCG captures every graph during a synthetic warmup forward at startup # so that serving never records a fresh graph. That requires # server-based warmup (a real warmup request issued at startup), not # request-based warmup which runs no forward until the first request. if self.enable_breakable_cuda_graph and self.disagg_role == RoleType.MONOLITHIC: self.warmup = True self.server_warmup = True if self.disagg_role != RoleType.MONOLITHIC: self.server_warmup = False if not self.warmup: self.server_warmup = False self.warmup_mode = ( "off" if not self.warmup else "server" if self.server_warmup else "request" ) @staticmethod def _require_port(port: int, name: str) -> None: """Raise if *port* is occupied (used under ``--strict-ports``).""" if not is_port_available(port): raise RuntimeError( f"{name} port {port} is unavailable and --strict-ports is enabled. " f"Either use a different port or disable --strict-ports." ) def _adjust_network_ports(self): # Disagg role instances (encoder/denoiser/decoder) don't serve HTTP, # so skip settling the HTTP port to avoid unnecessary port collisions. needs_http = self.disagg_role in ( RoleType.MONOLITHIC, RoleType.SERVER, ) if self.strict_ports: requested_ports = [] if needs_http: requested_ports.append((self.port, "HTTP")) requested_ports.append((self.scheduler_port, "Scheduler")) if self.master_port is not None: requested_ports.append((self.master_port, "Master")) seen_ports: dict[int, str] = {} for port, name in requested_ports: if port in seen_ports: raise RuntimeError( f"{name} port {port} duplicates {seen_ports[port]} port and " "--strict-ports is enabled." ) seen_ports[port] = name self._require_port(port, name) else: settled_ports: set[int] = set() if needs_http: self.port = self.settle_port(self.port) settled_ports.add(self.port) initial_scheduler_port = self.scheduler_port + ( random.randint(0, 100) if self.scheduler_port == 5555 else 0 ) self.scheduler_port = self.settle_port( initial_scheduler_port, avoid=settled_ports ) settled_ports.add(self.scheduler_port) if self.master_port is not None: self.master_port = self.settle_port( self.master_port, 37, avoid=settled_ports ) def _adjust_parallelism(self): sp_unspecified = self.sp_degree is None ulysses_unspecified = self.ulysses_degree is None ring_unspecified = self.ring_degree is None cfg_unspecified = self.enable_cfg_parallel is None if self.tp_size is None: self.tp_size = 1 if current_platform.is_cpu() and self.tp_size > 1: # CPU platform reuse num_gpus to represent num cpu numa nodes as devices self.num_gpus = self.tp_size if self.hsdp_shard_dim is None: self.hsdp_shard_dim = self.num_gpus # --cfg-parallel-size takes precedence over --enable-cfg-parallel bool. if self.cfg_parallel_degree is not None: if self.cfg_parallel_degree == 1: self.enable_cfg_parallel = False elif self.cfg_parallel_degree > 1: self.enable_cfg_parallel = True cfg_unspecified = False # Auto-enable CFG parallel when user hasn't set any parallelism flags # and there are enough GPUs. Only auto-enable for models whose default # SamplingParams use classifier-free guidance (negative_prompt is not None), # because non-CFG models (e.g. FLUX) crash when CFG parallel splits ranks. if cfg_unspecified: deployment_config = self.pipeline_config.get_model_deployment_config() auto_cfg_parallel_degree = deployment_config.get_auto_cfg_parallel_degree( self.num_gpus ) if auto_cfg_parallel_degree < 1: self.enable_cfg_parallel = False else: cfg_group_size = self.dp_size * self.tp_size * auto_cfg_parallel_degree if ( self.performance_mode != "manual" and deployment_config.auto_enable_cfg_parallel and self.num_gpus >= 2 and self.num_gpus % cfg_group_size == 0 and sp_unspecified and ulysses_unspecified and ring_unspecified and self._model_default_uses_cfg() ): self.cfg_parallel_degree = auto_cfg_parallel_degree self.enable_cfg_parallel = auto_cfg_parallel_degree > 1 if self.enable_cfg_parallel: logger.info( "Automatically enabled CFG parallel at degree %d for %d GPUs. " "Use --sp-degree / --ulysses-degree to use sequence " "parallelism instead.", self.cfg_parallel_degree, self.num_gpus, ) else: logger.info( "Automatically disabled CFG parallel for %d GPUs based on model deployment config.", self.num_gpus, ) else: self.enable_cfg_parallel = False # Resolve cfg_parallel_degree to a concrete int now that enable_cfg_parallel is settled. if self.cfg_parallel_degree is None: self.cfg_parallel_degree = 2 if self.enable_cfg_parallel else 1 # adjust sp_degree: allocate all remaining GPUs after TP and DP if self.sp_degree is None: num_gpus_per_group = self.dp_size * self.tp_size if self.enable_cfg_parallel: num_gpus_per_group *= self.cfg_parallel_degree if self.num_gpus % num_gpus_per_group == 0: self.sp_degree = self.num_gpus // num_gpus_per_group else: # Will be validated later self.sp_degree = 1 if ( self.ulysses_degree is None and self.ring_degree is None and self.sp_degree != 1 ): self.ulysses_degree = self.sp_degree logger.info( f"Automatically set ulysses_degree=sp_degree={self.ulysses_degree} for best performance" ) if self.ulysses_degree is None: self.ulysses_degree = 1 logger.debug( f"Ulysses degree not set, using default value {self.ulysses_degree}" ) if self.ring_degree is None: self.ring_degree = 1 logger.debug(f"Ring degree not set, using default value {self.ring_degree}") def _model_default_uses_cfg(self) -> bool: """ Check whether the model uses classifier-free guidance by default. CFG is active when *both* ``negative_prompt is not None`` and ``guidance_scale > 1``. """ from sglang.multimodal_gen.registry import get_model_info model_info = get_model_info(self.model_path, self.backend, self.model_id) if model_info is None: return False default_params = model_info.sampling_param_cls() return ( getattr(default_params, "negative_prompt", None) is not None and getattr(default_params, "guidance_scale", 0) > 1.0 ) @staticmethod def _is_ltx23_model_path(model_path: str | None) -> bool: if not model_path: return False normalized = model_path.lower() return any( token in normalized for token in ( "lightricks/ltx-2.3", "models--lightricks--ltx-2.3", "lightricks__ltx-2.3", ) ) def _adjust_platform_specific(self): if current_platform.is_mps(): self.use_fsdp_inference = False self.dit_layerwise_offload = False self.layerwise_offload_components = None if ( self.dit_cpu_offload or self.text_encoder_cpu_offload or self.image_encoder_cpu_offload or self.vae_cpu_offload ): logger.warning( "Disabling component CPU offload on MPS because CPU-to-MPS " "module relocation can produce invalid diffusion outputs." ) self.dit_cpu_offload = False self.text_encoder_cpu_offload = False self.image_encoder_cpu_offload = False self.vae_cpu_offload = False def is_arg_explicitly_set(self, arg_name: str) -> bool: return arg_name in self._explicit_arg_names def should_configure_layerwise_offload_for_lazy_component( self, component_name: str ) -> bool: """Return whether a lazy-loaded component should try layerwise offload. Lazy components are loaded after the normal pipeline-wide configuration pass, so they should only attempt layerwise configuration when their component name is covered by the selected layerwise scope. """ component_names = normalize_layerwise_offload_components( self.layerwise_offload_components ) if not component_names: return False if LAYERWISE_OFFLOAD_ALL_COMPONENTS in component_names: return True return layerwise_component_matches_any_selection( component_name, component_names ) @property def is_dit_layerwise_offload_selected(self) -> bool: """returns if dit is selected to be layerwise-offload""" component_names = self.layerwise_offload_components return bool( component_names and "dit_cpu_offload" in cpu_offload_flags_for_layerwise_components(component_names) ) def _adjust_layerwise_offload_components(self): explicitly_set_component_names = normalize_layerwise_offload_components( self.layerwise_offload_components ) if self.dit_layerwise_offload: if explicitly_set_component_names is None: explicitly_set_component_names = [LAYERWISE_OFFLOAD_DIT_GROUP] elif LAYERWISE_OFFLOAD_DIT_GROUP not in explicitly_set_component_names: explicitly_set_component_names = [ LAYERWISE_OFFLOAD_DIT_GROUP, *explicitly_set_component_names, ] if explicitly_set_component_names is not None: self.layerwise_offload_components = explicitly_set_component_names self._disable_non_dit_cpu_offload_for_layerwise_components( explicitly_set_component_names ) return def _disable_non_dit_cpu_offload_for_layerwise_components( self, component_names: list[str] ) -> None: # non-DiT layerwise offload replaces the corresponding component-level CPU offload flag_names = cpu_offload_flags_for_layerwise_components(component_names) disabled_flag_names: list[str] = [] if ( "text_encoder_cpu_offload" in flag_names and self.text_encoder_cpu_offload is not False ): self.text_encoder_cpu_offload = False disabled_flag_names.append("text_encoder_cpu_offload") if ( "image_encoder_cpu_offload" in flag_names and self.image_encoder_cpu_offload is not False ): self.image_encoder_cpu_offload = False disabled_flag_names.append("image_encoder_cpu_offload") if "vae_cpu_offload" in flag_names and self.vae_cpu_offload is not False: self.vae_cpu_offload = False disabled_flag_names.append("vae_cpu_offload") explicit_disabled_flag_names = [ flag_name for flag_name in disabled_flag_names if self.is_arg_explicitly_set(flag_name) ] if explicit_disabled_flag_names: logger.info( "Ignoring explicit CPU-offload flags because layerwise offload " "manages the same component weights: %s", ", ".join( f"{flag_name}=False" for flag_name in explicit_disabled_flag_names ), ) def _adjust_autocast(self): if self.disable_autocast is None: self.disable_autocast = not self.pipeline_config.enable_autocast def _parse_attention_backend_config(self, config_str: str) -> dict[str, Any]: """parse attention backend config from string.""" if not config_str: return {} # 1. treat as file path if os.path.exists(config_str): if config_str.endswith((".yaml", ".yml")): with open(config_str, "r") as f: return yaml.safe_load(f) elif config_str.endswith(".json"): with open(config_str, "r") as f: return json.load(f) # 2. treat as JSON string try: return json.loads(config_str) except json.JSONDecodeError: pass # 3. treat as k=v pairs (simple implementation). e.g., "sparsity=0.5,enable_x=true" try: config = {} pairs = config_str.split(",") for pair in pairs: k, v = pair.split("=", 1) k = k.strip() v = v.strip() if v.lower() == "true": v = True elif v.lower() == "false": v = False elif v.replace(".", "", 1).isdigit(): v = float(v) if "." in v else int(v) config[k] = v return config except Exception: raise ValueError(f"Could not parse attention backend config: {config_str}") def __post_init__(self): # configure logger before use configure_logger(server_args=self) # Convert string disagg_role to enum (from CLI/config) if isinstance(self.disagg_role, str): self.disagg_role = RoleType.from_string(self.disagg_role) self.gpu_ids = normalize_gpu_ids(self.gpu_ids) # 1. adjust parameters self._adjust_parameters() # 2. Validate parameters self._validate_parameters() # log clean server_args try: safe_args = _sanitize_for_logging(self, key_hint="server_args") logger.info("server_args: %s", json.dumps(safe_args, ensure_ascii=False)) except Exception: # Fallback to default repr if sanitization fails logger.info(f"server_args: {self}") @staticmethod def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: # Model and path configuration parser.add_argument( "--model-path", type=str, help="The path of the model weights. This can be a local folder or a Hugging Face repo ID.", ) parser.add_argument( "--model-id", type=str, default=ServerArgs.model_id, help=( "Override the model ID used for config resolution. " "Useful when --model-path is a local directory whose name does not match " "any registered HF repo name. Should be the repo name portion of the HF ID " "(e.g. 'Qwen-Image' for 'Qwen/Qwen-Image')." ), ) parser.add_argument( "--pipeline", "--pipeline-class-name", dest="pipeline_class_name", type=str, default=ServerArgs.pipeline_class_name, help=( "Advanced override for pipeline class selection from the model registry " "or model_index.json. Must match a registered pipeline_name." ), ) # attention parser.add_argument( "--attention-backend", type=str, default=None, help=( "The attention backend to use. For SGLang-native pipelines, use " "values like fa, torch_sdpa, sage_attn, etc. For diffusers pipelines, " "use diffusers attention backend names such as flash, _flash_3_hub, " "sage, or xformers." ), ) parser.add_argument( "--attention-backend-config", type=str, default=None, help="Configuration for the attention backend. Can be a JSON string, a path to a JSON/YAML file, or key=value pairs.", ) parser.add_argument( "--component-attention-backends", type=str, default=None, help=( "Per-component attention backend overrides for native pipelines. " "Use component names from model_index.json, e.g. " "'text_encoder=torch_sdpa,transformer=fa'." ), ) parser.add_argument( "--cache-dit-config", type=str, default=ServerArgs.cache_dit_config, help="Path to a Cache-DiT YAML/JSON config. Enables cache-dit for diffusers backend.", ) # HuggingFace specific parameters parser.add_argument( "--trust-remote-code", action=StoreBoolean, default=ServerArgs.trust_remote_code, help="Trust remote code when loading HuggingFace models", ) parser.add_argument( "--revision", type=str, default=ServerArgs.revision, help="The specific model version to use (can be a branch name, tag name, or commit id)", ) parser.add_argument( "--performance-mode", "--mode", type=str, choices=PERFORMANCE_MODES, default=ServerArgs.performance_mode, help=( "Preset for performance and memory defaults. " "'manual' keeps performance-related server args under explicit user control, no adjustment is made; " "'auto' keeps safe defaults and applies high-confidence FSDP/CFG improvements; " "'speed' favors GPU-resident execution for lower latency and higher throughput, and may OOM; " "'memory' favors lower GPU memory usage; " "Explicit offload/FSDP/parallelism flags take precedence." ), ) # Parallelism parser.add_argument( "--num-gpus", type=int, default=ServerArgs.num_gpus, help="The number of GPUs to use.", ) parser.add_argument( "--base-gpu-id", type=int, default=ServerArgs.base_gpu_id, help="The starting GPU ID for this instance. Used with --disagg-role " "to place role instances on specific GPUs without CUDA_VISIBLE_DEVICES.", ) parser.add_argument( "--gpu-ids", nargs="+", default=None, help=( "Physical GPU IDs for this instance, e.g. --gpu-ids 0 1 6 7 " "or --gpu-ids 0,1,6,7. Overrides --base-gpu-id for standalone " "disagg roles." ), ) parser.add_argument( "--tp-size", type=int, default=None, help="The tensor parallelism size. Defaults to 1 if not specified.", ) parser.add_argument( "--sp-degree", type=int, default=None, help="The sequence parallelism size. If not specified, will use all remaining GPUs after accounting for TP and DP.", ) parser.add_argument( "--ulysses-degree", type=int, default=ServerArgs.ulysses_degree, help="Ulysses sequence parallel degree. Used in attention layer.", ) parser.add_argument( "--ring-degree", type=int, default=ServerArgs.ring_degree, help="Ring sequence parallel degree. Used in attention layer.", ) parser.add_argument( "--enable-cfg-parallel", action=StoreBoolean, default=None, help="Enable cfg parallel at degree 2. Auto-enabled when num_gpus >= 2 and no SP flags are set. Use false to disable it explicitly.", ) parser.add_argument( "--cfg-parallel-size", dest="cfg_parallel_degree", type=int, default=None, help=( "Number of GPUs per CFG parallel group (1 = disabled, N > 1 = enabled at degree N). " "Supersedes --enable-cfg-parallel. Allows 4-branch CFG parallel (e.g., --cfg-parallel-size 4) " "for models with cond + neg + perturbed + modality branches." ), ) parser.add_argument( "--data-parallel-size", "--dp-size", "--dp", dest="dp_size", type=int, default=ServerArgs.dp_size, help="The data parallelism size.", ) parser.add_argument( "--hsdp-replicate-dim", type=int, default=ServerArgs.hsdp_replicate_dim, help="The data parallelism size.", ) parser.add_argument( "--hsdp-shard-dim", type=int, default=None, help="The data parallelism shards. Defaults to num_gpus if not specified.", ) parser.add_argument( "--dist-timeout", type=int, default=ServerArgs.dist_timeout, help="Timeout for torch.distributed operations in seconds. " "Increase this value if you encounter 'Connection closed by peer' errors after the service is idle. ", ) ServerArgs.add_disagg_cli_args(parser) # Prompt text file for batch processing parser.add_argument( "--prompt-file-path", type=str, default=ServerArgs.prompt_file_path, help="Path to a text file containing prompts (one per line) for batch processing", ) parser.add_argument( "--mask-strategy-file-path", type=str, help="Path to mask strategy JSON file for STA", ) parser.add_argument( "--enable-torch-compile", action=StoreBoolean, default=ServerArgs.enable_torch_compile, help="Use torch.compile to speed up diffusion hot paths. " + "When no warmup mode is configured, this enables server warmup " + "so first real requests do not pay compile latency. " + "However, will likely cause precision drifts. See (https://github.com/pytorch/pytorch/issues/145213)", ) parser.add_argument( "--offload-during-compile", action=StoreBoolean, default=ServerArgs.offload_during_compile, help="Offload components during the torch.compile warmup (the DiT layerwise) so max-autotune fits on tighter-memory GPUs, then restore the configured residency for serving. Skipped when the DiT is already layerwise-offloaded, or under cache-dit / FSDP.", ) parser.add_argument( "--enable-breakable-cuda-graph", action=StoreBoolean, default=ServerArgs.enable_breakable_cuda_graph, help="Capture the DiT forward as breakable CUDA graph segments " "(split at attention; SP all-to-all / dynamic attention stay " "eager) to cut per-kernel launch overhead. Mutually exclusive " "with --enable-torch-compile and Cache-DiT (BCG takes priority). " "Requires --warmup-resolutions; all of them are captured at warmup.", ) parser.add_argument( "--bcg-text-buckets", type=int, nargs="+", default=ServerArgs.bcg_text_buckets, help="Prompt sequence-length padding budget for breakable CUDA " "graph. Prompt-conditioning is padded up to the smallest bucket " "that fits so different prompt lengths reuse one captured graph; " "warmup captures one graph per bucket. Defaults to " f"{' '.join(map(str, DEFAULT_BCG_TEXT_BUCKETS))}. " "Replaces the legacy SGLANG_BCG_TEXT_BUCKETS env var.", ) parser.add_argument( "--enable-layerwise-nvtx-marker", action=StoreBoolean, default=ServerArgs.enable_layerwise_nvtx_marker, help="Enable layerwise NVTX markers for profiling with Nsight Systems. " "Adds NVTX ranges around each pipeline stage, the denoising loop, " "every denoising step, the predict_noise / scheduler_step " "sub-operations, and every transformer submodule forward (recursive). " "Warmup steps are excluded to keep captured traces clean.", ) # warmup parser.add_argument( "--warmup-mode", type=str, choices=list(WARMUP_MODES), default=ServerArgs.warmup_mode, help=( "Warmup mode (canonical knob). One of: " "`off` (no warmup); `request` (request-based: warm on real " "incoming requests); `server` (server-based: a synthetic warmup " "request right after the server is ready, before traffic). " "Takes precedence over the deprecated --warmup/--server-warmup. " "`sglang serve` defaults to `server`; other entrypoints default " "to request-based when warmup is enabled. When enabled, look for " "the line ending with `(with warmup excluded)` for actual " "processing time." ), ) parser.add_argument( "--warmup", action=StoreBoolean, default=ServerArgs.warmup, help=( "[DEPRECATED: use --warmup-mode] Perform warmup before normal " "traffic. Maps to --warmup-mode request (or server, combined " "with --server-warmup). Recommended when benchmarking for fair " "comparison and best performance." ), ) parser.add_argument( "--warmup-resolutions", type=str, nargs="+", default=ServerArgs.warmup_resolutions, help="Specify explicit warmup resolutions. e.g., `--warmup-resolutions 256x256 720x720`", ) parser.add_argument( "--warmup-steps", type=int, default=ServerArgs.warmup_steps, help="The number of warmup steps to perform for each resolution.", ) parser.add_argument( "--server-warmup", action=StoreBoolean, default=ServerArgs.server_warmup, help=( "[DEPRECATED: use --warmup-mode server] Send a synthetic warmup " "request after the server is ready (server-based warmup)." ), ) # layerwise offload parser.add_argument( "--dit-cpu-offload", action=StoreBoolean, help="Use CPU offload for DiT inference. Enable if run out of memory with FSDP.", ) parser.add_argument( "--dit-layerwise-offload", action=StoreBoolean, default=ServerArgs.dit_layerwise_offload, help="Enable layerwise CPU offload with async H2D prefetch overlap for DiTs. " "It selects only the DiT layerwise group. Cannot be used together with cache-dit " "(SGLANG_CACHE_DIT_ENABLED) or use_fsdp_inference. May be combined with " "--dit-cpu-offload, in which case DiT weights stay on host memory and only the " "layers needed for the current step are brought on-device (lowest peak GPU memory).", ) parser.add_argument( "--layerwise-offload-components", "--layerwise-offload-modules", type=str, nargs="+", default=ServerArgs.layerwise_offload_components, help="Select pipeline components for layerwise offload. " "Use dit to select the DiT layerwise group, default for the default group " "(currently text_encoder, image_encoder, and vae), " "or all to select every layerwise-offloadable component. " "This option does not imply --dit-layerwise-offload. Example: " "--layerwise-offload-components text_encoder image_encoder vae.", ) parser.add_argument( "--dit-offload-prefetch-size", type=float, default=ServerArgs.dit_offload_prefetch_size, help="The size of prefetch for dit-layerwise-offload. If the value is between 0.0 and 1.0, it is treated as a ratio of the total number of layers. If the value is >= 1, it is treated as the absolute number of layers. 0.0 means prefetch 1 layer (lowest memory). Values above 0.5 might have peak memory close to no offload but worse performance.", ) # offload flags parser.add_argument( "--text-encoder-cpu-offload", action=StoreBoolean, help="Use CPU offload for text encoder. Enable if run out of memory.", ) parser.add_argument( "--image-encoder-cpu-offload", action=StoreBoolean, help="Use CPU offload for image encoder. Enable if run out of memory.", ) parser.add_argument( "--vae-cpu-offload", action=StoreBoolean, help="Use CPU offload for VAE. Enable if run out of memory.", ) parser.add_argument( "--use-fsdp-inference", action=StoreBoolean, help="Use FSDP inference to shard DiT weights across GPUs. For single-GPU memory pressure, prefer CPU or layerwise offload.", ) parser.add_argument( "--pin-cpu-memory", action=StoreBoolean, help='Pin memory for CPU offload. Only added as a temp workaround if it throws "CUDA error: invalid argument". ' "Should be enabled in almost all cases", ) parser.add_argument( "--ltx2-two-stage-device-mode", type=str, choices=LTX2_TWO_STAGE_DEVICE_MODE_CHOICES, default=ServerArgs.ltx2_two_stage_device_mode, help=( "LTX-2.3 two-stage device residency mode: " "'original' keeps official two-stage semantics without premerged stage2, " "'resident' keeps both transformers resident on GPU. " "'snapshot' is deprecated, treated as 'original', and may be " "removed after two release cycles. " "Default is auto: resident on H200/high-memory CUDA GPUs, otherwise original." ), ) parser.add_argument( "--disable-autocast", action=StoreBoolean, help="Disable autocast for denoising loop and vae decoding in pipeline sampling", ) # quantization parser.add_argument( "--quantization", type=str, default=ServerArgs.quantization, help=( "Quantization method for the transformer. If omitted, the method is " "auto-detected from the checkpoint config or safetensors metadata when " "possible. Use this flag to override auto-detection. " "Online (post-load) quantization from a BF16/FP16 checkpoint " "is supported for 'fp8' and 'mxfp4'. Other methods " "('modelopt', 'modelopt_fp8', 'modelopt_fp4', 'mxfp8', " "'mxfp4_npu', 'modelslim') require a pre-quantized checkpoint. " "Note: 'mxfp4' targets ROCm + MI350+ (gfx95x); " "'mxfp4_npu' / 'mxfp8' target Ascend NPU (A5 series for mxfp4_npu)." ), ) parser.add_argument( "--quantization-ignored-layers", type=str, nargs="+", default=ServerArgs.quantization_ignored_layers, help=( "Layer name patterns to keep unquantized during online quantization " "(fp8/mxfp4). Each pattern is matched against the layer prefix. " "Example: --quantization-ignored-layers img_mod txt_mod to_out" ), ) # Nunchaku SVDQuant quantization parameters NunchakuSVDQuantArgs.add_cli_args(parser) # Master port for distributed inference parser.add_argument( "--master-port", type=int, default=ServerArgs.master_port, help="Master port for distributed inference. If not set, a random free port will be used.", ) parser.add_argument( "--scheduler-port", type=int, default=ServerArgs.scheduler_port, help="Port for the scheduler server.", ) parser.add_argument( "--batching-mode", type=str, default=ServerArgs.batching_mode, choices=["dynamic"], help="Request batching scheduler mode. Currently only 'dynamic' is implemented.", ) parser.add_argument( "--batching-max-size", type=int, default=ServerArgs.batching_max_size, help="Maximum number of compatible generation requests to merge into one batch.", ) parser.add_argument( "--batching-delay-ms", type=float, default=ServerArgs.batching_delay_ms, help="Maximum time (in ms) to wait for forming a larger batch before dispatch.", ) parser.add_argument( "--batching-config", type=str, default=ServerArgs.batching_config, help=( "Optional JSON file with {'schema_version': 1, 'rules': [...]} " "batching admission rules that can cap model/resolution shapes " "below --batching-max-size." ), ) parser.add_argument( "--enable-batching-metrics", action="store_true", default=ServerArgs.enable_batching_metrics, help="Log periodic batch efficiency metrics such as realized batch size and queue wait time.", ) parser.add_argument( "--host", type=str, default=ServerArgs.host, help="Host for the HTTP API server.", ) parser.add_argument( "--port", type=int, default=ServerArgs.port, help="Port for the HTTP API server.", ) parser.add_argument( "--strict-ports", action=StoreBoolean, default=ServerArgs.strict_ports, help="If enabled, fail when requested ports are unavailable instead of auto-selecting.", ) parser.add_argument( "--webui", action=StoreBoolean, default=ServerArgs.webui, help="Whether to use webui for better display", ) parser.add_argument( "--webui-port", type=int, default=ServerArgs.webui_port, help="Whether to use webui for better display", ) parser.add_argument( "--output-path", type=str, default=ServerArgs.output_path, help='Directory path to save generated images/videos. Set to "" to disable persistent saving.', ) parser.add_argument( "--input-save-path", type=str, default=ServerArgs.input_save_path, help='Directory path to save uploaded input images/videos. Set to "" to disable persistent saving.', ) # LoRA parser.add_argument( "--lora-path", type=str, default=ServerArgs.lora_path, help="The path to the LoRA adapter weights (can be local file path or HF hub id) to launch with", ) parser.add_argument( "--lora-nickname", type=str, default=ServerArgs.lora_nickname, help="The nickname for the LoRA adapter to launch with", ) parser.add_argument( "--lora-scale", type=float, default=ServerArgs.lora_scale, help="LoRA scale for merging (e.g., 0.125 for Hyper-SD). Same as lora_scale in Diffusers", ) parser.add_argument( "--lora-merge-mode", type=str, choices=LORA_MERGE_MODES, default=ServerArgs.lora_merge_mode, help=( "How LoRA is applied: auto keeps static merge for regular weights " "and uses dynamic LoRA for FSDP-sharded weights to avoid full-gather; " "merge always merges into base weights; dynamic always applies LoRA at forward time." ), ) parser.add_argument( "--lora-weight-name", type=str, default=ServerArgs.lora_weight_name, help="Specific safetensors filename to load from a multi-file LoRA repo", ) # Add pipeline configuration arguments PipelineConfig.add_cli_args(parser) # Logging parser.add_argument( "--log-level", type=str, default=ServerArgs.log_level, help="The logging level of all loggers.", ) # Tracing parser.add_argument( "--enable-trace", action="store_true", default=False, help="Enable OpenTelemetry tracing.", ) parser.add_argument( "--otlp-traces-endpoint", type=str, default=ServerArgs.otlp_traces_endpoint, help="OTLP collector endpoint when --enable-trace is set. Format: :", ) parser.add_argument( "--log-requests", action="store_true", help="Log user-facing fields of all requests (default: False). " "Verbosity is controlled by --log-requests-level.", ) parser.add_argument( "--log-requests-level", type=int, default=ServerArgs.log_requests_level, choices=[0, 1, 2, 3], help="Verbosity level for request logging. " "0: Log request metadata only (request_id). " "1: Log metadata + sampling config (seed, steps, guidance, resolution, frames, fps, ...). " "2: Log metadata + sampling config + prompt/negative prompt (truncated to 2 KiB). " "3: Log metadata + sampling config + full prompt/negative prompt.", ) parser.add_argument( "--log-requests-format", type=str, default=ServerArgs.log_requests_format, choices=["text", "json"], help="Format for request logging: 'text' (human-readable) or 'json' (structured)", ) parser.add_argument( "--log-requests-target", type=str, nargs="+", default=ServerArgs.log_requests_target, help="Target(s) for request logging: 'stdout' and/or directory path(s) for file output. " "Can specify multiple targets, e.g., '--log-requests-target stdout /my/path'. ", ) parser.add_argument( "--uvicorn-access-log-exclude-prefixes", type=str, nargs="*", default=[], help="Exclude uvicorn access logs whose request path starts with any of these prefixes. " "Defaults to empty (disabled). " "Example: --uvicorn-access-log-exclude-prefixes /metrics /health", ) parser.add_argument( "--backend", type=str, choices=Backend.choices(), default=ServerArgs.backend.value, help="The model backend to use. 'auto' prefers sglang native and falls back to diffusers. " "'sglang' uses native optimized implementation. 'diffusers' uses vanilla diffusers pipeline.", ) # SGLang backend for encoder stage parser.add_argument( "--srt-encoder-url", type=str, default=ServerArgs.srt_encoder_url, help="Url of SGLang server for encoder stage", ) parser.add_argument( "--srt-encoder-connection-timeout", type=int, default=ServerArgs.srt_encoder_connect_timeout, help="Timeout (in seconds) for establishing the initial TCP connection to the SGLang encoder server. " "Default value is 3.05.", ) parser.add_argument( "--srt-encoder-timeout", type=int, default=ServerArgs.srt_encoder_timeout, help="Timeout (in seconds) for HTTP requests to the SGLang encoder server. " "Increase value if connection between diffusion server and AR model server is slow.", ) return parser def url(self): host = self.host if not host or host == "0.0.0.0": host = "127.0.0.1" elif host == "::": host = "::1" if is_valid_ipv6_address(host): return f"http://[{host}]:{self.port}" else: return f"http://{host}:{self.port}" @property def scheduler_endpoint(self): """ Internal endpoint for scheduler. Prefers the configured host but normalizes localhost -> 127.0.0.1 to avoid ZMQ issues. """ scheduler_host = self.host if scheduler_host is None or scheduler_host == "localhost": scheduler_host = "127.0.0.1" return f"tcp://{scheduler_host}:{self.scheduler_port}" def settle_port( self, port: int, port_inc: int = 42, max_attempts: int = 100, avoid: set[int] | None = None, ) -> int: """ Find an available port with retry logic. """ attempts = 0 original_port = port avoid = avoid or set() while attempts < max_attempts: if port not in avoid and is_port_available(port): if attempts > 0: logger.info( f"Port {original_port} was unavailable, using port {port} instead" ) return port attempts += 1 if port < 60000: port += port_inc else: # Wrap around with randomization to avoid collision port = 5000 + random.randint(0, 1000) raise RuntimeError( f"Failed to find available port after {max_attempts} attempts " f"(started from port {original_port})" ) @staticmethod def _extract_component_paths( unknown_args: list[str], ) -> tuple[dict[str, str], list[str]]: """ Extract dynamic component path args from unrecognised CLI args. Supported forms: - ``---path /path/to/component`` - ``--component-paths. /path/to/component`` (expanded from config) """ component_paths: dict[str, str] = {} remaining: list[str] = [] i = 0 while i < len(unknown_args): arg = unknown_args[i] key_part = arg.split("=", 1)[0] if "=" in arg else arg component = None if key_part.startswith("--component-paths."): component = key_part[len("--component-paths.") :].replace("-", "_") elif key_part.startswith("--component_paths."): component = key_part[len("--component_paths.") :].replace("-", "_") elif key_part.startswith("--") and key_part.endswith("-path"): component = key_part[2:-5].replace("-", "_") if component is not None: if "=" in arg: component_paths[component] = arg.split("=", 1)[1] elif i + 1 < len(unknown_args) and not unknown_args[i + 1].startswith( "-" ): i += 1 component_paths[component] = unknown_args[i] else: remaining.append(arg) i += 1 continue else: remaining.append(arg) i += 1 # canonicalize and validate for component, path in component_paths.items(): path = os.path.expanduser(path) component_paths[component] = path return component_paths, remaining @staticmethod def _extract_component_attention_backends( unknown_args: list[str], ) -> tuple[dict[str, str], list[str]]: component_attention_backends: dict[str, str] = {} remaining: list[str] = [] i = 0 while i < len(unknown_args): arg = unknown_args[i] key_part = arg.split("=", 1)[0] if "=" in arg else arg component = None if key_part.startswith("--component-attention-backends."): component = key_part[len("--component-attention-backends.") :].replace( "-", "_" ) elif key_part.startswith("--component_attention_backends."): component = key_part[len("--component_attention_backends.") :].replace( "-", "_" ) if component is not None: if "=" in arg: component_attention_backends[component] = arg.split("=", 1)[1] elif i + 1 < len(unknown_args) and not unknown_args[i + 1].startswith( "-" ): i += 1 component_attention_backends[component] = unknown_args[i] else: remaining.append(arg) i += 1 continue else: remaining.append(arg) i += 1 return component_attention_backends, remaining @classmethod def from_cli_args( cls, args: argparse.Namespace, unknown_args: list[str] | None = None, default_args: dict[str, Any] | None = None, ) -> "ServerArgs": if unknown_args is None: unknown_args = [] # extract dynamic ---path from unknown args dynamic_paths, remaining = cls._extract_component_paths(unknown_args) dynamic_attention_backends, remaining = ( cls._extract_component_attention_backends(remaining) ) if remaining: raise SystemExit(f"error: unrecognized arguments: {' '.join(remaining)}") provided_args = cls.get_provided_args(args, unknown_args) explicit_arg_names = set(provided_args) # Handle config file config_file = provided_args.get("config") if config_file: config_args = cls.load_config_file(config_file) explicit_arg_names.update(config_args) provided_args = {**config_args, **provided_args} if default_args: for key, value in default_args.items(): provided_args.setdefault(key, value) if dynamic_paths: existing = dict(provided_args.get("component_paths") or {}) existing.update(dynamic_paths) provided_args["component_paths"] = existing explicit_arg_names.add("component_paths") if dynamic_attention_backends: existing = cls._parse_component_attention_backend_map( provided_args.get("component_attention_backends") ) existing.update(dynamic_attention_backends) provided_args["component_attention_backends"] = existing explicit_arg_names.add("component_attention_backends") provided_args["_explicit_arg_names"] = explicit_arg_names return cls.from_dict(provided_args) @classmethod def from_dict(cls, kwargs: dict[str, Any]) -> "ServerArgs": """Create a ServerArgs object from a dictionary.""" attrs = [attr.name for attr in dataclasses.fields(cls) if attr.init] server_args_kwargs: dict[str, Any] = {} explicit_arg_names = kwargs.get("_explicit_arg_names") if explicit_arg_names is None: explicit_arg_names = set(kwargs) component_paths = dict(kwargs.get("component_paths") or {}) if component_paths: server_args_kwargs["component_paths"] = component_paths server_args_kwargs["_explicit_arg_names"] = set(explicit_arg_names) for attr in attrs: if attr == "_explicit_arg_names": continue elif attr == "pipeline_config": pipeline_config = PipelineConfig.from_kwargs(kwargs) logger.debug(f"Using PipelineConfig: {type(pipeline_config)}") server_args_kwargs["pipeline_config"] = pipeline_config elif attr == "nunchaku_config": nunchaku_config = NunchakuSVDQuantArgs.from_dict(kwargs) server_args_kwargs["nunchaku_config"] = nunchaku_config elif attr in kwargs: server_args_kwargs[attr] = kwargs[attr] return cls(**server_args_kwargs) @staticmethod def load_config_file(config_file: str) -> dict[str, Any]: """Load a config file.""" if config_file.endswith(".json"): with open(config_file, "r") as f: return json.load(f) elif config_file.endswith((".yaml", ".yml")): try: import yaml except ImportError: raise ImportError( "Please install PyYAML to use YAML config files. " "`pip install pyyaml`" ) with open(config_file, "r") as f: return yaml.safe_load(f) else: raise ValueError(f"Unsupported config file format: {config_file}") @classmethod def from_kwargs(cls, **kwargs: Any) -> "ServerArgs": explicit_arg_names = set(kwargs) # Convert backend string to enum if necessary if "backend" in kwargs and isinstance(kwargs["backend"], str): kwargs["backend"] = Backend.from_string(kwargs["backend"]) # Convert disagg_role string to enum if necessary if "disagg_role" in kwargs and isinstance(kwargs["disagg_role"], str): kwargs["disagg_role"] = RoleType.from_string(kwargs["disagg_role"]) kwargs["pipeline_config"] = PipelineConfig.from_kwargs(kwargs) kwargs["_explicit_arg_names"] = explicit_arg_names return cls(**kwargs) @staticmethod def get_provided_args( args: argparse.Namespace, unknown_args: list[str] ) -> dict[str, Any]: """Get the arguments provided by the user.""" provided_args = {} # We need to check against the raw command-line arguments to see what was # explicitly provided by the user, vs. what's a default value from argparse. raw_argv = sys.argv + unknown_args # Create a set of argument names that were present on the command line. # This handles both styles: '--arg=value' and '--arg value'. provided_arg_names = set(getattr(args, "_sglang_explicit_arg_names", ())) for arg in raw_argv: if arg.startswith("--"): # For '--arg=value', this gets 'arg'; for '--arg', this also gets 'arg'. arg_name = arg.split("=", 1)[0].replace("-", "_").lstrip("_") provided_arg_names.add(arg_name) cli_aliases = { "cfg_parallel_size": "cfg_parallel_degree", "data_parallel_size": "dp_size", "dp": "dp_size", "layerwise_offload_modules": "layerwise_offload_components", "mode": "performance_mode", } for alias_name, dest_name in cli_aliases.items(): if alias_name in provided_arg_names: provided_arg_names.add(dest_name) # Populate provided_args if the argument from the namespace was on the command line. for k, v in vars(args).items(): if k.startswith("_sglang_"): continue if k in provided_arg_names: provided_args[k] = v return provided_args def _validate_pipeline(self): if self.pipeline_config is None: raise ValueError("pipeline_config is not set in ServerArgs") self.pipeline_config.check_pipeline_config() def _validate_offload(self): # validate dit_offload_prefetch_size if self.dit_offload_prefetch_size > 1 and ( isinstance(self.dit_offload_prefetch_size, float) and not self.dit_offload_prefetch_size.is_integer() ): self.dit_offload_prefetch_size = int( math.floor(self.dit_offload_prefetch_size) ) logger.info( f"Invalid --dit-offload-prefetch-size value passed, truncated to: {self.dit_offload_prefetch_size}" ) if 0.5 <= self.dit_offload_prefetch_size < 1.0: logger.info( "We do not recommend --dit-offload-prefetch-size to be between 0.5 and 1.0" ) # validate layerwise offload conflicts if envs.SGLANG_CACHE_DIT_ENABLED and self.use_fsdp_inference: if self.is_arg_explicitly_set("use_fsdp_inference"): raise ValueError( "FSDP inference cannot be enabled together with cache-dit. " "cache-dit wraps known DiT block structures, while FSDP wraps " "and shards modules before cache-dit can inspect them. " "Please disable --use-fsdp-inference or disable " "SGLANG_CACHE_DIT_ENABLED." ) logger.warning( "cache-dit is enabled, automatically disabling use_fsdp_inference." ) self.use_fsdp_inference = False if self.layerwise_offload_components: if self.dit_offload_prefetch_size < 0.0: raise ValueError("dit_offload_prefetch_size must be non-negative") is_dit_layerwise_offload_selected = self.is_dit_layerwise_offload_selected if self.use_fsdp_inference and is_dit_layerwise_offload_selected: logger.warning( "layerwise offload is selected for DiT components, automatically disabling use_fsdp_inference." ) self.use_fsdp_inference = False if envs.SGLANG_CACHE_DIT_ENABLED and is_dit_layerwise_offload_selected: raise ValueError( "DiT layerwise offload cannot be enabled together with cache-dit. " "cache-dit may reuse skipped blocks whose weights have been released by layerwise offload, " "causing shape mismatch errors. " "Please disable --dit-layerwise-offload, remove DiT from --layerwise-offload-components, " "or disable SGLANG_CACHE_DIT_ENABLED." ) if ( self.performance_mode == "memory" or self.is_arg_explicitly_set("layerwise_offload_components") or self.dit_layerwise_offload ): logger.info_once( "Using layerwise offload components: " f"{', '.join(self.layerwise_offload_components)}. " "This reduces peak GPU memory and can increase latency; use " "--performance-mode speed for GPU-resident defaults when memory allows." ) def _validate_parallelism(self): if self.sp_degree > self.num_gpus or self.num_gpus % self.sp_degree != 0: raise ValueError( f"num_gpus ({self.num_gpus}) must be >= and divisible by sp_degree ({self.sp_degree})" ) if ( self.hsdp_replicate_dim > self.num_gpus or self.num_gpus % self.hsdp_replicate_dim != 0 ): raise ValueError( f"num_gpus ({self.num_gpus}) must be >= and divisible by hsdp_replicate_dim ({self.hsdp_replicate_dim})" ) if ( self.hsdp_shard_dim > self.num_gpus or self.num_gpus % self.hsdp_shard_dim != 0 ): raise ValueError( f"num_gpus ({self.num_gpus}) must be >= and divisible by hsdp_shard_dim ({self.hsdp_shard_dim})" ) if self.num_gpus % self.dp_size != 0: raise ValueError( f"num_gpus ({self.num_gpus}) must be divisible by dp_size ({self.dp_size})" ) if self.dp_size < 1: raise ValueError("--dp-size must be a natural number") if self.dp_size > 1: raise ValueError("DP is not yet supported") num_gpus_per_group = self.dp_size * self.tp_size if self.enable_cfg_parallel: num_gpus_per_group *= self.cfg_parallel_degree if self.num_gpus % num_gpus_per_group != 0: raise ValueError( f"num_gpus ({self.num_gpus}) must be divisible by (dp_size * tp_size" f"{f' * {self.cfg_parallel_degree}' if self.enable_cfg_parallel else ''}" f") = {num_gpus_per_group}" ) if self.sp_degree != self.ring_degree * self.ulysses_degree: raise ValueError( f"sp_degree ({self.sp_degree}) must equal ring_degree * ulysses_degree " f"({self.ring_degree} * {self.ulysses_degree} = {self.ring_degree * self.ulysses_degree})" ) if os.getenv("SGLANG_CACHE_DIT_ENABLED", "").lower() == "true": has_sp = self.sp_degree > 1 has_tp = self.tp_size > 1 if has_sp and has_tp: logger.warning( "cache-dit is enabled with hybrid parallelism (SP + TP). " "Proceeding anyway (SGLang integration may support this mode)." ) def _validate_cfg_parallel(self): if self.enable_cfg_parallel and self.num_gpus == 1: raise ValueError( "CFG Parallelism is enabled via `--enable-cfg-parallel`, but num_gpus == 1" ) def _validate_batching(self): if self.batching_mode != "dynamic": raise ValueError("batching_mode must be one of: dynamic") if self.batching_max_size < 1: raise ValueError("batching_max_size must be >= 1") if self.batching_delay_ms < 0: raise ValueError("batching_delay_ms must be >= 0") def _set_default_attention_backend(self) -> None: """Configure ROCm defaults when users do not specify an attention backend.""" if current_platform.is_rocm(): default_backend = AttentionBackendEnum.AITER.name.lower() self.attention_backend = default_backend logger.info( "Attention backend not specified. Using '%s' by default on ROCm " "to match SGLang SRT defaults.", default_backend, ) @dataclasses.dataclass class PortArgs: # The ipc filename for scheduler (rank 0) to receive inputs from tokenizer (zmq) scheduler_input_ipc_name: str # The port for nccl initialization (torch.dist) nccl_port: int # The ipc filename for rpc call between Engine and Scheduler rpc_ipc_name: str # The ipc filename for Scheduler to send metrics metrics_ipc_name: str # Master port for distributed inference master_port: int | None = None @staticmethod def from_server_args( server_args: ServerArgs, dp_rank: Optional[int] = None ) -> "PortArgs": if server_args.nccl_port is None: nccl_port = server_args.scheduler_port + random.randint(100, 1000) while True: if is_port_available(nccl_port): break if nccl_port < 60000: nccl_port += 42 else: nccl_port -= 43 else: nccl_port = server_args.nccl_port # Normal case, use IPC within a single node return PortArgs( scheduler_input_ipc_name=f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}", nccl_port=nccl_port, rpc_ipc_name=f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}", metrics_ipc_name=f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}", master_port=server_args.master_port, ) _global_server_args = None def prepare_server_args(argv: list[str]) -> ServerArgs: """ Prepare the inference arguments from the command line arguments. """ parser = FlexibleArgumentParser() ServerArgs.add_cli_args(parser) raw_args, unknown_args = parser.parse_known_args(argv) server_args = ServerArgs.from_cli_args(raw_args, unknown_args) return server_args def set_global_server_args(server_args: ServerArgs): """ Set the global sgl_diffusion config for each process """ global _global_server_args _global_server_args = server_args def get_global_server_args() -> ServerArgs: if _global_server_args is None: # in ci, usually when we test custom ops/modules directly, # we don't set the sgl_diffusion config. In that case, we set a default # config. # TODO(will): may need to handle this for CI. raise ValueError("Global sgl_diffusion args is not set.") return _global_server_args