


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

__SCREAMING_SNAKE_CASE				:     Any		=     logging.get_logger(__name__)







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

														snake_case_     							=		LoRAAttnProcessor(hidden_size=_snake_case					,			cross_attention_dim=_snake_case					)
														snake_case_     							=		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   __A  (__snake_case    ,		__snake_case    ,		unittest.TestCase):





							'''simple docstring'''

							__lowercase:		Tuple				   =     UNetaDConditionModel
							__lowercase:		Optional[int]				   =     """sample"""



							@property
							def 		lowerCAmelCase							(					self				:   List[str]							)       ->Dict:



														"""simple docstring"""



														snake_case_     							=		4
														snake_case_     							=		4
														snake_case_     							=		4
														snake_case_     							=		(32, 32)

														snake_case_     							=		floats_tensor((batch_size, num_channels, num_frames) + sizes							).to(a_							)
														snake_case_     							=		torch.tensor([10]							).to(a_							)
														snake_case_     							=		floats_tensor((batch_size, 4, 32)							).to(a_							)

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



							@property
							def 		lowerCAmelCase							(					self				:   List[str]							)       ->Tuple:



														"""simple docstring"""



														return (4, 4, 32, 32)



							@property
							def 		lowerCAmelCase							(					self				:   List[str]							)       ->Union[str, Any]:



														"""simple docstring"""



														return (4, 4, 32, 32)



							def 		lowerCAmelCase							(					self				:   int							)       ->Optional[int]:



														"""simple docstring"""



														snake_case_     							=		{
														    '''block_out_channels''': (32, 64),
														    '''down_block_types''': (
														        '''CrossAttnDownBlock3D''',
														        '''DownBlock3D''',
														    ),
														    '''up_block_types''': ('''UpBlock3D''', '''CrossAttnUpBlock3D'''),
														    '''cross_attention_dim''': 32,
														    '''attention_head_dim''': 8,
														    '''out_channels''': 4,
														    '''in_channels''': 4,
														    '''layers_per_block''': 1,
														    '''sample_size''': 32,
														}
														snake_case_     							=		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 		lowerCAmelCase							(					self				:   Union[str, Any]							)       ->List[Any]:



														"""simple docstring"""



														snake_case_     							=		self.prepare_init_args_and_inputs_for_common()
														snake_case_     							=		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 		lowerCAmelCase							(					self				:   Union[str, Any]							)       ->Dict:



														"""simple docstring"""



														snake_case_     							=		self.prepare_init_args_and_inputs_for_common()

														snake_case_     							=		32

														snake_case_     							=		self.model_class(**a_							)
														model.to(a_							)
														model.eval()

														with torch.no_grad():
																					snake_case_     							=		model(**a_							)

																					if isinstance(a_		,							a_							):
																												snake_case_     							=		output.sample

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



							def 		lowerCAmelCase							(					self				:   List[str]							)       ->str:



														"""simple docstring"""



														snake_case_     							=		self.prepare_init_args_and_inputs_for_common()
														snake_case_     							=		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							)

																					snake_case_     							=		model(**a_							)
																					if isinstance(a_		,							a_							):
																												snake_case_     							=		first.sample

																					snake_case_     							=		model(**a_							)
																					if isinstance(a_		,							a_							):
																												snake_case_     							=		second.sample

														snake_case_     							=		first.cpu().numpy()
														snake_case_     							=		second.cpu().numpy()
														snake_case_     							=		out_a[~np.isnan(a_							)]
														snake_case_     							=		out_a[~np.isnan(a_							)]
														snake_case_     							=		np.amax(np.abs(out_a - out_a							)							)
														self.assertLessEqual(a_		,							1E-5							)



							def 		lowerCAmelCase							(					self				:   int							)       ->Any:



														"""simple docstring"""



														snake_case_     							=		self.prepare_init_args_and_inputs_for_common()

														snake_case_     							=		8

														snake_case_     							=		self.model_class(**a_							)
														model.to(a_							)
														model.eval()

														model.set_attention_slice("""auto"""							)
														with torch.no_grad():
																					snake_case_     							=		model(**a_							)
														assert output is not None

														model.set_attention_slice("""max"""							)
														with torch.no_grad():
																					snake_case_     							=		model(**a_							)
														assert output is not None

														model.set_attention_slice(2							)
														with torch.no_grad():
																					snake_case_     							=		model(**a_							)
														assert output is not None



							def 		lowerCAmelCase							(					self				:   Dict							)       ->Dict:



														"""simple docstring"""



														snake_case_     							=		self.prepare_init_args_and_inputs_for_common()

														snake_case_     							=		8

														snake_case_     							=		self.model_class(**a_							)
														model.to(a_							)

														with torch.no_grad():
																					snake_case_     							=		model(**a_							).sample

														snake_case_     							=		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():
																					snake_case_     							=		model(**a_		,							cross_attention_kwargs={"""scale""": 0.0}							).sample
																					snake_case_     							=		model(**a_		,							cross_attention_kwargs={"""scale""": 0.5}							).sample
																					snake_case_     							=		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 		lowerCAmelCase							(					self				:   int							)       ->List[str]:



														"""simple docstring"""



														snake_case_     							=		self.prepare_init_args_and_inputs_for_common()

														snake_case_     							=		8

														torch.manual_seed(0							)
														snake_case_     							=		self.model_class(**a_							)
														model.to(a_							)

														with torch.no_grad():
																					snake_case_     							=		model(**a_							).sample

														snake_case_     							=		create_lora_layers(a_							)
														model.set_attn_processor(a_							)

														with torch.no_grad():
																					snake_case_     							=		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							)
																					snake_case_     							=		self.model_class(**a_							)
																					new_model.to(a_							)
																					new_model.load_attn_procs(a_							)

														with torch.no_grad():
																					snake_case_     							=		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 		lowerCAmelCase							(					self				:   Any							)       ->Dict:



														"""simple docstring"""



														snake_case_     							=		self.prepare_init_args_and_inputs_for_common()

														snake_case_     							=		8

														torch.manual_seed(0							)
														snake_case_     							=		self.model_class(**a_							)
														model.to(a_							)

														with torch.no_grad():
																					snake_case_     							=		model(**a_							).sample

														snake_case_     							=		create_lora_layers(a_							)
														model.set_attn_processor(a_							)

														with torch.no_grad():
																					snake_case_     							=		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							)
																					snake_case_     							=		self.model_class(**a_							)
																					new_model.to(a_							)
																					new_model.load_attn_procs(a_							)

														with torch.no_grad():
																					snake_case_     							=		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 		lowerCAmelCase							(					self				:   List[str]							)       ->Tuple:



														"""simple docstring"""



														snake_case_     							=		self.prepare_init_args_and_inputs_for_common()

														snake_case_     							=		8

														torch.manual_seed(0							)
														snake_case_     							=		self.model_class(**a_							)
														model.to(a_							)

														snake_case_     							=		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							)
																					snake_case_     							=		self.model_class(**a_							)
																					new_model.to(a_							)
																					new_model.load_attn_procs(a_		,							weight_name="""pytorch_lora_weights.bin"""							)



							def 		lowerCAmelCase							(					self				:   int							)       ->int:



														"""simple docstring"""



														snake_case_     							=		self.prepare_init_args_and_inputs_for_common()

														snake_case_     							=		8

														torch.manual_seed(0							)
														snake_case_     							=		self.model_class(**a_							)
														model.to(a_							)

														snake_case_     							=		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							)
																					snake_case_     							=		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 		lowerCAmelCase							(					self				:   Dict							)       ->Tuple:



														"""simple docstring"""



														snake_case_     							=		self.prepare_init_args_and_inputs_for_common()

														snake_case_     							=		8

														torch.manual_seed(0							)
														snake_case_     							=		self.model_class(**a_							)
														model.to(a_							)

														with torch.no_grad():
																					snake_case_     							=		model(**a_							).sample

														snake_case_     							=		create_lora_layers(a_							)
														model.set_attn_processor(a_							)

														with torch.no_grad():
																					snake_case_     							=		model(**a_		,							cross_attention_kwargs={"""scale""": 0.0}							).sample

														model.set_attn_processor(AttnProcessor()							)

														with torch.no_grad():
																					snake_case_     							=		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 		lowerCAmelCase							(					self				:   List[Any]							)       ->List[Any]:



														"""simple docstring"""



														snake_case_     							=		self.prepare_init_args_and_inputs_for_common()

														snake_case_     							=		4

														torch.manual_seed(0							)
														snake_case_     							=		self.model_class(**a_							)
														model.to(a_							)
														snake_case_     							=		create_lora_layers(a_							)
														model.set_attn_processor(a_							)

														# default
														with torch.no_grad():
																					snake_case_     							=		model(**a_							).sample

																					model.enable_xformers_memory_efficient_attention()
																					snake_case_     							=		model(**a_							).sample

																					model.disable_xformers_memory_efficient_attention()
																					snake_case_     							=		model(**a_							).sample

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



							def 		lowerCAmelCase							(					self				:   Tuple							)       ->List[Any]:



														"""simple docstring"""



														snake_case_     							=		self.prepare_init_args_and_inputs_for_common()
														snake_case_     							=		32

														snake_case_     							=		self.model_class(**a_							)
														model.to(a_							)
														model.eval()

														with torch.no_grad():
																					snake_case_     							=		model(**a_							)[0]

														model.enable_forward_chunking()
														with torch.no_grad():
																					snake_case_     							=		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.







"""simple docstring"""
import unittest

from diffusers import __version__
from diffusers.utils import deprecate



class      _UpperCAmelCase					(				unittest.TestCase							):






		'''simple docstring'''

		lowerCamelCase__           ='.'.join([str(int(__version__.split('.'							)[0]							) + 1							)] + __version__.split('.'							)[1:]							)
		lowerCamelCase__           ='0.0.1'



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			__snake_case      :		int         =							{'''deprecated_arg''': 4}

			with self.assertWarns(a_   ) as warning:
				__snake_case      :		Union[str, Any]         =							deprecate('''deprecated_arg'''							,			self.higher_version							,			'''message'''							,			take_from=a_   )

			assert output == 4
			assert (
			    str(warning.warning   )
			    == f"""The `deprecated_arg` argument is deprecated and will be removed in version {self.higher_version}."""
			    " message"
			)



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			__snake_case      :		Tuple         =							{'''deprecated_arg''': 4}

			with self.assertWarns(a_   ) as warning:
				__snake_case      :		Dict         =							deprecate(('''deprecated_arg''', self.higher_version, '''message''')							,			take_from=a_   )

			assert output == 4
			assert (
			    str(warning.warning   )
			    == f"""The `deprecated_arg` argument is deprecated and will be removed in version {self.higher_version}."""
			    " message"
			)



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			__snake_case      :		Union[str, Any]         =							{'''deprecated_arg_1''': 4, '''deprecated_arg_2''': 8}
			with self.assertWarns(a_   ) as warning:
				__snake_case      ,				__snake_case      :		List[Any]         =							deprecate(
				    ('''deprecated_arg_1''', self.higher_version, '''Hey''')							,			('''deprecated_arg_2''', self.higher_version, '''Hey''')							,			take_from=a_							,			)
			assert output_a == 4
			assert output_a == 8
			assert (
			    str(warning.warnings[0].message   )
			    == "The `deprecated_arg_1` argument is deprecated and will be removed in version"
			    f""" {self.higher_version}. Hey"""
			)
			assert (
			    str(warning.warnings[1].message   )
			    == "The `deprecated_arg_2` argument is deprecated and will be removed in version"
			    f""" {self.higher_version}. Hey"""
			)



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			__snake_case      :		int         =							{'''deprecated_arg''': 4}

			with self.assertRaises(a_   ) as error:
				deprecate(('''wrong_arg''', self.higher_version, '''message''')							,			take_from=a_   )

			assert "test_deprecate_function_incorrect_arg in" in str(error.exception   )
			assert "line" in str(error.exception   )
			assert "got an unexpected keyword argument `deprecated_arg`" in str(error.exception   )



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			with self.assertWarns(a_   ) as warning:
				deprecate(('''deprecated_arg''', self.higher_version, '''message''')   )

			assert (
			    str(warning.warning   )
			    == f"""`deprecated_arg` is deprecated and will be removed in version {self.higher_version}. message"""
			)



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			with self.assertWarns(a_   ) as warning:
				deprecate(
				    ('''deprecated_arg_1''', self.higher_version, '''Hey''')							,			('''deprecated_arg_2''', self.higher_version, '''Hey''')							,			)
			assert (
			    str(warning.warnings[0].message   )
			    == f"""`deprecated_arg_1` is deprecated and will be removed in version {self.higher_version}. Hey"""
			)
			assert (
			    str(warning.warnings[1].message   )
			    == f"""`deprecated_arg_2` is deprecated and will be removed in version {self.higher_version}. Hey"""
			)



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			class      _UpperCAmelCase					:






					'''simple docstring'''

					lowerCamelCase__           =5

			with self.assertWarns(a_   ) as warning:
				__snake_case      :		List[Any]         =							deprecate(('''arg''', self.higher_version, '''message''')							,			take_from=Args()   )

			assert arg == 5
			assert (
			    str(warning.warning   )
			    == f"""The `arg` attribute is deprecated and will be removed in version {self.higher_version}. message"""
			)



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			class      _UpperCAmelCase					:






					'''simple docstring'''

					lowerCamelCase__           =5
					lowerCamelCase__           =7

			with self.assertWarns(a_   ) as warning:
				__snake_case      ,				__snake_case      :		List[str]         =							deprecate(
				    ('''arg''', self.higher_version, '''message''')							,			('''foo''', self.higher_version, '''message''')							,			('''does not exist''', self.higher_version, '''message''')							,			take_from=Args()							,			)

			assert arg_a == 5
			assert arg_a == 7
			assert (
			    str(warning.warning   )
			    == f"""The `arg` attribute is deprecated and will be removed in version {self.higher_version}. message"""
			)
			assert (
			    str(warning.warnings[0].message   )
			    == f"""The `arg` attribute is deprecated and will be removed in version {self.higher_version}. message"""
			)
			assert (
			    str(warning.warnings[1].message   )
			    == f"""The `foo` attribute is deprecated and will be removed in version {self.higher_version}. message"""
			)



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			__snake_case      :		str         =							{'''deprecated_arg''': 4}

			with self.assertRaises(a_   ) as error:
				deprecate(('''wrong_arg''', self.lower_version, '''message''')							,			take_from=a_   )

			assert (
			    str(error.exception   )
			    == "The deprecation tuple ('wrong_arg', '0.0.1', 'message') should be removed since diffusers' version"
			    f""" {__version__} is >= {self.lower_version}"""
			)



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			with self.assertWarns(a_   ) as warning:
				deprecate(('''deprecated_arg''', self.higher_version, '''This message is better!!!''')							,			standard_warn=a_   )

			assert str(warning.warning   ) == "This message is better!!!"



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			with self.assertWarns(a_   ) as warning:
				deprecate(('''deprecated_arg''', self.higher_version, '''This message is better!!!''')							,			standard_warn=a_   )
			assert str(warning.warning   ) == "This message is better!!!"
			assert "diffusers/tests/others/test_utils.py" in warning.filename





