
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    _a				(__SCREAMING_SNAKE_CASE  ,					unittest.TestCase):





       """simple docstring"""




       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      )    ->			Dict:
          return 32






       @property
       def        UpperCamelCase				(							self      )    ->			Optional[Any]:
          return 32






       @property
       def        UpperCamelCase				(							self      )    ->			Any:
          return self.time_input_dim






       @property
       def        UpperCamelCase				(							self      )    ->			Tuple:
          return self.time_input_dim * 4






       @property
       def        UpperCamelCase				(							self      )    ->			Dict:
          return 1_00






       @property
       def        UpperCamelCase				(							self      )    ->			List[str]:
          _SCREAMING_SNAKE_CASE      =				XLMRobertaTokenizerFast.from_pretrained("""YiYiXu/tiny-random-mclip-base"""      )
          return tokenizer






       @property
       def        UpperCamelCase				(							self      )    ->			Any:
          torch.manual_seed(0      )
          _SCREAMING_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=10_05       ,							)

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

          return text_encoder






       @property
       def        UpperCamelCase				(							self      )    ->			Dict:
          torch.manual_seed(0      )

          _SCREAMING_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,
          }

          _SCREAMING_SNAKE_CASE      =				UNetaDConditionModel(**__snake_case      )
          return model






       @property
       def        UpperCamelCase				(							self      )    ->			Dict:
          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      )    ->			Union[str, Any]:
          torch.manual_seed(0      )
          _SCREAMING_SNAKE_CASE      =				VQModel(**self.dummy_movq_kwargs      )
          return model






       def        UpperCamelCase				(							self      )    ->			int:
          _SCREAMING_SNAKE_CASE      =				self.dummy_text_encoder
          _SCREAMING_SNAKE_CASE      =				self.dummy_tokenizer
          _SCREAMING_SNAKE_CASE      =				self.dummy_unet
          _SCREAMING_SNAKE_CASE      =				self.dummy_movq

          _SCREAMING_SNAKE_CASE      =				{
              '''num_train_timesteps''': 10_00,
              '''beta_schedule''': '''linear''',
              '''beta_start''': 0.0_0085,
              '''beta_end''': 0.012,
              '''clip_sample''': False,
              '''set_alpha_to_one''': False,
              '''steps_offset''': 0,
              '''prediction_type''': '''epsilon''',
              '''thresholding''': False,
          }

          _SCREAMING_SNAKE_CASE      =				DDIMScheduler(**__snake_case      )

          _SCREAMING_SNAKE_CASE      =				{
              '''text_encoder''': text_encoder,
              '''tokenizer''': tokenizer,
              '''unet''': unet,
              '''scheduler''': scheduler,
              '''movq''': movq,
          }

          return components






       def        UpperCamelCase				(							self       ,							A__       ,							A__=0      )    ->			Tuple:
          _SCREAMING_SNAKE_CASE      =				floats_tensor((1, self.cross_attention_dim)       ,							rng=random.Random(__snake_case      )      ).to(__snake_case      )
          _SCREAMING_SNAKE_CASE      =				floats_tensor((1, self.cross_attention_dim)       ,							rng=random.Random(seed + 1      )      ).to(__snake_case      )
          # create init_image
          _SCREAMING_SNAKE_CASE      =				floats_tensor((1, 3, 64, 64)       ,							rng=random.Random(__snake_case      )      ).to(__snake_case      )
          _SCREAMING_SNAKE_CASE      =				image.cpu().permute(0       ,							2       ,							3       ,							1      )[0]
          _SCREAMING_SNAKE_CASE      =				Image.fromarray(np.uinta(__snake_case      )      ).convert("""RGB"""      ).resize((2_56, 2_56)      )

          if str(__snake_case      ).startswith("""mps"""      ):
             _SCREAMING_SNAKE_CASE      =				torch.manual_seed(__snake_case      )
          else:
             _SCREAMING_SNAKE_CASE      =				torch.Generator(device=__snake_case      ).manual_seed(__snake_case      )
          _SCREAMING_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        UpperCamelCase				(							self      )    ->			Any:
          _SCREAMING_SNAKE_CASE      =				'''cpu'''

          _SCREAMING_SNAKE_CASE      =				self.get_dummy_components()

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

          pipe.set_progress_bar_config(disable=__snake_case      )

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

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

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

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

          _SCREAMING_SNAKE_CASE      =				np.array(
              [0.6147_4943, 0.607_3539, 0.4330_8544, 0.592_8269, 0.4749_3595, 0.4675_5973, 0.461_3838, 0.4536_8797, 0.5011_9233]      )
          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    _a				(unittest.TestCase):





       """simple docstring"""




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






       def        UpperCamelCase				(							self      )    ->			List[Any]:
          _SCREAMING_SNAKE_CASE      =				load_numpy(
              """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
              """/kandinsky/kandinsky_img2img_frog.npy"""      )

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

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

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

          pipeline.set_progress_bar_config(disable=__snake_case      )

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

          _SCREAMING_SNAKE_CASE      =				pipeline(
              __snake_case       ,							image=__snake_case       ,							image_embeds=__snake_case       ,							negative_image_embeds=__snake_case       ,							generator=__snake_case       ,							num_inference_steps=1_00       ,							height=7_68       ,							width=7_68       ,							strength=0.2       ,							output_type="""np"""       ,							)

          _SCREAMING_SNAKE_CASE      =				output.images[0]

          assert image.shape == (7_68, 7_68, 3)

          assert_mean_pixel_difference(__snake_case       ,							__snake_case      )








'''simple docstring'''



import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding

class    _a				(_lowerCamelCase):





       """simple docstring"""




       SCREAMING_SNAKE_CASE      =  ['image_processor', 'tokenizer']
       SCREAMING_SNAKE_CASE      =  'ChineseCLIPImageProcessor'
       SCREAMING_SNAKE_CASE      =  ('BertTokenizer', 'BertTokenizerFast')






       def __init__(							self       ,							A__=None       ,							A__=None       ,							**A__      )    ->			int:
          _SCREAMING_SNAKE_CASE      =				None
          if "feature_extractor" in kwargs:
             warnings.warn(
                 """The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"""
                 """ instead."""       ,							A__       ,							)
             _SCREAMING_SNAKE_CASE      =				kwargs.pop("""feature_extractor"""      )

          _SCREAMING_SNAKE_CASE      =				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__(A__       ,							A__      )
          _SCREAMING_SNAKE_CASE      =				self.image_processor






       def __call__(							self       ,							A__=None       ,							A__=None       ,							A__=None       ,							**A__      )    ->			Optional[int]:

          if text is None and images is None:
             raise ValueError("""You have to specify either text or images. Both cannot be none."""      )

          if text is not None:
             _SCREAMING_SNAKE_CASE      =				self.tokenizer(A__       ,							return_tensors=A__       ,							**A__      )

          if images is not None:
             _SCREAMING_SNAKE_CASE      =				self.image_processor(A__       ,							return_tensors=A__       ,							**A__      )

          if text is not None and images is not None:
             _SCREAMING_SNAKE_CASE      =				image_features.pixel_values
             return encoding
          elif text is not None:
             return encoding
          else:
             return BatchEncoding(data=dict(**A__      )       ,							tensor_type=A__      )






       def        UpperCamelCase				(							self       ,							*A__       ,							**A__      )    ->			Dict:
          return self.tokenizer.batch_decode(*A__       ,							**A__      )






       def        UpperCamelCase				(							self       ,							*A__       ,							**A__      )    ->			Optional[Any]:
          return self.tokenizer.decode(*A__       ,							**A__      )






       @property
       def        UpperCamelCase				(							self      )    ->			List[Any]:
          _SCREAMING_SNAKE_CASE      =				self.tokenizer.model_input_names
          _SCREAMING_SNAKE_CASE      =				self.image_processor.model_input_names
          return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names      )      )






       @property
       def        UpperCamelCase				(							self      )    ->			Optional[int]:
          warnings.warn(
              """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."""       ,							A__       ,							)
          return self.image_processor_class






