






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



      """simple docstring"""


      a    :					Optional[int]							=      KandinskyImgaImgPipeline
      a    :					List[Any]							=      ['prompt', 'image_embeds', 'negative_image_embeds', 'image']
      a    :					Union[str, Any]							=      [
          'prompt',
          'negative_prompt',
          'image_embeds',
          'negative_image_embeds',
          'image',
      ]
      a    :					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',
      ]
      a    :					Optional[int]							=      False
      @property
      def       UpperCAmelCase ( self   :     Optional[int]   )				-> List[Any]:
        return 32
      @property
      def       UpperCAmelCase ( self   :     Optional[int]   )				-> Any:
        return 32
      @property
      def       UpperCAmelCase ( self   :     List[Any]   )				-> Optional[Any]:
        return self.time_input_dim
      @property
      def       UpperCAmelCase ( self   :     List[Any]   )				-> Union[str, Any]:
        return self.time_input_dim * 4
      @property
      def       UpperCAmelCase ( self   :     Any   )				-> Any:
        return 100
      @property
      def       UpperCAmelCase ( self   :     int   )				-> str:
        __UpperCAmelCase				:      List[str]	 =    XLMRobertaTokenizerFast.from_pretrained("""YiYiXu/tiny-random-mclip-base"""   )
        return tokenizer
      @property
      def       UpperCAmelCase ( self   :     Optional[int]   )				-> List[str]:
        torch.manual_seed(0   )
        __UpperCAmelCase				:      Dict	 =    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				:      int	 =    MultilingualCLIP(__lowercase   )
        __UpperCAmelCase				:      List[Any]	 =    text_encoder.eval()

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

        __UpperCAmelCase				:      List[Any]	 =    {
            """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				:      Dict	 =    UNetaDConditionModel(**__lowercase   )
        return model
      @property
      def       UpperCAmelCase ( self   :     Union[str, 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   :     int   )				-> Optional[int]:
        torch.manual_seed(0   )
        __UpperCAmelCase				:      str	 =    VQModel(**self.dummy_movq_kwargs   )
        return model
      def       UpperCAmelCase ( self   :     Optional[Any]   )				-> int:
        __UpperCAmelCase				:      str	 =    self.dummy_text_encoder
        __UpperCAmelCase				:      List[Any]	 =    self.dummy_tokenizer
        __UpperCAmelCase				:      Tuple	 =    self.dummy_unet
        __UpperCAmelCase				:      Any	 =    self.dummy_movq

        __UpperCAmelCase				:      List[Any]	 =    {
            """num_train_timesteps""": 1000,
            """beta_schedule""": """linear""",
            """beta_start""": 0.00_085,
            """beta_end""": 0.012,
            """clip_sample""": False,
            """set_alpha_to_one""": False,
            """steps_offset""": 0,
            """prediction_type""": """epsilon""",
            """thresholding""": False,
        }

        __UpperCAmelCase				:      Dict	 =    DDIMScheduler(**__lowercase   )

        __UpperCAmelCase				:      List[Any]	 =    {
            """text_encoder""": text_encoder,
            """tokenizer""": tokenizer,
            """unet""": unet,
            """scheduler""": scheduler,
            """movq""": movq,
        }

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

        if str(__lowercase   ).startswith("""mps"""   ):
          __UpperCAmelCase				:      List[Any]	 =    torch.manual_seed(__lowercase   )
        else:
          __UpperCAmelCase				:      Tuple	 =    torch.Generator(device=__lowercase   ).manual_seed(__lowercase   )
        __UpperCAmelCase				:      Dict	 =    {
            """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[Any]   )				-> str:
        __UpperCAmelCase				:      Any	 =    """cpu"""

        __UpperCAmelCase				:      Union[str, Any]	 =    self.get_dummy_components()

        __UpperCAmelCase				:      Optional[Any]	 =    self.pipeline_class(**__lowercase   )
        __UpperCAmelCase				:      Any	 =    pipe.to(__lowercase   )

        pipe.set_progress_bar_config(disable=__lowercase   )

        __UpperCAmelCase				:      List[Any]	 =    pipe(**self.get_dummy_inputs(__lowercase   )   )
        __UpperCAmelCase				:      Tuple	 =    output.images

        __UpperCAmelCase				:      Optional[int]	 =    pipe(
            **self.get_dummy_inputs(__lowercase   )						, return_dict=__lowercase						, )[0]

        __UpperCAmelCase				:      Any	 =    image[0, -3:, -3:, -1]
        __UpperCAmelCase				:      Union[str, Any]	 =    image_from_tuple[0, -3:, -3:, -1]

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

        __UpperCAmelCase				:      int	 =    np.array(
            [0.61_474_943, 0.6_073_539, 0.43_308_544, 0.5_928_269, 0.47_493_595, 0.46_755_973, 0.4_613_838, 0.45_368_797, 0.50_119_233]   )
        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   :     int   )				-> Optional[int]:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()




      def       UpperCAmelCase ( self   :     str   )				-> str:
        __UpperCAmelCase				:      Union[str, Any]	 =    load_numpy(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
            """/kandinsky/kandinsky_img2img_frog.npy"""   )

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

        __UpperCAmelCase				:      Tuple	 =    KandinskyPriorPipeline.from_pretrained(
            """kandinsky-community/kandinsky-2-1-prior"""						, torch_dtype=torch.floataa   )
        pipe_prior.to(__lowercase   )

        __UpperCAmelCase				:      List[str]	 =    KandinskyImgaImgPipeline.from_pretrained(
            """kandinsky-community/kandinsky-2-1"""						, torch_dtype=torch.floataa   )
        __UpperCAmelCase				:      str	 =    pipeline.to(__lowercase   )

        pipeline.set_progress_bar_config(disable=__lowercase   )

        __UpperCAmelCase				:      str	 =    torch.Generator(device="""cpu"""   ).manual_seed(0   )
        __UpperCAmelCase ,					__UpperCAmelCase				:      Union[str, Any]	 =    pipe_prior(
            __lowercase						, generator=__lowercase						, num_inference_steps=5						, negative_prompt=""""""						, ).to_tuple()

        __UpperCAmelCase				:      List[Any]	 =    pipeline(
            __lowercase						, image=__lowercase						, image_embeds=__lowercase						, negative_image_embeds=__lowercase						, generator=__lowercase						, num_inference_steps=100						, height=768						, width=768						, strength=0.2						, output_type="""np"""						, )

        __UpperCAmelCase				:      int	 =    output.images[0]

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

        assert_mean_pixel_difference(__lowercase						, __lowercase   )









'''simple docstring'''

import warnings

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






class   UpperCAmelCase      (			UpperCamelCase__							):
 __lowercase  						=							["""image_processor""", """tokenizer"""]
 __lowercase  						=							"""CLIPImageProcessor"""
 __lowercase  						=							("""CLIPTokenizer""", """CLIPTokenizerFast""")
 def __init__(							self :List[Any]		,			lowercase_ :List[str]=None		,			lowercase_ :Any=None		,			**lowercase_ :Union[str, Any]				)->				Optional[Any]:
     A__							=     None
     if "feature_extractor" in kwargs:
         warnings.warn(
             "The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"
             " instead."		,			lowercase_		,			)
         A__							=     kwargs.pop("feature_extractor"				)

     A__							=     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__(lowercase_		,			lowercase_				)
 def __call__(							self :Optional[Any]		,			lowercase_ :Optional[int]=None		,			lowercase_ :Dict=None		,			lowercase_ :List[Any]=None		,			**lowercase_ :Optional[Any]				)->				Union[str, Any]:

     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:
         A__							=     self.tokenizer(lowercase_		,			return_tensors=lowercase_		,			**lowercase_				)

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

     if text is not None and images is not None:
         A__							=     image_features.pixel_values
         return encoding
     elif text is not None:
         return encoding
     else:
         return BatchEncoding(data=dict(**lowercase_				)		,			tensor_type=lowercase_				)
 def 	UpperCAmelCase_				(							self :int		,			*lowercase_ :Tuple		,			**lowercase_ :Any				)->				Dict:
     return self.tokenizer.batch_decode(*lowercase_		,			**lowercase_				)
 def 	UpperCAmelCase_				(							self :Any		,			*lowercase_ :Any		,			**lowercase_ :Any				)->				List[str]:
     return self.tokenizer.decode(*lowercase_		,			**lowercase_				)
 @property
 def 	UpperCAmelCase_				(							self :Union[str, Any]				)->				Optional[int]:
     A__							=     self.tokenizer.model_input_names
     A__							=     self.image_processor.model_input_names
     return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names				)				)
 @property
 def 	UpperCAmelCase_				(							self :List[Any]				)->				str:
     warnings.warn(
         "`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."		,			lowercase_		,			)
     return self.image_processor_class



 @property
 def 	UpperCAmelCase_				(							self :List[str]				)->				List[Any]:
     warnings.warn(
         "`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."		,			lowercase_		,			)
     return self.image_processor



