



'''simple docstring'''



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,
    BertEmbeddings,
    BertLayer,
    BertPooler,
    BertPreTrainedModel,
)







def _lowercase      (     __A				):



 '''simple docstring'''



 __UpperCamelCase 					=     torch.exp(a__				)
 __UpperCamelCase 					=     torch.sum(a__							,dim=1				)  # sum of exp(x_i)
 __UpperCamelCase 					=     torch.sum(x * exp_x							,dim=1				)  # sum of x_i * exp(x_i)
 return torch.log(a__				) - B / A
class 				UpperCAmelCase__     (			nn.Module):



 def __init__(     self						,		lowercase	)     ->							Optional[int]:
  super().__init__()
  __UpperCamelCase 					=     config.output_attentions
  __UpperCamelCase 					=     config.output_hidden_states
  __UpperCamelCase 					=     nn.ModuleList([BertLayer(__SCREAMING_SNAKE_CASE	) for _ in range(config.num_hidden_layers	)]	)
  __UpperCamelCase 					=     nn.ModuleList([BertHighway(__SCREAMING_SNAKE_CASE	) for _ in range(config.num_hidden_layers	)]	)

  __UpperCamelCase 					=     [-1 for _ in range(config.num_hidden_layers	)]



 def       __lowerCamelCase (     self						,		lowercase	)     ->							int:




  if (type(__SCREAMING_SNAKE_CASE	) is float) or (type(__SCREAMING_SNAKE_CASE	) is int):
   for i in range(len(self.early_exit_entropy	)	):
    __UpperCamelCase 					=     x
  else:
   __UpperCamelCase 					=     x



 def       __lowerCamelCase (     self						,		lowercase	)     ->							Dict:
  __UpperCamelCase 					=     pooler.state_dict()
  for highway in self.highway:
   for name, param in highway.pooler.state_dict().items():
    param.copy_(loaded_model[name]	)




 def       __lowerCamelCase (     self						,		lowercase						,		lowercase=None						,		lowercase=None						,		lowercase=None						,		lowercase=None						,		)     ->							Optional[int]:
  __UpperCamelCase 					=     ()
  __UpperCamelCase 					=     ()
  __UpperCamelCase 					=     ()
  for i, layer_module in enumerate(self.layer	):
   if self.output_hidden_states:
    __UpperCamelCase 					=     all_hidden_states + (hidden_states,)

   __UpperCamelCase 					=     layer_module(
       __SCREAMING_SNAKE_CASE						,		__SCREAMING_SNAKE_CASE						,		head_mask[i]						,		__SCREAMING_SNAKE_CASE						,		__SCREAMING_SNAKE_CASE	)
   __UpperCamelCase 					=     layer_outputs[0]

   if self.output_attentions:
    __UpperCamelCase 					=     all_attentions + (layer_outputs[1],)

   __UpperCamelCase 					=     (hidden_states,)
   if self.output_hidden_states:
    __UpperCamelCase 					=     current_outputs + (all_hidden_states,)
   if self.output_attentions:
    __UpperCamelCase 					=     current_outputs + (all_attentions,)

   __UpperCamelCase 					=     self.highway[i](__SCREAMING_SNAKE_CASE	)
   # logits, pooled_output

   if not self.training:
    __UpperCamelCase 					=     highway_exit[0]
    __UpperCamelCase 					=     entropy(__SCREAMING_SNAKE_CASE	)
    __UpperCamelCase 					=     highway_exit + (highway_entropy,)  # logits, hidden_states(?), entropy
    __UpperCamelCase 					=     all_highway_exits + (highway_exit,)

    if highway_entropy < self.early_exit_entropy[i]:
     __UpperCamelCase 					=     (highway_logits,) + current_outputs[1:] + (all_highway_exits,)
     raise HighwayException(__SCREAMING_SNAKE_CASE						,		i + 1	)
   else:
    __UpperCamelCase 					=     all_highway_exits + (highway_exit,)

        # Add last layer
  if self.output_hidden_states:
   __UpperCamelCase 					=     all_hidden_states + (hidden_states,)

  __UpperCamelCase 					=     (hidden_states,)
  if self.output_hidden_states:
   __UpperCamelCase 					=     outputs + (all_hidden_states,)
  if self.output_attentions:
   __UpperCamelCase 					=     outputs + (all_attentions,)

  __UpperCamelCase 					=     outputs + (all_highway_exits,)
  return outputs  # last-layer hidden state, (all hidden states), (all attentions), all highway exits


@add_start_docstrings(
    '''The Bert Model transformer with early exiting (DeeBERT). '''							,	UpperCAmelCase_							,	)
class 				UpperCAmelCase__     (			UpperCAmelCase_):



 def __init__(     self						,		lowercase	)     ->							Optional[int]:
  super().__init__(__SCREAMING_SNAKE_CASE	)
  __UpperCamelCase 					=     config

  __UpperCamelCase 					=     BertEmbeddings(__SCREAMING_SNAKE_CASE	)
  __UpperCamelCase 					=     DeeBertEncoder(__SCREAMING_SNAKE_CASE	)
  __UpperCamelCase 					=     BertPooler(__SCREAMING_SNAKE_CASE	)

  self.init_weights()



 def       __lowerCamelCase (     self	)     ->							Optional[int]:
  self.encoder.init_highway_pooler(self.pooler	)



 def       __lowerCamelCase (     self	)     ->							str:
  return self.embeddings.word_embeddings



 def       __lowerCamelCase (     self						,		lowercase	)     ->							List[str]:
  __UpperCamelCase 					=     value



 def       __lowerCamelCase (     self						,		lowercase	)     ->							Optional[int]:




  for layer, heads in heads_to_prune.items():
   self.encoder.layer[layer].attention.prune_heads(__SCREAMING_SNAKE_CASE	)




 @add_start_docstrings_to_model_forward(__SCREAMING_SNAKE_CASE	)
 def       __lowerCamelCase (     self						,		lowercase=None						,		lowercase=None						,		lowercase=None						,		lowercase=None						,		lowercase=None						,		lowercase=None						,		lowercase=None						,		lowercase=None						,		)     ->							Optional[Any]:




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

  __UpperCamelCase 					=     input_ids.device if input_ids is not None else inputs_embeds.device

  if attention_mask is None:
   __UpperCamelCase 					=     torch.ones(__SCREAMING_SNAKE_CASE						,		device=__SCREAMING_SNAKE_CASE	)
  if encoder_attention_mask is None:
   __UpperCamelCase 					=     torch.ones(__SCREAMING_SNAKE_CASE						,		device=__SCREAMING_SNAKE_CASE	)
  if token_type_ids is None:
   __UpperCamelCase 					=     torch.zeros(__SCREAMING_SNAKE_CASE						,		dtype=torch.long						,		device=__SCREAMING_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 					=     self.get_extended_attention_mask(__SCREAMING_SNAKE_CASE						,		__SCREAMING_SNAKE_CASE						,		__SCREAMING_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 encoder_attention_mask.dim() == 3:
   __UpperCamelCase 					=     encoder_attention_mask[:, None, :, :]
  if encoder_attention_mask.dim() == 2:
   __UpperCamelCase 					=     encoder_attention_mask[:, None, None, :]

  __UpperCamelCase 					=     encoder_extended_attention_mask.to(
      dtype=next(self.parameters()	).dtype	)  # fp16 compatibility
  __UpperCamelCase 					=     (1.0 - encoder_extended_attention_mask) * -10_000.0

  # 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 					=     self.get_head_mask(__SCREAMING_SNAKE_CASE						,		self.config.num_hidden_layers	)

  __UpperCamelCase 					=     self.embeddings(
      input_ids=__SCREAMING_SNAKE_CASE						,		position_ids=__SCREAMING_SNAKE_CASE						,		token_type_ids=__SCREAMING_SNAKE_CASE						,		inputs_embeds=__SCREAMING_SNAKE_CASE	)
  __UpperCamelCase 					=     self.encoder(
      __SCREAMING_SNAKE_CASE						,		attention_mask=__SCREAMING_SNAKE_CASE						,		head_mask=__SCREAMING_SNAKE_CASE						,		encoder_hidden_states=__SCREAMING_SNAKE_CASE						,		encoder_attention_mask=__SCREAMING_SNAKE_CASE						,		)
  __UpperCamelCase 					=     encoder_outputs[0]
  __UpperCamelCase 					=     self.pooler(__SCREAMING_SNAKE_CASE	)

  __UpperCamelCase 					=     (
      sequence_output,
      pooled_output,
  ) + encoder_outputs[
      1:
  ]  # add hidden_states and attentions if they are here
  return outputs  # sequence_output, pooled_output, (hidden_states), (attentions), highway exits


class 				UpperCAmelCase__     (			UpperCAmelCase_):



 def __init__(     self						,		lowercase						,		lowercase	)     ->							Optional[int]:
  __UpperCamelCase 					=     message
  __UpperCamelCase 					=     exit_layer  # start from 1!


class 				UpperCAmelCase__     (			nn.Module):



 def __init__(     self						,		lowercase	)     ->							str:
  super().__init__()
  __UpperCamelCase 					=     BertPooler(__SCREAMING_SNAKE_CASE	)
  __UpperCamelCase 					=     nn.Dropout(config.hidden_dropout_prob	)
  __UpperCamelCase 					=     nn.Linear(config.hidden_size						,		config.num_labels	)




 def       __lowerCamelCase (     self						,		lowercase	)     ->							Optional[int]:
  __UpperCamelCase 					=     encoder_outputs[0]
  __UpperCamelCase 					=     self.pooler(__SCREAMING_SNAKE_CASE	)
  # "return" pooler_output

  # BertModel
  __UpperCamelCase 					=     (pooler_input, pooler_output) + encoder_outputs[1:]
  # "return" bmodel_output

  # Dropout and classification
  __UpperCamelCase 					=     bmodel_output[1]

  __UpperCamelCase 					=     self.dropout(__SCREAMING_SNAKE_CASE	)
  __UpperCamelCase 					=     self.classifier(__SCREAMING_SNAKE_CASE	)

  return logits, pooled_output


@add_start_docstrings(
    '''Bert Model (with early exiting - DeeBERT) with a classifier on top,\n    also takes care of multi-layer training. '''							,	UpperCAmelCase_							,	)
class 				UpperCAmelCase__     (			UpperCAmelCase_):



 def __init__(     self						,		lowercase	)     ->							Optional[Any]:
  super().__init__(__SCREAMING_SNAKE_CASE	)
  __UpperCamelCase 					=     config.num_labels
  __UpperCamelCase 					=     config.num_hidden_layers

  __UpperCamelCase 					=     DeeBertModel(__SCREAMING_SNAKE_CASE	)
  __UpperCamelCase 					=     nn.Dropout(config.hidden_dropout_prob	)
  __UpperCamelCase 					=     nn.Linear(config.hidden_size						,		self.config.num_labels	)

  self.init_weights()




 @add_start_docstrings_to_model_forward(__SCREAMING_SNAKE_CASE	)
 def       __lowerCamelCase (     self						,		lowercase=None						,		lowercase=None						,		lowercase=None						,		lowercase=None						,		lowercase=None						,		lowercase=None						,		lowercase=None						,		lowercase=-1						,		lowercase=False						,		)     ->							Any:
  __UpperCamelCase 					=     self.num_layers
  try:
   __UpperCamelCase 					=     self.bert(
       __SCREAMING_SNAKE_CASE						,		attention_mask=__SCREAMING_SNAKE_CASE						,		token_type_ids=__SCREAMING_SNAKE_CASE						,		position_ids=__SCREAMING_SNAKE_CASE						,		head_mask=__SCREAMING_SNAKE_CASE						,		inputs_embeds=__SCREAMING_SNAKE_CASE						,		)
   # sequence_output, pooled_output, (hidden_states), (attentions), highway exits

   __UpperCamelCase 					=     outputs[1]

   __UpperCamelCase 					=     self.dropout(__SCREAMING_SNAKE_CASE	)
   __UpperCamelCase 					=     self.classifier(__SCREAMING_SNAKE_CASE	)
   __UpperCamelCase 					=     (logits,) + outputs[2:]  # add hidden states and attention if they are here
  except HighwayException as e:
   __UpperCamelCase 					=     e.message
   __UpperCamelCase 					=     e.exit_layer
   __UpperCamelCase 					=     outputs[0]

  if not self.training:
   __UpperCamelCase 					=     entropy(__SCREAMING_SNAKE_CASE	)
   __UpperCamelCase 					=     []
   __UpperCamelCase 					=     []
  if labels is not None:
   if self.num_labels == 1:
    #  We are doing regression
    __UpperCamelCase 					=     MSELoss()
    __UpperCamelCase 					=     loss_fct(logits.view(-1	)						,		labels.view(-1	)	)
   else:
    __UpperCamelCase 					=     CrossEntropyLoss()
    __UpperCamelCase 					=     loss_fct(logits.view(-1						,		self.num_labels	)						,		labels.view(-1	)	)

   # work with highway exits
   __UpperCamelCase 					=     []
   for highway_exit in outputs[-1]:
    __UpperCamelCase 					=     highway_exit[0]
    if not self.training:
     highway_logits_all.append(__SCREAMING_SNAKE_CASE	)
     highway_entropy.append(highway_exit[2]	)
    if self.num_labels == 1:
     #  We are doing regression
     __UpperCamelCase 					=     MSELoss()
     __UpperCamelCase 					=     loss_fct(highway_logits.view(-1	)						,		labels.view(-1	)	)
    else:
     __UpperCamelCase 					=     CrossEntropyLoss()
     __UpperCamelCase 					=     loss_fct(highway_logits.view(-1						,		self.num_labels	)						,		labels.view(-1	)	)
    highway_losses.append(__SCREAMING_SNAKE_CASE	)

   if train_highway:
    __UpperCamelCase 					=     (sum(highway_losses[:-1]	),) + outputs
    # exclude the final highway, of course
   else:
    __UpperCamelCase 					=     (loss,) + outputs
  if not self.training:
   __UpperCamelCase 					=     outputs + ((original_entropy, highway_entropy), exit_layer)
   if output_layer >= 0:
    __UpperCamelCase 					=     (
        (outputs[0],) + (highway_logits_all[output_layer],) + outputs[2:]
    )  # use the highway of the last layer

  return outputs  # (loss), logits, (hidden_states), (attentions), (highway_exits)








'''simple docstring'''



import gc
import random
import unittest

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

from diffusers import AutoencoderKL, PNDMScheduler, StableDiffusionInpaintPipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()







class 						lowerCAmelCase__   (      a ,					a ,					a ,					unittest.TestCase					):



 """simple docstring"""



 lowerCAmelCase__  					=      StableDiffusionInpaintPipeline
 lowerCAmelCase__  					=      TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
 lowerCAmelCase__  					=      TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
 lowerCAmelCase__  					=      frozenset(
     []					)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
 lowerCAmelCase__  					=      frozenset([]					)
 def       UpperCAmelCase__		(		self					:       str			)  -> List[Any]:


      """simple docstring"""




      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      UNetaDConditionModel(
          block_out_channels=(32, 64)					,							layers_per_block=2					,							sample_size=32					,							in_channels=9					,							out_channels=4					,							down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")					,							up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")					,							cross_attention_dim=32					,							attention_head_dim=(2, 4)					,							use_linear_projection=__SCREAMING_SNAKE_CASE					,							)
      __SCREAMING_SNAKE_CASE     =      PNDMScheduler(skip_prk_steps=__SCREAMING_SNAKE_CASE			)
      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      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					,							sample_size=128					,							)
      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      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					,							hidden_act="""gelu"""					,							projection_dim=512					,							)
      __SCREAMING_SNAKE_CASE     =      CLIPTextModel(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""			)

      __SCREAMING_SNAKE_CASE     =      {
          """unet""": unet,
          """scheduler""": scheduler,
          """vae""": vae,
          """text_encoder""": text_encoder,
          """tokenizer""": tokenizer,
          """safety_checker""": None,
          """feature_extractor""": None,
      }
      return components
 def       UpperCAmelCase__		(		self					:       Any					,							__SCREAMING_SNAKE_CASE					:       Tuple					,							__SCREAMING_SNAKE_CASE					:       List[Any]=0			)  -> int:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      floats_tensor((1, 3, 32, 32)					,							rng=random.Random(__SCREAMING_SNAKE_CASE			)			).to(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      image.cpu().permute(0					,							2					,							3					,							1			)[0]
      __SCREAMING_SNAKE_CASE     =      Image.fromarray(np.uinta(__SCREAMING_SNAKE_CASE			)			).convert("""RGB"""			).resize((64, 64)			)
      __SCREAMING_SNAKE_CASE     =      Image.fromarray(np.uinta(image + 4			)			).convert("""RGB"""			).resize((64, 64)			)
      if str(__SCREAMING_SNAKE_CASE			).startswith("""mps"""			):
           __SCREAMING_SNAKE_CASE     =      torch.manual_seed(__SCREAMING_SNAKE_CASE			)
      else:
           __SCREAMING_SNAKE_CASE     =      torch.Generator(device=__SCREAMING_SNAKE_CASE			).manual_seed(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      {
          """prompt""": """A painting of a squirrel eating a burger""",
          """image""": init_image,
          """mask_image""": mask_image,
          """generator""": generator,
          """num_inference_steps""": 2,
          """guidance_scale""": 6.0,
          """output_type""": """numpy""",
      }
      return inputs
 def       UpperCAmelCase__		(		self					:       Union[str, Any]			)  -> int:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      """cpu"""  # ensure determinism for the device-dependent torch.Generator
      __SCREAMING_SNAKE_CASE     =      self.get_dummy_components()
      __SCREAMING_SNAKE_CASE     =      StableDiffusionInpaintPipeline(**__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      sd_pipe.to(__SCREAMING_SNAKE_CASE			)
      sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE			)

      __SCREAMING_SNAKE_CASE     =      self.get_dummy_inputs(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      sd_pipe(**__SCREAMING_SNAKE_CASE			).images
      __SCREAMING_SNAKE_CASE     =      image[0, -3:, -3:, -1]

      assert image.shape == (1, 64, 64, 3)
      __SCREAMING_SNAKE_CASE     =      np.array([0.4727, 0.5735, 0.3941, 0.5446, 0.5926, 0.4394, 0.5062, 0.4654, 0.4476]			)

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


 def       UpperCAmelCase__		(		self					:       Tuple			)  -> str:


      """simple docstring"""




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







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



 """simple docstring"""



 def       UpperCAmelCase__		(		self					:       List[Any]			)  -> str:


      """simple docstring"""




      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()
 def       UpperCAmelCase__		(		self					:       List[str]			)  -> str:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
          """/sd2-inpaint/init_image.png"""			)
      __SCREAMING_SNAKE_CASE     =      load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"""			)
      __SCREAMING_SNAKE_CASE     =      load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint"""
          """/yellow_cat_sitting_on_a_park_bench.npy"""			)

      __SCREAMING_SNAKE_CASE     =      """stabilityai/stable-diffusion-2-inpainting"""
      __SCREAMING_SNAKE_CASE     =      StableDiffusionInpaintPipeline.from_pretrained(__SCREAMING_SNAKE_CASE					,							safety_checker=__SCREAMING_SNAKE_CASE			)
      pipe.to(__SCREAMING_SNAKE_CASE			)
      pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE			)
      pipe.enable_attention_slicing()

      __SCREAMING_SNAKE_CASE     =      """Face of a yellow cat, high resolution, sitting on a park bench"""

      __SCREAMING_SNAKE_CASE     =      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      pipe(
          prompt=__SCREAMING_SNAKE_CASE					,							image=__SCREAMING_SNAKE_CASE					,							mask_image=__SCREAMING_SNAKE_CASE					,							generator=__SCREAMING_SNAKE_CASE					,							output_type="""np"""					,							)
      __SCREAMING_SNAKE_CASE     =      output.images[0]

      assert image.shape == (512, 512, 3)
      assert np.abs(expected_image - image			).max() < 9E-3
 def       UpperCAmelCase__		(		self					:       List[Any]			)  -> Optional[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
          """/sd2-inpaint/init_image.png"""			)
      __SCREAMING_SNAKE_CASE     =      load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"""			)
      __SCREAMING_SNAKE_CASE     =      load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint"""
          """/yellow_cat_sitting_on_a_park_bench_fp16.npy"""			)

      __SCREAMING_SNAKE_CASE     =      """stabilityai/stable-diffusion-2-inpainting"""
      __SCREAMING_SNAKE_CASE     =      StableDiffusionInpaintPipeline.from_pretrained(
          __SCREAMING_SNAKE_CASE					,							torch_dtype=torch.floataa					,							safety_checker=__SCREAMING_SNAKE_CASE					,							)
      pipe.to(__SCREAMING_SNAKE_CASE			)
      pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE			)
      pipe.enable_attention_slicing()

      __SCREAMING_SNAKE_CASE     =      """Face of a yellow cat, high resolution, sitting on a park bench"""

      __SCREAMING_SNAKE_CASE     =      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      pipe(
          prompt=__SCREAMING_SNAKE_CASE					,							image=__SCREAMING_SNAKE_CASE					,							mask_image=__SCREAMING_SNAKE_CASE					,							generator=__SCREAMING_SNAKE_CASE					,							output_type="""np"""					,							)
      __SCREAMING_SNAKE_CASE     =      output.images[0]

      assert image.shape == (512, 512, 3)
      assert np.abs(expected_image - image			).max() < 5E-1


 def       UpperCAmelCase__		(		self					:       Tuple			)  -> Any:


      """simple docstring"""




      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.reset_peak_memory_stats()

      __SCREAMING_SNAKE_CASE     =      load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
          """/sd2-inpaint/init_image.png"""			)
      __SCREAMING_SNAKE_CASE     =      load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"""			)

      __SCREAMING_SNAKE_CASE     =      """stabilityai/stable-diffusion-2-inpainting"""
      __SCREAMING_SNAKE_CASE     =      PNDMScheduler.from_pretrained(__SCREAMING_SNAKE_CASE					,							subfolder="""scheduler"""			)
      __SCREAMING_SNAKE_CASE     =      StableDiffusionInpaintPipeline.from_pretrained(
          __SCREAMING_SNAKE_CASE					,							safety_checker=__SCREAMING_SNAKE_CASE					,							scheduler=__SCREAMING_SNAKE_CASE					,							torch_dtype=torch.floataa					,							)
      pipe.to(__SCREAMING_SNAKE_CASE			)
      pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE			)
      pipe.enable_attention_slicing(1			)
      pipe.enable_sequential_cpu_offload()

      __SCREAMING_SNAKE_CASE     =      """Face of a yellow cat, high resolution, sitting on a park bench"""

      __SCREAMING_SNAKE_CASE     =      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      pipe(
          prompt=__SCREAMING_SNAKE_CASE					,							image=__SCREAMING_SNAKE_CASE					,							mask_image=__SCREAMING_SNAKE_CASE					,							generator=__SCREAMING_SNAKE_CASE					,							num_inference_steps=2					,							output_type="""np"""					,							)

      __SCREAMING_SNAKE_CASE     =      torch.cuda.max_memory_allocated()
      # make sure that less than 2.65 GB is allocated
      assert mem_bytes < 2.65 * 10**9






