






'''simple docstring'''





import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import OPTConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
						import os

						# The slow tests are often failing with OOM error on GPU
						# This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
						# but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
						_UpperCamelCase		:		List[Any]								='platform'

						import jax
						import jax.numpy as jnp

						from transformers import FlaxOPTForCausalLM, FlaxOPTModel, GPTaTokenizer
def __UpperCAmelCase (			A			:	List[Any]	,   A			:	List[str]	,   A			:	Union[str, Any]=None	,   A			:	Optional[int]=None						)  -> Tuple:
			if attention_mask is None:
						UpperCAmelCase_     :   Dict       =   np.where(input_ids != config.pad_token_id	,   1	,   0						)
			return {
			    "input_ids": input_ids,
			    "attention_mask": attention_mask,
			}


@require_flax
class 				snake_case__     :




					def __init__(       self       :						int							,				_A       :						Union[str, Any]							,				_A       :						Tuple=13							,				_A       :						Any=7							,				_A       :						Union[str, Any]=True							,				_A       :						Dict=False							,				_A       :						Union[str, Any]=99							,				_A       :						int=16							,				_A       :						Union[str, Any]=2							,				_A       :						Any=4							,				_A       :						int=4							,				_A       :						Optional[int]="gelu"							,				_A       :						Optional[int]=0.1							,				_A       :						Optional[int]=0.1							,				_A       :						Union[str, Any]=20							,				_A       :						Dict=2							,				_A       :						List[str]=1							,				_A       :						Any=0							,				_A       :						Union[str, Any]=16							,				_A       :						Optional[Any]=16							,				_A       :						Optional[Any]=0.02							,				)  ->    Dict:
								UpperCAmelCase_     :   Tuple       =   parent
								UpperCAmelCase_     :   Optional[Any]       =   batch_size
								UpperCAmelCase_     :   Any       =   seq_length
								UpperCAmelCase_     :   List[Any]       =   is_training
								UpperCAmelCase_     :   Union[str, Any]       =   use_labels
								UpperCAmelCase_     :   Optional[Any]       =   vocab_size
								UpperCAmelCase_     :   str       =   hidden_size
								UpperCAmelCase_     :   Union[str, Any]       =   num_hidden_layers
								UpperCAmelCase_     :   Any       =   num_attention_heads
								UpperCAmelCase_     :   Dict       =   intermediate_size
								UpperCAmelCase_     :   List[Any]       =   hidden_act
								UpperCAmelCase_     :   Any       =   hidden_dropout_prob
								UpperCAmelCase_     :   int       =   attention_probs_dropout_prob
								UpperCAmelCase_     :   int       =   max_position_embeddings
								UpperCAmelCase_     :   Optional[Any]       =   eos_token_id
								UpperCAmelCase_     :   str       =   pad_token_id
								UpperCAmelCase_     :   str       =   bos_token_id
								UpperCAmelCase_     :   List[Any]       =   embed_dim
								UpperCAmelCase_     :   Optional[int]       =   word_embed_proj_dim
								UpperCAmelCase_     :   str       =   initializer_range
								UpperCAmelCase_     :   Tuple       =   False




					def 						A     (       self       :						int     )  ->    Tuple:
								UpperCAmelCase_     :   Tuple       =   np.clip(ids_tensor([self.batch_size, self.seq_length - 1]							,				self.vocab_size     )							,				3							,				self.vocab_size     )
								UpperCAmelCase_     :   Any       =   np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)							,				dtype=np.intaa     ))							,				-1     )

								UpperCAmelCase_     :   Optional[Any]       =   OPTConfig(
								    vocab_size=self.vocab_size							,				hidden_size=self.hidden_size							,				num_hidden_layers=self.num_hidden_layers							,				num_attention_heads=self.num_attention_heads							,				ffn_dim=self.intermediate_size							,				dropout=self.hidden_dropout_prob							,				attention_dropout=self.attention_probs_dropout_prob							,				max_position_embeddings=self.max_position_embeddings							,				eos_token_id=self.eos_token_id							,				bos_token_id=self.bos_token_id							,				pad_token_id=self.pad_token_id							,				embed_dim=self.embed_dim							,				is_encoder_decoder=_A							,				word_embed_proj_dim=self.word_embed_proj_dim							,				initializer_range=self.initializer_range							,				use_cache=_A							,				)
								UpperCAmelCase_     :   List[str]       =   prepare_opt_inputs_dict(_A							,				_A     )
								return config, inputs_dict




					def 						A     (       self       :						str     )  ->    List[str]:
								UpperCAmelCase_     :   Tuple       =   self.prepare_config_and_inputs()
								return config, inputs_dict




					def 						A     (       self       :						Dict							,				_A       :						Dict							,				_A       :						Tuple							,				_A       :						Optional[int]     )  ->    str:
								UpperCAmelCase_     :   str       =   20
								UpperCAmelCase_     :   Tuple       =   model_class_name(_A     )

								UpperCAmelCase_     :   Optional[Any]       =   inputs_dict['''input_ids''']
								UpperCAmelCase_     :   List[str]       =   inputs_dict['''attention_mask''']

								UpperCAmelCase_     :   Optional[int]       =   model.init_cache(input_ids.shape[0]							,				_A     )
								UpperCAmelCase_     :   str       =   jnp.ones((input_ids.shape[0], max_length)							,				dtype='''i4'''     )

								UpperCAmelCase_     :   Any       =   jnp.broadcast_to(
								    jnp.arange(input_ids.shape[-1] - 1     )[None, :]							,				(input_ids.shape[0], input_ids.shape[-1] - 1)							,				)
								UpperCAmelCase_     :   Optional[Any]       =   model(
								    input_ids[:, :-1]							,				attention_mask=_A							,				past_key_values=_A							,				position_ids=_A							,				)

								UpperCAmelCase_     :   Dict       =   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]							,				dtype='''i4'''     )
								UpperCAmelCase_     :   Optional[int]       =   model(
								    input_ids[:, -1:]							,				attention_mask=_A							,				past_key_values=outputs_cache.past_key_values							,				position_ids=_A							,				)

								UpperCAmelCase_     :   List[Any]       =   model(_A     )

								UpperCAmelCase_     :   str       =   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])     )     )
								self.parent.assertTrue(diff < 1e-3							,				msg=F"Max diff is {diff}"     )





					def 						A     (       self       :						List[str]							,				_A       :						str							,				_A       :						Any							,				_A       :						List[str]     )  ->    Union[str, Any]:
								UpperCAmelCase_     :   Optional[int]       =   20
								UpperCAmelCase_     :   List[Any]       =   model_class_name(_A     )

								UpperCAmelCase_     :   Union[str, Any]       =   (
								    inputs_dict['''input_ids'''],
								    inputs_dict['''attention_mask'''],
								)

								UpperCAmelCase_     :   Union[str, Any]       =   jnp.concatenate(
								    [
								        attention_mask,
								        jnp.zeros((attention_mask.shape[0], max_length - attention_mask.shape[1])     ),
								    ]							,				axis=-1							,				)

								UpperCAmelCase_     :   Union[str, Any]       =   model.init_cache(input_ids.shape[0]							,				_A     )
								UpperCAmelCase_     :   Dict       =   jnp.broadcast_to(
								    jnp.arange(input_ids.shape[-1] - 1     )[None, :]							,				(input_ids.shape[0], input_ids.shape[-1] - 1)							,				)

								UpperCAmelCase_     :   Tuple       =   model(
								    input_ids[:, :-1]							,				attention_mask=_A							,				past_key_values=_A							,				position_ids=_A							,				)
								UpperCAmelCase_     :   Dict       =   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]							,				dtype='''i4'''     )
								UpperCAmelCase_     :   Any       =   model(
								    input_ids[:, -1:]							,				past_key_values=outputs_cache.past_key_values							,				attention_mask=_A							,				position_ids=_A							,				)

								UpperCAmelCase_     :   List[Any]       =   model(_A							,				attention_mask=_A     )

								UpperCAmelCase_     :   Dict       =   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])     )     )
								self.parent.assertTrue(diff < 1e-3							,				msg=F"Max diff is {diff}"     )






@require_flax
class 				snake_case__     (			UpperCamelCase      ,							unittest.TestCase      ,							UpperCamelCase):
					a_       =			(FlaxOPTModel, FlaxOPTForCausalLM) if is_flax_available() else ()
					a_       =			() if is_flax_available() else ()




					def 						A     (       self       :						Optional[int]     )  ->    Tuple:
								UpperCAmelCase_     :   Optional[Any]       =   FlaxOPTModelTester(self     )




					def 						A     (       self       :						Tuple     )  ->    Any:
								UpperCAmelCase_     :   List[str]       =   self.model_tester.prepare_config_and_inputs()
								for model_class in self.all_model_classes:
											self.model_tester.check_use_cache_forward(_A							,				_A							,				_A     )




					def 						A     (       self       :						str     )  ->    Optional[Any]:
								UpperCAmelCase_     :   Optional[Any]       =   self.model_tester.prepare_config_and_inputs()
								for model_class in self.all_model_classes:
											self.model_tester.check_use_cache_forward_with_attn_mask(_A							,				_A							,				_A     )





					@slow
					def 						A     (       self       :						List[Any]     )  ->    Dict:
								for model_class_name in self.all_model_classes:
											UpperCAmelCase_     :   Any       =   model_class_name.from_pretrained('''facebook/opt-125m'''     )
											UpperCAmelCase_     :   List[Any]       =   np.ones((1, 1)     ) * model.config.eos_token_id
											UpperCAmelCase_     :   Tuple       =   model(_A     )
											self.assertIsNotNone(_A     )






@require_sentencepiece
@require_flax
class 				snake_case__     (			unittest.TestCase):




					@slow
					def 						A     (       self       :						Optional[int]     )  ->    Optional[Any]:
								UpperCAmelCase_     :   Optional[int]       =   FlaxOPTModel.from_pretrained('''facebook/opt-350m'''     )
								UpperCAmelCase_     :   Tuple       =   jnp.array([[0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69, 4_60_78, 15_88, 2]]     )
								UpperCAmelCase_     :   str       =   model(input_ids=_A     ).last_hidden_state
								UpperCAmelCase_     :   int       =   (1, 11, 5_12)
								self.assertEqual(output.shape							,				_A     )
								UpperCAmelCase_     :   Any       =   jnp.array(
								    [[-0.2_867, -1.9_256, -0.3_062], [-1.2_711, -0.1_337, -0.1_897], [0.4_109, 0.1_187, -1.3_142]]     )
								self.assertTrue(jnp.allclose(output[:, :3, :3]							,				_A							,				atol=4e-2     )     )






@require_flax
@slow
class 				snake_case__     (			unittest.TestCase):




					def 						A     (       self       :						int     )  ->    Union[str, Any]:
								super().setUp()
								UpperCAmelCase_     :   str       =   '''facebook/opt-350m'''





					def 						A     (       self       :						List[Any]     )  ->    Optional[Any]:
								UpperCAmelCase_     :   Optional[Any]       =   FlaxOPTForCausalLM.from_pretrained(self.path_model     )
								UpperCAmelCase_     :   Tuple       =   GPTaTokenizer.from_pretrained(self.path_model     )

								UpperCAmelCase_     :   Tuple       =   [
								    '''Today is a beautiful day and I want to''',
								    '''In the city of''',
								    '''Paris is the capital of France and''',
								    '''Computers and mobile phones have taken''',
								]
								# verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
								UpperCAmelCase_     :   Dict       =   tokenizer(_A							,				return_tensors='''jax'''							,				padding=_A							,				add_special_tokens=_A     )
								UpperCAmelCase_     :   Tuple       =   model(inputs.input_ids							,				attention_mask=inputs.attention_mask     )[0].mean(axis=-1     )
								UpperCAmelCase_     :   List[Any]       =   jnp.array(
								    [
								        [1.3_851, -13.8_923, -10.5_229, -10.7_533, -0.2_309, -10.2_384, -0.5_365, -9.0_947, -5.1_670],
								        [-4.7_073, -10.6_276, -3.9_415, -21.5_242, -0.2_822, -0.2_822, -0.2_822, -0.2_822, -0.2_822],
								        [0.6_247, -3.4_229, -8.9_179, -1.4_297, -14.1_650, 1.4_146, -9.0_218, -0.2_703, -0.2_703],
								        [6.4_783, -1.9_913, -10.7_926, -2.3_336, 1.5_092, -0.9_974, -6.8_213, 1.3_477, 1.3_477],
								    ]     )
								self.assertTrue(jnp.allclose(_A							,				_A							,				atol=4e-2     )     )

								UpperCAmelCase_     :   List[Any]       =   jax.jit(_A     )
								UpperCAmelCase_     :   Union[str, Any]       =   model(inputs.input_ids							,				attention_mask=inputs.attention_mask     )[0].mean(axis=-1     )
								self.assertTrue(jnp.allclose(_A							,				_A							,				atol=4e-2     )     )






@require_flax
@slow
class 				snake_case__     (			unittest.TestCase):




					@property
					def 						A     (       self       :						Dict     )  ->    Any:
								return [
								    "Today is a beautiful day and I want",
								    "In the city of",
								    "Paris is the capital of France and",
								    "Computers and mobile phones have taken",
								]




					def 						A     (       self       :						List[str]     )  ->    Optional[int]:
								UpperCAmelCase_     :   Any       =   '''facebook/opt-125m'''

								UpperCAmelCase_     :   Optional[int]       =   [
								    '''Today is a beautiful day and I want to''',
								    '''In the city of New York, the city''',
								    '''Paris is the capital of France and the capital''',
								    '''Computers and mobile phones have taken over the''',
								]

								UpperCAmelCase_     :   List[Any]       =   []

								UpperCAmelCase_     :   Tuple       =   FlaxOPTForCausalLM.from_pretrained(_A     )
								UpperCAmelCase_     :   Tuple       =   GPTaTokenizer.from_pretrained(_A     )

								for prompt in self.prompts:
											UpperCAmelCase_     :   Dict       =   tokenizer(_A							,				return_tensors='''jax'''     ).input_ids

											UpperCAmelCase_     :   List[Any]       =   model.generate(_A							,				max_length=10     )
											UpperCAmelCase_     :   str       =   generated_ids[0]

											UpperCAmelCase_     :   Any       =   tokenizer.batch_decode(_A							,				skip_special_tokens=_A     )
											predicted_outputs += generated_string

								self.assertListEqual(_A							,				_A     )




					def 						A     (       self       :						Dict     )  ->    str:
								UpperCAmelCase_     :   str       =   '''facebook/opt-350m'''

								UpperCAmelCase_     :   str       =   [
								    '''Today is a beautiful day and I want to''',
								    '''In the city of San Francisco, the city''',
								    '''Paris is the capital of France and the capital''',
								    '''Computers and mobile phones have taken over the''',
								]

								UpperCAmelCase_     :   Union[str, Any]       =   []
								UpperCAmelCase_     :   Any       =   FlaxOPTForCausalLM.from_pretrained(_A     )
								UpperCAmelCase_     :   Dict       =   GPTaTokenizer.from_pretrained(_A     )

								for prompt in self.prompts:
											UpperCAmelCase_     :   str       =   tokenizer(_A							,				return_tensors='''jax'''     ).input_ids

											UpperCAmelCase_     :   Dict       =   model.generate(_A							,				max_length=10     )
											UpperCAmelCase_     :   List[Any]       =   generated_ids[0]

											UpperCAmelCase_     :   Optional[int]       =   tokenizer.batch_decode(_A							,				skip_special_tokens=_A     )
											predicted_outputs += generated_string

								self.assertListEqual(_A							,				_A     )




					def 						A     (       self       :						List[str]     )  ->    str:
								UpperCAmelCase_     :   Union[str, Any]       =   '''facebook/opt-125m'''
								UpperCAmelCase_     :   int       =   [
								    '''Today is a beautiful day and I want to thank''',
								    '''In the city of Rome Canaver Canaver Canaver Canaver''',
								]
								UpperCAmelCase_     :   Union[str, Any]       =   FlaxOPTForCausalLM.from_pretrained(_A     )
								UpperCAmelCase_     :   int       =   GPTaTokenizer.from_pretrained(_A     )
								UpperCAmelCase_     :   Optional[Any]       =   tokenizer(
								    [
								        '''Today is a beautiful day and I want to''',
								        '''In the city of''',
								    ]							,				return_tensors='''jax'''							,				padding=_A							,				)

								UpperCAmelCase_     :   str       =   jax.jit(model.generate     )

								UpperCAmelCase_     :   str       =   jit_generate(inputs['''input_ids''']							,				attention_mask=inputs['''attention_mask''']     ).sequences

								UpperCAmelCase_     :   int       =   tokenizer.batch_decode(_A							,				skip_special_tokens=_A     )

								self.assertIsNotNone(_A							,				_A     )





					def 						A     (       self       :						Any     )  ->    List[Any]:
								UpperCAmelCase_     :   Dict       =   '''facebook/opt-350m'''

								UpperCAmelCase_     :   List[str]       =   GPTaTokenizer.from_pretrained(_A     )
								UpperCAmelCase_     :   str       =   FlaxOPTForCausalLM.from_pretrained(_A     )

								UpperCAmelCase_     :   int       =   '''left'''

								# use different length sentences to test batching
								UpperCAmelCase_     :   List[str]       =   [
								    '''Hello, my dog is a little''',
								    '''Today, I''',
								]

								UpperCAmelCase_     :   List[Any]       =   tokenizer(_A							,				return_tensors='''jax'''							,				padding=_A     )
								UpperCAmelCase_     :   Dict       =   inputs['''input_ids''']

								UpperCAmelCase_     :   int       =   model.generate(input_ids=_A							,				attention_mask=inputs['''attention_mask''']							,				trace=_A     )

								UpperCAmelCase_     :   int       =   tokenizer(sentences[0]							,				return_tensors='''jax'''     ).input_ids
								UpperCAmelCase_     :   Optional[int]       =   model.generate(input_ids=_A     )

								UpperCAmelCase_     :   Dict       =   inputs_non_padded.shape[-1] - inputs['''attention_mask'''][-1].sum()
								UpperCAmelCase_     :   List[str]       =   tokenizer(sentences[1]							,				return_tensors='''jax'''     ).input_ids
								UpperCAmelCase_     :   Tuple       =   model.generate(input_ids=_A							,				max_length=model.config.max_length - num_paddings     )

								UpperCAmelCase_     :   Dict       =   tokenizer.batch_decode(outputs[0]							,				skip_special_tokens=_A     )
								UpperCAmelCase_     :   List[str]       =   tokenizer.decode(output_non_padded[0][0]							,				skip_special_tokens=_A     )
								UpperCAmelCase_     :   Tuple       =   tokenizer.decode(output_padded[0][0]							,				skip_special_tokens=_A     )

								UpperCAmelCase_     :   Dict       =   [
								    '''Hello, my dog is a little bit of a dork.\nI\'m a little bit''',
								    '''Today, I was in the middle of a conversation with a friend about the''',
								]
								self.assertListEqual(_A							,				_A     )
								self.assertListEqual(_A							,				[non_padded_sentence, padded_sentence]     )










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


