1
0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-23 09:09:01 +03:00

support loading pipeline from transformer style (flat) repo (#14096)

* Support root-hosted weights + add _allow_files download hook

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
YiYi Xu
2026-07-02 08:04:31 -10:00
committed by GitHub
parent 761d72b19c
commit 5dc153bc4a
4 changed files with 81 additions and 0 deletions

View File

@@ -69,6 +69,19 @@ DUMMY_MODULES_FOLDER = "diffusers.utils"
TRANSFORMERS_DUMMY_MODULES_FOLDER = "transformers.utils"
CONNECTED_PIPES_KEYS = ["prior"]
# Auxiliary (non-weight) files a transformers component saves next to its weights. Repos with a flat,
# transformers-style layout host a component's files at the repo root instead of in a subfolder, where the
# folder-based allow patterns of `DiffusionPipeline.download` would miss them. Root-hosted weights and
# `config.json` are matched by their own patterns, so only these auxiliary filenames need listing.
# Currently the set needed by DiffusionGemma — extend as new flat-layout pipelines require it.
TRANSFORMERS_COMPONENT_AUX_FILES = [
"chat_template.jinja",
"generation_config.json",
"processor_config.json",
"tokenizer.json",
"tokenizer_config.json",
]
logger = logging.get_logger(__name__)
LOADABLE_CLASSES = {
@@ -136,6 +149,8 @@ def is_safetensors_compatible(filenames, passed_components=None, folder_names=No
)
passed_components = passed_components or []
# only weight files matter for safetensors compatibility
filenames = filter_model_files(filenames)
if folder_names:
filenames = {f for f in filenames if os.path.split(f)[0] in folder_names}

View File

@@ -82,6 +82,7 @@ from .pipeline_loading_utils import (
CONNECTED_PIPES_KEYS,
CUSTOM_PIPELINE_FILE_NAME,
LOADABLE_CLASSES,
TRANSFORMERS_COMPONENT_AUX_FILES,
_download_dduf_file,
_fetch_class_library_tuple,
_get_custom_components_and_folders,
@@ -1751,6 +1752,11 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin):
p for p in allow_patterns if not (len(p.split("/")) == 2 and p.split("/")[0] in passed_components)
]
# Repos with a flat, transformers-style layout host a component's files at the repo root instead of
# in a subfolder, where the folder-based allow patterns above miss its auxiliary files (root-hosted
# weights are already included via `model_filenames`, root `config.json` via `CONFIG_NAME`).
allow_patterns += TRANSFORMERS_COMPONENT_AUX_FILES
if pipeline_class._load_connected_pipes:
allow_patterns.append("README.md")

View File

@@ -218,6 +218,44 @@ class IsSafetensorsCompatibleTests(unittest.TestCase):
]
self.assertFalse(is_safetensors_compatible(filenames))
def test_diffusers_is_compatible_no_components_safetensors(self):
filenames = [
"diffusion_pytorch_model.safetensors",
]
self.assertTrue(is_safetensors_compatible(filenames))
def test_diffusers_is_compatible_no_components_safetensors_only_variants(self):
filenames = [
"diffusion_pytorch_model.fp16.safetensors",
]
self.assertTrue(is_safetensors_compatible(filenames, variant="fp16"))
def test_transformers_is_compatible_weightless_subfolder(self):
# transformers-style flat layout: transformers-named weights at the root + a weight-less subfolder
filenames = [
"model.safetensors",
"scheduler/scheduler_config.json",
]
self.assertTrue(is_safetensors_compatible(filenames))
def test_transformers_is_not_compatible_weightless_subfolder(self):
# same flat layout but only .bin weights at the root -> not safetensors compatible
filenames = [
"pytorch_model.bin",
"scheduler/scheduler_config.json",
]
self.assertFalse(is_safetensors_compatible(filenames))
def test_transformers_is_compatible_sharded_root_weights(self):
# sharded transformers-style weights at the repo root (e.g. DiffusionGemma's model-00001-of-00011.safetensors)
filenames = [
"model-00001-of-00002.safetensors",
"model-00002-of-00002.safetensors",
"model.safetensors.index.json",
"scheduler/scheduler_config.json",
]
self.assertTrue(is_safetensors_compatible(filenames))
def test_is_compatible_mixed_variants(self):
filenames = [
"unet/diffusion_pytorch_model.fp16.safetensors",

View File

@@ -2083,6 +2083,28 @@ class PipelineSlowTests(unittest.TestCase):
# is not downloaded, but all the expected ones
assert not os.path.isfile(os.path.join(snapshot_dir, "big_array.npy"))
def test_download_flat_transformers_style_repo(self):
# Repos with a flat, transformers-style layout host a component's files at the repo root instead of in a
# subfolder (here `model` and `processor`; only `scheduler/` has a folder). The download patterns must
# pick up the transformers auxiliary files at the root, while unrelated root files are still skipped.
model_id = "hf-internal-testing/tiny-flat-transformers-style-pipe"
with tempfile.TemporaryDirectory() as tmpdirname:
snapshot_dir = DiffusionPipeline.download(model_id, cache_dir=tmpdirname, force_download=True)
assert os.path.isfile(os.path.join(snapshot_dir, "model.safetensors"))
assert os.path.isfile(os.path.join(snapshot_dir, CONFIG_NAME))
for aux_file in [
"tokenizer.json",
"tokenizer_config.json",
"processor_config.json",
"chat_template.jinja",
"generation_config.json",
]:
assert os.path.isfile(os.path.join(snapshot_dir, aux_file))
assert os.path.isfile(os.path.join(snapshot_dir, "scheduler", SCHEDULER_CONFIG_NAME))
# unrelated root files are still not downloaded
assert not os.path.isfile(os.path.join(snapshot_dir, "big_array.npy"))
def test_warning_unused_kwargs(self):
model_id = "hf-internal-testing/unet-pipeline-dummy"
logger = logging.get_logger("diffusers.pipelines")