chore: import upstream snapshot with attribution

python/sglang/multimodal_gen/runtime/utils/model_overlay.py

@@ -0,0 +1,780 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import glob
+import hashlib
+import importlib.util
+import json
+import os
+import shutil
+from typing import Any, Callable, cast
+
+from huggingface_hub.errors import (
+    LocalEntryNotFoundError,
+    RepositoryNotFoundError,
+    RevisionNotFoundError,
+)
+from requests.exceptions import ConnectionError as RequestsConnectionError
+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__)
+
+# Built-in diffusion model overlay registry.
+BUILTIN_MODEL_OVERLAY_REGISTRY: dict[str, dict[str, Any]] = {
+    "Lightricks/LTX-2.3": {
+        "overlay_repo_id": "MickJ/LTX-2.3-overlay",
+        "overlay_revision": "e0cc94f279ec16bb87c230134d40319f6ce40c5e",
+    },
+    "jdopensource/JoyAI-Echo": {
+        "overlay_repo_id": "Niehen6174/JoyAI-Echo-overlay",
+        "overlay_revision": "0a19f315c96532b7a5f61bcd765d1fefdd83dc7d",
+    },
+    "Efficient-Large-Model/SANA-WM_bidirectional": {
+        "overlay_repo_id": "sjmshsh/SANA-WM_bidirectional-overlay",
+        "overlay_revision": "e611beacbcc0cf33c676306ae0eb89f149e044ad",
+    },
+    "Efficient-Large-Model/SANA-WM_streaming": {
+        "overlay_repo_id": "AgainstEntropy/SANA-WM_streaming-overlay",
+        "overlay_revision": "62c6840871ecc3559189047513ba0670e1bf62e7",
+    },
+}
+
+
+MODEL_OVERLAY_METADATA_PATTERNS = [
+    "*.json",
+    "*.md",
+    "*.py",
+    "*.txt",
+    "**/*.json",
+    "**/*.md",
+    "**/*.py",
+    "**/*.txt",
+]
+
+_MATERIALIZED_WEIGHT_SUFFIXES = (".safetensors", ".bin", ".pth", ".pt")
+_MATERIALIZED_CONFIG_ONLY_COMPONENTS = {
+    "feature_extractor",
+    "image_processor",
+    "processor",
+    "scheduler",
+    "tokenizer",
+    "tokenizer_2",
+}
+
+_MODEL_OVERLAY_REGISTRY_CACHE: dict[str, dict[str, Any]] | None = None
+
+
+def _compute_overlay_fingerprint(overlay_dir: str) -> str:
+    hasher = hashlib.sha256()
+    for root, dir_names, file_names in os.walk(overlay_dir):
+        dir_names[:] = sorted(
+            d for d in dir_names if d != "__pycache__" and not d.endswith(".egg-info")
+        )
+        for file_name in sorted(file_names):
+            if file_name.endswith((".safetensors", ".bin", ".pth", ".pt")):
+                continue
+            file_path = os.path.join(root, file_name)
+            rel_path = os.path.relpath(file_path, overlay_dir).replace(os.sep, "/")
+            hasher.update(rel_path.encode("utf-8"))
+            with open(file_path, "rb") as f:
+                hasher.update(hashlib.sha256(f.read()).digest())
+    return hasher.hexdigest()
+
+
+def _resolve_bundled_overlay_dir(overlay_spec: dict[str, Any]) -> str | None:
+    bundled_overlay_subdir = overlay_spec.get("bundled_overlay_subdir")
+    if not bundled_overlay_subdir:
+        return None
+    bundled_overlay_dir = os.path.abspath(
+        os.path.join(
+            os.path.dirname(__file__),
+            "..",
+            "..",
+            "model_overlays",
+            str(bundled_overlay_subdir),
+        )
+    )
+    if not os.path.isdir(bundled_overlay_dir):
+        return None
+    if load_overlay_manifest_if_present(bundled_overlay_dir) is None:
+        return None
+    return bundled_overlay_dir
+
+
+def get_diffusion_cache_root() -> str:
+    return os.path.expanduser(
+        os.getenv("SGLANG_DIFFUSION_CACHE_ROOT", "~/.cache/sgl_diffusion")
+    )
+
+
+def clear_model_overlay_registry_cache() -> None:
+    global _MODEL_OVERLAY_REGISTRY_CACHE
+    _MODEL_OVERLAY_REGISTRY_CACHE = None
+
+
+def _load_model_overlay_registry() -> dict[str, dict[str, Any]]:
+    global _MODEL_OVERLAY_REGISTRY_CACHE
+    if _MODEL_OVERLAY_REGISTRY_CACHE is not None:
+        return _MODEL_OVERLAY_REGISTRY_CACHE
+
+    # Built-in registry is the stable default path; env only overrides it.
+    normalized = _normalize_model_overlay_registry(BUILTIN_MODEL_OVERLAY_REGISTRY)
+
+    env_registry = load_diffusion_overlay_registry_from_env()
+    if not env_registry:
+        _MODEL_OVERLAY_REGISTRY_CACHE = normalized
+        return _MODEL_OVERLAY_REGISTRY_CACHE
+
+    normalized.update(_normalize_model_overlay_registry(env_registry))
+    _MODEL_OVERLAY_REGISTRY_CACHE = normalized
+    return _MODEL_OVERLAY_REGISTRY_CACHE
+
+
+def _normalize_model_overlay_registry(
+    payload: dict[str, Any],
+) -> dict[str, dict[str, Any]]:
+    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}
+            continue
+        if not isinstance(spec, dict):
+            raise ValueError(
+                "Overlay registry values must be either strings or JSON objects"
+            )
+        overlay_repo_id = spec.get("overlay_repo_id")
+        if not overlay_repo_id:
+            raise ValueError(
+                f"Overlay registry entry for {source_model_id!r} is missing overlay_repo_id"
+            )
+        normalized[source_model_id] = dict(spec)
+    return normalized
+
+
+def resolve_model_overlay(model_name_or_path: str) -> dict[str, Any] | None:
+    registry = _load_model_overlay_registry()
+    return registry.get(model_name_or_path)
+
+
+def resolve_model_overlay_target(
+    model_name_or_path: str,
+) -> tuple[str, dict[str, Any]] | None:
+    registry = _load_model_overlay_registry()
+
+    exact = registry.get(model_name_or_path)
+    if exact is not None:
+        return model_name_or_path, exact
+
+    if os.path.exists(model_name_or_path):
+        # Local source dirs do not have a repo id, so match them by basename.
+        base_name = os.path.basename(os.path.normpath(model_name_or_path))
+        normalized_path = (
+            os.path.normpath(model_name_or_path).lower().replace(os.sep, "/")
+        )
+        for source_model_id, spec in registry.items():
+            if base_name == source_model_id.rsplit("/", 1)[-1]:
+                return source_model_id, spec
+            cache_repo_fragment = (
+                f"models--{source_model_id.lower().replace('/', '--')}"
+            )
+            if cache_repo_fragment in normalized_path:
+                return source_model_id, spec
+
+    return None
+
+
+def load_overlay_manifest_if_present(overlay_dir: str) -> dict[str, Any] | None:
+    overlay_manifest_path = os.path.join(
+        overlay_dir, "_overlay", "overlay_manifest.json"
+    )
+    if not os.path.exists(overlay_manifest_path):
+        return None
+    with open(overlay_manifest_path, encoding="utf-8") as f:
+        manifest = cast(dict[str, Any], json.load(f))
+    return manifest
+
+
+def load_model_index_from_dir(model_dir: str) -> dict[str, Any]:
+    model_index_path = os.path.join(model_dir, "model_index.json")
+    if not os.path.exists(model_index_path):
+        raise ValueError(f"model_index.json not found under {model_dir}")
+    with open(model_index_path, encoding="utf-8") as f:
+        config = cast(dict[str, Any], json.load(f))
+    if "_class_name" not in config or "_diffusers_version" not in config:
+        raise ValueError(f"Invalid model_index.json under {model_dir}")
+    config["pipeline_name"] = config["_class_name"]
+    return config
+
+
+def _component_has_weight_file(component_dir: str) -> bool:
+    for root, _, file_names in os.walk(component_dir):
+        if any(
+            file_name.endswith(_MATERIALIZED_WEIGHT_SUFFIXES)
+            and os.path.isfile(os.path.join(root, file_name))
+            for file_name in file_names
+        ):
+            return True
+    return False
+
+
+def _materialized_overlay_has_component_weights(model_dir: str) -> bool:
+    model_index = load_model_index_from_dir(model_dir)
+    for component_name, entry in model_index.items():
+        if (
+            component_name.startswith("_")
+            or component_name == "pipeline_name"
+            or component_name in _MATERIALIZED_CONFIG_ONLY_COMPONENTS
+            or not isinstance(entry, list)
+        ):
+            continue
+        component_dir = os.path.join(model_dir, component_name)
+        if not os.path.isdir(component_dir) or not _component_has_weight_file(
+            component_dir
+        ):
+            logger.warning(
+                "Materialized overlay cache for %s is missing weights for component %s",
+                model_dir,
+                component_name,
+            )
+            return False
+    return True
+
+
+def _materialized_overlay_cache_complete(
+    final_dir: str,
+    marker_path: str,
+    verify_diffusers_model_complete_fn: Callable[[str], bool],
+) -> bool:
+    return (
+        verify_diffusers_model_complete_fn(final_dir)
+        and os.path.exists(marker_path)
+        and _materialized_overlay_has_component_weights(final_dir)
+    )
+
+
+def _ensure_dir(path: str) -> None:
+    os.makedirs(path, exist_ok=True)
+
+
+def _find_missing_required_paths(
+    root_dir: str, required_paths: list[str] | tuple[str, ...]
+) -> list[str]:
+    missing: list[str] = []
+    for rel_path in required_paths:
+        if not os.path.exists(os.path.join(root_dir, rel_path)):
+            missing.append(rel_path)
+    return missing
+
+
+def _link_or_copy_file(src: str, dst: str) -> None:
+    src = os.path.realpath(src)
+    _ensure_dir(os.path.dirname(dst))
+    if os.path.lexists(dst):
+        os.remove(dst)
+    try:
+        os.link(src, dst)
+        return
+    except OSError:
+        pass
+    try:
+        os.symlink(src, dst)
+        return
+    except OSError:
+        pass
+    shutil.copy2(src, dst)
+
+
+def _copytree_link_or_copy(src_dir: str, dst_dir: str) -> None:
+    for root, _, files in os.walk(src_dir):
+        rel_root = os.path.relpath(root, src_dir)
+        target_root = dst_dir if rel_root == "." else os.path.join(dst_dir, rel_root)
+        _ensure_dir(target_root)
+        for file_name in files:
+            src_file = os.path.join(root, file_name)
+            dst_file = os.path.join(target_root, file_name)
+            _link_or_copy_file(src_file, dst_file)
+
+
+def ensure_overlay_source_dir_complete(
+    *,
+    source_model_id: str,
+    source_dir: str,
+    manifest: dict[str, Any],
+    local_dir: str | None,
+    allow_patterns: list[str] | None,
+    download: bool,
+    snapshot_download_fn: Callable[..., str],
+) -> str:
+    required_source_files = cast(
+        list[str], list(manifest.get("required_source_files", []))
+    )
+    if not required_source_files:
+        return source_dir
+
+    # Metadata-only overlays often need a partial source snapshot. Re-download
+    # only when the current source dir is missing required files.
+    missing_paths = _find_missing_required_paths(source_dir, required_source_files)
+    if not missing_paths:
+        return source_dir
+
+    if not download:
+        raise ValueError(
+            f"Overlay source model {source_model_id} is missing required files "
+            f"{missing_paths} and download=False."
+        )
+
+    logger.warning(
+        "Overlay source model %s is missing required files %s. "
+        "Re-downloading source snapshot.",
+        source_model_id,
+        missing_paths,
+    )
+    source_allow_patterns = manifest.get("source_allow_patterns")
+    effective_allow_patterns = (
+        cast(list[str] | None, source_allow_patterns)
+        if source_allow_patterns is not None
+        else allow_patterns
+    )
+    with get_lock(source_model_id).acquire(poll_interval=2):
+        source_dir = snapshot_download_fn(
+            repo_id=source_model_id,
+            ignore_patterns=["*.onnx", "*.msgpack"],
+            allow_patterns=effective_allow_patterns,
+            local_dir=local_dir,
+            max_workers=8,
+            force_download=True,
+        )
+    missing_after_redownload = _find_missing_required_paths(
+        source_dir, required_source_files
+    )
+    if missing_after_redownload:
+        raise ValueError(
+            f"Overlay source model {source_model_id} is still missing required files "
+            f"{missing_after_redownload} after re-download."
+        )
+    return str(source_dir)
+
+
+def resolve_direct_overlay_repo(
+    model_name_or_path: str,
+    *,
+    hf_hub_download_fn: Callable[..., str],
+) -> tuple[dict[str, Any], str, dict[str, Any]] | None:
+    if os.path.exists(model_name_or_path):
+        manifest = load_overlay_manifest_if_present(model_name_or_path)
+        if manifest is None:
+            return None
+        source_model_id = manifest.get("source_model_id")
+        if not source_model_id:
+            raise ValueError(
+                f"Overlay repo {model_name_or_path} is missing source_model_id in _overlay/overlay_manifest.json"
+            )
+        overlay_spec = {
+            "overlay_repo_id": model_name_or_path,
+            "overlay_revision": "local",
+        }
+        return overlay_spec, model_name_or_path, manifest
+
+    try:
+        manifest_path = hf_hub_download_fn(
+            repo_id=model_name_or_path,
+            filename="_overlay/overlay_manifest.json",
+        )
+        overlay_dir = os.path.dirname(os.path.dirname(manifest_path))
+    except (
+        RepositoryNotFoundError,
+        RevisionNotFoundError,
+        LocalEntryNotFoundError,
+        RequestsConnectionError,
+        RequestException,
+    ):
+        return None
+    except Exception:
+        return None
+
+    manifest = load_overlay_manifest_if_present(overlay_dir)
+    if manifest is None:
+        return None
+    source_model_id = manifest.get("source_model_id")
+    if not source_model_id:
+        raise ValueError(
+            f"Overlay repo {model_name_or_path} is missing source_model_id in _overlay/overlay_manifest.json"
+        )
+    overlay_spec = {
+        "overlay_repo_id": model_name_or_path,
+        "overlay_revision": "main",
+    }
+    return overlay_spec, overlay_dir, manifest
+
+
+def download_overlay_metadata(
+    source_model_id: str,
+    overlay_spec: dict[str, Any],
+    *,
+    snapshot_download_fn: Callable[..., str],
+) -> str:
+    bundled_overlay_dir = _resolve_bundled_overlay_dir(overlay_spec)
+    if bundled_overlay_dir is not None:
+        logger.info(
+            "Using bundled overlay metadata for %s from %s",
+            source_model_id,
+            bundled_overlay_dir,
+        )
+        return bundled_overlay_dir
+
+    overlay_repo_id = str(overlay_spec["overlay_repo_id"])
+    if os.path.exists(overlay_repo_id):
+        logger.info(
+            "Using local overlay metadata for %s from %s",
+            source_model_id,
+            overlay_repo_id,
+        )
+        return overlay_repo_id
+    revision = overlay_spec.get("overlay_revision")
+    logger.info(
+        "Downloading overlay metadata for %s from %s",
+        source_model_id,
+        overlay_repo_id,
+    )
+    return str(
+        snapshot_download_fn(
+            repo_id=overlay_repo_id,
+            allow_patterns=MODEL_OVERLAY_METADATA_PATTERNS,
+            revision=revision,
+            max_workers=4,
+        )
+    )
+
+
+def _apply_overlay_file_mappings(
+    *,
+    source_dir: str,
+    output_dir: str,
+    file_mappings: list[dict[str, Any]],
+) -> None:
+    for mapping in file_mappings:
+        mapping_type = mapping.get("type", "file")
+        src_rel = mapping.get("src")
+        if not src_rel:
+            raise ValueError(f"Overlay file mapping is missing src: {mapping}")
+        src_path = os.path.join(source_dir, src_rel)
+        if mapping_type == "tree":
+            if not os.path.isdir(src_path):
+                raise ValueError(f"Tree mapping source does not exist: {src_path}")
+            dst_dir = os.path.join(output_dir, str(mapping.get("dst_dir", src_rel)))
+            _copytree_link_or_copy(src_path, dst_dir)
+            continue
+        if mapping_type == "glob":
+            matched = glob.glob(src_path, recursive=True)
+            if not matched:
+                raise ValueError(f"Glob mapping matched no files: {src_path}")
+            for matched_path in matched:
+                if os.path.isdir(matched_path):
+                    continue
+                rel_path = os.path.relpath(matched_path, source_dir)
+                dst_path = os.path.join(output_dir, rel_path)
+                _link_or_copy_file(matched_path, dst_path)
+            continue
+
+        if not os.path.isfile(src_path):
+            raise ValueError(f"File mapping source does not exist: {src_path}")
+        dst_rel = str(mapping.get("dst", os.path.basename(src_rel)))
+        dst_path = os.path.join(output_dir, dst_rel)
+        _link_or_copy_file(src_path, dst_path)
+
+
+def _run_overlay_custom_materializer(
+    *,
+    overlay_dir: str,
+    source_dir: str,
+    output_dir: str,
+    manifest: dict[str, Any],
+) -> None:
+    custom_materializer = manifest.get("custom_materializer")
+    if not custom_materializer:
+        return
+    script_path = os.path.join(overlay_dir, str(custom_materializer))
+    if not os.path.exists(script_path):
+        raise ValueError(f"Custom materializer script not found: {script_path}")
+
+    spec = importlib.util.spec_from_file_location(
+        "_sglang_overlay_materializer", script_path
+    )
+    if spec is None or spec.loader is None:
+        raise ValueError(f"Failed to import custom materializer: {script_path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+
+    materialize_fn = getattr(module, "materialize", None)
+    if materialize_fn is None:
+        raise ValueError(
+            f"Custom materializer {script_path} must define materialize(...)"
+        )
+
+    materialize_fn(
+        overlay_dir=overlay_dir,
+        source_dir=source_dir,
+        output_dir=output_dir,
+        manifest=manifest,
+    )
+
+
+def materialize_overlay_model(
+    *,
+    source_model_id: str,
+    overlay_spec: dict[str, Any],
+    overlay_dir: str,
+    source_dir: str,
+    verify_diffusers_model_complete_fn: Callable[[str], bool],
+) -> str:
+    overlay_manifest_path = os.path.join(
+        overlay_dir, "_overlay", "overlay_manifest.json"
+    )
+    if not os.path.exists(overlay_manifest_path):
+        raise ValueError(
+            f"Overlay repo for {source_model_id} is missing _overlay/overlay_manifest.json"
+        )
+
+    with open(overlay_manifest_path, encoding="utf-8") as f:
+        manifest = cast(dict[str, Any], json.load(f))
+
+    materializer_version = str(manifest.get("materializer_version", "v1"))
+    overlay_repo_id = str(overlay_spec["overlay_repo_id"])
+    overlay_revision = str(overlay_spec.get("overlay_revision", "main"))
+    overlay_fingerprint = _compute_overlay_fingerprint(overlay_dir)
+    cache_key = hashlib.sha256(
+        json.dumps(
+            {
+                "source_model_id": source_model_id,
+                "overlay_repo_id": overlay_repo_id,
+                "overlay_revision": overlay_revision,
+                "materializer_version": materializer_version,
+                "overlay_fingerprint": overlay_fingerprint,
+            },
+            sort_keys=True,
+        ).encode("utf-8")
+    ).hexdigest()[:16]
+    cache_root = os.path.join(get_diffusion_cache_root(), "materialized_models")
+    _ensure_dir(cache_root)
+    safe_name = source_model_id.replace("/", "__")
+    final_dir = os.path.join(cache_root, f"{safe_name}-{cache_key}")
+    marker_path = os.path.join(final_dir, ".sglang_overlay_materialized.json")
+    if _materialized_overlay_cache_complete(
+        final_dir, marker_path, verify_diffusers_model_complete_fn
+    ):
+        return final_dir
+
+    lock_name = (
+        f"overlay-materialize::{source_model_id}::{overlay_repo_id}::{overlay_revision}"
+    )
+    with get_lock(lock_name).acquire(poll_interval=2):
+        if _materialized_overlay_cache_complete(
+            final_dir, marker_path, verify_diffusers_model_complete_fn
+        ):
+            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,
+            ignore=shutil.ignore_patterns("*.safetensors", "*.bin", "*.pth", "*.pt"),
+        )
+
+        overlay_hidden_dir = os.path.join(tmp_dir, "_overlay")
+        if os.path.isdir(overlay_hidden_dir):
+            shutil.rmtree(overlay_hidden_dir)
+
+        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,
+            output_dir=tmp_dir,
+            manifest=manifest,
+        )
+
+        with open(marker_path.replace(final_dir, tmp_dir), "w", encoding="utf-8") as f:
+            json.dump(
+                {
+                    "source_model_id": source_model_id,
+                    "source_dir": source_dir,
+                    "overlay_repo_id": overlay_repo_id,
+                    "overlay_revision": overlay_revision,
+                    "materializer_version": materializer_version,
+                    "overlay_fingerprint": overlay_fingerprint,
+                },
+                f,
+                indent=2,
+                sort_keys=True,
+            )
+
+        os.replace(tmp_dir, final_dir)
+        logger.info("Overlay materialization finished: %s", final_dir)
+
+    return final_dir
+
+
+def maybe_load_overlay_model_index(
+    model_name_or_path: str,
+    *,
+    snapshot_download_fn: Callable[..., str],
+    hf_hub_download_fn: Callable[..., str],
+) -> dict[str, Any] | None:
+    if os.path.exists(model_name_or_path):
+        # A local overlay repo already contains the model_index we need.
+        if load_overlay_manifest_if_present(model_name_or_path) is not None:
+            return load_model_index_from_dir(model_name_or_path)
+
+    overlay_target = resolve_model_overlay_target(model_name_or_path)
+    if overlay_target is not None:
+        # Registry-mapped source model ids first resolve to overlay metadata.
+        source_model_id, overlay_spec = overlay_target
+        overlay_dir = download_overlay_metadata(
+            source_model_id,
+            overlay_spec,
+            snapshot_download_fn=snapshot_download_fn,
+        )
+        return load_model_index_from_dir(overlay_dir)
+
+    direct_overlay = resolve_direct_overlay_repo(
+        model_name_or_path, hf_hub_download_fn=hf_hub_download_fn
+    )
+    if direct_overlay is None:
+        return None
+
+    _, overlay_dir, _ = direct_overlay
+    return load_model_index_from_dir(overlay_dir)
+
+
+def maybe_resolve_overlay_model_path(
+    model_name_or_path: str,
+    *,
+    local_dir: str | None,
+    download: bool,
+    allow_patterns: list[str] | None,
+    snapshot_download_fn: Callable[..., str],
+    hf_hub_download_fn: Callable[..., str],
+    verify_diffusers_model_complete_fn: Callable[[str], bool],
+    base_model_download_fn: Callable[..., str],
+) -> str | None:
+    overlay_target = resolve_model_overlay_target(model_name_or_path)
+    if overlay_target is not None:
+        source_model_id, overlay_spec = overlay_target
+        overlay_dir = download_overlay_metadata(
+            source_model_id,
+            overlay_spec,
+            snapshot_download_fn=snapshot_download_fn,
+        )
+        manifest = load_overlay_manifest_if_present(overlay_dir)
+        if manifest is None:
+            # Full diffusers overlays do not need materialization.
+            return base_model_download_fn(
+                str(overlay_spec["overlay_repo_id"]),
+                local_dir=local_dir,
+                download=download,
+                allow_patterns=allow_patterns,
+                force_diffusers_model=True,
+                skip_overlay_resolution=True,
+            )
+        source_allow_patterns = cast(
+            list[str] | None, manifest.get("source_allow_patterns")
+        )
+        # For local source paths, reuse the directory directly instead of
+        # round-tripping through snapshot_download.
+        source_dir = (
+            model_name_or_path
+            if os.path.exists(model_name_or_path)
+            else base_model_download_fn(
+                source_model_id,
+                local_dir=local_dir,
+                download=download,
+                allow_patterns=source_allow_patterns or allow_patterns,
+                force_diffusers_model=False,
+                skip_overlay_resolution=True,
+            )
+        )
+        source_dir = ensure_overlay_source_dir_complete(
+            source_model_id=source_model_id,
+            source_dir=source_dir,
+            manifest=manifest,
+            local_dir=local_dir,
+            allow_patterns=allow_patterns,
+            download=download,
+            snapshot_download_fn=snapshot_download_fn,
+        )
+        return materialize_overlay_model(
+            source_model_id=source_model_id,
+            overlay_spec=overlay_spec,
+            overlay_dir=overlay_dir,
+            source_dir=source_dir,
+            verify_diffusers_model_complete_fn=verify_diffusers_model_complete_fn,
+        )
+
+    direct_overlay = resolve_direct_overlay_repo(
+        model_name_or_path, hf_hub_download_fn=hf_hub_download_fn
+    )
+    if direct_overlay is None:
+        return None
+
+    overlay_spec, overlay_dir, manifest = direct_overlay
+    source_model_id = str(manifest["source_model_id"])
+    # Direct overlay repos are always metadata-only; they need the original
+    # source weights before they can be materialized into a diffusers-like dir.
+    source_allow_patterns = cast(
+        list[str] | None, manifest.get("source_allow_patterns")
+    )
+    source_dir = base_model_download_fn(
+        source_model_id,
+        local_dir=local_dir,
+        download=download,
+        allow_patterns=source_allow_patterns or allow_patterns,
+        force_diffusers_model=False,
+        skip_overlay_resolution=True,
+    )
+    source_dir = ensure_overlay_source_dir_complete(
+        source_model_id=source_model_id,
+        source_dir=source_dir,
+        manifest=manifest,
+        local_dir=local_dir,
+        allow_patterns=allow_patterns,
+        download=download,
+        snapshot_download_fn=snapshot_download_fn,
+    )
+    return materialize_overlay_model(
+        source_model_id=source_model_id,
+        overlay_spec=overlay_spec,
+        overlay_dir=overlay_dir,
+        source_dir=source_dir,
+        verify_diffusers_model_complete_fn=verify_diffusers_model_complete_fn,
+    )