
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   :					List[Any]    =			logging.get_logger(__name__)


def  A			(   lowercase       ,			lowercase = True					)  ->    Union[str, Any]:
		'''simple docstring'''

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

				UpperCamelCase            =	LoRAAttnProcessor(hidden_size=lowercase       ,			cross_attention_dim=lowercase					)
				UpperCamelCase            =	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 							lowercase   (						_SCREAMING_SNAKE_CASE    ,  _SCREAMING_SNAKE_CASE    ,  unittest.TestCase  ):
		__lowercase			:			Any      	=							UNetaDConditionModel
		__lowercase			:			List[Any]      	=							"sample"



		@property
		def __UpperCamelCase					(					self			)    ->   int:
				"""simple docstring"""
				UpperCamelCase            =	4
				UpperCamelCase            =	4
				UpperCamelCase            =	4
				UpperCamelCase            =	(32, 32)

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

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



		@property
		def __UpperCamelCase					(					self			)    ->   int:
				"""simple docstring"""
				return (4, 4, 32, 32)



		@property
		def __UpperCamelCase					(					self			)    ->   Any:
				"""simple docstring"""
				return (4, 4, 32, 32)



		def __UpperCamelCase					(					self			)    ->   List[str]:
				"""simple docstring"""
				UpperCamelCase            =	{
				    '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,
				}
				UpperCamelCase            =	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			)    ->   List[Any]:
				"""simple docstring"""
				UpperCamelCase      ,    UpperCamelCase            =	self.prepare_init_args_and_inputs_for_common()
				UpperCamelCase            =	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			)    ->   List[Any]:
				"""simple docstring"""
				UpperCamelCase      ,    UpperCamelCase            =	self.prepare_init_args_and_inputs_for_common()

				UpperCamelCase            =	32

				UpperCamelCase            =	self.model_class(**A_			)
				model.to(A_			)
				model.eval()

				with torch.no_grad():
						UpperCamelCase            =	model(**A_			)

						if isinstance(A_							, A_			):
								UpperCamelCase            =	output.sample

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



		def __UpperCamelCase					(					self			)    ->   int:
				"""simple docstring"""
				UpperCamelCase      ,    UpperCamelCase            =	self.prepare_init_args_and_inputs_for_common()
				UpperCamelCase            =	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			)

						UpperCamelCase            =	model(**A_			)
						if isinstance(A_							, A_			):
								UpperCamelCase            =	first.sample

						UpperCamelCase            =	model(**A_			)
						if isinstance(A_							, A_			):
								UpperCamelCase            =	second.sample

				UpperCamelCase            =	first.cpu().numpy()
				UpperCamelCase            =	second.cpu().numpy()
				UpperCamelCase            =	out_a[~np.isnan(A_			)]
				UpperCamelCase            =	out_a[~np.isnan(A_			)]
				UpperCamelCase            =	np.amax(np.abs(out_a - out_a			)			)
				self.assertLessEqual(A_							, 1e-5			)



		def __UpperCamelCase					(					self			)    ->   Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase      ,    UpperCamelCase            =	self.prepare_init_args_and_inputs_for_common()

				UpperCamelCase            =	8

				UpperCamelCase            =	self.model_class(**A_			)
				model.to(A_			)
				model.eval()

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

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

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



		def __UpperCamelCase					(					self			)    ->   Tuple:
				"""simple docstring"""
				UpperCamelCase      ,    UpperCamelCase            =	self.prepare_init_args_and_inputs_for_common()

				UpperCamelCase            =	8

				UpperCamelCase            =	self.model_class(**A_			)
				model.to(A_			)

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

				UpperCamelCase            =	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():
						UpperCamelCase            =	model(**A_							, cross_attention_kwargs={'scale': 0.0}			).sample
						UpperCamelCase            =	model(**A_							, cross_attention_kwargs={'scale': 0.5}			).sample
						UpperCamelCase            =	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			)    ->   Optional[Any]:
				"""simple docstring"""
				UpperCamelCase      ,    UpperCamelCase            =	self.prepare_init_args_and_inputs_for_common()

				UpperCamelCase            =	8

				torch.manual_seed(0			)
				UpperCamelCase            =	self.model_class(**A_			)
				model.to(A_			)

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

				UpperCamelCase            =	create_lora_layers(A_			)
				model.set_attn_processor(A_			)

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

				with torch.no_grad():
						UpperCamelCase            =	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[int]:
				"""simple docstring"""
				UpperCamelCase      ,    UpperCamelCase            =	self.prepare_init_args_and_inputs_for_common()

				UpperCamelCase            =	8

				torch.manual_seed(0			)
				UpperCamelCase            =	self.model_class(**A_			)
				model.to(A_			)

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

				UpperCamelCase            =	create_lora_layers(A_			)
				model.set_attn_processor(A_			)

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

				with torch.no_grad():
						UpperCamelCase            =	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[Any]:
				"""simple docstring"""
				# enable deterministic behavior for gradient checkpointing
				UpperCamelCase      ,    UpperCamelCase            =	self.prepare_init_args_and_inputs_for_common()

				UpperCamelCase            =	8

				torch.manual_seed(0			)
				UpperCamelCase            =	self.model_class(**A_			)
				model.to(A_			)

				UpperCamelCase            =	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			)
						UpperCamelCase            =	self.model_class(**A_			)
						new_model.to(A_			)
						new_model.load_attn_procs(A_							, weight_name='pytorch_lora_weights.bin'			)



		def __UpperCamelCase					(					self			)    ->   Any:
				"""simple docstring"""
				UpperCamelCase      ,    UpperCamelCase            =	self.prepare_init_args_and_inputs_for_common()

				UpperCamelCase            =	8

				torch.manual_seed(0			)
				UpperCamelCase            =	self.model_class(**A_			)
				model.to(A_			)

				UpperCamelCase            =	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			)
						UpperCamelCase            =	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			)    ->   Any:
				"""simple docstring"""
				UpperCamelCase      ,    UpperCamelCase            =	self.prepare_init_args_and_inputs_for_common()

				UpperCamelCase            =	8

				torch.manual_seed(0			)
				UpperCamelCase            =	self.model_class(**A_			)
				model.to(A_			)

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

				UpperCamelCase            =	create_lora_layers(A_			)
				model.set_attn_processor(A_			)

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

				model.set_attn_processor(AttnProcessor()			)

				with torch.no_grad():
						UpperCamelCase            =	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[int]:
				"""simple docstring"""
				# enable deterministic behavior for gradient checkpointing
				UpperCamelCase      ,    UpperCamelCase            =	self.prepare_init_args_and_inputs_for_common()

				UpperCamelCase            =	4

				torch.manual_seed(0			)
				UpperCamelCase            =	self.model_class(**A_			)
				model.to(A_			)
				UpperCamelCase            =	create_lora_layers(A_			)
				model.set_attn_processor(A_			)

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

						model.enable_xformers_memory_efficient_attention()
						UpperCamelCase            =	model(**A_			).sample

						model.disable_xformers_memory_efficient_attention()
						UpperCamelCase            =	model(**A_			).sample

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



		def __UpperCamelCase					(					self			)    ->   Optional[Any]:
				"""simple docstring"""
				UpperCamelCase      ,    UpperCamelCase            =	self.prepare_init_args_and_inputs_for_common()
				UpperCamelCase            =	32

				UpperCamelCase            =	self.model_class(**A_			)
				model.to(A_			)
				model.eval()

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

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


def  A			(   lowercase					)  ->    int:
		'''simple docstring'''

		return int(lowercase					) if x - int(lowercase					) >= 0 else int(lowercase					) - 1


if __name__ == "__main__":
					import doctest

					doctest.testmod()
