






import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyInpaintPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
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

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()







class      UpperCAmelCase_		(   lowercase,      unittest.TestCase					):



      """simple docstring"""
      UpperCamelCase_							:   Dict 	=KandinskyInpaintPipeline
      UpperCamelCase_							:   List[Any] 	=['prompt', 'image_embeds', 'negative_image_embeds', 'image', 'mask_image']
      UpperCamelCase_							:   Dict 	=[
          'prompt',
          'negative_prompt',
          'image_embeds',
          'negative_image_embeds',
          'image',
          'mask_image',
      ]
      UpperCamelCase_							:   Optional[Any] 	=[
          'generator',
          'height',
          'width',
          'latents',
          'guidance_scale',
          'negative_prompt',
          'num_inference_steps',
          'return_dict',
          'guidance_scale',
          'num_images_per_prompt',
          'output_type',
          'return_dict',
      ]
      UpperCamelCase_							:   Any 	=False
      @property
      def        UpperCAmelCase		(							self					)							-> Union[str, Any]:
          return 32
      @property
      def        UpperCAmelCase		(							self					)							-> Tuple:
          return 32
      @property
      def        UpperCAmelCase		(							self					)							-> Dict:
          return self.time_input_dim
      @property
      def        UpperCAmelCase		(							self					)							-> int:
          return self.time_input_dim * 4
      @property
      def        UpperCAmelCase		(							self					)							-> Optional[int]:
          return 100
      @property
      def        UpperCAmelCase		(							self					)							-> str:
          UpperCamelCase			:Tuple      				=  XLMRobertaTokenizerFast.from_pretrained('''YiYiXu/tiny-random-mclip-base'''					)
          return tokenizer
      @property
      def        UpperCAmelCase		(							self					)							-> Optional[int]:
          torch.manual_seed(0					)
          UpperCamelCase			:Union[str, Any]      				=  MCLIPConfig(
              numDims=self.cross_attention_dim       ,							transformerDimensions=self.text_embedder_hidden_size       ,							hidden_size=self.text_embedder_hidden_size       ,							intermediate_size=37       ,							num_attention_heads=4       ,							num_hidden_layers=5       ,							vocab_size=1005       ,							)

          UpperCamelCase			:List[Any]      				=  MultilingualCLIP(__UpperCAmelCase					)
          UpperCamelCase			:Tuple      				=  text_encoder.eval()

          return text_encoder
      @property
      def        UpperCAmelCase		(							self					)							-> Any:
          torch.manual_seed(0					)

          UpperCamelCase			:int      				=  {
              '''in_channels''': 9,
              # Out channels is double in channels because predicts mean and variance
              '''out_channels''': 8,
              '''addition_embed_type''': '''text_image''',
              '''down_block_types''': ('''ResnetDownsampleBlock2D''', '''SimpleCrossAttnDownBlock2D'''),
              '''up_block_types''': ('''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D'''),
              '''mid_block_type''': '''UNetMidBlock2DSimpleCrossAttn''',
              '''block_out_channels''': (self.block_out_channels_a, self.block_out_channels_a * 2),
              '''layers_per_block''': 1,
              '''encoder_hid_dim''': self.text_embedder_hidden_size,
              '''encoder_hid_dim_type''': '''text_image_proj''',
              '''cross_attention_dim''': self.cross_attention_dim,
              '''attention_head_dim''': 4,
              '''resnet_time_scale_shift''': '''scale_shift''',
              '''class_embed_type''': None,
          }

          UpperCamelCase			:int      				=  UNetaDConditionModel(**__UpperCAmelCase					)
          return model
      @property
      def        UpperCAmelCase		(							self					)							-> Optional[Any]:
          return {
              "block_out_channels": [32, 64],
              "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
              "in_channels": 3,
              "latent_channels": 4,
              "layers_per_block": 1,
              "norm_num_groups": 8,
              "norm_type": "spatial",
              "num_vq_embeddings": 12,
              "out_channels": 3,
              "up_block_types": [
                  "AttnUpDecoderBlock2D",
                  "UpDecoderBlock2D",
              ],
              "vq_embed_dim": 4,
          }
      @property
      def        UpperCAmelCase		(							self					)							-> int:
          torch.manual_seed(0					)
          UpperCamelCase			:Tuple      				=  VQModel(**self.dummy_movq_kwargs					)
          return model
      def        UpperCAmelCase		(							self					)							-> Dict:
          UpperCamelCase			:List[str]      				=  self.dummy_text_encoder
          UpperCamelCase			:Any      				=  self.dummy_tokenizer
          UpperCamelCase			:Dict      				=  self.dummy_unet
          UpperCamelCase			:Optional[int]      				=  self.dummy_movq

          UpperCamelCase			:Any      				=  DDIMScheduler(
              num_train_timesteps=1000       ,							beta_schedule='''linear'''       ,							beta_start=0.0_0085       ,							beta_end=0.012       ,							clip_sample=__UpperCAmelCase       ,							set_alpha_to_one=__UpperCAmelCase       ,							steps_offset=1       ,							prediction_type='''epsilon'''       ,							thresholding=__UpperCAmelCase       ,							)

          UpperCamelCase			:Tuple      				=  {
              '''text_encoder''': text_encoder,
              '''tokenizer''': tokenizer,
              '''unet''': unet,
              '''scheduler''': scheduler,
              '''movq''': movq,
          }

          return components
      def        UpperCAmelCase		(							self       ,							SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_=0					)							-> str:
          UpperCamelCase			:Optional[int]      				=  floats_tensor((1, self.cross_attention_dim)       ,							rng=random.Random(__UpperCAmelCase					)					).to(__UpperCAmelCase					)
          UpperCamelCase			:str      				=  floats_tensor((1, self.cross_attention_dim)       ,							rng=random.Random(seed + 1					)					).to(__UpperCAmelCase					)
          # create init_image
          UpperCamelCase			:Any      				=  floats_tensor((1, 3, 64, 64)       ,							rng=random.Random(__UpperCAmelCase					)					).to(__UpperCAmelCase					)
          UpperCamelCase			:Dict      				=  image.cpu().permute(0       ,							2       ,							3       ,							1					)[0]
          UpperCamelCase			:str      				=  Image.fromarray(np.uinta(__UpperCAmelCase					)					).convert('''RGB'''					).resize((256, 256)					)
          # create mask
          UpperCamelCase			:Dict      				=  np.ones((64, 64)       ,							dtype=np.floataa					)
          UpperCamelCase			:Union[str, Any]      				=  0

          if str(__UpperCAmelCase					).startswith('''mps'''					):
              UpperCamelCase			:List[Any]      				=  torch.manual_seed(__UpperCAmelCase					)
          else:
              UpperCamelCase			:Optional[Any]      				=  torch.Generator(device=__UpperCAmelCase					).manual_seed(__UpperCAmelCase					)
          UpperCamelCase			:Dict      				=  {
              '''prompt''': '''horse''',
              '''image''': init_image,
              '''mask_image''': mask,
              '''image_embeds''': image_embeds,
              '''negative_image_embeds''': negative_image_embeds,
              '''generator''': generator,
              '''height''': 64,
              '''width''': 64,
              '''num_inference_steps''': 2,
              '''guidance_scale''': 4.0,
              '''output_type''': '''np''',
          }
          return inputs
      def        UpperCAmelCase		(							self					)							-> Optional[Any]:
          UpperCamelCase			:Any      				=  '''cpu'''

          UpperCamelCase			:Dict      				=  self.get_dummy_components()

          UpperCamelCase			:Any      				=  self.pipeline_class(**__UpperCAmelCase					)
          UpperCamelCase			:int      				=  pipe.to(__UpperCAmelCase					)

          pipe.set_progress_bar_config(disable=__UpperCAmelCase					)

          UpperCamelCase			:Union[str, Any]      				=  pipe(**self.get_dummy_inputs(__UpperCAmelCase					)					)
          UpperCamelCase			:int      				=  output.images

          UpperCamelCase			:List[Any]      				=  pipe(
              **self.get_dummy_inputs(__UpperCAmelCase					)       ,							return_dict=__UpperCAmelCase       ,							)[0]

          UpperCamelCase			:Tuple      				=  image[0, -3:, -3:, -1]
          UpperCamelCase			:str      				=  image_from_tuple[0, -3:, -3:, -1]

          print(F'''image.shape {image.shape}'''					)

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

          UpperCamelCase			:int      				=  np.array(
              [0.832_6919, 0.7379_0467, 0.2091_8581, 0.930_9612, 0.551_1791, 0.4371_3328, 0.551_3321, 0.4992_2934, 0.5949_7786]					)

          assert (
              np.abs(image_slice.flatten() - expected_slice					).max() < 1e-2
          ), F''' expected_slice {expected_slice}, but got {image_slice.flatten()}'''
          assert (
              np.abs(image_from_tuple_slice.flatten() - expected_slice					).max() < 1e-2
          ), F''' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'''




      def        UpperCAmelCase		(							self					)							-> Any:
          super().test_inference_batch_single_identical(expected_max_diff=3e-3					)







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



      """simple docstring"""
      def        UpperCAmelCase		(							self					)							-> str:
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()




      def        UpperCAmelCase		(							self					)							-> List[str]:
          UpperCamelCase			:List[str]      				=  load_numpy(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
              '''/kandinsky/kandinsky_inpaint_cat_with_hat_fp16.npy'''					)

          UpperCamelCase			:List[str]      				=  load_image(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/kandinsky/cat.png'''					)
          UpperCamelCase			:Optional[int]      				=  np.ones((768, 768)       ,							dtype=np.floataa					)
          UpperCamelCase			:List[str]      				=  0

          UpperCamelCase			:Any      				=  '''a hat'''

          UpperCamelCase			:Optional[Any]      				=  KandinskyPriorPipeline.from_pretrained(
              '''kandinsky-community/kandinsky-2-1-prior'''       ,							torch_dtype=torch.floataa					)
          pipe_prior.to(__UpperCAmelCase					)

          UpperCamelCase			:Optional[Any]      				=  KandinskyInpaintPipeline.from_pretrained(
              '''kandinsky-community/kandinsky-2-1-inpaint'''       ,							torch_dtype=torch.floataa					)
          UpperCamelCase			:Dict      				=  pipeline.to(__UpperCAmelCase					)
          pipeline.set_progress_bar_config(disable=__UpperCAmelCase					)

          UpperCamelCase			:Tuple      				=  torch.Generator(device='''cpu'''					).manual_seed(0					)
          UpperCamelCase				,  UpperCamelCase			:Optional[Any]      				=  pipe_prior(
              __UpperCAmelCase       ,							generator=__UpperCAmelCase       ,							num_inference_steps=5       ,							negative_prompt=''''''       ,							).to_tuple()

          UpperCamelCase			:Dict      				=  pipeline(
              __UpperCAmelCase       ,							image=__UpperCAmelCase       ,							mask_image=__UpperCAmelCase       ,							image_embeds=__UpperCAmelCase       ,							negative_image_embeds=__UpperCAmelCase       ,							generator=__UpperCAmelCase       ,							num_inference_steps=100       ,							height=768       ,							width=768       ,							output_type='''np'''       ,							)

          UpperCamelCase			:List[str]      				=  output.images[0]

          assert image.shape == (768, 768, 3)

          assert_mean_pixel_difference(__UpperCAmelCase       ,							__UpperCAmelCase					)













import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import PaddingStrategy, logging
from .tokenization_realm import RealmTokenizer


UpperCAmelCase__							  =						logging.get_logger(__name__)

UpperCAmelCase__							  =						{"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}

UpperCAmelCase__							  =						{
    "vocab_file": {
        "google/realm-cc-news-pretrained-embedder": (
            "https://huggingface.co/google/realm-cc-news-pretrained-embedder/resolve/main/vocab.txt"
        ),
        "google/realm-cc-news-pretrained-encoder": (
            "https://huggingface.co/google/realm-cc-news-pretrained-encoder/resolve/main/vocab.txt"
        ),
        "google/realm-cc-news-pretrained-scorer": (
            "https://huggingface.co/google/realm-cc-news-pretrained-scorer/resolve/main/vocab.txt"
        ),
        "google/realm-cc-news-pretrained-openqa": (
            "https://huggingface.co/google/realm-cc-news-pretrained-openqa/aresolve/main/vocab.txt"
        ),
        "google/realm-orqa-nq-openqa": "https://huggingface.co/google/realm-orqa-nq-openqa/resolve/main/vocab.txt",
        "google/realm-orqa-nq-reader": "https://huggingface.co/google/realm-orqa-nq-reader/resolve/main/vocab.txt",
        "google/realm-orqa-wq-openqa": "https://huggingface.co/google/realm-orqa-wq-openqa/resolve/main/vocab.txt",
        "google/realm-orqa-wq-reader": "https://huggingface.co/google/realm-orqa-wq-reader/resolve/main/vocab.txt",
    },
    "tokenizer_file": {
        "google/realm-cc-news-pretrained-embedder": (
            "https://huggingface.co/google/realm-cc-news-pretrained-embedder/resolve/main/tokenizer.jsont"
        ),
        "google/realm-cc-news-pretrained-encoder": (
            "https://huggingface.co/google/realm-cc-news-pretrained-encoder/resolve/main/tokenizer.json"
        ),
        "google/realm-cc-news-pretrained-scorer": (
            "https://huggingface.co/google/realm-cc-news-pretrained-scorer/resolve/main/tokenizer.json"
        ),
        "google/realm-cc-news-pretrained-openqa": (
            "https://huggingface.co/google/realm-cc-news-pretrained-openqa/aresolve/main/tokenizer.json"
        ),
        "google/realm-orqa-nq-openqa": (
            "https://huggingface.co/google/realm-orqa-nq-openqa/resolve/main/tokenizer.json"
        ),
        "google/realm-orqa-nq-reader": (
            "https://huggingface.co/google/realm-orqa-nq-reader/resolve/main/tokenizer.json"
        ),
        "google/realm-orqa-wq-openqa": (
            "https://huggingface.co/google/realm-orqa-wq-openqa/resolve/main/tokenizer.json"
        ),
        "google/realm-orqa-wq-reader": (
            "https://huggingface.co/google/realm-orqa-wq-reader/resolve/main/tokenizer.json"
        ),
    },
}

UpperCAmelCase__							  =						{
    "google/realm-cc-news-pretrained-embedder": 512,
    "google/realm-cc-news-pretrained-encoder": 512,
    "google/realm-cc-news-pretrained-scorer": 512,
    "google/realm-cc-news-pretrained-openqa": 512,
    "google/realm-orqa-nq-openqa": 512,
    "google/realm-orqa-nq-reader": 512,
    "google/realm-orqa-wq-openqa": 512,
    "google/realm-orqa-wq-reader": 512,
}

UpperCAmelCase__							  =						{
    "google/realm-cc-news-pretrained-embedder": {"do_lower_case": True},
    "google/realm-cc-news-pretrained-encoder": {"do_lower_case": True},
    "google/realm-cc-news-pretrained-scorer": {"do_lower_case": True},
    "google/realm-cc-news-pretrained-openqa": {"do_lower_case": True},
    "google/realm-orqa-nq-openqa": {"do_lower_case": True},
    "google/realm-orqa-nq-reader": {"do_lower_case": True},
    "google/realm-orqa-wq-openqa": {"do_lower_case": True},
    "google/realm-orqa-wq-reader": {"do_lower_case": True},
}






class 						lowercase_   (      lowercase  ):






      '''simple docstring'''




      __snake_case	    =    VOCAB_FILES_NAMES
      __snake_case	    =    PRETRAINED_VOCAB_FILES_MAP
      __snake_case	    =    PRETRAINED_INIT_CONFIGURATION
      __snake_case	    =    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      __snake_case	    =    RealmTokenizer







      def __init__( self :		List[Any]					,   __UpperCAmelCase :		Union[str, Any]=None					,   __UpperCAmelCase :		Tuple=None					,   __UpperCAmelCase :		List[Any]=True					,   __UpperCAmelCase :		Any="[UNK]"					,   __UpperCAmelCase :		Dict="[SEP]"					,   __UpperCAmelCase :		List[str]="[PAD]"					,   __UpperCAmelCase :		str="[CLS]"					,   __UpperCAmelCase :		Optional[int]="[MASK]"					,   __UpperCAmelCase :		str=True					,   __UpperCAmelCase :		List[str]=None					,   **__UpperCAmelCase :		str					,   )       ->Optional[int]:




       """simple docstring"""


       super().__init__(
           __UpperCAmelCase					,   tokenizer_file=__UpperCAmelCase					,   do_lower_case=__UpperCAmelCase					,   unk_token=__UpperCAmelCase					,   sep_token=__UpperCAmelCase					,   pad_token=__UpperCAmelCase					,   cls_token=__UpperCAmelCase					,   mask_token=__UpperCAmelCase					,   tokenize_chinese_chars=__UpperCAmelCase					,   strip_accents=__UpperCAmelCase					,   **__UpperCAmelCase					,   )

       a						    = json.loads(self.backend_tokenizer.normalizer.__getstate__()    )
       if (
           normalizer_state.get('''lowercase'''					,   __UpperCAmelCase    ) != do_lower_case
           or normalizer_state.get('''strip_accents'''					,   __UpperCAmelCase    ) != strip_accents
           or normalizer_state.get('''handle_chinese_chars'''					,   __UpperCAmelCase    ) != tokenize_chinese_chars
       ):
        a						    = getattr(__UpperCAmelCase					,   normalizer_state.pop('''type'''    )    )
        a						    = do_lower_case
        a						    = strip_accents
        a						    = tokenize_chinese_chars
        a						    = normalizer_class(**__UpperCAmelCase    )

       a						    = do_lower_case







      def  __lowerCAmelCase   ( self :		str					,   __UpperCAmelCase :		List[Any]					,   **__UpperCAmelCase :		List[str]    )       ->List[Any]:




       """simple docstring"""


       a						    = PaddingStrategy.MAX_LENGTH

       a						    = text
       a						    = kwargs.pop('''text_pair'''					,   __UpperCAmelCase    )
       a						    = kwargs.pop('''return_tensors'''					,   __UpperCAmelCase    )

       a						    = {
           '''input_ids''': [],
           '''attention_mask''': [],
           '''token_type_ids''': [],
       }

       for idx, candidate_text in enumerate(__UpperCAmelCase    ):
        if batch_text_pair is not None:
         a						    = batch_text_pair[idx]
        else:
         a						    = None

        a						    = super().__call__(__UpperCAmelCase					,   __UpperCAmelCase					,   return_tensors=__UpperCAmelCase					,   **__UpperCAmelCase    )

        a						    = encoded_candidates.get('''input_ids'''    )
        a						    = encoded_candidates.get('''attention_mask'''    )
        a						    = encoded_candidates.get('''token_type_ids'''    )

        if encoded_input_ids is not None:
         output_data["input_ids"].append(__UpperCAmelCase    )
        if encoded_attention_mask is not None:
         output_data["attention_mask"].append(__UpperCAmelCase    )
        if encoded_token_type_ids is not None:
         output_data["token_type_ids"].append(__UpperCAmelCase    )

       a						    = {key: item for key, item in output_data.items() if len(__UpperCAmelCase    ) != 0}

       return BatchEncoding(__UpperCAmelCase					,   tensor_type=__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Tuple					,   __UpperCAmelCase :		Optional[int]					,   __UpperCAmelCase :		Dict=None    )       ->str:




       """simple docstring"""


       a						    = [self.cls_token_id] + token_ids_a + [self.sep_token_id]

       if token_ids_a:
        output += token_ids_a + [self.sep_token_id]

       return output







      def  __lowerCAmelCase   ( self :		str					,   __UpperCAmelCase :		List[int]					,   __UpperCAmelCase :		Optional[List[int]] = None    )       ->List[int]:




       """simple docstring"""


       a						    = [self.sep_token_id]
       a						    = [self.cls_token_id]
       if token_ids_a is None:
        return len(cls + token_ids_a + sep    ) * [0]
       return len(cls + token_ids_a + sep    ) * [0] + len(token_ids_a + sep    ) * [1]







      def  __lowerCAmelCase   ( self :		Dict					,   __UpperCAmelCase :		str					,   __UpperCAmelCase :		Optional[str] = None    )       ->Tuple[str]:




       """simple docstring"""


       a						    = self._tokenizer.model.save(__UpperCAmelCase					,   name=__UpperCAmelCase    )
       return tuple(__UpperCAmelCase    )


