




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

_lowerCamelCase					:     Dict				       =    logging.get_logger(__name__)

def 						lowercase_ (						_UpperCAmelCase					, _UpperCAmelCase = True		):




	"""simple docstring"""






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

		A_					:			Optional[Any]						=						LoRAAttnProcessor(hidden_size=_UpperCAmelCase					, cross_attention_dim=_UpperCAmelCase		)
		A_					:			Tuple						=						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 (      __UpperCAmelCase       ,		__UpperCAmelCase       ,		unittest.TestCase):
				__lowerCAmelCase      :      int					    =				UNetaDConditionModel
				__lowerCAmelCase      :      List[str]					    =				"""sample"""







				@property
				def    a_					(							self				:   Dict    ):



					"""simple docstring"""

					A_					:			Optional[int]						=						4
					A_					:			Dict						=						4
					A_					:			List[Any]						=						4
					A_					:			Optional[int]						=						(32, 32)

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

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







				@property
				def    a_					(							self				:   Tuple    ):



					"""simple docstring"""

					return (4, 4, 32, 32)







				@property
				def    a_					(							self				:   str    ):



					"""simple docstring"""

					return (4, 4, 32, 32)







				def    a_					(							self				:   List[Any]    ):



					"""simple docstring"""

					A_					:			int						=						{
					    '''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,
					}
					A_					:			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    a_					(							self				:   int    ):



					"""simple docstring"""

					A_				,	A_					:			str						=						self.prepare_init_args_and_inputs_for_common()
					A_					:			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    a_					(							self				:   str    ):



					"""simple docstring"""

					A_				,	A_					:			Any						=						self.prepare_init_args_and_inputs_for_common()

					A_					:			Optional[Any]						=						32

					A_					:			int						=						self.model_class(**_lowerCamelCase    )
					model.to(_lowerCamelCase    )
					model.eval()

					with torch.no_grad():
						A_					:			Dict						=						model(**_lowerCamelCase    )

						if isinstance(_lowerCamelCase ,       _lowerCamelCase    ):
							A_					:			Optional[int]						=						output.sample

					self.assertIsNotNone(_lowerCamelCase    )
					A_					:			Union[str, Any]						=						inputs_dict['''sample'''].shape
					self.assertEqual(output.shape ,       _lowerCamelCase ,       '''Input and output shapes do not match'''    )







				def    a_					(							self				:   List[str]    ):



					"""simple docstring"""

					A_				,	A_					:			List[str]						=						self.prepare_init_args_and_inputs_for_common()
					A_					:			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    )

						A_					:			Optional[Any]						=						model(**_lowerCamelCase    )
						if isinstance(_lowerCamelCase ,       _lowerCamelCase    ):
							A_					:			Union[str, Any]						=						first.sample

						A_					:			Optional[int]						=						model(**_lowerCamelCase    )
						if isinstance(_lowerCamelCase ,       _lowerCamelCase    ):
							A_					:			List[str]						=						second.sample

					A_					:			str						=						first.cpu().numpy()
					A_					:			Dict						=						second.cpu().numpy()
					A_					:			Union[str, Any]						=						out_a[~np.isnan(_lowerCamelCase    )]
					A_					:			Optional[Any]						=						out_a[~np.isnan(_lowerCamelCase    )]
					A_					:			List[Any]						=						np.amax(np.abs(out_a - out_a    )    )
					self.assertLessEqual(_lowerCamelCase ,       1E-5    )







				def    a_					(							self				:   Any    ):



					"""simple docstring"""

					A_				,	A_					:			List[str]						=						self.prepare_init_args_and_inputs_for_common()

					A_					:			str						=						8

					A_					:			Optional[Any]						=						self.model_class(**_lowerCamelCase    )
					model.to(_lowerCamelCase    )
					model.eval()

					model.set_attention_slice('''auto'''    )
					with torch.no_grad():
						A_					:			List[Any]						=						model(**_lowerCamelCase    )
					assert output is not None

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

					model.set_attention_slice(2    )
					with torch.no_grad():
						A_					:			Optional[Any]						=						model(**_lowerCamelCase    )
					assert output is not None







				def    a_					(							self				:   List[Any]    ):



					"""simple docstring"""

					A_				,	A_					:			List[str]						=						self.prepare_init_args_and_inputs_for_common()

					A_					:			List[str]						=						8

					A_					:			str						=						self.model_class(**_lowerCamelCase    )
					model.to(_lowerCamelCase    )

					with torch.no_grad():
						A_					:			List[Any]						=						model(**_lowerCamelCase    ).sample

					A_					:			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():
						A_					:			str						=						model(**_lowerCamelCase ,       cross_attention_kwargs={'''scale''': 0.0}    ).sample
						A_					:			Tuple						=						model(**_lowerCamelCase ,       cross_attention_kwargs={'''scale''': 0.5}    ).sample
						A_					:			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    a_					(							self				:   Dict    ):



					"""simple docstring"""

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

					A_					:			List[str]						=						8

					torch.manual_seed(0    )
					A_					:			Any						=						self.model_class(**_lowerCamelCase    )
					model.to(_lowerCamelCase    )

					with torch.no_grad():
						A_					:			List[Any]						=						model(**_lowerCamelCase    ).sample

					A_					:			List[Any]						=						create_lora_layers(_lowerCamelCase    )
					model.set_attn_processor(_lowerCamelCase    )

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

					with torch.no_grad():
						A_					:			Optional[int]						=						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    a_					(							self				:   List[Any]    ):



					"""simple docstring"""

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

					A_					:			Union[str, Any]						=						8

					torch.manual_seed(0    )
					A_					:			List[Any]						=						self.model_class(**_lowerCamelCase    )
					model.to(_lowerCamelCase    )

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

					A_					:			Tuple						=						create_lora_layers(_lowerCamelCase    )
					model.set_attn_processor(_lowerCamelCase    )

					with torch.no_grad():
						A_					:			Optional[int]						=						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    )
						A_					:			List[Any]						=						self.model_class(**_lowerCamelCase    )
						new_model.to(_lowerCamelCase    )
						new_model.load_attn_procs(_lowerCamelCase    )

					with torch.no_grad():
						A_					:			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    a_					(							self				:   Tuple    ):



					"""simple docstring"""

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

					A_					:			Dict						=						8

					torch.manual_seed(0    )
					A_					:			int						=						self.model_class(**_lowerCamelCase    )
					model.to(_lowerCamelCase    )

					A_					:			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    )
						A_					:			Optional[int]						=						self.model_class(**_lowerCamelCase    )
						new_model.to(_lowerCamelCase    )
						new_model.load_attn_procs(_lowerCamelCase ,       weight_name='''pytorch_lora_weights.bin'''    )







				def    a_					(							self				:   Dict    ):



					"""simple docstring"""

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

					A_					:			Tuple						=						8

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

					A_					:			str						=						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    )
						A_					:			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    a_					(							self				:   Union[str, Any]    ):



					"""simple docstring"""

					A_				,	A_					:			Dict						=						self.prepare_init_args_and_inputs_for_common()

					A_					:			Dict						=						8

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

					with torch.no_grad():
						A_					:			List[Any]						=						model(**_lowerCamelCase    ).sample

					A_					:			List[Any]						=						create_lora_layers(_lowerCamelCase    )
					model.set_attn_processor(_lowerCamelCase    )

					with torch.no_grad():
						A_					:			List[str]						=						model(**_lowerCamelCase ,       cross_attention_kwargs={'''scale''': 0.0}    ).sample

					model.set_attn_processor(AttnProcessor()    )

					with torch.no_grad():
						A_					:			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    a_					(							self				:   Optional[int]    ):



					"""simple docstring"""

					A_				,	A_					:			List[str]						=						self.prepare_init_args_and_inputs_for_common()

					A_					:			List[Any]						=						4

					torch.manual_seed(0    )
					A_					:			str						=						self.model_class(**_lowerCamelCase    )
					model.to(_lowerCamelCase    )
					A_					:			str						=						create_lora_layers(_lowerCamelCase    )
					model.set_attn_processor(_lowerCamelCase    )

					# default
					with torch.no_grad():
						A_					:			Union[str, Any]						=						model(**_lowerCamelCase    ).sample

						model.enable_xformers_memory_efficient_attention()
						A_					:			Dict						=						model(**_lowerCamelCase    ).sample

						model.disable_xformers_memory_efficient_attention()
						A_					:			Tuple						=						model(**_lowerCamelCase    ).sample

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







				def    a_					(							self				:   Optional[int]    ):



					"""simple docstring"""

					A_				,	A_					:			int						=						self.prepare_init_args_and_inputs_for_common()
					A_					:			Union[str, Any]						=						32

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

					with torch.no_grad():
						A_					:			List[Any]						=						model(**_lowerCamelCase    )[0]

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











"""simple docstring"""

import enum
import inspect
import warnings
from functools import wraps
from typing import Callable, Optional

from .logging import get_logger


_lowerCamelCase					:     Optional[int]				       =    set()
_lowerCamelCase					:     Union[str, Any]				       =    get_logger(__name__)

def 						lowercase_ (						_UpperCAmelCase = None		):




	"""simple docstring"""






	def decorator(_UpperCAmelCase		):
		global _emitted_deprecation_warnings

		if inspect.isclass(_UpperCAmelCase		):
			A_					:			Union[str, Any]						=						deprecated_class_or_function.__init__
			A_					:			int						=						deprecated_class_or_function.__name__
		else:
			A_					:			Union[str, Any]						=						deprecated_class_or_function
			A_					:			List[Any]						=						deprecated_function.__name__
			# Support deprecating __init__ class method: class name instead
			A_					:			List[Any]						=						name if name != '''__init__''' else deprecated_function.__qualname__.split('''.'''		)[-2]

		A_					:			str						=						(
		    f"""{name} is deprecated and will be removed in the next major version of datasets.""" + f""" {help_message}"""
		    if help_message
		    else ''''''
		)

		@wraps(_UpperCAmelCase		)
		def wrapper(*_UpperCAmelCase					, **_UpperCAmelCase		):
			A_					:			Any						=						hash(_UpperCAmelCase		)
			if func_hash not in _emitted_deprecation_warnings:
				warnings.warn(_UpperCAmelCase					, category=_UpperCAmelCase					, stacklevel=2		)
				_emitted_deprecation_warnings.add(_UpperCAmelCase		)
			return deprecated_function(*_UpperCAmelCase					, **_UpperCAmelCase		)

		A_					:			Any						=						'''deprecated'''

		if inspect.isclass(_UpperCAmelCase		):
			A_					:			List[Any]						=						wrapper
			return deprecated_class_or_function
		else:
			return wrapper

	return decorator

class     lowercase (      enum.EnumMeta):







				def __getattribute__(							cls				:   Any ,       _lowerCamelCase				:   List[Any]    ):



					"""simple docstring"""

					A_					:			str						=						super().__getattribute__(_lowerCamelCase    )
					if isinstance(_lowerCamelCase ,       enum.Enum    ) and obj._on_access:
						obj._on_access()
					return obj







				def __getitem__(							cls				:   Any ,       _lowerCamelCase				:   Union[str, Any]    ):



					"""simple docstring"""

					A_					:			Optional[Any]						=						super().__getitem__(_lowerCamelCase    )
					if member._on_access:
						member._on_access()
					return member







				def __call__(							cls				:   str ,       _lowerCamelCase				:   str ,       _lowerCamelCase				:   List[Any]=None ,       *, _lowerCamelCase				:   int=None ,       _lowerCamelCase				:   str=None ,       _lowerCamelCase				:   Dict=None ,       _lowerCamelCase				:   Tuple=1    ):



					"""simple docstring"""

					A_					:			Dict						=						super().__call__(_lowerCamelCase ,       _lowerCamelCase ,       module=_lowerCamelCase ,       qualname=_lowerCamelCase ,       type=_lowerCamelCase ,       start=_lowerCamelCase    )
					if isinstance(_lowerCamelCase ,       enum.Enum    ) and obj._on_access:
						obj._on_access()
					return obj

class     lowercase (      enum.Enum       ,		metaclass=__UpperCAmelCase):







				def __new__(							cls				:   Optional[Any] ,       _lowerCamelCase				:   int    ):



					"""simple docstring"""

					A_					:			Tuple						=						object.__new__(cls    )
					A_					:			Union[str, Any]						=						value
					A_					:			Union[str, Any]						=						member.deprecate
					return member







				@property
				def    a_					(							self				:   Tuple    ):



					"""simple docstring"""

					return ""







				def    a_					(							self				:   Any    ):



					"""simple docstring"""

					A_					:			Optional[int]						=						F""" {self.help_message}""" if self.help_message else ''''''
					warnings.warn(
					    F"""'{self.__objclass__.__name__}' is deprecated and will be removed in the next major version of datasets."""
					    + help_message ,       _lowerCamelCase ,       stacklevel=3 ,       )





