






"""simple docstring"""







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    lowerCamelCase__   (			lowerCAmelCase__       ,	unittest.TestCase					):
    SCREAMING_SNAKE_CASE						=     KandinskyImgaImgPipeline
    SCREAMING_SNAKE_CASE						=     ['''prompt''', '''image_embeds''', '''negative_image_embeds''', '''image''']
    SCREAMING_SNAKE_CASE						=     [
        '''prompt''',
        '''negative_prompt''',
        '''image_embeds''',
        '''negative_image_embeds''',
        '''image''',
    ]
    SCREAMING_SNAKE_CASE						=     [
        '''generator''',
        '''height''',
        '''width''',
        '''strength''',
        '''guidance_scale''',
        '''negative_prompt''',
        '''num_inference_steps''',
        '''return_dict''',
        '''guidance_scale''',
        '''num_images_per_prompt''',
        '''output_type''',
        '''return_dict''',
    ]
    SCREAMING_SNAKE_CASE						=     False






    @property
    def 							_UpperCamelCase				(					self   ):
          return 32






    @property
    def 							_UpperCamelCase				(					self   ):
          return 32






    @property
    def 							_UpperCamelCase				(					self   ):
          return self.time_input_dim






    @property
    def 							_UpperCamelCase				(					self   ):
          return self.time_input_dim * 4






    @property
    def 							_UpperCamelCase				(					self   ):
          return 100






    @property
    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			XLMRobertaTokenizerFast.from_pretrained("""YiYiXu/tiny-random-mclip-base"""   )
          return tokenizer






    @property
    def 							_UpperCamelCase				(					self   ):
          torch.manual_seed(0   )
          UpperCAmelCase							     =			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=1_005      ,)

          UpperCAmelCase							     =			MultilingualCLIP(__lowerCAmelCase   )
          UpperCAmelCase							     =			text_encoder.eval()

          return text_encoder






    @property
    def 							_UpperCamelCase				(					self   ):
          torch.manual_seed(0   )

          UpperCAmelCase							     =			{
              """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,
          }

          UpperCAmelCase							     =			UNetaDConditionModel(**__lowerCAmelCase   )
          return model






    @property
    def 							_UpperCamelCase				(					self   ):
          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   ):
          torch.manual_seed(0   )
          UpperCAmelCase							     =			VQModel(**self.dummy_movq_kwargs   )
          return model






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			self.dummy_text_encoder
          UpperCAmelCase							     =			self.dummy_tokenizer
          UpperCAmelCase							     =			self.dummy_unet
          UpperCAmelCase							     =			self.dummy_movq

          UpperCAmelCase							     =			{
              """num_train_timesteps""": 1_000,
              """beta_schedule""": """linear""",
              """beta_start""": 0.00085,
              """beta_end""": 0.012,
              """clip_sample""": False,
              """set_alpha_to_one""": False,
              """steps_offset""": 0,
              """prediction_type""": """epsilon""",
              """thresholding""": False,
          }

          UpperCAmelCase							     =			DDIMScheduler(**__lowerCAmelCase   )

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

          return components






    def 							_UpperCamelCase				(					self      ,A      ,A=0   ):
          UpperCAmelCase							     =			floats_tensor((1, self.cross_attention_dim)      ,rng=random.Random(__lowerCAmelCase   )   ).to(__lowerCAmelCase   )
          UpperCAmelCase							     =			floats_tensor((1, self.cross_attention_dim)      ,rng=random.Random(seed + 1   )   ).to(__lowerCAmelCase   )
          # create init_image
          UpperCAmelCase							     =			floats_tensor((1, 3, 64, 64)      ,rng=random.Random(__lowerCAmelCase   )   ).to(__lowerCAmelCase   )
          UpperCAmelCase							     =			image.cpu().permute(0      ,2      ,3      ,1   )[0]
          UpperCAmelCase							     =			Image.fromarray(np.uinta(__lowerCAmelCase   )   ).convert("""RGB"""   ).resize((256, 256)   )

          if str(__lowerCAmelCase   ).startswith("""mps"""   ):
                UpperCAmelCase							     =			torch.manual_seed(__lowerCAmelCase   )
          else:
                UpperCAmelCase							     =			torch.Generator(device=__lowerCAmelCase   ).manual_seed(__lowerCAmelCase   )
          UpperCAmelCase							     =			{
              """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 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			"""cpu"""

          UpperCAmelCase							     =			self.get_dummy_components()

          UpperCAmelCase							     =			self.pipeline_class(**__lowerCAmelCase   )
          UpperCAmelCase							     =			pipe.to(__lowerCAmelCase   )

          pipe.set_progress_bar_config(disable=__lowerCAmelCase   )

          UpperCAmelCase							     =			pipe(**self.get_dummy_inputs(__lowerCAmelCase   )   )
          UpperCAmelCase							     =			output.images

          UpperCAmelCase							     =			pipe(
              **self.get_dummy_inputs(__lowerCAmelCase   )      ,return_dict=__lowerCAmelCase      ,)[0]

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

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

          UpperCAmelCase							     =			np.array(
              [0.61474943, 0.6073539, 0.43308544, 0.5928269, 0.47493595, 0.46755973, 0.4613838, 0.45368797, 0.50119233]   )
          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    lowerCamelCase__   (			unittest.TestCase					):






    def 							_UpperCamelCase				(					self   ):
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			load_numpy(
              """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
              """/kandinsky/kandinsky_img2img_frog.npy"""   )

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

          UpperCAmelCase							     =			KandinskyPriorPipeline.from_pretrained(
              """kandinsky-community/kandinsky-2-1-prior"""      ,torch_dtype=torch.floataa   )
          pipe_prior.to(__lowerCAmelCase   )

          UpperCAmelCase							     =			KandinskyImgaImgPipeline.from_pretrained(
              """kandinsky-community/kandinsky-2-1"""      ,torch_dtype=torch.floataa   )
          UpperCAmelCase							     =			pipeline.to(__lowerCAmelCase   )

          pipeline.set_progress_bar_config(disable=__lowerCAmelCase   )

          UpperCAmelCase							     =			torch.Generator(device="""cpu"""   ).manual_seed(0   )
          UpperCAmelCase     ,   UpperCAmelCase							     =			pipe_prior(
              __lowerCAmelCase      ,generator=__lowerCAmelCase      ,num_inference_steps=5      ,negative_prompt=""""""      ,).to_tuple()

          UpperCAmelCase							     =			pipeline(
              __lowerCAmelCase      ,image=__lowerCAmelCase      ,image_embeds=__lowerCAmelCase      ,negative_image_embeds=__lowerCAmelCase      ,generator=__lowerCAmelCase      ,num_inference_steps=100      ,height=768      ,width=768      ,strength=0.2      ,output_type="""np"""      ,)

          UpperCAmelCase							     =			output.images[0]

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

          assert_mean_pixel_difference(__lowerCAmelCase      ,__lowerCAmelCase   )




'''simple docstring'''




import warnings
from typing import List, Optional, Union

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType






class 							a__(   lowerCAmelCase__      ):


			'''simple docstring'''






			UpperCAmelCase_     :	str						=			['''image_processor''', '''tokenizer''']
			UpperCAmelCase_     :	str						=			'''ViltImageProcessor'''
			UpperCAmelCase_     :	Any						=			('''BertTokenizer''', '''BertTokenizerFast''')



			def __init__(  self			,       __lowerCAmelCase=None			,       __lowerCAmelCase=None			,       **__lowerCAmelCase):




					"""simple docstring"""


					lowerCAmelCase     =    None
					if "feature_extractor" in kwargs:
							warnings.warn(
							    """The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"""
							    """ instead."""			,       __lowerCAmelCase			,       )
							lowerCAmelCase     =    kwargs.pop("""feature_extractor""")

					lowerCAmelCase     =    image_processor if image_processor is not None else feature_extractor
					if image_processor is None:
							raise ValueError("""You need to specify an `image_processor`.""")
					if tokenizer is None:
							raise ValueError("""You need to specify a `tokenizer`.""")

					super().__init__(__lowerCAmelCase			,       __lowerCAmelCase)
					lowerCAmelCase     =    self.image_processor



			def __call__(  self			,       __lowerCAmelCase			,       __lowerCAmelCase = None			,       __lowerCAmelCase = True			,       __lowerCAmelCase = False			,       __lowerCAmelCase = None			,       __lowerCAmelCase = None			,       __lowerCAmelCase = 0			,       __lowerCAmelCase = None			,       __lowerCAmelCase = None			,       __lowerCAmelCase = None			,       __lowerCAmelCase = False			,       __lowerCAmelCase = False			,       __lowerCAmelCase = False			,       __lowerCAmelCase = False			,       __lowerCAmelCase = True			,       __lowerCAmelCase = None			,       **__lowerCAmelCase			,       ):




					"""simple docstring"""


					lowerCAmelCase     =    self.tokenizer(
					    text=__lowerCAmelCase			,       add_special_tokens=__lowerCAmelCase			,       padding=__lowerCAmelCase			,       truncation=__lowerCAmelCase			,       max_length=__lowerCAmelCase			,       stride=__lowerCAmelCase			,       pad_to_multiple_of=__lowerCAmelCase			,       return_token_type_ids=__lowerCAmelCase			,       return_attention_mask=__lowerCAmelCase			,       return_overflowing_tokens=__lowerCAmelCase			,       return_special_tokens_mask=__lowerCAmelCase			,       return_offsets_mapping=__lowerCAmelCase			,       return_length=__lowerCAmelCase			,       verbose=__lowerCAmelCase			,       return_tensors=__lowerCAmelCase			,       **__lowerCAmelCase			,       )
					# add pixel_values + pixel_mask
					lowerCAmelCase     =    self.image_processor(__lowerCAmelCase			,       return_tensors=__lowerCAmelCase)
					encoding.update(__lowerCAmelCase)

					return encoding



			def 				a_			(  self			,       *__lowerCAmelCase			,       **__lowerCAmelCase):




					"""simple docstring"""


					return self.tokenizer.batch_decode(*__lowerCAmelCase			,       **__lowerCAmelCase)



			def 				a_			(  self			,       *__lowerCAmelCase			,       **__lowerCAmelCase):




					"""simple docstring"""


					return self.tokenizer.decode(*__lowerCAmelCase			,       **__lowerCAmelCase)



			@property
			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    self.tokenizer.model_input_names
					lowerCAmelCase     =    self.image_processor.model_input_names
					return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))



			@property
			def 				a_			(  self):




					"""simple docstring"""


					warnings.warn(
					    """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."""			,       __lowerCAmelCase			,       )
					return self.image_processor_class







			@property
			def 				a_			(  self):




					"""simple docstring"""


					warnings.warn(
					    """`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."""			,       __lowerCAmelCase			,       )
					return self.image_processor




