


import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()

class 							a						(unittest.TestCase ):




							"""simple docstring"""


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

							@property
							def        __snake_case   (				self			:    Union[str, Any]							)      ->       Tuple:
														__snake_case						:     List[Any]									=  1
														__snake_case						:     Any									=  3
														__snake_case						:     str									=  (32, 32)

														__snake_case						:     Optional[Any]									=  floats_tensor((batch_size, num_channels) + sizes       ,      rng=random.Random(0							)							).to(lowerCamelCase							)
														return image

							@property
							def        __snake_case   (				self			:    List[Any]							)      ->       Union[str, Any]:
														torch.manual_seed(0							)
														__snake_case						:     int									=  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       ,      )
														return model

							@property
							def        __snake_case   (				self			:    List[str]							)      ->       Optional[Any]:
														torch.manual_seed(0							)
														__snake_case						:     Optional[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       ,      )
														return model

							@property
							def        __snake_case   (				self			:    int							)      ->       Any:
														torch.manual_seed(0							)
														__snake_case						:     Optional[Any]									=  RobertaSeriesConfig(
														    hidden_size=32       ,      project_dim=32       ,      intermediate_size=37       ,      layer_norm_eps=1E-05       ,      num_attention_heads=4       ,      num_hidden_layers=5       ,      pad_token_id=1       ,      vocab_size=5006       ,      )
														return RobertaSeriesModelWithTransformation(lowerCamelCase							)

							@property
							def        __snake_case   (				self			:    int							)      ->       Any:
														def extract(*lowerCamelCase			:    int       ,      **lowerCamelCase			:    Any							):

																					class 							a						:




																												"""simple docstring"""


																												def __init__(				self			:    Tuple							)      ->       int:
																																			__snake_case						:     List[Any]									=  torch.ones([0]							)






																												def        __snake_case   (				self			:    List[Any]       ,      lowerCamelCase			:    List[str]							)      ->       Dict:
																																			self.pixel_values.to(lowerCamelCase							)
																																			return self
																					return Out()

														return extract

							def        __snake_case   (				self			:    Tuple							)      ->       Tuple:
														__snake_case						:     List[Any]									=  "cpu"  # ensure determinism for the device-dependent torch.Generator
														__snake_case						:     Optional[Any]									=  self.dummy_cond_unet
														__snake_case						:     Dict									=  PNDMScheduler(skip_prk_steps=lowerCamelCase							)
														__snake_case						:     List[Any]									=  self.dummy_vae
														__snake_case						:     Optional[int]									=  self.dummy_text_encoder
														__snake_case						:     int									=  XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta"							)
														__snake_case						:     List[Any]									=  77

														__snake_case						:     int									=  self.dummy_image.to(lowerCamelCase							)
														__snake_case						:     str									=  init_image / 2 + 0.5

														# make sure here that pndm scheduler skips prk
														__snake_case						:     Dict									=  AltDiffusionImgaImgPipeline(
														    unet=lowerCamelCase       ,      scheduler=lowerCamelCase       ,      vae=lowerCamelCase       ,      text_encoder=lowerCamelCase       ,      tokenizer=lowerCamelCase       ,      safety_checker=lowerCamelCase       ,      feature_extractor=self.dummy_extractor       ,      )
														__snake_case						:     List[str]									=  VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor       ,      do_normalize=lowerCamelCase							)
														__snake_case						:     Any									=  alt_pipe.to(lowerCamelCase							)
														alt_pipe.set_progress_bar_config(disable=lowerCamelCase							)

														__snake_case						:     str									=  "A painting of a squirrel eating a burger"
														__snake_case						:     Optional[int]									=  torch.Generator(device=lowerCamelCase							).manual_seed(0							)
														__snake_case						:     Any									=  alt_pipe(
														    [prompt]       ,      generator=lowerCamelCase       ,      guidance_scale=6.0       ,      num_inference_steps=2       ,      output_type="np"       ,      image=lowerCamelCase       ,      )

														__snake_case						:     Dict									=  output.images

														__snake_case						:     Tuple									=  torch.Generator(device=lowerCamelCase							).manual_seed(0							)
														__snake_case						:     Dict									=  alt_pipe(
														    [prompt]       ,      generator=lowerCamelCase       ,      guidance_scale=6.0       ,      num_inference_steps=2       ,      output_type="np"       ,      image=lowerCamelCase       ,      return_dict=lowerCamelCase       ,      )[0]

														__snake_case						:     int									=  image[0, -3:, -3:, -1]
														__snake_case						:     int									=  image_from_tuple[0, -3:, -3:, -1]

														assert image.shape == (1, 32, 32, 3)
														__snake_case						:     Tuple									=  np.array([0.44_27, 0.37_31, 0.42_49, 0.49_41, 0.45_46, 0.41_48, 0.41_93, 0.46_66, 0.44_99]							)

														assert np.abs(image_slice.flatten() - expected_slice							).max() < 5E-3
														assert np.abs(image_from_tuple_slice.flatten() - expected_slice							).max() < 5E-3

							@unittest.skipIf(torch_device != "cuda"       ,      "This test requires a GPU"							)
							def        __snake_case   (				self			:    int							)      ->       Optional[int]:
														__snake_case						:     Optional[int]									=  self.dummy_cond_unet
														__snake_case						:     Union[str, Any]									=  PNDMScheduler(skip_prk_steps=lowerCamelCase							)
														__snake_case						:     Optional[int]									=  self.dummy_vae
														__snake_case						:     Optional[int]									=  self.dummy_text_encoder
														__snake_case						:     Union[str, Any]									=  XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta"							)
														__snake_case						:     List[str]									=  77

														__snake_case						:     Optional[int]									=  self.dummy_image.to(lowerCamelCase							)

														# put models in fp16
														__snake_case						:     str									=  unet.half()
														__snake_case						:     Dict									=  vae.half()
														__snake_case						:     int									=  bert.half()

														# make sure here that pndm scheduler skips prk
														__snake_case						:     int									=  AltDiffusionImgaImgPipeline(
														    unet=lowerCamelCase       ,      scheduler=lowerCamelCase       ,      vae=lowerCamelCase       ,      text_encoder=lowerCamelCase       ,      tokenizer=lowerCamelCase       ,      safety_checker=lowerCamelCase       ,      feature_extractor=self.dummy_extractor       ,      )
														__snake_case						:     List[str]									=  VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor       ,      do_normalize=lowerCamelCase							)
														__snake_case						:     Dict									=  alt_pipe.to(lowerCamelCase							)
														alt_pipe.set_progress_bar_config(disable=lowerCamelCase							)

														__snake_case						:     int									=  "A painting of a squirrel eating a burger"
														__snake_case						:     int									=  torch.manual_seed(0							)
														__snake_case						:     int									=  alt_pipe(
														    [prompt]       ,      generator=lowerCamelCase       ,      num_inference_steps=2       ,      output_type="np"       ,      image=lowerCamelCase       ,      ).images

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






							@unittest.skipIf(torch_device != "cuda"       ,      "This test requires a GPU"							)
							def        __snake_case   (				self			:    Any							)      ->       List[str]:
														__snake_case						:     str									=  load_image(
														    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
														    "/img2img/sketch-mountains-input.jpg"							)
														# resize to resolution that is divisible by 8 but not 16 or 32
														__snake_case						:     Any									=  init_image.resize((760, 504)							)

														__snake_case						:     Optional[Any]									=  "BAAI/AltDiffusion"
														__snake_case						:     List[str]									=  AltDiffusionImgaImgPipeline.from_pretrained(
														    lowerCamelCase       ,      safety_checker=lowerCamelCase       ,      )
														pipe.to(lowerCamelCase							)
														pipe.set_progress_bar_config(disable=lowerCamelCase							)
														pipe.enable_attention_slicing()

														__snake_case						:     Dict									=  "A fantasy landscape, trending on artstation"

														__snake_case						:     int									=  torch.manual_seed(0							)
														__snake_case						:     Dict									=  pipe(
														    prompt=lowerCamelCase       ,      image=lowerCamelCase       ,      strength=0.75       ,      guidance_scale=7.5       ,      generator=lowerCamelCase       ,      output_type="np"       ,      )
														__snake_case						:     Optional[Any]									=  output.images[0]

														__snake_case						:     Dict									=  image[255:258, 383:386, -1]

														assert image.shape == (504, 760, 3)
														__snake_case						:     Any									=  np.array([0.93_58, 0.93_97, 0.95_99, 0.99_01, 1.00_00, 1.00_00, 0.98_82, 1.00_00, 1.00_00]							)

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

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




							"""simple docstring"""


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






							def        __snake_case   (				self			:    Optional[Any]							)      ->       int:
														__snake_case						:     Optional[int]									=  load_image(
														    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
														    "/img2img/sketch-mountains-input.jpg"							)
														__snake_case						:     str									=  init_image.resize((768, 512)							)
														__snake_case						:     Dict									=  load_numpy(
														    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy"							)

														__snake_case						:     Any									=  "BAAI/AltDiffusion"
														__snake_case						:     Any									=  AltDiffusionImgaImgPipeline.from_pretrained(
														    lowerCamelCase       ,      safety_checker=lowerCamelCase       ,      )
														pipe.to(lowerCamelCase							)
														pipe.set_progress_bar_config(disable=lowerCamelCase							)
														pipe.enable_attention_slicing()

														__snake_case						:     Optional[Any]									=  "A fantasy landscape, trending on artstation"

														__snake_case						:     Optional[Any]									=  torch.manual_seed(0							)
														__snake_case						:     Optional[Any]									=  pipe(
														    prompt=lowerCamelCase       ,      image=lowerCamelCase       ,      strength=0.75       ,      guidance_scale=7.5       ,      generator=lowerCamelCase       ,      output_type="np"       ,      )
														__snake_case						:     Union[str, Any]									=  output.images[0]

														assert image.shape == (512, 768, 3)
														# img2img is flaky across GPUs even in fp32, so using MAE here
														assert np.abs(expected_image - image							).max() < 1E-2





'''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 							=  logging.getLogger(__name__)







class 							__magic_name__	(       lowerCAmelCase ):

     """simple docstring"""






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


          '''simple docstring'''





          __A	:       List[Any]      =							self.layer[current_layer](lowerCamelCase			,     lowerCamelCase			,     head_mask[current_layer]		)

          __A	:       Tuple      =							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.'    ,		lowerCAmelCase    ,		)
class 							__magic_name__	(       lowerCAmelCase ):

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase		):


          '''simple docstring'''





          super().__init__(lowerCamelCase		)

          __A	:       Any      =							BertEncoderWithPabee(lowerCamelCase		)

          self.init_weights()
          __A	:       Optional[Any]      =							0
          __A	:       Union[str, Any]      =							0
          __A	:       Dict      =							0

          __A	:       List[Any]      =							0


     def lowerCAmelCase__  (   self			,     lowerCamelCase		):


          '''simple docstring'''





          __A	:       Any      =							threshold


     def lowerCAmelCase__  (   self			,     lowerCamelCase		):


          '''simple docstring'''





          __A	:       List[Any]      =							patience


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       List[str]      =							0
          __A	:       int      =							0


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Optional[Any]      =							self.inference_layers_num / self.inference_instances_num
          __A	:       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(lowerCamelCase		)


     @add_start_docstrings_to_model_forward(lowerCamelCase		)
     def lowerCAmelCase__  (   self			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=False			,     ):


          '''simple docstring'''





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

          __A	:       List[str]      =							input_ids.device if input_ids is not None else inputs_embeds.device

          if attention_mask is None:
               __A	:       Any      =							torch.ones(lowerCamelCase			,     device=lowerCamelCase		)
          if token_type_ids is None:
               __A	:       List[Any]      =							torch.zeros(lowerCamelCase			,     dtype=torch.long			,     device=lowerCamelCase		)

          # 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(lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase		)

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

          __A	:       Tuple      =							self.embeddings(
              input_ids=lowerCamelCase			,     position_ids=lowerCamelCase			,     token_type_ids=lowerCamelCase			,     inputs_embeds=lowerCamelCase		)
          __A	:       str      =							embedding_output

          if self.training:
               __A	:       Dict      =							[]
               for i in range(self.config.num_hidden_layers		):
                    __A	:       str      =							self.encoder.adaptive_forward(
                        lowerCamelCase			,     current_layer=lowerCamelCase			,     attention_mask=lowerCamelCase			,     head_mask=lowerCamelCase		)

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

                    __A	:       List[str]      =							self.pooler(lowerCamelCase		)
                    __A	:       str      =							output_layers[i](lowerCamelCase		)
                    if regression:
                         __A	:       int      =							logits.detach()
                         if patient_result is not None:
                              __A	:       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:
                              __A	:       Any      =							0
                    else:
                         __A	:       str      =							logits.detach().argmax(dim=1		)
                         if patient_result is not None:
                              __A	:       List[str]      =							patient_result.detach().argmax(dim=1		)
                         if (patient_result is not None) and torch.all(labels.eq(lowerCamelCase		)		):
                              patient_counter += 1
                         else:
                              __A	:       Union[str, Any]      =							0

                    __A	:       int      =							logits
                    if patient_counter == self.patience:
                         break
               __A	:       List[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\n    the pooled output) e.g. for GLUE tasks. '    ,		lowerCAmelCase    ,		)
class 							__magic_name__	(       lowerCAmelCase ):

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase		):


          '''simple docstring'''





          super().__init__(lowerCamelCase		)
          __A	:       Union[str, Any]      =							config.num_labels

          __A	:       Union[str, Any]      =							BertModelWithPabee(lowerCamelCase		)
          __A	:       Union[str, Any]      =							nn.Dropout(config.hidden_dropout_prob		)
          __A	:       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(lowerCamelCase		)
     def lowerCAmelCase__  (   self			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     ):


          '''simple docstring'''





          __A	:       Optional[Any]      =							self.bert(
              input_ids=lowerCamelCase			,     attention_mask=lowerCamelCase			,     token_type_ids=lowerCamelCase			,     position_ids=lowerCamelCase			,     head_mask=lowerCamelCase			,     inputs_embeds=lowerCamelCase			,     output_dropout=self.dropout			,     output_layers=self.classifiers			,     regression=self.num_labels == 1			,     )

          __A	:       Any      =							(logits[-1],)

          if labels is not None:
               __A	:       Optional[Any]      =							None
               __A	:       int      =							0
               for ix, logits_item in enumerate(lowerCamelCase		):
                    if self.num_labels == 1:
                         #  We are doing regression
                         __A	:       Union[str, Any]      =							MSELoss()
                         __A	:       str      =							loss_fct(logits_item.view(-1		)			,     labels.view(-1		)		)
                    else:
                         __A	:       Optional[int]      =							CrossEntropyLoss()
                         __A	:       Optional[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	:       Tuple      =							(total_loss / total_weights,) + outputs

          return outputs

