



'''simple docstring'''






import collections
import types

import numpy as np

from ..utils import (
    add_end_docstrings,
    is_tensorflow_probability_available,
    is_tf_available,
    is_torch_available,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, ArgumentHandler, Dataset, Pipeline, PipelineException


if is_torch_available():
						import torch

						from ..models.auto.modeling_auto import (
						    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
						    MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING,
						)

if is_tf_available() and is_tensorflow_probability_available():
						import tensorflow as tf
						import tensorflow_probability as tfp

						from ..models.auto.modeling_tf_auto import (
						    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
						    TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING,
						)

class 	_lowerCAmelCase     (    __A       ):






		'''simple docstring'''
		def __call__(						self							:       int      ,		UpperCamelCase_							:       Optional[Any]=None      ,		UpperCamelCase_							:       int=None      ,		**UpperCamelCase_							:       Optional[int]		)      ->			Optional[Any]:







									'''simple docstring'''





									requires_backends(self      ,		'''pandas'''		)
									import pandas as pd

									if table is None:
																raise ValueError('''Keyword argument `table` cannot be None.'''		)
									elif query is None:
																if isinstance(UpperCamelCase_      ,		UpperCamelCase_		) and table.get('''query'''		) is not None and table.get('''table'''		) is not None:
																							_lowercase							:		Tuple				=   [table]
																elif isinstance(UpperCamelCase_      ,		UpperCamelCase_		) and len(UpperCamelCase_		) > 0:
																							if not all(isinstance(UpperCamelCase_      ,		UpperCamelCase_		) for d in table		):
																														raise ValueError(
																														    F"Keyword argument `table` should be a list of dict, but is {(type(UpperCamelCase_		) for d in table)}"		)

																							if table[0].get('''query'''		) is not None and table[0].get('''table'''		) is not None:
																														_lowercase							:		List[Any]				=   table
																							else:
																														raise ValueError(
																														    '''If keyword argument `table` is a list of dictionaries, each dictionary should have a `table`'''
																														    F" and `query` key, but only dictionary has keys {table[0].keys()} `table` and `query` keys."		)
																elif Dataset is not None and isinstance(UpperCamelCase_      ,		UpperCamelCase_		) or isinstance(UpperCamelCase_      ,		types.GeneratorType		):
																							return table
																else:
																							raise ValueError(
																							    '''Invalid input. Keyword argument `table` should be either of type `dict` or `list`, but '''
																							    F"is {type(UpperCamelCase_		)})"		)
									else:
																_lowercase							:		List[str]				=   [{'''table''': table, '''query''': query}]

									for tqa_pipeline_input in tqa_pipeline_inputs:
																if not isinstance(tqa_pipeline_input['''table''']      ,		pd.DataFrame		):
																							if tqa_pipeline_input["table"] is None:
																														raise ValueError('''Table cannot be None.'''		)

																							_lowercase							:		Dict				=   pd.DataFrame(tqa_pipeline_input['''table''']		)

									return tqa_pipeline_inputs




@add_end_docstrings(__A       )
class 	_lowerCAmelCase     (    __A       ):






		'''simple docstring'''
		snake_case_							= 'table,query'
		def __init__(						self							:       str      ,		UpperCamelCase_							:       Any=TableQuestionAnsweringArgumentHandler()      ,		*UpperCamelCase_							:       List[Any]      ,		**UpperCamelCase_							:       str		)      ->			List[Any]:







									'''simple docstring'''





									super().__init__(*UpperCamelCase_      ,		**UpperCamelCase_		)
									_lowercase							:		Optional[Any]				=   args_parser

									self.check_model_type(
									    dict(
									        TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING.items()
									        + TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.items()		)
									    if self.framework == '''tf'''
									    else dict(
									        MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING.items() + MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.items()		)		)

									_lowercase							:		List[str]				=   bool(getattr(self.model.config      ,		'''aggregation_labels'''      ,		UpperCamelCase_		)		) and bool(
									    getattr(self.model.config      ,		'''num_aggregation_labels'''      ,		UpperCamelCase_		)		)
									_lowercase							:		Tuple				=   '''tapas''' if hasattr(self.model.config      ,		'''aggregation_labels'''		) else None
		def 						__lowercase      (						self							:       Union[str, Any]      ,		**UpperCamelCase_							:       str		)      ->			List[Any]:







									'''simple docstring'''





									return self.model(**UpperCamelCase_		)
		def 						__lowercase      (						self							:       List[Any]      ,		**UpperCamelCase_							:       Optional[Any]		)      ->			Union[str, Any]:







									'''simple docstring'''





									if self.framework == "pt":
																_lowercase							:		Tuple				=   []
																_lowercase							:		Dict				=   []
																_lowercase							:		Dict				=   None
																_lowercase							:		str				=   inputs['''input_ids'''].shape[0]

																_lowercase							:		str				=   inputs['''input_ids'''].to(self.device		)
																_lowercase							:		str				=   inputs['''attention_mask'''].to(self.device		)
																_lowercase							:		Optional[Any]				=   inputs['''token_type_ids'''].to(self.device		)
																_lowercase							:		List[str]				=   None

																for index in range(UpperCamelCase_		):
																							# If sequences have already been processed, the token type IDs will be created according to the previous
																							# answer.
																							if prev_answers is not None:
																														_lowercase							:		Optional[Any]				=   token_type_ids_example[:, 3]  # shape (seq_len,)
																														_lowercase							:		Optional[int]				=   np.zeros_like(prev_labels_example.cpu().numpy()		)  # shape (seq_len,)

																														_lowercase							:		Optional[Any]				=   token_type_ids[index]  # shape (seq_len, 7)
																														for i in range(model_labels.shape[0]		):
																																					_lowercase							:		List[Any]				=   token_type_ids_example[:, 0].tolist()[i]
																																					_lowercase							:		Optional[int]				=   token_type_ids_example[:, 1].tolist()[i] - 1
																																					_lowercase							:		List[str]				=   token_type_ids_example[:, 2].tolist()[i] - 1

																																					if row_id >= 0 and col_id >= 0 and segment_id == 1:
																																												_lowercase							:		Union[str, Any]				=   int(prev_answers[(col_id, row_id)]		)

																														_lowercase							:		Optional[Any]				=   torch.from_numpy(UpperCamelCase_		).type(torch.long		).to(self.device		)

																							_lowercase							:		Dict				=   input_ids[index]
																							_lowercase							:		Tuple				=   attention_mask[index]  # shape (seq_len,)
																							_lowercase							:		Optional[int]				=   token_type_ids[index]  # shape (seq_len, 7)
																							_lowercase							:		int				=   self.model(
																							    input_ids=input_ids_example.unsqueeze(0		)      ,		attention_mask=attention_mask_example.unsqueeze(0		)      ,		token_type_ids=token_type_ids_example.unsqueeze(0		)      ,		)
																							_lowercase							:		Dict				=   outputs.logits

																							if self.aggregate:
																														all_aggregations.append(outputs.logits_aggregation		)

																							all_logits.append(UpperCamelCase_		)

																							_lowercase							:		str				=   torch.distributions.Bernoulli(logits=UpperCamelCase_		)
																							_lowercase							:		Dict				=   dist_per_token.probs * attention_mask_example.type(torch.floataa		).to(
																							    dist_per_token.probs.device		)

																							_lowercase							:		List[Any]				=   collections.defaultdict(UpperCamelCase_		)
																							for i, p in enumerate(probabilities.squeeze().tolist()		):
																														_lowercase							:		str				=   token_type_ids_example[:, 0].tolist()[i]
																														_lowercase							:		str				=   token_type_ids_example[:, 1].tolist()[i] - 1
																														_lowercase							:		Optional[int]				=   token_type_ids_example[:, 2].tolist()[i] - 1
																														if col >= 0 and row >= 0 and segment_id == 1:
																																					coords_to_probs[(col, row)].append(UpperCamelCase_		)

																							_lowercase							:		Dict				=   {key: np.array(coords_to_probs[key]		).mean() > 0.5 for key in coords_to_probs}

																_lowercase							:		List[Any]				=   torch.cat(tuple(UpperCamelCase_		)      ,		0		)

																return (logits_batch,) if not self.aggregate else (logits_batch, torch.cat(tuple(UpperCamelCase_		)      ,		0		))
									else:
																_lowercase							:		Any				=   []
																_lowercase							:		Dict				=   []
																_lowercase							:		List[str]				=   None
																_lowercase							:		Any				=   inputs['''input_ids'''].shape[0]

																_lowercase							:		Optional[Any]				=   inputs['''input_ids''']
																_lowercase							:		Dict				=   inputs['''attention_mask''']
																_lowercase							:		str				=   inputs['''token_type_ids'''].numpy()
																_lowercase							:		Optional[int]				=   None

																for index in range(UpperCamelCase_		):
																							# If sequences have already been processed, the token type IDs will be created according to the previous
																							# answer.
																							if prev_answers is not None:
																														_lowercase							:		int				=   token_type_ids_example[:, 3]  # shape (seq_len,)
																														_lowercase							:		str				=   np.zeros_like(UpperCamelCase_      ,		dtype=np.intaa		)  # shape (seq_len,)

																														_lowercase							:		List[str]				=   token_type_ids[index]  # shape (seq_len, 7)
																														for i in range(model_labels.shape[0]		):
																																					_lowercase							:		Tuple				=   token_type_ids_example[:, 0].tolist()[i]
																																					_lowercase							:		Optional[Any]				=   token_type_ids_example[:, 1].tolist()[i] - 1
																																					_lowercase							:		str				=   token_type_ids_example[:, 2].tolist()[i] - 1

																																					if row_id >= 0 and col_id >= 0 and segment_id == 1:
																																												_lowercase							:		Optional[int]				=   int(prev_answers[(col_id, row_id)]		)

																														_lowercase							:		Any				=   model_labels

																							_lowercase							:		Union[str, Any]				=   input_ids[index]
																							_lowercase							:		Dict				=   attention_mask[index]  # shape (seq_len,)
																							_lowercase							:		List[str]				=   token_type_ids[index]  # shape (seq_len, 7)
																							_lowercase							:		Union[str, Any]				=   self.model(
																							    input_ids=np.expand_dims(UpperCamelCase_      ,		axis=0		)      ,		attention_mask=np.expand_dims(UpperCamelCase_      ,		axis=0		)      ,		token_type_ids=np.expand_dims(UpperCamelCase_      ,		axis=0		)      ,		)
																							_lowercase							:		Optional[Any]				=   outputs.logits

																							if self.aggregate:
																														all_aggregations.append(outputs.logits_aggregation		)

																							all_logits.append(UpperCamelCase_		)

																							_lowercase							:		List[Any]				=   tfp.distributions.Bernoulli(logits=UpperCamelCase_		)
																							_lowercase							:		Tuple				=   dist_per_token.probs_parameter() * tf.cast(UpperCamelCase_      ,		tf.floataa		)

																							_lowercase							:		Tuple				=   collections.defaultdict(UpperCamelCase_		)
																							_lowercase							:		List[str]				=   token_type_ids_example
																							for i, p in enumerate(tf.squeeze(UpperCamelCase_		).numpy().tolist()		):
																														_lowercase							:		str				=   token_type_ids_example[:, 0].tolist()[i]
																														_lowercase							:		List[Any]				=   token_type_ids_example[:, 1].tolist()[i] - 1
																														_lowercase							:		Optional[Any]				=   token_type_ids_example[:, 2].tolist()[i] - 1
																														if col >= 0 and row >= 0 and segment_id == 1:
																																					coords_to_probs[(col, row)].append(UpperCamelCase_		)

																							_lowercase							:		Any				=   {key: np.array(coords_to_probs[key]		).mean() > 0.5 for key in coords_to_probs}

																_lowercase							:		Tuple				=   tf.concat(tuple(UpperCamelCase_		)      ,		0		)

																return (logits_batch,) if not self.aggregate else (logits_batch, tf.concat(tuple(UpperCamelCase_		)      ,		0		))
		def __call__(						self							:       Tuple      ,		*UpperCamelCase_							:       Optional[int]      ,		**UpperCamelCase_							:       int		)      ->			List[str]:







									'''simple docstring'''





									_lowercase							:		Tuple				=   self._args_parser(*UpperCamelCase_      ,		**UpperCamelCase_		)

									_lowercase							:		int				=   super().__call__(UpperCamelCase_      ,		**UpperCamelCase_		)
									if len(UpperCamelCase_		) == 1:
																return results[0]
									return results
		def 						__lowercase      (						self							:       List[Any]      ,		UpperCamelCase_							:       Tuple=None      ,		UpperCamelCase_							:       List[Any]=None      ,		UpperCamelCase_							:       Optional[Any]=None      ,		**UpperCamelCase_							:       Optional[Any]		)      ->			Optional[Any]:







									'''simple docstring'''





									_lowercase							:		Tuple				=   {}
									if padding is not None:
																_lowercase							:		Any				=   padding
									if truncation is not None:
																_lowercase							:		Tuple				=   truncation

									_lowercase							:		Tuple				=   {}
									if sequential is not None:
																_lowercase							:		Optional[int]				=   sequential
									return preprocess_params, forward_params, {}
		def 						__lowercase      (						self							:       Union[str, Any]      ,		UpperCamelCase_							:       Union[str, Any]      ,		UpperCamelCase_							:       List[str]=None      ,		UpperCamelCase_							:       Any=True      ,		UpperCamelCase_							:       Dict=None		)      ->			List[Any]:







									'''simple docstring'''





									if truncation is None:
																if self.type == "tapas":
																							_lowercase							:		Union[str, Any]				=   '''drop_rows_to_fit'''
																else:
																							_lowercase							:		str				=   '''do_not_truncate'''

									_lowercase			, _lowercase							:		Any				=   pipeline_input['''table'''], pipeline_input['''query''']
									if table.empty:
																raise ValueError('''table is empty'''		)
									if query is None or query == "":
																raise ValueError('''query is empty'''		)
									_lowercase							:		Any				=   self.tokenizer(UpperCamelCase_      ,		UpperCamelCase_      ,		return_tensors=self.framework      ,		truncation=UpperCamelCase_      ,		padding=UpperCamelCase_		)
									_lowercase							:		str				=   table
									return inputs
		def 						__lowercase      (						self							:       Dict      ,		UpperCamelCase_							:       str      ,		UpperCamelCase_							:       Union[str, Any]=False		)      ->			Any:







									'''simple docstring'''





									_lowercase							:		Dict				=   model_inputs.pop('''table'''		)

									if self.type == "tapas":
																if sequential:
																							_lowercase							:		Union[str, Any]				=   self.sequential_inference(**UpperCamelCase_		)
																else:
																							_lowercase							:		Dict				=   self.batch_inference(**UpperCamelCase_		)
									else:
																_lowercase							:		int				=   self.model.generate(**UpperCamelCase_		)
									_lowercase							:		List[Any]				=   {'''model_inputs''': model_inputs, '''table''': table, '''outputs''': outputs}
									return model_outputs





		def 						__lowercase      (						self							:       List[str]      ,		UpperCamelCase_							:       int		)      ->			List[Any]:







									'''simple docstring'''





									_lowercase							:		str				=   model_outputs['''model_inputs''']
									_lowercase							:		List[Any]				=   model_outputs['''table''']
									_lowercase							:		int				=   model_outputs['''outputs''']
									if self.type == "tapas":
																if self.aggregate:
																							_lowercase			, _lowercase							:		Dict				=   outputs[:2]
																							_lowercase							:		Optional[int]				=   self.tokenizer.convert_logits_to_predictions(UpperCamelCase_      ,		UpperCamelCase_      ,		UpperCamelCase_		)
																							_lowercase			, _lowercase							:		Union[str, Any]				=   predictions
																							_lowercase							:		Tuple				=   {i: self.model.config.aggregation_labels[pred] for i, pred in enumerate(UpperCamelCase_		)}

																							_lowercase							:		Dict				=   self.model.config.no_aggregation_label_index
																							_lowercase							:		Tuple				=   {
																							    i: aggregators[i] + ''' > ''' for i, pred in enumerate(UpperCamelCase_		) if pred != no_agg_label_index
																							}
																else:
																							_lowercase							:		Optional[int]				=   outputs[0]
																							_lowercase							:		Optional[Any]				=   self.tokenizer.convert_logits_to_predictions(UpperCamelCase_      ,		UpperCamelCase_		)
																							_lowercase							:		str				=   predictions[0]
																							_lowercase							:		Optional[int]				=   {}
																							_lowercase							:		Optional[Any]				=   {}
																_lowercase							:		Optional[int]				=   []
																for index, coordinates in enumerate(UpperCamelCase_		):
																							_lowercase							:		List[str]				=   [table.iat[coordinate] for coordinate in coordinates]
																							_lowercase							:		Union[str, Any]				=   aggregators.get(UpperCamelCase_      ,		''''''		)
																							_lowercase							:		str				=   aggregators_prefix.get(UpperCamelCase_      ,		''''''		)
																							_lowercase							:		List[Any]				=   {
																							    '''answer''': aggregator_prefix + ''', '''.join(UpperCamelCase_		),
																							    '''coordinates''': coordinates,
																							    '''cells''': [table.iat[coordinate] for coordinate in coordinates],
																							}
																							if aggregator:
																														_lowercase							:		Union[str, Any]				=   aggregator

																							answers.append(UpperCamelCase_		)
																if len(UpperCamelCase_		) == 0:
																							raise PipelineException('''Empty answer'''		)
									else:
																_lowercase							:		Optional[int]				=   [{'''answer''': answer} for answer in self.tokenizer.batch_decode(UpperCamelCase_      ,		skip_special_tokens=UpperCamelCase_		)]

									return answers if len(UpperCamelCase_		) > 1 else answers[0]








'''simple docstring'''






import gc
import os
import tempfile
import unittest

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNetaDConditionModel
from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
from diffusers.models.attention_processor import (
    Attention,
    AttnProcessor,
    AttnProcessor2_0,
    LoRAAttnProcessor,
    LoRAAttnProcessor2_0,
    LoRAXFormersAttnProcessor,
    XFormersAttnProcessor,
)
from diffusers.utils import TEXT_ENCODER_ATTN_MODULE, floats_tensor, torch_device





def 				_SCREAMING_SNAKE_CASE(					snake_case_    :     nn.Module  )		->Any:






							'''simple docstring'''
							_lowercase							:		Optional[Any]				=   {}
							for name in unet.attn_processors.keys():
														_lowercase							:		str				=   None if name.endswith('''attn1.processor'''  ) else unet.config.cross_attention_dim
														if name.startswith('''mid_block'''  ):
																					_lowercase							:		Dict				=   unet.config.block_out_channels[-1]
														elif name.startswith('''up_blocks'''  ):
																					_lowercase							:		Any				=   int(name[len('''up_blocks.'''  )]  )
																					_lowercase							:		Any				=   list(reversed(unet.config.block_out_channels  )  )[block_id]
														elif name.startswith('''down_blocks'''  ):
																					_lowercase							:		int				=   int(name[len('''down_blocks.'''  )]  )
																					_lowercase							:		int				=   unet.config.block_out_channels[block_id]
														_lowercase							:		Any				=   (
														    LoRAAttnProcessor2_0 if hasattr(snake_case_ ,       '''scaled_dot_product_attention'''  ) else LoRAAttnProcessor
														)
														_lowercase							:		Tuple				=   lora_attn_processor_class(
														    hidden_size=snake_case_ ,       cross_attention_dim=snake_case_  )
							_lowercase							:		Any				=   AttnProcsLayers(snake_case_  )
							return lora_attn_procs, unet_lora_layers





def 				_SCREAMING_SNAKE_CASE(					snake_case_    :     nn.Module  )		->List[Any]:






							'''simple docstring'''
							_lowercase							:		int				=   {}
							_lowercase							:		str				=   (
							    LoRAAttnProcessor2_0 if hasattr(snake_case_ ,       '''scaled_dot_product_attention'''  ) else LoRAAttnProcessor
							)
							for name, module in text_encoder.named_modules():
														if name.endswith(snake_case_  ):
																					_lowercase							:		Optional[int]				=   lora_attn_processor_class(
																					    hidden_size=module.out_proj.out_features ,       cross_attention_dim=snake_case_  )
							return text_lora_attn_procs





def 				_SCREAMING_SNAKE_CASE(					snake_case_    :     nn.Module  )		->Union[str, Any]:






							'''simple docstring'''
							_lowercase							:		List[str]				=   create_text_encoder_lora_attn_procs(snake_case_  )
							_lowercase							:		List[Any]				=   AttnProcsLayers(snake_case_  )
							return text_encoder_lora_layers





def 				_SCREAMING_SNAKE_CASE(					snake_case_    :     Dict ,       snake_case_    :     str=False  )		->Optional[int]:






							'''simple docstring'''
							for _, attn_proc in text_lora_attn_procs.items():
														# set up.weights
														for layer_name, layer_module in attn_proc.named_modules():
																					if layer_name.endswith('''_lora'''  ):
																												_lowercase							:		Tuple				=   (
																												    torch.randn_like(layer_module.up.weight  )
																												    if randn_weight
																												    else torch.zeros_like(layer_module.up.weight  )
																												)
																												_lowercase							:		Any				=   torch.nn.Parameter(snake_case_  )




class 	_lowerCAmelCase     (    unittest.TestCase       ):






		'''simple docstring'''
		def 						__lowercase      (						self							:       Dict		)      ->			Tuple:







									'''simple docstring'''





									torch.manual_seed(0		)
									_lowercase							:		Any				=   UNetaDConditionModel(
									    block_out_channels=(32, 64)      ,		layers_per_block=2      ,		sample_size=32      ,		in_channels=4      ,		out_channels=4      ,		down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')      ,		up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')      ,		cross_attention_dim=32      ,		)
									_lowercase							:		Optional[Any]				=   DDIMScheduler(
									    beta_start=0.00_085      ,		beta_end=0.012      ,		beta_schedule='''scaled_linear'''      ,		clip_sample=UpperCamelCase_      ,		set_alpha_to_one=UpperCamelCase_      ,		steps_offset=1      ,		)
									torch.manual_seed(0		)
									_lowercase							:		int				=   AutoencoderKL(
									    block_out_channels=[32, 64]      ,		in_channels=3      ,		out_channels=3      ,		down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']      ,		up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']      ,		latent_channels=4      ,		)
									_lowercase							:		Any				=   CLIPTextConfig(
									    bos_token_id=0      ,		eos_token_id=2      ,		hidden_size=32      ,		intermediate_size=37      ,		layer_norm_eps=1e-05      ,		num_attention_heads=4      ,		num_hidden_layers=5      ,		pad_token_id=1      ,		vocab_size=1_000      ,		)
									_lowercase							:		List[Any]				=   CLIPTextModel(UpperCamelCase_		)
									_lowercase							:		Union[str, Any]				=   CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

									_lowercase			, _lowercase							:		Optional[int]				=   create_unet_lora_layers(UpperCamelCase_		)
									_lowercase							:		Optional[int]				=   create_text_encoder_lora_layers(UpperCamelCase_		)

									_lowercase							:		Dict				=   {
									    '''unet''': unet,
									    '''scheduler''': scheduler,
									    '''vae''': vae,
									    '''text_encoder''': text_encoder,
									    '''tokenizer''': tokenizer,
									    '''safety_checker''': None,
									    '''feature_extractor''': None,
									}
									_lowercase							:		List[str]				=   {
									    '''unet_lora_layers''': unet_lora_layers,
									    '''text_encoder_lora_layers''': text_encoder_lora_layers,
									    '''unet_lora_attn_procs''': unet_lora_attn_procs,
									}
									return pipeline_components, lora_components
		def 						__lowercase      (						self							:       Optional[Any]		)      ->			Optional[Any]:







									'''simple docstring'''





									_lowercase							:		Optional[Any]				=   1
									_lowercase							:		List[Any]				=   10
									_lowercase							:		Union[str, Any]				=   4
									_lowercase							:		Union[str, Any]				=   (32, 32)

									_lowercase							:		Optional[Any]				=   torch.manual_seed(0		)
									_lowercase							:		Tuple				=   floats_tensor((batch_size, num_channels) + sizes		)
									_lowercase							:		str				=   torch.randint(1      ,		UpperCamelCase_      ,		size=(batch_size, sequence_length)      ,		generator=UpperCamelCase_		)

									_lowercase							:		Optional[Any]				=   {
									    '''prompt''': '''A painting of a squirrel eating a burger''',
									    '''generator''': generator,
									    '''num_inference_steps''': 2,
									    '''guidance_scale''': 6.0,
									    '''output_type''': '''numpy''',
									}

									return noise, input_ids, pipeline_inputs

									# copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb
		def 						__lowercase      (						self							:       Union[str, Any]		)      ->			Any:







									'''simple docstring'''





									_lowercase							:		Optional[Any]				=   77

									_lowercase							:		str				=   torch.randint(2      ,		56      ,		size=(1, max_seq_length)      ,		generator=torch.manual_seed(0		)		)

									_lowercase							:		List[str]				=   {}
									_lowercase							:		Optional[Any]				=   inputs
									return prepared_inputs
		def 						__lowercase      (						self							:       List[str]      ,		UpperCamelCase_							:       str		)      ->			List[str]:







									'''simple docstring'''





									_lowercase			, _lowercase							:		Tuple				=   self.get_dummy_components()
									LoraLoaderMixin.save_lora_weights(
									    save_directory=UpperCamelCase_      ,		unet_lora_layers=lora_components['''unet_lora_layers''']      ,		text_encoder_lora_layers=lora_components['''text_encoder_lora_layers''']      ,		)
									self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase_      ,		'''pytorch_lora_weights.bin'''		)		)		)
		def 						__lowercase      (						self							:       List[str]		)      ->			Union[str, Any]:







									'''simple docstring'''





									_lowercase			, _lowercase							:		Any				=   self.get_dummy_components()
									_lowercase							:		Tuple				=   StableDiffusionPipeline(**UpperCamelCase_		)
									_lowercase							:		Optional[Any]				=   sd_pipe.to(UpperCamelCase_		)
									sd_pipe.set_progress_bar_config(disable=UpperCamelCase_		)

									_lowercase			, _lowercase			, _lowercase							:		Dict				=   self.get_dummy_inputs()

									_lowercase							:		Optional[int]				=   sd_pipe(**UpperCamelCase_		).images
									_lowercase							:		Optional[Any]				=   original_images[0, -3:, -3:, -1]

									with tempfile.TemporaryDirectory() as tmpdirname:
																LoraLoaderMixin.save_lora_weights(
																    save_directory=UpperCamelCase_      ,		unet_lora_layers=lora_components['''unet_lora_layers''']      ,		text_encoder_lora_layers=lora_components['''text_encoder_lora_layers''']      ,		)
																self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase_      ,		'''pytorch_lora_weights.bin'''		)		)		)
																sd_pipe.load_lora_weights(UpperCamelCase_		)

									_lowercase							:		Dict				=   sd_pipe(**UpperCamelCase_		).images
									_lowercase							:		Any				=   lora_images[0, -3:, -3:, -1]

									# Outputs shouldn't match.
									self.assertFalse(torch.allclose(torch.from_numpy(UpperCamelCase_		)      ,		torch.from_numpy(UpperCamelCase_		)		)		)
		def 						__lowercase      (						self							:       List[str]		)      ->			Optional[int]:







									'''simple docstring'''





									_lowercase			, _lowercase							:		Tuple				=   self.get_dummy_components()
									_lowercase							:		int				=   StableDiffusionPipeline(**UpperCamelCase_		)
									_lowercase							:		Dict				=   sd_pipe.to(UpperCamelCase_		)
									sd_pipe.set_progress_bar_config(disable=UpperCamelCase_		)

									_lowercase			, _lowercase			, _lowercase							:		str				=   self.get_dummy_inputs()

									_lowercase							:		Optional[Any]				=   sd_pipe(**UpperCamelCase_		).images
									_lowercase							:		str				=   original_images[0, -3:, -3:, -1]

									with tempfile.TemporaryDirectory() as tmpdirname:
																LoraLoaderMixin.save_lora_weights(
																    save_directory=UpperCamelCase_      ,		unet_lora_layers=lora_components['''unet_lora_layers''']      ,		text_encoder_lora_layers=lora_components['''text_encoder_lora_layers''']      ,		safe_serialization=UpperCamelCase_      ,		)
																self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase_      ,		'''pytorch_lora_weights.safetensors'''		)		)		)
																sd_pipe.load_lora_weights(UpperCamelCase_		)

									_lowercase							:		int				=   sd_pipe(**UpperCamelCase_		).images
									_lowercase							:		int				=   lora_images[0, -3:, -3:, -1]

									# Outputs shouldn't match.
									self.assertFalse(torch.allclose(torch.from_numpy(UpperCamelCase_		)      ,		torch.from_numpy(UpperCamelCase_		)		)		)
		def 						__lowercase      (						self							:       Union[str, Any]		)      ->			Tuple:







									'''simple docstring'''





									_lowercase			, _lowercase							:		str				=   self.get_dummy_components()
									_lowercase							:		Optional[Any]				=   lora_components['''unet_lora_attn_procs''']
									_lowercase							:		Optional[int]				=   StableDiffusionPipeline(**UpperCamelCase_		)
									_lowercase							:		List[str]				=   sd_pipe.to(UpperCamelCase_		)
									sd_pipe.set_progress_bar_config(disable=UpperCamelCase_		)

									_lowercase			, _lowercase			, _lowercase							:		Any				=   self.get_dummy_inputs()

									_lowercase							:		Any				=   sd_pipe(**UpperCamelCase_		).images
									_lowercase							:		Union[str, Any]				=   original_images[0, -3:, -3:, -1]

									with tempfile.TemporaryDirectory() as tmpdirname:
																_lowercase							:		Union[str, Any]				=   sd_pipe.unet
																unet.set_attn_processor(UpperCamelCase_		)
																unet.save_attn_procs(UpperCamelCase_		)
																self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase_      ,		'''pytorch_lora_weights.bin'''		)		)		)
																sd_pipe.load_lora_weights(UpperCamelCase_		)

									_lowercase							:		List[str]				=   sd_pipe(**UpperCamelCase_		).images
									_lowercase							:		Tuple				=   lora_images[0, -3:, -3:, -1]

									# Outputs shouldn't match.
									self.assertFalse(torch.allclose(torch.from_numpy(UpperCamelCase_		)      ,		torch.from_numpy(UpperCamelCase_		)		)		)
		def 						__lowercase      (						self							:       Union[str, Any]		)      ->			Tuple:







									'''simple docstring'''





									_lowercase			, _lowercase							:		Union[str, Any]				=   self.get_dummy_components()
									_lowercase							:		Any				=   StableDiffusionPipeline(**UpperCamelCase_		)

									_lowercase							:		int				=   self.get_dummy_tokens()

									# inference without lora
									_lowercase							:		Dict				=   pipe.text_encoder(**UpperCamelCase_		)[0]
									assert outputs_without_lora.shape == (1, 77, 32)

									# create lora_attn_procs with zeroed out up.weights
									_lowercase							:		Any				=   create_text_encoder_lora_attn_procs(pipe.text_encoder		)
									set_lora_up_weights(UpperCamelCase_      ,		randn_weight=UpperCamelCase_		)

									# monkey patch
									pipe._modify_text_encoder(UpperCamelCase_		)

									# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
									del text_attn_procs
									gc.collect()

									# inference with lora
									_lowercase							:		Dict				=   pipe.text_encoder(**UpperCamelCase_		)[0]
									assert outputs_with_lora.shape == (1, 77, 32)

									assert torch.allclose(
									    UpperCamelCase_      ,		UpperCamelCase_		), "lora_up_weight are all zero, so the lora outputs should be the same to without lora outputs"

									# create lora_attn_procs with randn up.weights
									_lowercase							:		str				=   create_text_encoder_lora_attn_procs(pipe.text_encoder		)
									set_lora_up_weights(UpperCamelCase_      ,		randn_weight=UpperCamelCase_		)

									# monkey patch
									pipe._modify_text_encoder(UpperCamelCase_		)

									# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
									del text_attn_procs
									gc.collect()

									# inference with lora
									_lowercase							:		List[str]				=   pipe.text_encoder(**UpperCamelCase_		)[0]
									assert outputs_with_lora.shape == (1, 77, 32)

									assert not torch.allclose(
									    UpperCamelCase_      ,		UpperCamelCase_		), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"
		def 						__lowercase      (						self							:       List[Any]		)      ->			Union[str, Any]:







									'''simple docstring'''





									_lowercase			, _lowercase							:		Optional[Any]				=   self.get_dummy_components()
									_lowercase							:		int				=   StableDiffusionPipeline(**UpperCamelCase_		)

									_lowercase							:		Tuple				=   self.get_dummy_tokens()

									# inference without lora
									_lowercase							:		Any				=   pipe.text_encoder(**UpperCamelCase_		)[0]
									assert outputs_without_lora.shape == (1, 77, 32)

									# create lora_attn_procs with randn up.weights
									_lowercase							:		Tuple				=   create_text_encoder_lora_attn_procs(pipe.text_encoder		)
									set_lora_up_weights(UpperCamelCase_      ,		randn_weight=UpperCamelCase_		)

									# monkey patch
									pipe._modify_text_encoder(UpperCamelCase_		)

									# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
									del text_attn_procs
									gc.collect()

									# inference with lora
									_lowercase							:		Tuple				=   pipe.text_encoder(**UpperCamelCase_		)[0]
									assert outputs_with_lora.shape == (1, 77, 32)

									assert not torch.allclose(
									    UpperCamelCase_      ,		UpperCamelCase_		), "lora outputs should be different to without lora outputs"

									# remove monkey patch
									pipe._remove_text_encoder_monkey_patch()

									# inference with removed lora
									_lowercase							:		Dict				=   pipe.text_encoder(**UpperCamelCase_		)[0]
									assert outputs_without_lora_removed.shape == (1, 77, 32)

									assert torch.allclose(
									    UpperCamelCase_      ,		UpperCamelCase_		), "remove lora monkey patch should restore the original outputs"
		def 						__lowercase      (						self							:       Union[str, Any]		)      ->			Optional[Any]:







									'''simple docstring'''





									_lowercase			, _lowercase							:		str				=   self.get_dummy_components()
									_lowercase							:		Optional[Any]				=   StableDiffusionPipeline(**UpperCamelCase_		)
									_lowercase							:		str				=   sd_pipe.to(UpperCamelCase_		)
									sd_pipe.set_progress_bar_config(disable=UpperCamelCase_		)

									_lowercase			, _lowercase			, _lowercase							:		str				=   self.get_dummy_inputs()

									with tempfile.TemporaryDirectory() as tmpdirname:
																LoraLoaderMixin.save_lora_weights(
																    save_directory=UpperCamelCase_      ,		unet_lora_layers=lora_components['''unet_lora_layers''']      ,		text_encoder_lora_layers=lora_components['''text_encoder_lora_layers''']      ,		)
																self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase_      ,		'''pytorch_lora_weights.bin'''		)		)		)
																sd_pipe.load_lora_weights(UpperCamelCase_		)

									_lowercase							:		Dict				=   sd_pipe(**UpperCamelCase_		).images
									_lowercase							:		List[Any]				=   lora_images[0, -3:, -3:, -1]

									_lowercase							:		List[str]				=   sd_pipe(**UpperCamelCase_      ,		cross_attention_kwargs={'''scale''': 0.5}		).images
									_lowercase							:		Tuple				=   lora_images_with_scale[0, -3:, -3:, -1]

									# Outputs shouldn't match.
									self.assertFalse(
									    torch.allclose(torch.from_numpy(UpperCamelCase_		)      ,		torch.from_numpy(UpperCamelCase_		)		)		)
		def 						__lowercase      (						self							:       str		)      ->			Optional[int]:







									'''simple docstring'''





									with tempfile.TemporaryDirectory() as tmpdirname:
																self.create_lora_weight_file(UpperCamelCase_		)

																_lowercase			, _lowercase							:		Union[str, Any]				=   self.get_dummy_components()
																_lowercase							:		Tuple				=   StableDiffusionPipeline(**UpperCamelCase_		)
																_lowercase							:		List[Any]				=   sd_pipe.to(UpperCamelCase_		)
																sd_pipe.set_progress_bar_config(disable=UpperCamelCase_		)

																# check if vanilla attention processors are used
																for _, module in sd_pipe.unet.named_modules():
																							if isinstance(UpperCamelCase_      ,		UpperCamelCase_		):
																														self.assertIsInstance(module.processor      ,		(AttnProcessor, AttnProcessor2_0)		)

            # load LoRA weight file
																sd_pipe.load_lora_weights(UpperCamelCase_		)

																# check if lora attention processors are used
																for _, module in sd_pipe.unet.named_modules():
																							if isinstance(UpperCamelCase_      ,		UpperCamelCase_		):
																														_lowercase							:		Optional[Any]				=   (
																														    LoRAAttnProcessor2_0 if hasattr(UpperCamelCase_      ,		'''scaled_dot_product_attention'''		) else LoRAAttnProcessor
																														)
																														self.assertIsInstance(module.processor      ,		UpperCamelCase_		)
		@unittest.skipIf(torch_device != '''cuda'''      ,		'''This test is supposed to run on GPU'''		)
		def 						__lowercase      (						self							:       str		)      ->			Tuple:







									'''simple docstring'''





									with tempfile.TemporaryDirectory() as tmpdirname:
																self.create_lora_weight_file(UpperCamelCase_		)

																_lowercase			, _lowercase							:		Dict				=   self.get_dummy_components()
																_lowercase							:		Dict				=   StableDiffusionPipeline(**UpperCamelCase_		)
																_lowercase							:		str				=   sd_pipe.to(UpperCamelCase_		)
																sd_pipe.set_progress_bar_config(disable=UpperCamelCase_		)

																# enable XFormers
																sd_pipe.enable_xformers_memory_efficient_attention()

																# check if xFormers attention processors are used
																for _, module in sd_pipe.unet.named_modules():
																							if isinstance(UpperCamelCase_      ,		UpperCamelCase_		):
																														self.assertIsInstance(module.processor      ,		UpperCamelCase_		)

            # load LoRA weight file
																sd_pipe.load_lora_weights(UpperCamelCase_		)

																# check if lora attention processors are used
																for _, module in sd_pipe.unet.named_modules():
																							if isinstance(UpperCamelCase_      ,		UpperCamelCase_		):
																														self.assertIsInstance(module.processor      ,		UpperCamelCase_		)





		@unittest.skipIf(torch_device != '''cuda'''      ,		'''This test is supposed to run on GPU'''		)
		def 						__lowercase      (						self							:       List[str]		)      ->			str:







									'''simple docstring'''





									_lowercase			, _lowercase							:		Dict				=   self.get_dummy_components()
									_lowercase							:		Dict				=   StableDiffusionPipeline(**UpperCamelCase_		)
									_lowercase							:		Union[str, Any]				=   sd_pipe.to(UpperCamelCase_		)
									sd_pipe.set_progress_bar_config(disable=UpperCamelCase_		)

									_lowercase			, _lowercase			, _lowercase							:		Dict				=   self.get_dummy_inputs()

									# enable XFormers
									sd_pipe.enable_xformers_memory_efficient_attention()

									_lowercase							:		Union[str, Any]				=   sd_pipe(**UpperCamelCase_		).images
									_lowercase							:		Dict				=   original_images[0, -3:, -3:, -1]

									with tempfile.TemporaryDirectory() as tmpdirname:
																LoraLoaderMixin.save_lora_weights(
																    save_directory=UpperCamelCase_      ,		unet_lora_layers=lora_components['''unet_lora_layers''']      ,		text_encoder_lora_layers=lora_components['''text_encoder_lora_layers''']      ,		)
																self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase_      ,		'''pytorch_lora_weights.bin'''		)		)		)
																sd_pipe.load_lora_weights(UpperCamelCase_		)

									_lowercase							:		List[str]				=   sd_pipe(**UpperCamelCase_		).images
									_lowercase							:		List[Any]				=   lora_images[0, -3:, -3:, -1]

									# Outputs shouldn't match.
									self.assertFalse(torch.allclose(torch.from_numpy(UpperCamelCase_		)      ,		torch.from_numpy(UpperCamelCase_		)		)		)



