From 667389c50f2a2bb62af55df4b3585bb0642f3aff Mon Sep 17 00:00:00 2001 From: Emil Bogomolov Date: Tue, 18 Aug 2026 01:05:05 -0700 Subject: [PATCH] [diffusion] chore: filter transformer safetensors by index.json to drop duplicate shard variants (#35107) Co-authored-by: Emil Bogomolov --- .../runtime/loader/transformer_load_utils.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index 56f89aefa..9967cd38e 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py @@ -14,6 +14,7 @@ from functools import partial from typing import Callable, Optional import torch +from diffusers.utils import SAFE_WEIGHTS_INDEX_NAME from torch import nn from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import ( @@ -21,6 +22,9 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config i _patch_nunchaku_scales, ) from sglang.multimodal_gen.runtime.loader.utils import _list_safetensors_files +from sglang.multimodal_gen.runtime.loader.weight_utils import ( + filter_duplicate_safetensors_files, +) from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( COMPONENT_OFFLOAD, ComponentResidencyError, @@ -471,6 +475,18 @@ def resolve_transformer_safetensors_to_load( else: safetensors_list = _list_safetensors_files(component_model_path) + if safetensors_list: + # Diffusers repos occasionally ship more than one shard split for the + # same checkpoint (e.g. a 4-way and an 8-way split side by side). The + # index file is the authoritative source for which files belong to + # the checkpoint that was actually exported; anything else is a + # leftover sibling variant. + safetensors_list = filter_duplicate_safetensors_files( + safetensors_list, + os.path.dirname(safetensors_list[0]), + SAFE_WEIGHTS_INDEX_NAME, + ) + safetensors_list = _prefer_mixed_safetensors_files(safetensors_list) safetensors_list = _filter_duplicate_precision_variant_safetensors(safetensors_list)