


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

__SCREAMING_SNAKE_CASE     :   Tuple					=						logging.get_logger(__name__)



def 						__a      (  a,		a = True					):





       """simple docstring"""







       _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=_A,		cross_attention_dim=_A					)
              _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 __snake_case		(					_UpperCAmelCase				,       _UpperCAmelCase				,       unittest.TestCase      ):







   """simple docstring"""




   lowerCAmelCase_				:   Optional[int]     	=			UNetaDConditionModel
   lowerCAmelCase_				:   Optional[int]     	=			'sample'




   @property
   def SCREAMING_SNAKE_CASE_      (		self							:Tuple       ):
          _a          =				4
          _a          =				4
          _a          =				4
          _a          =				(32, 32)

          _a          =				floats_tensor((batch_size, num_channels, num_frames) + sizes       ).to(UpperCamelCase__       )
          _a          =				torch.tensor([10]       ).to(UpperCamelCase__       )
          _a          =				floats_tensor((batch_size, 4, 32)       ).to(UpperCamelCase__       )

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




   @property
   def SCREAMING_SNAKE_CASE_      (		self							:str       ):
          return (4, 4, 32, 32)




   @property
   def SCREAMING_SNAKE_CASE_      (		self							:Dict       ):
          return (4, 4, 32, 32)




   def SCREAMING_SNAKE_CASE_      (		self							:Any       ):
          _a          =				{
              "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          =				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							:Tuple       ):
          _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							:List[str]       ):
          _a   , _a          =				self.prepare_init_args_and_inputs_for_common()

          _a          =				32

          _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							:str       ):
          _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							:str       ):
          _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							: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							: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.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							: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__       )
                 new_model.load_attn_procs(UpperCamelCase__				,  weight_name="pytorch_lora_weights.bin"       )




   def SCREAMING_SNAKE_CASE_      (		self							: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__       )

          _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							: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.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							:Optional[Any]       ):
          _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							:Tuple       ):
          _a   , _a          =				self.prepare_init_args_and_inputs_for_common()
          _a          =				32

          _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.




"""simple docstring"""




from ...configuration_utils import PretrainedConfig
from ...utils import logging


__SCREAMING_SNAKE_CASE					=						logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE					=						{
    """bert-base-uncased""": """https://huggingface.co/bert-base-uncased/resolve/main/config.json""",
    # QDQBERT models can be loaded from any BERT checkpoint, available at https://huggingface.co/models?filter=bert
}





class __snake_case		(					_SCREAMING_SNAKE_CASE      ):







   """simple docstring"""




   lowerCAmelCase_				:   Tuple     	=			'qdqbert'




   def __init__(		self							:List[Any]				,  UpperCamelCase__							:List[Any]=30_522				,  UpperCamelCase__							:Optional[int]=768				,  UpperCamelCase__							:str=12				,  UpperCamelCase__							:Tuple=12				,  UpperCamelCase__							:Any=3_072				,  UpperCamelCase__							:Union[str, Any]="gelu"				,  UpperCamelCase__							:Dict=0.1				,  UpperCamelCase__							:Any=0.1				,  UpperCamelCase__							:Dict=512				,  UpperCamelCase__							:Optional[int]=2				,  UpperCamelCase__							:Any=0.02				,  UpperCamelCase__							:Optional[int]=1E-12				,  UpperCamelCase__							:str=True				,  UpperCamelCase__							:str=1				,  UpperCamelCase__							:List[Any]=0				,  UpperCamelCase__							:List[str]=2				,  **UpperCamelCase__							:Optional[Any]				,  ):
          super().__init__(pad_token_id=UpperCamelCase__				,  bos_token_id=UpperCamelCase__				,  eos_token_id=UpperCamelCase__				,  **UpperCamelCase__       )

          _a          =				vocab_size
          _a          =				max_position_embeddings
          _a          =				hidden_size
          _a          =				num_hidden_layers
          _a          =				num_attention_heads
          _a          =				intermediate_size
          _a          =				hidden_act
          _a          =				hidden_dropout_prob
          _a          =				attention_probs_dropout_prob
          _a          =				initializer_range
          _a          =				type_vocab_size
          _a          =				layer_norm_eps
          _a          =				use_cache
