test_modeling_common.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.

importcopy
importgc
importinspect
importjson
importos
importre
importtempfile
importtraceback
importunittest
importunittest.mockasmock
importuuid
importwarnings
fromcollectionsimportdefaultdict
fromtypingimportDict, List, Optional, Tuple, Union

importnumpyasnp
importrequests_mock
importtorch
importtorch.nnasnn
fromaccelerate.utils.modelingimport_get_proper_dtype, compute_module_sizes, dtype_byte_size
fromhuggingface_hubimportModelCard, delete_repo, snapshot_download
fromhuggingface_hub.utilsimportis_jinja_available
fromparameterizedimportparameterized
fromrequests.exceptionsimportHTTPError

fromdiffusers.modelsimportSD3Transformer2DModel, UNet2DConditionModel
fromdiffusers.models.attention_processorimport (
AttnProcessor,
AttnProcessor2_0,
AttnProcessorNPU,
XFormersAttnProcessor,
)
fromdiffusers.models.auto_modelimportAutoModel
fromdiffusers.training_utilsimportEMAModel
fromdiffusers.utilsimport (
SAFE_WEIGHTS_INDEX_NAME,
WEIGHTS_INDEX_NAME,
is_peft_available,
is_torch_npu_available,
is_xformers_available,
logging,
)
fromdiffusers.utils.hub_utilsimport_add_variant
fromdiffusers.utils.testing_utilsimport (
CaptureLogger,
backend_empty_cache,
backend_max_memory_allocated,
backend_reset_peak_memory_stats,
backend_synchronize,
floats_tensor,
get_python_version,
is_torch_compile,
numpy_cosine_similarity_distance,
require_peft_backend,
require_peft_version_greater,
require_torch_2,
require_torch_accelerator,
require_torch_accelerator_with_training,
require_torch_gpu,
require_torch_multi_accelerator,
run_test_in_subprocess,
slow,
torch_all_close,
torch_device,
)
fromdiffusers.utils.torch_utilsimportget_torch_cuda_device_capability

from ..others.test_utilsimportTOKEN, USER, is_staging_test


ifis_peft_available():
frompeft.tuners.tuners_utilsimportBaseTunerLayer


defcaculate_expected_num_shards(index_map_path):
withopen(index_map_path) asf:
weight_map_dict=json.load(f)["weight_map"]
first_key=list(weight_map_dict.keys())[0]
weight_loc=weight_map_dict[first_key] # e.g., diffusion_pytorch_model-00001-of-00002.safetensors
expected_num_shards=int(weight_loc.split("-")[-1].split(".")[0])
returnexpected_num_shards


defcheck_if_lora_correctly_set(model) ->bool:
"""
 Checks if the LoRA layers are correctly set with peft
 """
formoduleinmodel.modules():
ifisinstance(module, BaseTunerLayer):
returnTrue
returnFalse


# Will be run via run_test_in_subprocess
def_test_from_save_pretrained_dynamo(in_queue, out_queue, timeout):
error=None
try:
init_dict, model_class=in_queue.get(timeout=timeout)

model=model_class(**init_dict)
model.to(torch_device)
model=torch.compile(model)

withtempfile.TemporaryDirectory() astmpdirname:
model.save_pretrained(tmpdirname, safe_serialization=False)
new_model=model_class.from_pretrained(tmpdirname)
new_model.to(torch_device)

assertnew_model.__class__==model_class
exceptException:
error=f"{traceback.format_exc()}"

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


defnamed_persistent_module_tensors(
module: nn.Module,
recurse: bool=False,
):
"""
 A helper function that gathers all the tensors (parameters + persistent buffers) of a given module.

 Args:
 module (`torch.nn.Module`):
 The module we want the tensors on.
 recurse (`bool`, *optional`, defaults to `False`):
 Whether or not to go look in every submodule or just return the direct parameters and buffers.
 """
yieldfrommodule.named_parameters(recurse=recurse)

fornamed_bufferinmodule.named_buffers(recurse=recurse):
name, _=named_buffer
# Get parent by splitting on dots and traversing the model
parent=module
if"."inname:
parent_name=name.rsplit(".", 1)[0]
forpartinparent_name.split("."):
parent=getattr(parent, part)
name=name.split(".")[-1]
ifnamenotinparent._non_persistent_buffers_set:
yieldnamed_buffer


defcompute_module_persistent_sizes(
model: nn.Module,
dtype: Optional[Union[str, torch.device]] =None,
special_dtypes: Optional[Dict[str, Union[str, torch.device]]] =None,
):
"""
 Compute the size of each submodule of a given model (parameters + persistent buffers).
 """
ifdtypeisnotNone:
dtype=_get_proper_dtype(dtype)
dtype_size=dtype_byte_size(dtype)
ifspecial_dtypesisnotNone:
special_dtypes= {key: _get_proper_dtype(dtyp) forkey, dtypinspecial_dtypes.items()}
special_dtypes_size= {key: dtype_byte_size(dtyp) forkey, dtypinspecial_dtypes.items()}
module_sizes=defaultdict(int)

module_list= []

module_list=named_persistent_module_tensors(model, recurse=True)

forname, tensorinmodule_list:
ifspecial_dtypesisnotNoneandnameinspecial_dtypes:
size=tensor.numel() *special_dtypes_size[name]
elifdtypeisNone:
size=tensor.numel() *dtype_byte_size(tensor.dtype)
elifstr(tensor.dtype).startswith(("torch.uint", "torch.int", "torch.bool")):
# According to the code in set_module_tensor_to_device, these types won't be converted
# so use their original size here
size=tensor.numel() *dtype_byte_size(tensor.dtype)
else:
size=tensor.numel() *min(dtype_size, dtype_byte_size(tensor.dtype))
name_parts=name.split(".")
foridxinrange(len(name_parts) +1):
module_sizes[".".join(name_parts[:idx])] +=size

returnmodule_sizes


defcast_maybe_tensor_dtype(maybe_tensor, current_dtype, target_dtype):
iftorch.is_tensor(maybe_tensor):
returnmaybe_tensor.to(target_dtype) ifmaybe_tensor.dtype==current_dtypeelsemaybe_tensor
ifisinstance(maybe_tensor, dict):
return {k: cast_maybe_tensor_dtype(v, current_dtype, target_dtype) fork, vinmaybe_tensor.items()}
ifisinstance(maybe_tensor, list):
return [cast_maybe_tensor_dtype(v, current_dtype, target_dtype) forvinmaybe_tensor]
returnmaybe_tensor


classModelUtilsTest(unittest.TestCase):
deftearDown(self):
super().tearDown()

deftest_missing_key_loading_warning_message(self):
withself.assertLogs("diffusers.models.modeling_utils", level="WARNING") aslogs:
UNet2DConditionModel.from_pretrained("hf-internal-testing/stable-diffusion-broken", subfolder="unet")

# make sure that error message states what keys are missing
assert"conv_out.bias"in" ".join(logs.output)

@parameterized.expand(
 [
 ("hf-internal-testing/tiny-stable-diffusion-pipe-variants-all-kinds", "unet", False),
 ("hf-internal-testing/tiny-stable-diffusion-pipe-variants-all-kinds", "unet", True),
 ("hf-internal-testing/tiny-sd-unet-with-sharded-ckpt", None, False),
 ("hf-internal-testing/tiny-sd-unet-with-sharded-ckpt", None, True),
 ]
 )
deftest_variant_sharded_ckpt_legacy_format_raises_warning(self, repo_id, subfolder, use_local):
defload_model(path):
kwargs= {"variant": "fp16"}
ifsubfolder:
kwargs["subfolder"] =subfolder
returnUNet2DConditionModel.from_pretrained(path, **kwargs)

withself.assertWarns(FutureWarning) aswarning:
ifuse_local:
withtempfile.TemporaryDirectory() astmpdirname:
tmpdirname=snapshot_download(repo_id=repo_id)
_=load_model(tmpdirname)
else:
_=load_model(repo_id)

warning_message=str(warning.warnings[0].message)
self.assertIn("This serialization format is now deprecated to standardize the serialization", warning_message)

# Local tests are already covered down below.
@parameterized.expand(
 [
 ("hf-internal-testing/tiny-sd-unet-sharded-latest-format", None, "fp16"),
 ("hf-internal-testing/tiny-sd-unet-sharded-latest-format-subfolder", "unet", "fp16"),
 ("hf-internal-testing/tiny-sd-unet-sharded-no-variants", None, None),
 ("hf-internal-testing/tiny-sd-unet-sharded-no-variants-subfolder", "unet", None),
 ]
 )
deftest_variant_sharded_ckpt_loads_from_hub(self, repo_id, subfolder, variant=None):
defload_model():
kwargs= {}
ifvariant:
kwargs["variant"] =variant
ifsubfolder:
kwargs["subfolder"] =subfolder
returnUNet2DConditionModel.from_pretrained(repo_id, **kwargs)

assertload_model()

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_model=UNet2DConditionModel.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-torch", subfolder="unet"
 )

# 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.
model=UNet2DConditionModel.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-torch", subfolder="unet", local_files_only=True
 )

forp1, p2inzip(orig_model.parameters(), model.parameters()):
ifp1.data.ne(p2.data).sum() >0:
assertFalse, "Parameters not the same!"

@unittest.skip("Flaky behaviour on CI. Re-enable after migrating to new runners")
@unittest.skipIf(torch_device=="mps", reason="Test not supported for MPS.")
deftest_one_request_upon_cached(self):
use_safetensors=False

withtempfile.TemporaryDirectory() astmpdirname:
withrequests_mock.mock(real_http=True) asm:
UNet2DConditionModel.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-torch",
subfolder="unet",
cache_dir=tmpdirname,
use_safetensors=use_safetensors,
 )

download_requests= [r.methodforrinm.request_history]
assertdownload_requests.count("HEAD") ==3, (
"3 HEAD requests one for config, one for model, and one for shard index file."
 )
assertdownload_requests.count("GET") ==2, "2 GET requests one for config, one for model"

withrequests_mock.mock(real_http=True) asm:
UNet2DConditionModel.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-torch",
subfolder="unet",
cache_dir=tmpdirname,
use_safetensors=use_safetensors,
 )

cache_requests= [r.methodforrinm.request_history]
assert"HEAD"==cache_requests[0] andlen(cache_requests) ==2, (
"We should call only `model_info` to check for commit hash and knowing if shard index is present."
 )

deftest_weight_overwrite(self):
withtempfile.TemporaryDirectory() astmpdirname, self.assertRaises(ValueError) aserror_context:
UNet2DConditionModel.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-torch",
subfolder="unet",
cache_dir=tmpdirname,
in_channels=9,
 )

# make sure that error message states what keys are missing
assert"Cannot load"instr(error_context.exception)

withtempfile.TemporaryDirectory() astmpdirname:
model=UNet2DConditionModel.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-torch",
subfolder="unet",
cache_dir=tmpdirname,
in_channels=9,
low_cpu_mem_usage=False,
ignore_mismatched_sizes=True,
 )

assertmodel.config.in_channels==9

@require_torch_accelerator
deftest_keep_modules_in_fp32(self):
r"""
 A simple tests to check if the modules under `_keep_in_fp32_modules` are kept in fp32 when we load the model in fp16/bf16
 Also ensures if inference works.
 """
fp32_modules=SD3Transformer2DModel._keep_in_fp32_modules

fortorch_dtypein [torch.bfloat16, torch.float16]:
SD3Transformer2DModel._keep_in_fp32_modules= ["proj_out"]

model=SD3Transformer2DModel.from_pretrained(
"hf-internal-testing/tiny-sd3-pipe", subfolder="transformer", torch_dtype=torch_dtype
 ).to(torch_device)

forname, moduleinmodel.named_modules():
ifisinstance(module, torch.nn.Linear):
ifnameinmodel._keep_in_fp32_modules:
self.assertTrue(module.weight.dtype==torch.float32)
else:
self.assertTrue(module.weight.dtype==torch_dtype)

defget_dummy_inputs():
batch_size=2
num_channels=4
height=width=embedding_dim=32
pooled_embedding_dim=embedding_dim*2
sequence_length=154

hidden_states=torch.randn((batch_size, num_channels, height, width)).to(torch_device)
encoder_hidden_states=torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device)
pooled_prompt_embeds=torch.randn((batch_size, pooled_embedding_dim)).to(torch_device)
timestep=torch.randint(0, 1000, size=(batch_size,)).to(torch_device)

return {
"hidden_states": hidden_states,
"encoder_hidden_states": encoder_hidden_states,
"pooled_projections": pooled_prompt_embeds,
"timestep": timestep,
 }

# test if inference works.
withtorch.no_grad() andtorch.amp.autocast(torch_device, dtype=torch_dtype):
input_dict_for_transformer=get_dummy_inputs()
model_inputs= {
k: v.to(device=torch_device) fork, vininput_dict_for_transformer.items() ifnotisinstance(v, bool)
 }
model_inputs.update({k: vfork, vininput_dict_for_transformer.items() ifknotinmodel_inputs})
_=model(**model_inputs)

SD3Transformer2DModel._keep_in_fp32_modules=fp32_modules


classUNetTesterMixin:
deftest_forward_with_norm_groups(self):
init_dict, inputs_dict=self.prepare_init_args_and_inputs_for_common()

init_dict["norm_num_groups"] =16
init_dict["block_out_channels"] = (16, 32)

model=self.model_class(**init_dict)
model.to(torch_device)
model.eval()

withtorch.no_grad():
output=model(**inputs_dict)

ifisinstance(output, dict):
output=output.to_tuple()[0]

self.assertIsNotNone(output)
expected_shape=inputs_dict["sample"].shape
self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")


classModelTesterMixin:
main_input_name=None# overwrite in model specific tester class
base_precision=1e-3
forward_requires_fresh_args=False
model_split_percents= [0.5, 0.7, 0.9]
uses_custom_attn_processor=False

defcheck_device_map_is_respected(self, model, device_map):
forparam_name, paraminmodel.named_parameters():
# Find device in device_map
whilelen(param_name) >0andparam_namenotindevice_map:
param_name=".".join(param_name.split(".")[:-1])
ifparam_namenotindevice_map:
raiseValueError("device map is incomplete, it does not contain any device for `param_name`.")

param_device=device_map[param_name]
ifparam_devicein ["cpu", "disk"]:
self.assertEqual(param.device, torch.device("meta"))
else:
self.assertEqual(param.device, torch.device(param_device))

deftest_from_save_pretrained(self, expected_max_diff=5e-5):
ifself.forward_requires_fresh_args:
model=self.model_class(**self.init_dict)
else:
init_dict, inputs_dict=self.prepare_init_args_and_inputs_for_common()
model=self.model_class(**init_dict)

ifhasattr(model, "set_default_attn_processor"):
model.set_default_attn_processor()
model.to(torch_device)
model.eval()

withtempfile.TemporaryDirectory() astmpdirname:
model.save_pretrained(tmpdirname, safe_serialization=False)
new_model=self.model_class.from_pretrained(tmpdirname)
ifhasattr(new_model, "set_default_attn_processor"):
new_model.set_default_attn_processor()
new_model.to(torch_device)

withtorch.no_grad():
ifself.forward_requires_fresh_args:
image=model(**self.inputs_dict(0))
else:
image=model(**inputs_dict)

ifisinstance(image, dict):
image=image.to_tuple()[0]

ifself.forward_requires_fresh_args:
new_image=new_model(**self.inputs_dict(0))
else:
new_image=new_model(**inputs_dict)

ifisinstance(new_image, dict):
new_image=new_image.to_tuple()[0]

max_diff= (image-new_image).abs().max().item()
self.assertLessEqual(max_diff, expected_max_diff, "Models give different forward passes")

deftest_getattr_is_correct(self):
init_dict, inputs_dict=self.prepare_init_args_and_inputs_for_common()
model=self.model_class(**init_dict)

# save some things to test
model.dummy_attribute=5
model.register_to_config(test_attribute=5)

logger=logging.get_logger("diffusers.models.modeling_utils")
# 30 for warning
logger.setLevel(30)
withCaptureLogger(logger) ascap_logger:
asserthasattr(model, "dummy_attribute")
assertgetattr(model, "dummy_attribute") ==5
assertmodel.dummy_attribute==5

# no warning should be thrown
assertcap_logger.out==""

logger=logging.get_logger("diffusers.models.modeling_utils")
# 30 for warning
logger.setLevel(30)
withCaptureLogger(logger) ascap_logger:
asserthasattr(model, "save_pretrained")
fn=model.save_pretrained
fn_1=getattr(model, "save_pretrained")

assertfn==fn_1
# no warning should be thrown
assertcap_logger.out==""

# warning should be thrown
withself.assertWarns(FutureWarning):
assertmodel.test_attribute==5

withself.assertWarns(FutureWarning):
assertgetattr(model, "test_attribute") ==5

withself.assertRaises(AttributeError) aserror:
model.does_not_exist

assertstr(error.exception) ==f"'{type(model).__name__}' object has no attribute 'does_not_exist'"

@unittest.skipIf(
torch_device!="npu"ornotis_torch_npu_available(),
reason="torch npu flash attention is only available with NPU and `torch_npu` installed",
 )
deftest_set_torch_npu_flash_attn_processor_determinism(self):
torch.use_deterministic_algorithms(False)
ifself.forward_requires_fresh_args:
model=self.model_class(**self.init_dict)
else:
init_dict, inputs_dict=self.prepare_init_args_and_inputs_for_common()
model=self.model_class(**init_dict)
model.to(torch_device)

ifnothasattr(model, "set_attn_processor"):
# If not has `set_attn_processor`, skip test
return

model.set_default_attn_processor()
assertall(type(proc) ==AttnProcessorNPUforprocinmodel.attn_processors.values())
withtorch.no_grad():
ifself.forward_requires_fresh_args:
output=model(**self.inputs_dict(0))[0]
else:
output=model(**inputs_dict)[0]

model.enable_npu_flash_attention()
assertall(type(proc) ==AttnProcessorNPUforprocinmodel.attn_processors.values())
withtorch.no_grad():
ifself.forward_requires_fresh_args:
output_2=model(**self.inputs_dict(0))[0]
else:
output_2=model(**inputs_dict)[0]

model.set_attn_processor(AttnProcessorNPU())
assertall(type(proc) ==AttnProcessorNPUforprocinmodel.attn_processors.values())
withtorch.no_grad():
ifself.forward_requires_fresh_args:
output_3=model(**self.inputs_dict(0))[0]
else:
output_3=model(**inputs_dict)[0]

torch.use_deterministic_algorithms(True)

asserttorch.allclose(output, output_2, atol=self.base_precision)
asserttorch.allclose(output, output_3, atol=self.base_precision)
asserttorch.allclose(output_2, output_3, atol=self.base_precision)

@unittest.skipIf(
torch_device!="cuda"ornotis_xformers_available(),
reason="XFormers attention is only available with CUDA and `xformers` installed",
 )
deftest_set_xformers_attn_processor_for_determinism(self):
torch.use_deterministic_algorithms(False)
ifself.forward_requires_fresh_args:
model=self.model_class(**self.init_dict)
else:
init_dict, inputs_dict=self.prepare_init_args_and_inputs_for_common()
model=self.model_class(**init_dict)
model.to(torch_device)

ifnothasattr(model, "set_attn_processor"):
# If not has `set_attn_processor`, skip test
return

ifnothasattr(model, "set_default_attn_processor"):
# If not has `set_attn_processor`, skip test
return

model.set_default_attn_processor()
assertall(type(proc) ==AttnProcessorforprocinmodel.attn_processors.values())
withtorch.no_grad():
ifself.forward_requires_fresh_args:
output=model(**self.inputs_dict(0))[0]
else:
output=model(**inputs_dict)[0]

model.enable_xformers_memory_efficient_attention()
assertall(type(proc) ==XFormersAttnProcessorforprocinmodel.attn_processors.values())
withtorch.no_grad():
ifself.forward_requires_fresh_args:
output_2=model(**self.inputs_dict(0))[0]
else:
output_2=model(**inputs_dict)[0]

model.set_attn_processor(XFormersAttnProcessor())
assertall(type(proc) ==XFormersAttnProcessorforprocinmodel.attn_processors.values())
withtorch.no_grad():
ifself.forward_requires_fresh_args:
output_3=model(**self.inputs_dict(0))[0]
else:
output_3=model(**inputs_dict)[0]

torch.use_deterministic_algorithms(True)

asserttorch.allclose(output, output_2, atol=self.base_precision)
asserttorch.allclose(output, output_3, atol=self.base_precision)
asserttorch.allclose(output_2, output_3, atol=self.base_precision)

@require_torch_accelerator
deftest_set_attn_processor_for_determinism(self):
ifself.uses_custom_attn_processor:
return

torch.use_deterministic_algorithms(False)
ifself.forward_requires_fresh_args:
model=self.model_class(**self.init_dict)
else:
init_dict, inputs_dict=self.prepare_init_args_and_inputs_for_common()
model=self.model_class(**init_dict)

model.to(torch_device)

ifnothasattr(model, "set_attn_processor"):
# If not has `set_attn_processor`, skip test
return

assertall(type(proc) ==AttnProcessor2_0forprocinmodel.attn_processors.values())
withtorch.no_grad():
ifself.forward_requires_fresh_args:
output_1=model(**self.inputs_dict(0))[0]
else:
output_1=model(**inputs_dict)[0]

model.set_default_attn_processor()
assertall(type(proc) ==AttnProcessorforprocinmodel.attn_processors.values())
withtorch.no_grad():
ifself.forward_requires_fresh_args:
output_2=model(**self.inputs_dict(0))[0]
else:
output_2=model(**inputs_dict)[0]

model.set_attn_processor(AttnProcessor2_0())
assertall(type(proc) ==AttnProcessor2_0forprocinmodel.attn_processors.values())
withtorch.no_grad():
ifself.forward_requires_fresh_args:
output_4=model(**self.inputs_dict(0))[0]
else:
output_4=model(**inputs_dict)[0]

model.set_attn_processor(AttnProcessor())
assertall(type(proc) ==AttnProcessorforprocinmodel.attn_processors.values())
withtorch.no_grad():
ifself.forward_requires_fresh_args:
output_5=model(**self.inputs_dict(0))[0]
else:
output_5=model(**inputs_dict)[0]

torch.use_deterministic_algorithms(True)

# make sure that outputs match
asserttorch.allclose(output_2, output_1, atol=self.base_precision)
asserttorch.allclose(output_2, output_4, atol=self.base_precision)
asserttorch.allclose(output_2, output_5, atol=self.base_precision)

deftest_from_save_pretrained_variant(self, expected_max_diff=5e-5):
ifself.forward_requires_fresh_args:
model=self.model_class(**self.init_dict)
else:
init_dict, inputs_dict=self.prepare_init_args_and_inputs_for_common()
model=self.model_class(**init_dict)

ifhasattr(model, "set_default_attn_processor"):
model.set_default_attn_processor()

model.to(torch_device)
model.eval()

withtempfile.TemporaryDirectory() astmpdirname:
model.save_pretrained(tmpdirname, variant="fp16", safe_serialization=False)
new_model=self.model_class.from_pretrained(tmpdirname, variant="fp16")
ifhasattr(new_model, "set_default_attn_processor"):
new_model.set_default_attn_processor()

# non-variant cannot be loaded
withself.assertRaises(OSError) aserror_context:
self.model_class.from_pretrained(tmpdirname)

# make sure that error message states what keys are missing
assert"Error no file named diffusion_pytorch_model.bin found in directory"instr(error_context.exception)

new_model.to(torch_device)

withtorch.no_grad():
ifself.forward_requires_fresh_args:
image=model(**self.inputs_dict(0))
else:
image=model(**inputs_dict)
ifisinstance(image, dict):
image=image.to_tuple()[0]

ifself.forward_requires_fresh_args:
new_image=new_model(**self.inputs_dict(0))
else:
new_image=new_model(**inputs_dict)

ifisinstance(new_image, dict):
new_image=new_image.to_tuple()[0]

max_diff= (image-new_image).abs().max().item()
self.assertLessEqual(max_diff, expected_max_diff, "Models give different forward passes")

@is_torch_compile
@require_torch_2
@unittest.skipIf(
get_python_version== (3, 12),
reason="Torch Dynamo isn't yet supported for Python 3.12.",
 )
deftest_from_save_pretrained_dynamo(self):
init_dict, _=self.prepare_init_args_and_inputs_for_common()
inputs= [init_dict, self.model_class]
run_test_in_subprocess(test_case=self, target_func=_test_from_save_pretrained_dynamo, inputs=inputs)

deftest_from_save_pretrained_dtype(self):
init_dict, inputs_dict=self.prepare_init_args_and_inputs_for_common()

model=self.model_class(**init_dict)
model.to(torch_device)
model.eval()

fordtypein [torch.float32, torch.float16, torch.bfloat16]:
iftorch_device=="mps"anddtype==torch.bfloat16:
continue
withtempfile.TemporaryDirectory() astmpdirname:
model.to(dtype)
model.save_pretrained(tmpdirname, safe_serialization=False)
new_model=self.model_class.from_pretrained(tmpdirname, low_cpu_mem_usage=True, torch_dtype=dtype)
assertnew_model.dtype==dtype
if (
hasattr(self.model_class, "_keep_in_fp32_modules")
andself.model_class._keep_in_fp32_modulesisNone
 ):
new_model=self.model_class.from_pretrained(
tmpdirname, low_cpu_mem_usage=False, torch_dtype=dtype
 )
assertnew_model.dtype==dtype

deftest_determinism(self, expected_max_diff=1e-5):
ifself.forward_requires_fresh_args:
model=self.model_class(**self.init_dict)
else:
init_dict, inputs_dict=self.prepare_init_args_and_inputs_for_common()
model=self.model_class(**init_dict)
model.to(torch_device)
model.eval()

withtorch.no_grad():
ifself.forward_requires_fresh_args:
first=model(**self.inputs_dict(0))
else:
first=model(**inputs_dict)
ifisinstance(first, dict):
first=first.to_tuple()[0]

ifself.forward_requires_fresh_args:
second=model(**self.inputs_dict(0))
else:
second=model(**inputs_dict)
ifisinstance(second, dict):
second=second.to_tuple()[0]

out_1=first.cpu().numpy()
out_2=second.cpu().numpy()
out_1=out_1[~np.isnan(out_1)]
out_2=out_2[~np.isnan(out_2)]
max_diff=np.amax(np.abs(out_1-out_2))
self.assertLessEqual(max_diff, expected_max_diff)

deftest_output(self, expected_output_shape=None):
init_dict, inputs_dict=self.prepare_init_args_and_inputs_for_common()
model=self.model_class(**init_dict)
model.to(torch_device)
model.eval()

withtorch.no_grad():
output=model(**inputs_dict)

ifisinstance(output, dict):
output=output.to_tuple()[0]

self.assertIsNotNone(output)

# input & output have to have the same shape
input_tensor=inputs_dict[self.main_input_name]

ifexpected_output_shapeisNone:
expected_shape=input_tensor.shape
self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
else:
self.assertEqual(output.shape, expected_output_shape, "Input and output shapes do not match")

deftest_model_from_pretrained(self):
init_dict, inputs_dict=self.prepare_init_args_and_inputs_for_common()

model=self.model_class(**init_dict)
model.to(torch_device)
model.eval()

# test if the model can be loaded from the config
# and has all the expected shape
withtempfile.TemporaryDirectory() astmpdirname:
model.save_pretrained(tmpdirname, safe_serialization=False)
new_model=self.model_class.from_pretrained(tmpdirname)
new_model.to(torch_device)
new_model.eval()

# check if all parameters shape are the same
forparam_nameinmodel.state_dict().keys():
param_1=model.state_dict()[param_name]
param_2=new_model.state_dict()[param_name]
self.assertEqual(param_1.shape, param_2.shape)

withtorch.no_grad():
output_1=model(**inputs_dict)

ifisinstance(output_1, dict):
output_1=output_1.to_tuple()[0]

output_2=new_model(**inputs_dict)

ifisinstance(output_2, dict):
output_2=output_2.to_tuple()[0]

self.assertEqual(output_1.shape, output_2.shape)

@require_torch_accelerator_with_training
deftest_training(self):
init_dict, inputs_dict=self.prepare_init_args_and_inputs_for_common()

model=self.model_class(**init_dict)
model.to(torch_device)
model.train()
output=model(**inputs_dict)

ifisinstance(output, dict):
output=output.to_tuple()[0]

input_tensor=inputs_dict[self.main_input_name]
noise=torch.randn((input_tensor.shape[0],) +self.output_shape).to(torch_device)
loss=torch.nn.functional.mse_loss(output, noise)
loss.backward()

@require_torch_accelerator_with_training
deftest_ema_training(self):
init_dict, inputs_dict=self.prepare_init_args_and_inputs_for_common()

model=self.model_class(**init_dict)
model.to(torch_device)
model.train()
ema_model=EMAModel(model.parameters())

output=model(**inputs_dict)

ifisinstance(output, dict):
output=output.to_tuple()[0]

input_tensor=inputs_dict[self.main_input_name]
noise=torch.randn((input_tensor.shape[0],) +self.output_shape).to(torch_device)
loss=torch.nn.functional.mse_loss(output, noise)
loss.backward()
ema_model.step(model.parameters())

deftest_outputs_equivalence(self):
defset_nan_tensor_to_zero(t):
# Temporary fallback until `aten::_index_put_impl_` is implemented in mps
# Track progress in https://github.com/pytorch/pytorch/issues/77764
device=t.device
ifdevice.type=="mps":
t=t.to("cpu")
t[t!=t] =0
returnt.to(device)

defrecursive_check(tuple_object, dict_object):
ifisinstance(tuple_object, (List, Tuple)):
fortuple_iterable_value, dict_iterable_valueinzip(tuple_object, dict_object.values()):
recursive_check(tuple_iterable_value, dict_iterable_value)
elifisinstance(tuple_object, Dict):
fortuple_iterable_value, dict_iterable_valueinzip(tuple_object.values(), dict_object.values()):
recursive_check(tuple_iterable_value, dict_iterable_value)
eliftuple_objectisNone:
return
else:
self.assertTrue(
torch.allclose(
set_nan_tensor_to_zero(tuple_object), set_nan_tensor_to_zero(dict_object), atol=1e-5
 ),
msg=(
"Tuple and dict output are not equal. Difference:"
f" {torch.max(torch.abs(tuple_object-dict_object))}. Tuple has `nan`:"
f" {torch.isnan(tuple_object).any()} and `inf`: {torch.isinf(tuple_object)}. Dict has"
f" `nan`: {torch.isnan(dict_object).any()} and `inf`: {torch.isinf(dict_object)}."
 ),
 )

ifself.forward_requires_fresh_args:
model=self.model_class(**self.init_dict)
else:
init_dict, inputs_dict=self.prepare_init_args_and_inputs_for_common()
model=self.model_class(**init_dict)

model.to(torch_device)
model.eval()

withtorch.no_grad():
ifself.forward_requires_fresh_args:
outputs_dict=model(**self.inputs_dict(0))
outputs_tuple=model(**self.inputs_dict(0), return_dict=False)
else:
outputs_dict=model(**inputs_dict)
outputs_tuple=model(**inputs_dict, return_dict=False)

recursive_check(outputs_tuple, outputs_dict)

@require_torch_accelerator_with_training
deftest_enable_disable_gradient_checkpointing(self):
ifnotself.model_class._supports_gradient_checkpointing:
return# Skip test if model does not support gradient checkpointing

init_dict, _=self.prepare_init_args_and_inputs_for_common()

# at init model should have gradient checkpointing disabled
model=self.model_class(**init_dict)
self.assertFalse(model.is_gradient_checkpointing)

# check enable works
model.enable_gradient_checkpointing()
self.assertTrue(model.is_gradient_checkpointing)

# check disable works
model.disable_gradient_checkpointing()
self.assertFalse(model.is_gradient_checkpointing)

@require_torch_accelerator_with_training
deftest_effective_gradient_checkpointing(self, loss_tolerance=1e-5, param_grad_tol=5e-5, skip: set[str] = {}):
ifnotself.model_class._supports_gradient_checkpointing:
return# Skip test if model does not support gradient checkpointing

# enable deterministic behavior for gradient checkpointing
init_dict, inputs_dict=self.prepare_init_args_and_inputs_for_common()
inputs_dict_copy=copy.deepcopy(inputs_dict)
torch.manual_seed(0)
model=self.model_class(**init_dict)
model.to(torch_device)

assertnotmodel.is_gradient_checkpointingandmodel.training

out=model(**inputs_dict).sample
# run the backwards pass on the model. For backwards pass, for simplicity purpose,
# we won't calculate the loss and rather backprop on out.sum()
model.zero_grad()

labels=torch.randn_like(out)
loss= (out-labels).mean()
loss.backward()

# re-instantiate the model now enabling gradient checkpointing
torch.manual_seed(0)
model_2=self.model_class(**init_dict)
# clone model
model_2.load_state_dict(model.state_dict())
model_2.to(torch_device)
model_2.enable_gradient_checkpointing()

assertmodel_2.is_gradient_checkpointingandmodel_2.training

out_2=model_2(**inputs_dict_copy).sample
# run the backwards pass on the model. For backwards pass, for simplicity purpose,
# we won't calculate the loss and rather backprop on out.sum()
model_2.zero_grad()
loss_2= (out_2-labels).mean()
loss_2.backward()

# compare the output and parameters gradients
self.assertTrue((loss-loss_2).abs() <loss_tolerance)
named_params=dict(model.named_parameters())
named_params_2=dict(model_2.named_parameters())