[diffusion] chore: filter transformer safetensors by index.json to drop duplicate shard variants (#35107)
Co-authored-by: Emil Bogomolov <zetyquickly@googlemail.com>
This commit is contained in:
co-authored by
Emil Bogomolov
parent
0111b29031
commit
667389c50f
@@ -14,6 +14,7 @@ from functools import partial
|
|||||||
from typing import Callable, Optional
|
from typing import Callable, Optional
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
from diffusers.utils import SAFE_WEIGHTS_INDEX_NAME
|
||||||
from torch import nn
|
from torch import nn
|
||||||
|
|
||||||
from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
|
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,
|
_patch_nunchaku_scales,
|
||||||
)
|
)
|
||||||
from sglang.multimodal_gen.runtime.loader.utils import _list_safetensors_files
|
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 (
|
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
|
||||||
COMPONENT_OFFLOAD,
|
COMPONENT_OFFLOAD,
|
||||||
ComponentResidencyError,
|
ComponentResidencyError,
|
||||||
@@ -471,6 +475,18 @@ def resolve_transformer_safetensors_to_load(
|
|||||||
else:
|
else:
|
||||||
safetensors_list = _list_safetensors_files(component_model_path)
|
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 = _prefer_mixed_safetensors_files(safetensors_list)
|
||||||
safetensors_list = _filter_duplicate_precision_variant_safetensors(safetensors_list)
|
safetensors_list = _filter_duplicate_precision_variant_safetensors(safetensors_list)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user