






'''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.albert.modeling_albert import (
    ALBERT_INPUTS_DOCSTRING,
    ALBERT_START_DOCSTRING,
    AlbertModel,
    AlbertPreTrainedModel,
    AlbertTransformer,
)


UpperCamelCase_					=				logging.getLogger(__name__)




class 					__UpperCAmelCase (		__A	):

      '''simple docstring'''





      def 							lowerCamelCase     (					self  ,		_UpperCAmelCase  ,		_UpperCAmelCase  ,		_UpperCAmelCase=None  ,		_UpperCAmelCase=None						):

             if current_layer == 0:
                    UpperCAmelCase__	:   Optional[Any]												=	self.embedding_hidden_mapping_in(__snake_case						)
             else:
                    UpperCAmelCase__	:   Optional[int]												=	hidden_states[0]

             UpperCAmelCase__	:   Dict												=	int(self.config.num_hidden_layers / self.config.num_hidden_groups						)

             # Index of the hidden group
             UpperCAmelCase__	:   Tuple												=	int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)						)

             UpperCAmelCase__	:   Tuple												=	self.albert_layer_groups[group_idx](
                 __snake_case  ,		__snake_case  ,		head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]  ,		)
             UpperCAmelCase__	:   str												=	layer_group_output[0]

             return (hidden_states,)




@add_start_docstrings(
    "The bare ALBERT Model transformer with PABEE outputting raw hidden-states without any specific head on top."      ,				__A      ,				)
class 					__UpperCAmelCase (		__A	):

      '''simple docstring'''





      def __init__(					self  ,		_UpperCAmelCase						):
             super().__init__(__snake_case						)

             UpperCAmelCase__	:   Any												=	AlbertTransformerWithPabee(__snake_case						)

             self.init_weights()
             UpperCAmelCase__	:   int												=	0
             UpperCAmelCase__	:   Optional[int]												=	0
             UpperCAmelCase__	:   List[str]												=	0

             UpperCAmelCase__	:   List[Any]												=	0




      def 							lowerCamelCase     (					self  ,		_UpperCAmelCase						):
             UpperCAmelCase__	:   Optional[int]												=	threshold




      def 							lowerCamelCase     (					self  ,		_UpperCAmelCase						):
             UpperCAmelCase__	:   Tuple												=	patience




      def 							lowerCamelCase     (					self						):
             UpperCAmelCase__	:   List[Any]												=	0
             UpperCAmelCase__	:   Dict												=	0




      def 							lowerCamelCase     (					self						):
             UpperCAmelCase__	:   List[str]												=	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 							lowerCamelCase     (					self  ,		_UpperCAmelCase=None  ,		_UpperCAmelCase=None  ,		_UpperCAmelCase=None  ,		_UpperCAmelCase=None  ,		_UpperCAmelCase=None  ,		_UpperCAmelCase=None  ,		_UpperCAmelCase=None  ,		_UpperCAmelCase=None  ,		_UpperCAmelCase=False  ,		):

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

             UpperCAmelCase__	:   str												=	input_ids.device if input_ids is not None else inputs_embeds.device

             if attention_mask is None:
                    UpperCAmelCase__	:   Tuple												=	torch.ones(__snake_case  ,		device=__snake_case						)
             if token_type_ids is None:
                    UpperCAmelCase__	:   int												=	torch.zeros(__snake_case  ,		dtype=torch.long  ,		device=__snake_case						)

             UpperCAmelCase__	:   Dict												=	attention_mask.unsqueeze(1						).unsqueeze(2						)
             UpperCAmelCase__	:   Union[str, Any]												=	extended_attention_mask.to(dtype=self.dtype						)  # fp16 compatibility
             UpperCAmelCase__	:   int												=	(1.0 - extended_attention_mask) * -1_0000.0
             UpperCAmelCase__	:   Tuple												=	self.get_head_mask(__snake_case  ,		self.config.num_hidden_layers						)

             UpperCAmelCase__	:   List[str]												=	self.embeddings(
                 __snake_case  ,		position_ids=__snake_case  ,		token_type_ids=__snake_case  ,		inputs_embeds=__snake_case						)
             UpperCAmelCase__	:   str												=	embedding_output

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

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

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

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

             return res




@add_start_docstrings(
    "Albert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of\n    the pooled output) e.g. for GLUE tasks. "      ,				__A      ,				)
class 					__UpperCAmelCase (		__A	):

      '''simple docstring'''





      def __init__(					self  ,		_UpperCAmelCase						):
             super().__init__(__snake_case						)
             UpperCAmelCase__	:   str												=	config.num_labels

             UpperCAmelCase__	:   List[str]												=	AlbertModelWithPabee(__snake_case						)
             UpperCAmelCase__	:   int												=	nn.Dropout(config.classifier_dropout_prob						)
             UpperCAmelCase__	:   int												=	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 							lowerCamelCase     (					self  ,		_UpperCAmelCase=None  ,		_UpperCAmelCase=None  ,		_UpperCAmelCase=None  ,		_UpperCAmelCase=None  ,		_UpperCAmelCase=None  ,		_UpperCAmelCase=None  ,		_UpperCAmelCase=None  ,		):
             UpperCAmelCase__	:   List[Any]												=	self.albert(
                 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__	:   Dict												=	(logits[-1],)

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

             return outputs






import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPImageProcessor, CLIPVisionConfig

from diffusers import AutoencoderKL, PaintByExamplePipeline, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.paint_by_example import PaintByExampleImageEncoder
from diffusers.utils import floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import IMAGE_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, IMAGE_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()






class 							_a(     __A						, unittest.TestCase     ):
		lowerCamelCase__  :Optional[int] 					=		PaintByExamplePipeline
		lowerCamelCase__  :Union[str, Any] 					=		IMAGE_GUIDED_IMAGE_INPAINTING_PARAMS
		lowerCamelCase__  :Optional[Any] 					=		IMAGE_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
		lowerCamelCase__  :int 					=		frozenset([]     )  # TO_DO: update the image_prams once refactored VaeImageProcessor.preprocess







		def 							lowercase							(     self					)   ->   Dict:


							'''simple docstring'''

							torch.manual_seed(0					)
							_snake_case							:       Union[str, Any]												= UNetaDConditionModel(
							    block_out_channels=(3_2, 6_4)			,    layers_per_block=2			,    sample_size=3_2			,    in_channels=9			,    out_channels=4			,    down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")			,    up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")			,    cross_attention_dim=3_2			,    )
							_snake_case							:       int												= PNDMScheduler(skip_prk_steps=__snake_case					)
							torch.manual_seed(0					)
							_snake_case							:       Any												= AutoencoderKL(
							    block_out_channels=[3_2, 6_4]			,    in_channels=3			,    out_channels=3			,    down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]			,    up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]			,    latent_channels=4			,    )
							torch.manual_seed(0					)
							_snake_case							:       Dict												= CLIPVisionConfig(
							    hidden_size=3_2			,    projection_dim=3_2			,    intermediate_size=3_7			,    layer_norm_eps=1E-05			,    num_attention_heads=4			,    num_hidden_layers=5			,    image_size=3_2			,    patch_size=4			,    )
							_snake_case							:       Any												= PaintByExampleImageEncoder(__snake_case			,    proj_size=3_2					)
							_snake_case							:       Dict												= CLIPImageProcessor(crop_size=3_2			,    size=3_2					)

							_snake_case							:       Any												= {
							    "unet": unet,
							    "scheduler": scheduler,
							    "vae": vae,
							    "image_encoder": image_encoder,
							    "safety_checker": None,
							    "feature_extractor": feature_extractor,
							}
							return components







		def 							lowercase							(     self			,    __snake_case					)   ->   Tuple:


							'''simple docstring'''

							_snake_case							:       List[str]												= np.array(image.convert("RGB"					)					)
							_snake_case							:       str												= image[None].transpose(0			,    3			,    1			,    2					)
							_snake_case							:       Dict												= torch.from_numpy(__snake_case					).to(dtype=torch.floataa					) / 1_27.5 - 1.0
							return image







		def 							lowercase							(     self			,    __snake_case="cpu"			,    __snake_case=0					)   ->   Any:


							'''simple docstring'''

							_snake_case							:       Tuple												= floats_tensor((1, 3, 3_2, 3_2)			,    rng=random.Random(__snake_case					)					).to(__snake_case					)
							_snake_case							:       Union[str, Any]												= image.cpu().permute(0			,    2			,    3			,    1					)[0]
							_snake_case							:       List[Any]												= Image.fromarray(np.uinta(__snake_case					)					).convert("RGB"					).resize((6_4, 6_4)					)
							_snake_case							:       List[Any]												= Image.fromarray(np.uinta(image + 4					)					).convert("RGB"					).resize((6_4, 6_4)					)
							_snake_case							:       Any												= Image.fromarray(np.uinta(__snake_case					)					).convert("RGB"					).resize((3_2, 3_2)					)

							if str(__snake_case					).startswith("mps"					):
												_snake_case							:       Union[str, Any]												= torch.manual_seed(__snake_case					)
							else:
												_snake_case							:       List[Any]												= torch.Generator(device=__snake_case					).manual_seed(__snake_case					)
							_snake_case							:       Optional[int]												= {
							    "example_image": example_image,
							    "image": init_image,
							    "mask_image": mask_image,
							    "generator": generator,
							    "num_inference_steps": 2,
							    "guidance_scale": 6.0,
							    "output_type": "numpy",
							}
							return inputs







		def 							lowercase							(     self					)   ->   int:


							'''simple docstring'''

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

							# make sure here that pndm scheduler skips prk
							_snake_case							:       int												= PaintByExamplePipeline(**__snake_case					)
							_snake_case							:       List[Any]												= pipe.to("cpu"					)
							pipe.set_progress_bar_config(disable=__snake_case					)

							_snake_case							:       List[str]												= self.get_dummy_inputs()
							_snake_case							:       int												= pipe(**__snake_case					)
							_snake_case							:       Dict												= output.images

							_snake_case							:       Optional[int]												= image[0, -3:, -3:, -1]

							assert image.shape == (1, 6_4, 6_4, 3)
							_snake_case							:       int												= np.array([0.46_86, 0.56_87, 0.40_07, 0.52_18, 0.57_41, 0.44_82, 0.49_40, 0.46_29, 0.45_03]					)

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







		def 							lowercase							(     self					)   ->   List[Any]:


							'''simple docstring'''

							_snake_case							:       Optional[int]												= "cpu"
							_snake_case							:       str												= self.get_dummy_inputs()
							inputs.pop("mask_image"					)
							_snake_case							:       Union[str, Any]												= self.convert_to_pt(inputs.pop("image"					)					)
							_snake_case							:       Union[str, Any]												= image.clamp(0			,    1					) / 2

							# make sure here that pndm scheduler skips prk
							_snake_case							:       Dict												= PaintByExamplePipeline(**self.get_dummy_components()					)
							_snake_case							:       Union[str, Any]												= pipe.to(__snake_case					)
							pipe.set_progress_bar_config(disable=__snake_case					)

							_snake_case							:       int												= pipe(image=__snake_case			,    mask_image=mask_image[:, 0]			,    **__snake_case					)
							_snake_case							:       Optional[int]												= output.images

							_snake_case							:       Any												= image.cpu().permute(0			,    2			,    3			,    1					)[0]
							_snake_case							:       List[str]												= mask_image.cpu().permute(0			,    2			,    3			,    1					)[0]

							_snake_case							:       Any												= Image.fromarray(np.uinta(__snake_case					)					).convert("RGB"					)
							_snake_case							:       int												= Image.fromarray(np.uinta(__snake_case					)					).convert("RGB"					)

							_snake_case							:       Optional[int]												= pipe(**self.get_dummy_inputs()					)
							_snake_case							:       Dict												= output.images

							assert out_a.shape == (1, 6_4, 6_4, 3)
							assert np.abs(out_a.flatten() - out_a.flatten()					).max() < 5E-2







		def 							lowercase							(     self					)   ->   Union[str, Any]:


							'''simple docstring'''

							super().test_inference_batch_single_identical(expected_max_diff=3E-3					)






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







		def 							lowercase							(     self					)   ->   Union[str, Any]:


							'''simple docstring'''

							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()







		def 							lowercase							(     self					)   ->   int:


							'''simple docstring'''

							_snake_case							:       Tuple												= load_image(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
							    "/paint_by_example/dog_in_bucket.png"					)
							_snake_case							:       Optional[Any]												= load_image(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
							    "/paint_by_example/mask.png"					)
							_snake_case							:       Optional[int]												= load_image(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
							    "/paint_by_example/panda.jpg"					)

							_snake_case							:       Any												= PaintByExamplePipeline.from_pretrained("Fantasy-Studio/Paint-by-Example"					)
							_snake_case							:       Any												= pipe.to(__snake_case					)
							pipe.set_progress_bar_config(disable=__snake_case					)

							_snake_case							:       Union[str, Any]												= torch.manual_seed(3_2_1					)
							_snake_case							:       int												= pipe(
							    image=__snake_case			,    mask_image=__snake_case			,    example_image=__snake_case			,    generator=__snake_case			,    guidance_scale=5.0			,    num_inference_steps=5_0			,    output_type="np"			,    )

							_snake_case							:       str												= output.images

							_snake_case							:       Optional[int]												= image[0, -3:, -3:, -1]

							assert image.shape == (1, 5_1_2, 5_1_2, 3)
							_snake_case							:       int												= np.array([0.48_34, 0.48_11, 0.48_74, 0.51_22, 0.50_81, 0.51_44, 0.52_91, 0.52_90, 0.53_74]					)

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