






import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class        lowerCAmelCase_				(      lowercase						,      lowercase						,      unittest.TestCase	):







				"""simple docstring"""




				_snake_case				:   Dict        =			StableDiffusionSAGPipeline
				_snake_case				:   Union[str, Any]        =			TEXT_TO_IMAGE_PARAMS
				_snake_case				:   int        =			TEXT_TO_IMAGE_BATCH_PARAMS
				_snake_case				:   List[str]        =			TEXT_TO_IMAGE_IMAGE_PARAMS
				_snake_case				:   Any        =			TEXT_TO_IMAGE_IMAGE_PARAMS
				_snake_case				:   Union[str, Any]        =			False







				def 				__a   ( self     :List[Any]     ):
					torch.manual_seed(0     )
					UpperCamelCase__ :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  ,     )
					UpperCamelCase__ :Any						=      DDIMScheduler(
					    beta_start=0.0_0085  ,     beta_end=0.012  ,     beta_schedule="""scaled_linear"""  ,     clip_sample=lowerCamelCase__  ,     set_alpha_to_one=lowerCamelCase__  ,     )
					torch.manual_seed(0     )
					UpperCamelCase__ :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  ,     )
					torch.manual_seed(0     )
					UpperCamelCase__ :str						=      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=10_00  ,     )
					UpperCamelCase__ :int						=      CLIPTextModel(lowerCamelCase__     )
					UpperCamelCase__ :str						=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     )

					UpperCamelCase__ :Dict						=      {
					    """unet""": unet,
					    """scheduler""": scheduler,
					    """vae""": vae,
					    """text_encoder""": text_encoder,
					    """tokenizer""": tokenizer,
					    """safety_checker""": None,
					    """feature_extractor""": None,
					}
					return components







				def 				__a   ( self     :List[Any]  ,     lowerCamelCase__     :Dict  ,     lowerCamelCase__     :Union[str, Any]=0     ):
					if str(lowerCamelCase__     ).startswith("""mps"""     ):
						UpperCamelCase__ :Dict						=      torch.manual_seed(lowerCamelCase__     )
					else:
						UpperCamelCase__ :int						=      torch.Generator(device=lowerCamelCase__     ).manual_seed(lowerCamelCase__     )
					UpperCamelCase__ :Any						=      {
					    """prompt""": """.""",
					    """generator""": generator,
					    """num_inference_steps""": 2,
					    """guidance_scale""": 1.0,
					    """sag_scale""": 1.0,
					    """output_type""": """numpy""",
					}
					return inputs







				def 				__a   ( self     :Optional[Any]     ):
					super().test_inference_batch_single_identical(expected_max_diff=3e-3     )

@slow
@require_torch_gpu
class        lowerCAmelCase_				(      unittest.TestCase	):







				"""simple docstring"""




				def 				__a   ( self     :List[Any]     ):
					# clean up the VRAM after each test
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()







				def 				__a   ( self     :int     ):
					UpperCamelCase__ :int						=      StableDiffusionSAGPipeline.from_pretrained("""CompVis/stable-diffusion-v1-4"""     )
					UpperCamelCase__ :Optional[int]						=      sag_pipe.to(lowerCamelCase__     )
					sag_pipe.set_progress_bar_config(disable=lowerCamelCase__     )

					UpperCamelCase__ :str						=      """."""
					UpperCamelCase__ :Dict						=      torch.manual_seed(0     )
					UpperCamelCase__ :str						=      sag_pipe(
					    [prompt]  ,     generator=lowerCamelCase__  ,     guidance_scale=7.5  ,     sag_scale=1.0  ,     num_inference_steps=20  ,     output_type="""np"""     )

					UpperCamelCase__ :Optional[int]						=      output.images

					UpperCamelCase__ :Optional[Any]						=      image[0, -3:, -3:, -1]

					assert image.shape == (1, 5_12, 5_12, 3)
					UpperCamelCase__ :Optional[Any]						=      np.array([0.1568, 0.1738, 0.1695, 0.1693, 0.1507, 0.1705, 0.1547, 0.1751, 0.1949]     )

					assert np.abs(image_slice.flatten() - expected_slice     ).max() < 5e-2







				def 				__a   ( self     :Dict     ):
					UpperCamelCase__ :Tuple						=      StableDiffusionSAGPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""     )
					UpperCamelCase__ :Dict						=      sag_pipe.to(lowerCamelCase__     )
					sag_pipe.set_progress_bar_config(disable=lowerCamelCase__     )

					UpperCamelCase__ :Dict						=      """."""
					UpperCamelCase__ :List[str]						=      torch.manual_seed(0     )
					UpperCamelCase__ :List[Any]						=      sag_pipe(
					    [prompt]  ,     generator=lowerCamelCase__  ,     guidance_scale=7.5  ,     sag_scale=1.0  ,     num_inference_steps=20  ,     output_type="""np"""     )

					UpperCamelCase__ :List[str]						=      output.images

					UpperCamelCase__ :Tuple						=      image[0, -3:, -3:, -1]

					assert image.shape == (1, 5_12, 5_12, 3)
					UpperCamelCase__ :List[str]						=      np.array([0.3459, 0.2876, 0.2537, 0.3002, 0.2671, 0.2160, 0.3026, 0.2262, 0.2371]     )

					assert np.abs(image_slice.flatten() - expected_slice     ).max() < 5e-2







				def 				__a   ( self     :str     ):
					UpperCamelCase__ :str						=      StableDiffusionSAGPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""     )
					UpperCamelCase__ :Tuple						=      sag_pipe.to(lowerCamelCase__     )
					sag_pipe.set_progress_bar_config(disable=lowerCamelCase__     )

					UpperCamelCase__ :str						=      """."""
					UpperCamelCase__ :int						=      torch.manual_seed(0     )
					UpperCamelCase__ :int						=      sag_pipe(
					    [prompt]  ,     width=7_68  ,     height=5_12  ,     generator=lowerCamelCase__  ,     guidance_scale=7.5  ,     sag_scale=1.0  ,     num_inference_steps=20  ,     output_type="""np"""  ,     )

					UpperCamelCase__ :str						=      output.images

					assert image.shape == (1, 5_12, 7_68, 3)


import logging

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

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEncoder,
    BertModel,
    BertPreTrainedModel,
)


__snake_case				:      Optional[int]   		=			logging.getLogger(__name__)



class        A__(a_							):
						"""simple docstring"""






						def 			UpperCamelCase__					(    self  ,      _lowercase  ,      _lowercase  ,      _lowercase=None  ,      _lowercase=None    )     -> Tuple:
									a_						:						str 							= self.layer[current_layer](_lowercase  ,      _lowercase  ,      head_mask[current_layer]    )

									a_						:						str 							= layer_outputs[0]

									return hidden_states







@add_start_docstrings(
    '''The bare Bert Model transformer with PABEE outputting raw hidden-states without any specific head on top.''',  a_,  )
class        A__(a_							):
						"""simple docstring"""






						def __init__(    self  ,      _lowercase    )     -> List[str]:
									super().__init__(_lowercase    )

									a_						:						Tuple 							= BertEncoderWithPabee(_lowercase    )

									self.init_weights()
									a_						:						int 							= 0
									a_						:						Any 							= 0
									a_						:						Tuple 							= 0

									a_						:						Optional[int] 							= 0





						def 			UpperCamelCase__					(    self  ,      _lowercase    )     -> Tuple:
									a_						:						Dict 							= threshold





						def 			UpperCamelCase__					(    self  ,      _lowercase    )     -> List[Any]:
									a_						:						Optional[int] 							= patience





						def 			UpperCamelCase__					(    self    )     -> Dict:
									a_						:						str 							= 0
									a_						:						Optional[int] 							= 0





						def 			UpperCamelCase__					(    self    )     -> List[Any]:
									a_						:						Union[str, Any] 							= self.inference_layers_num / self.inference_instances_num
									a_						:						Optional[int] 							= (
									    F'''*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ='''
									    F''' {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***'''
									)
									print(_lowercase    )





						@add_start_docstrings_to_model_forward(_lowercase    )
						def 			UpperCamelCase__					(    self  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=False  ,      )     -> str:

									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_						:						Dict 							= input_ids.size()
									elif inputs_embeds is not None:
												a_						:						Dict 							= inputs_embeds.size()[:-1]
									else:
												raise ValueError("""You have to specify either input_ids or inputs_embeds"""    )

									a_						:						Optional[Any] 							= input_ids.device if input_ids is not None else inputs_embeds.device

									if attention_mask is None:
												a_						:						Tuple 							= torch.ones(_lowercase  ,      device=_lowercase    )
									if token_type_ids is None:
												a_						:						List[str] 							= torch.zeros(_lowercase  ,      dtype=torch.long  ,      device=_lowercase    )

									# We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
									# ourselves in which case we just need to make it broadcastable to all heads.
									a_						:						torch.Tensor 							= self.get_extended_attention_mask(_lowercase  ,      _lowercase  ,      _lowercase    )

									# If a 2D ou 3D attention mask is provided for the cross-attention
									# we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
									if self.config.is_decoder and encoder_hidden_states is not None:
												a_							,					a_							,					a_						:						int 							= encoder_hidden_states.size()
												a_						:						List[str] 							= (encoder_batch_size, encoder_sequence_length)
												if encoder_attention_mask is None:
															a_						:						Tuple 							= torch.ones(_lowercase  ,      device=_lowercase    )
												a_						:						List[Any] 							= self.invert_attention_mask(_lowercase    )
									else:
												a_						:						Optional[Any] 							= None

									# Prepare head mask if needed
									# 1.0 in head_mask indicate we keep the head
									# attention_probs has shape bsz x n_heads x N x N
									# input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
									# and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
									a_						:						List[Any] 							= self.get_head_mask(_lowercase  ,      self.config.num_hidden_layers    )

									a_						:						List[str] 							= self.embeddings(
									    input_ids=_lowercase  ,      position_ids=_lowercase  ,      token_type_ids=_lowercase  ,      inputs_embeds=_lowercase    )
									a_						:						List[Any] 							= embedding_output

									if self.training:
												a_						:						Any 							= []
												for i in range(self.config.num_hidden_layers    ):
															a_						:						int 							= self.encoder.adaptive_forward(
															    _lowercase  ,      current_layer=_lowercase  ,      attention_mask=_lowercase  ,      head_mask=_lowercase    )

															a_						:						List[Any] 							= self.pooler(_lowercase    )
															a_						:						Optional[int] 							= output_layers[i](output_dropout(_lowercase    )    )
															res.append(_lowercase    )
									elif self.patience == 0:  # Use all layers for inference
												a_						:						Union[str, Any] 							= self.encoder(
												    _lowercase  ,      attention_mask=_lowercase  ,      head_mask=_lowercase  ,      encoder_hidden_states=_lowercase  ,      encoder_attention_mask=_lowercase  ,      )
												a_						:						Union[str, Any] 							= self.pooler(encoder_outputs[0]    )
												a_						:						List[str] 							= [output_layers[self.config.num_hidden_layers - 1](_lowercase    )]
									else:
												a_						:						Any 							= 0
												a_						:						Dict 							= None
												a_						:						Tuple 							= 0
												for i in range(self.config.num_hidden_layers    ):
															calculated_layer_num += 1
															a_						:						Optional[Any] 							= self.encoder.adaptive_forward(
															    _lowercase  ,      current_layer=_lowercase  ,      attention_mask=_lowercase  ,      head_mask=_lowercase    )

															a_						:						Optional[int] 							= self.pooler(_lowercase    )
															a_						:						int 							= output_layers[i](_lowercase    )
															if regression:
																		a_						:						Dict 							= logits.detach()
																		if patient_result is not None:
																					a_						:						Optional[Any] 							= patient_result.detach()
																		if (patient_result is not None) and torch.abs(patient_result - labels    ) < self.regression_threshold:
																					patient_counter += 1
																		else:
																					a_						:						Optional[Any] 							= 0
															else:
																		a_						:						str 							= logits.detach().argmax(dim=1    )
																		if patient_result is not None:
																					a_						:						str 							= patient_result.detach().argmax(dim=1    )
																		if (patient_result is not None) and torch.all(labels.eq(_lowercase    )    ):
																					patient_counter += 1
																		else:
																					a_						:						Tuple 							= 0

															a_						:						Union[str, Any] 							= logits
															if patient_counter == self.patience:
																		break
												a_						:						str 							= [patient_result]
												self.inference_layers_num += calculated_layer_num
												self.inference_instances_num += 1

									return res







@add_start_docstrings(
    '''Bert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of
    the pooled output) e.g. for GLUE tasks. ''',  a_,  )
class        A__(a_							):
						"""simple docstring"""






						def __init__(    self  ,      _lowercase    )     -> str:
									super().__init__(_lowercase    )
									a_						:						str 							= config.num_labels

									a_						:						Optional[Any] 							= BertModelWithPabee(_lowercase    )
									a_						:						int 							= nn.Dropout(config.hidden_dropout_prob    )
									a_						:						str 							= nn.ModuleList(
									    [nn.Linear(config.hidden_size  ,      self.config.num_labels    ) for _ in range(config.num_hidden_layers    )]    )

									self.init_weights()





						@add_start_docstrings_to_model_forward(_lowercase    )
						def 			UpperCamelCase__					(    self  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      )     -> Tuple:
									a_						:						Optional[Any] 							= self.bert(
									    input_ids=_lowercase  ,      attention_mask=_lowercase  ,      token_type_ids=_lowercase  ,      position_ids=_lowercase  ,      head_mask=_lowercase  ,      inputs_embeds=_lowercase  ,      output_dropout=self.dropout  ,      output_layers=self.classifiers  ,      regression=self.num_labels == 1  ,      )

									a_						:						Optional[Any] 							= (logits[-1],)

									if labels is not None:
												a_						:						int 							= None
												a_						:						Union[str, Any] 							= 0
												for ix, logits_item in enumerate(_lowercase    ):
															if self.num_labels == 1:
																		#  We are doing regression
																		a_						:						Any 							= MSELoss()
																		a_						:						Union[str, Any] 							= loss_fct(logits_item.view(-1    )  ,      labels.view(-1    )    )
															else:
																		a_						:						Any 							= CrossEntropyLoss()
																		a_						:						int 							= loss_fct(logits_item.view(-1  ,      self.num_labels    )  ,      labels.view(-1    )    )
															if total_loss is None:
																		a_						:						str 							= loss
															else:
																		total_loss += loss * (ix + 1)
															total_weights += ix + 1
												a_						:						Any 							= (total_loss / total_weights,) + outputs

									return outputs


