sgl-project--sglang/python/sglang/multimodal_gen/runtime/utils/model_overlay.py

# 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,
    )