



import os
import tempfile
import unittest

import numpy as np
import torch

from diffusers.models import ModelMixin, UNetaDConditionModel
from diffusers.models.attention_processor import AttnProcessor, LoRAAttnProcessor
from diffusers.utils import (
    floats_tensor,
    logging,
    skip_mps,
    torch_device,
)
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()

A__  :  Tuple       = logging.get_logger(__name__)


def       UpperCamelCase(  __UpperCamelCase							:		Optional[int]					,__UpperCamelCase							:		bool = True			):
		lowerCAmelCase_      :					Tuple     =			{}
		for name in model.attn_processors.keys():
				lowerCAmelCase_      :					Dict     =			name.endswith('''attn2.processor'''			) and not (
				    name.startswith('''transformer_in'''			) or '''temp_attentions''' in name.split('''.'''			)
				)
				lowerCAmelCase_      :					Optional[Any]     =			model.config.cross_attention_dim if has_cross_attention else None
				if name.startswith('''mid_block'''			):
						lowerCAmelCase_      :					int     =			model.config.block_out_channels[-1]
				elif name.startswith('''up_blocks'''			):
						lowerCAmelCase_      :					List[str]     =			int(name[len('''up_blocks.'''			)]			)
						lowerCAmelCase_      :					Optional[Any]     =			list(reversed(model.config.block_out_channels			)			)[block_id]
				elif name.startswith('''down_blocks'''			):
						lowerCAmelCase_      :					str     =			int(name[len('''down_blocks.'''			)]			)
						lowerCAmelCase_      :					Optional[Any]     =			model.config.block_out_channels[block_id]
				elif name.startswith('''transformer_in'''			):
						# Note that the `8 * ...` comes from: https://github.com/huggingface/diffusers/blob/7139f0e874f10b2463caa8cbd585762a309d12d6/src/diffusers/models/unet_3d_condition.py#L148
						lowerCAmelCase_      :					int     =			8 * model.config.attention_head_dim

				lowerCAmelCase_      :					str     =			LoRAAttnProcessor(hidden_size=__UpperCamelCase					,cross_attention_dim=__UpperCamelCase			)
				lowerCAmelCase_      :					List[str]     =			lora_attn_procs[name].to(model.device			)

				if mock_weights:
						# add 1 to weights to mock trained weights
						with torch.no_grad():
								lora_attn_procs[name].to_q_lora.up.weight += 1
								lora_attn_procs[name].to_k_lora.up.weight += 1
								lora_attn_procs[name].to_v_lora.up.weight += 1
								lora_attn_procs[name].to_out_lora.up.weight += 1

		return lora_attn_procs





@skip_mps
class     __snake_case						(					UpperCamelCase_						,UpperCamelCase_						,unittest.TestCase						):
							_a      				=			UNetaDConditionModel
							_a      				=			'''sample'''




							@property
							def        UpperCAmelCase__      (	self					:	Union[str, Any]):
									lowerCAmelCase_      :					Dict     =			4
									lowerCAmelCase_      :					Dict     =			4
									lowerCAmelCase_      :					Union[str, Any]     =			4
									lowerCAmelCase_      :					Union[str, Any]     =			(3_2, 3_2)

									lowerCAmelCase_      :					int     =			floats_tensor((batch_size, num_channels, num_frames) + sizes).to(A_)
									lowerCAmelCase_      :					int     =			torch.tensor([1_0]).to(A_)
									lowerCAmelCase_      :					List[str]     =			floats_tensor((batch_size, 4, 3_2)).to(A_)

									return {"sample": noise, "timestep": time_step, "encoder_hidden_states": encoder_hidden_states}




							@property
							def        UpperCAmelCase__      (	self					:	Dict):
									return (4, 4, 3_2, 3_2)




							@property
							def        UpperCAmelCase__      (	self					:	List[str]):
									return (4, 4, 3_2, 3_2)




							def        UpperCAmelCase__      (	self					:	Optional[Any]):
									lowerCAmelCase_      :					List[str]     =			{
									    '''block_out_channels''': (3_2, 6_4),
									    '''down_block_types''': (
									        '''CrossAttnDownBlock3D''',
									        '''DownBlock3D''',
									    ),
									    '''up_block_types''': ('''UpBlock3D''', '''CrossAttnUpBlock3D'''),
									    '''cross_attention_dim''': 3_2,
									    '''attention_head_dim''': 8,
									    '''out_channels''': 4,
									    '''in_channels''': 4,
									    '''layers_per_block''': 1,
									    '''sample_size''': 3_2,
									}
									lowerCAmelCase_      :					Optional[Any]     =			self.dummy_input
									return init_dict, inputs_dict




							@unittest.skipIf(
							    torch_device != '''cuda''' or not is_xformers_available()    ,    reason='''XFormers attention is only available with CUDA and `xformers` installed'''    ,    )
							def        UpperCAmelCase__      (	self					:	str):
									lowerCAmelCase_ ,  lowerCAmelCase_      :					Union[str, Any]     =			self.prepare_init_args_and_inputs_for_common()
									lowerCAmelCase_      :					Tuple     =			self.model_class(**A_)

									model.enable_xformers_memory_efficient_attention()

									assert (
									    model.mid_block.attentions[0].transformer_blocks[0].attna.processor.__class__.__name__
									    == "XFormersAttnProcessor"
									), "xformers is not enabled"




							def        UpperCAmelCase__      (	self					:	Union[str, Any]):
									lowerCAmelCase_ ,  lowerCAmelCase_      :					Optional[int]     =			self.prepare_init_args_and_inputs_for_common()

									lowerCAmelCase_      :					List[Any]     =			3_2

									lowerCAmelCase_      :					int     =			self.model_class(**A_)
									model.to(A_)
									model.eval()

									with torch.no_grad():
											lowerCAmelCase_      :					List[str]     =			model(**A_)

											if isinstance(A_    ,    A_):
													lowerCAmelCase_      :					Optional[Any]     =			output.sample

									self.assertIsNotNone(A_)
									lowerCAmelCase_      :					int     =			inputs_dict['''sample'''].shape
									self.assertEqual(output.shape    ,    A_    ,    '''Input and output shapes do not match''')




							def        UpperCAmelCase__      (	self					:	int):
									lowerCAmelCase_ ,  lowerCAmelCase_      :					Union[str, Any]     =			self.prepare_init_args_and_inputs_for_common()
									lowerCAmelCase_      :					Union[str, Any]     =			self.model_class(**A_)
									model.to(A_)
									model.eval()

									with torch.no_grad():
											# Warmup pass when using mps (see #372)
											if torch_device == "mps" and isinstance(A_    ,    A_):
													model(**self.dummy_input)

											lowerCAmelCase_      :					List[Any]     =			model(**A_)
											if isinstance(A_    ,    A_):
													lowerCAmelCase_      :					Dict     =			first.sample

											lowerCAmelCase_      :					Tuple     =			model(**A_)
											if isinstance(A_    ,    A_):
													lowerCAmelCase_      :					List[str]     =			second.sample

									lowerCAmelCase_      :					Optional[Any]     =			first.cpu().numpy()
									lowerCAmelCase_      :					int     =			second.cpu().numpy()
									lowerCAmelCase_      :					Any     =			out_a[~np.isnan(A_)]
									lowerCAmelCase_      :					Dict     =			out_a[~np.isnan(A_)]
									lowerCAmelCase_      :					Any     =			np.amax(np.abs(out_a - out_a))
									self.assertLessEqual(A_    ,    1e-5)




							def        UpperCAmelCase__      (	self					:	Union[str, Any]):
									lowerCAmelCase_ ,  lowerCAmelCase_      :					int     =			self.prepare_init_args_and_inputs_for_common()

									lowerCAmelCase_      :					str     =			8

									lowerCAmelCase_      :					List[Any]     =			self.model_class(**A_)
									model.to(A_)
									model.eval()

									model.set_attention_slice('''auto''')
									with torch.no_grad():
											lowerCAmelCase_      :					str     =			model(**A_)
									assert output is not None

									model.set_attention_slice('''max''')
									with torch.no_grad():
											lowerCAmelCase_      :					Dict     =			model(**A_)
									assert output is not None

									model.set_attention_slice(2)
									with torch.no_grad():
											lowerCAmelCase_      :					List[str]     =			model(**A_)
									assert output is not None




							def        UpperCAmelCase__      (	self					:	Tuple):
									lowerCAmelCase_ ,  lowerCAmelCase_      :					List[Any]     =			self.prepare_init_args_and_inputs_for_common()

									lowerCAmelCase_      :					Tuple     =			8

									lowerCAmelCase_      :					Any     =			self.model_class(**A_)
									model.to(A_)

									with torch.no_grad():
											lowerCAmelCase_      :					Tuple     =			model(**A_).sample

									lowerCAmelCase_      :					str     =			create_lora_layers(A_)

									# make sure we can set a list of attention processors
									model.set_attn_processor(A_)
									model.to(A_)

									# test that attn processors can be set to itself
									model.set_attn_processor(model.attn_processors)

									with torch.no_grad():
											lowerCAmelCase_      :					Optional[Any]     =			model(**A_    ,    cross_attention_kwargs={'''scale''': 0.0}).sample
											lowerCAmelCase_      :					str     =			model(**A_    ,    cross_attention_kwargs={'''scale''': 0.5}).sample
											lowerCAmelCase_      :					List[Any]     =			model(**A_    ,    cross_attention_kwargs={'''scale''': 0.5}).sample

									assert (samplea - samplea).abs().max() < 3e-3
									assert (samplea - samplea).abs().max() < 3e-3

									# sample 2 and sample 3 should be different
									assert (samplea - samplea).abs().max() > 3e-3




							def        UpperCAmelCase__      (	self					:	Any):
									lowerCAmelCase_ ,  lowerCAmelCase_      :					Optional[Any]     =			self.prepare_init_args_and_inputs_for_common()

									lowerCAmelCase_      :					str     =			8

									torch.manual_seed(0)
									lowerCAmelCase_      :					Optional[int]     =			self.model_class(**A_)
									model.to(A_)

									with torch.no_grad():
											lowerCAmelCase_      :					List[Any]     =			model(**A_).sample

									lowerCAmelCase_      :					Any     =			create_lora_layers(A_)
									model.set_attn_processor(A_)

									with torch.no_grad():
											lowerCAmelCase_      :					Any     =			model(**A_    ,    cross_attention_kwargs={'''scale''': 0.5}).sample

									with tempfile.TemporaryDirectory() as tmpdirname:
											model.save_attn_procs(A_)
											self.assertTrue(os.path.isfile(os.path.join(A_    ,    '''pytorch_lora_weights.bin''')))
											torch.manual_seed(0)
											lowerCAmelCase_      :					int     =			self.model_class(**A_)
											new_model.to(A_)
											new_model.load_attn_procs(A_)

									with torch.no_grad():
											lowerCAmelCase_      :					List[str]     =			new_model(**A_    ,    cross_attention_kwargs={'''scale''': 0.5}).sample

									assert (sample - new_sample).abs().max() < 1e-3

									# LoRA and no LoRA should NOT be the same
									assert (sample - old_sample).abs().max() > 1e-4




							def        UpperCAmelCase__      (	self					:	Optional[Any]):
									lowerCAmelCase_ ,  lowerCAmelCase_      :					str     =			self.prepare_init_args_and_inputs_for_common()

									lowerCAmelCase_      :					Optional[int]     =			8

									torch.manual_seed(0)
									lowerCAmelCase_      :					Optional[Any]     =			self.model_class(**A_)
									model.to(A_)

									with torch.no_grad():
											lowerCAmelCase_      :					Optional[int]     =			model(**A_).sample

									lowerCAmelCase_      :					Dict     =			create_lora_layers(A_)
									model.set_attn_processor(A_)

									with torch.no_grad():
											lowerCAmelCase_      :					str     =			model(**A_    ,    cross_attention_kwargs={'''scale''': 0.5}).sample

									with tempfile.TemporaryDirectory() as tmpdirname:
											model.save_attn_procs(A_    ,    safe_serialization=A_)
											self.assertTrue(os.path.isfile(os.path.join(A_    ,    '''pytorch_lora_weights.safetensors''')))
											torch.manual_seed(0)
											lowerCAmelCase_      :					List[Any]     =			self.model_class(**A_)
											new_model.to(A_)
											new_model.load_attn_procs(A_)

									with torch.no_grad():
											lowerCAmelCase_      :					List[Any]     =			new_model(**A_    ,    cross_attention_kwargs={'''scale''': 0.5}).sample

									assert (sample - new_sample).abs().max() < 3e-3

									# LoRA and no LoRA should NOT be the same
									assert (sample - old_sample).abs().max() > 1e-4




							def        UpperCAmelCase__      (	self					:	List[str]):
									# enable deterministic behavior for gradient checkpointing
									lowerCAmelCase_ ,  lowerCAmelCase_      :					List[str]     =			self.prepare_init_args_and_inputs_for_common()

									lowerCAmelCase_      :					List[str]     =			8

									torch.manual_seed(0)
									lowerCAmelCase_      :					int     =			self.model_class(**A_)
									model.to(A_)

									lowerCAmelCase_      :					Dict     =			create_lora_layers(A_    ,    mock_weights=A_)
									model.set_attn_processor(A_)
									# Saving as torch, properly reloads with directly filename
									with tempfile.TemporaryDirectory() as tmpdirname:
											model.save_attn_procs(A_)
											self.assertTrue(os.path.isfile(os.path.join(A_    ,    '''pytorch_lora_weights.bin''')))
											torch.manual_seed(0)
											lowerCAmelCase_      :					Union[str, Any]     =			self.model_class(**A_)
											new_model.to(A_)
											new_model.load_attn_procs(A_    ,    weight_name='''pytorch_lora_weights.bin''')




							def        UpperCAmelCase__      (	self					:	Union[str, Any]):
									lowerCAmelCase_ ,  lowerCAmelCase_      :					Tuple     =			self.prepare_init_args_and_inputs_for_common()

									lowerCAmelCase_      :					Any     =			8

									torch.manual_seed(0)
									lowerCAmelCase_      :					Any     =			self.model_class(**A_)
									model.to(A_)

									lowerCAmelCase_      :					Dict     =			create_lora_layers(A_    ,    mock_weights=A_)
									model.set_attn_processor(A_)
									# Saving as torch, properly reloads with directly filename
									with tempfile.TemporaryDirectory() as tmpdirname:
											model.save_attn_procs(A_)
											self.assertTrue(os.path.isfile(os.path.join(A_    ,    '''pytorch_lora_weights.bin''')))
											torch.manual_seed(0)
											lowerCAmelCase_      :					Any     =			self.model_class(**A_)
											new_model.to(A_)
											with self.assertRaises(A_) as e:
													new_model.load_attn_procs(A_    ,    use_safetensors=A_)
											self.assertIn('''Error no file named pytorch_lora_weights.safetensors'''    ,    str(e.exception))




							def        UpperCAmelCase__      (	self					:	List[str]):
									lowerCAmelCase_ ,  lowerCAmelCase_      :					int     =			self.prepare_init_args_and_inputs_for_common()

									lowerCAmelCase_      :					Dict     =			8

									torch.manual_seed(0)
									lowerCAmelCase_      :					List[Any]     =			self.model_class(**A_)
									model.to(A_)

									with torch.no_grad():
											lowerCAmelCase_      :					List[str]     =			model(**A_).sample

									lowerCAmelCase_      :					Optional[int]     =			create_lora_layers(A_)
									model.set_attn_processor(A_)

									with torch.no_grad():
											lowerCAmelCase_      :					int     =			model(**A_    ,    cross_attention_kwargs={'''scale''': 0.0}).sample

									model.set_attn_processor(AttnProcessor())

									with torch.no_grad():
											lowerCAmelCase_      :					Any     =			model(**A_).sample

									assert (sample - new_sample).abs().max() < 1e-4
									assert (sample - old_sample).abs().max() < 3e-3




							@unittest.skipIf(
							    torch_device != '''cuda''' or not is_xformers_available()    ,    reason='''XFormers attention is only available with CUDA and `xformers` installed'''    ,    )
							def        UpperCAmelCase__      (	self					:	Optional[Any]):
									# enable deterministic behavior for gradient checkpointing
									lowerCAmelCase_ ,  lowerCAmelCase_      :					List[Any]     =			self.prepare_init_args_and_inputs_for_common()

									lowerCAmelCase_      :					Dict     =			4

									torch.manual_seed(0)
									lowerCAmelCase_      :					Tuple     =			self.model_class(**A_)
									model.to(A_)
									lowerCAmelCase_      :					List[Any]     =			create_lora_layers(A_)
									model.set_attn_processor(A_)

									# default
									with torch.no_grad():
											lowerCAmelCase_      :					Tuple     =			model(**A_).sample

											model.enable_xformers_memory_efficient_attention()
											lowerCAmelCase_      :					Union[str, Any]     =			model(**A_).sample

											model.disable_xformers_memory_efficient_attention()
											lowerCAmelCase_      :					Optional[int]     =			model(**A_).sample

									assert (sample - on_sample).abs().max() < 1e-4
									assert (sample - off_sample).abs().max() < 1e-4







							def        UpperCAmelCase__      (	self					:	Dict):
									lowerCAmelCase_ ,  lowerCAmelCase_      :					Optional[Any]     =			self.prepare_init_args_and_inputs_for_common()
									lowerCAmelCase_      :					Any     =			3_2

									lowerCAmelCase_      :					Optional[int]     =			self.model_class(**A_)
									model.to(A_)
									model.eval()

									with torch.no_grad():
											lowerCAmelCase_      :					Any     =			model(**A_)[0]

									model.enable_forward_chunking()
									with torch.no_grad():
											lowerCAmelCase_      :					List[str]     =			model(**A_)[0]

									self.assertEqual(output.shape    ,    output_a.shape    ,    '''Shape doesn\'t match''')
									assert np.abs(output.cpu() - output_a.cpu()).max() < 1e-2


# (todo: sayakpaul) implement SLOW tests.










import random
import unittest

import torch

from diffusers import IFImgaImgPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin



@skip_mps
class     __snake_case						(					UpperCamelCase_						,UpperCamelCase_						,unittest.TestCase						):
							_a      				=			IFImgaImgPipeline
							_a      				=			TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''width''', '''height'''}
							_a      				=			TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
							_a      				=			PipelineTesterMixin.required_optional_params - {'''latents'''}




							def        UpperCAmelCase__      (	self					:	Any):
									return self._get_dummy_components()




							def        UpperCAmelCase__      (	self					:	str    ,    A_					:	Optional[int]    ,    A_					:	List[Any]=0):
									if str(A_).startswith('''mps'''):
											lowerCAmelCase_      :					Dict     =			torch.manual_seed(A_)
									else:
											lowerCAmelCase_      :					List[Any]     =			torch.Generator(device=A_).manual_seed(A_)

									lowerCAmelCase_      :					List[Any]     =			floats_tensor((1, 3, 3_2, 3_2)    ,    rng=random.Random(A_)).to(A_)

									lowerCAmelCase_      :					Any     =			{
									    '''prompt''': '''A painting of a squirrel eating a burger''',
									    '''image''': image,
									    '''generator''': generator,
									    '''num_inference_steps''': 2,
									    '''output_type''': '''numpy''',
									}

									return inputs




							def        UpperCAmelCase__      (	self					:	Union[str, Any]):
									self._test_save_load_optional_components()




							@unittest.skipIf(
							    torch_device != '''cuda''' or not is_xformers_available()    ,    reason='''XFormers attention is only available with CUDA and `xformers` installed'''    ,    )
							def        UpperCAmelCase__      (	self					:	List[str]):
									self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3)




							@unittest.skipIf(torch_device != '''cuda'''    ,    reason='''float16 requires CUDA''')
							def        UpperCAmelCase__      (	self					:	Dict):
									# Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
									super().test_save_load_floataa(expected_max_diff=1e-1)




							@unittest.skipIf(torch_device != '''cuda'''    ,    reason='''float16 requires CUDA''')
							def        UpperCAmelCase__      (	self					:	Tuple):
									super().test_floataa_inference(expected_max_diff=1e-1)




							def        UpperCAmelCase__      (	self					:	str):
									self._test_attention_slicing_forward_pass(expected_max_diff=1e-2)




							def        UpperCAmelCase__      (	self					:	int):
									self._test_save_load_local()







							def        UpperCAmelCase__      (	self					:	Union[str, Any]):
									self._test_inference_batch_single_identical(
									    expected_max_diff=1e-2    ,    )





