[Bugfix] Prevent LTX-2.5 from silently loading unindexed Diffusers shards (#6234)

Signed-off-by: mglyn <1203789601@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -5,6 +5,7 @@
 Tests for the DiffusersPipelineLoader.
 """
 
+import json
 from types import SimpleNamespace
 
 import pytest
@@ -102,6 +103,138 @@ def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
     )
 
 
+@pytest.mark.parametrize("offline", [False, True])
+def test_prepare_weights_honors_component_index_and_explicit_override(tmp_path, mocker, offline):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    snapshot = tmp_path / "snapshot"
+    transformer = snapshot / "transformer"
+    transformer.mkdir(parents=True)
+    indexed_files = [f"diffusion_pytorch_model-{index:05d}-of-00002.safetensors" for index in (1, 2)]
+    stale_file = "diffusion_pytorch_model-00001-of-00008.safetensors"
+    index_path = transformer / "diffusion_pytorch_model.safetensors.index.json"
+    index_path.write_text(
+        json.dumps({"weight_map": {f"weight.{index}": filename for index, filename in enumerate(indexed_files)}})
+    )
+    for filename in indexed_files + [stale_file]:
+        (transformer / filename).touch()
+
+    mocker.patch.object(loader_mod.huggingface_hub.constants, "HF_HUB_OFFLINE", offline)
+
+    def download_index(*, filename, **_kwargs):
+        if filename == "transformer/diffusion_pytorch_model.safetensors.index.json":
+            return str(index_path)
+        raise loader_mod.huggingface_hub.errors.EntryNotFoundError(filename)
+
+    hub_api = mocker.Mock()
+    hub_api.hf_hub_download.side_effect = download_index
+    mocker.patch.object(loader_mod, "hf_api", return_value=hub_api)
+    indexed_download = mocker.patch.object(loader_mod, "download_weights_from_hf_specific", return_value=str(snapshot))
+    generic_download = mocker.patch.object(loader_mod, "download_weights_from_hf", return_value=str(snapshot))
+
+    loader = _make_loader_with_weights([])
+    cache_dir = str(tmp_path / "cache")
+    loader.load_config.download_dir = cache_dir
+    folder, files, use_safetensors = loader._prepare_weights(
+        "org/model",
+        subfolder="transformer",
+        revision="revision",
+        fall_back_to_pt=True,
+        allow_patterns_overrides=None,
+    )
+
+    assert folder == str(transformer)
+    assert [str(transformer / filename) for filename in indexed_files] == files
+    assert use_safetensors
+    assert hub_api.hf_hub_download.call_count == 2
+    hub_api.hf_hub_download.assert_any_call(
+        repo_id="org/model",
+        filename="transformer/diffusion_pytorch_model.safetensors.index.json",
+        cache_dir=cache_dir,
+        revision="revision",
+        local_files_only=offline,
+    )
+    indexed_download.assert_called_once_with(
+        model_name_or_path="org/model",
+        cache_dir=cache_dir,
+        allow_patterns=[f"transformer/{filename}" for filename in indexed_files],
+        revision="revision",
+        ignore_patterns=loader.load_config.ignore_patterns,
+        require_all=True,
+    )
+
+    _, override_files, _ = loader._prepare_weights(
+        "org/model",
+        subfolder="transformer",
+        revision="revision",
+        fall_back_to_pt=True,
+        allow_patterns_overrides=[stale_file],
+    )
+    assert override_files == [str(transformer / stale_file)]
+    generic_download.assert_called_once_with(
+        "org/model",
+        cache_dir,
+        [stale_file],
+        "revision",
+        subfolder="transformer",
+        ignore_patterns=loader.load_config.ignore_patterns,
+    )
+
+
+def test_prepare_local_weights_honors_component_index(tmp_path):
+    transformer = tmp_path / "transformer"
+    transformer.mkdir()
+    indexed_file = "diffusion_pytorch_model-00001-of-00001.safetensors"
+    stale_file = "diffusion_pytorch_model-00001-of-00008.safetensors"
+    (transformer / "diffusion_pytorch_model.safetensors.index.json").write_text(
+        json.dumps({"weight_map": {"weight": indexed_file}})
+    )
+    for filename in (indexed_file, stale_file):
+        (transformer / filename).touch()
+
+    folder, files, use_safetensors = _make_loader_with_weights([])._prepare_weights(
+        tmp_path,
+        subfolder="transformer",
+        revision=None,
+        fall_back_to_pt=True,
+        allow_patterns_overrides=None,
+    )
+
+    assert folder == str(transformer)
+    assert files == [str(transformer / indexed_file)]
+    assert use_safetensors
+
+
+def test_prepare_weights_rejects_polluted_offline_cache_without_index(tmp_path, mocker):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    snapshot = tmp_path / "snapshot"
+    transformer = snapshot / "transformer"
+    transformer.mkdir(parents=True)
+    for shard_count in (4, 8):
+        for shard in range(1, shard_count + 1):
+            (transformer / f"diffusion_pytorch_model-{shard:05d}-of-{shard_count:05d}.safetensors").touch()
+
+    mocker.patch.object(loader_mod.huggingface_hub.constants, "HF_HUB_OFFLINE", True)
+    hub_api = mocker.Mock()
+    hub_api.hf_hub_download.side_effect = loader_mod.huggingface_hub.errors.EntryNotFoundError("index is not cached")
+    mocker.patch.object(loader_mod, "hf_api", return_value=hub_api)
+    mocker.patch.object(loader_mod, "download_weights_from_hf", return_value=str(snapshot))
+
+    loader = _make_loader_with_weights([])
+    with pytest.raises(ValueError, match="conflicting shard totals"):
+        loader._prepare_weights(
+            "org/model",
+            subfolder="transformer",
+            revision="revision",
+            fall_back_to_pt=True,
+            allow_patterns_overrides=None,
+        )
+
+    assert hub_api.hf_hub_download.call_count == 2
+    assert all(call.kwargs["local_files_only"] for call in hub_api.hf_hub_download.call_args_list)
+
+
 def test_strict_check_only_validates_source_prefix_parameters():
     model = _DummyPipelineModel(source_prefix="transformer.")
     loader = _make_loader_with_weights(["transformer.weight"])

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -3,6 +3,7 @@
 import contextlib
 import dataclasses
 import glob
+import json
 import os
 import re
 import time
@@ -10,21 +11,20 @@ from collections.abc import Generator, Iterable, Sequence
 from pathlib import Path
 from typing import cast
 
+import huggingface_hub
 import torch
 from torch import nn
 from vllm.config.load import LoadConfig
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
 from vllm.model_executor.model_loader.weight_utils import (
-    download_safetensors_index_file_from_hf,
     download_weights_from_hf,
-    filter_duplicate_safetensors_files,
     filter_files_not_needed_for_inference,
     maybe_download_from_modelscope,
     multi_thread_safetensors_weights_iterator,
     safetensors_weights_iterator,
 )
-from vllm.transformers_utils.repo_utils import file_exists
+from vllm.transformers_utils.repo_utils import hf_api
 from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.utils.torch_utils import set_default_torch_dtype
 
@@ -42,6 +42,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
 from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 from vllm_omni.diffusion.registry import initialize_model
+from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 
 
 # download_gguf was removed from upstream vLLM (commit 6635279d8).
@@ -88,6 +89,23 @@ def _natural_sort_key(filepath: str) -> list:
 DIFFUSION_MODEL_WEIGHTS_INDEX = "diffusion_pytorch_model.safetensors.index.json"
 TRANSFORMER_WEIGHTS_INDEX = "model.safetensors.index.json"
 INDEX_FILES = [DIFFUSION_MODEL_WEIGHTS_INDEX, TRANSFORMER_WEIGHTS_INDEX]
+SHARDED_SAFETENSORS_PATTERN = re.compile(r"^(?P<family>.+)-\d+-of-(?P<count>\d+)\.safetensors$")
+
+
+def _validate_unindexed_safetensors_layout(weight_files: Sequence[str]) -> None:
+    """Reject ambiguous shard families when no index is available."""
+    shard_counts: dict[str, set[int]] = {}
+    for weight_file in weight_files:
+        match = SHARDED_SAFETENSORS_PATTERN.fullmatch(os.path.basename(weight_file))
+        if match is not None:
+            shard_counts.setdefault(match.group("family"), set()).add(int(match.group("count")))
+
+    conflicts = {family: sorted(counts) for family, counts in shard_counts.items() if len(counts) > 1}
+    if conflicts:
+        raise ValueError(
+            "Ambiguous unindexed safetensors checkpoint with conflicting shard totals: "
+            f"{conflicts}. Refusing to load potentially stale checkpoint shards."
+        )
 
 
 def _resolve_custom_pipeline_cls(custom_pipeline_name: str | type | None) -> type:
@@ -148,6 +166,60 @@ class DiffusersPipelineLoader:
         self.host_weight_plan = None
         return plan
 
+    @staticmethod
+    def _repo_relative_path(subfolder: str | None, filename: str) -> str:
+        if subfolder is None:
+            return filename
+        prefix = f"{subfolder.rstrip('/')}/"
+        return filename if filename.startswith(prefix) else f"{prefix}{filename.lstrip('/')}"
+
+    def _resolve_weight_index(
+        self,
+        model_name_or_path: Path | str,
+        subfolder: str | None,
+        revision: str | None,
+    ) -> list[str] | None:
+        """Resolve an index and return its authoritative shard manifest."""
+        is_local = os.path.isdir(model_name_or_path)
+        index_paths: list[tuple[str, Path]] = []
+        for index_file in INDEX_FILES:
+            repo_index_path = self._repo_relative_path(subfolder, index_file)
+            if is_local:
+                index_path = Path(model_name_or_path) / repo_index_path
+                if index_path.is_file():
+                    index_paths.append((index_file, index_path))
+                continue
+
+            try:
+                index_path = hf_api().hf_hub_download(
+                    repo_id=str(model_name_or_path),
+                    filename=repo_index_path,
+                    cache_dir=self.load_config.download_dir,
+                    revision=revision,
+                    local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
+                )
+            except huggingface_hub.errors.EntryNotFoundError:
+                continue
+            index_paths.append((index_file, Path(index_path)))
+
+        if len(index_paths) > 1:
+            raise ValueError(
+                f"Multiple index files found in {model_name_or_path} with subfolder {subfolder}: "
+                f"{[index_file for index_file, _ in index_paths]}"
+            )
+        if not index_paths:
+            return None
+
+        index_file, index_path = index_paths[0]
+        with open(index_path) as f:
+            index = json.load(f)
+        weight_map = index.get("weight_map") if isinstance(index, dict) else None
+        if not isinstance(weight_map, dict) or not weight_map:
+            raise ValueError(f"Weight index {index_file} must contain a non-empty `weight_map`")
+        if not all(isinstance(filename, str) and filename for filename in weight_map.values()):
+            raise ValueError(f"Weight index {index_file} contains an invalid shard filename")
+        return sorted(set(weight_map.values()))
+
     def _prepare_weights(
         self,
         model_name_or_path: Path | str,
@@ -164,17 +236,11 @@ class DiffusersPipelineLoader:
         is_local = os.path.isdir(model_name_or_path)
         load_format = self.load_config.load_format
         use_safetensors = False
-        possible_index_files = [
-            f"{subfolder}/{index_file}" if subfolder is not None else index_file for index_file in INDEX_FILES
-        ]
-        available_index_file = [
-            f for f in possible_index_files if file_exists(model_name_or_path, f, revision=revision)
-        ]
-        if len(available_index_file) > 1:
-            raise ValueError(
-                f"Multiple index files found in {model_name_or_path} with subfolder {subfolder}: {available_index_file}"
-            )
-        index_file = available_index_file[0] if available_index_file else ""
+        indexed_weight_files = (
+            self._resolve_weight_index(model_name_or_path, subfolder, revision)
+            if allow_patterns_overrides is None
+            else None
+        )
 
         # only hf is supported currently
         if load_format == "auto":
@@ -192,7 +258,16 @@ class DiffusersPipelineLoader:
         if allow_patterns_overrides is not None:
             allow_patterns = allow_patterns_overrides
 
-        if not is_local:
+        if not is_local and indexed_weight_files is not None:
+            hf_folder = download_weights_from_hf_specific(
+                model_name_or_path=str(model_name_or_path),
+                cache_dir=self.load_config.download_dir,
+                allow_patterns=[self._repo_relative_path(subfolder, filename) for filename in indexed_weight_files],
+                revision=revision,
+                ignore_patterns=self.load_config.ignore_patterns,
+                require_all=True,
+            )
+        elif not is_local:
             hf_folder = download_weights_from_hf(
                 model_name_or_path,
                 self.load_config.download_dir,
@@ -207,31 +282,24 @@ class DiffusersPipelineLoader:
         if subfolder is not None:
             hf_folder = os.path.join(hf_folder, subfolder)
 
-        hf_weights_files: list[str] = []
-        for pattern in allow_patterns:
-            hf_weights_files += glob.glob(os.path.join(hf_folder, pattern))
-            if len(hf_weights_files) > 0:
-                # Decide by actual files rather than pattern name (patterns may include subfolders).
-                use_safetensors = any(f.endswith(".safetensors") for f in hf_weights_files)
-                break
-
-        if use_safetensors:
-            # For models like Mistral-7B-Instruct-v0.3
-            # there are both sharded safetensors files and a consolidated
-            # safetensors file. Using both breaks.
-            # Here, we download the `model.safetensors.index.json` and filter
-            # any files not found in the index.
-            if not is_local:
-                download_safetensors_index_file_from_hf(
-                    model_name_or_path,
-                    index_file,
-                    cache_dir=self.load_config.download_dir,
-                    subfolder=subfolder,
-                    revision=revision,
-                )
-            hf_weights_files = filter_duplicate_safetensors_files(hf_weights_files, hf_folder, index_file)
+        if indexed_weight_files is not None:
+            hf_weights_files = [os.path.join(hf_folder, filename) for filename in indexed_weight_files]
+            missing_files = [filename for filename in hf_weights_files if not os.path.isfile(filename)]
+            if missing_files:
+                raise FileNotFoundError(f"Weight files referenced in index but missing: {missing_files}")
+            use_safetensors = any(filename.endswith(".safetensors") for filename in hf_weights_files)
         else:
-            hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files)
+            hf_weights_files = []
+            for pattern in allow_patterns:
+                hf_weights_files += glob.glob(os.path.join(hf_folder, pattern))
+                if hf_weights_files:
+                    # Decide by actual files rather than pattern name (patterns may include subfolders).
+                    use_safetensors = any(f.endswith(".safetensors") for f in hf_weights_files)
+                    break
+            if use_safetensors:
+                _validate_unindexed_safetensors_layout(hf_weights_files)
+            else:
+                hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files)
 
         if len(hf_weights_files) == 0:
             raise RuntimeError(f"Cannot find any model weights with `{model_name_or_path}`")