test_pipelines.py

# coding=utf-8
# Copyright 2024 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

importgc
importjson
importos
importrandom
importre
importshutil
importsys
importtempfile
importtraceback
importunittest
importunittest.mockasmock
importwarnings

importnumpyasnp
importPIL.Image
importrequests_mock
importsafetensors.torch
importtorch
importtorch.nnasnn
fromhuggingface_hubimportsnapshot_download
fromparameterizedimportparameterized
fromPILimportImage
fromrequests.exceptionsimportHTTPError
fromtransformersimportCLIPImageProcessor, CLIPModel, CLIPTextConfig, CLIPTextModel, CLIPTokenizer

fromdiffusersimport (
AutoencoderKL,
ConfigMixin,
DDIMPipeline,
DDIMScheduler,
DDPMPipeline,
DDPMScheduler,
DiffusionPipeline,
DPMSolverMultistepScheduler,
EulerAncestralDiscreteScheduler,
EulerDiscreteScheduler,
LMSDiscreteScheduler,
ModelMixin,
PNDMScheduler,
StableDiffusionImg2ImgPipeline,
StableDiffusionInpaintPipelineLegacy,
StableDiffusionPipeline,
UNet2DConditionModel,
UNet2DModel,
UniPCMultistepScheduler,
logging,
)
fromdiffusers.pipelines.pipeline_utilsimport_get_pipeline_class
fromdiffusers.schedulers.scheduling_utilsimportSCHEDULER_CONFIG_NAME
fromdiffusers.utilsimport (
CONFIG_NAME,
WEIGHTS_NAME,
)
fromdiffusers.utils.testing_utilsimport (
CaptureLogger,
backend_empty_cache,
enable_full_determinism,
floats_tensor,
get_python_version,
get_tests_dir,
is_torch_compile,
load_numpy,
nightly,
require_compel,
require_flax,
require_hf_hub_version_greater,
require_onnxruntime,
require_peft_backend,
require_peft_version_greater,
require_torch_2,
require_torch_accelerator,
require_transformers_version_greater,
run_test_in_subprocess,
slow,
torch_device,
)
fromdiffusers.utils.torch_utilsimportis_compiled_module


enable_full_determinism()


# Will be run via run_test_in_subprocess
def_test_from_save_pretrained_dynamo(in_queue, out_queue, timeout):
error=None
try:
# 1. Load models
model=UNet2DModel(
block_out_channels=(32, 64),
layers_per_block=2,
sample_size=32,
in_channels=3,
out_channels=3,
down_block_types=("DownBlock2D", "AttnDownBlock2D"),
up_block_types=("AttnUpBlock2D", "UpBlock2D"),
 )
model=torch.compile(model)
scheduler=DDPMScheduler(num_train_timesteps=10)

ddpm=DDPMPipeline(model, scheduler)

# previous diffusers versions stripped compilation off
# compiled modules
assertis_compiled_module(ddpm.unet)

ddpm.to(torch_device)
ddpm.set_progress_bar_config(disable=None)

withtempfile.TemporaryDirectory() astmpdirname:
ddpm.save_pretrained(tmpdirname)
new_ddpm=DDPMPipeline.from_pretrained(tmpdirname)
new_ddpm.to(torch_device)

generator=torch.Generator(device=torch_device).manual_seed(0)
image=ddpm(generator=generator, num_inference_steps=5, output_type="np").images

generator=torch.Generator(device=torch_device).manual_seed(0)
new_image=new_ddpm(generator=generator, num_inference_steps=5, output_type="np").images

assertnp.abs(image-new_image).max() <1e-5, "Models don't give the same forward pass"
exceptException:
error=f"{traceback.format_exc()}"

results= {"error": error}
out_queue.put(results, timeout=timeout)
out_queue.join()


classCustomEncoder(ModelMixin, ConfigMixin):
def__init__(self):
super().__init__()
self.linear=nn.Linear(3, 3)


classCustomPipeline(DiffusionPipeline):
def__init__(self, encoder: CustomEncoder, scheduler: DDIMScheduler):
super().__init__()
self.register_modules(encoder=encoder, scheduler=scheduler)


classDownloadTests(unittest.TestCase):
@unittest.skip("Flaky behaviour on CI. Re-enable after migrating to new runners")
deftest_one_request_upon_cached(self):
# TODO: For some reason this test fails on MPS where no HEAD call is made.
iftorch_device=="mps":
return

withtempfile.TemporaryDirectory() astmpdirname:
withrequests_mock.mock(real_http=True) asm:
DiffusionPipeline.download("hf-internal-testing/tiny-stable-diffusion-pipe", cache_dir=tmpdirname)

download_requests= [r.methodforrinm.request_history]
assertdownload_requests.count("HEAD") ==15, "15 calls to files"
assertdownload_requests.count("GET") ==17, "15 calls to files + model_info + model_index.json"
assertlen(download_requests) ==32, (
"2 calls per file (15 files) + send_telemetry, model_info and model_index.json"
 )

withrequests_mock.mock(real_http=True) asm:
DiffusionPipeline.download(
"hf-internal-testing/tiny-stable-diffusion-pipe", safety_checker=None, cache_dir=tmpdirname
 )

cache_requests= [r.methodforrinm.request_history]
assertcache_requests.count("HEAD") ==1, "model_index.json is only HEAD"
assertcache_requests.count("GET") ==1, "model info is only GET"
assertlen(cache_requests) ==2, (
"We should call only `model_info` to check for _commit hash and `send_telemetry`"
 )

deftest_less_downloads_passed_object(self):
withtempfile.TemporaryDirectory() astmpdirname:
cached_folder=DiffusionPipeline.download(
"hf-internal-testing/tiny-stable-diffusion-pipe", safety_checker=None, cache_dir=tmpdirname
 )

# make sure safety checker is not downloaded
assert"safety_checker"notinos.listdir(cached_folder)

# make sure rest is downloaded
assert"unet"inos.listdir(cached_folder)
assert"tokenizer"inos.listdir(cached_folder)
assert"vae"inos.listdir(cached_folder)
assert"model_index.json"inos.listdir(cached_folder)
assert"scheduler"inos.listdir(cached_folder)
assert"feature_extractor"inos.listdir(cached_folder)

@unittest.skip("Flaky behaviour on CI. Re-enable after migrating to new runners")
deftest_less_downloads_passed_object_calls(self):
# TODO: For some reason this test fails on MPS where no HEAD call is made.
iftorch_device=="mps":
return

withtempfile.TemporaryDirectory() astmpdirname:
withrequests_mock.mock(real_http=True) asm:
DiffusionPipeline.download(
"hf-internal-testing/tiny-stable-diffusion-pipe", safety_checker=None, cache_dir=tmpdirname
 )

download_requests= [r.methodforrinm.request_history]
# 15 - 2 because no call to config or model file for `safety_checker`
assertdownload_requests.count("HEAD") ==13, "13 calls to files"
# 17 - 2 because no call to config or model file for `safety_checker`
assertdownload_requests.count("GET") ==15, "13 calls to files + model_info + model_index.json"
assertlen(download_requests) ==28, (
"2 calls per file (13 files) + send_telemetry, model_info and model_index.json"
 )

withrequests_mock.mock(real_http=True) asm:
DiffusionPipeline.download(
"hf-internal-testing/tiny-stable-diffusion-pipe", safety_checker=None, cache_dir=tmpdirname
 )

cache_requests= [r.methodforrinm.request_history]
assertcache_requests.count("HEAD") ==1, "model_index.json is only HEAD"
assertcache_requests.count("GET") ==1, "model info is only GET"
assertlen(cache_requests) ==2, (
"We should call only `model_info` to check for _commit hash and `send_telemetry`"
 )

deftest_download_only_pytorch(self):
withtempfile.TemporaryDirectory() astmpdirname:
# pipeline has Flax weights
tmpdirname=DiffusionPipeline.download(
"hf-internal-testing/tiny-stable-diffusion-pipe", safety_checker=None, cache_dir=tmpdirname
 )

all_root_files= [t[-1] fortinos.walk(os.path.join(tmpdirname))]
files= [itemforsublistinall_root_filesforiteminsublist]

# None of the downloaded files should be a flax file even if we have some here:
# https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe/blob/main/unet/diffusion_flax_model.msgpack
assertnotany(f.endswith(".msgpack") forfinfiles)
# We need to never convert this tiny model to safetensors for this test to pass
assertnotany(f.endswith(".safetensors") forfinfiles)

deftest_force_safetensors_error(self):
withtempfile.TemporaryDirectory() astmpdirname:
# pipeline has Flax weights
withself.assertRaises(EnvironmentError):
tmpdirname=DiffusionPipeline.download(
"hf-internal-testing/tiny-stable-diffusion-pipe-no-safetensors",
safety_checker=None,
cache_dir=tmpdirname,
use_safetensors=True,
 )

deftest_download_safetensors(self):
withtempfile.TemporaryDirectory() astmpdirname:
# pipeline has Flax weights
tmpdirname=DiffusionPipeline.download(
"hf-internal-testing/tiny-stable-diffusion-pipe-safetensors",
safety_checker=None,
cache_dir=tmpdirname,
 )

all_root_files= [t[-1] fortinos.walk(os.path.join(tmpdirname))]
files= [itemforsublistinall_root_filesforiteminsublist]

# None of the downloaded files should be a pytorch file even if we have some here:
# https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe/blob/main/unet/diffusion_flax_model.msgpack
assertnotany(f.endswith(".bin") forfinfiles)

deftest_download_safetensors_index(self):
forvariantin ["fp16", None]:
withtempfile.TemporaryDirectory() astmpdirname:
tmpdirname=DiffusionPipeline.download(
"hf-internal-testing/tiny-stable-diffusion-pipe-indexes",
cache_dir=tmpdirname,
use_safetensors=True,
variant=variant,
 )

all_root_files= [t[-1] fortinos.walk(os.path.join(tmpdirname))]
files= [itemforsublistinall_root_filesforiteminsublist]

# None of the downloaded files should be a safetensors file even if we have some here:
# https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe-indexes/tree/main/text_encoder
ifvariantisNone:
assertnotany("fp16"infforfinfiles)
else:
model_files= [fforfinfilesif"safetensors"inf]
assertall("fp16"infforfinmodel_files)

assertlen([fforfinfilesif".safetensors"inf]) ==8
assertnotany(".bin"infforfinfiles)

deftest_download_bin_index(self):
forvariantin ["fp16", None]:
withtempfile.TemporaryDirectory() astmpdirname:
tmpdirname=DiffusionPipeline.download(
"hf-internal-testing/tiny-stable-diffusion-pipe-indexes",
cache_dir=tmpdirname,
use_safetensors=False,
variant=variant,
 )

all_root_files= [t[-1] fortinos.walk(os.path.join(tmpdirname))]
files= [itemforsublistinall_root_filesforiteminsublist]

# None of the downloaded files should be a safetensors file even if we have some here:
# https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe-indexes/tree/main/text_encoder
ifvariantisNone:
assertnotany("fp16"infforfinfiles)
else:
model_files= [fforfinfilesif"bin"inf]
assertall("fp16"infforfinmodel_files)

assertlen([fforfinfilesif".bin"inf]) ==8
assertnotany(".safetensors"infforfinfiles)

deftest_download_no_openvino_by_default(self):
withtempfile.TemporaryDirectory() astmpdirname:
tmpdirname=DiffusionPipeline.download(
"hf-internal-testing/tiny-stable-diffusion-open-vino",
cache_dir=tmpdirname,
 )

all_root_files= [t[-1] fortinos.walk(os.path.join(tmpdirname))]
files= [itemforsublistinall_root_filesforiteminsublist]

# make sure that by default no openvino weights are downloaded
assertall((f.endswith(".json") orf.endswith(".bin") orf.endswith(".txt")) forfinfiles)
assertnotany("openvino_"infforfinfiles)

deftest_download_no_onnx_by_default(self):
withtempfile.TemporaryDirectory() astmpdirname:
tmpdirname=DiffusionPipeline.download(
"hf-internal-testing/tiny-stable-diffusion-xl-pipe",
cache_dir=tmpdirname,
use_safetensors=False,
 )

all_root_files= [t[-1] fortinos.walk(os.path.join(tmpdirname))]
files= [itemforsublistinall_root_filesforiteminsublist]

# make sure that by default no onnx weights are downloaded for non-ONNX pipelines
assertall((f.endswith(".json") orf.endswith(".bin") orf.endswith(".txt")) forfinfiles)
assertnotany((f.endswith(".onnx") orf.endswith(".pb")) forfinfiles)

@require_onnxruntime
deftest_download_onnx_by_default_for_onnx_pipelines(self):
withtempfile.TemporaryDirectory() astmpdirname:
tmpdirname=DiffusionPipeline.download(
"hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline",
cache_dir=tmpdirname,
 )

all_root_files= [t[-1] fortinos.walk(os.path.join(tmpdirname))]
files= [itemforsublistinall_root_filesforiteminsublist]

# make sure that by default onnx weights are downloaded for ONNX pipelines
assertany((f.endswith(".json") orf.endswith(".bin") orf.endswith(".txt")) forfinfiles)
assertany((f.endswith(".onnx")) forfinfiles)
assertany((f.endswith(".pb")) forfinfiles)

deftest_download_no_safety_checker(self):
prompt="hello"
pipe=StableDiffusionPipeline.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-torch", safety_checker=None
 )
pipe=pipe.to(torch_device)
generator=torch.manual_seed(0)
out=pipe(prompt, num_inference_steps=2, generator=generator, output_type="np").images

pipe_2=StableDiffusionPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-torch")
pipe_2=pipe_2.to(torch_device)
generator=torch.manual_seed(0)
out_2=pipe_2(prompt, num_inference_steps=2, generator=generator, output_type="np").images

assertnp.max(np.abs(out-out_2)) <1e-3

deftest_load_no_safety_checker_explicit_locally(self):
prompt="hello"
pipe=StableDiffusionPipeline.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-torch", safety_checker=None
 )
pipe=pipe.to(torch_device)
generator=torch.manual_seed(0)
out=pipe(prompt, num_inference_steps=2, generator=generator, output_type="np").images

withtempfile.TemporaryDirectory() astmpdirname:
pipe.save_pretrained(tmpdirname)
pipe_2=StableDiffusionPipeline.from_pretrained(tmpdirname, safety_checker=None)
pipe_2=pipe_2.to(torch_device)

generator=torch.manual_seed(0)

out_2=pipe_2(prompt, num_inference_steps=2, generator=generator, output_type="np").images

assertnp.max(np.abs(out-out_2)) <1e-3

deftest_load_no_safety_checker_default_locally(self):
prompt="hello"
pipe=StableDiffusionPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-torch")
pipe=pipe.to(torch_device)

generator=torch.manual_seed(0)
out=pipe(prompt, num_inference_steps=2, generator=generator, output_type="np").images

withtempfile.TemporaryDirectory() astmpdirname:
pipe.save_pretrained(tmpdirname)
pipe_2=StableDiffusionPipeline.from_pretrained(tmpdirname)
pipe_2=pipe_2.to(torch_device)

generator=torch.manual_seed(0)

out_2=pipe_2(prompt, num_inference_steps=2, generator=generator, output_type="np").images

assertnp.max(np.abs(out-out_2)) <1e-3

deftest_cached_files_are_used_when_no_internet(self):
# A mock response for an HTTP head request to emulate server down
response_mock=mock.Mock()
response_mock.status_code=500
response_mock.headers= {}
response_mock.raise_for_status.side_effect=HTTPError
response_mock.json.return_value= {}

# Download this model to make sure it's in the cache.
orig_pipe=DiffusionPipeline.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-torch", safety_checker=None
 )
orig_comps= {k: vfork, vinorig_pipe.components.items() ifhasattr(v, "parameters")}

# Under the mock environment we get a 500 error when trying to reach the model.
withmock.patch("requests.request", return_value=response_mock):
# Download this model to make sure it's in the cache.
pipe=DiffusionPipeline.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-torch", safety_checker=None
 )
comps= {k: vfork, vinpipe.components.items() ifhasattr(v, "parameters")}

form1, m2inzip(orig_comps.values(), comps.values()):
forp1, p2inzip(m1.parameters(), m2.parameters()):
ifp1.data.ne(p2.data).sum() >0:
assertFalse, "Parameters not the same!"

deftest_local_files_only_are_used_when_no_internet(self):
# A mock response for an HTTP head request to emulate server down
response_mock=mock.Mock()
response_mock.status_code=500
response_mock.headers= {}
response_mock.raise_for_status.side_effect=HTTPError
response_mock.json.return_value= {}

# first check that with local files only the pipeline can only be used if cached
withself.assertRaises(FileNotFoundError):
withtempfile.TemporaryDirectory() astmpdirname:
orig_pipe=DiffusionPipeline.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-torch", local_files_only=True, cache_dir=tmpdirname
 )

# now download
orig_pipe=DiffusionPipeline.download("hf-internal-testing/tiny-stable-diffusion-torch")

# make sure it can be loaded with local_files_only
orig_pipe=DiffusionPipeline.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-torch", local_files_only=True
 )
orig_comps= {k: vfork, vinorig_pipe.components.items() ifhasattr(v, "parameters")}

# Under the mock environment we get a 500 error when trying to connect to the internet.
# Make sure it works local_files_only only works here!
withmock.patch("requests.request", return_value=response_mock):
# Download this model to make sure it's in the cache.
pipe=DiffusionPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-torch")
comps= {k: vfork, vinpipe.components.items() ifhasattr(v, "parameters")}

form1, m2inzip(orig_comps.values(), comps.values()):
forp1, p2inzip(m1.parameters(), m2.parameters()):
ifp1.data.ne(p2.data).sum() >0:
assertFalse, "Parameters not the same!"

deftest_download_from_variant_folder(self):
foruse_safetensorsin [False, True]:
other_format=".bin"ifuse_safetensorselse".safetensors"
withtempfile.TemporaryDirectory() astmpdirname:
tmpdirname=StableDiffusionPipeline.download(
"hf-internal-testing/stable-diffusion-all-variants",
cache_dir=tmpdirname,
use_safetensors=use_safetensors,
 )
all_root_files= [t[-1] fortinos.walk(tmpdirname)]
files= [itemforsublistinall_root_filesforiteminsublist]

# None of the downloaded files should be a variant file even if we have some here:
# https://huggingface.co/hf-internal-testing/stable-diffusion-all-variants/tree/main/unet
assertlen(files) ==15, f"We should only download 15 files, not {len(files)}"
assertnotany(f.endswith(other_format) forfinfiles)
# no variants
assertnotany(len(f.split(".")) ==3forfinfiles)

deftest_download_variant_all(self):
foruse_safetensorsin [False, True]:
other_format=".bin"ifuse_safetensorselse".safetensors"
this_format=".safetensors"ifuse_safetensorselse".bin"
variant="fp16"

withtempfile.TemporaryDirectory() astmpdirname:
tmpdirname=StableDiffusionPipeline.download(
"hf-internal-testing/stable-diffusion-all-variants",
cache_dir=tmpdirname,
variant=variant,
use_safetensors=use_safetensors,
 )
all_root_files= [t[-1] fortinos.walk(tmpdirname)]
files= [itemforsublistinall_root_filesforiteminsublist]

# None of the downloaded files should be a non-variant file even if we have some here:
# https://huggingface.co/hf-internal-testing/stable-diffusion-all-variants/tree/main/unet
assertlen(files) ==15, f"We should only download 15 files, not {len(files)}"
# unet, vae, text_encoder, safety_checker
assertlen([fforfinfilesiff.endswith(f"{variant}{this_format}")]) ==4
# all checkpoints should have variant ending
assertnotany(f.endswith(this_format) andnotf.endswith(f"{variant}{this_format}") forfinfiles)
assertnotany(f.endswith(other_format) forfinfiles)

deftest_download_variant_partly(self):
foruse_safetensorsin [False, True]:
other_format=".bin"ifuse_safetensorselse".safetensors"
this_format=".safetensors"ifuse_safetensorselse".bin"
variant="no_ema"

withtempfile.TemporaryDirectory() astmpdirname:
tmpdirname=StableDiffusionPipeline.download(
"hf-internal-testing/stable-diffusion-all-variants",
cache_dir=tmpdirname,
variant=variant,
use_safetensors=use_safetensors,
 )
all_root_files= [t[-1] fortinos.walk(tmpdirname)]
files= [itemforsublistinall_root_filesforiteminsublist]

unet_files=os.listdir(os.path.join(tmpdirname, "unet"))

# Some of the downloaded files should be a non-variant file, check:
# https://huggingface.co/hf-internal-testing/stable-diffusion-all-variants/tree/main/unet
assertlen(files) ==15, f"We should only download 15 files, not {len(files)}"
# only unet has "no_ema" variant
assertf"diffusion_pytorch_model.{variant}{this_format}"inunet_files
assertlen([fforfinfilesiff.endswith(f"{variant}{this_format}")]) ==1
# vae, safety_checker and text_encoder should have no variant
assertsum(f.endswith(this_format) andnotf.endswith(f"{variant}{this_format}") forfinfiles) ==3
assertnotany(f.endswith(other_format) forfinfiles)

deftest_download_variants_with_sharded_checkpoints(self):
# Here we test for downloading of "variant" files belonging to the `unet` and
# the `text_encoder`. Their checkpoints can be sharded.
foruse_safetensorsin [True, False]:
forvariantin ["fp16", None]:
withtempfile.TemporaryDirectory() astmpdirname:
tmpdirname=DiffusionPipeline.download(
"hf-internal-testing/tiny-stable-diffusion-pipe-variants-right-format",
safety_checker=None,
cache_dir=tmpdirname,
variant=variant,
use_safetensors=use_safetensors,
 )

all_root_files= [t[-1] fortinos.walk(os.path.join(tmpdirname))]
files= [itemforsublistinall_root_filesforiteminsublist]

# Check for `model_ext` and `variant`.
model_ext=".safetensors"ifuse_safetensorselse".bin"
unexpected_ext=".bin"ifuse_safetensorselse".safetensors"
model_files= [fforfinfilesiff.endswith(model_ext)]
assertnotany(f.endswith(unexpected_ext) forfinfiles)
assertall(variantinfforfinmodel_filesiff.endswith(model_ext) andvariantisnotNone)

deftest_download_legacy_variants_with_sharded_ckpts_raises_warning(self):
repo_id="hf-internal-testing/tiny-stable-diffusion-pipe-variants-all-kinds"
logger=logging.get_logger("diffusers.pipelines.pipeline_utils")
deprecated_warning_msg="Warning: The repository contains sharded checkpoints for variant"

foris_localin [True, False]:
withCaptureLogger(logger) ascap_logger:
withtempfile.TemporaryDirectory() astmpdirname:
local_repo_id=repo_id
ifis_local:
local_repo_id=snapshot_download(repo_id, cache_dir=tmpdirname)

_=DiffusionPipeline.from_pretrained(
local_repo_id,
safety_checker=None,
variant="fp16",
use_safetensors=True,
 )
assertdeprecated_warning_msginstr(cap_logger), "Deprecation warning not found in logs"

deftest_download_safetensors_only_variant_exists_for_model(self):
variant=None
use_safetensors=True

# text encoder is missing no variant weights, so the following can't work
withtempfile.TemporaryDirectory() astmpdirname:
withself.assertRaises(OSError) aserror_context:
tmpdirname=StableDiffusionPipeline.from_pretrained(
"hf-internal-testing/stable-diffusion-broken-variants",
cache_dir=tmpdirname,
variant=variant,
use_safetensors=use_safetensors,
 )
assert"Error no file name"instr(error_context.exception)

# text encoder has fp16 variants so we can load it
withtempfile.TemporaryDirectory() astmpdirname:
tmpdirname=StableDiffusionPipeline.download(
"hf-internal-testing/stable-diffusion-broken-variants",
use_safetensors=use_safetensors,
cache_dir=tmpdirname,
variant="fp16",
 )
all_root_files= [t[-1] fortinos.walk(tmpdirname)]
files= [itemforsublistinall_root_filesforiteminsublist]
# None of the downloaded files should be a non-variant file even if we have some here:
# https://huggingface.co/hf-internal-testing/stable-diffusion-broken-variants/tree/main/unet
assertlen(files) ==15, f"We should only download 15 files, not {len(files)}"

deftest_download_bin_only_variant_exists_for_model(self):
variant=None
use_safetensors=False

# text encoder is missing Non-variant weights, so the following can't work
withtempfile.TemporaryDirectory() astmpdirname:
withself.assertRaises(OSError) aserror_context:
tmpdirname=StableDiffusionPipeline.from_pretrained(
"hf-internal-testing/stable-diffusion-broken-variants",
cache_dir=tmpdirname,
variant=variant,
use_safetensors=use_safetensors,
 )
assert"Error no file name"instr(error_context.exception)

# text encoder has fp16 variants so we can load it
withtempfile.TemporaryDirectory() astmpdirname:
tmpdirname=StableDiffusionPipeline.download(
"hf-internal-testing/stable-diffusion-broken-variants",
use_safetensors=use_safetensors,
cache_dir=tmpdirname,
variant="fp16",
 )
all_root_files= [t[-1] fortinos.walk(tmpdirname)]
files= [itemforsublistinall_root_filesforiteminsublist]
# None of the downloaded files should be a non-variant file even if we have some here:
# https://huggingface.co/hf-internal-testing/stable-diffusion-broken-variants/tree/main/unet
assertlen(files) ==15, f"We should only download 15 files, not {len(files)}"

deftest_download_safetensors_variant_does_not_exist_for_model(self):
variant="no_ema"
use_safetensors=True

# text encoder is missing no_ema variant weights, so the following can't work
withtempfile.TemporaryDirectory() astmpdirname:
withself.assertRaises(OSError) aserror_context:
tmpdirname=StableDiffusionPipeline.from_pretrained(
"hf-internal-testing/stable-diffusion-broken-variants",
cache_dir=tmpdirname,
variant=variant,
use_safetensors=use_safetensors,
 )

assert"Error no file name"instr(error_context.exception)

deftest_download_bin_variant_does_not_exist_for_model(self):
variant="no_ema"
use_safetensors=False

# text encoder is missing no_ema variant weights, so the following can't work
withtempfile.TemporaryDirectory() astmpdirname:
withself.assertRaises(OSError) aserror_context:
tmpdirname=StableDiffusionPipeline.from_pretrained(
"hf-internal-testing/stable-diffusion-broken-variants",
cache_dir=tmpdirname,
variant=variant,
use_safetensors=use_safetensors,
 )
assert"Error no file name"instr(error_context.exception)

deftest_local_save_load_index(self):
prompt="hello"
forvariantin [None, "fp16"]:
foruse_safein [True, False]:
pipe=StableDiffusionPipeline.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-pipe-indexes",
variant=variant,
use_safetensors=use_safe,
safety_checker=None,
 )
pipe=pipe.to(torch_device)
generator=torch.manual_seed(0)
out=pipe(prompt, num_inference_steps=2, generator=generator, output_type="np").images

withtempfile.TemporaryDirectory() astmpdirname:
pipe.save_pretrained(tmpdirname, variant=variant, safe_serialization=use_safe)
pipe_2=StableDiffusionPipeline.from_pretrained(
tmpdirname, safe_serialization=use_safe, variant=variant
 )
pipe_2=pipe_2.to(torch_device)

generator=torch.manual_seed(0)

out_2=pipe_2(prompt, num_inference_steps=2, generator=generator, output_type="np").images

assertnp.max(np.abs(out-out_2)) <1e-3

deftest_text_inversion_download(self):
pipe=StableDiffusionPipeline.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-torch", safety_checker=None
 )
pipe=pipe.to(torch_device)

num_tokens=len(pipe.tokenizer)

# single token load local
withtempfile.TemporaryDirectory() astmpdirname:
ten= {"<*>": torch.ones((32,))}
torch.save(ten, os.path.join(tmpdirname, "learned_embeds.bin"))

pipe.load_textual_inversion(tmpdirname)

token=pipe.tokenizer.convert_tokens_to_ids("<*>")
asserttoken==num_tokens, "Added token must be at spot `num_tokens`"
assertpipe.text_encoder.get_input_embeddings().weight[-1].sum().item() ==32
assertpipe._maybe_convert_prompt("<*>", pipe.tokenizer) =="<*>"

prompt="hey <*>"
out=pipe(prompt, num_inference_steps=1, output_type="np").images
assertout.shape== (1, 128, 128, 3)

# single token load local with weight name
withtempfile.TemporaryDirectory() astmpdirname:
ten= {"<**>": 2*torch.ones((1, 32))}
torch.save(ten, os.path.join(tmpdirname, "learned_embeds.bin"))

pipe.load_textual_inversion(tmpdirname, weight_name="learned_embeds.bin")

token=pipe.tokenizer.convert_tokens_to_ids("<**>")
asserttoken==num_tokens+1, "Added token must be at spot `num_tokens`"
assertpipe.text_encoder.get_input_embeddings().weight[-1].sum().item() ==64
assertpipe._maybe_convert_prompt("<**>", pipe.tokenizer) =="<**>"

prompt="hey <**>"
out=pipe(prompt, num_inference_steps=1, output_type="np").images
assertout.shape== (1, 128, 128, 3)

# multi token load
withtempfile.TemporaryDirectory() astmpdirname:
ten= {"<***>": torch.cat([3*torch.ones((1, 32)), 4*torch.ones((1, 32)), 5*torch.ones((1, 32))])}
torch.save(ten, os.path.join(tmpdirname, "learned_embeds.bin"))

pipe.load_textual_inversion(tmpdirname)

token=pipe.tokenizer.convert_tokens_to_ids("<***>")
token_1=pipe.tokenizer.convert_tokens_to_ids("<***>_1")
token_2=pipe.tokenizer.convert_tokens_to_ids("<***>_2")

asserttoken==num_tokens+2, "Added token must be at spot `num_tokens`"
asserttoken_1==num_tokens+3, "Added token must be at spot `num_tokens`"
asserttoken_2==num_tokens+4, "Added token must be at spot `num_tokens`"
assertpipe.text_encoder.get_input_embeddings().weight[-3].sum().item() ==96
assertpipe.text_encoder.get_input_embeddings().weight[-2].sum().item() ==128
assertpipe.text_encoder.get_input_embeddings().weight[-1].sum().item() ==160
assertpipe._maybe_convert_prompt("<***>", pipe.tokenizer) =="<***> <***>_1 <***>_2"

prompt="hey <***>"
out=pipe(prompt, num_inference_steps=1, output_type="np").images
assertout.shape== (1, 128, 128, 3)

# multi token load a1111
withtempfile.TemporaryDirectory() astmpdirname:
ten= {
"string_to_param": {
"*": torch.cat([3*torch.ones((1, 32)), 4*torch.ones((1, 32)), 5*torch.ones((1, 32))])
 },
"name": "<****>",
 }
torch.save(ten, os.path.join(tmpdirname, "a1111.bin"))

pipe.load_textual_inversion(tmpdirname, weight_name="a1111.bin")

token=pipe.tokenizer.convert_tokens_to_ids("<****>")
token_1=pipe.tokenizer.convert_tokens_to_ids("<****>_1")
token_2=pipe.tokenizer.convert_tokens_to_ids("<****>_2")

asserttoken==num_tokens+5, "Added token must be at spot `num_tokens`"
asserttoken_1==num_tokens+6, "Added token must be at spot `num_tokens`"
asserttoken_2==num_tokens+7, "Added token must be at spot `num_tokens`"
assertpipe.text_encoder.get_input_embeddings().weight[-3].sum().item() ==96
assertpipe.text_encoder.get_input_embeddings().weight[-2].sum().item() ==128
assertpipe.text_encoder.get_input_embeddings().weight[-1].sum().item() ==160
assertpipe._maybe_convert_prompt("<****>", pipe.tokenizer) =="<****> <****>_1 <****>_2"

prompt="hey <****>"
out=pipe(prompt, num_inference_steps=1, output_type="np").images
assertout.shape== (1, 128, 128, 3)

# multi embedding load
withtempfile.TemporaryDirectory() astmpdirname1:
withtempfile.TemporaryDirectory() astmpdirname2:
ten= {"<*****>": torch.ones((32,))}
torch.save(ten, os.path.join(tmpdirname1, "learned_embeds.bin"))

ten= {"<******>": 2*torch.ones((1, 32))}
torch.save(ten, os.path.join(tmpdirname2, "learned_embeds.bin"))

pipe.load_textual_inversion([tmpdirname1, tmpdirname2])

token=pipe.tokenizer.convert_tokens_to_ids("<*****>")
asserttoken==num_tokens+8, "Added token must be at spot `num_tokens`"
assertpipe.text_encoder.get_input_embeddings().weight[-2].sum().item() ==32
assertpipe._maybe_convert_prompt("<*****>", pipe.tokenizer) =="<*****>"

token=pipe.tokenizer.convert_tokens_to_ids("<******>")
asserttoken==num_tokens+9, "Added token must be at spot `num_tokens`"
assertpipe.text_encoder.get_input_embeddings().weight[-1].sum().item() ==64
assertpipe._maybe_convert_prompt("<******>", pipe.tokenizer) =="<******>"

prompt="hey <*****> <******>"
out=pipe(prompt, num_inference_steps=1, output_type="np").images
assertout.shape== (1, 128, 128, 3)

# single token state dict load
ten= {"<x>": torch.ones((32,))}
pipe.load_textual_inversion(ten)

token=pipe.tokenizer.convert_tokens_to_ids("<x>")
asserttoken==num_tokens+10, "Added token must be at spot `num_tokens`"
assertpipe.text_encoder.get_input_embeddings().weight[-1].sum().item() ==32
assertpipe._maybe_convert_prompt("<x>", pipe.tokenizer) =="<x>"

prompt="hey <x>"
out=pipe(prompt, num_inference_steps=1, output_type="np").images
assertout.shape== (1, 128, 128, 3)

# multi embedding state dict load
ten1= {"<xxxxx>": torch.ones((32,))}
ten2= {"<xxxxxx>": 2*torch.ones((1, 32))}

pipe.load_textual_inversion([ten1, ten2])

token=pipe.tokenizer.convert_tokens_to_ids("<xxxxx>")
asserttoken==num_tokens+11, "Added token must be at spot `num_tokens`"
assertpipe.text_encoder.get_input_embeddings().weight[-2].sum().item() ==32
assertpipe._maybe_convert_prompt("<xxxxx>", pipe.tokenizer) =="<xxxxx>"

token=pipe.tokenizer.convert_tokens_to_ids("<xxxxxx>")
asserttoken==num_tokens+12, "Added token must be at spot `num_tokens`"
assertpipe.text_encoder.get_input_embeddings().weight[-1].sum().item() ==64
assertpipe._maybe_convert_prompt("<xxxxxx>", pipe.tokenizer) =="<xxxxxx>"

prompt="hey <xxxxx> <xxxxxx>"
out=pipe(prompt, num_inference_steps=1, output_type="np").images
assertout.shape== (1, 128, 128, 3)

# auto1111 multi-token state dict load
ten= {
"string_to_param": {
"*": torch.cat([3*torch.ones((1, 32)), 4*torch.ones((1, 32)), 5*torch.ones((1, 32))])
 },
"name": "<xxxx>",
 }

pipe.load_textual_inversion(ten)

token=pipe.tokenizer.convert_tokens_to_ids("<xxxx>")
token_1=pipe.tokenizer.convert_tokens_to_ids("<xxxx>_1")
token_2=pipe.tokenizer.convert_tokens_to_ids("<xxxx>_2")

asserttoken==num_tokens+13, "Added token must be at spot `num_tokens`"
asserttoken_1==num_tokens+14, "Added token must be at spot `num_tokens`"
asserttoken_2==num_tokens+15, "Added token must be at spot `num_tokens`"
assertpipe.text_encoder.get_input_embeddings().weight[-3].sum().item() ==96
assertpipe.text_encoder.get_input_embeddings().weight[-2].sum().item() ==128
assertpipe.text_encoder.get_input_embeddings().weight[-1].sum().item() ==160
assertpipe._maybe_convert_prompt("<xxxx>", pipe.tokenizer) =="<xxxx> <xxxx>_1 <xxxx>_2"

prompt="hey <xxxx>"
out=pipe(prompt, num_inference_steps=1, output_type="np").images
assertout.shape== (1, 128, 128, 3)

# multiple references to multi embedding
ten= {"<cat>": torch.ones(3, 32)}
pipe.load_textual_inversion(ten)

assert (
pipe._maybe_convert_prompt("<cat> <cat>", pipe.tokenizer) =="<cat> <cat>_1 <cat>_2 <cat> <cat>_1 <cat>_2"
 )

prompt="hey <cat> <cat>"
out=pipe(prompt, num_inference_steps=1, output_type="np").images
assertout.shape== (1, 128, 128, 3)

deftest_text_inversion_multi_tokens(self):
pipe1=StableDiffusionPipeline.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-torch", safety_checker=None
 )
pipe1=pipe1.to(torch_device)

token1, token2="<*>", "<**>"
ten1=torch.ones((32,))
ten2=torch.ones((32,)) *2

num_tokens=len(pipe1.tokenizer)

pipe1.load_textual_inversion(ten1, token=token1)
pipe1.load_textual_inversion(ten2, token=token2)
emb1=pipe1.text_encoder.get_input_embeddings().weight

pipe2=StableDiffusionPipeline.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-torch", safety_checker=None
 )
pipe2=pipe2.to(torch_device)
pipe2.load_textual_inversion([ten1, ten2], token=[token1, token2])
emb2=pipe2.text_encoder.get_input_embeddings().weight

pipe3=StableDiffusionPipeline.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-torch", safety_checker=None
 )
pipe3=pipe3.to(torch_device)
pipe3.load_textual_inversion(torch.stack([ten1, ten2], dim=0), token=[token1, token2])
emb3=pipe3.text_encoder.get_input_embeddings().weight

assertlen(pipe1.tokenizer) ==len(pipe2.tokenizer) ==len(pipe3.tokenizer) ==num_tokens+2
assert (
pipe1.tokenizer.convert_tokens_to_ids(token1)
==pipe2.tokenizer.convert_tokens_to_ids(token1)
==pipe3.tokenizer.convert_tokens_to_ids(token1)
==num_tokens
 )
assert (
pipe1.tokenizer.convert_tokens_to_ids(token2)
==pipe2.tokenizer.convert_tokens_to_ids(token2)
==pipe3.tokenizer.convert_tokens_to_ids(token2)
==num_tokens+1
 )
assertemb1[num_tokens].sum().item() ==emb2[num_tokens].sum().item() ==emb3[num_tokens].sum().item()
assert (
emb1[num_tokens+1].sum().item() ==emb2[num_tokens+1].sum().item() ==emb3[num_tokens+1].sum().item()
 )

deftest_textual_inversion_unload(self):
pipe1=StableDiffusionPipeline.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-torch", safety_checker=None
 )
pipe1=pipe1.to(torch_device)
orig_tokenizer_size=len(pipe1.tokenizer)
orig_emb_size=len(pipe1.text_encoder.get_input_embeddings().weight)

token="<*>"
ten=torch.ones((32,))
pipe1.load_textual_inversion(ten, token=token)
pipe1.unload_textual_inversion()
pipe1.load_textual_inversion(ten, token=token)
pipe1.unload_textual_inversion()

final_tokenizer_size=len(pipe1.tokenizer)
final_emb_size=len(pipe1.text_encoder.get_input_embeddings().weight)
# both should be restored to original size
assertfinal_tokenizer_size==orig_tokenizer_size
assertfinal_emb_size==orig_emb_size

deftest_download_ignore_files(self):
# Check https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe-ignore-files/blob/72f58636e5508a218c6b3f60550dc96445547817/model_index.json#L4
withtempfile.TemporaryDirectory() astmpdirname:
# pipeline has Flax weights
tmpdirname=DiffusionPipeline.download("hf-internal-testing/tiny-stable-diffusion-pipe-ignore-files")
all_root_files= [t[-1] fortinos.walk(os.path.join(tmpdirname))]
files= [itemforsublistinall_root_filesforiteminsublist]

# None of the downloaded files should be a pytorch file even if we have some here:
# https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe/blob/main/unet/diffusion_flax_model.msgpack
assertnotany(fin ["vae/diffusion_pytorch_model.bin", "text_encoder/config.json"] forfinfiles)
assertlen(files) ==14

deftest_download_dduf_with_custom_pipeline_raises_error(self):
withself.assertRaises(NotImplementedError):
_=DiffusionPipeline.download(
"DDUF/tiny-flux-dev-pipe-dduf", dduf_file="fluxpipeline.dduf", custom_pipeline="my_pipeline"
 )

deftest_download_dduf_with_connected_pipeline_raises_error(self):
withself.assertRaises(NotImplementedError):
_=DiffusionPipeline.download(
"DDUF/tiny-flux-dev-pipe-dduf", dduf_file="fluxpipeline.dduf", load_connected_pipeline=True