[diffusion] refactor: resolve lora weight sources deterministically (#35774)

docs/docs/sglang-diffusion/api/cli.mdx

@@ -79,7 +79,7 @@ Use `sglang generate --help` and `sglang serve --help` for the full argument lis
 - `--model-variant {NAME}`: semantic checkpoint variant to load when one model repository contains multiple weight partitions. The pipeline maps this stable name to the repository layout before loading; for example, MiniMax-H3 accepts `fl2va` and `ref2va`. This is a server/load-time choice, unlike a request's `task`.
 - `--minimax-h3-adaln-cache-path {FILE}`: advanced MiniMax-H3-only inference cache. It replaces the checkpoint's AdaLN projection weights with precomputed outputs and only accepts requests whose exact FP32 timestep plan is included in the cache. It requires unquantized weights and the matching model variant.
 - `--model-subfolder {PATH}`: advanced direct override for a component subfolder inside the model repository. Prefer `--model-variant` when the pipeline exposes semantic routing. If both are supplied, they must resolve to the same weight partition.
-- `--lora-path {PATH}` and `--lora-nickname {NAME}`: load a LoRA adapter
+- `--lora-path {PATH}` and `--lora-nickname {NAME}`: load a LoRA adapter from a local path, Hugging Face repo/subfolder, or exact Hub file URL
 - `--lora-weight-name {FILE}`: select one adapter file from a repository that contains multiple LoRA revisions. The Hub download is filtered to that file plus JSON metadata, so unused weights are not downloaded.
 - `--lora-alpha {N}`: supply the training alpha when a single-file adapter omits both per-layer alpha tensors and `adapter_config.json`. Do not set it when the adapter already records alpha metadata.
 - `--lora-merge-mode {auto|merge|dynamic}`: choose how LoRA is applied. `auto` statically merges regular weights and uses dynamic LoRA for FSDP-sharded weights to avoid full-gather peaks.

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

@@ -52,6 +52,7 @@ from sglang.multimodal_gen.runtime.utils.model_overlay import (
 from sglang.multimodal_gen.runtime.utils.quantization_utils import (
     normalize_flat_modelopt_quant_config,
 )
+from sglang.multimodal_gen.runtime.weights.source import resolve_weight
 from sglang.srt.environ import envs
 from sglang.srt.utils.hf_transformers import check_gguf_file
 from sglang.utils import is_in_ci
@@ -616,42 +617,66 @@ def maybe_download_lora(
     Returns:
         Local path to the model
     """
-    # Repositories often publish several adapter revisions side by side.  If a
-    # filename is pinned, do not download every weight before selecting it.
-    # Keep JSON metadata so PEFT's lora_alpha remains available.
-    allow_patterns = (
-        ["*.json", weight_name, f"**/{weight_name}"]
-        if weight_name is not None
-        else ["*.json", "*.safetensors", "*.bin"]
-    )
+    if envs.SGLANG_USE_MODELSCOPE.get():
+        allow_patterns = (
+            ["*.json", weight_name, f"**/{weight_name}"]
+            if weight_name is not None
+            else ["*.json", "*.safetensors", "*.bin"]
+        )
+        local_path = maybe_download_model(
+            model_name_or_path,
+            local_dir,
+            download,
+            is_lora=True,
+            allow_patterns=allow_patterns,
+        )
+        if os.path.isfile(local_path):
+            return local_path
+        if weight_name is not None:
+            target = os.path.join(local_path, weight_name)
+            if not os.path.isfile(target):
+                raise FileNotFoundError(
+                    f"Specified lora_weight_name '{weight_name}' not found in "
+                    f"{local_path}"
+                )
+            return target
+        guessed = _best_guess_weight_name(local_path, file_extension=".safetensors")
+        if guessed is None and current_platform.is_rocm():
+            guessed = _best_guess_weight_name(
+                model_name_or_path, file_extension=".safetensors"
+            )
+        return os.path.join(local_path, guessed)
 
+    resolved_weight = resolve_weight(model_name_or_path, weight_name=weight_name)
+    selected_file = resolved_weight.selected_file
+    if not selected_file.endswith(".safetensors"):
+        raise ValueError(
+            "Native diffusion LoRA loading requires a safetensors file, got "
+            f"{selected_file!r}"
+        )
+
+    source = resolved_weight.inventory.source
+    if source.kind == "local":
+        assert source.local_path is not None
+        if os.path.isfile(source.local_path):
+            return source.local_path
+        return os.path.join(source.local_path, selected_file)
+
+    assert source.repo_id is not None
     local_path = maybe_download_model(
-        model_name_or_path,
+        source.repo_id,
         local_dir,
         download,
         is_lora=True,
-        allow_patterns=allow_patterns,
+        allow_patterns=["*.json", selected_file],
+        revision=resolved_weight.inventory.resolved_revision or source.revision,
     )
-    # return directly if local_path is a file
-    if os.path.isfile(local_path):
-        return local_path
-
-    if weight_name is not None:
-        target = os.path.join(local_path, weight_name)
-        if not os.path.isfile(target):
-            raise FileNotFoundError(
-                f"Specified lora_weight_name '{weight_name}' not found in {local_path}"
-            )
-        return target
-
-    guessed = _best_guess_weight_name(local_path, file_extension=".safetensors")
-    # AMD workaround: PR 15813 changed from model_name_or_path to local_path,
-    # which can return None. Fall back to original behavior on ROCm.
-    if guessed is None and current_platform.is_rocm():
-        guessed = _best_guess_weight_name(
-            model_name_or_path, file_extension=".safetensors"
+    target = os.path.join(local_path, selected_file)
+    if not os.path.isfile(target):
+        raise FileNotFoundError(
+            f"Resolved LoRA weight {selected_file!r} was not downloaded to {local_path}"
         )
-    return os.path.join(local_path, guessed)
+    return target
 
 
 def verify_model_config_and_directory(model_path: str) -> dict[str, Any]:

python/sglang/multimodal_gen/runtime/weights/__init__.py

@@ -0,0 +1 @@
+"""Weight source contracts shared by runtime loaders."""

python/sglang/multimodal_gen/runtime/weights/source.py

@@ -0,0 +1,266 @@
+"""Resolve weight sources for runtime loaders."""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from pathlib import Path, PurePosixPath
+from typing import Literal
+from urllib.parse import unquote, urlparse
+
+from huggingface_hub import HfApi
+from huggingface_hub.utils import validate_repo_id
+
+WeightSourceKind = Literal["local", "huggingface"]
+
+_WEIGHT_SUFFIXES = (".safetensors", ".gguf", ".bin", ".pt", ".pth", ".ckpt")
+
+
+@dataclass(frozen=True)
+class WeightSource:
+    original: str
+    kind: WeightSourceKind
+    local_path: str | None = None
+    repo_id: str | None = None
+    revision: str | None = None
+    subfolder: str | None = None
+    filename: str | None = None
+
+
+@dataclass(frozen=True)
+class WeightInventory:
+    source: WeightSource
+    resolved_revision: str | None
+    files: tuple[str, ...]
+
+
+@dataclass(frozen=True)
+class ResolvedWeight:
+    inventory: WeightInventory
+    selected_file: str
+
+
+def _validate_relative_hub_path(path: str, field_name: str) -> str:
+    normalized = str(PurePosixPath(path))
+    pure_path = PurePosixPath(normalized)
+    if not path or pure_path.is_absolute() or ".." in pure_path.parts:
+        raise ValueError(f"Invalid Hugging Face {field_name}: {path!r}")
+    return normalized
+
+
+def _merge_revision(url_revision: str | None, revision: str | None) -> str | None:
+    if url_revision is not None and revision is not None and url_revision != revision:
+        raise ValueError(
+            f"Weight URL pins revision {url_revision!r}, which conflicts with "
+            f"revision {revision!r}"
+        )
+    return url_revision or revision
+
+
+def _parse_huggingface_url(source: str, revision: str | None) -> WeightSource:
+    parsed = urlparse(source)
+    if parsed.netloc.lower() not in ("huggingface.co", "www.huggingface.co"):
+        raise ValueError(
+            "Only huggingface.co weight URLs are supported; use a local path "
+            "or an owner/repo reference for other sources"
+        )
+
+    raw_parts = [part for part in parsed.path.split("/") if part]
+    if raw_parts and raw_parts[0] in ("datasets", "spaces"):
+        raise ValueError("Diffusion weights must come from a Hugging Face model repo")
+    if len(raw_parts) < 2:
+        raise ValueError(f"Hugging Face weight URL has no model repo: {source!r}")
+
+    repo_id = "/".join(unquote(part) for part in raw_parts[:2])
+    validate_repo_id(repo_id)
+    action = raw_parts[2] if len(raw_parts) > 2 else None
+    if action is None:
+        return WeightSource(
+            original=source,
+            kind="huggingface",
+            repo_id=repo_id,
+            revision=revision,
+        )
+    if action not in ("tree", "blob", "resolve") or len(raw_parts) < 4:
+        raise ValueError(f"Unsupported Hugging Face weight URL: {source!r}")
+
+    url_revision = unquote(raw_parts[3])
+    selected_revision = _merge_revision(url_revision, revision)
+    tail = "/".join(unquote(part) for part in raw_parts[4:])
+    if action == "tree":
+        subfolder = _validate_relative_hub_path(tail, "subfolder") if tail else None
+        return WeightSource(
+            original=source,
+            kind="huggingface",
+            repo_id=repo_id,
+            revision=selected_revision,
+            subfolder=subfolder,
+        )
+    if not tail:
+        raise ValueError(f"Hugging Face weight URL has no filename: {source!r}")
+    return WeightSource(
+        original=source,
+        kind="huggingface",
+        repo_id=repo_id,
+        revision=selected_revision,
+        filename=_validate_relative_hub_path(tail, "filename"),
+    )
+
+
+def parse_weight_source(
+    source: str,
+    *,
+    revision: str | None = None,
+) -> WeightSource:
+    """Parse local paths, Hub repo IDs, subfolders, and exact Hub URLs."""
+    expanded = os.path.expanduser(source)
+    parsed = urlparse(source)
+    if parsed.scheme in ("http", "https"):
+        return _parse_huggingface_url(source, revision)
+
+    looks_local = (
+        os.path.exists(expanded)
+        or os.path.isabs(expanded)
+        or source.startswith(("./", "../", "~"))
+    )
+    if looks_local:
+        return WeightSource(
+            original=source,
+            kind="local",
+            local_path=os.path.abspath(expanded),
+        )
+
+    parts = source.split("/")
+    if len(parts) < 2 or not all(parts[:2]):
+        raise ValueError(
+            f"Weight source {source!r} is neither a local path nor an "
+            "owner/repo Hugging Face reference"
+        )
+    repo_id = "/".join(parts[:2])
+    validate_repo_id(repo_id)
+    tail = "/".join(parts[2:]) or None
+    filename = (
+        _validate_relative_hub_path(tail, "filename")
+        if tail is not None and tail.lower().endswith(_WEIGHT_SUFFIXES)
+        else None
+    )
+    subfolder = tail if filename is None else None
+    if subfolder is not None:
+        subfolder = _validate_relative_hub_path(subfolder, "subfolder")
+    return WeightSource(
+        original=source,
+        kind="huggingface",
+        repo_id=repo_id,
+        revision=revision,
+        subfolder=subfolder,
+        filename=filename,
+    )
+
+
+def _filter_inventory_files(
+    files: tuple[str, ...], source: WeightSource
+) -> tuple[str, ...]:
+    if source.filename is not None:
+        selected = tuple(path for path in files if path == source.filename)
+        if not selected:
+            raise FileNotFoundError(
+                f"Weight file {source.filename!r} was not found in {source.repo_id}"
+            )
+        return selected
+    if source.subfolder is None:
+        return files
+    prefix = source.subfolder.rstrip("/") + "/"
+    selected = tuple(path for path in files if path.startswith(prefix))
+    if not selected:
+        raise FileNotFoundError(
+            f"Weight subfolder {source.subfolder!r} was not found in {source.repo_id}"
+        )
+    return selected
+
+
+def resolve_weight_inventory(source: WeightSource) -> WeightInventory:
+    """List source files and pin a remote source to an immutable revision."""
+    if source.kind == "local":
+        assert source.local_path is not None
+        local_path = Path(source.local_path)
+        if not local_path.exists():
+            raise FileNotFoundError(f"Weight path does not exist: {local_path}")
+        if local_path.is_file():
+            files = (local_path.name,)
+        else:
+            files = tuple(
+                path.relative_to(local_path).as_posix()
+                for path in sorted(local_path.rglob("*"))
+                if path.is_file()
+            )
+        return WeightInventory(
+            source=source,
+            resolved_revision=None,
+            files=files,
+        )
+
+    assert source.repo_id is not None
+    model_info = HfApi().model_info(
+        source.repo_id,
+        revision=source.revision,
+    )
+    files = tuple(sibling.rfilename for sibling in model_info.siblings)
+    return WeightInventory(
+        source=source,
+        resolved_revision=model_info.sha,
+        files=_filter_inventory_files(files, source),
+    )
+
+
+def _select_named_file(candidates: tuple[str, ...], weight_name: str) -> str:
+    exact = tuple(path for path in candidates if path == weight_name)
+    if exact:
+        return exact[0]
+    basename_matches = tuple(
+        path for path in candidates if PurePosixPath(path).name == weight_name
+    )
+    if len(basename_matches) == 1:
+        return basename_matches[0]
+    if not basename_matches:
+        raise FileNotFoundError(f"Requested weight {weight_name!r} was not found")
+    raise ValueError(
+        f"Weight name {weight_name!r} matches multiple files: "
+        f"{list(basename_matches)}"
+    )
+
+
+def select_weight_file(
+    inventory: WeightInventory, weight_name: str | None = None
+) -> str:
+    """Select weights deterministically; never guess among independent files."""
+    candidates = tuple(
+        path for path in inventory.files if path.lower().endswith(_WEIGHT_SUFFIXES)
+    )
+    if inventory.source.filename is not None:
+        return inventory.files[0]
+    if weight_name is not None:
+        return _select_named_file(candidates, weight_name)
+    if len(candidates) == 1:
+        return candidates[0]
+    if not candidates:
+        raise FileNotFoundError("Source contains no recognized weight files")
+    raise ValueError(
+        "Source contains multiple independent weight files; select one with "
+        f"an exact file URL or weight name. Candidates: {list(candidates)}"
+    )
+
+
+def resolve_weight(
+    source: str,
+    *,
+    revision: str | None = None,
+    weight_name: str | None = None,
+) -> ResolvedWeight:
+    """Resolve one weight file without downloading its tensor payload."""
+    parsed_source = parse_weight_source(source, revision=revision)
+    inventory = resolve_weight_inventory(parsed_source)
+    selected_file = select_weight_file(inventory, weight_name)
+    return ResolvedWeight(
+        inventory=inventory,
+        selected_file=selected_file,
+    )

python/sglang/multimodal_gen/test/unit/test_lora_pipeline.py

@@ -142,20 +142,71 @@ def test_lora_alpha_override_updates_cached_adapter_scale():
     assert layer.lora_alpha == 8
 
 
-def test_pinned_lora_weight_limits_snapshot_download(tmp_path):
+def test_lora_tree_url_selects_one_pinned_weight(tmp_path):
     weight_name = "adapter-v4.safetensors"
-    weight_path = tmp_path / weight_name
+    adapter_dir = tmp_path / "adapters"
+    adapter_dir.mkdir()
+    weight_path = adapter_dir / weight_name
     weight_path.touch()
+    model_info = SimpleNamespace(
+        sha="immutable-sha",
+        siblings=[
+            SimpleNamespace(rfilename="adapters/adapter-v3.safetensors"),
+            SimpleNamespace(rfilename="adapters/adapter-v4.safetensors"),
+        ],
+    )
 
     download_target = (
         "sglang.multimodal_gen.runtime.utils.hf_diffusers_utils.maybe_download_model"
     )
-    with patch(download_target, return_value=str(tmp_path)) as download:
-        actual = maybe_download_lora("org/multi-adapter", weight_name=weight_name)
+    with (
+        patch(
+            "sglang.multimodal_gen.runtime.weights.source.HfApi.model_info",
+            return_value=model_info,
+        ),
+        patch(download_target, return_value=str(tmp_path)) as download,
+    ):
+        actual = maybe_download_lora(
+            "https://huggingface.co/org/multi-adapter/tree/main/adapters",
+            weight_name=weight_name,
+        )
+
+    assert actual == str(weight_path)
+    assert download.call_args.args[0] == "org/multi-adapter"
+    assert download.call_args.kwargs["revision"] == "immutable-sha"
+    assert download.call_args.kwargs["allow_patterns"] == [
+        "*.json",
+        f"adapters/{weight_name}",
+    ]
+
+
+def test_lora_exact_file_url_needs_no_weight_name(tmp_path):
+    weight_path = tmp_path / "adapter.safetensors"
+    weight_path.touch()
+    model_info = SimpleNamespace(
+        sha="immutable-sha",
+        siblings=[
+            SimpleNamespace(rfilename="adapter.safetensors"),
+            SimpleNamespace(rfilename="other.safetensors"),
+        ],
+    )
+
+    download_target = (
+        "sglang.multimodal_gen.runtime.utils.hf_diffusers_utils.maybe_download_model"
+    )
+    with (
+        patch(
+            "sglang.multimodal_gen.runtime.weights.source.HfApi.model_info",
+            return_value=model_info,
+        ),
+        patch(download_target, return_value=str(tmp_path)) as download,
+    ):
+        actual = maybe_download_lora(
+            "https://huggingface.co/org/multi-adapter/resolve/main/adapter.safetensors"
+        )
 
     assert actual == str(weight_path)
     assert download.call_args.kwargs["allow_patterns"] == [
         "*.json",
-        weight_name,
-        f"**/{weight_name}",
+        "adapter.safetensors",
     ]

python/sglang/multimodal_gen/test/unit/test_weight_source.py

@@ -0,0 +1,82 @@
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from sglang.multimodal_gen.runtime.weights.source import (
+    parse_weight_source,
+    resolve_weight,
+    resolve_weight_inventory,
+)
+
+
+def test_parse_weight_source_accepts_repo_subfolder_and_exact_url():
+    subfolder = parse_weight_source("owner/repo/text_encoder", revision="v1")
+    repo_file = parse_weight_source("owner/repo/adapter.safetensors")
+    exact_file = parse_weight_source(
+        "https://huggingface.co/owner/repo/resolve/main/weights/model.safetensors"
+    )
+
+    assert subfolder.repo_id == "owner/repo"
+    assert subfolder.subfolder == "text_encoder"
+    assert subfolder.revision == "v1"
+    assert repo_file.filename == "adapter.safetensors"
+    assert repo_file.subfolder is None
+    assert exact_file.repo_id == "owner/repo"
+    assert exact_file.revision == "main"
+    assert exact_file.filename == "weights/model.safetensors"
+
+
+def test_parse_weight_source_rejects_conflicting_url_revision():
+    with pytest.raises(ValueError, match="conflicts with revision"):
+        parse_weight_source(
+            "https://huggingface.co/owner/repo/tree/main/transformer",
+            revision="v2",
+        )
+
+
+def test_resolve_local_inventory_lists_files_without_loading_tensors(tmp_path):
+    component = tmp_path / "component"
+    component.mkdir()
+    (component / "config.json").write_text("{}")
+    (component / "model.safetensors").write_bytes(b"header-only-fixture")
+
+    inventory = resolve_weight_inventory(parse_weight_source(str(component)))
+
+    assert inventory.resolved_revision is None
+    assert list(inventory.files) == [
+        "config.json",
+        "model.safetensors",
+    ]
+
+
+def test_resolve_remote_inventory_pins_revision_and_filters_subfolder():
+    source = parse_weight_source("owner/repo/text_encoder", revision="main")
+    model_info = SimpleNamespace(
+        sha="immutable-sha",
+        siblings=[
+            SimpleNamespace(rfilename="text_encoder/config.json"),
+            SimpleNamespace(rfilename="text_encoder/model.safetensors"),
+            SimpleNamespace(rfilename="vae/config.json"),
+        ],
+    )
+
+    with patch(
+        "sglang.multimodal_gen.runtime.weights.source.HfApi.model_info",
+        return_value=model_info,
+    ):
+        inventory = resolve_weight_inventory(source)
+
+    assert inventory.resolved_revision == "immutable-sha"
+    assert inventory.files == (
+        "text_encoder/config.json",
+        "text_encoder/model.safetensors",
+    )
+
+
+def test_weight_source_rejects_ambiguous_files(tmp_path):
+    (tmp_path / "a.safetensors").write_bytes(b"a")
+    (tmp_path / "b.safetensors").write_bytes(b"b")
+
+    with pytest.raises(ValueError, match="multiple independent weight files"):
+        resolve_weight(str(tmp_path))