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

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

										__a							=   LoRAAttnProcessor(hidden_size=__lowerCamelCase      ,				cross_attention_dim=__lowerCamelCase						)
										__a							=   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		):
			A__				:   Any								=			UNetaDConditionModel
			A__				:   List[Any]								=			'sample'



			@property
			def __SCREAMING_SNAKE_CASE							(						self    ) ->    List[str]:
								__a							=   4
								__a							=   4
								__a							=   4
								__a							=   (3_2, 3_2)

								__a							=   floats_tensor((batch_size, num_channels, num_frames) + sizes    ).to(UpperCAmelCase    )
								__a							=   torch.tensor([1_0]    ).to(UpperCAmelCase    )
								__a							=   floats_tensor((batch_size, 4, 3_2)    ).to(UpperCAmelCase    )

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



			@property
			def __SCREAMING_SNAKE_CASE							(						self    ) ->    Dict:
								return (4, 4, 3_2, 3_2)



			@property
			def __SCREAMING_SNAKE_CASE							(						self    ) ->    Any:
								return (4, 4, 3_2, 3_2)



			def __SCREAMING_SNAKE_CASE							(						self    ) ->    Optional[int]:
								__a							=   {
								    '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,
								}
								__a							=   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 __SCREAMING_SNAKE_CASE							(						self    ) ->    List[Any]:
								__a					,      __a							=   self.prepare_init_args_and_inputs_for_common()
								__a							=   self.model_class(**UpperCAmelCase    )

								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 __SCREAMING_SNAKE_CASE							(						self    ) ->    Optional[Any]:
								__a					,      __a							=   self.prepare_init_args_and_inputs_for_common()

								__a							=   3_2

								__a							=   self.model_class(**UpperCAmelCase    )
								model.to(UpperCAmelCase    )
								model.eval()

								with torch.no_grad():
													__a							=   model(**UpperCAmelCase    )

													if isinstance(UpperCAmelCase      ,						UpperCAmelCase    ):
																		__a							=   output.sample

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



			def __SCREAMING_SNAKE_CASE							(						self    ) ->    Optional[int]:
								__a					,      __a							=   self.prepare_init_args_and_inputs_for_common()
								__a							=   self.model_class(**UpperCAmelCase    )
								model.to(UpperCAmelCase    )
								model.eval()

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

													__a							=   model(**UpperCAmelCase    )
													if isinstance(UpperCAmelCase      ,						UpperCAmelCase    ):
																		__a							=   first.sample

													__a							=   model(**UpperCAmelCase    )
													if isinstance(UpperCAmelCase      ,						UpperCAmelCase    ):
																		__a							=   second.sample

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



			def __SCREAMING_SNAKE_CASE							(						self    ) ->    Union[str, Any]:
								__a					,      __a							=   self.prepare_init_args_and_inputs_for_common()

								__a							=   8

								__a							=   self.model_class(**UpperCAmelCase    )
								model.to(UpperCAmelCase    )
								model.eval()

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

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

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



			def __SCREAMING_SNAKE_CASE							(						self    ) ->    List[str]:
								__a					,      __a							=   self.prepare_init_args_and_inputs_for_common()

								__a							=   8

								__a							=   self.model_class(**UpperCAmelCase    )
								model.to(UpperCAmelCase    )

								with torch.no_grad():
													__a							=   model(**UpperCAmelCase    ).sample

								__a							=   create_lora_layers(UpperCAmelCase    )

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

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

								with torch.no_grad():
													__a							=   model(**UpperCAmelCase      ,						cross_attention_kwargs={'scale': 0.0}    ).sample
													__a							=   model(**UpperCAmelCase      ,						cross_attention_kwargs={'scale': 0.5}    ).sample
													__a							=   model(**UpperCAmelCase      ,						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 __SCREAMING_SNAKE_CASE							(						self    ) ->    List[Any]:
								__a					,      __a							=   self.prepare_init_args_and_inputs_for_common()

								__a							=   8

								torch.manual_seed(0    )
								__a							=   self.model_class(**UpperCAmelCase    )
								model.to(UpperCAmelCase    )

								with torch.no_grad():
													__a							=   model(**UpperCAmelCase    ).sample

								__a							=   create_lora_layers(UpperCAmelCase    )
								model.set_attn_processor(UpperCAmelCase    )

								with torch.no_grad():
													__a							=   model(**UpperCAmelCase      ,						cross_attention_kwargs={'scale': 0.5}    ).sample

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

								with torch.no_grad():
													__a							=   new_model(**UpperCAmelCase      ,						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 __SCREAMING_SNAKE_CASE							(						self    ) ->    Dict:
								__a					,      __a							=   self.prepare_init_args_and_inputs_for_common()

								__a							=   8

								torch.manual_seed(0    )
								__a							=   self.model_class(**UpperCAmelCase    )
								model.to(UpperCAmelCase    )

								with torch.no_grad():
													__a							=   model(**UpperCAmelCase    ).sample

								__a							=   create_lora_layers(UpperCAmelCase    )
								model.set_attn_processor(UpperCAmelCase    )

								with torch.no_grad():
													__a							=   model(**UpperCAmelCase      ,						cross_attention_kwargs={'scale': 0.5}    ).sample

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

								with torch.no_grad():
													__a							=   new_model(**UpperCAmelCase      ,						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 __SCREAMING_SNAKE_CASE							(						self    ) ->    Any:
								# enable deterministic behavior for gradient checkpointing
								__a					,      __a							=   self.prepare_init_args_and_inputs_for_common()

								__a							=   8

								torch.manual_seed(0    )
								__a							=   self.model_class(**UpperCAmelCase    )
								model.to(UpperCAmelCase    )

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



			def __SCREAMING_SNAKE_CASE							(						self    ) ->    str:
								__a					,      __a							=   self.prepare_init_args_and_inputs_for_common()

								__a							=   8

								torch.manual_seed(0    )
								__a							=   self.model_class(**UpperCAmelCase    )
								model.to(UpperCAmelCase    )

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



			def __SCREAMING_SNAKE_CASE							(						self    ) ->    Union[str, Any]:
								__a					,      __a							=   self.prepare_init_args_and_inputs_for_common()

								__a							=   8

								torch.manual_seed(0    )
								__a							=   self.model_class(**UpperCAmelCase    )
								model.to(UpperCAmelCase    )

								with torch.no_grad():
													__a							=   model(**UpperCAmelCase    ).sample

								__a							=   create_lora_layers(UpperCAmelCase    )
								model.set_attn_processor(UpperCAmelCase    )

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

								model.set_attn_processor(AttnProcessor()    )

								with torch.no_grad():
													__a							=   model(**UpperCAmelCase    ).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 __SCREAMING_SNAKE_CASE							(						self    ) ->    Tuple:
								# enable deterministic behavior for gradient checkpointing
								__a					,      __a							=   self.prepare_init_args_and_inputs_for_common()

								__a							=   4

								torch.manual_seed(0    )
								__a							=   self.model_class(**UpperCAmelCase    )
								model.to(UpperCAmelCase    )
								__a							=   create_lora_layers(UpperCAmelCase    )
								model.set_attn_processor(UpperCAmelCase    )

								# default
								with torch.no_grad():
													__a							=   model(**UpperCAmelCase    ).sample

													model.enable_xformers_memory_efficient_attention()
													__a							=   model(**UpperCAmelCase    ).sample

													model.disable_xformers_memory_efficient_attention()
													__a							=   model(**UpperCAmelCase    ).sample

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







			def __SCREAMING_SNAKE_CASE							(						self    ) ->    Dict:
								__a					,      __a							=   self.prepare_init_args_and_inputs_for_common()
								__a							=   3_2

								__a							=   self.model_class(**UpperCAmelCase    )
								model.to(UpperCAmelCase    )
								model.eval()

								with torch.no_grad():
													__a							=   model(**UpperCAmelCase    )[0]

								model.enable_forward_chunking()
								with torch.no_grad():
													__a							=   model(**UpperCAmelCase    )[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 	UpperCAmelCase_    (	__UpperCamelCase	):
    if not isinstance(__UpperCamelCase,				__UpperCamelCase	):
        raise ValueError("""double_factorial() only accepts integral values"""	)
    if n < 0:
        raise ValueError("""double_factorial() not defined for negative values"""	)
    return 1 if n <= 1 else n * double_factorial(n - 2	)


if __name__ == "__main__":
    import doctest

    doctest.testmod()






