test_pipelines_common.py

importgc
importinspect
importjson
importos
importtempfile
importunittest
importuuid
fromtypingimportAny, Callable, Dict, Union

importnumpyasnp
importPIL.Image
importtorch
importtorch.nnasnn
fromhuggingface_hubimportModelCard, delete_repo
fromhuggingface_hub.utilsimportis_jinja_available
fromtransformersimportCLIPTextConfig, CLIPTextModel, CLIPTokenizer

importdiffusers
fromdiffusersimport (
AsymmetricAutoencoderKL,
AutoencoderKL,
AutoencoderTiny,
ConsistencyDecoderVAE,
DDIMScheduler,
DiffusionPipeline,
FasterCacheConfig,
KolorsPipeline,
PyramidAttentionBroadcastConfig,
StableDiffusionPipeline,
StableDiffusionXLPipeline,
UNet2DConditionModel,
apply_faster_cache,
)
fromdiffusers.hooksimportapply_group_offloading
fromdiffusers.hooks.faster_cacheimportFasterCacheBlockHook, FasterCacheDenoiserHook
fromdiffusers.hooks.pyramid_attention_broadcastimportPyramidAttentionBroadcastHook
fromdiffusers.image_processorimportVaeImageProcessor
fromdiffusers.loadersimportFluxIPAdapterMixin, IPAdapterMixin
fromdiffusers.models.attention_processorimportAttnProcessor
fromdiffusers.models.controlnets.controlnet_xsimportUNetControlNetXSModel
fromdiffusers.models.unets.unet_3d_conditionimportUNet3DConditionModel
fromdiffusers.models.unets.unet_i2vgen_xlimportI2VGenXLUNet
fromdiffusers.models.unets.unet_motion_modelimportUNetMotionModel
fromdiffusers.pipelines.pipeline_utilsimportStableDiffusionMixin
fromdiffusers.schedulersimportKarrasDiffusionSchedulers
fromdiffusers.utilsimportlogging
fromdiffusers.utils.import_utilsimportis_xformers_available
fromdiffusers.utils.source_code_parsing_utilsimportReturnNameVisitor
fromdiffusers.utils.testing_utilsimport (
CaptureLogger,
backend_empty_cache,
require_accelerate_version_greater,
require_accelerator,
require_hf_hub_version_greater,
require_torch,
require_torch_accelerator,
require_transformers_version_greater,
skip_mps,
torch_device,
)

from ..models.autoencoders.vaeimport (
get_asym_autoencoder_kl_config,
get_autoencoder_kl_config,
get_autoencoder_tiny_config,
get_consistency_vae_config,
)
from ..models.transformers.test_models_transformer_fluximportcreate_flux_ip_adapter_state_dict
from ..models.unets.test_models_unet_2d_conditionimport (
create_ip_adapter_faceid_state_dict,
create_ip_adapter_state_dict,
)
from ..others.test_utilsimportTOKEN, USER, is_staging_test


defto_np(tensor):
ifisinstance(tensor, torch.Tensor):
tensor=tensor.detach().cpu().numpy()

returntensor


defcheck_same_shape(tensor_list):
shapes= [tensor.shapefortensorintensor_list]
returnall(shape==shapes[0] forshapeinshapes[1:])


defcheck_qkv_fusion_matches_attn_procs_length(model, original_attn_processors):
current_attn_processors=model.attn_processors
returnlen(current_attn_processors) ==len(original_attn_processors)


defcheck_qkv_fusion_processors_exist(model):
current_attn_processors=model.attn_processors
proc_names= [v.__class__.__name__for_, vincurrent_attn_processors.items()]
returnall(p.startswith("Fused") forpinproc_names)


classSDFunctionTesterMixin:
"""
 This mixin is designed to be used with PipelineTesterMixin and unittest.TestCase classes.
 It provides a set of common tests for PyTorch pipeline that inherit from StableDiffusionMixin, e.g. vae_slicing, vae_tiling, freeu, etc.
 """

deftest_vae_slicing(self, image_count=4):
device="cpu"# ensure determinism for the device-dependent torch.Generator
components=self.get_dummy_components()
# components["scheduler"] = LMSDiscreteScheduler.from_config(components["scheduler"].config)
pipe=self.pipeline_class(**components)
pipe=pipe.to(device)
pipe.set_progress_bar_config(disable=None)

inputs=self.get_dummy_inputs(device)
inputs["prompt"] = [inputs["prompt"]] *image_count
if"image"ininputs: # fix batch size mismatch in I2V_Gen pipeline
inputs["image"] = [inputs["image"]] *image_count
output_1=pipe(**inputs)

# make sure sliced vae decode yields the same result
pipe.enable_vae_slicing()
inputs=self.get_dummy_inputs(device)
inputs["prompt"] = [inputs["prompt"]] *image_count
if"image"ininputs:
inputs["image"] = [inputs["image"]] *image_count
inputs["return_dict"] =False
output_2=pipe(**inputs)

assertnp.abs(output_2[0].flatten() -output_1[0].flatten()).max() <1e-2

deftest_vae_tiling(self):
components=self.get_dummy_components()

# make sure here that pndm scheduler skips prk
if"safety_checker"incomponents:
components["safety_checker"] =None
pipe=self.pipeline_class(**components)
pipe=pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)

inputs=self.get_dummy_inputs(torch_device)
inputs["return_dict"] =False

# Test that tiled decode at 512x512 yields the same result as the non-tiled decode
output_1=pipe(**inputs)[0]

# make sure tiled vae decode yields the same result
pipe.enable_vae_tiling()
inputs=self.get_dummy_inputs(torch_device)
inputs["return_dict"] =False
output_2=pipe(**inputs)[0]

assertnp.abs(to_np(output_2) -to_np(output_1)).max() <5e-1

# test that tiled decode works with various shapes
shapes= [(1, 4, 73, 97), (1, 4, 65, 49)]
withtorch.no_grad():
forshapeinshapes:
zeros=torch.zeros(shape).to(torch_device)
pipe.vae.decode(zeros)

# MPS currently doesn't support ComplexFloats, which are required for FreeU - see https://github.com/huggingface/diffusers/issues/7569.
@skip_mps
deftest_freeu(self):
components=self.get_dummy_components()
pipe=self.pipeline_class(**components)
pipe=pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)

# Normal inference
inputs=self.get_dummy_inputs(torch_device)
inputs["return_dict"] =False
inputs["output_type"] ="np"
output=pipe(**inputs)[0]

# FreeU-enabled inference
pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.2, b2=1.4)
inputs=self.get_dummy_inputs(torch_device)
inputs["return_dict"] =False
inputs["output_type"] ="np"
output_freeu=pipe(**inputs)[0]

# FreeU-disabled inference
pipe.disable_freeu()
freeu_keys= {"s1", "s2", "b1", "b2"}
forupsample_blockinpipe.unet.up_blocks:
forkeyinfreeu_keys:
assertgetattr(upsample_block, key) isNone, f"Disabling of FreeU should have set {key} to None."

inputs=self.get_dummy_inputs(torch_device)
inputs["return_dict"] =False
inputs["output_type"] ="np"
output_no_freeu=pipe(**inputs)[0]

assertnotnp.allclose(output[0, -3:, -3:, -1], output_freeu[0, -3:, -3:, -1]), (
"Enabling of FreeU should lead to different results."
 )
assertnp.allclose(output, output_no_freeu, atol=1e-2), (
f"Disabling of FreeU should lead to results similar to the default pipeline results but Max Abs Error={np.abs(output_no_freeu-output).max()}."
 )

deftest_fused_qkv_projections(self):
device="cpu"# ensure determinism for the device-dependent torch.Generator
components=self.get_dummy_components()
pipe=self.pipeline_class(**components)
pipe=pipe.to(device)
pipe.set_progress_bar_config(disable=None)

inputs=self.get_dummy_inputs(device)
inputs["return_dict"] =False
image=pipe(**inputs)[0]
original_image_slice=image[0, -3:, -3:, -1]

pipe.fuse_qkv_projections()
for_, componentinpipe.components.items():
if (
isinstance(component, nn.Module)
andhasattr(component, "original_attn_processors")
andcomponent.original_attn_processorsisnotNone
 ):
assertcheck_qkv_fusion_processors_exist(component), (
"Something wrong with the fused attention processors. Expected all the attention processors to be fused."
 )
assertcheck_qkv_fusion_matches_attn_procs_length(component, component.original_attn_processors), (
"Something wrong with the attention processors concerning the fused QKV projections."
 )

inputs=self.get_dummy_inputs(device)
inputs["return_dict"] =False
image_fused=pipe(**inputs)[0]
image_slice_fused=image_fused[0, -3:, -3:, -1]

pipe.unfuse_qkv_projections()
inputs=self.get_dummy_inputs(device)
inputs["return_dict"] =False
image_disabled=pipe(**inputs)[0]
image_slice_disabled=image_disabled[0, -3:, -3:, -1]

assertnp.allclose(original_image_slice, image_slice_fused, atol=1e-2, rtol=1e-2), (
"Fusion of QKV projections shouldn't affect the outputs."
 )
assertnp.allclose(image_slice_fused, image_slice_disabled, atol=1e-2, rtol=1e-2), (
"Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled."
 )
assertnp.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), (
"Original outputs should match when fused QKV projections are disabled."
 )


classIPAdapterTesterMixin:
"""
 This mixin is designed to be used with PipelineTesterMixin and unittest.TestCase classes.
 It provides a set of common tests for pipelines that support IP Adapters.
 """

deftest_pipeline_signature(self):
parameters=inspect.signature(self.pipeline_class.__call__).parameters

assertissubclass(self.pipeline_class, IPAdapterMixin)
self.assertIn(
"ip_adapter_image",
parameters,
"`ip_adapter_image` argument must be supported by the `__call__` method",
 )
self.assertIn(
"ip_adapter_image_embeds",
parameters,
"`ip_adapter_image_embeds` argument must be supported by the `__call__` method",
 )

def_get_dummy_image_embeds(self, cross_attention_dim: int=32):
returntorch.randn((2, 1, cross_attention_dim), device=torch_device)

def_get_dummy_faceid_image_embeds(self, cross_attention_dim: int=32):
returntorch.randn((2, 1, 1, cross_attention_dim), device=torch_device)

def_get_dummy_masks(self, input_size: int=64):
_masks=torch.zeros((1, 1, input_size, input_size), device=torch_device)
_masks[0, :, :, : int(input_size/2)] =1
return_masks

def_modify_inputs_for_ip_adapter_test(self, inputs: Dict[str, Any]):
parameters=inspect.signature(self.pipeline_class.__call__).parameters
if"image"inparameters.keys() and"strength"inparameters.keys():
inputs["num_inference_steps"] =4

inputs["output_type"] ="np"
inputs["return_dict"] =False
returninputs

deftest_ip_adapter(self, expected_max_diff: float=1e-4, expected_pipe_slice=None):
r"""Tests for IP-Adapter.

 The following scenarios are tested:
 - Single IP-Adapter with scale=0 should produce same output as no IP-Adapter.
 - Multi IP-Adapter with scale=0 should produce same output as no IP-Adapter.
 - Single IP-Adapter with scale!=0 should produce different output compared to no IP-Adapter.
 - Multi IP-Adapter with scale!=0 should produce different output compared to no IP-Adapter.
 """
# Raising the tolerance for this test when it's run on a CPU because we
# compare against static slices and that can be shaky (with a VVVV low probability).
expected_max_diff=9e-4iftorch_device=="cpu"elseexpected_max_diff

components=self.get_dummy_components()
pipe=self.pipeline_class(**components).to(torch_device)
pipe.set_progress_bar_config(disable=None)
cross_attention_dim=pipe.unet.config.get("cross_attention_dim", 32)

# forward pass without ip adapter
inputs=self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
ifexpected_pipe_sliceisNone:
output_without_adapter=pipe(**inputs)[0]
else:
output_without_adapter=expected_pipe_slice

# 1. Single IP-Adapter test cases
adapter_state_dict=create_ip_adapter_state_dict(pipe.unet)
pipe.unet._load_ip_adapter_weights(adapter_state_dict)

# forward pass with single ip adapter, but scale=0 which should have no effect
inputs=self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)]
pipe.set_ip_adapter_scale(0.0)
output_without_adapter_scale=pipe(**inputs)[0]
ifexpected_pipe_sliceisnotNone:
output_without_adapter_scale=output_without_adapter_scale[0, -3:, -3:, -1].flatten()

# forward pass with single ip adapter, but with scale of adapter weights
inputs=self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)]
pipe.set_ip_adapter_scale(42.0)
output_with_adapter_scale=pipe(**inputs)[0]
ifexpected_pipe_sliceisnotNone:
output_with_adapter_scale=output_with_adapter_scale[0, -3:, -3:, -1].flatten()

max_diff_without_adapter_scale=np.abs(output_without_adapter_scale-output_without_adapter).max()
max_diff_with_adapter_scale=np.abs(output_with_adapter_scale-output_without_adapter).max()

self.assertLess(
max_diff_without_adapter_scale,
expected_max_diff,
"Output without ip-adapter must be same as normal inference",
 )
self.assertGreater(
max_diff_with_adapter_scale, 1e-2, "Output with ip-adapter must be different from normal inference"
 )

# 2. Multi IP-Adapter test cases
adapter_state_dict_1=create_ip_adapter_state_dict(pipe.unet)
adapter_state_dict_2=create_ip_adapter_state_dict(pipe.unet)
pipe.unet._load_ip_adapter_weights([adapter_state_dict_1, adapter_state_dict_2])

# forward pass with multi ip adapter, but scale=0 which should have no effect
inputs=self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] *2
pipe.set_ip_adapter_scale([0.0, 0.0])
output_without_multi_adapter_scale=pipe(**inputs)[0]
ifexpected_pipe_sliceisnotNone:
output_without_multi_adapter_scale=output_without_multi_adapter_scale[0, -3:, -3:, -1].flatten()

# forward pass with multi ip adapter, but with scale of adapter weights
inputs=self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] *2
pipe.set_ip_adapter_scale([42.0, 42.0])
output_with_multi_adapter_scale=pipe(**inputs)[0]
ifexpected_pipe_sliceisnotNone:
output_with_multi_adapter_scale=output_with_multi_adapter_scale[0, -3:, -3:, -1].flatten()

max_diff_without_multi_adapter_scale=np.abs(
output_without_multi_adapter_scale-output_without_adapter
 ).max()
max_diff_with_multi_adapter_scale=np.abs(output_with_multi_adapter_scale-output_without_adapter).max()
self.assertLess(
max_diff_without_multi_adapter_scale,
expected_max_diff,
"Output without multi-ip-adapter must be same as normal inference",
 )
self.assertGreater(
max_diff_with_multi_adapter_scale,
1e-2,
"Output with multi-ip-adapter scale must be different from normal inference",
 )

deftest_ip_adapter_cfg(self, expected_max_diff: float=1e-4):
parameters=inspect.signature(self.pipeline_class.__call__).parameters

if"guidance_scale"notinparameters:
return

components=self.get_dummy_components()
pipe=self.pipeline_class(**components).to(torch_device)
pipe.set_progress_bar_config(disable=None)
cross_attention_dim=pipe.unet.config.get("cross_attention_dim", 32)

adapter_state_dict=create_ip_adapter_state_dict(pipe.unet)
pipe.unet._load_ip_adapter_weights(adapter_state_dict)
pipe.set_ip_adapter_scale(1.0)

# forward pass with CFG not applied
inputs=self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)[0].unsqueeze(0)]
inputs["guidance_scale"] =1.0
out_no_cfg=pipe(**inputs)[0]

# forward pass with CFG applied
inputs=self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)]
inputs["guidance_scale"] =7.5
out_cfg=pipe(**inputs)[0]

assertout_cfg.shape==out_no_cfg.shape

deftest_ip_adapter_masks(self, expected_max_diff: float=1e-4):
components=self.get_dummy_components()
pipe=self.pipeline_class(**components).to(torch_device)
pipe.set_progress_bar_config(disable=None)
cross_attention_dim=pipe.unet.config.get("cross_attention_dim", 32)
sample_size=pipe.unet.config.get("sample_size", 32)
block_out_channels=pipe.vae.config.get("block_out_channels", [128, 256, 512, 512])
input_size=sample_size* (2** (len(block_out_channels) -1))

# forward pass without ip adapter
inputs=self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
output_without_adapter=pipe(**inputs)[0]
output_without_adapter=output_without_adapter[0, -3:, -3:, -1].flatten()

adapter_state_dict=create_ip_adapter_state_dict(pipe.unet)
pipe.unet._load_ip_adapter_weights(adapter_state_dict)

# forward pass with single ip adapter and masks, but scale=0 which should have no effect
inputs=self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)]
inputs["cross_attention_kwargs"] = {"ip_adapter_masks": [self._get_dummy_masks(input_size)]}
pipe.set_ip_adapter_scale(0.0)
output_without_adapter_scale=pipe(**inputs)[0]
output_without_adapter_scale=output_without_adapter_scale[0, -3:, -3:, -1].flatten()

# forward pass with single ip adapter and masks, but with scale of adapter weights
inputs=self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)]
inputs["cross_attention_kwargs"] = {"ip_adapter_masks": [self._get_dummy_masks(input_size)]}
pipe.set_ip_adapter_scale(42.0)
output_with_adapter_scale=pipe(**inputs)[0]
output_with_adapter_scale=output_with_adapter_scale[0, -3:, -3:, -1].flatten()

max_diff_without_adapter_scale=np.abs(output_without_adapter_scale-output_without_adapter).max()
max_diff_with_adapter_scale=np.abs(output_with_adapter_scale-output_without_adapter).max()

self.assertLess(
max_diff_without_adapter_scale,
expected_max_diff,
"Output without ip-adapter must be same as normal inference",
 )
self.assertGreater(
max_diff_with_adapter_scale, 1e-3, "Output with ip-adapter must be different from normal inference"
 )

deftest_ip_adapter_faceid(self, expected_max_diff: float=1e-4):
components=self.get_dummy_components()
pipe=self.pipeline_class(**components).to(torch_device)
pipe.set_progress_bar_config(disable=None)
cross_attention_dim=pipe.unet.config.get("cross_attention_dim", 32)

# forward pass without ip adapter
inputs=self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
output_without_adapter=pipe(**inputs)[0]
output_without_adapter=output_without_adapter[0, -3:, -3:, -1].flatten()

adapter_state_dict=create_ip_adapter_faceid_state_dict(pipe.unet)
pipe.unet._load_ip_adapter_weights(adapter_state_dict)

# forward pass with single ip adapter, but scale=0 which should have no effect
inputs=self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
inputs["ip_adapter_image_embeds"] = [self._get_dummy_faceid_image_embeds(cross_attention_dim)]
pipe.set_ip_adapter_scale(0.0)
output_without_adapter_scale=pipe(**inputs)[0]
output_without_adapter_scale=output_without_adapter_scale[0, -3:, -3:, -1].flatten()

# forward pass with single ip adapter, but with scale of adapter weights
inputs=self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
inputs["ip_adapter_image_embeds"] = [self._get_dummy_faceid_image_embeds(cross_attention_dim)]
pipe.set_ip_adapter_scale(42.0)
output_with_adapter_scale=pipe(**inputs)[0]
output_with_adapter_scale=output_with_adapter_scale[0, -3:, -3:, -1].flatten()

max_diff_without_adapter_scale=np.abs(output_without_adapter_scale-output_without_adapter).max()
max_diff_with_adapter_scale=np.abs(output_with_adapter_scale-output_without_adapter).max()

self.assertLess(
max_diff_without_adapter_scale,
expected_max_diff,
"Output without ip-adapter must be same as normal inference",
 )
self.assertGreater(
max_diff_with_adapter_scale, 1e-3, "Output with ip-adapter must be different from normal inference"
 )


classFluxIPAdapterTesterMixin:
"""
 This mixin is designed to be used with PipelineTesterMixin and unittest.TestCase classes.
 It provides a set of common tests for pipelines that support IP Adapters.
 """

deftest_pipeline_signature(self):
parameters=inspect.signature(self.pipeline_class.__call__).parameters

assertissubclass(self.pipeline_class, FluxIPAdapterMixin)
self.assertIn(
"ip_adapter_image",
parameters,
"`ip_adapter_image` argument must be supported by the `__call__` method",
 )
self.assertIn(
"ip_adapter_image_embeds",
parameters,
"`ip_adapter_image_embeds` argument must be supported by the `__call__` method",
 )

def_get_dummy_image_embeds(self, image_embed_dim: int=768):
returntorch.randn((1, 1, image_embed_dim), device=torch_device)

def_modify_inputs_for_ip_adapter_test(self, inputs: Dict[str, Any]):
inputs["negative_prompt"] =""
inputs["true_cfg_scale"] =4.0
inputs["output_type"] ="np"
inputs["return_dict"] =False
returninputs

deftest_ip_adapter(self, expected_max_diff: float=1e-4, expected_pipe_slice=None):
r"""Tests for IP-Adapter.

 The following scenarios are tested:
 - Single IP-Adapter with scale=0 should produce same output as no IP-Adapter.
 - Multi IP-Adapter with scale=0 should produce same output as no IP-Adapter.
 - Single IP-Adapter with scale!=0 should produce different output compared to no IP-Adapter.
 - Multi IP-Adapter with scale!=0 should produce different output compared to no IP-Adapter.
 """
# Raising the tolerance for this test when it's run on a CPU because we
# compare against static slices and that can be shaky (with a VVVV low probability).
expected_max_diff=9e-4iftorch_device=="cpu"elseexpected_max_diff

components=self.get_dummy_components()
pipe=self.pipeline_class(**components).to(torch_device)
pipe.set_progress_bar_config(disable=None)
image_embed_dim=pipe.transformer.config.pooled_projection_dim

# forward pass without ip adapter
inputs=self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
ifexpected_pipe_sliceisNone:
output_without_adapter=pipe(**inputs)[0]
else:
output_without_adapter=expected_pipe_slice

# 1. Single IP-Adapter test cases
adapter_state_dict=create_flux_ip_adapter_state_dict(pipe.transformer)
pipe.transformer._load_ip_adapter_weights(adapter_state_dict)

# forward pass with single ip adapter, but scale=0 which should have no effect
inputs=self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)]
inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)]
pipe.set_ip_adapter_scale(0.0)
output_without_adapter_scale=pipe(**inputs)[0]
ifexpected_pipe_sliceisnotNone:
output_without_adapter_scale=output_without_adapter_scale[0, -3:, -3:, -1].flatten()

# forward pass with single ip adapter, but with scale of adapter weights
inputs=self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)]
inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)]
pipe.set_ip_adapter_scale(42.0)
output_with_adapter_scale=pipe(**inputs)[0]
ifexpected_pipe_sliceisnotNone:
output_with_adapter_scale=output_with_adapter_scale[0, -3:, -3:, -1].flatten()

max_diff_without_adapter_scale=np.abs(output_without_adapter_scale-output_without_adapter).max()
max_diff_with_adapter_scale=np.abs(output_with_adapter_scale-output_without_adapter).max()

self.assertLess(
max_diff_without_adapter_scale,
expected_max_diff,
"Output without ip-adapter must be same as normal inference",
 )
self.assertGreater(
max_diff_with_adapter_scale, 1e-2, "Output with ip-adapter must be different from normal inference"
 )

# 2. Multi IP-Adapter test cases
adapter_state_dict_1=create_flux_ip_adapter_state_dict(pipe.transformer)
adapter_state_dict_2=create_flux_ip_adapter_state_dict(pipe.transformer)
pipe.transformer._load_ip_adapter_weights([adapter_state_dict_1, adapter_state_dict_2])

# forward pass with multi ip adapter, but scale=0 which should have no effect
inputs=self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] *2
inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] *2
pipe.set_ip_adapter_scale([0.0, 0.0])
output_without_multi_adapter_scale=pipe(**inputs)[0]
ifexpected_pipe_sliceisnotNone:
output_without_multi_adapter_scale=output_without_multi_adapter_scale[0, -3:, -3:, -1].flatten()

# forward pass with multi ip adapter, but with scale of adapter weights
inputs=self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] *2
inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] *2
pipe.set_ip_adapter_scale([42.0, 42.0])
output_with_multi_adapter_scale=pipe(**inputs)[0]
ifexpected_pipe_sliceisnotNone:
output_with_multi_adapter_scale=output_with_multi_adapter_scale[0, -3:, -3:, -1].flatten()

max_diff_without_multi_adapter_scale=np.abs(
output_without_multi_adapter_scale-output_without_adapter
 ).max()
max_diff_with_multi_adapter_scale=np.abs(output_with_multi_adapter_scale-output_without_adapter).max()
self.assertLess(
max_diff_without_multi_adapter_scale,
expected_max_diff,
"Output without multi-ip-adapter must be same as normal inference",
 )
self.assertGreater(
max_diff_with_multi_adapter_scale,
1e-2,
"Output with multi-ip-adapter scale must be different from normal inference",
 )


classPipelineLatentTesterMixin:
"""
 This mixin is designed to be used with PipelineTesterMixin and unittest.TestCase classes.
 It provides a set of common tests for PyTorch pipeline that has vae, e.g.
 equivalence of different input and output types, etc.
 """

@property
defimage_params(self) ->frozenset:
raiseNotImplementedError(
"You need to set the attribute `image_params` in the child test class. "
"`image_params` are tested for if all accepted input image types (i.e. `pt`,`pil`,`np`) are producing same results"
 )

@property
defimage_latents_params(self) ->frozenset:
raiseNotImplementedError(
"You need to set the attribute `image_latents_params` in the child test class. "
"`image_latents_params` are tested for if passing latents directly are producing same results"
 )

defget_dummy_inputs_by_type(self, device, seed=0, input_image_type="pt", output_type="np"):
inputs=self.get_dummy_inputs(device, seed)

defconvert_to_pt(image):
ifisinstance(image, torch.Tensor):
input_image=image
elifisinstance(image, np.ndarray):
input_image=VaeImageProcessor.numpy_to_pt(image)
elifisinstance(image, PIL.Image.Image):
input_image=VaeImageProcessor.pil_to_numpy(image)
input_image=VaeImageProcessor.numpy_to_pt(input_image)
else:
raiseValueError(f"unsupported input_image_type {type(image)}")
returninput_image

defconvert_pt_to_type(image, input_image_type):
ifinput_image_type=="pt":
input_image=image
elifinput_image_type=="np":
input_image=VaeImageProcessor.pt_to_numpy(image)
elifinput_image_type=="pil":
input_image=VaeImageProcessor.pt_to_numpy(image)
input_image=VaeImageProcessor.numpy_to_pil(input_image)
else:
raiseValueError(f"unsupported input_image_type {input_image_type}.")
returninput_image

forimage_paraminself.image_params:
ifimage_paramininputs.keys():
inputs[image_param] =convert_pt_to_type(
convert_to_pt(inputs[image_param]).to(device), input_image_type
 )

inputs["output_type"] =output_type

returninputs

deftest_pt_np_pil_outputs_equivalent(self, expected_max_diff=1e-4):
self._test_pt_np_pil_outputs_equivalent(expected_max_diff=expected_max_diff)

def_test_pt_np_pil_outputs_equivalent(self, expected_max_diff=1e-4, input_image_type="pt"):
components=self.get_dummy_components()
pipe=self.pipeline_class(**components)
pipe=pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)

output_pt=pipe(
**self.get_dummy_inputs_by_type(torch_device, input_image_type=input_image_type, output_type="pt")
 )[0]
output_np=pipe(
**self.get_dummy_inputs_by_type(torch_device, input_image_type=input_image_type, output_type="np")
 )[0]
output_pil=pipe(
**self.get_dummy_inputs_by_type(torch_device, input_image_type=input_image_type, output_type="pil")
 )[0]

max_diff=np.abs(output_pt.cpu().numpy().transpose(0, 2, 3, 1) -output_np).max()
self.assertLess(
max_diff, expected_max_diff, "`output_type=='pt'` generate different results from `output_type=='np'`"
 )

max_diff=np.abs(np.array(output_pil[0]) - (output_np*255).round()).max()
self.assertLess(max_diff, 2.0, "`output_type=='pil'` generate different results from `output_type=='np'`")

deftest_pt_np_pil_inputs_equivalent(self):
iflen(self.image_params) ==0:
return

components=self.get_dummy_components()
pipe=self.pipeline_class(**components)
pipe=pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)

out_input_pt=pipe(**self.get_dummy_inputs_by_type(torch_device, input_image_type="pt"))[0]
out_input_np=pipe(**self.get_dummy_inputs_by_type(torch_device, input_image_type="np"))[0]
out_input_pil=pipe(**self.get_dummy_inputs_by_type(torch_device, input_image_type="pil"))[0]

max_diff=np.abs(out_input_pt-out_input_np).max()
self.assertLess(max_diff, 1e-4, "`input_type=='pt'` generate different result from `input_type=='np'`")
max_diff=np.abs(out_input_pil-out_input_np).max()
self.assertLess(max_diff, 1e-2, "`input_type=='pt'` generate different result from `input_type=='np'`")

deftest_latents_input(self):
iflen(self.image_latents_params) ==0:
return

components=self.get_dummy_components()
pipe=self.pipeline_class(**components)
pipe.image_processor=VaeImageProcessor(do_resize=False, do_normalize=False)
pipe=pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)

out=pipe(**self.get_dummy_inputs_by_type(torch_device, input_image_type="pt"))[0]

vae=components["vae"]
inputs=self.get_dummy_inputs_by_type(torch_device, input_image_type="pt")
generator=inputs["generator"]
forimage_paraminself.image_latents_params:
ifimage_paramininputs.keys():
inputs[image_param] = (
vae.encode(inputs[image_param]).latent_dist.sample(generator) *vae.config.scaling_factor
 )
out_latents_inputs=pipe(**inputs)[0]

max_diff=np.abs(out-out_latents_inputs).max()
self.assertLess(max_diff, 1e-4, "passing latents as image input generate different result from passing image")

deftest_multi_vae(self):
components=self.get_dummy_components()
pipe=self.pipeline_class(**components)
pipe=pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)

block_out_channels=pipe.vae.config.block_out_channels
norm_num_groups=pipe.vae.config.norm_num_groups

vae_classes= [AutoencoderKL, AsymmetricAutoencoderKL, ConsistencyDecoderVAE, AutoencoderTiny]
configs= [
get_autoencoder_kl_config(block_out_channels, norm_num_groups),
get_asym_autoencoder_kl_config(block_out_channels, norm_num_groups),
get_consistency_vae_config(block_out_channels, norm_num_groups),
get_autoencoder_tiny_config(block_out_channels),
 ]

out_np=pipe(**self.get_dummy_inputs_by_type(torch_device, input_image_type="np"))[0]

forvae_cls, configinzip(vae_classes, configs):
vae=vae_cls(**config)
vae=vae.to(torch_device)
components["vae"] =vae
vae_pipe=self.pipeline_class(**components)
out_vae_np=vae_pipe(**self.get_dummy_inputs_by_type(torch_device, input_image_type="np"))[0]

assertout_vae_np.shape==out_np.shape


@require_torch
classPipelineFromPipeTesterMixin:
@property
deforiginal_pipeline_class(self):
if"xl"inself.pipeline_class.__name__.lower():
original_pipeline_class=StableDiffusionXLPipeline
elif"kolors"inself.pipeline_class.__name__.lower():
original_pipeline_class=KolorsPipeline
else:
original_pipeline_class=StableDiffusionPipeline

returnoriginal_pipeline_class

defget_dummy_inputs_pipe(self, device, seed=0):
inputs=self.get_dummy_inputs(device, seed=seed)
inputs["output_type"] ="np"
inputs["return_dict"] =False
returninputs

defget_dummy_inputs_for_pipe_original(self, device, seed=0):
inputs= {}
fork, vinself.get_dummy_inputs_pipe(device, seed=seed).items():
ifkinset(inspect.signature(self.original_pipeline_class.__call__).parameters.keys()):
inputs[k] =v
returninputs

deftest_from_pipe_consistent_config(self):
ifself.original_pipeline_class==StableDiffusionPipeline:
original_repo="hf-internal-testing/tiny-stable-diffusion-pipe"
original_kwargs= {"requires_safety_checker": False}
elifself.original_pipeline_class==StableDiffusionXLPipeline:
original_repo="hf-internal-testing/tiny-stable-diffusion-xl-pipe"
original_kwargs= {"requires_aesthetics_score": True, "force_zeros_for_empty_prompt": False}
elifself.original_pipeline_class==KolorsPipeline:
original_repo="hf-internal-testing/tiny-kolors-pipe"
original_kwargs= {"force_zeros_for_empty_prompt": False}
else:
raiseValueError(
"original_pipeline_class must be either StableDiffusionPipeline or StableDiffusionXLPipeline"
 )

# create original_pipeline_class(sd/sdxl)
pipe_original=self.original_pipeline_class.from_pretrained(original_repo, **original_kwargs)

# original_pipeline_class(sd/sdxl) -> pipeline_class
pipe_components=self.get_dummy_components()
pipe_additional_components= {}
forname, componentinpipe_components.items():
ifnamenotinpipe_original.components:
pipe_additional_components[name] =component

pipe=self.pipeline_class.from_pipe(pipe_original, **pipe_additional_components)

# pipeline_class -> original_pipeline_class(sd/sdxl)
original_pipe_additional_components= {}
forname, componentinpipe_original.components.items():
ifnamenotinpipe.componentsornotisinstance(component, pipe.components[name].__class__):
original_pipe_additional_components[name] =component

pipe_original_2=self.original_pipeline_class.from_pipe(pipe, **original_pipe_additional_components)

# compare the config
original_config= {k: vfork, vinpipe_original.config.items() ifnotk.startswith("_")}
original_config_2= {k: vfork, vinpipe_original_2.config.items() ifnotk.startswith("_")}
assertoriginal_config_2==original_config

deftest_from_pipe_consistent_forward_pass(self, expected_max_diff=1e-3):
components=self.get_dummy_components()
original_expected_modules, _=self.original_pipeline_class._get_signature_keys(self.original_pipeline_class)

# pipeline components that are also expected to be in the original pipeline
original_pipe_components= {}
# additional components that are not in the pipeline, but expected in the original pipeline
original_pipe_additional_components= {}
# additional components that are in the pipeline, but not expected in the original pipeline
current_pipe_additional_components= {}

forname, componentincomponents.items():
ifnameinoriginal_expected_modules:
original_pipe_components[name] =component
else:
current_pipe_additional_components[name] =component
fornameinoriginal_expected_modules:
ifnamenotinoriginal_pipe_components:
ifnameinself.original_pipeline_class._optional_components:
original_pipe_additional_components[name] =None
else:
raiseValueError(f"missing required module for {self.original_pipeline_class.__class__}: {name}")

pipe_original=self.original_pipeline_class(**original_pipe_components, **original_pipe_additional_components)
forcomponentinpipe_original.components.values():
ifhasattr(component, "set_default_attn_processor"):
component.set_default_attn_processor()
pipe_original.to(torch_device)
pipe_original.set_progress_bar_config(disable=None)
inputs=self.get_dummy_inputs_for_pipe_original(torch_device)
output_original=pipe_original(**inputs)[0]

pipe=self.pipeline_class(**components)
forcomponentinpipe.components.values():
ifhasattr(component, "set_default_attn_processor"):
component.set_default_attn_processor()
pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
inputs=self.get_dummy_inputs_pipe(torch_device)
output=pipe(**inputs)[0]

pipe_from_original=self.pipeline_class.from_pipe(pipe_original, **current_pipe_additional_components)
pipe_from_original.to(torch_device)
pipe_from_original.set_progress_bar_config(disable=None)
inputs=self.get_dummy_inputs_pipe(torch_device)
output_from_original=pipe_from_original(**inputs)[0]

max_diff=np.abs(to_np(output) -to_np(output_from_original)).max()
self.assertLess(
max_diff,
expected_max_diff,
"The outputs of the pipelines created with `from_pipe` and `__init__` are different.",
 )

inputs=self.get_dummy_inputs_for_pipe_original(torch_device)
output_original_2=pipe_original(**inputs)[0]

max_diff=np.abs(to_np(output_original) -to_np(output_original_2)).max()
self.assertLess(max_diff, expected_max_diff, "`from_pipe` should not change the output of original pipeline.")

forcomponentinpipe_original.components.values():
ifhasattr(component, "attn_processors"):
assertall(type(proc) ==AttnProcessorforprocincomponent.attn_processors.values()), (
"`from_pipe` changed the attention processor in original pipeline."
 )

@require_accelerator
@require_accelerate_version_greater("0.14.0")
deftest_from_pipe_consistent_forward_pass_cpu_offload(self, expected_max_diff=1e-3):
components=self.get_dummy_components()
pipe=self.pipeline_class(**components)
forcomponentinpipe.components.values():
ifhasattr(component, "set_default_attn_processor"):
component.set_default_attn_processor()
pipe.enable_model_cpu_offload(device=torch_device)
pipe.set_progress_bar_config(disable=None)
inputs=self.get_dummy_inputs_pipe(torch_device)
output=pipe(**inputs)[0]

original_expected_modules, _=self.original_pipeline_class._get_signature_keys(self.original_pipeline_class)
# pipeline components that are also expected to be in the original pipeline
original_pipe_components= {}
# additional components that are not in the pipeline, but expected in the original pipeline
original_pipe_additional_components= {}
# additional components that are in the pipeline, but not expected in the original pipeline
current_pipe_additional_components= {}
forname, componentincomponents.items():
ifnameinoriginal_expected_modules:
original_pipe_components[name] =component
else:
current_pipe_additional_components[name] =component
fornameinoriginal_expected_modules:
ifnamenotinoriginal_pipe_components:
ifnameinself.original_pipeline_class._optional_components:
original_pipe_additional_components[name] =None
else:
raiseValueError(f"missing required module for {self.original_pipeline_class.__class__}: {name}")

pipe_original=self.original_pipeline_class(**original_pipe_components, **original_pipe_additional_components)
forcomponentinpipe_original.components.values():
ifhasattr(component, "set_default_attn_processor"):
component.set_default_attn_processor()
pipe_original.set_progress_bar_config(disable=None)

pipe_from_original=self.pipeline_class.from_pipe(pipe_original, **current_pipe_additional_components)
forcomponentinpipe_from_original.components.values():
ifhasattr(component, "set_default_attn_processor"):
component.set_default_attn_processor()

pipe_from_original.enable_model_cpu_offload(device=torch_device)
pipe_from_original.set_progress_bar_config(disable=None)
inputs=self.get_dummy_inputs_pipe(torch_device)
output_from_original=pipe_from_original(**inputs)[0]

max_diff=np.abs(to_np(output) -to_np(output_from_original)).max()
self.assertLess(
max_diff,
expected_max_diff,
"The outputs of the pipelines created with `from_pipe` and `__init__` are different.",
 )


@require_torch
classPipelineKarrasSchedulerTesterMixin:
"""
 This mixin is designed to be used with unittest.TestCase classes.
 It provides a set of common tests for each PyTorch pipeline that makes use of KarrasDiffusionSchedulers
 equivalence of dict and tuple outputs, etc.
 """

deftest_karras_schedulers_shape(
self, num_inference_steps_for_strength=4, num_inference_steps_for_strength_for_iterations=5
 ):
components=self.get_dummy_components()
pipe=self.pipeline_class(**components)

# make sure that PNDM does not need warm-up
pipe.scheduler.register_to_config(skip_prk_steps=True)

pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
inputs=self.get_dummy_inputs(torch_device)
inputs["num_inference_steps"] =2

if"strength"ininputs:
inputs["num_inference_steps"] =num_inference_steps_for_strength
inputs["strength"] =0.5

outputs= []
forscheduler_enuminKarrasDiffusionSchedulers:
if"KDPM2"inscheduler_enum.name: