






"""simple docstring"""



import unittest

from transformers import ElectraConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import (
						    MODEL_FOR_PRETRAINING_MAPPING,
						    ElectraForCausalLM,
						    ElectraForMaskedLM,
						    ElectraForMultipleChoice,
						    ElectraForPreTraining,
						    ElectraForQuestionAnswering,
						    ElectraForSequenceClassification,
						    ElectraForTokenClassification,
						    ElectraModel,
						)
						from transformers.models.electra.modeling_electra import ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST






class 				__SCREAMING_SNAKE_CASE    :


							def __init__(      self   :    Any      ,     snake_case   :    Union[str, Any]      ,     snake_case   :    int=13      ,     snake_case   :    Dict=7      ,     snake_case   :    Tuple=True      ,     snake_case   :    List[str]=True      ,     snake_case   :    Tuple=True      ,     snake_case   :    List[str]=True      ,     snake_case   :    Dict=99      ,     snake_case   :    Dict=32      ,     snake_case   :    Union[str, Any]=5      ,     snake_case   :    int=4      ,     snake_case   :    Optional[Any]=37      ,     snake_case   :    Optional[int]="gelu"      ,     snake_case   :    List[str]=0.1      ,     snake_case   :    int=0.1      ,     snake_case   :    Tuple=512      ,     snake_case   :    Tuple=16      ,     snake_case   :    Any=2      ,     snake_case   :    int=0.02      ,     snake_case   :    Dict=3      ,     snake_case   :    Dict=4      ,     snake_case   :    List[str]=None      ,     ):
													'''simple docstring'''





													A__    :	List[Any]			=	parent
													A__    :	List[Any]			=	batch_size
													A__    :	Optional[int]			=	seq_length
													A__    :	Optional[Any]			=	is_training
													A__    :	List[str]			=	use_input_mask
													A__    :	Union[str, Any]			=	use_token_type_ids
													A__    :	Any			=	use_labels
													A__    :	int			=	vocab_size
													A__    :	int			=	hidden_size
													A__    :	Union[str, Any]			=	num_hidden_layers
													A__    :	Optional[int]			=	num_attention_heads
													A__    :	List[str]			=	intermediate_size
													A__    :	Dict			=	hidden_act
													A__    :	Optional[int]			=	hidden_dropout_prob
													A__    :	Any			=	attention_probs_dropout_prob
													A__    :	List[str]			=	max_position_embeddings
													A__    :	str			=	type_vocab_size
													A__    :	Optional[Any]			=	type_sequence_label_size
													A__    :	Tuple			=	initializer_range
													A__    :	int			=	num_labels
													A__    :	str			=	num_choices
													A__    :	Optional[Any]			=	scope


							def     _UpperCamelCase							(      self   :    Dict				):
													'''simple docstring'''





													A__    :	Optional[int]			=	ids_tensor([self.batch_size, self.seq_length]      ,     self.vocab_size				)

													A__    :	Optional[Any]			=	None
													if self.use_input_mask:
																			A__    :	Union[str, Any]			=	random_attention_mask([self.batch_size, self.seq_length]				)

													A__    :	Tuple			=	None
													if self.use_token_type_ids:
																			A__    :	List[Any]			=	ids_tensor([self.batch_size, self.seq_length]      ,     self.type_vocab_size				)

													A__    :	List[str]			=	None
													A__    :	Union[str, Any]			=	None
													A__    :	Union[str, Any]			=	None
													if self.use_labels:
																			A__    :	List[Any]			=	ids_tensor([self.batch_size]      ,     self.type_sequence_label_size				)
																			A__    :	List[Any]			=	ids_tensor([self.batch_size, self.seq_length]      ,     self.num_labels				)
																			A__    :	str			=	ids_tensor([self.batch_size]      ,     self.num_choices				)
																			A__    :	List[Any]			=	ids_tensor([self.batch_size, self.seq_length]      ,     1				)

													A__    :	Union[str, Any]			=	self.get_config()

													return (
													    config,
													    input_ids,
													    token_type_ids,
													    input_mask,
													    sequence_labels,
													    token_labels,
													    choice_labels,
													    fake_token_labels,
													)


							def     _UpperCamelCase							(      self   :    int				):
													'''simple docstring'''





													return ElectraConfig(
													    vocab_size=self.vocab_size      ,     hidden_size=self.hidden_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      ,     is_decoder=snake_case      ,     initializer_range=self.initializer_range      ,     )


							def     _UpperCamelCase							(      self   :    Dict				):
													'''simple docstring'''





													(
													    (
													    A__
													)	,	(
													    A__
													)	,	(
													    A__
													)	,	(
													    A__
													)	,	(
													    A__
													)	,	(
													    A__
													)	,	(
													    A__
													)	,	(
													    A__
													)	,	
													)    :	Union[str, Any]			=	self.prepare_config_and_inputs()

													A__    :	Optional[Any]			=	True
													A__    :	Optional[int]			=	floats_tensor([self.batch_size, self.seq_length, self.hidden_size]				)
													A__    :	Dict			=	ids_tensor([self.batch_size, self.seq_length]      ,     vocab_size=2				)

													return (
													    config,
													    input_ids,
													    token_type_ids,
													    input_mask,
													    sequence_labels,
													    token_labels,
													    choice_labels,
													    encoder_hidden_states,
													    encoder_attention_mask,
													)


							def     _UpperCamelCase							(      self   :    Tuple      ,     snake_case   :    str      ,     snake_case   :    str      ,     snake_case   :    int      ,     snake_case   :    Tuple      ,     snake_case   :    int      ,     snake_case   :    str      ,     snake_case   :    Tuple      ,     snake_case   :    Optional[int]      ,     ):
													'''simple docstring'''





													A__    :	Optional[int]			=	ElectraModel(config=snake_case				)
													model.to(snake_case				)
													model.eval()
													A__    :	int			=	model(snake_case      ,     attention_mask=snake_case      ,     token_type_ids=snake_case				)
													A__    :	List[str]			=	model(snake_case      ,     token_type_ids=snake_case				)
													A__    :	List[Any]			=	model(snake_case				)
													self.parent.assertEqual(result.last_hidden_state.shape      ,     (self.batch_size, self.seq_length, self.hidden_size)				)


							def     _UpperCamelCase							(      self   :    Optional[int]      ,     snake_case   :    List[str]      ,     snake_case   :    Optional[int]      ,     snake_case   :    Tuple      ,     snake_case   :    Dict      ,     snake_case   :    Any      ,     snake_case   :    Tuple      ,     snake_case   :    List[str]      ,     snake_case   :    str      ,     snake_case   :    Union[str, Any]      ,     ):
													'''simple docstring'''





													A__    :	Any			=	True
													A__    :	Tuple			=	ElectraModel(snake_case				)
													model.to(snake_case				)
													model.eval()
													A__    :	Any			=	model(
													    snake_case      ,     attention_mask=snake_case      ,     token_type_ids=snake_case      ,     encoder_hidden_states=snake_case      ,     encoder_attention_mask=snake_case      ,     )
													A__    :	Optional[Any]			=	model(
													    snake_case      ,     attention_mask=snake_case      ,     token_type_ids=snake_case      ,     encoder_hidden_states=snake_case      ,     )
													A__    :	Any			=	model(snake_case      ,     attention_mask=snake_case      ,     token_type_ids=snake_case				)
													self.parent.assertEqual(result.last_hidden_state.shape      ,     (self.batch_size, self.seq_length, self.hidden_size)				)


							def     _UpperCamelCase							(      self   :    str      ,     snake_case   :    Optional[Any]      ,     snake_case   :    List[Any]      ,     snake_case   :    List[Any]      ,     snake_case   :    Any      ,     snake_case   :    List[str]      ,     snake_case   :    List[Any]      ,     snake_case   :    Optional[Any]      ,     snake_case   :    Dict      ,     ):
													'''simple docstring'''





													A__    :	Optional[int]			=	ElectraForMaskedLM(config=snake_case				)
													model.to(snake_case				)
													model.eval()
													A__    :	List[str]			=	model(snake_case      ,     attention_mask=snake_case      ,     token_type_ids=snake_case      ,     labels=snake_case				)
													self.parent.assertEqual(result.logits.shape      ,     (self.batch_size, self.seq_length, self.vocab_size)				)


							def     _UpperCamelCase							(      self   :    Optional[Any]      ,     snake_case   :    Dict      ,     snake_case   :    Any      ,     snake_case   :    List[Any]      ,     snake_case   :    List[str]      ,     snake_case   :    List[str]      ,     snake_case   :    List[Any]      ,     snake_case   :    Any      ,     snake_case   :    int      ,     snake_case   :    str      ,     ):
													'''simple docstring'''





													A__    :	int			=	ElectraForCausalLM(config=snake_case				)
													model.to(snake_case				)
													model.eval()
													A__    :	Optional[Any]			=	model(snake_case      ,     attention_mask=snake_case      ,     token_type_ids=snake_case      ,     labels=snake_case				)
													self.parent.assertEqual(result.logits.shape      ,     (self.batch_size, self.seq_length, self.vocab_size)				)


							def     _UpperCamelCase							(      self   :    Optional[Any]      ,     snake_case   :    Union[str, Any]      ,     snake_case   :    Optional[int]      ,     snake_case   :    str      ,     snake_case   :    str      ,     snake_case   :    Optional[Any]      ,     snake_case   :    str      ,     snake_case   :    int      ,     snake_case   :    Union[str, Any]      ,     ):
													'''simple docstring'''





													A__    :	Tuple			=	self.num_labels
													A__    :	Optional[Any]			=	ElectraForTokenClassification(config=snake_case				)
													model.to(snake_case				)
													model.eval()
													A__    :	str			=	model(snake_case      ,     attention_mask=snake_case      ,     token_type_ids=snake_case      ,     labels=snake_case				)
													self.parent.assertEqual(result.logits.shape      ,     (self.batch_size, self.seq_length, self.num_labels)				)


							def     _UpperCamelCase							(      self   :    List[str]      ,     snake_case   :    Any      ,     snake_case   :    Union[str, Any]      ,     snake_case   :    str      ,     snake_case   :    Any      ,     snake_case   :    List[str]      ,     snake_case   :    Any      ,     snake_case   :    Union[str, Any]      ,     snake_case   :    Union[str, Any]      ,     ):
													'''simple docstring'''





													A__    :	Any			=	self.num_labels
													A__    :	int			=	ElectraForPreTraining(config=snake_case				)
													model.to(snake_case				)
													model.eval()
													A__    :	Union[str, Any]			=	model(snake_case      ,     attention_mask=snake_case      ,     token_type_ids=snake_case      ,     labels=snake_case				)
													self.parent.assertEqual(result.logits.shape      ,     (self.batch_size, self.seq_length)				)


							def     _UpperCamelCase							(      self   :    Dict      ,     snake_case   :    int      ,     snake_case   :    Optional[Any]      ,     snake_case   :    List[str]      ,     snake_case   :    Dict      ,     snake_case   :    Dict      ,     snake_case   :    Optional[int]      ,     snake_case   :    str      ,     snake_case   :    int      ,     ):
													'''simple docstring'''





													A__    :	List[Any]			=	self.num_labels
													A__    :	Optional[int]			=	ElectraForSequenceClassification(snake_case				)
													model.to(snake_case				)
													model.eval()
													A__    :	List[Any]			=	model(snake_case      ,     attention_mask=snake_case      ,     token_type_ids=snake_case      ,     labels=snake_case				)
													self.parent.assertEqual(result.logits.shape      ,     (self.batch_size, self.num_labels)				)


							def     _UpperCamelCase							(      self   :    List[Any]      ,     snake_case   :    List[Any]      ,     snake_case   :    Optional[Any]      ,     snake_case   :    Dict      ,     snake_case   :    Tuple      ,     snake_case   :    Any      ,     snake_case   :    Any      ,     snake_case   :    str      ,     snake_case   :    List[Any]      ,     ):
													'''simple docstring'''





													A__    :	Optional[Any]			=	ElectraForQuestionAnswering(config=snake_case				)
													model.to(snake_case				)
													model.eval()
													A__    :	Dict			=	model(
													    snake_case      ,     attention_mask=snake_case      ,     token_type_ids=snake_case      ,     start_positions=snake_case      ,     end_positions=snake_case      ,     )
													self.parent.assertEqual(result.start_logits.shape      ,     (self.batch_size, self.seq_length)				)
													self.parent.assertEqual(result.end_logits.shape      ,     (self.batch_size, self.seq_length)				)


							def     _UpperCamelCase							(      self   :    Tuple      ,     snake_case   :    Optional[int]      ,     snake_case   :    Tuple      ,     snake_case   :    Optional[Any]      ,     snake_case   :    Union[str, Any]      ,     snake_case   :    Tuple      ,     snake_case   :    Dict      ,     snake_case   :    Tuple      ,     snake_case   :    Optional[Any]      ,     ):
													'''simple docstring'''





													A__    :	Dict			=	self.num_choices
													A__    :	List[Any]			=	ElectraForMultipleChoice(config=snake_case				)
													model.to(snake_case				)
													model.eval()
													A__    :	str			=	input_ids.unsqueeze(1				).expand(-1      ,     self.num_choices      ,     -1				).contiguous()
													A__    :	Dict			=	token_type_ids.unsqueeze(1				).expand(-1      ,     self.num_choices      ,     -1				).contiguous()
													A__    :	Optional[int]			=	input_mask.unsqueeze(1				).expand(-1      ,     self.num_choices      ,     -1				).contiguous()
													A__    :	Optional[Any]			=	model(
													    snake_case      ,     attention_mask=snake_case      ,     token_type_ids=snake_case      ,     labels=snake_case      ,     )
													self.parent.assertEqual(result.logits.shape      ,     (self.batch_size, self.num_choices)				)


							def     _UpperCamelCase							(      self   :    List[str]				):
													'''simple docstring'''





													A__    :	Optional[int]			=	self.prepare_config_and_inputs()
													(
													    (
													    A__
													)	,	(
													    A__
													)	,	(
													    A__
													)	,	(
													    A__
													)	,	(
													    A__
													)	,	(
													    A__
													)	,	(
													    A__
													)	,	(
													    A__
													)	,	
													)    :	List[Any]			=	config_and_inputs
													A__    :	Dict			=	{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
													return config, inputs_dict






@require_torch
class 				__SCREAMING_SNAKE_CASE    (					UpperCamelCase							,		UpperCamelCase							,		unittest.TestCase		):
							snake_case_  =				(
							    (
							        ElectraModel,
							        ElectraForPreTraining,
							        ElectraForMaskedLM,
							        ElectraForCausalLM,
							        ElectraForMultipleChoice,
							        ElectraForTokenClassification,
							        ElectraForSequenceClassification,
							        ElectraForQuestionAnswering,
							    )
							    if is_torch_available()
							    else ()
							)
							snake_case_  =				(
							    {
							        'feature-extraction': ElectraModel,
							        'fill-mask': ElectraForMaskedLM,
							        'question-answering': ElectraForQuestionAnswering,
							        'text-classification': ElectraForSequenceClassification,
							        'text-generation': ElectraForCausalLM,
							        'token-classification': ElectraForTokenClassification,
							        'zero-shot': ElectraForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							snake_case_  =				True


							def     _UpperCamelCase							(      self   :    Tuple      ,     snake_case   :    List[str]      ,     snake_case   :    List[Any]      ,     snake_case   :    str=False				):
													'''simple docstring'''





													A__    :	Optional[int]			=	super()._prepare_for_class(snake_case      ,     snake_case      ,     return_labels=snake_case				)

													if return_labels:
																			if model_class in get_values(snake_case				):
																									A__    :	List[str]			=	torch.zeros(
																									    (self.model_tester.batch_size, self.model_tester.seq_length)      ,     dtype=torch.long      ,     device=snake_case				)
													return inputs_dict


							def     _UpperCamelCase							(      self   :    List[Any]				):
													'''simple docstring'''





													A__    :	Optional[Any]			=	ElectraModelTester(self				)
													A__    :	List[str]			=	ConfigTester(self      ,     config_class=snake_case      ,     hidden_size=37				)


							def     _UpperCamelCase							(      self   :    Any				):
													'''simple docstring'''





													self.config_tester.run_common_tests()


							def     _UpperCamelCase							(      self   :    Union[str, Any]				):
													'''simple docstring'''





													A__    :	Optional[int]			=	self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_electra_model(*snake_case				)


							def     _UpperCamelCase							(      self   :    Union[str, Any]				):
													'''simple docstring'''





													A__    :	Optional[int]			=	self.model_tester.prepare_config_and_inputs_for_decoder()
													self.model_tester.create_and_check_electra_model_as_decoder(*snake_case				)


							def     _UpperCamelCase							(      self   :    List[str]				):
													'''simple docstring'''





													A__    :	Optional[int]			=	self.model_tester.prepare_config_and_inputs()
													for type in ["absolute", "relative_key", "relative_key_query"]:
																			A__    :	Any			=	type
																			self.model_tester.create_and_check_electra_model(*snake_case				)


							def     _UpperCamelCase							(      self   :    List[Any]				):
													'''simple docstring'''





													A__    :	List[Any]			=	self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_electra_for_masked_lm(*snake_case				)


							def     _UpperCamelCase							(      self   :    Tuple				):
													'''simple docstring'''





													A__    :	List[str]			=	self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_electra_for_token_classification(*snake_case				)


							def     _UpperCamelCase							(      self   :    Optional[Any]				):
													'''simple docstring'''





													A__    :	Union[str, Any]			=	self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_electra_for_pretraining(*snake_case				)


							def     _UpperCamelCase							(      self   :    List[str]				):
													'''simple docstring'''





													A__    :	Union[str, Any]			=	self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_electra_for_sequence_classification(*snake_case				)


							def     _UpperCamelCase							(      self   :    Optional[Any]				):
													'''simple docstring'''





													A__    :	int			=	self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_electra_for_question_answering(*snake_case				)


							def     _UpperCamelCase							(      self   :    int				):
													'''simple docstring'''





													A__    :	Tuple			=	self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_electra_for_multiple_choice(*snake_case				)


							@slow
							def     _UpperCamelCase							(      self   :    Optional[Any]				):
													'''simple docstring'''





													for model_name in ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			A__    :	Optional[Any]			=	ElectraModel.from_pretrained(snake_case				)
																			self.assertIsNotNone(snake_case				)


							def     _UpperCamelCase							(      self   :    Tuple				):
													'''simple docstring'''





													A__    :	Any			=	self.model_tester.prepare_config_and_inputs_for_decoder()
													self.model_tester.create_and_check_electra_for_causal_lm(*snake_case				)






@require_torch
class 				__SCREAMING_SNAKE_CASE    (					unittest.TestCase		):


							@slow
							def     _UpperCamelCase							(      self   :    List[Any]				):
													'''simple docstring'''





													A__    :	Optional[int]			=	ElectraModel.from_pretrained("""google/electra-small-discriminator"""				)
													A__    :	str			=	torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]				)
													A__    :	str			=	torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]				)
													A__    :	str			=	model(snake_case      ,     attention_mask=snake_case				)[0]
													A__    :	str			=	torch.Size((1, 11, 256)				)
													self.assertEqual(output.shape      ,     snake_case				)
													A__    :	Optional[Any]			=	torch.tensor(
													    [[[0.4471, 0.6821, -0.3265], [0.4627, 0.5255, -0.3668], [0.4532, 0.3313, -0.4344]]]				)

													self.assertTrue(torch.allclose(output[:, 1:4, 1:4]      ,     snake_case      ,     atol=1e-4				)				)














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




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

												A__    :	Optional[Any]			=	LoRAAttnProcessor(hidden_size=UpperCAmelCase__,    cross_attention_dim=UpperCAmelCase__   )
												A__    :	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 				__SCREAMING_SNAKE_CASE    (					UpperCamelCase							,		UpperCamelCase							,		unittest.TestCase		):
							snake_case_  =				UNetaDConditionModel
							snake_case_  =				'sample'


							@property
							def     _UpperCamelCase							(      self   :    str				):
													'''simple docstring'''





													A__    :	Dict			=	4
													A__    :	str			=	4
													A__    :	str			=	4
													A__    :	Union[str, Any]			=	(32, 32)

													A__    :	List[Any]			=	floats_tensor((batch_size, num_channels, num_frames) + sizes				).to(snake_case				)
													A__    :	Any			=	torch.tensor([10]				).to(snake_case				)
													A__    :	Optional[Any]			=	floats_tensor((batch_size, 4, 32)				).to(snake_case				)

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


							@property
							def     _UpperCamelCase							(      self   :    Union[str, Any]				):
													'''simple docstring'''





													return (4, 4, 32, 32)


							@property
							def     _UpperCamelCase							(      self   :    Optional[int]				):
													'''simple docstring'''





													return (4, 4, 32, 32)


							def     _UpperCamelCase							(      self   :    int				):
													'''simple docstring'''





													A__    :	Tuple			=	{
													    """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__    :	Optional[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     _UpperCamelCase							(      self   :    int				):
													'''simple docstring'''





													A__	,	A__    :	Optional[Any]			=	self.prepare_init_args_and_inputs_for_common()
													A__    :	Dict			=	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     _UpperCamelCase							(      self   :    List[str]				):
													'''simple docstring'''





													A__	,	A__    :	Optional[int]			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	Tuple			=	32

													A__    :	List[Any]			=	self.model_class(**snake_case				)
													model.to(snake_case				)
													model.eval()

													with torch.no_grad():
																			A__    :	Dict			=	model(**snake_case				)

																			if isinstance(snake_case      ,     snake_case				):
																									A__    :	int			=	output.sample

													self.assertIsNotNone(snake_case				)
													A__    :	Any			=	inputs_dict["""sample"""].shape
													self.assertEqual(output.shape      ,     snake_case      ,     """Input and output shapes do not match"""				)


							def     _UpperCamelCase							(      self   :    int				):
													'''simple docstring'''





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

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

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

													A__    :	Optional[Any]			=	first.cpu().numpy()
													A__    :	Tuple			=	second.cpu().numpy()
													A__    :	List[Any]			=	out_a[~np.isnan(snake_case				)]
													A__    :	Dict			=	out_a[~np.isnan(snake_case				)]
													A__    :	str			=	np.amax(np.abs(out_a - out_a				)				)
													self.assertLessEqual(snake_case      ,     1e-5				)


							def     _UpperCamelCase							(      self   :    Any				):
													'''simple docstring'''





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

													A__    :	List[str]			=	8

													A__    :	Optional[Any]			=	self.model_class(**snake_case				)
													model.to(snake_case				)
													model.eval()

													model.set_attention_slice("""auto"""				)
													with torch.no_grad():
																			A__    :	List[str]			=	model(**snake_case				)
													assert output is not None

													model.set_attention_slice("""max"""				)
													with torch.no_grad():
																			A__    :	Union[str, Any]			=	model(**snake_case				)
													assert output is not None

													model.set_attention_slice(2				)
													with torch.no_grad():
																			A__    :	Dict			=	model(**snake_case				)
													assert output is not None


							def     _UpperCamelCase							(      self   :    str				):
													'''simple docstring'''





													A__	,	A__    :	Dict			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	List[Any]			=	8

													A__    :	Union[str, Any]			=	self.model_class(**snake_case				)
													model.to(snake_case				)

													with torch.no_grad():
																			A__    :	Optional[int]			=	model(**snake_case				).sample

													A__    :	Dict			=	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():
																			A__    :	int			=	model(**snake_case      ,     cross_attention_kwargs={"""scale""": 0.0}				).sample
																			A__    :	Optional[Any]			=	model(**snake_case      ,     cross_attention_kwargs={"""scale""": 0.5}				).sample
																			A__    :	List[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     _UpperCamelCase							(      self   :    str				):
													'''simple docstring'''





													A__	,	A__    :	Dict			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	Optional[Any]			=	8

													torch.manual_seed(0				)
													A__    :	str			=	self.model_class(**snake_case				)
													model.to(snake_case				)

													with torch.no_grad():
																			A__    :	Any			=	model(**snake_case				).sample

													A__    :	str			=	create_lora_layers(snake_case				)
													model.set_attn_processor(snake_case				)

													with torch.no_grad():
																			A__    :	Union[str, Any]			=	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				)
																			A__    :	Union[str, Any]			=	self.model_class(**snake_case				)
																			new_model.to(snake_case				)
																			new_model.load_attn_procs(snake_case				)

													with torch.no_grad():
																			A__    :	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     _UpperCamelCase							(      self   :    List[Any]				):
													'''simple docstring'''





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

													A__    :	Optional[Any]			=	8

													torch.manual_seed(0				)
													A__    :	Any			=	self.model_class(**snake_case				)
													model.to(snake_case				)

													with torch.no_grad():
																			A__    :	Tuple			=	model(**snake_case				).sample

													A__    :	Union[str, Any]			=	create_lora_layers(snake_case				)
													model.set_attn_processor(snake_case				)

													with torch.no_grad():
																			A__    :	Tuple			=	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				)
																			A__    :	Optional[Any]			=	self.model_class(**snake_case				)
																			new_model.to(snake_case				)
																			new_model.load_attn_procs(snake_case				)

													with torch.no_grad():
																			A__    :	Any			=	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     _UpperCamelCase							(      self   :    List[Any]				):
													'''simple docstring'''





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

													A__    :	Any			=	8

													torch.manual_seed(0				)
													A__    :	str			=	self.model_class(**snake_case				)
													model.to(snake_case				)

													A__    :	int			=	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				)
																			A__    :	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     _UpperCamelCase							(      self   :    Union[str, Any]				):
													'''simple docstring'''





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

													A__    :	Optional[Any]			=	8

													torch.manual_seed(0				)
													A__    :	Any			=	self.model_class(**snake_case				)
													model.to(snake_case				)

													A__    :	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				)
																			A__    :	Optional[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     _UpperCamelCase							(      self   :    List[str]				):
													'''simple docstring'''





													A__	,	A__    :	Dict			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	Tuple			=	8

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

													with torch.no_grad():
																			A__    :	Dict			=	model(**snake_case				).sample

													A__    :	Union[str, Any]			=	create_lora_layers(snake_case				)
													model.set_attn_processor(snake_case				)

													with torch.no_grad():
																			A__    :	Any			=	model(**snake_case      ,     cross_attention_kwargs={"""scale""": 0.0}				).sample

													model.set_attn_processor(AttnProcessor()				)

													with torch.no_grad():
																			A__    :	str			=	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     _UpperCamelCase							(      self   :    List[str]				):
													'''simple docstring'''





													A__	,	A__    :	List[str]			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	List[Any]			=	4

													torch.manual_seed(0				)
													A__    :	Tuple			=	self.model_class(**snake_case				)
													model.to(snake_case				)
													A__    :	List[str]			=	create_lora_layers(snake_case				)
													model.set_attn_processor(snake_case				)

													# default
													with torch.no_grad():
																			A__    :	Optional[Any]			=	model(**snake_case				).sample

																			model.enable_xformers_memory_efficient_attention()
																			A__    :	List[str]			=	model(**snake_case				).sample

																			model.disable_xformers_memory_efficient_attention()
																			A__    :	Union[str, Any]			=	model(**snake_case				).sample

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


							def     _UpperCamelCase							(      self   :    Any				):
													'''simple docstring'''





													A__	,	A__    :	Optional[Any]			=	self.prepare_init_args_and_inputs_for_common()
													A__    :	List[Any]			=	32

													A__    :	Any			=	self.model_class(**snake_case				)
													model.to(snake_case				)
													model.eval()

													with torch.no_grad():
																			A__    :	List[Any]			=	model(**snake_case				)[0]

													model.enable_forward_chunking()
													with torch.no_grad():
																			A__    :	Union[str, Any]			=	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.






