






'''simple docstring'''





import unittest

from transformers import LayoutLMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

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


if is_torch_available():
						import torch

						from transformers import (
						    LayoutLMForMaskedLM,
						    LayoutLMForQuestionAnswering,
						    LayoutLMForSequenceClassification,
						    LayoutLMForTokenClassification,
						    LayoutLMModel,
						)


class 				snake_case__     :




					def __init__(       self       :						Optional[int]							,				_A       :						List[str]							,				_A       :						Union[str, Any]=13							,				_A       :						Union[str, Any]=7							,				_A       :						Tuple=True							,				_A       :						Dict=True							,				_A       :						Optional[int]=True							,				_A       :						List[Any]=True							,				_A       :						Optional[int]=99							,				_A       :						List[str]=32							,				_A       :						List[Any]=5							,				_A       :						Dict=4							,				_A       :						Union[str, Any]=37							,				_A       :						Dict="gelu"							,				_A       :						Optional[Any]=0.1							,				_A       :						Dict=0.1							,				_A       :						str=5_12							,				_A       :						Tuple=16							,				_A       :						str=2							,				_A       :						str=0.02							,				_A       :						Tuple=3							,				_A       :						Optional[Any]=4							,				_A       :						str=None							,				_A       :						Tuple=10_00							,				)  ->    Optional[int]:
								UpperCAmelCase_     :   Union[str, Any]       =   parent
								UpperCAmelCase_     :   Tuple       =   batch_size
								UpperCAmelCase_     :   List[str]       =   seq_length
								UpperCAmelCase_     :   Union[str, Any]       =   is_training
								UpperCAmelCase_     :   Optional[Any]       =   use_input_mask
								UpperCAmelCase_     :   List[str]       =   use_token_type_ids
								UpperCAmelCase_     :   int       =   use_labels
								UpperCAmelCase_     :   Dict       =   vocab_size
								UpperCAmelCase_     :   List[str]       =   hidden_size
								UpperCAmelCase_     :   str       =   num_hidden_layers
								UpperCAmelCase_     :   str       =   num_attention_heads
								UpperCAmelCase_     :   Dict       =   intermediate_size
								UpperCAmelCase_     :   Union[str, Any]       =   hidden_act
								UpperCAmelCase_     :   Any       =   hidden_dropout_prob
								UpperCAmelCase_     :   Union[str, Any]       =   attention_probs_dropout_prob
								UpperCAmelCase_     :   Optional[int]       =   max_position_embeddings
								UpperCAmelCase_     :   Any       =   type_vocab_size
								UpperCAmelCase_     :   Tuple       =   type_sequence_label_size
								UpperCAmelCase_     :   Optional[Any]       =   initializer_range
								UpperCAmelCase_     :   Union[str, Any]       =   num_labels
								UpperCAmelCase_     :   Tuple       =   num_choices
								UpperCAmelCase_     :   int       =   scope
								UpperCAmelCase_     :   int       =   range_bbox




					def 						A     (       self       :						str     )  ->    Optional[Any]:
								UpperCAmelCase_     :   Any       =   ids_tensor([self.batch_size, self.seq_length]							,				self.vocab_size     )

								UpperCAmelCase_     :   str       =   ids_tensor([self.batch_size, self.seq_length, 4]							,				self.range_bbox     )
								# Ensure that bbox is legal
								for i in range(bbox.shape[0]     ):
											for j in range(bbox.shape[1]     ):
														if bbox[i, j, 3] < bbox[i, j, 1]:
																	UpperCAmelCase_     :   Optional[Any]       =   bbox[i, j, 3]
																	UpperCAmelCase_     :   Optional[Any]       =   bbox[i, j, 1]
																	UpperCAmelCase_     :   Tuple       =   t
														if bbox[i, j, 2] < bbox[i, j, 0]:
																	UpperCAmelCase_     :   Any       =   bbox[i, j, 2]
																	UpperCAmelCase_     :   str       =   bbox[i, j, 0]
																	UpperCAmelCase_     :   Tuple       =   t

								UpperCAmelCase_     :   Union[str, Any]       =   None
								if self.use_input_mask:
											UpperCAmelCase_     :   Optional[int]       =   ids_tensor([self.batch_size, self.seq_length]							,				vocab_size=2     )

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

								UpperCAmelCase_     :   int       =   None
								UpperCAmelCase_     :   str       =   None
								UpperCAmelCase_     :   Dict       =   None
								if self.use_labels:
											UpperCAmelCase_     :   Any       =   ids_tensor([self.batch_size]							,				self.type_sequence_label_size     )
											UpperCAmelCase_     :   Any       =   ids_tensor([self.batch_size, self.seq_length]							,				self.num_labels     )
											UpperCAmelCase_     :   List[str]       =   ids_tensor([self.batch_size]							,				self.num_choices     )

								UpperCAmelCase_     :   str       =   self.get_config()

								return config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels




					def 						A     (       self       :						List[Any]     )  ->    Tuple:
								return LayoutLMConfig(
								    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							,				initializer_range=self.initializer_range							,				)




					def 						A     (       self       :						Any							,				_A       :						Union[str, Any]							,				_A       :						Union[str, Any]							,				_A       :						str							,				_A       :						Dict							,				_A       :						Any							,				_A       :						Any							,				_A       :						Tuple							,				_A       :						Dict     )  ->    str:
								UpperCAmelCase_     :   Any       =   LayoutLMModel(config=_A     )
								model.to(_A     )
								model.eval()
								UpperCAmelCase_     :   List[str]       =   model(_A							,				_A							,				attention_mask=_A							,				token_type_ids=_A     )
								UpperCAmelCase_     :   Any       =   model(_A							,				_A							,				token_type_ids=_A     )
								UpperCAmelCase_     :   Tuple       =   model(_A							,				_A     )
								self.parent.assertEqual(result.last_hidden_state.shape							,				(self.batch_size, self.seq_length, self.hidden_size)     )
								self.parent.assertEqual(result.pooler_output.shape							,				(self.batch_size, self.hidden_size)     )




					def 						A     (       self       :						Optional[int]							,				_A       :						Optional[int]							,				_A       :						Tuple							,				_A       :						Dict							,				_A       :						List[str]							,				_A       :						Dict							,				_A       :						Union[str, Any]							,				_A       :						Dict							,				_A       :						Dict     )  ->    Dict:
								UpperCAmelCase_     :   Any       =   LayoutLMForMaskedLM(config=_A     )
								model.to(_A     )
								model.eval()
								UpperCAmelCase_     :   Union[str, Any]       =   model(_A							,				_A							,				attention_mask=_A							,				token_type_ids=_A							,				labels=_A     )
								self.parent.assertEqual(result.logits.shape							,				(self.batch_size, self.seq_length, self.vocab_size)     )




					def 						A     (       self       :						Optional[int]							,				_A       :						List[str]							,				_A       :						Optional[Any]							,				_A       :						Optional[int]							,				_A       :						int							,				_A       :						int							,				_A       :						int							,				_A       :						int							,				_A       :						Union[str, Any]     )  ->    Optional[int]:
								UpperCAmelCase_     :   List[Any]       =   self.num_labels
								UpperCAmelCase_     :   Any       =   LayoutLMForSequenceClassification(_A     )
								model.to(_A     )
								model.eval()
								UpperCAmelCase_     :   Any       =   model(
								    _A							,				_A							,				attention_mask=_A							,				token_type_ids=_A							,				labels=_A     )
								self.parent.assertEqual(result.logits.shape							,				(self.batch_size, self.num_labels)     )




					def 						A     (       self       :						Union[str, Any]							,				_A       :						int							,				_A       :						List[Any]							,				_A       :						Any							,				_A       :						Any							,				_A       :						Optional[int]							,				_A       :						List[Any]							,				_A       :						Union[str, Any]							,				_A       :						Dict     )  ->    int:
								UpperCAmelCase_     :   List[Any]       =   self.num_labels
								UpperCAmelCase_     :   int       =   LayoutLMForTokenClassification(config=_A     )
								model.to(_A     )
								model.eval()
								UpperCAmelCase_     :   int       =   model(_A							,				_A							,				attention_mask=_A							,				token_type_ids=_A							,				labels=_A     )
								self.parent.assertEqual(result.logits.shape							,				(self.batch_size, self.seq_length, self.num_labels)     )




					def 						A     (       self       :						Any							,				_A       :						int							,				_A       :						Dict							,				_A       :						List[Any]							,				_A       :						List[Any]							,				_A       :						Any							,				_A       :						Union[str, Any]							,				_A       :						str							,				_A       :						List[str]     )  ->    Optional[Any]:
								UpperCAmelCase_     :   Union[str, Any]       =   LayoutLMForQuestionAnswering(config=_A     )
								model.to(_A     )
								model.eval()
								UpperCAmelCase_     :   Any       =   model(
								    _A							,				bbox=_A							,				attention_mask=_A							,				token_type_ids=_A							,				start_positions=_A							,				end_positions=_A							,				)
								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 						A     (       self       :						int     )  ->    Union[str, Any]:
								UpperCAmelCase_     :   Union[str, Any]       =   self.prepare_config_and_inputs()
								(
								    UpperCAmelCase_   
								)     :   Dict       =   config_and_inputs
								UpperCAmelCase_     :   Any       =   {
								    "input_ids": input_ids,
								    "bbox": bbox,
								    "token_type_ids": token_type_ids,
								    "attention_mask": input_mask,
								}
								return config, inputs_dict






@require_torch
class 				snake_case__     (			A__      ,							A__      ,							unittest.TestCase):
					a_       =			(
					    (
					        LayoutLMModel,
					        LayoutLMForMaskedLM,
					        LayoutLMForSequenceClassification,
					        LayoutLMForTokenClassification,
					        LayoutLMForQuestionAnswering,
					    )
					    if is_torch_available()
					    else None
					)
					a_       =			(
					    {
					        "document-question-answering": LayoutLMForQuestionAnswering,
					        "feature-extraction": LayoutLMModel,
					        "fill-mask": LayoutLMForMaskedLM,
					        "text-classification": LayoutLMForSequenceClassification,
					        "token-classification": LayoutLMForTokenClassification,
					        "zero-shot": LayoutLMForSequenceClassification,
					    }
					    if is_torch_available()
					    else {}
					)
					a_       =			True




					def 						A     (       self       :						List[Any]     )  ->    Tuple:
								UpperCAmelCase_     :   List[str]       =   LayoutLMModelTester(self     )
								UpperCAmelCase_     :   str       =   ConfigTester(self							,				config_class=_A							,				hidden_size=37     )




					def 						A     (       self       :						str     )  ->    List[Any]:
								self.config_tester.run_common_tests()




					def 						A     (       self       :						Union[str, Any]     )  ->    str:
								UpperCAmelCase_     :   Dict       =   self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_A     )




					def 						A     (       self       :						Union[str, Any]     )  ->    Union[str, Any]:
								UpperCAmelCase_     :   Optional[Any]       =   self.model_tester.prepare_config_and_inputs()
								for type in ["absolute", "relative_key", "relative_key_query"]:
											UpperCAmelCase_     :   Optional[int]       =   type
											self.model_tester.create_and_check_model(*_A     )




					def 						A     (       self       :						List[str]     )  ->    Dict:
								UpperCAmelCase_     :   Dict       =   self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*_A     )




					def 						A     (       self       :						Dict     )  ->    List[Any]:
								UpperCAmelCase_     :   List[Any]       =   self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_sequence_classification(*_A     )




					def 						A     (       self       :						Union[str, Any]     )  ->    Dict:
								UpperCAmelCase_     :   Union[str, Any]       =   self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*_A     )




					def 						A     (       self       :						Dict     )  ->    Dict:
								UpperCAmelCase_     :   int       =   self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_question_answering(*_A     )





					@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''     )
					def 						A     (       self       :						Any     )  ->    Any:
								pass






def __UpperCAmelCase (			)  -> Tuple:
			UpperCAmelCase_     :   Tuple       =   torch.tensor([[1_0_1,1_0_1_9,1_0_1_4,1_0_1_6,1_0_3_7,1_2_8_4_9,4_7_4_7,1_0_0_4,1_4_2_4_6,2_2_7_8,5_4_3_9,4_5_2_4,5_0_0_2,2_9_3_0,2_1_9_3,2_9_3_0,4_3_4_1,3_2_0_8,1_0_0_5,1_0_5_5,2_1_7_1,2_8_4_8,1_1_3_0_0,3_5_3_1,1_0_2],[1_0_1,4_0_7_0,4_0_3_4,7_0_2_0,1_0_2_4,3_0_5_8,1_0_1_5,1_0_1_3,2_8_6_1,1_0_1_3,6_0_7_0,1_9_2_7_4,2_7_7_2,6_2_0_5,2_7_8_1_4,1_6_1_4_7,1_6_1_4_7,4_3_4_3,2_0_4_7,1_0_2_8_3,1_0_9_6_9,1_4_3_8_9,1_0_1_2,2_3_3_8,1_0_2]]	,   device=A						)  # noqa: E231
			UpperCAmelCase_     :   Optional[int]       =   torch.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],]	,   device=A						)  # noqa: E231
			UpperCAmelCase_     :   Dict       =   torch.tensor([[[0,0,0,0],[4_2_3,2_3_7,4_4_0,2_5_1],[4_2_7,2_7_2,4_4_1,2_8_7],[4_1_9,1_1_5,4_3_7,1_2_9],[9_6_1,8_8_5,9_9_2,9_1_2],[2_5_6,3_8,3_3_0,5_8],[2_5_6,3_8,3_3_0,5_8],[3_3_6,4_2,3_5_3,5_7],[3_6_0,3_9,4_0_1,5_6],[3_6_0,3_9,4_0_1,5_6],[4_1_1,3_9,4_7_1,5_9],[4_7_9,4_1,5_2_8,5_9],[5_3_3,3_9,6_3_0,6_0],[6_7,1_1_3,1_3_4,1_3_1],[1_4_1,1_1_5,2_0_9,1_3_2],[6_8,1_4_9,1_3_3,1_6_6],[1_4_1,1_4_9,1_8_7,1_6_4],[1_9_5,1_4_8,2_8_7,1_6_5],[1_9_5,1_4_8,2_8_7,1_6_5],[1_9_5,1_4_8,2_8_7,1_6_5],[2_9_5,1_4_8,3_4_9,1_6_5],[4_4_1,1_4_9,4_9_2,1_6_6],[4_9_7,1_4_9,5_4_6,1_6_4],[6_4,2_0_1,1_2_5,2_1_8],[1_0_0_0,1_0_0_0,1_0_0_0,1_0_0_0]],[[0,0,0,0],[6_6_2,1_5_0,7_5_4,1_6_6],[6_6_5,1_9_9,7_4_2,2_1_1],[5_1_9,2_1_3,5_5_4,2_2_8],[5_1_9,2_1_3,5_5_4,2_2_8],[1_3_4,4_3_3,1_8_7,4_5_4],[1_3_0,4_6_7,2_0_4,4_8_0],[1_3_0,4_6_7,2_0_4,4_8_0],[1_3_0,4_6_7,2_0_4,4_8_0],[1_3_0,4_6_7,2_0_4,4_8_0],[1_3_0,4_6_7,2_0_4,4_8_0],[3_1_4,4_6_9,3_7_6,4_8_2],[5_0_4,6_8_4,5_8_2,7_0_6],[9_4_1,8_2_5,9_7_3,9_0_0],[9_4_1,8_2_5,9_7_3,9_0_0],[9_4_1,8_2_5,9_7_3,9_0_0],[9_4_1,8_2_5,9_7_3,9_0_0],[6_1_0,7_4_9,6_5_2,7_6_5],[1_3_0,6_5_9,1_6_8,6_7_2],[1_7_6,6_5_7,2_3_7,6_7_2],[2_3_8,6_5_7,3_1_2,6_7_2],[4_4_3,6_5_3,6_2_8,6_7_2],[4_4_3,6_5_3,6_2_8,6_7_2],[7_1_6,3_0_1,8_2_5,3_1_7],[1_0_0_0,1_0_0_0,1_0_0_0,1_0_0_0]]]	,   device=A						)  # noqa: E231
			UpperCAmelCase_     :   List[Any]       =   torch.tensor([[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]]	,   device=A						)  # noqa: E231
			# these are sequence labels (i.e. at the token level)
			UpperCAmelCase_     :   List[Any]       =   torch.tensor([[-1_0_0,1_0,1_0,1_0,9,1,-1_0_0,7,7,-1_0_0,7,7,4,2,5,2,8,8,-1_0_0,-1_0_0,5,0,3,2,-1_0_0],[-1_0_0,1_2,1_2,1_2,-1_0_0,1_2,1_0,-1_0_0,-1_0_0,-1_0_0,-1_0_0,1_0,1_2,9,-1_0_0,-1_0_0,-1_0_0,1_0,1_0,1_0,9,1_2,-1_0_0,1_0,-1_0_0]]	,   device=A						)  # noqa: E231
			# fmt: on

			return input_ids, attention_mask, bbox, token_type_ids, labels







@require_torch
class 				snake_case__     (			unittest.TestCase):




					@slow
					def 						A     (       self       :						Dict     )  ->    Dict:
								UpperCAmelCase_     :   List[Any]       =   LayoutLMModel.from_pretrained('''microsoft/layoutlm-base-uncased'''     ).to(_A     )

								UpperCAmelCase_     :   Union[str, Any]       =   prepare_layoutlm_batch_inputs()

								# forward pass
								UpperCAmelCase_     :   Tuple       =   model(input_ids=_A							,				bbox=_A							,				attention_mask=_A							,				token_type_ids=_A     )

								# test the sequence output on [0, :3, :3]
								UpperCAmelCase_     :   Optional[int]       =   torch.tensor(
								    [[0.1_785, -0.1_947, -0.0_425], [-0.3_254, -0.2_807, 0.2_553], [-0.5_391, -0.3_322, 0.3_364]]							,				device=_A							,				)

								self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]							,				_A							,				atol=1e-3     )     )

								# test the pooled output on [1, :3]
								UpperCAmelCase_     :   List[Any]       =   torch.tensor([-0.6_580, -0.0_214, 0.8_552]							,				device=_A     )

								self.assertTrue(torch.allclose(outputs.pooler_output[1, :3]							,				_A							,				atol=1e-3     )     )




					@slow
					def 						A     (       self       :						Optional[Any]     )  ->    Union[str, Any]:
								UpperCAmelCase_     :   str       =   LayoutLMForSequenceClassification.from_pretrained('''microsoft/layoutlm-base-uncased'''							,				num_labels=2     ).to(
								    _A     )

								UpperCAmelCase_     :   Dict       =   prepare_layoutlm_batch_inputs()

								# forward pass
								UpperCAmelCase_     :   Tuple       =   model(
								    input_ids=_A							,				bbox=_A							,				attention_mask=_A							,				token_type_ids=_A							,				labels=torch.tensor([1, 1]							,				device=_A     )							,				)

								# test whether we get a loss as a scalar
								UpperCAmelCase_     :   Dict       =   outputs.loss
								UpperCAmelCase_     :   str       =   torch.Size([]     )
								self.assertEqual(loss.shape							,				_A     )

								# test the shape of the logits
								UpperCAmelCase_     :   int       =   outputs.logits
								UpperCAmelCase_     :   Tuple       =   torch.Size((2, 2)     )
								self.assertEqual(logits.shape							,				_A     )




					@slow
					def 						A     (       self       :						Dict     )  ->    str:
								UpperCAmelCase_     :   Tuple       =   LayoutLMForTokenClassification.from_pretrained('''microsoft/layoutlm-base-uncased'''							,				num_labels=13     ).to(
								    _A     )

								UpperCAmelCase_     :   Tuple       =   prepare_layoutlm_batch_inputs()

								# forward pass
								UpperCAmelCase_     :   Any       =   model(
								    input_ids=_A							,				bbox=_A							,				attention_mask=_A							,				token_type_ids=_A							,				labels=_A     )

								# test the loss calculation to be around 2.65
								# expected_loss = torch.tensor(2.65, device=torch_device)

								# The loss is currently somewhat random and can vary between 0.1-0.3 atol.
								# self.assertTrue(torch.allclose(outputs.loss, expected_loss, atol=0.1))

								# test the shape of the logits
								UpperCAmelCase_     :   str       =   outputs.logits
								UpperCAmelCase_     :   int       =   torch.Size((2, 25, 13)     )
								self.assertEqual(logits.shape							,				_A     )





					@slow
					def 						A     (       self       :						Any     )  ->    List[Any]:
								UpperCAmelCase_     :   Optional[int]       =   LayoutLMForQuestionAnswering.from_pretrained('''microsoft/layoutlm-base-uncased'''     ).to(_A     )

								UpperCAmelCase_     :   str       =   prepare_layoutlm_batch_inputs()

								# forward pass
								UpperCAmelCase_     :   Optional[int]       =   model(input_ids=_A							,				bbox=_A							,				attention_mask=_A							,				token_type_ids=_A     )

								# test the shape of the logits
								UpperCAmelCase_     :   Optional[Any]       =   torch.Size((2, 25)     )
								self.assertEqual(outputs.start_logits.shape							,				_A     )
								self.assertEqual(outputs.end_logits.shape							,				_A     )










'''simple docstring'''





import os
import tempfile
import unittest

import numpy as np
import torch

from diffusers.models import ModelMixin, UNetaDConditionModel
from diffusers.models.attention_processor import AttnProcessor, LoRAAttnProcessor
from diffusers.utils import (
    floats_tensor,
    logging,
    skip_mps,
    torch_device,
)
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()

_UpperCamelCase		:		Dict								=			logging.get_logger(__name__)
def __UpperCAmelCase (			A			:	List[Any]	,   A			:	bool = True						)  -> List[Any]:
			UpperCAmelCase_     :   Any       =   {}
			for name in model.attn_processors.keys():
						UpperCAmelCase_     :   List[Any]       =   name.endswith('''attn2.processor'''						) and not (
						    name.startswith('''transformer_in'''						) or '''temp_attentions''' in name.split('''.'''						)
						)
						UpperCAmelCase_     :   Optional[Any]       =   model.config.cross_attention_dim if has_cross_attention else None
						if name.startswith('''mid_block'''						):
									UpperCAmelCase_     :   Dict       =   model.config.block_out_channels[-1]
						elif name.startswith('''up_blocks'''						):
									UpperCAmelCase_     :   Any       =   int(name[len('''up_blocks.'''						)]						)
									UpperCAmelCase_     :   int       =   list(reversed(model.config.block_out_channels						)						)[block_id]
						elif name.startswith('''down_blocks'''						):
									UpperCAmelCase_     :   Dict       =   int(name[len('''down_blocks.'''						)]						)
									UpperCAmelCase_     :   Any       =   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_     :   int       =   8 * model.config.attention_head_dim

						UpperCAmelCase_     :   Optional[int]       =   LoRAAttnProcessor(hidden_size=A	,   cross_attention_dim=A						)
						UpperCAmelCase_     :   Tuple       =   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):
					a_       =			UNetaDConditionModel
					a_       =			"sample"




					@property
					def 						A     (       self       :						Tuple     )  ->    int:
								UpperCAmelCase_     :   Optional[Any]       =   4
								UpperCAmelCase_     :   int       =   4
								UpperCAmelCase_     :   List[Any]       =   4
								UpperCAmelCase_     :   int       =   (32, 32)

								UpperCAmelCase_     :   Optional[int]       =   floats_tensor((batch_size, num_channels, num_frames) + sizes     ).to(_A     )
								UpperCAmelCase_     :   Any       =   torch.tensor([10]     ).to(_A     )
								UpperCAmelCase_     :   Optional[Any]       =   floats_tensor((batch_size, 4, 32)     ).to(_A     )

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




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




					@property
					def 						A     (       self       :						Union[str, Any]     )  ->    Union[str, Any]:
								return (4, 4, 32, 32)




					def 						A     (       self       :						Dict     )  ->    List[Any]:
								UpperCAmelCase_     :   List[str]       =   {
								    '''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_     :   Union[str, 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 						A     (       self       :						Union[str, Any]     )  ->    Optional[int]:
								UpperCAmelCase_  ,     UpperCAmelCase_     :   Optional[int]       =   self.prepare_init_args_and_inputs_for_common()
								UpperCAmelCase_     :   List[str]       =   self.model_class(**_A     )

								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 						A     (       self       :						Optional[Any]     )  ->    List[str]:
								UpperCAmelCase_  ,     UpperCAmelCase_     :   Tuple       =   self.prepare_init_args_and_inputs_for_common()

								UpperCAmelCase_     :   str       =   32

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

								with torch.no_grad():
											UpperCAmelCase_     :   Any       =   model(**_A     )

											if isinstance(_A							,				_A     ):
														UpperCAmelCase_     :   Optional[Any]       =   output.sample

								self.assertIsNotNone(_A     )
								UpperCAmelCase_     :   int       =   inputs_dict['''sample'''].shape
								self.assertEqual(output.shape							,				_A							,				'''Input and output shapes do not match'''     )




					def 						A     (       self       :						Union[str, Any]     )  ->    Optional[int]:
								UpperCAmelCase_  ,     UpperCAmelCase_     :   Optional[Any]       =   self.prepare_init_args_and_inputs_for_common()
								UpperCAmelCase_     :   Optional[Any]       =   self.model_class(**_A     )
								model.to(_A     )
								model.eval()

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

											UpperCAmelCase_     :   List[Any]       =   model(**_A     )
											if isinstance(_A							,				_A     ):
														UpperCAmelCase_     :   List[Any]       =   first.sample

											UpperCAmelCase_     :   Tuple       =   model(**_A     )
											if isinstance(_A							,				_A     ):
														UpperCAmelCase_     :   Union[str, Any]       =   second.sample

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




					def 						A     (       self       :						str     )  ->    str:
								UpperCAmelCase_  ,     UpperCAmelCase_     :   Optional[Any]       =   self.prepare_init_args_and_inputs_for_common()

								UpperCAmelCase_     :   Union[str, Any]       =   8

								UpperCAmelCase_     :   Tuple       =   self.model_class(**_A     )
								model.to(_A     )
								model.eval()

								model.set_attention_slice('''auto'''     )
								with torch.no_grad():
											UpperCAmelCase_     :   Dict       =   model(**_A     )
								assert output is not None

								model.set_attention_slice('''max'''     )
								with torch.no_grad():
											UpperCAmelCase_     :   Union[str, Any]       =   model(**_A     )
								assert output is not None

								model.set_attention_slice(2     )
								with torch.no_grad():
											UpperCAmelCase_     :   int       =   model(**_A     )
								assert output is not None




					def 						A     (       self       :						List[Any]     )  ->    Any:
								UpperCAmelCase_  ,     UpperCAmelCase_     :   List[Any]       =   self.prepare_init_args_and_inputs_for_common()

								UpperCAmelCase_     :   Dict       =   8

								UpperCAmelCase_     :   List[Any]       =   self.model_class(**_A     )
								model.to(_A     )

								with torch.no_grad():
											UpperCAmelCase_     :   int       =   model(**_A     ).sample

								UpperCAmelCase_     :   Any       =   create_lora_layers(_A     )

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

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

								with torch.no_grad():
											UpperCAmelCase_     :   int       =   model(**_A							,				cross_attention_kwargs={'''scale''': 0.0}     ).sample
											UpperCAmelCase_     :   int       =   model(**_A							,				cross_attention_kwargs={'''scale''': 0.5}     ).sample
											UpperCAmelCase_     :   int       =   model(**_A							,				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 						A     (       self       :						Union[str, Any]     )  ->    Tuple:
								UpperCAmelCase_  ,     UpperCAmelCase_     :   List[str]       =   self.prepare_init_args_and_inputs_for_common()

								UpperCAmelCase_     :   str       =   8

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

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

								UpperCAmelCase_     :   int       =   create_lora_layers(_A     )
								model.set_attn_processor(_A     )

								with torch.no_grad():
											UpperCAmelCase_     :   Optional[int]       =   model(**_A							,				cross_attention_kwargs={'''scale''': 0.5}     ).sample

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

								with torch.no_grad():
											UpperCAmelCase_     :   Dict       =   new_model(**_A							,				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 						A     (       self       :						Tuple     )  ->    Union[str, Any]:
								UpperCAmelCase_  ,     UpperCAmelCase_     :   Dict       =   self.prepare_init_args_and_inputs_for_common()

								UpperCAmelCase_     :   List[str]       =   8

								torch.manual_seed(0     )
								UpperCAmelCase_     :   str       =   self.model_class(**_A     )
								model.to(_A     )

								with torch.no_grad():
											UpperCAmelCase_     :   Dict       =   model(**_A     ).sample

								UpperCAmelCase_     :   Tuple       =   create_lora_layers(_A     )
								model.set_attn_processor(_A     )

								with torch.no_grad():
											UpperCAmelCase_     :   Dict       =   model(**_A							,				cross_attention_kwargs={'''scale''': 0.5}     ).sample

								with tempfile.TemporaryDirectory() as tmpdirname:
											model.save_attn_procs(_A							,				safe_serialization=_A     )
											self.assertTrue(os.path.isfile(os.path.join(_A							,				'''pytorch_lora_weights.safetensors'''     )     )     )
											torch.manual_seed(0     )
											UpperCAmelCase_     :   Dict       =   self.model_class(**_A     )
											new_model.to(_A     )
											new_model.load_attn_procs(_A     )

								with torch.no_grad():
											UpperCAmelCase_     :   Union[str, Any]       =   new_model(**_A							,				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 						A     (       self       :						Union[str, Any]     )  ->    Dict:
								# enable deterministic behavior for gradient checkpointing
								UpperCAmelCase_  ,     UpperCAmelCase_     :   List[Any]       =   self.prepare_init_args_and_inputs_for_common()

								UpperCAmelCase_     :   Any       =   8

								torch.manual_seed(0     )
								UpperCAmelCase_     :   Optional[Any]       =   self.model_class(**_A     )
								model.to(_A     )

								UpperCAmelCase_     :   str       =   create_lora_layers(_A							,				mock_weights=_A     )
								model.set_attn_processor(_A     )
								# Saving as torch, properly reloads with directly filename
								with tempfile.TemporaryDirectory() as tmpdirname:
											model.save_attn_procs(_A     )
											self.assertTrue(os.path.isfile(os.path.join(_A							,				'''pytorch_lora_weights.bin'''     )     )     )
											torch.manual_seed(0     )
											UpperCAmelCase_     :   Union[str, Any]       =   self.model_class(**_A     )
											new_model.to(_A     )
											new_model.load_attn_procs(_A							,				weight_name='''pytorch_lora_weights.bin'''     )




					def 						A     (       self       :						Optional[Any]     )  ->    Optional[int]:
								UpperCAmelCase_  ,     UpperCAmelCase_     :   Optional[int]       =   self.prepare_init_args_and_inputs_for_common()

								UpperCAmelCase_     :   List[Any]       =   8

								torch.manual_seed(0     )
								UpperCAmelCase_     :   Dict       =   self.model_class(**_A     )
								model.to(_A     )

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




					def 						A     (       self       :						int     )  ->    Tuple:
								UpperCAmelCase_  ,     UpperCAmelCase_     :   Any       =   self.prepare_init_args_and_inputs_for_common()

								UpperCAmelCase_     :   Dict       =   8

								torch.manual_seed(0     )
								UpperCAmelCase_     :   Optional[Any]       =   self.model_class(**_A     )
								model.to(_A     )

								with torch.no_grad():
											UpperCAmelCase_     :   str       =   model(**_A     ).sample

								UpperCAmelCase_     :   str       =   create_lora_layers(_A     )
								model.set_attn_processor(_A     )

								with torch.no_grad():
											UpperCAmelCase_     :   Optional[int]       =   model(**_A							,				cross_attention_kwargs={'''scale''': 0.0}     ).sample

								model.set_attn_processor(AttnProcessor()     )

								with torch.no_grad():
											UpperCAmelCase_     :   int       =   model(**_A     ).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 						A     (       self       :						str     )  ->    Optional[Any]:
								# enable deterministic behavior for gradient checkpointing
								UpperCAmelCase_  ,     UpperCAmelCase_     :   Any       =   self.prepare_init_args_and_inputs_for_common()

								UpperCAmelCase_     :   Dict       =   4

								torch.manual_seed(0     )
								UpperCAmelCase_     :   List[Any]       =   self.model_class(**_A     )
								model.to(_A     )
								UpperCAmelCase_     :   List[str]       =   create_lora_layers(_A     )
								model.set_attn_processor(_A     )

								# default
								with torch.no_grad():
											UpperCAmelCase_     :   List[str]       =   model(**_A     ).sample

											model.enable_xformers_memory_efficient_attention()
											UpperCAmelCase_     :   List[str]       =   model(**_A     ).sample

											model.disable_xformers_memory_efficient_attention()
											UpperCAmelCase_     :   Any       =   model(**_A     ).sample

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





					def 						A     (       self       :						Tuple     )  ->    Union[str, Any]:
								UpperCAmelCase_  ,     UpperCAmelCase_     :   Optional[int]       =   self.prepare_init_args_and_inputs_for_common()
								UpperCAmelCase_     :   int       =   32

								UpperCAmelCase_     :   Any       =   self.model_class(**_A     )
								model.to(_A     )
								model.eval()

								with torch.no_grad():
											UpperCAmelCase_     :   Tuple       =   model(**_A     )[0]

								model.enable_forward_chunking()
								with torch.no_grad():
											UpperCAmelCase_     :   Optional[Any]       =   model(**_A     )[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.


