



'''simple docstring'''




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()

UpperCamelCase__	:     Dict    =			logging.get_logger(__name__)


def     lowerCAmelCase_       (     _lowerCamelCase:       Optional[Any]	,				_lowerCamelCase:       bool = True ):
		__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		{}
		for name in model.attn_processors.keys():
				__SCREAMING_SNAKE_CASE	:		Dict   	=		name.endswith("""attn2.processor""" ) and not (
				    name.startswith("""transformer_in""" ) or """temp_attentions""" in name.split(""".""" )
				)
				__SCREAMING_SNAKE_CASE	:		Any   	=		model.config.cross_attention_dim if has_cross_attention else None
				if name.startswith("""mid_block""" ):
						__SCREAMING_SNAKE_CASE	:		Dict   	=		model.config.block_out_channels[-1]
				elif name.startswith("""up_blocks""" ):
						__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		int(name[len("""up_blocks.""" )] )
						__SCREAMING_SNAKE_CASE	:		str   	=		list(reversed(model.config.block_out_channels ) )[block_id]
				elif name.startswith("""down_blocks""" ):
						__SCREAMING_SNAKE_CASE	:		Tuple   	=		int(name[len("""down_blocks.""" )] )
						__SCREAMING_SNAKE_CASE	:		int   	=		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
						__SCREAMING_SNAKE_CASE	:		List[str]   	=		8 * model.config.attention_head_dim

				__SCREAMING_SNAKE_CASE	:		Dict   	=		LoRAAttnProcessor(hidden_size=_lowerCamelCase	,				cross_attention_dim=_lowerCamelCase )
				__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		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 						_UpperCamelCase  (							lowerCamelCase__   , lowerCamelCase__   , unittest.TestCase   ):


			'''simple docstring'''

			_A				:		Optional[Any]         =		UNetaDConditionModel
			_A				:		Union[str, Any]         =		'''sample'''

			@property
			def UpperCamelCase__ (							self					:			Any						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		int   	=		4
					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		4
					__SCREAMING_SNAKE_CASE	:		str   	=		4
					__SCREAMING_SNAKE_CASE	:		Tuple   	=		(3_2, 3_2)

					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		floats_tensor((batch_size, num_channels, num_frames) + sizes						).to(lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		torch.tensor([1_0]						).to(lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		floats_tensor((batch_size, 4, 3_2)						).to(lowerCAmelCase__						)

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

			@property
			def UpperCamelCase__ (							self					:			Optional[Any]						):





					"""simple docstring"""






					return (4, 4, 3_2, 3_2)

			@property
			def UpperCamelCase__ (							self					:			Any						):





					"""simple docstring"""






					return (4, 4, 3_2, 3_2)

			def UpperCamelCase__ (							self					:			Union[str, Any]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		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,
					}
					__SCREAMING_SNAKE_CASE	:		str   	=		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					:			Optional[Any]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	,						__SCREAMING_SNAKE_CASE	:		Tuple   	=		self.prepare_init_args_and_inputs_for_common()
					__SCREAMING_SNAKE_CASE	:		int   	=		self.model_class(**lowerCAmelCase__						)

					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					:			Optional[int]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	,						__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		self.prepare_init_args_and_inputs_for_common()

					__SCREAMING_SNAKE_CASE	:		Dict   	=		3_2

					__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		self.model_class(**lowerCAmelCase__						)
					model.to(lowerCAmelCase__						)
					model.eval()

					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		model(**lowerCAmelCase__						)

							if isinstance(lowerCAmelCase__							,					lowerCAmelCase__						):
									__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		output.sample

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

			def UpperCamelCase__ (							self					:			int						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	,						__SCREAMING_SNAKE_CASE	:		Tuple   	=		self.prepare_init_args_and_inputs_for_common()
					__SCREAMING_SNAKE_CASE	:		List[str]   	=		self.model_class(**lowerCAmelCase__						)
					model.to(lowerCAmelCase__						)
					model.eval()

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

							__SCREAMING_SNAKE_CASE	:		str   	=		model(**lowerCAmelCase__						)
							if isinstance(lowerCAmelCase__							,					lowerCAmelCase__						):
									__SCREAMING_SNAKE_CASE	:		Any   	=		first.sample

							__SCREAMING_SNAKE_CASE	:		List[str]   	=		model(**lowerCAmelCase__						)
							if isinstance(lowerCAmelCase__							,					lowerCAmelCase__						):
									__SCREAMING_SNAKE_CASE	:		List[str]   	=		second.sample

					__SCREAMING_SNAKE_CASE	:		Dict   	=		first.cpu().numpy()
					__SCREAMING_SNAKE_CASE	:		List[str]   	=		second.cpu().numpy()
					__SCREAMING_SNAKE_CASE	:		List[str]   	=		out_a[~np.isnan(lowerCAmelCase__						)]
					__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		out_a[~np.isnan(lowerCAmelCase__						)]
					__SCREAMING_SNAKE_CASE	:		Tuple   	=		np.amax(np.abs(out_a - out_a						)						)
					self.assertLessEqual(lowerCAmelCase__							,					1E-5						)

			def UpperCamelCase__ (							self					:			int						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	,						__SCREAMING_SNAKE_CASE	:		List[Any]   	=		self.prepare_init_args_and_inputs_for_common()

					__SCREAMING_SNAKE_CASE	:		int   	=		8

					__SCREAMING_SNAKE_CASE	:		int   	=		self.model_class(**lowerCAmelCase__						)
					model.to(lowerCAmelCase__						)
					model.eval()

					model.set_attention_slice("""auto"""						)
					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		List[str]   	=		model(**lowerCAmelCase__						)
					assert output is not None

					model.set_attention_slice("""max"""						)
					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		Any   	=		model(**lowerCAmelCase__						)
					assert output is not None

					model.set_attention_slice(2						)
					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		int   	=		model(**lowerCAmelCase__						)
					assert output is not None

			def UpperCamelCase__ (							self					:			Dict						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	,						__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		self.prepare_init_args_and_inputs_for_common()

					__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		8

					__SCREAMING_SNAKE_CASE	:		Tuple   	=		self.model_class(**lowerCAmelCase__						)
					model.to(lowerCAmelCase__						)

					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		List[str]   	=		model(**lowerCAmelCase__						).sample

					__SCREAMING_SNAKE_CASE	:		Tuple   	=		create_lora_layers(lowerCAmelCase__						)

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

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

					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		model(**lowerCAmelCase__							,					cross_attention_kwargs={"""scale""": 0.0}						).sample
							__SCREAMING_SNAKE_CASE	:		Dict   	=		model(**lowerCAmelCase__							,					cross_attention_kwargs={"""scale""": 0.5}						).sample
							__SCREAMING_SNAKE_CASE	:		int   	=		model(**lowerCAmelCase__							,					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					:			List[str]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	,						__SCREAMING_SNAKE_CASE	:		int   	=		self.prepare_init_args_and_inputs_for_common()

					__SCREAMING_SNAKE_CASE	:		Tuple   	=		8

					torch.manual_seed(0						)
					__SCREAMING_SNAKE_CASE	:		str   	=		self.model_class(**lowerCAmelCase__						)
					model.to(lowerCAmelCase__						)

					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		Tuple   	=		model(**lowerCAmelCase__						).sample

					__SCREAMING_SNAKE_CASE	:		str   	=		create_lora_layers(lowerCAmelCase__						)
					model.set_attn_processor(lowerCAmelCase__						)

					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		model(**lowerCAmelCase__							,					cross_attention_kwargs={"""scale""": 0.5}						).sample

					with tempfile.TemporaryDirectory() as tmpdirname:
							model.save_attn_procs(lowerCAmelCase__						)
							self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__							,					"""pytorch_lora_weights.bin"""						)						)						)
							torch.manual_seed(0						)
							__SCREAMING_SNAKE_CASE	:		List[Any]   	=		self.model_class(**lowerCAmelCase__						)
							new_model.to(lowerCAmelCase__						)
							new_model.load_attn_procs(lowerCAmelCase__						)

					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		new_model(**lowerCAmelCase__							,					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					:			Any						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	,						__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		self.prepare_init_args_and_inputs_for_common()

					__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		8

					torch.manual_seed(0						)
					__SCREAMING_SNAKE_CASE	:		List[str]   	=		self.model_class(**lowerCAmelCase__						)
					model.to(lowerCAmelCase__						)

					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		Any   	=		model(**lowerCAmelCase__						).sample

					__SCREAMING_SNAKE_CASE	:		List[Any]   	=		create_lora_layers(lowerCAmelCase__						)
					model.set_attn_processor(lowerCAmelCase__						)

					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		Any   	=		model(**lowerCAmelCase__							,					cross_attention_kwargs={"""scale""": 0.5}						).sample

					with tempfile.TemporaryDirectory() as tmpdirname:
							model.save_attn_procs(lowerCAmelCase__							,					safe_serialization=lowerCAmelCase__						)
							self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__							,					"""pytorch_lora_weights.safetensors"""						)						)						)
							torch.manual_seed(0						)
							__SCREAMING_SNAKE_CASE	:		str   	=		self.model_class(**lowerCAmelCase__						)
							new_model.to(lowerCAmelCase__						)
							new_model.load_attn_procs(lowerCAmelCase__						)

					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		new_model(**lowerCAmelCase__							,					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					:			int						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	,						__SCREAMING_SNAKE_CASE	:		List[Any]   	=		self.prepare_init_args_and_inputs_for_common()

					__SCREAMING_SNAKE_CASE	:		Dict   	=		8

					torch.manual_seed(0						)
					__SCREAMING_SNAKE_CASE	:		Dict   	=		self.model_class(**lowerCAmelCase__						)
					model.to(lowerCAmelCase__						)

					__SCREAMING_SNAKE_CASE	:		int   	=		create_lora_layers(lowerCAmelCase__							,					mock_weights=lowerCAmelCase__						)
					model.set_attn_processor(lowerCAmelCase__						)
					# Saving as torch, properly reloads with directly filename
					with tempfile.TemporaryDirectory() as tmpdirname:
							model.save_attn_procs(lowerCAmelCase__						)
							self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__							,					"""pytorch_lora_weights.bin"""						)						)						)
							torch.manual_seed(0						)
							__SCREAMING_SNAKE_CASE	:		List[Any]   	=		self.model_class(**lowerCAmelCase__						)
							new_model.to(lowerCAmelCase__						)
							new_model.load_attn_procs(lowerCAmelCase__							,					weight_name="""pytorch_lora_weights.bin"""						)

			def UpperCamelCase__ (							self					:			Optional[int]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	,						__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		self.prepare_init_args_and_inputs_for_common()

					__SCREAMING_SNAKE_CASE	:		str   	=		8

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

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

			def UpperCamelCase__ (							self					:			Dict						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	,						__SCREAMING_SNAKE_CASE	:		str   	=		self.prepare_init_args_and_inputs_for_common()

					__SCREAMING_SNAKE_CASE	:		str   	=		8

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

					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		model(**lowerCAmelCase__						).sample

					__SCREAMING_SNAKE_CASE	:		Dict   	=		create_lora_layers(lowerCAmelCase__						)
					model.set_attn_processor(lowerCAmelCase__						)

					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		int   	=		model(**lowerCAmelCase__							,					cross_attention_kwargs={"""scale""": 0.0}						).sample

					model.set_attn_processor(AttnProcessor()						)

					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		Tuple   	=		model(**lowerCAmelCase__						).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					:			List[Any]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	,						__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		self.prepare_init_args_and_inputs_for_common()

					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		4

					torch.manual_seed(0						)
					__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		self.model_class(**lowerCAmelCase__						)
					model.to(lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		create_lora_layers(lowerCAmelCase__						)
					model.set_attn_processor(lowerCAmelCase__						)

					# default
					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		model(**lowerCAmelCase__						).sample

							model.enable_xformers_memory_efficient_attention()
							__SCREAMING_SNAKE_CASE	:		Dict   	=		model(**lowerCAmelCase__						).sample

							model.disable_xformers_memory_efficient_attention()
							__SCREAMING_SNAKE_CASE	:		List[str]   	=		model(**lowerCAmelCase__						).sample

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



			def UpperCamelCase__ (							self					:			Dict						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	,						__SCREAMING_SNAKE_CASE	:		Dict   	=		self.prepare_init_args_and_inputs_for_common()
					__SCREAMING_SNAKE_CASE	:		List[Any]   	=		3_2

					__SCREAMING_SNAKE_CASE	:		List[Any]   	=		self.model_class(**lowerCAmelCase__						)
					model.to(lowerCAmelCase__						)
					model.eval()

					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		Any   	=		model(**lowerCAmelCase__						)[0]

					model.enable_forward_chunking()
					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		List[str]   	=		model(**lowerCAmelCase__						)[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.




'''simple docstring'''




UpperCamelCase__	:     List[Any]    =			10_00


def     lowerCAmelCase_       (     _lowerCamelCase:       float ):
		return x * x


def     lowerCAmelCase_       (     _lowerCamelCase:       Any	,				_lowerCamelCase:       float	,				_lowerCamelCase:       float	,				_lowerCamelCase:       int = 4 ):
		assert callable(
		    _lowerCamelCase ), F"the function(object) passed should be callable your input : {function}"
		assert isinstance(_lowerCamelCase	,				(float, int) ), F"a should be float or integer your input : {a}"
		assert isinstance(function(_lowerCamelCase )	,				(float, int) ), (
		    "the function should return integer or float return type of your function, "
		    F"{type(_lowerCamelCase )}"
		)
		assert isinstance(_lowerCamelCase	,				(float, int) ), F"b should be float or integer your input : {b}"
		assert (
		    isinstance(_lowerCamelCase	,				_lowerCamelCase ) and precision > 0
		), F"precision should be positive integer your input : {precision}"

		# just applying the formula of simpson for approximate integration written in
		# mentioned article in first comment of this file and above this function

		__SCREAMING_SNAKE_CASE	:		List[Any]   	=		(b - a) / N_STEPS
		__SCREAMING_SNAKE_CASE	:		List[Any]   	=		function(_lowerCamelCase ) + function(_lowerCamelCase )

		for i in range(1	,				_lowerCamelCase ):
				__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		a + h * i
				result += function(_lowerCamelCase ) * (4 if i % 2 else 2)

		result *= h / 3
		return round(_lowerCamelCase	,				_lowerCamelCase )


if __name__ == "__main__":
						import doctest

						doctest.testmod()