




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    :   Dict       =logging.get_logger(__name__)

def 			lowerCAmelCase_			(  lowerCamelCase:       Any     ,				lowerCamelCase:       bool = True						)					->			Union[str, Any]:


      '''simple docstring'''




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

            lowerCamelCase__        =				LoRAAttnProcessor(hidden_size=lowerCamelCase     ,				cross_attention_dim=lowerCamelCase						)
            lowerCamelCase__        =				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			(		_lowerCAmelCase			,						_lowerCAmelCase			,						unittest.TestCase		):
     """simple docstring"""

     _snake_case:	Tuple     =UNetaDConditionModel
     _snake_case:	Optional[Any]     ="sample"




     @property
     def  __snake_case   (    self  :   List[str]			)				->       Dict:
           lowerCamelCase__        =				4
           lowerCamelCase__        =				4
           lowerCamelCase__        =				4
           lowerCamelCase__        =				(3_2, 3_2)

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

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




     @property
     def  __snake_case   (    self  :   Tuple			)				->       List[str]:
           return (4, 4, 3_2, 3_2)




     @property
     def  __snake_case   (    self  :   Dict			)				->       Optional[int]:
           return (4, 4, 3_2, 3_2)




     def  __snake_case   (    self  :   Optional[int]			)				->       Optional[int]:
           lowerCamelCase__        =				{
               '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,
           }
           lowerCamelCase__        =				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  __snake_case   (    self  :   Union[str, Any]			)				->       Tuple:
           lowerCamelCase__  ,		lowerCamelCase__        =				self.prepare_init_args_and_inputs_for_common()
           lowerCamelCase__        =				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  __snake_case   (    self  :   int			)				->       Optional[int]:
           lowerCamelCase__  ,		lowerCamelCase__        =				self.prepare_init_args_and_inputs_for_common()

           lowerCamelCase__        =				3_2

           lowerCamelCase__        =				self.model_class(**A__			)
           model.to(A__			)
           model.eval()

           with torch.no_grad():
                 lowerCamelCase__        =				model(**A__			)

                 if isinstance(A__	,  A__			):
                       lowerCamelCase__        =				output.sample

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




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

                 lowerCamelCase__        =				model(**A__			)
                 if isinstance(A__	,  A__			):
                       lowerCamelCase__        =				first.sample

                 lowerCamelCase__        =				model(**A__			)
                 if isinstance(A__	,  A__			):
                       lowerCamelCase__        =				second.sample

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




     def  __snake_case   (    self  :   Optional[int]			)				->       Optional[int]:
           lowerCamelCase__  ,		lowerCamelCase__        =				self.prepare_init_args_and_inputs_for_common()

           lowerCamelCase__        =				8

           lowerCamelCase__        =				self.model_class(**A__			)
           model.to(A__			)
           model.eval()

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

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

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




     def  __snake_case   (    self  :   Union[str, Any]			)				->       Tuple:
           lowerCamelCase__  ,		lowerCamelCase__        =				self.prepare_init_args_and_inputs_for_common()

           lowerCamelCase__        =				8

           lowerCamelCase__        =				self.model_class(**A__			)
           model.to(A__			)

           with torch.no_grad():
                 lowerCamelCase__        =				model(**A__			).sample

           lowerCamelCase__        =				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():
                 lowerCamelCase__        =				model(**A__	,  cross_attention_kwargs={'scale': 0.0}			).sample
                 lowerCamelCase__        =				model(**A__	,  cross_attention_kwargs={'scale': 0.5}			).sample
                 lowerCamelCase__        =				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  __snake_case   (    self  :   int			)				->       Any:
           lowerCamelCase__  ,		lowerCamelCase__        =				self.prepare_init_args_and_inputs_for_common()

           lowerCamelCase__        =				8

           torch.manual_seed(0			)
           lowerCamelCase__        =				self.model_class(**A__			)
           model.to(A__			)

           with torch.no_grad():
                 lowerCamelCase__        =				model(**A__			).sample

           lowerCamelCase__        =				create_lora_layers(A__			)
           model.set_attn_processor(A__			)

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

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

           lowerCamelCase__        =				8

           torch.manual_seed(0			)
           lowerCamelCase__        =				self.model_class(**A__			)
           model.to(A__			)

           with torch.no_grad():
                 lowerCamelCase__        =				model(**A__			).sample

           lowerCamelCase__        =				create_lora_layers(A__			)
           model.set_attn_processor(A__			)

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

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

           lowerCamelCase__        =				8

           torch.manual_seed(0			)
           lowerCamelCase__        =				self.model_class(**A__			)
           model.to(A__			)

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




     def  __snake_case   (    self  :   List[Any]			)				->       Any:
           lowerCamelCase__  ,		lowerCamelCase__        =				self.prepare_init_args_and_inputs_for_common()

           lowerCamelCase__        =				8

           torch.manual_seed(0			)
           lowerCamelCase__        =				self.model_class(**A__			)
           model.to(A__			)

           lowerCamelCase__        =				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			)
                 lowerCamelCase__        =				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  __snake_case   (    self  :   Tuple			)				->       Tuple:
           lowerCamelCase__  ,		lowerCamelCase__        =				self.prepare_init_args_and_inputs_for_common()

           lowerCamelCase__        =				8

           torch.manual_seed(0			)
           lowerCamelCase__        =				self.model_class(**A__			)
           model.to(A__			)

           with torch.no_grad():
                 lowerCamelCase__        =				model(**A__			).sample

           lowerCamelCase__        =				create_lora_layers(A__			)
           model.set_attn_processor(A__			)

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

           model.set_attn_processor(AttnProcessor()			)

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

           lowerCamelCase__        =				4

           torch.manual_seed(0			)
           lowerCamelCase__        =				self.model_class(**A__			)
           model.to(A__			)
           lowerCamelCase__        =				create_lora_layers(A__			)
           model.set_attn_processor(A__			)

           # default
           with torch.no_grad():
                 lowerCamelCase__        =				model(**A__			).sample

                 model.enable_xformers_memory_efficient_attention()
                 lowerCamelCase__        =				model(**A__			).sample

                 model.disable_xformers_memory_efficient_attention()
                 lowerCamelCase__        =				model(**A__			).sample

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




     def  __snake_case   (    self  :   Optional[int]			)				->       Union[str, Any]:
           lowerCamelCase__  ,		lowerCamelCase__        =				self.prepare_init_args_and_inputs_for_common()
           lowerCamelCase__        =				3_2

           lowerCamelCase__        =				self.model_class(**A__			)
           model.to(A__			)
           model.eval()

           with torch.no_grad():
                 lowerCamelCase__        =				model(**A__			)[0]

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





import importlib.metadata
import sys

from transformers.testing_utils import TestCasePlus
from transformers.utils.versions import require_version, require_version_core


__SCREAMING_SNAKE_CASE    :   int       =importlib.metadata.version("""numpy""")
__SCREAMING_SNAKE_CASE    :   Optional[int]       =""".""".join([str(x) for x in sys.version_info[:3]])






class 			A			(		_lowerCAmelCase		):
     """simple docstring"""

     def  __snake_case   (    self  :   List[Any]			)				->       Optional[int]:
           # lt + different version strings
           require_version_core('numpy<1000.4.5'			)
           require_version_core('numpy<1000.4'			)
           require_version_core('numpy<1000'			)

           # le
           require_version_core('numpy<=1000.4.5'			)
           require_version_core(F"""numpy<={numpy_ver}"""			)

           # eq
           require_version_core(F"""numpy=={numpy_ver}"""			)

           # ne
           require_version_core('numpy!=1000.4.5'			)

           # ge
           require_version_core('numpy>=1.0'			)
           require_version_core('numpy>=1.0.0'			)
           require_version_core(F"""numpy>={numpy_ver}"""			)

           # gt
           require_version_core('numpy>1.0.0'			)

           # mix
           require_version_core('numpy>1.0.0,<1000'			)

           # requirement w/o version
           require_version_core('numpy'			)

           # unmet requirements due to version conflict
           for req in ["numpy==1.0.0", "numpy>=1000.0.0", F"""numpy<{numpy_ver}"""]:
                 try:
                       require_version_core(A__			)
                 except ImportError as e:
                       self.assertIn(F"""{req} is required"""	,  str(A__			)			)
                       self.assertIn('but found'	,  str(A__			)			)

        # unmet requirements due to missing module
           for req in ["numpipypie>1", "numpipypie2"]:
                 try:
                       require_version_core(A__			)
                 except importlib.metadata.PackageNotFoundError as e:
                       self.assertIn(F"""The '{req}' distribution was not found and is required by this application"""	,  str(A__			)			)
                       self.assertIn('Try: pip install transformers -U'	,  str(A__			)			)

        # bogus requirements formats:
        # 1. whole thing
           for req in ["numpy??1.0.0", "numpy1.0.0"]:
                 try:
                       require_version_core(A__			)
                 except ValueError as e:
                       self.assertIn('requirement needs to be in the pip package format'	,  str(A__			)			)
        # 2. only operators
           for req in ["numpy=1.0.0", "numpy == 1.00", "numpy<>1.0.0", "numpy><1.00", "numpy>>1.0.0"]:
                 try:
                       require_version_core(A__			)
                 except ValueError as e:
                       self.assertIn('need one of '	,  str(A__			)			)




     def  __snake_case   (    self  :   str			)				->       Optional[int]:
           # matching requirement
           require_version('python>=3.6.0'			)

           # not matching requirements
           for req in ["python>9.9.9", "python<3.0.0"]:
                 try:
                       require_version_core(A__			)
                 except ImportError as e:
                       self.assertIn(F"""{req} is required"""	,  str(A__			)			)
                       self.assertIn(F"""but found python=={python_ver}"""	,  str(A__			)			)