



'''simple docstring'''






import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEImgaImgPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import floats_tensor, load_image, load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference

class        SCREAMING_SNAKE_CASE(							__lowercase							,	unittest.TestCase    ):







	"""simple docstring"""






	lowerCamelCase__    				= ShapEImgaImgPipeline
	lowerCamelCase__    				= ["""image"""]
	lowerCamelCase__    				= ["""image"""]
	lowerCamelCase__    				= [
	    """num_images_per_prompt""",
	    """num_inference_steps""",
	    """generator""",
	    """latents""",
	    """guidance_scale""",
	    """frame_size""",
	    """output_type""",
	    """return_dict""",
	]
	lowerCamelCase__    				= False







	@property
	def   A			(		self    :  List[str]	)    ->	Dict:
								return 32







	@property
	def   A			(		self    :  Tuple	)    ->	Tuple:
								return 32







	@property
	def   A			(		self    :  str	)    ->	Tuple:
								return self.time_input_dim * 4







	@property
	def   A			(		self    :  str	)    ->	Optional[int]:
								return 8







	@property
	def   A			(		self    :  Dict	)    ->	List[Any]:
								torch.manual_seed(0	)
								UpperCAmelCase				: Tuple											=	CLIPVisionConfig(
								    hidden_size=self.text_embedder_hidden_size	,    image_size=64	,    projection_dim=self.text_embedder_hidden_size	,    intermediate_size=37	,    num_attention_heads=4	,    num_channels=3	,    num_hidden_layers=5	,    patch_size=1	,    )

								UpperCAmelCase				: str											=	CLIPVisionModel(__snake_case	)
								return model







	@property
	def   A			(		self    :  Optional[int]	)    ->	Optional[Any]:
								UpperCAmelCase				: Any											=	CLIPImageProcessor(
								    crop_size=224	,    do_center_crop=__snake_case	,    do_normalize=__snake_case	,    do_resize=__snake_case	,    image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]	,    image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]	,    resample=3	,    size=224	,    )

								return image_processor







	@property
	def   A			(		self    :  Dict	)    ->	Optional[Any]:
								torch.manual_seed(0	)

								UpperCAmelCase				: Tuple											=	{
								    "num_attention_heads": 2,
								    "attention_head_dim": 16,
								    "embedding_dim": self.time_input_dim,
								    "num_embeddings": 32,
								    "embedding_proj_dim": self.text_embedder_hidden_size,
								    "time_embed_dim": self.time_embed_dim,
								    "num_layers": 1,
								    "clip_embed_dim": self.time_input_dim * 2,
								    "additional_embeddings": 0,
								    "time_embed_act_fn": "gelu",
								    "norm_in_type": "layer",
								    "embedding_proj_norm_type": "layer",
								    "encoder_hid_proj_type": None,
								    "added_emb_type": None,
								}

								UpperCAmelCase				: Union[str, Any]											=	PriorTransformer(**__snake_case	)
								return model







	@property
	def   A			(		self    :  str	)    ->	Tuple:
								torch.manual_seed(0	)

								UpperCAmelCase				: Optional[Any]											=	{
								    "param_shapes": (
								        (self.renderer_dim, 93),
								        (self.renderer_dim, 8),
								        (self.renderer_dim, 8),
								        (self.renderer_dim, 8),
								    ),
								    "d_latent": self.time_input_dim,
								    "d_hidden": self.renderer_dim,
								    "n_output": 12,
								    "background": (
								        0.1,
								        0.1,
								        0.1,
								    ),
								}
								UpperCAmelCase				: Any											=	ShapERenderer(**__snake_case	)
								return model







	def   A			(		self    :  int	)    ->	Dict:
								UpperCAmelCase				: Optional[int]											=	self.dummy_prior
								UpperCAmelCase				: Tuple											=	self.dummy_image_encoder
								UpperCAmelCase				: List[str]											=	self.dummy_image_processor
								UpperCAmelCase				: Optional[Any]											=	self.dummy_renderer

								UpperCAmelCase				: str											=	HeunDiscreteScheduler(
								    beta_schedule='''exp'''	,    num_train_timesteps=1024	,    prediction_type='''sample'''	,    use_karras_sigmas=__snake_case	,    clip_sample=__snake_case	,    clip_sample_range=1.0	,    )
								UpperCAmelCase				: List[Any]											=	{
								    "prior": prior,
								    "image_encoder": image_encoder,
								    "image_processor": image_processor,
								    "renderer": renderer,
								    "scheduler": scheduler,
								}

								return components







	def   A			(		self    :  int	,    __snake_case    :  List[str]	,    __snake_case    :  List[Any]=0	)    ->	Optional[Any]:
								UpperCAmelCase				: Any											=	floats_tensor((1, 3, 64, 64)	,    rng=random.Random(__snake_case	)	).to(__snake_case	)

								if str(__snake_case	).startswith('''mps'''	):
															UpperCAmelCase				: str											=	torch.manual_seed(__snake_case	)
								else:
															UpperCAmelCase				: str											=	torch.Generator(device=__snake_case	).manual_seed(__snake_case	)
								UpperCAmelCase				: Any											=	{
								    "image": input_image,
								    "generator": generator,
								    "num_inference_steps": 1,
								    "frame_size": 32,
								    "output_type": "np",
								}
								return inputs







	def   A			(		self    :  Any	)    ->	Optional[int]:
								UpperCAmelCase				: str											=	"cpu"

								UpperCAmelCase				: Optional[Any]											=	self.get_dummy_components()

								UpperCAmelCase				: Union[str, Any]											=	self.pipeline_class(**__snake_case	)
								UpperCAmelCase				: List[Any]											=	pipe.to(__snake_case	)

								pipe.set_progress_bar_config(disable=__snake_case	)

								UpperCAmelCase				: List[str]											=	pipe(**self.get_dummy_inputs(__snake_case	)	)
								UpperCAmelCase				: int											=	output.images[0]
								UpperCAmelCase				: List[str]											=	image[0, -3:, -3:, -1]

								assert image.shape == (20, 32, 32, 3)

								UpperCAmelCase				: Optional[int]											=	np.array(
								    [
								        0.00_03_92_16,
								        0.00_03_92_16,
								        0.00_03_92_16,
								        0.00_03_92_16,
								        0.00_03_92_16,
								        0.00_03_92_16,
								        0.00_03_92_16,
								        0.00_03_92_16,
								        0.00_03_92_16,
								    ]	)

								assert np.abs(image_slice.flatten() - expected_slice	).max() < 1E-2







	def   A			(		self    :  str	)    ->	Dict:
								# NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
								self._test_inference_batch_consistent(batch_sizes=[1, 2]	)







	def   A			(		self    :  List[Any]	)    ->	Union[str, Any]:
								UpperCAmelCase				: List[Any]											=	torch_device == "cpu"
								UpperCAmelCase				: Any											=	True
								self._test_inference_batch_single_identical(
								    batch_size=2	,    test_max_difference=__snake_case	,    relax_max_difference=__snake_case	,    )







	def   A			(		self    :  List[str]	)    ->	Any:
								UpperCAmelCase				: Dict											=	self.get_dummy_components()
								UpperCAmelCase				: Optional[Any]											=	self.pipeline_class(**__snake_case	)
								UpperCAmelCase				: Any											=	pipe.to(__snake_case	)
								pipe.set_progress_bar_config(disable=__snake_case	)

								UpperCAmelCase				: Union[str, Any]											=	1
								UpperCAmelCase				: Optional[Any]											=	2

								UpperCAmelCase				: List[Any]											=	self.get_dummy_inputs(__snake_case	)

								for key in inputs.keys():
															if key in self.batch_params:
																						UpperCAmelCase				: Any											=	batch_size * [inputs[key]]

								UpperCAmelCase				: Any											=	pipe(**__snake_case	,    num_images_per_prompt=__snake_case	)[0]

								assert images.shape[0] == batch_size * num_images_per_prompt

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







	"""simple docstring"""






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







	def   A			(		self    :  Any	)    ->	Any:
								UpperCAmelCase				: List[str]											=	load_image(
								    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/shap_e/corgi.png'''	)
								UpperCAmelCase				: Dict											=	load_numpy(
								    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
								    '''/shap_e/test_shap_e_img2img_out.npy'''	)
								UpperCAmelCase				: List[str]											=	ShapEImgaImgPipeline.from_pretrained('''openai/shap-e-img2img'''	)
								UpperCAmelCase				: int											=	pipe.to(__snake_case	)
								pipe.set_progress_bar_config(disable=__snake_case	)

								UpperCAmelCase				: Tuple											=	torch.Generator(device=__snake_case	).manual_seed(0	)

								UpperCAmelCase				: Union[str, Any]											=	pipe(
								    __snake_case	,    generator=__snake_case	,    guidance_scale=3.0	,    num_inference_steps=64	,    frame_size=64	,    output_type='''np'''	,    ).images[0]

								assert images.shape == (20, 64, 64, 3)

								assert_mean_pixel_difference(__snake_case	,    __snake_case	)










'''simple docstring'''






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,
)


UpperCamelCase__:     List[str]										=    logging.getLogger(__name__)

class        SCREAMING_SNAKE_CASE(							A__    ):







	"""simple docstring"""






	def   A			(		self    :  str	,    __snake_case    :  Tuple	,    __snake_case    :  int	,    __snake_case    :  List[Any]=None	,    __snake_case    :  List[Any]=None	)    ->	Optional[int]:
								UpperCAmelCase				: Union[str, Any]											=	self.layer[current_layer](__snake_case	,    __snake_case	,    head_mask[current_layer]	)

								UpperCAmelCase				: Optional[Any]											=	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        SCREAMING_SNAKE_CASE(							A__    ):







	"""simple docstring"""






	def __init__(		self    :  int	,    __snake_case    :  Any	)    ->	Any:
								super().__init__(__snake_case	)

								UpperCAmelCase				: Any											=	BertEncoderWithPabee(__snake_case	)

								self.init_weights()
								UpperCAmelCase				: Optional[Any]											=	0
								UpperCAmelCase				: List[str]											=	0
								UpperCAmelCase				: Tuple											=	0

								UpperCAmelCase				: List[str]											=	0







	def   A			(		self    :  List[str]	,    __snake_case    :  List[str]	)    ->	List[Any]:
								UpperCAmelCase				: Union[str, Any]											=	threshold







	def   A			(		self    :  str	,    __snake_case    :  List[str]	)    ->	Optional[Any]:
								UpperCAmelCase				: List[str]											=	patience







	def   A			(		self    :  Dict	)    ->	str:
								UpperCAmelCase				: List[Any]											=	0
								UpperCAmelCase				: Dict											=	0







	def   A			(		self    :  Optional[Any]	)    ->	Optional[int]:
								UpperCAmelCase				: Dict											=	self.inference_layers_num / self.inference_instances_num
								UpperCAmelCase				: 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(__snake_case	)







	@add_start_docstrings_to_model_forward(__snake_case	)
	def   A			(		self    :  Union[str, Any]	,    __snake_case    :  Optional[Any]=None	,    __snake_case    :  str=None	,    __snake_case    :  int=None	,    __snake_case    :  Optional[Any]=None	,    __snake_case    :  Tuple=None	,    __snake_case    :  Tuple=None	,    __snake_case    :  Tuple=None	,    __snake_case    :  Optional[Any]=None	,    __snake_case    :  Union[str, Any]=None	,    __snake_case    :  Optional[int]=None	,    __snake_case    :  Optional[int]=False	,    )    ->	List[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:
															UpperCAmelCase				: Union[str, Any]											=	input_ids.size()
								elif inputs_embeds is not None:
															UpperCAmelCase				: List[Any]											=	inputs_embeds.size()[:-1]
								else:
															raise ValueError('''You have to specify either input_ids or inputs_embeds'''	)

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

								if attention_mask is None:
															UpperCAmelCase				: Optional[Any]											=	torch.ones(__snake_case	,    device=__snake_case	)
								if token_type_ids is None:
															UpperCAmelCase				: Tuple											=	torch.zeros(__snake_case	,    dtype=torch.long	,    device=__snake_case	)

								# 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.
								UpperCAmelCase				: torch.Tensor											=	self.get_extended_attention_mask(__snake_case	,    __snake_case	,    __snake_case	)

								# 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:
															UpperCAmelCase							,	UpperCAmelCase							,	UpperCAmelCase				: Any											=	encoder_hidden_states.size()
															UpperCAmelCase				: List[Any]											=	(encoder_batch_size, encoder_sequence_length)
															if encoder_attention_mask is None:
																						UpperCAmelCase				: str											=	torch.ones(__snake_case	,    device=__snake_case	)
															UpperCAmelCase				: Tuple											=	self.invert_attention_mask(__snake_case	)
								else:
															UpperCAmelCase				: List[str]											=	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]
								UpperCAmelCase				: List[Any]											=	self.get_head_mask(__snake_case	,    self.config.num_hidden_layers	)

								UpperCAmelCase				: Any											=	self.embeddings(
								    input_ids=__snake_case	,    position_ids=__snake_case	,    token_type_ids=__snake_case	,    inputs_embeds=__snake_case	)
								UpperCAmelCase				: Optional[int]											=	embedding_output

								if self.training:
															UpperCAmelCase				: str											=	[]
															for i in range(self.config.num_hidden_layers	):
																						UpperCAmelCase				: Optional[int]											=	self.encoder.adaptive_forward(
																						    __snake_case	,    current_layer=__snake_case	,    attention_mask=__snake_case	,    head_mask=__snake_case	)

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

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

																						UpperCAmelCase				: List[str]											=	logits
																						if patient_counter == self.patience:
																													break
															UpperCAmelCase				: List[Any]											=	[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        SCREAMING_SNAKE_CASE(							A__    ):







	"""simple docstring"""






	def __init__(		self    :  str	,    __snake_case    :  str	)    ->	Optional[Any]:
								super().__init__(__snake_case	)
								UpperCAmelCase				: Optional[int]											=	config.num_labels

								UpperCAmelCase				: Optional[Any]											=	BertModelWithPabee(__snake_case	)
								UpperCAmelCase				: Tuple											=	nn.Dropout(config.hidden_dropout_prob	)
								UpperCAmelCase				: Dict											=	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(__snake_case	)
	def   A			(		self    :  Any	,    __snake_case    :  Dict=None	,    __snake_case    :  str=None	,    __snake_case    :  List[Any]=None	,    __snake_case    :  List[Any]=None	,    __snake_case    :  Union[str, Any]=None	,    __snake_case    :  int=None	,    __snake_case    :  str=None	,    )    ->	Union[str, Any]:
								UpperCAmelCase				: Optional[int]											=	self.bert(
								    input_ids=__snake_case	,    attention_mask=__snake_case	,    token_type_ids=__snake_case	,    position_ids=__snake_case	,    head_mask=__snake_case	,    inputs_embeds=__snake_case	,    output_dropout=self.dropout	,    output_layers=self.classifiers	,    regression=self.num_labels == 1	,    )

								UpperCAmelCase				: List[str]											=	(logits[-1],)

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

								return outputs





