





import gc
import random
import unittest

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

from diffusers import DDIMScheduler, KandinskyImgaImgPipeline, 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 		SCREAMING_SNAKE_CASE__    (	__lowercase						,				unittest.TestCase  ):
      A_						:    int							=   KandinskyImgaImgPipeline
      A_						:    Any							=   ['''prompt''', '''image_embeds''', '''negative_image_embeds''', '''image''']
      A_						:    int							=   [
          '''prompt''',
          '''negative_prompt''',
          '''image_embeds''',
          '''negative_image_embeds''',
          '''image''',
      ]
      A_						:    List[Any]							=   [
          '''generator''',
          '''height''',
          '''width''',
          '''strength''',
          '''guidance_scale''',
          '''negative_prompt''',
          '''num_inference_steps''',
          '''return_dict''',
          '''guidance_scale''',
          '''num_images_per_prompt''',
          '''output_type''',
          '''return_dict''',
      ]
      A_						:    Any							=   False


      @property
      def        a			(self	:				Optional[Any]				):





         """simple docstring"""



         return 32


      @property
      def        a			(self	:				str				):





         """simple docstring"""



         return 32


      @property
      def        a			(self	:				Tuple				):





         """simple docstring"""



         return self.time_input_dim


      @property
      def        a			(self	:				List[str]				):





         """simple docstring"""



         return self.time_input_dim * 4


      @property
      def        a			(self	:				Dict				):





         """simple docstring"""



         return 100


      @property
      def        a			(self	:				List[str]				):





         """simple docstring"""



         __snake_case								=       XLMRobertaTokenizerFast.from_pretrained('''YiYiXu/tiny-random-mclip-base'''				)
         return tokenizer


      @property
      def        a			(self	:				Union[str, Any]				):





         """simple docstring"""



         torch.manual_seed(0				)
         __snake_case								=       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    ,      )

         __snake_case								=       MultilingualCLIP(a__				)
         __snake_case								=       text_encoder.eval()

         return text_encoder


      @property
      def        a			(self	:				Tuple				):





         """simple docstring"""



         torch.manual_seed(0				)

         __snake_case								=       {
             """in_channels""": 4,
             # 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,
         }

         __snake_case								=       UNetaDConditionModel(**a__				)
         return model


      @property
      def        a			(self	:				str				):





         """simple docstring"""



         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        a			(self	:				Optional[Any]				):





         """simple docstring"""



         torch.manual_seed(0				)
         __snake_case								=       VQModel(**self.dummy_movq_kwargs				)
         return model


      def        a			(self	:				Tuple				):





         """simple docstring"""



         __snake_case								=       self.dummy_text_encoder
         __snake_case								=       self.dummy_tokenizer
         __snake_case								=       self.dummy_unet
         __snake_case								=       self.dummy_movq

         __snake_case								=       {
             """num_train_timesteps""": 1000,
             """beta_schedule""": """linear""",
             """beta_start""": 0.0_0_0_8_5,
             """beta_end""": 0.0_1_2,
             """clip_sample""": False,
             """set_alpha_to_one""": False,
             """steps_offset""": 0,
             """prediction_type""": """epsilon""",
             """thresholding""": False,
         }

         __snake_case								=       DDIMScheduler(**a__				)

         __snake_case								=       {
             """text_encoder""": text_encoder,
             """tokenizer""": tokenizer,
             """unet""": unet,
             """scheduler""": scheduler,
             """movq""": movq,
         }

         return components


      def        a			(self	:				str    ,      a__	:				str    ,      a__	:				Union[str, Any]=0				):





         """simple docstring"""



         __snake_case								=       floats_tensor((1, self.cross_attention_dim)    ,      rng=random.Random(a__				)				).to(a__				)
         __snake_case								=       floats_tensor((1, self.cross_attention_dim)    ,      rng=random.Random(seed + 1				)				).to(a__				)
         # create init_image
         __snake_case								=       floats_tensor((1, 3, 64, 64)    ,      rng=random.Random(a__				)				).to(a__				)
         __snake_case								=       image.cpu().permute(0    ,      2    ,      3    ,      1				)[0]
         __snake_case								=       Image.fromarray(np.uinta(a__				)				).convert('''RGB'''				).resize((256, 256)				)

         if str(a__				).startswith('''mps'''				):
            __snake_case								=       torch.manual_seed(a__				)
         else:
            __snake_case								=       torch.Generator(device=a__				).manual_seed(a__				)
         __snake_case								=       {
             """prompt""": """horse""",
             """image""": init_image,
             """image_embeds""": image_embeds,
             """negative_image_embeds""": negative_image_embeds,
             """generator""": generator,
             """height""": 64,
             """width""": 64,
             """num_inference_steps""": 10,
             """guidance_scale""": 7.0,
             """strength""": 0.2,
             """output_type""": """np""",
         }
         return inputs


      def        a			(self	:				int				):





         """simple docstring"""



         __snake_case								=       """cpu"""

         __snake_case								=       self.get_dummy_components()

         __snake_case								=       self.pipeline_class(**a__				)
         __snake_case								=       pipe.to(a__				)

         pipe.set_progress_bar_config(disable=a__				)

         __snake_case								=       pipe(**self.get_dummy_inputs(a__				)				)
         __snake_case								=       output.images

         __snake_case								=       pipe(
             **self.get_dummy_inputs(a__				)    ,      return_dict=a__    ,      )[0]

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

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

         __snake_case								=       np.array(
             [0.6_1_4_7_4_9_4_3, 0.6_0_7_3_5_3_9, 0.4_3_3_0_8_5_4_4, 0.5_9_2_8_2_6_9, 0.4_7_4_9_3_5_9_5, 0.4_6_7_5_5_9_7_3, 0.4_6_1_3_8_3_8, 0.4_5_3_6_8_7_9_7, 0.5_0_1_1_9_2_3_3]				)
         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()}"""






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


      def        a			(self	:				List[str]				):





         """simple docstring"""



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


      def        a			(self	:				Optional[int]				):





         """simple docstring"""



         __snake_case								=       load_numpy(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/kandinsky/kandinsky_img2img_frog.npy'''				)

         __snake_case								=       load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/kandinsky/cat.png'''				)
         __snake_case								=       """A red cartoon frog, 4k"""

         __snake_case								=       KandinskyPriorPipeline.from_pretrained(
             '''kandinsky-community/kandinsky-2-1-prior'''    ,      torch_dtype=torch.floataa				)
         pipe_prior.to(a__				)

         __snake_case								=       KandinskyImgaImgPipeline.from_pretrained(
             '''kandinsky-community/kandinsky-2-1'''    ,      torch_dtype=torch.floataa				)
         __snake_case								=       pipeline.to(a__				)

         pipeline.set_progress_bar_config(disable=a__				)

         __snake_case								=       torch.Generator(device='''cpu'''				).manual_seed(0				)
         __snake_case								=       pipe_prior(
             a__    ,      generator=a__    ,      num_inference_steps=5    ,      negative_prompt=''''''    ,      ).to_tuple()

         __snake_case								=       pipeline(
             a__    ,      image=a__    ,      image_embeds=a__    ,      negative_image_embeds=a__    ,      generator=a__    ,      num_inference_steps=100    ,      height=768    ,      width=768    ,      strength=0.2    ,      output_type='''np'''    ,      )

         __snake_case								=       output.images[0]

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

         assert_mean_pixel_difference(a__    ,      a__				)










import json
from typing import TYPE_CHECKING, List, Optional, Tuple

from tokenizers import pre_tokenizers

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging


if TYPE_CHECKING:
    from transformers.pipelines.conversational import Conversation


snake_case_           =      logging.get_logger(__name__)

snake_case_           =      {'tokenizer_file': 'tokenizer.json'}

snake_case_           =      {
    'tokenizer_file': {
        'bigscience/tokenizer': 'https://huggingface.co/bigscience/tokenizer/blob/main/tokenizer.json',
        'bigscience/bloom-560m': 'https://huggingface.co/bigscience/bloom-560m/blob/main/tokenizer.json',
        'bigscience/bloom-1b1': 'https://huggingface.co/bigscience/bloom-1b1/blob/main/tokenizer.json',
        'bigscience/bloom-1b7': 'https://huggingface.co/bigscience/bloom-1b7/blob/main/tokenizer.json',
        'bigscience/bloom-3b': 'https://huggingface.co/bigscience/bloom-3b/blob/main/tokenizer.json',
        'bigscience/bloom-7b1': 'https://huggingface.co/bigscience/bloom-7b1/blob/main/tokenizer.json',
        'bigscience/bloom': 'https://huggingface.co/bigscience/bloom/blob/main/tokenizer.json',
    },
}

class 		SCREAMING_SNAKE_CASE__    (	_UpperCAmelCase  ):
      A_						:    List[str]							=   VOCAB_FILES_NAMES
      A_						:    Optional[Any]							=   PRETRAINED_VOCAB_FILES_MAP
      A_						:    str							=   ['input_ids', 'attention_mask']
      A_						:    Optional[Any]							=   None


      def __init__(self	:				Optional[int]    ,      a__	:				int=None    ,      a__	:				str=None    ,      a__	:				Any=None    ,      a__	:				List[Any]="<unk>"    ,      a__	:				List[Any]="<s>"    ,      a__	:				Optional[int]="</s>"    ,      a__	:				List[str]="<pad>"    ,      a__	:				Union[str, Any]=False    ,      a__	:				str=False    ,      **a__	:				Optional[Any]    ,      ):





         """simple docstring"""



         super().__init__(
             a__    ,      a__    ,      tokenizer_file=a__    ,      unk_token=a__    ,      bos_token=a__    ,      eos_token=a__    ,      pad_token=a__    ,      add_prefix_space=a__    ,      clean_up_tokenization_spaces=a__    ,      **a__    ,      )
         __snake_case								=       json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()				)
         if pre_tok_state.get('''add_prefix_space'''    ,      a__				) != add_prefix_space:
            __snake_case								=       getattr(a__    ,      pre_tok_state.pop('''type'''				)				)
            __snake_case								=       add_prefix_space
            __snake_case								=       pre_tok_class(**a__				)

         __snake_case								=       add_prefix_space


      def        a			(self	:				int    ,      *a__	:				Tuple    ,      **a__	:				Optional[Any]				):





         """simple docstring"""



         __snake_case								=       kwargs.get('''is_split_into_words'''    ,      a__				)
         if not (self.add_prefix_space or not is_split_into_words):
            raise Exception(
                f"""You need to instantiate {self.__class__.__name__} with add_prefix_space=True to use it with"""
                ''' pretokenized inputs.'''				)

         return super()._batch_encode_plus(*a__    ,      **a__				)


      def        a			(self	:				List[str]    ,      *a__	:				List[str]    ,      **a__	:				List[str]				):





         """simple docstring"""



         __snake_case								=       kwargs.get('''is_split_into_words'''    ,      a__				)

         if not (self.add_prefix_space or not is_split_into_words):
            raise Exception(
                f"""You need to instantiate {self.__class__.__name__} with add_prefix_space=True to use it with"""
                ''' pretokenized inputs.'''				)

         return super()._encode_plus(*a__    ,      **a__				)


      def        a			(self	:				List[Any]    ,      a__	:				str    ,      a__	:				Optional[str] = None				):





         """simple docstring"""



         __snake_case								=       self._tokenizer.model.save(a__    ,      name=a__				)
         return tuple(a__				)


      def        a			(self	:				Tuple    ,      a__	:				"Conversation"				):





         """simple docstring"""



         __snake_case								=       []
         for is_user, text in conversation.iter_texts():
            input_ids.extend(self.encode(a__    ,      add_special_tokens=a__				) + [self.eos_token_id]				)

         if len(a__				) > self.model_max_length:
            __snake_case								=       input_ids[-self.model_max_length :]
         return input_ids

