




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



def   lowerCAmelCase						(						_lowerCAmelCase							:				List[str]  ,							_lowerCAmelCase							:				bool = True		):




 """simple docstring"""







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

  UpperCAmelCase__								=  LoRAAttnProcessor(hidden_size=_lowerCAmelCase  ,							cross_attention_dim=_lowerCAmelCase		)
  UpperCAmelCase__								=  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      _UpperCamelCase (			lowerCAmelCase		,    lowerCAmelCase		,    unittest.TestCase						):
      UpperCAmelCase_            =						UNetaDConditionModel
      UpperCAmelCase_            =						"""sample"""
      @property
      def 		UpperCAmelCase_  (     self				:str )       ->     Optional[Any]:
       UpperCAmelCase__								=  4
       UpperCAmelCase__								=  4
       UpperCAmelCase__								=  4
       UpperCAmelCase__								=  (32, 32)

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

       return {"sample": noise, "timestep": time_step, "encoder_hidden_states": encoder_hidden_states}
      @property
      def 		UpperCAmelCase_  (     self				:List[str] )       ->     Optional[Any]:
       return (4, 4, 32, 32)
      @property
      def 		UpperCAmelCase_  (     self				:Optional[Any] )       ->     List[Any]:
       return (4, 4, 32, 32)
      def 		UpperCAmelCase_  (     self				:List[Any] )       ->     Dict:
       UpperCAmelCase__								=  {
           "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,
       }
       UpperCAmelCase__								=  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 		UpperCAmelCase_  (     self				:Tuple )       ->     int:
       UpperCAmelCase__	,	UpperCAmelCase__								=  self.prepare_init_args_and_inputs_for_common()
       UpperCAmelCase__								=  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 		UpperCAmelCase_  (     self				:Tuple )       ->     Optional[int]:
       UpperCAmelCase__	,	UpperCAmelCase__								=  self.prepare_init_args_and_inputs_for_common()

       UpperCAmelCase__								=  32

       UpperCAmelCase__								=  self.model_class(**lowerCamelCase )
       model.to(lowerCamelCase )
       model.eval()

       with torch.no_grad():
        UpperCAmelCase__								=  model(**lowerCamelCase )

        if isinstance(lowerCamelCase	, lowerCamelCase ):
         UpperCAmelCase__								=  output.sample

       self.assertIsNotNone(lowerCamelCase )
       UpperCAmelCase__								=  inputs_dict["sample"].shape
       self.assertEqual(output.shape	, lowerCamelCase	, "Input and output shapes do not match" )
      def 		UpperCAmelCase_  (     self				:Tuple )       ->     Union[str, Any]:
       UpperCAmelCase__	,	UpperCAmelCase__								=  self.prepare_init_args_and_inputs_for_common()
       UpperCAmelCase__								=  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 )

        UpperCAmelCase__								=  model(**lowerCamelCase )
        if isinstance(lowerCamelCase	, lowerCamelCase ):
         UpperCAmelCase__								=  first.sample

        UpperCAmelCase__								=  model(**lowerCamelCase )
        if isinstance(lowerCamelCase	, lowerCamelCase ):
         UpperCAmelCase__								=  second.sample

       UpperCAmelCase__								=  first.cpu().numpy()
       UpperCAmelCase__								=  second.cpu().numpy()
       UpperCAmelCase__								=  out_a[~np.isnan(lowerCamelCase )]
       UpperCAmelCase__								=  out_a[~np.isnan(lowerCamelCase )]
       UpperCAmelCase__								=  np.amax(np.abs(out_a - out_a ) )
       self.assertLessEqual(lowerCamelCase	, 1e-5 )
      def 		UpperCAmelCase_  (     self				:int )       ->     Optional[Any]:
       UpperCAmelCase__	,	UpperCAmelCase__								=  self.prepare_init_args_and_inputs_for_common()

       UpperCAmelCase__								=  8

       UpperCAmelCase__								=  self.model_class(**lowerCamelCase )
       model.to(lowerCamelCase )
       model.eval()

       model.set_attention_slice("auto" )
       with torch.no_grad():
        UpperCAmelCase__								=  model(**lowerCamelCase )
       assert output is not None

       model.set_attention_slice("max" )
       with torch.no_grad():
        UpperCAmelCase__								=  model(**lowerCamelCase )
       assert output is not None

       model.set_attention_slice(2 )
       with torch.no_grad():
        UpperCAmelCase__								=  model(**lowerCamelCase )
       assert output is not None
      def 		UpperCAmelCase_  (     self				:Any )       ->     List[str]:
       UpperCAmelCase__	,	UpperCAmelCase__								=  self.prepare_init_args_and_inputs_for_common()

       UpperCAmelCase__								=  8

       UpperCAmelCase__								=  self.model_class(**lowerCamelCase )
       model.to(lowerCamelCase )

       with torch.no_grad():
        UpperCAmelCase__								=  model(**lowerCamelCase ).sample

       UpperCAmelCase__								=  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():
        UpperCAmelCase__								=  model(**lowerCamelCase	, cross_attention_kwargs={"scale": 0.0} ).sample
        UpperCAmelCase__								=  model(**lowerCamelCase	, cross_attention_kwargs={"scale": 0.5} ).sample
        UpperCAmelCase__								=  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 		UpperCAmelCase_  (     self				:Union[str, Any] )       ->     List[str]:
       UpperCAmelCase__	,	UpperCAmelCase__								=  self.prepare_init_args_and_inputs_for_common()

       UpperCAmelCase__								=  8

       torch.manual_seed(0 )
       UpperCAmelCase__								=  self.model_class(**lowerCamelCase )
       model.to(lowerCamelCase )

       with torch.no_grad():
        UpperCAmelCase__								=  model(**lowerCamelCase ).sample

       UpperCAmelCase__								=  create_lora_layers(lowerCamelCase )
       model.set_attn_processor(lowerCamelCase )

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

       with torch.no_grad():
        UpperCAmelCase__								=  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 		UpperCAmelCase_  (     self				:Optional[int] )       ->     int:
       UpperCAmelCase__	,	UpperCAmelCase__								=  self.prepare_init_args_and_inputs_for_common()

       UpperCAmelCase__								=  8

       torch.manual_seed(0 )
       UpperCAmelCase__								=  self.model_class(**lowerCamelCase )
       model.to(lowerCamelCase )

       with torch.no_grad():
        UpperCAmelCase__								=  model(**lowerCamelCase ).sample

       UpperCAmelCase__								=  create_lora_layers(lowerCamelCase )
       model.set_attn_processor(lowerCamelCase )

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

       with torch.no_grad():
        UpperCAmelCase__								=  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 		UpperCAmelCase_  (     self				:int )       ->     int:
       # enable deterministic behavior for gradient checkpointing
       UpperCAmelCase__	,	UpperCAmelCase__								=  self.prepare_init_args_and_inputs_for_common()

       UpperCAmelCase__								=  8

       torch.manual_seed(0 )
       UpperCAmelCase__								=  self.model_class(**lowerCamelCase )
       model.to(lowerCamelCase )

       UpperCAmelCase__								=  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 )
        UpperCAmelCase__								=  self.model_class(**lowerCamelCase )
        new_model.to(lowerCamelCase )
        new_model.load_attn_procs(lowerCamelCase	, weight_name="pytorch_lora_weights.bin" )
      def 		UpperCAmelCase_  (     self				:int )       ->     Any:
       UpperCAmelCase__	,	UpperCAmelCase__								=  self.prepare_init_args_and_inputs_for_common()

       UpperCAmelCase__								=  8

       torch.manual_seed(0 )
       UpperCAmelCase__								=  self.model_class(**lowerCamelCase )
       model.to(lowerCamelCase )

       UpperCAmelCase__								=  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 )
        UpperCAmelCase__								=  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 		UpperCAmelCase_  (     self				:List[Any] )       ->     Dict:
       UpperCAmelCase__	,	UpperCAmelCase__								=  self.prepare_init_args_and_inputs_for_common()

       UpperCAmelCase__								=  8

       torch.manual_seed(0 )
       UpperCAmelCase__								=  self.model_class(**lowerCamelCase )
       model.to(lowerCamelCase )

       with torch.no_grad():
        UpperCAmelCase__								=  model(**lowerCamelCase ).sample

       UpperCAmelCase__								=  create_lora_layers(lowerCamelCase )
       model.set_attn_processor(lowerCamelCase )

       with torch.no_grad():
        UpperCAmelCase__								=  model(**lowerCamelCase	, cross_attention_kwargs={"scale": 0.0} ).sample

       model.set_attn_processor(AttnProcessor() )

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

       UpperCAmelCase__								=  4

       torch.manual_seed(0 )
       UpperCAmelCase__								=  self.model_class(**lowerCamelCase )
       model.to(lowerCamelCase )
       UpperCAmelCase__								=  create_lora_layers(lowerCamelCase )
       model.set_attn_processor(lowerCamelCase )

       # default
       with torch.no_grad():
        UpperCAmelCase__								=  model(**lowerCamelCase ).sample

        model.enable_xformers_memory_efficient_attention()
        UpperCAmelCase__								=  model(**lowerCamelCase ).sample

        model.disable_xformers_memory_efficient_attention()
        UpperCAmelCase__								=  model(**lowerCamelCase ).sample

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







      def 		UpperCAmelCase_  (     self				:Tuple )       ->     str:
       UpperCAmelCase__	,	UpperCAmelCase__								=  self.prepare_init_args_and_inputs_for_common()
       UpperCAmelCase__								=  32

       UpperCAmelCase__								=  self.model_class(**lowerCamelCase )
       model.to(lowerCamelCase )
       model.eval()

       with torch.no_grad():
        UpperCAmelCase__								=  model(**lowerCamelCase )[0]

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











from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


_lowerCAmelCase  :				Dict						=				logging.get_logger(__name__)

_lowerCAmelCase  :				Optional[int]						=				{
    "google/rembert": "https://huggingface.co/google/rembert/resolve/main/config.json",
    # See all RemBERT models at https://huggingface.co/models?filter=rembert
}
class      _UpperCamelCase (			lowerCAmelCase						):
      UpperCAmelCase_            =						"""rembert"""
      def __init__(     self				:Union[str, Any]	, lowerCamelCase				:Union[str, Any]=25_0300	, lowerCamelCase				:Optional[Any]=1152	, lowerCamelCase				:Optional[int]=32	, lowerCamelCase				:Union[str, Any]=18	, lowerCamelCase				:str=256	, lowerCamelCase				:Any=1664	, lowerCamelCase				:Any=4608	, lowerCamelCase				:List[str]="gelu"	, lowerCamelCase				:Optional[Any]=0.0	, lowerCamelCase				:Optional[int]=0.0	, lowerCamelCase				:Dict=0.1	, lowerCamelCase				:List[str]=512	, lowerCamelCase				:int=2	, lowerCamelCase				:str=0.02	, lowerCamelCase				:Optional[int]=1e-12	, lowerCamelCase				:List[Any]=True	, lowerCamelCase				:Optional[Any]=0	, lowerCamelCase				:Optional[Any]=312	, lowerCamelCase				:int=313	, **lowerCamelCase				:Dict	, )       ->     Any:
       super().__init__(pad_token_id=lowerCamelCase	, bos_token_id=lowerCamelCase	, eos_token_id=lowerCamelCase	, **lowerCamelCase )

       UpperCAmelCase__								=  vocab_size
       UpperCAmelCase__								=  input_embedding_size
       UpperCAmelCase__								=  output_embedding_size
       UpperCAmelCase__								=  max_position_embeddings
       UpperCAmelCase__								=  hidden_size
       UpperCAmelCase__								=  num_hidden_layers
       UpperCAmelCase__								=  num_attention_heads
       UpperCAmelCase__								=  intermediate_size
       UpperCAmelCase__								=  hidden_act
       UpperCAmelCase__								=  hidden_dropout_prob
       UpperCAmelCase__								=  attention_probs_dropout_prob
       UpperCAmelCase__								=  classifier_dropout_prob
       UpperCAmelCase__								=  initializer_range
       UpperCAmelCase__								=  type_vocab_size
       UpperCAmelCase__								=  layer_norm_eps
       UpperCAmelCase__								=  use_cache
       UpperCAmelCase__								=  False



class      _UpperCamelCase (			lowerCAmelCase						):
      @property
      def 		UpperCAmelCase_  (     self				:List[Any] )       ->     Mapping[str, Mapping[int, str]]:
       if self.task == "multiple-choice":
        UpperCAmelCase__								=  {0: "batch", 1: "choice", 2: "sequence"}
       else:
        UpperCAmelCase__								=  {0: "batch", 1: "sequence"}
       return OrderedDict(
           [
               ("input_ids", dynamic_axis),
               ("attention_mask", dynamic_axis),
               ("token_type_ids", dynamic_axis),
           ] )







      @property
      def 		UpperCAmelCase_  (     self				:Union[str, Any] )       ->     float:
       return 1e-4





