From b76730701b66dc611fb7c8a4232aabb898d32895 Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 30 Mar 2026 19:45:34 +0800 Subject: [PATCH] [diffusion] feat: enhance overlay mechanism (#21648) --- docs/diffusion/api/cli.md | 5 ++ python/sglang/cli/utils.py | 27 +++++++--- python/sglang/multimodal_gen/registry.py | 25 +++------- .../runtime/utils/hf_diffusers_utils.py | 26 ++++------ .../runtime/utils/model_overlay.py | 42 +++++++++------- python/sglang/utils.py | 50 +++++++++++++++++++ 6 files changed, 112 insertions(+), 63 deletions(-) diff --git a/docs/diffusion/api/cli.md b/docs/diffusion/api/cli.md index 8f2bba5ca..a4caaca7d 100644 --- a/docs/diffusion/api/cli.md +++ b/docs/diffusion/api/cli.md @@ -32,6 +32,11 @@ Notes: 1. `SGLANG_DIFFUSION_MODEL_OVERLAY_REGISTRY` is only an optional override for development and debugging. It accepts either a JSON object or a path to a JSON file, and can extend or replace built-in entries for the current process. +2. On the first load, SGLang will: + - download overlay metadata from the overlay repo + - download the required files from the original source repo + - materialize a local standard component repo under `~/.cache/sgl_diffusion/materialized_models/` +3. Later loads reuse the materialized local repo. The materialized repo is what the runtime loads as a normal componentized model directory. ## Quick Start diff --git a/python/sglang/cli/utils.py b/python/sglang/cli/utils.py index 1d867eccf..60fb10d92 100644 --- a/python/sglang/cli/utils.py +++ b/python/sglang/cli/utils.py @@ -5,10 +5,24 @@ import subprocess from functools import lru_cache from sglang.srt.environ import envs +from sglang.utils import ( + has_diffusion_overlay_registry_match, + is_known_non_diffusers_diffusion_model, + load_diffusion_overlay_registry_from_env, +) logger = logging.getLogger(__name__) +@lru_cache(maxsize=1) +def _load_overlay_registry() -> dict: + return load_diffusion_overlay_registry_from_env() + + +def _is_overlay_diffusion_model(model_path: str) -> bool: + return has_diffusion_overlay_registry_match(model_path, _load_overlay_registry()) + + def _is_diffusers_model_dir(model_dir: str) -> bool: """Check if a local directory contains a valid diffusers model_index.json.""" config_path = os.path.join(model_dir, "model_index.json") @@ -29,19 +43,16 @@ def get_is_diffusion_model(model_path: str) -> bool: Returns False on any failure (network error, 404, offline mode, etc.) so that the caller falls through to the standard LLM server path. """ - try: - from sglang.multimodal_gen.registry import ( - is_known_non_diffusers_multimodal_model, - ) - except ImportError: - is_known_non_diffusers_multimodal_model = lambda _: False + if _is_overlay_diffusion_model(model_path): + # short-circuit, if applicable for the overlay mechanism (diffusion-only) + return True if os.path.isdir(model_path): if _is_diffusers_model_dir(model_path): return True - return is_known_non_diffusers_multimodal_model(model_path) + return is_known_non_diffusers_diffusion_model(model_path) - if is_known_non_diffusers_multimodal_model(model_path): + if is_known_non_diffusers_diffusion_model(model_path): return True try: diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 1e3aac6a4..b9e03c3d9 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -122,9 +122,9 @@ from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ) from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( maybe_download_model_index, - verify_model_config_and_directory, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.utils import KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS logger = init_logger(__name__) @@ -327,10 +327,7 @@ def _get_config_info( return _CONFIG_REGISTRY.get(model_id) # 3. Use detectors - if os.path.exists(model_path): - config = verify_model_config_and_directory(model_path) - else: - config = maybe_download_model_index(model_path) + config = maybe_download_model_index(model_path) pipeline_name = config.get("_class_name", "").lower() matched_model_names = [] @@ -499,10 +496,7 @@ def get_model_info( else: # Try to get from model_index.json try: - if os.path.exists(model_path): - config = verify_model_config_and_directory(model_path) - else: - config = maybe_download_model_index(model_path) + config = maybe_download_model_index(model_path) except Exception as e: logger.error(f"Could not read model config for '{model_path}': {e}") if backend == Backend.AUTO: @@ -876,25 +870,18 @@ def _register_configs(): _register_configs() -# Known non-diffusers multimodal model patterns -# Maps pattern -> pipeline_name for models that don't have model_index.json -_NON_DIFFUSERS_MULTIMODAL_PATTERNS: Dict[str, str] = { - "hunyuan3d": "Hunyuan3D2Pipeline", - "flux.2-dev-nvfp4": "Flux2NvfpPipeline", -} - - def is_known_non_diffusers_multimodal_model(model_path: str) -> bool: model_path_lower = model_path.lower() return any( - pattern in model_path_lower for pattern in _NON_DIFFUSERS_MULTIMODAL_PATTERNS + pattern in model_path_lower + for pattern in KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS ) def get_non_diffusers_pipeline_name(model_path: str) -> Optional[str]: """Get the pipeline name for a known non-diffusers model.""" model_path_lower = model_path.lower() - for pattern, pipeline_name in _NON_DIFFUSERS_MULTIMODAL_PATTERNS.items(): + for pattern, pipeline_name in KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS.items(): if pattern in model_path_lower: return pipeline_name return None diff --git a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py index 36516ddc4..f12f30411 100644 --- a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py @@ -489,15 +489,16 @@ def maybe_download_model_index(model_name_or_path: str) -> dict[str, Any]: from huggingface_hub.errors import EntryNotFoundError - # If it's a local path, verify it directly + overlay_config = maybe_load_overlay_model_index( + model_name_or_path, + snapshot_download_fn=snapshot_download, + hf_hub_download_fn=hf_hub_download, + ) + if overlay_config is not None: + return overlay_config + + # If it's a local path, verify it directly. if os.path.exists(model_name_or_path): - overlay_config = maybe_load_overlay_model_index( - model_name_or_path, - snapshot_download_fn=snapshot_download, - hf_hub_download_fn=hf_hub_download, - ) - if overlay_config is not None: - return overlay_config try: return verify_model_config_and_directory(model_name_or_path) except ValueError: @@ -509,15 +510,6 @@ def maybe_download_model_index(model_name_or_path: str) -> dict[str, Any]: return config raise - # return resolved overlay config if applicable - overlay_config = maybe_load_overlay_model_index( - model_name_or_path, - snapshot_download_fn=snapshot_download, - hf_hub_download_fn=hf_hub_download, - ) - if overlay_config is not None: - return overlay_config - # For remote models, download just the model_index.json try: with tempfile.TemporaryDirectory() as tmp_dir: diff --git a/python/sglang/multimodal_gen/runtime/utils/model_overlay.py b/python/sglang/multimodal_gen/runtime/utils/model_overlay.py index 5085db56f..eabc2e9dd 100644 --- a/python/sglang/multimodal_gen/runtime/utils/model_overlay.py +++ b/python/sglang/multimodal_gen/runtime/utils/model_overlay.py @@ -19,6 +19,7 @@ from requests.exceptions import RequestException from sglang.multimodal_gen.runtime.loader.weight_utils import get_lock from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.utils import load_diffusion_overlay_registry_from_env logger = init_logger(__name__) @@ -60,28 +61,12 @@ def _load_model_overlay_registry() -> dict[str, dict[str, Any]]: # Built-in registry is the stable default path; env only overrides it. normalized = _normalize_model_overlay_registry(BUILTIN_MODEL_OVERLAY_REGISTRY) - raw_value = os.getenv("SGLANG_DIFFUSION_MODEL_OVERLAY_REGISTRY", "").strip() - if not raw_value: + env_registry = load_diffusion_overlay_registry_from_env() + if not env_registry: _MODEL_OVERLAY_REGISTRY_CACHE = normalized return _MODEL_OVERLAY_REGISTRY_CACHE - try: - if raw_value.startswith("{"): - payload = json.loads(raw_value) - else: - with open(os.path.expanduser(raw_value), encoding="utf-8") as f: - payload = json.load(f) - except Exception as exc: - raise ValueError( - "Failed to parse SGLANG_DIFFUSION_MODEL_OVERLAY_REGISTRY" - ) from exc - - if not isinstance(payload, dict): - raise ValueError( - "SGLANG_DIFFUSION_MODEL_OVERLAY_REGISTRY must be a JSON object" - ) - - normalized.update(_normalize_model_overlay_registry(payload)) + normalized.update(_normalize_model_overlay_registry(env_registry)) _MODEL_OVERLAY_REGISTRY_CACHE = normalized return _MODEL_OVERLAY_REGISTRY_CACHE @@ -462,11 +447,23 @@ def materialize_overlay_model( ): return final_dir + logger.info( + "Materializing overlay model for %s into %s", + source_model_id, + final_dir, + ) + logger.info( + "Overlay source repo: %s, overlay repo: %s@%s", + source_model_id, + overlay_repo_id, + overlay_revision, + ) tmp_dir = final_dir + ".tmp" if os.path.exists(tmp_dir): shutil.rmtree(tmp_dir) if os.path.exists(final_dir): shutil.rmtree(final_dir) + logger.info("Copying overlay metadata into temporary materialized directory") shutil.copytree( overlay_dir, tmp_dir, @@ -479,11 +476,17 @@ def materialize_overlay_model( file_mappings = manifest.get("file_mappings", []) if file_mappings: + logger.info("Applying %d overlay file mappings", len(file_mappings)) _apply_overlay_file_mappings( source_dir=source_dir, output_dir=tmp_dir, file_mappings=cast(list[dict[str, Any]], file_mappings), ) + if manifest.get("custom_materializer"): + logger.info( + "Running custom overlay materializer: %s", + manifest["custom_materializer"], + ) _run_overlay_custom_materializer( overlay_dir=overlay_dir, source_dir=source_dir, @@ -506,6 +509,7 @@ def materialize_overlay_model( ) os.replace(tmp_dir, final_dir) + logger.info("Overlay materialization finished: %s", final_dir) return final_dir diff --git a/python/sglang/utils.py b/python/sglang/utils.py index cd7fbe721..7cc322a23 100644 --- a/python/sglang/utils.py +++ b/python/sglang/utils.py @@ -31,6 +31,56 @@ from sglang.srt.environ import envs logger = logging.getLogger(__name__) +KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS: dict[str, str] = { + "hunyuan3d": "Hunyuan3D2Pipeline", + "flux.2-dev-nvfp4": "Flux2NvfpPipeline", +} + + +def load_diffusion_overlay_registry_from_env() -> dict[str, dict[str, Any]]: + raw_value = os.getenv("SGLANG_DIFFUSION_MODEL_OVERLAY_REGISTRY", "").strip() + if not raw_value: + return {} + + if raw_value.startswith("{"): + payload = json.loads(raw_value) + else: + with open(os.path.expanduser(raw_value), encoding="utf-8") as f: + payload = json.load(f) + + if not isinstance(payload, dict): + return {} + + normalized: dict[str, dict[str, Any]] = {} + for source_model_id, spec in payload.items(): + if isinstance(spec, str): + normalized[source_model_id] = {"overlay_repo_id": spec} + elif isinstance(spec, dict) and spec.get("overlay_repo_id"): + normalized[source_model_id] = dict(spec) + return normalized + + +def has_diffusion_overlay_registry_match( + model_path: str, registry: dict[str, dict[str, Any]] | None = None +) -> bool: + registry = ( + load_diffusion_overlay_registry_from_env() if registry is None else registry + ) + if model_path in registry: + return True + if not os.path.exists(model_path): + return False + base_name = os.path.basename(os.path.normpath(model_path)) + return any(base_name == key.rsplit("/", 1)[-1] for key in registry) + + +def is_known_non_diffusers_diffusion_model(model_path: str) -> bool: + model_path_lower = model_path.lower() + return any( + pattern in model_path_lower + for pattern in KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS + ) + def execute_once(func): has_run = None