






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

_UpperCamelCase							:       Optional[int]       					=		logging.get_logger(__name__)




def      snake_case    (A_ :List[str]      ,     A_ :Union[str, Any] = True ):







  '''simple docstring'''

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

    a	:     List[Any]							=  LoRAAttnProcessor(hidden_size=__UpperCAmelCase      ,     cross_attention_dim=__UpperCAmelCase )
    a	:     List[Any]							=  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 		snake_case    (    _snake_case  ,					_snake_case  ,					unittest.TestCase			):
    __magic_name__       	=	UNetaDConditionModel
    __magic_name__       	=	"sample"





    @property
    def 							lowerCamelCase__ (				self      :   Tuple     ):


      '''simple docstring'''




      a	:     Tuple							=  4
      a	:     int							=  4
      a	:     Optional[Any]							=  4
      a	:     List[str]							=  (3_2, 3_2)

      a	:     int							=  floats_tensor((batch_size, num_channels, num_frames) + sizes     ).to(_lowercase     )
      a	:     Optional[int]							=  torch.tensor([1_0]     ).to(_lowercase     )
      a	:     str							=  floats_tensor((batch_size, 4, 3_2)     ).to(_lowercase     )

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





    @property
    def 							lowerCamelCase__ (				self      :   List[Any]     ):


      '''simple docstring'''




      return (4, 4, 3_2, 3_2)





    @property
    def 							lowerCamelCase__ (				self      :   str     ):


      '''simple docstring'''




      return (4, 4, 3_2, 3_2)





    def 							lowerCamelCase__ (				self      :   int     ):


      '''simple docstring'''




      a	:     List[Any]							=  {
          '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	:     Optional[int]							=  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 							lowerCamelCase__ (				self      :   Any     ):


      '''simple docstring'''




      a,	a	:     List[str]							=  self.prepare_init_args_and_inputs_for_common()
      a	:     int							=  self.model_class(**_lowercase     )

      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 							lowerCamelCase__ (				self      :   int     ):


      '''simple docstring'''




      a,	a	:     Tuple							=  self.prepare_init_args_and_inputs_for_common()

      a	:     Any							=  3_2

      a	:     str							=  self.model_class(**_lowercase     )
      model.to(_lowercase     )
      model.eval()

      with torch.no_grad():
        a	:     str							=  model(**_lowercase     )

        if isinstance(_lowercase							,						_lowercase     ):
          a	:     List[str]							=  output.sample

      self.assertIsNotNone(_lowercase     )
      a	:     Any							=  inputs_dict['sample'].shape
      self.assertEqual(output.shape							,						_lowercase							,						'Input and output shapes do not match'     )





    def 							lowerCamelCase__ (				self      :   Union[str, Any]     ):


      '''simple docstring'''




      a,	a	:     Any							=  self.prepare_init_args_and_inputs_for_common()
      a	:     List[str]							=  self.model_class(**_lowercase     )
      model.to(_lowercase     )
      model.eval()

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

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

        a	:     List[str]							=  model(**_lowercase     )
        if isinstance(_lowercase							,						_lowercase     ):
          a	:     List[str]							=  second.sample

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





    def 							lowerCamelCase__ (				self      :   List[str]     ):


      '''simple docstring'''




      a,	a	:     Union[str, Any]							=  self.prepare_init_args_and_inputs_for_common()

      a	:     Dict							=  8

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

      model.set_attention_slice('auto'     )
      with torch.no_grad():
        a	:     List[str]							=  model(**_lowercase     )
      assert output is not None

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

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





    def 							lowerCamelCase__ (				self      :   str     ):


      '''simple docstring'''




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

      a	:     Dict							=  8

      a	:     Any							=  self.model_class(**_lowercase     )
      model.to(_lowercase     )

      with torch.no_grad():
        a	:     List[str]							=  model(**_lowercase     ).sample

      a	:     Any							=  create_lora_layers(_lowercase     )

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

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

      with torch.no_grad():
        a	:     Dict							=  model(**_lowercase							,						cross_attention_kwargs={'scale': 0.0}     ).sample
        a	:     int							=  model(**_lowercase							,						cross_attention_kwargs={'scale': 0.5}     ).sample
        a	:     Tuple							=  model(**_lowercase							,						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 							lowerCamelCase__ (				self      :   Dict     ):


      '''simple docstring'''




      a,	a	:     Tuple							=  self.prepare_init_args_and_inputs_for_common()

      a	:     Any							=  8

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

      with torch.no_grad():
        a	:     Optional[int]							=  model(**_lowercase     ).sample

      a	:     List[str]							=  create_lora_layers(_lowercase     )
      model.set_attn_processor(_lowercase     )

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

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

      with torch.no_grad():
        a	:     Optional[Any]							=  new_model(**_lowercase							,						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 							lowerCamelCase__ (				self      :   Tuple     ):


      '''simple docstring'''




      a,	a	:     Any							=  self.prepare_init_args_and_inputs_for_common()

      a	:     Union[str, Any]							=  8

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

      with torch.no_grad():
        a	:     str							=  model(**_lowercase     ).sample

      a	:     Optional[Any]							=  create_lora_layers(_lowercase     )
      model.set_attn_processor(_lowercase     )

      with torch.no_grad():
        a	:     Union[str, Any]							=  model(**_lowercase							,						cross_attention_kwargs={'scale': 0.5}     ).sample

      with tempfile.TemporaryDirectory() as tmpdirname:
        model.save_attn_procs(_lowercase							,						safe_serialization=_lowercase     )
        self.assertTrue(os.path.isfile(os.path.join(_lowercase							,						'pytorch_lora_weights.safetensors'     )     )     )
        torch.manual_seed(0     )
        a	:     List[Any]							=  self.model_class(**_lowercase     )
        new_model.to(_lowercase     )
        new_model.load_attn_procs(_lowercase     )

      with torch.no_grad():
        a	:     Union[str, Any]							=  new_model(**_lowercase							,						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 							lowerCamelCase__ (				self      :   Dict     ):


      '''simple docstring'''




      a,	a	:     Union[str, Any]							=  self.prepare_init_args_and_inputs_for_common()

      a	:     Optional[Any]							=  8

      torch.manual_seed(0     )
      a	:     int							=  self.model_class(**_lowercase     )
      model.to(_lowercase     )

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





    def 							lowerCamelCase__ (				self      :   List[str]     ):


      '''simple docstring'''




      a,	a	:     str							=  self.prepare_init_args_and_inputs_for_common()

      a	:     Tuple							=  8

      torch.manual_seed(0     )
      a	:     int							=  self.model_class(**_lowercase     )
      model.to(_lowercase     )

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





    def 							lowerCamelCase__ (				self      :   str     ):


      '''simple docstring'''




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

      a	:     List[str]							=  8

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

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

      a	:     List[Any]							=  create_lora_layers(_lowercase     )
      model.set_attn_processor(_lowercase     )

      with torch.no_grad():
        a	:     Optional[Any]							=  model(**_lowercase							,						cross_attention_kwargs={'scale': 0.0}     ).sample

      model.set_attn_processor(AttnProcessor()     )

      with torch.no_grad():
        a	:     Optional[int]							=  model(**_lowercase     ).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 							lowerCamelCase__ (				self      :   str     ):


      '''simple docstring'''




      a,	a	:     Tuple							=  self.prepare_init_args_and_inputs_for_common()

      a	:     Any							=  4

      torch.manual_seed(0     )
      a	:     Optional[Any]							=  self.model_class(**_lowercase     )
      model.to(_lowercase     )
      a	:     str							=  create_lora_layers(_lowercase     )
      model.set_attn_processor(_lowercase     )

      # default
      with torch.no_grad():
        a	:     Optional[int]							=  model(**_lowercase     ).sample

        model.enable_xformers_memory_efficient_attention()
        a	:     Any							=  model(**_lowercase     ).sample

        model.disable_xformers_memory_efficient_attention()
        a	:     Optional[int]							=  model(**_lowercase     ).sample

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





    def 							lowerCamelCase__ (				self      :   Tuple     ):


      '''simple docstring'''




      a,	a	:     str							=  self.prepare_init_args_and_inputs_for_common()
      a	:     List[Any]							=  3_2

      a	:     Tuple							=  self.model_class(**_lowercase     )
      model.to(_lowercase     )
      model.eval()

      with torch.no_grad():
        a	:     List[str]							=  model(**_lowercase     )[0]

      model.enable_forward_chunking()
      with torch.no_grad():
        a	:     List[Any]							=  model(**_lowercase     )[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.









from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


UpperCamelCase_      			=  {
    "configuration_mobilenet_v1": [
        "MOBILENET_V1_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "MobileNetV1Config",
        "MobileNetV1OnnxConfig",
    ],
}

try:
							if not is_vision_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCamelCase_      			=  ["MobileNetV1FeatureExtractor"]
							UpperCamelCase_      			=  ["MobileNetV1ImageProcessor"]

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCamelCase_      			=  [
							    "MOBILENET_V1_PRETRAINED_MODEL_ARCHIVE_LIST",
							    "MobileNetV1ForImageClassification",
							    "MobileNetV1Model",
							    "MobileNetV1PreTrainedModel",
							    "load_tf_weights_in_mobilenet_v1",
							]


if TYPE_CHECKING:
							from .configuration_mobilenet_va import (
							    MOBILENET_V1_PRETRAINED_CONFIG_ARCHIVE_MAP,
							    MobileNetVaConfig,
							    MobileNetVaOnnxConfig,
							)

							try:
														if not is_vision_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .feature_extraction_mobilenet_va import MobileNetVaFeatureExtractor
														from .image_processing_mobilenet_va import MobileNetVaImageProcessor

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_mobilenet_va import (
														    MOBILENET_V1_PRETRAINED_MODEL_ARCHIVE_LIST,
														    MobileNetVaForImageClassification,
														    MobileNetVaModel,
														    MobileNetVaPreTrainedModel,
														    load_tf_weights_in_mobilenet_va,
														)


else:
							import sys

							UpperCamelCase_      			=  _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





