





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

A__       :     str          =logging.get_logger(__name__)
def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase = True						):





						"""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 				UpperCAmelCase  (    snake_case_							,	snake_case_							,	unittest.TestCase     ):
						_lowercase:    int		  =     UNetaDConditionModel
						_lowercase:    Optional[int]		  =     '''sample'''




						@property
						def        lowercase__					(    self    :				Any						)       ->   int:
												_lowerCAmelCase							  =							4
												_lowerCAmelCase							  =							4
												_lowerCAmelCase							  =							4
												_lowerCAmelCase							  =							(32, 32)

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

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




						@property
						def        lowercase__					(    self    :				Dict						)       ->   Any:
												return (4, 4, 32, 32)




						@property
						def        lowercase__					(    self    :				List[str]						)       ->   Optional[int]:
												return (4, 4, 32, 32)




						def        lowercase__					(    self    :				Optional[Any]						)       ->   Optional[Any]:
												_lowerCAmelCase							  =							{
												    """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,
												}
												_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        lowercase__					(    self    :				List[str]						)       ->   int:
												_lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()
												_lowerCAmelCase							  =							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        lowercase__					(    self    :				str						)       ->   int:
												_lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

												_lowerCAmelCase							  =							32

												_lowerCAmelCase							  =							self.model_class(**__snake_case						)
												model.to(__snake_case						)
												model.eval()

												with torch.no_grad():
																		_lowerCAmelCase							  =							model(**__snake_case						)

																		if isinstance(__snake_case    ,    __snake_case						):
																								_lowerCAmelCase							  =							output.sample

												self.assertIsNotNone(__snake_case						)
												_lowerCAmelCase							  =							inputs_dict["""sample"""].shape
												self.assertEqual(output.shape    ,    __snake_case    ,    """Input and output shapes do not match"""						)




						def        lowercase__					(    self    :				Any						)       ->   Any:
												_lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()
												_lowerCAmelCase							  =							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						)

																		_lowerCAmelCase							  =							model(**__snake_case						)
																		if isinstance(__snake_case    ,    __snake_case						):
																								_lowerCAmelCase							  =							first.sample

																		_lowerCAmelCase							  =							model(**__snake_case						)
																		if isinstance(__snake_case    ,    __snake_case						):
																								_lowerCAmelCase							  =							second.sample

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




						def        lowercase__					(    self    :				Any						)       ->   List[str]:
												_lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

												_lowerCAmelCase							  =							8

												_lowerCAmelCase							  =							self.model_class(**__snake_case						)
												model.to(__snake_case						)
												model.eval()

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

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

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




						def        lowercase__					(    self    :				str						)       ->   Any:
												_lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

												_lowerCAmelCase							  =							8

												_lowerCAmelCase							  =							self.model_class(**__snake_case						)
												model.to(__snake_case						)

												with torch.no_grad():
																		_lowerCAmelCase							  =							model(**__snake_case						).sample

												_lowerCAmelCase							  =							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():
																		_lowerCAmelCase							  =							model(**__snake_case    ,    cross_attention_kwargs={"""scale""": 0.0}						).sample
																		_lowerCAmelCase							  =							model(**__snake_case    ,    cross_attention_kwargs={"""scale""": 0.5}						).sample
																		_lowerCAmelCase							  =							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        lowercase__					(    self    :				Tuple						)       ->   int:
												_lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

												_lowerCAmelCase							  =							8

												torch.manual_seed(0						)
												_lowerCAmelCase							  =							self.model_class(**__snake_case						)
												model.to(__snake_case						)

												with torch.no_grad():
																		_lowerCAmelCase							  =							model(**__snake_case						).sample

												_lowerCAmelCase							  =							create_lora_layers(__snake_case						)
												model.set_attn_processor(__snake_case						)

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

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

												_lowerCAmelCase							  =							8

												torch.manual_seed(0						)
												_lowerCAmelCase							  =							self.model_class(**__snake_case						)
												model.to(__snake_case						)

												with torch.no_grad():
																		_lowerCAmelCase							  =							model(**__snake_case						).sample

												_lowerCAmelCase							  =							create_lora_layers(__snake_case						)
												model.set_attn_processor(__snake_case						)

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

												with torch.no_grad():
																		_lowerCAmelCase							  =							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        lowercase__					(    self    :				Dict						)       ->   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(**__snake_case						)
												model.to(__snake_case						)

												_lowerCAmelCase							  =							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						)
																		_lowerCAmelCase							  =							self.model_class(**__snake_case						)
																		new_model.to(__snake_case						)
																		new_model.load_attn_procs(__snake_case    ,    weight_name="""pytorch_lora_weights.bin"""						)




						def        lowercase__					(    self    :				Optional[int]						)       ->   List[Any]:
												_lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

												_lowerCAmelCase							  =							8

												torch.manual_seed(0						)
												_lowerCAmelCase							  =							self.model_class(**__snake_case						)
												model.to(__snake_case						)

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

												_lowerCAmelCase							  =							8

												torch.manual_seed(0						)
												_lowerCAmelCase							  =							self.model_class(**__snake_case						)
												model.to(__snake_case						)

												with torch.no_grad():
																		_lowerCAmelCase							  =							model(**__snake_case						).sample

												_lowerCAmelCase							  =							create_lora_layers(__snake_case						)
												model.set_attn_processor(__snake_case						)

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

												model.set_attn_processor(AttnProcessor()						)

												with torch.no_grad():
																		_lowerCAmelCase							  =							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        lowercase__					(    self    :				Union[str, Any]						)       ->   Tuple:
												# 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(**__snake_case						)
												model.to(__snake_case						)
												_lowerCAmelCase							  =							create_lora_layers(__snake_case						)
												model.set_attn_processor(__snake_case						)

												# default
												with torch.no_grad():
																		_lowerCAmelCase							  =							model(**__snake_case						).sample

																		model.enable_xformers_memory_efficient_attention()
																		_lowerCAmelCase							  =							model(**__snake_case						).sample

																		model.disable_xformers_memory_efficient_attention()
																		_lowerCAmelCase							  =							model(**__snake_case						).sample

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




						def        lowercase__					(    self    :				Dict						)       ->   Optional[int]:
												_lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()
												_lowerCAmelCase							  =							32

												_lowerCAmelCase							  =							self.model_class(**__snake_case						)
												model.to(__snake_case						)
												model.eval()

												with torch.no_grad():
																		_lowerCAmelCase							  =							model(**__snake_case						)[0]

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













'''simple docstring'''






import unittest

import numpy as np

from transformers import ElectraConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
      from transformers.models.electra.modeling_flax_electra import (
          FlaxElectraForCausalLM,
          FlaxElectraForMaskedLM,
          FlaxElectraForMultipleChoice,
          FlaxElectraForPreTraining,
          FlaxElectraForQuestionAnswering,
          FlaxElectraForSequenceClassification,
          FlaxElectraForTokenClassification,
          FlaxElectraModel,
      )




class 				UpperCAmelCase  (    unittest.TestCase     ):




      def __init__(    self    :				str    ,    __snake_case    :				str    ,    __snake_case    :				str=13    ,    __snake_case    :				Optional[Any]=7    ,    __snake_case    :				str=True    ,    __snake_case    :				Any=True    ,    __snake_case    :				List[str]=True    ,    __snake_case    :				Any=True    ,    __snake_case    :				List[str]=99    ,    __snake_case    :				Dict=24    ,    __snake_case    :				Any=32    ,    __snake_case    :				List[str]=5    ,    __snake_case    :				Any=4    ,    __snake_case    :				Optional[Any]=37    ,    __snake_case    :				Union[str, Any]="gelu"    ,    __snake_case    :				Dict=0.1    ,    __snake_case    :				int=0.1    ,    __snake_case    :				Dict=5_12    ,    __snake_case    :				List[Any]=16    ,    __snake_case    :				Dict=2    ,    __snake_case    :				str=0.02    ,    __snake_case    :				Any=4    ,    )       ->   Union[str, Any]:
            _lowerCAmelCase							  =							parent
            _lowerCAmelCase							  =							batch_size
            _lowerCAmelCase							  =							seq_length
            _lowerCAmelCase							  =							is_training
            _lowerCAmelCase							  =							use_attention_mask
            _lowerCAmelCase							  =							use_token_type_ids
            _lowerCAmelCase							  =							use_labels
            _lowerCAmelCase							  =							vocab_size
            _lowerCAmelCase							  =							hidden_size
            _lowerCAmelCase							  =							embedding_size
            _lowerCAmelCase							  =							num_hidden_layers
            _lowerCAmelCase							  =							num_attention_heads
            _lowerCAmelCase							  =							intermediate_size
            _lowerCAmelCase							  =							hidden_act
            _lowerCAmelCase							  =							hidden_dropout_prob
            _lowerCAmelCase							  =							attention_probs_dropout_prob
            _lowerCAmelCase							  =							max_position_embeddings
            _lowerCAmelCase							  =							type_vocab_size
            _lowerCAmelCase							  =							type_sequence_label_size
            _lowerCAmelCase							  =							initializer_range
            _lowerCAmelCase							  =							num_choices




      def        lowercase__					(    self    :				Optional[Any]						)       ->   Optional[Any]:
            _lowerCAmelCase							  =							ids_tensor([self.batch_size, self.seq_length]    ,    self.vocab_size						)

            _lowerCAmelCase							  =							None
            if self.use_attention_mask:
                  _lowerCAmelCase							  =							random_attention_mask([self.batch_size, self.seq_length]						)

            _lowerCAmelCase							  =							None
            if self.use_token_type_ids:
                  _lowerCAmelCase							  =							ids_tensor([self.batch_size, self.seq_length]    ,    self.type_vocab_size						)

            _lowerCAmelCase							  =							ElectraConfig(
                vocab_size=self.vocab_size    ,    hidden_size=self.hidden_size    ,    embedding_size=self.embedding_size    ,    num_hidden_layers=self.num_hidden_layers    ,    num_attention_heads=self.num_attention_heads    ,    intermediate_size=self.intermediate_size    ,    hidden_act=self.hidden_act    ,    hidden_dropout_prob=self.hidden_dropout_prob    ,    attention_probs_dropout_prob=self.attention_probs_dropout_prob    ,    max_position_embeddings=self.max_position_embeddings    ,    type_vocab_size=self.type_vocab_size    ,    initializer_range=self.initializer_range    ,    )

            return config, input_ids, token_type_ids, attention_mask




      def        lowercase__					(    self    :				Any						)       ->   Optional[int]:
            _lowerCAmelCase							  =							self.prepare_config_and_inputs()
            _lowerCAmelCase  ,							_lowerCAmelCase  ,							_lowerCAmelCase  ,							_lowerCAmelCase							  =							config_and_inputs
            _lowerCAmelCase							  =							{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
            return config, inputs_dict




@require_flax
class 				UpperCAmelCase  (    snake_case_							,	unittest.TestCase     ):
      _lowercase:    str		  =     True

      _lowercase:    Optional[Any]		  =     (
          (
              FlaxElectraModel,
              FlaxElectraForCausalLM,
              FlaxElectraForMaskedLM,
              FlaxElectraForPreTraining,
              FlaxElectraForTokenClassification,
              FlaxElectraForQuestionAnswering,
              FlaxElectraForMultipleChoice,
              FlaxElectraForSequenceClassification,
          )
          if is_flax_available()
          else ()
      )




      def        lowercase__					(    self    :				Tuple						)       ->   Optional[int]:
            _lowerCAmelCase							  =							FlaxElectraModelTester(self						)




      @slow
      def        lowercase__					(    self    :				List[Any]						)       ->   str:
            for model_class_name in self.all_model_classes:
                  if model_class_name == FlaxElectraForMaskedLM:
                        _lowerCAmelCase							  =							model_class_name.from_pretrained("""google/electra-small-generator"""						)
                  else:
                        _lowerCAmelCase							  =							model_class_name.from_pretrained("""google/electra-small-discriminator"""						)
                  _lowerCAmelCase							  =							model(np.ones((1, 1)						)						)
                  self.assertIsNotNone(__snake_case						)






