Skip to content

Commit a9477bb

Browse files
[Pipeline download] Improve pipeline download for index and passed co… (#2980)
* [Pipeline download] Improve pipeline download for index and passed components * correct * add more tests * up
1 parent ee20d1f commit a9477bb

File tree

2 files changed

+221
-38
lines changed

2 files changed

+221
-38
lines changed

src/diffusers/pipelines/pipeline_utils.py

Lines changed: 96 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,7 @@ class AudioPipelineOutput(BaseOutput):
134134
audios: np.ndarray
135135

136136

137-
def is_safetensors_compatible(filenames, variant=None) -> bool:
137+
def is_safetensors_compatible(filenames, variant=None, passed_components=None) -> bool:
138138
"""
139139
Checking for safetensors compatibility:
140140
- By default, all models are saved with the default pytorch serialization, so we use the list of default pytorch
@@ -150,9 +150,14 @@ def is_safetensors_compatible(filenames, variant=None) -> bool:
150150

151151
sf_filenames = set()
152152

153+
passed_components = passed_components or []
154+
153155
for filename in filenames:
154156
_, extension = os.path.splitext(filename)
155157

158+
if len(filename.split("/")) == 2 and filename.split("/")[0] in passed_components:
159+
continue
160+
156161
if extension == ".bin":
157162
pt_filenames.append(filename)
158163
elif extension == ".safetensors":
@@ -163,10 +168,8 @@ def is_safetensors_compatible(filenames, variant=None) -> bool:
163168
path, filename = os.path.split(filename)
164169
filename, extension = os.path.splitext(filename)
165170

166-
if filename == "pytorch_model":
167-
filename = "model"
168-
elif filename == f"pytorch_model.{variant}":
169-
filename = f"model.{variant}"
171+
if filename.startswith("pytorch_model"):
172+
filename = filename.replace("pytorch_model", "model")
170173
else:
171174
filename = filename
172175

@@ -196,24 +199,51 @@ def variant_compatible_siblings(filenames, variant=None) -> Union[List[os.PathLi
196199
weight_prefixes = [w.split(".")[0] for w in weight_names]
197200
# .bin, .safetensors, ...
198201
weight_suffixs = [w.split(".")[-1] for w in weight_names]
202+
# -00001-of-00002
203+
transformers_index_format = "\d{5}-of-\d{5}"
204+
205+
if variant is not None:
206+
# `diffusion_pytorch_model.fp16.bin` as well as `model.fp16-00001-of-00002.safetenstors`
207+
variant_file_re = re.compile(
208+
f"({'|'.join(weight_prefixes)})\.({variant}|{variant}-{transformers_index_format})\.({'|'.join(weight_suffixs)})$"
209+
)
210+
# `text_encoder/pytorch_model.bin.index.fp16.json`
211+
variant_index_re = re.compile(
212+
f"({'|'.join(weight_prefixes)})\.({'|'.join(weight_suffixs)})\.index\.{variant}\.json$"
213+
)
199214

200-
variant_file_regex = (
201-
re.compile(f"({'|'.join(weight_prefixes)})(.{variant}.)({'|'.join(weight_suffixs)})")
202-
if variant is not None
203-
else None
215+
# `diffusion_pytorch_model.bin` as well as `model-00001-of-00002.safetenstors`
216+
non_variant_file_re = re.compile(
217+
f"({'|'.join(weight_prefixes)})(-{transformers_index_format})?\.({'|'.join(weight_suffixs)})$"
204218
)
205-
non_variant_file_regex = re.compile(f"{'|'.join(weight_names)}")
219+
# `text_encoder/pytorch_model.bin.index.json`
220+
non_variant_index_re = re.compile(f"({'|'.join(weight_prefixes)})\.({'|'.join(weight_suffixs)})\.index\.json")
206221

207222
if variant is not None:
208-
variant_filenames = {f for f in filenames if variant_file_regex.match(f.split("/")[-1]) is not None}
223+
variant_weights = {f for f in filenames if variant_file_re.match(f.split("/")[-1]) is not None}
224+
variant_indexes = {f for f in filenames if variant_index_re.match(f.split("/")[-1]) is not None}
225+
variant_filenames = variant_weights | variant_indexes
209226
else:
210227
variant_filenames = set()
211228

212-
non_variant_filenames = {f for f in filenames if non_variant_file_regex.match(f.split("/")[-1]) is not None}
229+
non_variant_weights = {f for f in filenames if non_variant_file_re.match(f.split("/")[-1]) is not None}
230+
non_variant_indexes = {f for f in filenames if non_variant_index_re.match(f.split("/")[-1]) is not None}
231+
non_variant_filenames = non_variant_weights | non_variant_indexes
213232

233+
# all variant filenames will be used by default
214234
usable_filenames = set(variant_filenames)
235+
236+
def convert_to_variant(filename):
237+
if "index" in filename:
238+
variant_filename = filename.replace("index", f"index.{variant}")
239+
elif re.compile(f"^(.*?){transformers_index_format}").match(filename) is not None:
240+
variant_filename = f"{filename.split('-')[0]}.{variant}-{'-'.join(filename.split('-')[1:])}"
241+
else:
242+
variant_filename = f"{filename.split('.')[0]}.{variant}.{filename.split('.')[1]}"
243+
return variant_filename
244+
215245
for f in non_variant_filenames:
216-
variant_filename = f"{f.split('.')[0]}.{variant}.{f.split('.')[1]}"
246+
variant_filename = convert_to_variant(f)
217247
if variant_filename not in usable_filenames:
218248
usable_filenames.add(f)
219249

@@ -292,6 +322,27 @@ def get_class_obj_and_candidates(library_name, class_name, importable_classes, p
292322
return class_obj, class_candidates
293323

294324

325+
def _get_pipeline_class(class_obj, config, custom_pipeline=None, cache_dir=None, revision=None):
326+
if custom_pipeline is not None:
327+
if custom_pipeline.endswith(".py"):
328+
path = Path(custom_pipeline)
329+
# decompose into folder & file
330+
file_name = path.name
331+
custom_pipeline = path.parent.absolute()
332+
else:
333+
file_name = CUSTOM_PIPELINE_FILE_NAME
334+
335+
return get_class_from_dynamic_module(
336+
custom_pipeline, module_file=file_name, cache_dir=cache_dir, revision=revision
337+
)
338+
339+
if class_obj != DiffusionPipeline:
340+
return class_obj
341+
342+
diffusers_module = importlib.import_module(class_obj.__module__.split(".")[0])
343+
return getattr(diffusers_module, config["_class_name"])
344+
345+
295346
def load_sub_model(
296347
library_name: str,
297348
class_name: str,
@@ -779,7 +830,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
779830
device_map = kwargs.pop("device_map", None)
780831
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT)
781832
variant = kwargs.pop("variant", None)
782-
kwargs.pop("use_safetensors", None if is_safetensors_available() else False)
833+
use_safetensors = kwargs.pop("use_safetensors", None if is_safetensors_available() else False)
783834

784835
# 1. Download the checkpoints and configs
785836
# use snapshot download here to get it working from from_pretrained
@@ -794,8 +845,11 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
794845
use_auth_token=use_auth_token,
795846
revision=revision,
796847
from_flax=from_flax,
848+
use_safetensors=use_safetensors,
797849
custom_pipeline=custom_pipeline,
850+
custom_revision=custom_revision,
798851
variant=variant,
852+
**kwargs,
799853
)
800854
else:
801855
cached_folder = pretrained_model_name_or_path
@@ -810,29 +864,17 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
810864
for folder in os.listdir(cached_folder):
811865
folder_path = os.path.join(cached_folder, folder)
812866
is_folder = os.path.isdir(folder_path) and folder in config_dict
813-
variant_exists = is_folder and any(path.split(".")[1] == variant for path in os.listdir(folder_path))
867+
variant_exists = is_folder and any(
868+
p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)
869+
)
814870
if variant_exists:
815871
model_variants[folder] = variant
816872

817873
# 3. Load the pipeline class, if using custom module then load it from the hub
818874
# if we load from explicit class, let's use it
819-
if custom_pipeline is not None:
820-
if custom_pipeline.endswith(".py"):
821-
path = Path(custom_pipeline)
822-
# decompose into folder & file
823-
file_name = path.name
824-
custom_pipeline = path.parent.absolute()
825-
else:
826-
file_name = CUSTOM_PIPELINE_FILE_NAME
827-
828-
pipeline_class = get_class_from_dynamic_module(
829-
custom_pipeline, module_file=file_name, cache_dir=cache_dir, revision=custom_revision
830-
)
831-
elif cls != DiffusionPipeline:
832-
pipeline_class = cls
833-
else:
834-
diffusers_module = importlib.import_module(cls.__module__.split(".")[0])
835-
pipeline_class = getattr(diffusers_module, config_dict["_class_name"])
875+
pipeline_class = _get_pipeline_class(
876+
cls, config_dict, custom_pipeline=custom_pipeline, cache_dir=cache_dir, revision=custom_revision
877+
)
836878

837879
# DEPRECATED: To be removed in 1.0.0
838880
if pipeline_class.__name__ == "StableDiffusionInpaintPipeline" and version.parse(
@@ -1095,6 +1137,7 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]:
10951137
revision = kwargs.pop("revision", None)
10961138
from_flax = kwargs.pop("from_flax", False)
10971139
custom_pipeline = kwargs.pop("custom_pipeline", None)
1140+
custom_revision = kwargs.pop("custom_revision", None)
10981141
variant = kwargs.pop("variant", None)
10991142
use_safetensors = kwargs.pop("use_safetensors", None)
11001143

@@ -1153,7 +1196,7 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]:
11531196
# this enables downloading schedulers, tokenizers, ...
11541197
allow_patterns += [os.path.join(k, "*") for k in folder_names if k not in model_folder_names]
11551198
# also allow downloading config.json files with the model
1156-
allow_patterns += [os.path.join(k, "*.json") for k in model_folder_names]
1199+
allow_patterns += [os.path.join(k, "config.json") for k in model_folder_names]
11571200

11581201
allow_patterns += [
11591202
SCHEDULER_CONFIG_NAME,
@@ -1162,17 +1205,28 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]:
11621205
CUSTOM_PIPELINE_FILE_NAME,
11631206
]
11641207

1208+
# retrieve passed components that should not be downloaded
1209+
pipeline_class = _get_pipeline_class(
1210+
cls, config_dict, custom_pipeline=custom_pipeline, cache_dir=cache_dir, revision=custom_revision
1211+
)
1212+
expected_components, _ = cls._get_signature_keys(pipeline_class)
1213+
passed_components = [k for k in expected_components if k in kwargs]
1214+
11651215
if (
11661216
use_safetensors
11671217
and not allow_pickle
1168-
and not is_safetensors_compatible(model_filenames, variant=variant)
1218+
and not is_safetensors_compatible(
1219+
model_filenames, variant=variant, passed_components=passed_components
1220+
)
11691221
):
11701222
raise EnvironmentError(
11711223
f"Could not found the necessary `safetensors` weights in {model_filenames} (variant={variant})"
11721224
)
11731225
if from_flax:
11741226
ignore_patterns = ["*.bin", "*.safetensors", "*.onnx", "*.pb"]
1175-
elif use_safetensors and is_safetensors_compatible(model_filenames, variant=variant):
1227+
elif use_safetensors and is_safetensors_compatible(
1228+
model_filenames, variant=variant, passed_components=passed_components
1229+
):
11761230
ignore_patterns = ["*.bin", "*.msgpack"]
11771231

11781232
safetensors_variant_filenames = {f for f in variant_filenames if f.endswith(".safetensors")}
@@ -1194,6 +1248,13 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]:
11941248
f"\nA mixture of {variant} and non-{variant} filenames will be loaded.\nLoaded {variant} filenames:\n[{', '.join(bin_variant_filenames)}]\nLoaded non-{variant} filenames:\n[{', '.join(bin_model_filenames - bin_variant_filenames)}\nIf this behavior is not expected, please check your folder structure."
11951249
)
11961250

1251+
# Don't download any objects that are passed
1252+
allow_patterns = [
1253+
p for p in allow_patterns if not (len(p.split("/")) == 2 and p.split("/")[0] in passed_components)
1254+
]
1255+
# Don't download index files of forbidden patterns either
1256+
ignore_patterns = ignore_patterns + [f"{i}.index.*json" for i in ignore_patterns]
1257+
11971258
re_ignore_pattern = [re.compile(fnmatch.translate(p)) for p in ignore_patterns]
11981259
re_allow_pattern = [re.compile(fnmatch.translate(p)) for p in allow_patterns]
11991260

tests/test_pipelines.py

Lines changed: 125 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -78,9 +78,7 @@ def test_one_request_upon_cached(self):
7878

7979
with tempfile.TemporaryDirectory() as tmpdirname:
8080
with requests_mock.mock(real_http=True) as m:
81-
DiffusionPipeline.download(
82-
"hf-internal-testing/tiny-stable-diffusion-pipe", safety_checker=None, cache_dir=tmpdirname
83-
)
81+
DiffusionPipeline.download("hf-internal-testing/tiny-stable-diffusion-pipe", cache_dir=tmpdirname)
8482

8583
download_requests = [r.method for r in m.request_history]
8684
assert download_requests.count("HEAD") == 15, "15 calls to files"
@@ -101,6 +99,55 @@ def test_one_request_upon_cached(self):
10199
len(cache_requests) == 2
102100
), "We should call only `model_info` to check for _commit hash and `send_telemetry`"
103101

102+
def test_less_downloads_passed_object(self):
103+
with tempfile.TemporaryDirectory() as tmpdirname:
104+
cached_folder = DiffusionPipeline.download(
105+
"hf-internal-testing/tiny-stable-diffusion-pipe", safety_checker=None, cache_dir=tmpdirname
106+
)
107+
108+
# make sure safety checker is not downloaded
109+
assert "safety_checker" not in os.listdir(cached_folder)
110+
111+
# make sure rest is downloaded
112+
assert "unet" in os.listdir(cached_folder)
113+
assert "tokenizer" in os.listdir(cached_folder)
114+
assert "vae" in os.listdir(cached_folder)
115+
assert "model_index.json" in os.listdir(cached_folder)
116+
assert "scheduler" in os.listdir(cached_folder)
117+
assert "feature_extractor" in os.listdir(cached_folder)
118+
119+
def test_less_downloads_passed_object_calls(self):
120+
# TODO: For some reason this test fails on MPS where no HEAD call is made.
121+
if torch_device == "mps":
122+
return
123+
124+
with tempfile.TemporaryDirectory() as tmpdirname:
125+
with requests_mock.mock(real_http=True) as m:
126+
DiffusionPipeline.download(
127+
"hf-internal-testing/tiny-stable-diffusion-pipe", safety_checker=None, cache_dir=tmpdirname
128+
)
129+
130+
download_requests = [r.method for r in m.request_history]
131+
# 15 - 2 because no call to config or model file for `safety_checker`
132+
assert download_requests.count("HEAD") == 13, "13 calls to files"
133+
# 17 - 2 because no call to config or model file for `safety_checker`
134+
assert download_requests.count("GET") == 15, "13 calls to files + model_info + model_index.json"
135+
assert (
136+
len(download_requests) == 28
137+
), "2 calls per file (13 files) + send_telemetry, model_info and model_index.json"
138+
139+
with requests_mock.mock(real_http=True) as m:
140+
DiffusionPipeline.download(
141+
"hf-internal-testing/tiny-stable-diffusion-pipe", safety_checker=None, cache_dir=tmpdirname
142+
)
143+
144+
cache_requests = [r.method for r in m.request_history]
145+
assert cache_requests.count("HEAD") == 1, "model_index.json is only HEAD"
146+
assert cache_requests.count("GET") == 1, "model info is only GET"
147+
assert (
148+
len(cache_requests) == 2
149+
), "We should call only `model_info` to check for _commit hash and `send_telemetry`"
150+
104151
def test_download_only_pytorch(self):
105152
with tempfile.TemporaryDirectory() as tmpdirname:
106153
# pipeline has Flax weights
@@ -165,6 +212,54 @@ def test_download_safetensors(self):
165212
# https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe/blob/main/unet/diffusion_flax_model.msgpack
166213
assert not any(f.endswith(".bin") for f in files)
167214

215+
def test_download_safetensors_index(self):
216+
for variant in ["fp16", None]:
217+
with tempfile.TemporaryDirectory() as tmpdirname:
218+
tmpdirname = DiffusionPipeline.download(
219+
"hf-internal-testing/tiny-stable-diffusion-pipe-indexes",
220+
cache_dir=tmpdirname,
221+
use_safetensors=True,
222+
variant=variant,
223+
)
224+
225+
all_root_files = [t[-1] for t in os.walk(os.path.join(tmpdirname))]
226+
files = [item for sublist in all_root_files for item in sublist]
227+
228+
# None of the downloaded files should be a safetensors file even if we have some here:
229+
# https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe-indexes/tree/main/text_encoder
230+
if variant is None:
231+
assert not any("fp16" in f for f in files)
232+
else:
233+
model_files = [f for f in files if "safetensors" in f]
234+
assert all("fp16" in f for f in model_files)
235+
236+
assert len([f for f in files if ".safetensors" in f]) == 8
237+
assert not any(".bin" in f for f in files)
238+
239+
def test_download_bin_index(self):
240+
for variant in ["fp16", None]:
241+
with tempfile.TemporaryDirectory() as tmpdirname:
242+
tmpdirname = DiffusionPipeline.download(
243+
"hf-internal-testing/tiny-stable-diffusion-pipe-indexes",
244+
cache_dir=tmpdirname,
245+
use_safetensors=False,
246+
variant=variant,
247+
)
248+
249+
all_root_files = [t[-1] for t in os.walk(os.path.join(tmpdirname))]
250+
files = [item for sublist in all_root_files for item in sublist]
251+
252+
# None of the downloaded files should be a safetensors file even if we have some here:
253+
# https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe-indexes/tree/main/text_encoder
254+
if variant is None:
255+
assert not any("fp16" in f for f in files)
256+
else:
257+
model_files = [f for f in files if "bin" in f]
258+
assert all("fp16" in f for f in model_files)
259+
260+
assert len([f for f in files if ".bin" in f]) == 8
261+
assert not any(".safetensors" in f for f in files)
262+
168263
def test_download_no_safety_checker(self):
169264
prompt = "hello"
170265
pipe = StableDiffusionPipeline.from_pretrained(
@@ -362,6 +457,33 @@ def test_download_broken_variant(self):
362457

363458
diffusers.utils.import_utils._safetensors_available = True
364459

460+
def test_local_save_load_index(self):
461+
prompt = "hello"
462+
for variant in [None, "fp16"]:
463+
for use_safe in [True, False]:
464+
pipe = StableDiffusionPipeline.from_pretrained(
465+
"hf-internal-testing/tiny-stable-diffusion-pipe-indexes",
466+
variant=variant,
467+
use_safetensors=use_safe,
468+
safety_checker=None,
469+
)
470+
pipe = pipe.to(torch_device)
471+
generator = torch.manual_seed(0)
472+
out = pipe(prompt, num_inference_steps=2, generator=generator, output_type="numpy").images
473+
474+
with tempfile.TemporaryDirectory() as tmpdirname:
475+
pipe.save_pretrained(tmpdirname)
476+
pipe_2 = StableDiffusionPipeline.from_pretrained(
477+
tmpdirname, safe_serialization=use_safe, variant=variant
478+
)
479+
pipe_2 = pipe_2.to(torch_device)
480+
481+
generator = torch.manual_seed(0)
482+
483+
out_2 = pipe_2(prompt, num_inference_steps=2, generator=generator, output_type="numpy").images
484+
485+
assert np.max(np.abs(out - out_2)) < 1e-3
486+
365487
def test_text_inversion_download(self):
366488
pipe = StableDiffusionPipeline.from_pretrained(
367489
"hf-internal-testing/tiny-stable-diffusion-torch", safety_checker=None

0 commit comments

Comments
 (0)