

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   lowercase_   (  __lowercase       ,							unittest.TestCase					):
      UpperCamelCase_      :		Union[str, Any]										=    KandinskyImgaImgPipeline
      UpperCamelCase_      :		Optional[Any]										=    ["prompt", "image_embeds", "negative_image_embeds", "image"]
      UpperCamelCase_      :		Optional[int]										=    [
          "prompt",
          "negative_prompt",
          "image_embeds",
          "negative_image_embeds",
          "image",
      ]
      UpperCamelCase_      :		Union[str, Any]										=    [
          "generator",
          "height",
          "width",
          "strength",
          "guidance_scale",
          "negative_prompt",
          "num_inference_steps",
          "return_dict",
          "guidance_scale",
          "num_images_per_prompt",
          "output_type",
          "return_dict",
      ]
      UpperCamelCase_      :		Optional[int]										=    False



      @property
      def 						UpperCamelCase_		(      self				:    str		)   ->					Optional[int]:
        return 32



      @property
      def 						UpperCamelCase_		(      self				:    Dict		)   ->					int:
        return 32



      @property
      def 						UpperCamelCase_		(      self				:    int		)   ->					Tuple:
        return self.time_input_dim



      @property
      def 						UpperCamelCase_		(      self				:    Dict		)   ->					List[str]:
        return self.time_input_dim * 4



      @property
      def 						UpperCamelCase_		(      self				:    str		)   ->					str:
        return 100



      @property
      def 						UpperCamelCase_		(      self				:    Dict		)   ->					Optional[int]:
        _snake_case						=							XLMRobertaTokenizerFast.from_pretrained('''YiYiXu/tiny-random-mclip-base'''		)
        return tokenizer



      @property
      def 						UpperCamelCase_		(      self				:    Union[str, Any]		)   ->					int:
        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 						UpperCamelCase_		(      self				:    List[str]		)   ->					Tuple:
        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 						UpperCamelCase_		(      self				:    Any		)   ->					str:
        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				:    Optional[Any]		)   ->					Tuple:
        torch.manual_seed(0		)
        _snake_case						=							VQModel(**self.dummy_movq_kwargs		)
        return model



      def 						UpperCamelCase_		(      self				:    Dict		)   ->					Any:
        _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_0085,
            '''beta_end''': 0.012,
            '''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 						UpperCamelCase_		(      self				:    Union[str, Any]      , A__				:    str      , A__				:    Optional[int]=0		)   ->					int:
        _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 						UpperCamelCase_		(      self				:    List[str]		)   ->					str:
        _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.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   lowercase_   (  unittest.TestCase					):



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







      def 						UpperCamelCase_		(      self				:    Optional[Any]		)   ->					str:
        _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,					_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__		)




__A      		=     '''ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'''



def 				snake_case_(_UpperCamelCase					)							->     bytes:



  """simple docstring"""






  if not isinstance(_UpperCamelCase			,	_UpperCamelCase					):
    _snake_case						=							F"""a bytes-like object is required, not '{data.__class__.__name__}'"""
    raise TypeError(_UpperCamelCase					)

  _snake_case						=							''''''.join(bin(_UpperCamelCase					)[2:].zfill(8					) for byte in data					)

  _snake_case						=							len(_UpperCamelCase					) % 6 != 0

  if padding_needed:
    # The padding that will be added later
    _snake_case						=							b'''=''' * ((6 - len(_UpperCamelCase					) % 6) // 2)

    # Append binary_stream with arbitrary binary digits (0's by default) to make its
    # length a multiple of 6.
    binary_stream += "0" * (6 - len(_UpperCamelCase					) % 6)
  else:
    _snake_case						=							b''''''

  # Encode every 6 binary digits to their corresponding Base64 character
  return (
      "".join(
          B64_CHARSET[int(binary_stream[index : index + 6]			,	2					)]
          for index in range(0			,	len(_UpperCamelCase					)			,	6					)					).encode()
      + padding
  )



def 				snake_case_(_UpperCamelCase					)							->     bytes:



  """simple docstring"""






  if not isinstance(_UpperCamelCase			,	_UpperCamelCase					) and not isinstance(_UpperCamelCase			,	_UpperCamelCase					):
    _snake_case						=							(
        '''argument should be a bytes-like object or ASCII string, '''
        F"""not '{encoded_data.__class__.__name__}'"""
    )
    raise TypeError(_UpperCamelCase					)

  # In case encoded_data is a bytes-like object, make sure it contains only
  # ASCII characters so we convert it to a string object
  if isinstance(_UpperCamelCase			,	_UpperCamelCase					):
    try:
      _snake_case						=							encoded_data.decode('''utf-8'''					)
    except UnicodeDecodeError:
      raise ValueError('''base64 encoded data should only contain ASCII characters'''					)

  _snake_case						=							encoded_data.count('''='''					)

  # Check if the encoded string contains non base64 characters
  if padding:
    assert all(
        char in B64_CHARSET for char in encoded_data[:-padding]					), "Invalid base64 character(s) found."
  else:
    assert all(
        char in B64_CHARSET for char in encoded_data					), "Invalid base64 character(s) found."

  # Check the padding
  assert len(_UpperCamelCase					) % 4 == 0 and padding < 3, "Incorrect padding"

  if padding:
    # Remove padding if there is one
    _snake_case						=							encoded_data[:-padding]

    _snake_case						=							''''''.join(
        bin(B64_CHARSET.index(_UpperCamelCase					)					)[2:].zfill(6					) for char in encoded_data					)[: -padding * 2]
  else:
    _snake_case						=							''''''.join(
        bin(B64_CHARSET.index(_UpperCamelCase					)					)[2:].zfill(6					) for char in encoded_data					)

  _snake_case						=							[
      int(binary_stream[index : index + 8]			,	2					)
      for index in range(0			,	len(_UpperCamelCase					)			,	8					)
  ]

  return bytes(_UpperCamelCase					)


if __name__ == "__main__":
 import doctest

 doctest.testmod()

