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


def       UpperCAmelCase_					(  __lowerCAmelCase		,   __lowerCAmelCase = True		)	->  Union[str, Any]:
    __lowercase						:     str		      =    {}
    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						:     str		      =    model.config.cross_attention_dim if has_cross_attention else None
        if name.startswith('''mid_block'''		):
            __lowercase						:     Union[str, Any]		      =    model.config.block_out_channels[-1]
        elif name.startswith('''up_blocks'''		):
            __lowercase						:     Optional[Any]		      =    int(name[len('''up_blocks.'''		)]		)
            __lowercase						:     List[str]		      =    list(reversed(model.config.block_out_channels		)		)[block_id]
        elif name.startswith('''down_blocks'''		):
            __lowercase						:     Optional[int]		      =    int(name[len('''down_blocks.'''		)]		)
            __lowercase						:     Any		      =    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=__lowerCAmelCase		,   cross_attention_dim=__lowerCAmelCase		)
        __lowercase						:     Optional[int]		      =    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   __lowerCAmelCase		(		lowerCAmelCase_	,	lowerCAmelCase_	,	unittest.TestCase ):



      """simple docstring"""







      A__		:     Union[str, Any]       				=      UNetaDConditionModel
      A__		:     Optional[int]       				=      '''sample'''




      @property
      def 						snake_case_     ( self						:    Union[str, Any]			):
          __lowercase						:     List[str]		      =    4
          __lowercase						:     Optional[Any]		      =    4
          __lowercase						:     Any		      =    4
          __lowercase						:     List[str]		      =    (32, 32)

          __lowercase						:     Any		      =    floats_tensor((batch_size, num_channels, num_frames) + sizes			).to(_snake_case			)
          __lowercase						:     int		      =    torch.tensor([10]			).to(_snake_case			)
          __lowercase						:     Tuple		      =    floats_tensor((batch_size, 4, 32)			).to(_snake_case			)

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




      @property
      def 						snake_case_     ( self						:    Optional[Any]			):
          return (4, 4, 32, 32)




      @property
      def 						snake_case_     ( self						:    List[str]			):
          return (4, 4, 32, 32)




      def 						snake_case_     ( self						:    Dict			):
          __lowercase						:     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,
          }
          __lowercase						:     List[Any]		      =    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]			):
          __lowercase						:     Dict		      =    self.prepare_init_args_and_inputs_for_common()
          __lowercase						:     str		      =    self.model_class(**_snake_case			)

          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			):
          __lowercase						:     Any		      =    self.prepare_init_args_and_inputs_for_common()

          __lowercase						:     Tuple		      =    32

          __lowercase						:     Optional[int]		      =    self.model_class(**_snake_case			)
          model.to(_snake_case			)
          model.eval()

          with torch.no_grad():
              __lowercase						:     str		      =    model(**_snake_case			)

              if isinstance(_snake_case			,    _snake_case			):
                  __lowercase						:     Dict		      =    output.sample

          self.assertIsNotNone(_snake_case			)
          __lowercase						:     str		      =    inputs_dict['''sample'''].shape
          self.assertEqual(output.shape			,    _snake_case			,    '''Input and output shapes do not match'''			)




      def 						snake_case_     ( self						:    Dict			):
          __lowercase						:     Optional[Any]		      =    self.prepare_init_args_and_inputs_for_common()
          __lowercase						:     Any		      =    self.model_class(**_snake_case			)
          model.to(_snake_case			)
          model.eval()

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

              __lowercase						:     List[str]		      =    model(**_snake_case			)
              if isinstance(_snake_case			,    _snake_case			):
                  __lowercase						:     List[Any]		      =    first.sample

              __lowercase						:     Optional[Any]		      =    model(**_snake_case			)
              if isinstance(_snake_case			,    _snake_case			):
                  __lowercase						:     Optional[Any]		      =    second.sample

          __lowercase						:     Optional[Any]		      =    first.cpu().numpy()
          __lowercase						:     Optional[Any]		      =    second.cpu().numpy()
          __lowercase						:     Dict		      =    out_a[~np.isnan(_snake_case			)]
          __lowercase						:     Optional[Any]		      =    out_a[~np.isnan(_snake_case			)]
          __lowercase						:     Dict		      =    np.amax(np.abs(out_a - out_a			)			)
          self.assertLessEqual(_snake_case			,    1E-5			)




      def 						snake_case_     ( self						:    Dict			):
          __lowercase						:     Dict		      =    self.prepare_init_args_and_inputs_for_common()

          __lowercase						:     List[Any]		      =    8

          __lowercase						:     List[Any]		      =    self.model_class(**_snake_case			)
          model.to(_snake_case			)
          model.eval()

          model.set_attention_slice('''auto'''			)
          with torch.no_grad():
              __lowercase						:     Optional[int]		      =    model(**_snake_case			)
          assert output is not None

          model.set_attention_slice('''max'''			)
          with torch.no_grad():
              __lowercase						:     Optional[int]		      =    model(**_snake_case			)
          assert output is not None

          model.set_attention_slice(2			)
          with torch.no_grad():
              __lowercase						:     int		      =    model(**_snake_case			)
          assert output is not None




      def 						snake_case_     ( self						:    Optional[Any]			):
          __lowercase						:     Optional[Any]		      =    self.prepare_init_args_and_inputs_for_common()

          __lowercase						:     Union[str, Any]		      =    8

          __lowercase						:     List[Any]		      =    self.model_class(**_snake_case			)
          model.to(_snake_case			)

          with torch.no_grad():
              __lowercase						:     Any		      =    model(**_snake_case			).sample

          __lowercase						:     Tuple		      =    create_lora_layers(_snake_case			)

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

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

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

          __lowercase						:     Optional[int]		      =    8

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

          with torch.no_grad():
              __lowercase						:     Dict		      =    model(**_snake_case			).sample

          __lowercase						:     int		      =    create_lora_layers(_snake_case			)
          model.set_attn_processor(_snake_case			)

          with torch.no_grad():
              __lowercase						:     List[str]		      =    model(**_snake_case			,    cross_attention_kwargs={'''scale''': 0.5}			).sample

          with tempfile.TemporaryDirectory() as tmpdirname:
              model.save_attn_procs(_snake_case			)
              self.assertTrue(os.path.isfile(os.path.join(_snake_case			,    '''pytorch_lora_weights.bin'''			)			)			)
              torch.manual_seed(0			)
              __lowercase						:     Optional[int]		      =    self.model_class(**_snake_case			)
              new_model.to(_snake_case			)
              new_model.load_attn_procs(_snake_case			)

          with torch.no_grad():
              __lowercase						:     List[Any]		      =    new_model(**_snake_case			,    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						:    Dict			):
          __lowercase						:     Union[str, Any]		      =    self.prepare_init_args_and_inputs_for_common()

          __lowercase						:     Optional[Any]		      =    8

          torch.manual_seed(0			)
          __lowercase						:     str		      =    self.model_class(**_snake_case			)
          model.to(_snake_case			)

          with torch.no_grad():
              __lowercase						:     Optional[Any]		      =    model(**_snake_case			).sample

          __lowercase						:     Dict		      =    create_lora_layers(_snake_case			)
          model.set_attn_processor(_snake_case			)

          with torch.no_grad():
              __lowercase						:     int		      =    model(**_snake_case			,    cross_attention_kwargs={'''scale''': 0.5}			).sample

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

          with torch.no_grad():
              __lowercase						:     Dict		      =    new_model(**_snake_case			,    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						:    Optional[Any]			):
          # enable deterministic behavior for gradient checkpointing
          __lowercase						:     Dict		      =    self.prepare_init_args_and_inputs_for_common()

          __lowercase						:     Optional[int]		      =    8

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

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




      def 						snake_case_     ( self						:    Union[str, Any]			):
          __lowercase						:     List[str]		      =    self.prepare_init_args_and_inputs_for_common()

          __lowercase						:     Tuple		      =    8

          torch.manual_seed(0			)
          __lowercase						:     Dict		      =    self.model_class(**_snake_case			)
          model.to(_snake_case			)

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




      def 						snake_case_     ( self						:    List[Any]			):
          __lowercase						:     int		      =    self.prepare_init_args_and_inputs_for_common()

          __lowercase						:     Optional[int]		      =    8

          torch.manual_seed(0			)
          __lowercase						:     str		      =    self.model_class(**_snake_case			)
          model.to(_snake_case			)

          with torch.no_grad():
              __lowercase						:     Any		      =    model(**_snake_case			).sample

          __lowercase						:     Optional[Any]		      =    create_lora_layers(_snake_case			)
          model.set_attn_processor(_snake_case			)

          with torch.no_grad():
              __lowercase						:     Any		      =    model(**_snake_case			,    cross_attention_kwargs={'''scale''': 0.0}			).sample

          model.set_attn_processor(AttnProcessor()			)

          with torch.no_grad():
              __lowercase						:     Optional[int]		      =    model(**_snake_case			).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						:    int			):
          # enable deterministic behavior for gradient checkpointing
          __lowercase						:     List[str]		      =    self.prepare_init_args_and_inputs_for_common()

          __lowercase						:     Optional[Any]		      =    4

          torch.manual_seed(0			)
          __lowercase						:     Any		      =    self.model_class(**_snake_case			)
          model.to(_snake_case			)
          __lowercase						:     List[Any]		      =    create_lora_layers(_snake_case			)
          model.set_attn_processor(_snake_case			)

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

              model.enable_xformers_memory_efficient_attention()
              __lowercase						:     str		      =    model(**_snake_case			).sample

              model.disable_xformers_memory_efficient_attention()
              __lowercase						:     Union[str, Any]		      =    model(**_snake_case			).sample

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




      def 						snake_case_     ( self						:    Dict			):
          __lowercase						:     List[str]		      =    self.prepare_init_args_and_inputs_for_common()
          __lowercase						:     List[str]		      =    32

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

          with torch.no_grad():
              __lowercase						:     Any		      =    model(**_snake_case			)[0]

          model.enable_forward_chunking()
          with torch.no_grad():
              __lowercase						:     Optional[int]		      =    model(**_snake_case			)[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


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


class   __lowerCAmelCase		(		lowerCAmelCase_ ):



      """simple docstring"""







      def 						snake_case_     ( self						:    List[Any]			):
          # 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(_snake_case			)
              except ImportError as e:
                  self.assertIn(F'{req} is required'			,    str(_snake_case			)			)
                  self.assertIn('''but found'''			,    str(_snake_case			)			)

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

        # bogus requirements formats:
        # 1. whole thing
          for req in ["numpy??1.0.0", "numpy1.0.0"]:
              try:
                  require_version_core(_snake_case			)
              except ValueError as e:
                  self.assertIn('''requirement needs to be in the pip package format'''			,    str(_snake_case			)			)
        # 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(_snake_case			)
              except ValueError as e:
                  self.assertIn('''need one of '''			,    str(_snake_case			)			)




      def 						snake_case_     ( self						:    List[str]			):
          # 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(_snake_case			)
              except ImportError as e:
                  self.assertIn(F'{req} is required'			,    str(_snake_case			)			)
                  self.assertIn(F'but found python=={python_ver}'			,    str(_snake_case			)			)


