From aa69acbd090d4b6e021ed2bc22966fca03504572 Mon Sep 17 00:00:00 2001 From: Mu GuanLin <1203789601@qq.com> Date: Mon, 24 Aug 2026 20:57:24 +0800 Subject: [PATCH] [Bugfix] Prevent LTX-2.5 from silently loading unindexed Diffusers shards (#6234) Signed-off-by: mglyn <1203789601@qq.com> Co-authored-by: Hongsheng Liu --- .../model_loader/test_diffusers_loader.py | 133 ++++++++++++++++ .../model_loader/diffusers_loader.py | 146 +++++++++++++----- 2 files changed, 240 insertions(+), 39 deletions(-) diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 527942961..ce3c99418 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/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"]) diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 62bc0525f..a7a41088d 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/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.+)-\d+-of-(?P\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}`")