


'''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 A__       (						UpperCAmelCase__ ,     unittest.TestCase      ):
      __UpperCamelCase       :  List[str]	       =							KandinskyImgaImgPipeline
      __UpperCamelCase       :  Optional[int]	       =							["prompt", "image_embeds", "negative_image_embeds", "image"]
      __UpperCamelCase       :  Any	       =							[
          "prompt",
          "negative_prompt",
          "image_embeds",
          "negative_image_embeds",
          "image",
      ]
      __UpperCamelCase       :  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",
      ]
      __UpperCamelCase       :  Any	       =							False






      @property
      def __UpperCAmelCase       ( self			:Dict	)       ->					Optional[Any]:




             '''simple docstring'''
             return 3_2






      @property
      def __UpperCAmelCase       ( self			:List[str]	)       ->					Optional[int]:




             '''simple docstring'''
             return 3_2






      @property
      def __UpperCAmelCase       ( self			:Dict	)       ->					Dict:




             '''simple docstring'''
             return self.time_input_dim






      @property
      def __UpperCAmelCase       ( self			:Optional[Any]	)       ->					Any:




             '''simple docstring'''
             return self.time_input_dim * 4






      @property
      def __UpperCAmelCase       ( self			:Union[str, Any]	)       ->					List[str]:




             '''simple docstring'''
             return 1_0_0






      @property
      def __UpperCAmelCase       ( self			:str	)       ->					Any:




             '''simple docstring'''
             _a    : Tuple     					=XLMRobertaTokenizerFast.from_pretrained("""YiYiXu/tiny-random-mclip-base"""	)
             return tokenizer






      @property
      def __UpperCAmelCase       ( self			:Union[str, Any]	)       ->					Optional[int]:




             '''simple docstring'''
             torch.manual_seed(0	)
             _a    : List[Any]     					=MCLIPConfig(
                 numDims=self.cross_attention_dim     ,						transformerDimensions=self.text_embedder_hidden_size     ,						hidden_size=self.text_embedder_hidden_size     ,						intermediate_size=3_7     ,						num_attention_heads=4     ,						num_hidden_layers=5     ,						vocab_size=1_0_0_5     ,						)

             _a    : Dict     					=MultilingualCLIP(SCREAMING_SNAKE_CASE	)
             _a    : Optional[int]     					=text_encoder.eval()

             return text_encoder






      @property
      def __UpperCAmelCase       ( self			:Optional[int]	)       ->					Union[str, Any]:




             '''simple docstring'''
             torch.manual_seed(0	)

             _a    : Optional[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,
             }

             _a    : Optional[int]     					=UNetaDConditionModel(**SCREAMING_SNAKE_CASE	)
             return model






      @property
      def __UpperCAmelCase       ( self			:Optional[Any]	)       ->					Optional[int]:




             '''simple docstring'''
             return {
                 "block_out_channels": [3_2, 6_4],
                 "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": 1_2,
                 "out_channels": 3,
                 "up_block_types": [
                     "AttnUpDecoderBlock2D",
                     "UpDecoderBlock2D",
                 ],
                 "vq_embed_dim": 4,
             }






      @property
      def __UpperCAmelCase       ( self			:Tuple	)       ->					str:




             '''simple docstring'''
             torch.manual_seed(0	)
             _a    : int     					=VQModel(**self.dummy_movq_kwargs	)
             return model






      def __UpperCAmelCase       ( self			:str	)       ->					List[Any]:




             '''simple docstring'''
             _a    : Union[str, Any]     					=self.dummy_text_encoder
             _a    : List[Any]     					=self.dummy_tokenizer
             _a    : Optional[int]     					=self.dummy_unet
             _a    : str     					=self.dummy_movq

             _a    : int     					={
                 """num_train_timesteps""": 1_0_0_0,
                 """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,
             }

             _a    : Union[str, Any]     					=DDIMScheduler(**SCREAMING_SNAKE_CASE	)

             _a    : List[str]     					={
                 """text_encoder""": text_encoder,
                 """tokenizer""": tokenizer,
                 """unet""": unet,
                 """scheduler""": scheduler,
                 """movq""": movq,
             }

             return components






      def __UpperCAmelCase       ( self			:Dict     ,						SCREAMING_SNAKE_CASE			:Any     ,						SCREAMING_SNAKE_CASE			:List[Any]=0	)       ->					Tuple:




             '''simple docstring'''
             _a    : List[Any]     					=floats_tensor((1, self.cross_attention_dim)     ,						rng=random.Random(SCREAMING_SNAKE_CASE	)	).to(SCREAMING_SNAKE_CASE	)
             _a    : List[Any]     					=floats_tensor((1, self.cross_attention_dim)     ,						rng=random.Random(seed + 1	)	).to(SCREAMING_SNAKE_CASE	)
             # create init_image
             _a    : Union[str, Any]     					=floats_tensor((1, 3, 6_4, 6_4)     ,						rng=random.Random(SCREAMING_SNAKE_CASE	)	).to(SCREAMING_SNAKE_CASE	)
             _a    : Dict     					=image.cpu().permute(0     ,						2     ,						3     ,						1	)[0]
             _a    : Union[str, Any]     					=Image.fromarray(np.uinta(SCREAMING_SNAKE_CASE	)	).convert("""RGB"""	).resize((2_5_6, 2_5_6)	)

             if str(SCREAMING_SNAKE_CASE	).startswith("""mps"""	):
                    _a    : Optional[Any]     					=torch.manual_seed(SCREAMING_SNAKE_CASE	)
             else:
                    _a    : Tuple     					=torch.Generator(device=SCREAMING_SNAKE_CASE	).manual_seed(SCREAMING_SNAKE_CASE	)
             _a    : Dict     					={
                 """prompt""": """horse""",
                 """image""": init_image,
                 """image_embeds""": image_embeds,
                 """negative_image_embeds""": negative_image_embeds,
                 """generator""": generator,
                 """height""": 6_4,
                 """width""": 6_4,
                 """num_inference_steps""": 1_0,
                 """guidance_scale""": 7.0,
                 """strength""": 0.2,
                 """output_type""": """np""",
             }
             return inputs






      def __UpperCAmelCase       ( self			:Optional[int]	)       ->					Optional[Any]:




             '''simple docstring'''
             _a    : Any     					="""cpu"""

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

             _a    : Optional[Any]     					=self.pipeline_class(**SCREAMING_SNAKE_CASE	)
             _a    : int     					=pipe.to(SCREAMING_SNAKE_CASE	)

             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE	)

             _a    : str     					=pipe(**self.get_dummy_inputs(SCREAMING_SNAKE_CASE	)	)
             _a    : Optional[Any]     					=output.images

             _a    : Tuple     					=pipe(
                 **self.get_dummy_inputs(SCREAMING_SNAKE_CASE	)     ,						return_dict=SCREAMING_SNAKE_CASE     ,						)[0]

             _a    : str     					=image[0, -3:, -3:, -1]
             _a    : List[Any]     					=image_from_tuple[0, -3:, -3:, -1]

             assert image.shape == (1, 6_4, 6_4, 3)

             _a    : 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      ):






      def __UpperCAmelCase       ( self			:str	)       ->					Any:




             '''simple docstring'''
             # clean up the VRAM after each test
             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()






      def __UpperCAmelCase       ( self			:Tuple	)       ->					Optional[Any]:




             '''simple docstring'''
             _a    : Optional[Any]     					=load_numpy(
                 """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
                 """/kandinsky/kandinsky_img2img_frog.npy"""	)

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

             _a    : int     					=KandinskyPriorPipeline.from_pretrained(
                 """kandinsky-community/kandinsky-2-1-prior"""     ,						torch_dtype=torch.floataa	)
             pipe_prior.to(SCREAMING_SNAKE_CASE	)

             _a    : Optional[int]     					=KandinskyImgaImgPipeline.from_pretrained(
                 """kandinsky-community/kandinsky-2-1"""     ,						torch_dtype=torch.floataa	)
             _a    : int     					=pipeline.to(SCREAMING_SNAKE_CASE	)

             pipeline.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE	)

             _a    : Tuple     					=torch.Generator(device="""cpu"""	).manual_seed(0	)
             _a					,    _a    : Any     					=pipe_prior(
                 SCREAMING_SNAKE_CASE     ,						generator=SCREAMING_SNAKE_CASE     ,						num_inference_steps=5     ,						negative_prompt=""""""     ,						).to_tuple()

             _a    : int     					=pipeline(
                 SCREAMING_SNAKE_CASE     ,						image=SCREAMING_SNAKE_CASE     ,						image_embeds=SCREAMING_SNAKE_CASE     ,						negative_image_embeds=SCREAMING_SNAKE_CASE     ,						generator=SCREAMING_SNAKE_CASE     ,						num_inference_steps=1_0_0     ,						height=7_6_8     ,						width=7_6_8     ,						strength=0.2     ,						output_type="""np"""     ,						)

             _a    : Optional[Any]     					=output.images[0]

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

             assert_mean_pixel_difference(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)






'''simple docstring'''







from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


A__:    List[str]					    =       logging.get_logger(__name__)

A__:    Union[str, Any]					    =       {
    '''facebook/data2vec-text-base''': '''https://huggingface.co/data2vec/resolve/main/config.json''',
}







class A__       (						UpperCAmelCase__      ):
      __UpperCamelCase       :  int	       =							"data2vec-text"






      def __init__( self			:str     ,						SCREAMING_SNAKE_CASE			:Optional[Any]=3_0_5_2_2     ,						SCREAMING_SNAKE_CASE			:Any=7_6_8     ,						SCREAMING_SNAKE_CASE			:List[Any]=1_2     ,						SCREAMING_SNAKE_CASE			:List[str]=1_2     ,						SCREAMING_SNAKE_CASE			:Dict=3_0_7_2     ,						SCREAMING_SNAKE_CASE			:List[str]="gelu"     ,						SCREAMING_SNAKE_CASE			:Any=0.1     ,						SCREAMING_SNAKE_CASE			:List[str]=0.1     ,						SCREAMING_SNAKE_CASE			:int=5_1_2     ,						SCREAMING_SNAKE_CASE			:int=2     ,						SCREAMING_SNAKE_CASE			:Union[str, Any]=0.02     ,						SCREAMING_SNAKE_CASE			:Dict=1e-12     ,						SCREAMING_SNAKE_CASE			:int=1     ,						SCREAMING_SNAKE_CASE			:Dict=0     ,						SCREAMING_SNAKE_CASE			:List[Any]=2     ,						SCREAMING_SNAKE_CASE			:str="absolute"     ,						SCREAMING_SNAKE_CASE			:Tuple=True     ,						SCREAMING_SNAKE_CASE			:Union[str, Any]=None     ,						**SCREAMING_SNAKE_CASE			:Union[str, Any]     ,						)       ->					List[str]:




             '''simple docstring'''
             super().__init__(pad_token_id=SCREAMING_SNAKE_CASE     ,						bos_token_id=SCREAMING_SNAKE_CASE     ,						eos_token_id=SCREAMING_SNAKE_CASE     ,						**SCREAMING_SNAKE_CASE	)

             _a    : Optional[Any]     					=vocab_size
             _a    : Optional[Any]     					=hidden_size
             _a    : Any     					=num_hidden_layers
             _a    : List[str]     					=num_attention_heads
             _a    : Union[str, Any]     					=hidden_act
             _a    : Any     					=intermediate_size
             _a    : str     					=hidden_dropout_prob
             _a    : Optional[Any]     					=attention_probs_dropout_prob
             _a    : Optional[Any]     					=max_position_embeddings
             _a    : Union[str, Any]     					=type_vocab_size
             _a    : Tuple     					=initializer_range
             _a    : Optional[int]     					=layer_norm_eps
             _a    : Tuple     					=position_embedding_type
             _a    : int     					=use_cache
             _a    : List[str]     					=classifier_dropout







class A__       (						UpperCAmelCase__      ):






      @property
      def __UpperCAmelCase       ( self			:int	)       ->					Mapping[str, Mapping[int, str]]:




             '''simple docstring'''
             if self.task == "multiple-choice":
                    _a    : Tuple     					={0: """batch""", 1: """choice""", 2: """sequence"""}
             else:
                    _a    : List[Any]     					={0: """batch""", 1: """sequence"""}
             return OrderedDict(
                 [
                     ("""input_ids""", dynamic_axis),
                     ("""attention_mask""", dynamic_axis),
                 ]	)


