






"""simple docstring"""



import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from ....modeling_outputs import BaseModelOutputWithPooling, SequenceClassifierOutput
from ....modeling_utils import ModuleUtilsMixin
from ....utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings


A_     =				logging.get_logger(__name__)

A_     =				'''MMBTConfig'''






class 				__SCREAMING_SNAKE_CASE    (					nn.Module		):


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





													super().__init__()
													A__    :	Optional[int]			=	config
													A__    :	List[str]			=	encoder
													A__    :	List[str]			=	nn.Linear(config.modal_hidden_size      ,     config.hidden_size				)
													A__    :	Union[str, Any]			=	embeddings.position_embeddings
													A__    :	Any			=	embeddings.token_type_embeddings
													A__    :	Optional[int]			=	embeddings.word_embeddings
													A__    :	List[Any]			=	embeddings.LayerNorm
													A__    :	int			=	nn.Dropout(p=config.hidden_dropout_prob				)


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





													A__    :	Dict			=	self.proj_embeddings(self.encoder(snake_case				)				)
													A__    :	Tuple			=	token_embeddings.size(1				)

													if start_token is not None:
																			A__    :	Dict			=	self.word_embeddings(snake_case				)
																			seq_length += 1
																			A__    :	Any			=	torch.cat([start_token_embeds.unsqueeze(1				), token_embeddings]      ,     dim=1				)

													if end_token is not None:
																			A__    :	str			=	self.word_embeddings(snake_case				)
																			seq_length += 1
																			A__    :	int			=	torch.cat([token_embeddings, end_token_embeds.unsqueeze(1				)]      ,     dim=1				)

													if position_ids is None:
																			A__    :	Union[str, Any]			=	torch.arange(snake_case      ,     dtype=torch.long      ,     device=input_modal.device				)
																			A__    :	Tuple			=	position_ids.unsqueeze(0				).expand(input_modal.size(0				)      ,     snake_case				)

													if token_type_ids is None:
																			A__    :	int			=	torch.zeros(
																			    (input_modal.size(0				), seq_length)      ,     dtype=torch.long      ,     device=input_modal.device				)

													A__    :	List[str]			=	self.position_embeddings(snake_case				)
													A__    :	Any			=	self.token_type_embeddings(snake_case				)
													A__    :	Dict			=	token_embeddings + position_embeddings + token_type_embeddings
													A__    :	Dict			=	self.LayerNorm(snake_case				)
													A__    :	int			=	self.dropout(snake_case				)
													return embeddings




A_     =				r'''
    MMBT model was proposed in [Supervised Multimodal Bitransformers for Classifying Images and
    Text](https://github.com/facebookresearch/mmbt) by Douwe Kiela, Suvrat Bhooshan, Hamed Firooz, Davide Testuggine.
    It\'s a supervised multimodal bitransformer model that fuses information from text and other image encoders, and
    obtain state-of-the-art performance on various multimodal classification benchmark tasks.

    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
    etc.)

    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
    and behavior.

    Parameters:
        config ([`MMBTConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration.
        transformer (`nn.Module`): A text transformer that is used by MMBT.
            It should have embeddings, encoder, and pooler attributes.
        encoder (`nn.Module`): Encoder for the second modality.
            It should take in a batch of modal inputs and return k, n dimension embeddings.
'''

A_     =				r'''
    Args:
        input_modal (`torch.FloatTensor` of shape `(batch_size, ***)`):
            The other modality data. It will be the shape that the encoder for that type expects. e.g. With an Image
            Encoder, the shape would be (batch_size, channels, height, width)
        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
            Indices of input sequence tokens in the vocabulary. It does not expect [CLS] token to be added as it\'s
            appended to the end of other modality embeddings. Indices can be obtained using [`AutoTokenizer`]. See
            [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details.

            [What are input IDs?](../glossary#input-ids)
        modal_start_tokens (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Optional start token to be added to Other Modality Embedding. [CLS] Most commonly used for classification
            tasks.
        modal_end_tokens (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Optional end token to be added to Other Modality Embedding. [SEP] Most commonly used.
        attention_mask (*optional*) `torch.FloatTensor` of shape `(batch_size, sequence_length)`:
            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.

            [What are attention masks?](../glossary#attention-mask)
        token_type_ids (*optional*) `torch.LongTensor` of shape `(batch_size, sequence_length)`:
            Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
            1]`:

            - 0 corresponds to a *sentence A* token,
            - 1 corresponds to a *sentence B* token.

            [What are token type IDs?](../glossary#token-type-ids)
        modal_token_type_ids (*optional*) `torch.LongTensor` of shape `(batch_size, modal_sequence_length)`:
            Segment token indices to indicate different portions of the non-text modality. The embeddings from these
            tokens will be summed with the respective token embeddings for the non-text modality.
        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
            config.max_position_embeddings - 1]`.

            [What are position IDs?](../glossary#position-ids)
        modal_position_ids (`torch.LongTensor` of shape `(batch_size, modal_sequence_length)`, *optional*):
            Indices of positions of each input sequence tokens in the position embeddings for the non-text modality.
            Selected in the range `[0, config.max_position_embeddings - 1]`.

            [What are position IDs?](../glossary#position-ids)
        head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
            Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:

            - 1 indicates the head is **not masked**,
            - 0 indicates the head is **masked**.

        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, embedding_dim)`, *optional*):
            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
            model\'s internal embedding lookup matrix.
        encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
            the model is configured as a decoder.
        encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
            the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:

            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.

        output_attentions (`bool`, *optional*):
            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
            tensors for more detail.
        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
'''






@add_start_docstrings(
    'The bare MMBT Model outputting raw hidden-states without any specific head on top.'							,		UpperCamelCase							,		)
class 				__SCREAMING_SNAKE_CASE    (					nn.Module							,		UpperCamelCase		):


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





													super().__init__()
													A__    :	Tuple			=	config
													A__    :	List[str]			=	transformer
													A__    :	Optional[int]			=	ModalEmbeddings(snake_case      ,     snake_case      ,     transformer.embeddings				)


							@add_start_docstrings_to_model_forward(snake_case				)
							@replace_return_docstrings(output_type=snake_case      ,     config_class=_CONFIG_FOR_DOC				)
							def     _UpperCamelCase							(      self   :    str      ,     snake_case   :    Dict      ,     snake_case   :    Optional[int]=None      ,     snake_case   :    int=None      ,     snake_case   :    Optional[int]=None      ,     snake_case   :    int=None      ,     snake_case   :    Any=None      ,     snake_case   :    List[Any]=None      ,     snake_case   :    Tuple=None      ,     snake_case   :    Optional[int]=None      ,     snake_case   :    Optional[Any]=None      ,     snake_case   :    List[str]=None      ,     snake_case   :    Optional[Any]=None      ,     snake_case   :    List[str]=None      ,     snake_case   :    List[Any]=None      ,     snake_case   :    int=None      ,     snake_case   :    Tuple=None      ,     ):
													'''simple docstring'''





													A__    :	Dict			=	output_attentions if output_attentions is not None else self.config.output_attentions
													A__    :	List[str]			=	(
													    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
													)
													A__    :	Dict			=	return_dict if return_dict is not None else self.config.use_return_dict

													if input_ids is not None and inputs_embeds is not None:
																			raise ValueError("""You cannot specify both input_ids and inputs_embeds at the same time"""				)
													elif input_ids is not None:
																			A__    :	Any			=	input_ids.size()
													elif inputs_embeds is not None:
																			A__    :	Optional[Any]			=	inputs_embeds.size()[:-1]
													else:
																			raise ValueError("""You have to specify either input_ids or inputs_embeds"""				)

													A__    :	Tuple			=	input_ids.device if input_ids is not None else inputs_embeds.device

													A__    :	Tuple			=	self.modal_encoder(
													    snake_case      ,     start_token=snake_case      ,     end_token=snake_case      ,     position_ids=snake_case      ,     token_type_ids=snake_case      ,     )

													A__    :	List[Any]			=	modal_embeddings.size()[:-1]

													if token_type_ids is None:
																			A__    :	Tuple			=	torch.ones(snake_case      ,     dtype=torch.long      ,     device=snake_case				)

													A__    :	Optional[Any]			=	self.transformer.embeddings(
													    input_ids=snake_case      ,     position_ids=snake_case      ,     token_type_ids=snake_case      ,     inputs_embeds=snake_case				)

													A__    :	Any			=	torch.cat([modal_embeddings, txt_embeddings]      ,     1				)

													A__    :	Optional[int]			=	embedding_output.size()[:-1]

													if attention_mask is None:
																			A__    :	Tuple			=	torch.ones(snake_case      ,     device=snake_case				)
													else:
																			A__    :	int			=	torch.cat(
																			    [torch.ones(snake_case      ,     device=snake_case      ,     dtype=torch.long				), attention_mask]      ,     dim=1				)
													if encoder_attention_mask is None:
																			A__    :	str			=	torch.ones(snake_case      ,     device=snake_case				)
													else:
																			A__    :	Any			=	torch.cat(
																			    [torch.ones(snake_case      ,     device=snake_case				), encoder_attention_mask]      ,     dim=1				)

													A__    :	Any			=	self.get_extended_attention_mask(snake_case      ,     snake_case				)
													A__    :	Optional[int]			=	self.invert_attention_mask(snake_case				)
													A__    :	Any			=	self.get_head_mask(snake_case      ,     self.config.num_hidden_layers				)

													A__    :	int			=	self.transformer.encoder(
													    snake_case      ,     attention_mask=snake_case      ,     head_mask=snake_case      ,     encoder_hidden_states=snake_case      ,     encoder_attention_mask=snake_case      ,     output_attentions=snake_case      ,     output_hidden_states=snake_case      ,     return_dict=snake_case      ,     )

													A__    :	List[str]			=	encoder_outputs[0]
													A__    :	Any			=	self.transformer.pooler(snake_case				)

													if not return_dict:
																			return (sequence_output, pooled_output) + encoder_outputs[1:]

													return BaseModelOutputWithPooling(
													    last_hidden_state=snake_case      ,     pooler_output=snake_case      ,     hidden_states=encoder_outputs.hidden_states      ,     attentions=encoder_outputs.attentions      ,     )


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





													return self.embeddings.word_embeddings


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





													A__    :	Optional[Any]			=	value






@add_start_docstrings(
    '\n    MMBT Model with a sequence classification/regression head on top (a linear layer on top of the pooled output)\n    '							,		UpperCamelCase							,		UpperCamelCase							,		)
class 				__SCREAMING_SNAKE_CASE    (					nn.Module		):


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





													super().__init__()
													A__    :	str			=	config.num_labels

													A__    :	Union[str, Any]			=	MMBTModel(snake_case      ,     snake_case      ,     snake_case				)
													A__    :	Any			=	nn.Dropout(config.hidden_dropout_prob				)
													A__    :	Optional[int]			=	nn.Linear(config.hidden_size      ,     config.num_labels				)


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





													A__    :	Dict			=	return_dict if return_dict is not None else self.config.use_return_dict

													A__    :	Optional[Any]			=	self.mmbt(
													    input_modal=snake_case      ,     input_ids=snake_case      ,     modal_start_tokens=snake_case      ,     modal_end_tokens=snake_case      ,     attention_mask=snake_case      ,     token_type_ids=snake_case      ,     modal_token_type_ids=snake_case      ,     position_ids=snake_case      ,     modal_position_ids=snake_case      ,     head_mask=snake_case      ,     inputs_embeds=snake_case      ,     return_dict=snake_case      ,     )

													A__    :	Optional[int]			=	outputs[1]

													A__    :	Tuple			=	self.dropout(snake_case				)
													A__    :	Any			=	self.classifier(snake_case				)

													A__    :	List[str]			=	None
													if labels is not None:
																			if self.num_labels == 1:
																									#  We are doing regression
																									A__    :	int			=	MSELoss()
																									A__    :	Union[str, Any]			=	loss_fct(logits.view(-1				)      ,     labels.view(-1				)				)
																			else:
																									A__    :	int			=	CrossEntropyLoss()
																									A__    :	int			=	loss_fct(logits.view(-1      ,     self.num_labels				)      ,     labels.view(-1				)				)

													if not return_dict:
																			A__    :	Tuple			=	(logits,) + outputs[2:]
																			return ((loss,) + output) if loss is not None else output

													return SequenceClassifierOutput(
													    loss=snake_case      ,     logits=snake_case      ,     hidden_states=outputs.hidden_states      ,     attentions=outputs.attentions      ,     )














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






