






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

A:   List[str]						=					logging.get_logger(__name__)






def 						_UpperCAmelCase     (	a		:						Tuple    ,      a		:						bool = True				)      ->      Union[str, Any]:





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

												lowercase_   : Dict   	=      LoRAAttnProcessor(hidden_size=UpperCAmelCase_    ,      cross_attention_dim=UpperCAmelCase_				)
												lowercase_   : 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 				__magic_name__  (						UpperCAmelCase__,    UpperCAmelCase__,    unittest.TestCase					):







	"""simple docstring"""



	SCREAMING_SNAKE_CASE_       :   Union[str, Any]   =	UNetaDConditionModel
	SCREAMING_SNAKE_CASE_       :   Any   =	"sample"
	@property
	def 							lowerCamelCase__		(				self							)    ->	int:
							lowercase_   : Optional[int]   	=      4
							lowercase_   : List[str]   	=      4
							lowercase_   : Union[str, Any]   	=      4
							lowercase_   : Union[str, Any]   	=      (32, 32)

							lowercase_   : Optional[int]   	=      floats_tensor((batch_size, num_channels, num_frames) + sizes							).to(_lowerCAmelCase							)
							lowercase_   : int   	=      torch.tensor([10]							).to(_lowerCAmelCase							)
							lowercase_   : Optional[int]   	=      floats_tensor((batch_size, 4, 32)							).to(_lowerCAmelCase							)

							return {"sample": noise, "timestep": time_step, "encoder_hidden_states": encoder_hidden_states}
	@property
	def 							lowerCamelCase__		(				self							)    ->	Any:
							return (4, 4, 32, 32)
	@property
	def 							lowerCamelCase__		(				self							)    ->	Tuple:
							return (4, 4, 32, 32)
	def 							lowerCamelCase__		(				self							)    ->	Tuple:
							lowercase_   : Tuple   	=      {
							    '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,
							}
							lowercase_   : 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 							lowerCamelCase__		(				self							)    ->	List[Any]:
							lowercase_	,   lowercase_   : Tuple   	=      self.prepare_init_args_and_inputs_for_common()
							lowercase_   : Optional[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 							lowerCamelCase__		(				self							)    ->	int:
							lowercase_	,   lowercase_   : Optional[Any]   	=      self.prepare_init_args_and_inputs_for_common()

							lowercase_   : List[Any]   	=      32

							lowercase_   : Tuple   	=      self.model_class(**_lowerCAmelCase							)
							model.to(_lowerCAmelCase							)
							model.eval()

							with torch.no_grad():
													lowercase_   : List[str]   	=      model(**_lowerCAmelCase							)

													if isinstance(_lowerCAmelCase	,   _lowerCAmelCase							):
																			lowercase_   : Any   	=      output.sample

							self.assertIsNotNone(_lowerCAmelCase							)
							lowercase_   : Any   	=      inputs_dict['sample'].shape
							self.assertEqual(output.shape	,   _lowerCAmelCase	,   'Input and output shapes do not match'							)
	def 							lowerCamelCase__		(				self							)    ->	Any:
							lowercase_	,   lowercase_   : Dict   	=      self.prepare_init_args_and_inputs_for_common()
							lowercase_   : int   	=      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							)

													lowercase_   : int   	=      model(**_lowerCAmelCase							)
													if isinstance(_lowerCAmelCase	,   _lowerCAmelCase							):
																			lowercase_   : int   	=      first.sample

													lowercase_   : Optional[int]   	=      model(**_lowerCAmelCase							)
													if isinstance(_lowerCAmelCase	,   _lowerCAmelCase							):
																			lowercase_   : Tuple   	=      second.sample

							lowercase_   : int   	=      first.cpu().numpy()
							lowercase_   : Tuple   	=      second.cpu().numpy()
							lowercase_   : str   	=      out_a[~np.isnan(_lowerCAmelCase							)]
							lowercase_   : List[Any]   	=      out_a[~np.isnan(_lowerCAmelCase							)]
							lowercase_   : Dict   	=      np.amax(np.abs(out_a - out_a							)							)
							self.assertLessEqual(_lowerCAmelCase	,   1E-5							)
	def 							lowerCamelCase__		(				self							)    ->	Dict:
							lowercase_	,   lowercase_   : str   	=      self.prepare_init_args_and_inputs_for_common()

							lowercase_   : Optional[int]   	=      8

							lowercase_   : Any   	=      self.model_class(**_lowerCAmelCase							)
							model.to(_lowerCAmelCase							)
							model.eval()

							model.set_attention_slice('auto'							)
							with torch.no_grad():
													lowercase_   : Optional[Any]   	=      model(**_lowerCAmelCase							)
							assert output is not None

							model.set_attention_slice('max'							)
							with torch.no_grad():
													lowercase_   : Union[str, Any]   	=      model(**_lowerCAmelCase							)
							assert output is not None

							model.set_attention_slice(2							)
							with torch.no_grad():
													lowercase_   : Tuple   	=      model(**_lowerCAmelCase							)
							assert output is not None
	def 							lowerCamelCase__		(				self							)    ->	Any:
							lowercase_	,   lowercase_   : Optional[Any]   	=      self.prepare_init_args_and_inputs_for_common()

							lowercase_   : Dict   	=      8

							lowercase_   : Optional[int]   	=      self.model_class(**_lowerCAmelCase							)
							model.to(_lowerCAmelCase							)

							with torch.no_grad():
													lowercase_   : List[str]   	=      model(**_lowerCAmelCase							).sample

							lowercase_   : int   	=      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():
													lowercase_   : Tuple   	=      model(**_lowerCAmelCase	,   cross_attention_kwargs={'scale': 0.0}							).sample
													lowercase_   : Union[str, Any]   	=      model(**_lowerCAmelCase	,   cross_attention_kwargs={'scale': 0.5}							).sample
													lowercase_   : Optional[Any]   	=      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 							lowerCamelCase__		(				self							)    ->	Optional[Any]:
							lowercase_	,   lowercase_   : Any   	=      self.prepare_init_args_and_inputs_for_common()

							lowercase_   : str   	=      8

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

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

							lowercase_   : str   	=      create_lora_layers(_lowerCAmelCase							)
							model.set_attn_processor(_lowerCAmelCase							)

							with torch.no_grad():
													lowercase_   : List[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							)
													lowercase_   : Union[str, Any]   	=      self.model_class(**_lowerCAmelCase							)
													new_model.to(_lowerCAmelCase							)
													new_model.load_attn_procs(_lowerCAmelCase							)

							with torch.no_grad():
													lowercase_   : str   	=      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 							lowerCamelCase__		(				self							)    ->	int:
							lowercase_	,   lowercase_   : List[Any]   	=      self.prepare_init_args_and_inputs_for_common()

							lowercase_   : List[str]   	=      8

							torch.manual_seed(0							)
							lowercase_   : str   	=      self.model_class(**_lowerCAmelCase							)
							model.to(_lowerCAmelCase							)

							with torch.no_grad():
													lowercase_   : List[str]   	=      model(**_lowerCAmelCase							).sample

							lowercase_   : Tuple   	=      create_lora_layers(_lowerCAmelCase							)
							model.set_attn_processor(_lowerCAmelCase							)

							with torch.no_grad():
													lowercase_   : Dict   	=      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							)
													lowercase_   : Optional[Any]   	=      self.model_class(**_lowerCAmelCase							)
													new_model.to(_lowerCAmelCase							)
													new_model.load_attn_procs(_lowerCAmelCase							)

							with torch.no_grad():
													lowercase_   : Dict   	=      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 							lowerCamelCase__		(				self							)    ->	int:
							# enable deterministic behavior for gradient checkpointing
							lowercase_	,   lowercase_   : Union[str, Any]   	=      self.prepare_init_args_and_inputs_for_common()

							lowercase_   : int   	=      8

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

							lowercase_   : 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							)
													lowercase_   : str   	=      self.model_class(**_lowerCAmelCase							)
													new_model.to(_lowerCAmelCase							)
													new_model.load_attn_procs(_lowerCAmelCase	,   weight_name='pytorch_lora_weights.bin'							)
	def 							lowerCamelCase__		(				self							)    ->	int:
							lowercase_	,   lowercase_   : List[str]   	=      self.prepare_init_args_and_inputs_for_common()

							lowercase_   : str   	=      8

							torch.manual_seed(0							)
							lowercase_   : Union[str, Any]   	=      self.model_class(**_lowerCAmelCase							)
							model.to(_lowerCAmelCase							)

							lowercase_   : 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							)
													lowercase_   : List[Any]   	=      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 							lowerCamelCase__		(				self							)    ->	Optional[int]:
							lowercase_	,   lowercase_   : Tuple   	=      self.prepare_init_args_and_inputs_for_common()

							lowercase_   : List[str]   	=      8

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

							with torch.no_grad():
													lowercase_   : Optional[Any]   	=      model(**_lowerCAmelCase							).sample

							lowercase_   : Dict   	=      create_lora_layers(_lowerCAmelCase							)
							model.set_attn_processor(_lowerCAmelCase							)

							with torch.no_grad():
													lowercase_   : str   	=      model(**_lowerCAmelCase	,   cross_attention_kwargs={'scale': 0.0}							).sample

							model.set_attn_processor(AttnProcessor()							)

							with torch.no_grad():
													lowercase_   : List[str]   	=      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 							lowerCamelCase__		(				self							)    ->	Optional[int]:
							# enable deterministic behavior for gradient checkpointing
							lowercase_	,   lowercase_   : Dict   	=      self.prepare_init_args_and_inputs_for_common()

							lowercase_   : Tuple   	=      4

							torch.manual_seed(0							)
							lowercase_   : List[Any]   	=      self.model_class(**_lowerCAmelCase							)
							model.to(_lowerCAmelCase							)
							lowercase_   : str   	=      create_lora_layers(_lowerCAmelCase							)
							model.set_attn_processor(_lowerCAmelCase							)

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

													model.enable_xformers_memory_efficient_attention()
													lowercase_   : Any   	=      model(**_lowerCAmelCase							).sample

													model.disable_xformers_memory_efficient_attention()
													lowercase_   : List[Any]   	=      model(**_lowerCAmelCase							).sample

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







	def 							lowerCamelCase__		(				self							)    ->	int:
							lowercase_	,   lowercase_   : str   	=      self.prepare_init_args_and_inputs_for_common()
							lowercase_   : List[Any]   	=      32

							lowercase_   : Union[str, Any]   	=      self.model_class(**_lowerCAmelCase							)
							model.to(_lowerCAmelCase							)
							model.eval()

							with torch.no_grad():
													lowercase_   : Dict   	=      model(**_lowerCAmelCase							)[0]

							model.enable_forward_chunking()
							with torch.no_grad():
													lowercase_   : Union[str, Any]   	=      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.






import math
from timeit import timeit





def 							lowerCamelCase    (							UpperCAmelCase_	:      int      )->			int:

       """simple docstring"""




       a			   =0
       a			   =abs(UpperCAmelCase_      )
       while True:
              a			   =n // 10
              digits += 1
              if n == 0:
                     break
       return digits





def 							lowerCamelCase    (							UpperCAmelCase_	:      int      )->			int:

       """simple docstring"""




       return 1 if n == 0 else math.floor(math.log(abs(UpperCAmelCase_      )   ,       10      ) + 1      )





def 							lowerCamelCase    (							UpperCAmelCase_	:      int      )->			int:

       """simple docstring"""




       return len(str(abs(UpperCAmelCase_      )      )      )





def 							lowerCamelCase    (							)->			None:

       """simple docstring"""




       from collections.abc import Callable

       def benchmark_a_function(UpperCAmelCase_	:      Callable   ,       UpperCAmelCase_	:      int      ) -> None:
              a			   =F'''{func.__name__}({value})'''
              a			   =timeit(F'''__main__.{call}'''   ,       setup="""import __main__"""      )
              print(F'''{call}: {func(UpperCAmelCase_      )} -- {timing} seconds'''      )

       for value in (26_2144, 1125_8999_0684_2624, 126_7650_6002_2822_9401_4967_0320_5376):
              for func in (num_digits, num_digits_fast, num_digits_faster):
                     benchmark_a_function(UpperCAmelCase_   ,       UpperCAmelCase_      )
              print()


if __name__ == "__main__":
    import doctest

    doctest.testmod()
    benchmark()
