






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   (   _A						,				unittest.TestCase						):
      _a							:						Any										=							KandinskyImgaImgPipeline
      _a							:						List[Any]										=							['prompt', 'image_embeds', 'negative_image_embeds', 'image']
      _a							:						int										=							[
          'prompt',
          'negative_prompt',
          'image_embeds',
          'negative_image_embeds',
          'image',
      ]
      _a							:						int										=							[
          'generator',
          'height',
          'width',
          'strength',
          'guidance_scale',
          'negative_prompt',
          'num_inference_steps',
          'return_dict',
          'guidance_scale',
          'num_images_per_prompt',
          'output_type',
          'return_dict',
      ]
      _a							:						List[str]										=							False
      @property
      def 						lowercase__					(     self   ):
       return 3_2
      @property
      def 						lowercase__					(     self   ):
       return 3_2
      @property
      def 						lowercase__					(     self   ):
       return self.time_input_dim
      @property
      def 						lowercase__					(     self   ):
       return self.time_input_dim * 4
      @property
      def 						lowercase__					(     self   ):
       return 1_0_0
      @property
      def 						lowercase__					(     self   ):
       snake_case__  :    Dict							       =XLMRobertaTokenizerFast.from_pretrained("""YiYiXu/tiny-random-mclip-base"""   )
       return tokenizer
      @property
      def 						lowercase__					(     self   ):
       torch.manual_seed(0   )
       snake_case__  :    Optional[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				,			)

       snake_case__  :    Optional[int]							       =MultilingualCLIP(a   )
       snake_case__  :    Union[str, Any]							       =text_encoder.eval()

       return text_encoder
      @property
      def 						lowercase__					(     self   ):
       torch.manual_seed(0   )

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

       snake_case__  :    Union[str, Any]							       ={
           """num_train_timesteps""": 1_0_0_0,
           """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,
       }

       snake_case__  :    Union[str, Any]							       =DDIMScheduler(**a   )

       snake_case__  :    Any							       ={
           """text_encoder""": text_encoder,
           """tokenizer""": tokenizer,
           """unet""": unet,
           """scheduler""": scheduler,
           """movq""": movq,
       }

       return components
      def 						lowercase__					(     self				,			a				,			a=0   ):
       snake_case__  :    str							       =floats_tensor((1, self.cross_attention_dim)				,			rng=random.Random(a   )   ).to(a   )
       snake_case__  :    Tuple							       =floats_tensor((1, self.cross_attention_dim)				,			rng=random.Random(seed + 1   )   ).to(a   )
       # create init_image
       snake_case__  :    str							       =floats_tensor((1, 3, 6_4, 6_4)				,			rng=random.Random(a   )   ).to(a   )
       snake_case__  :    Tuple							       =image.cpu().permute(0				,			2				,			3				,			1   )[0]
       snake_case__  :    Dict							       =Image.fromarray(np.uinta(a   )   ).convert("""RGB"""   ).resize((2_5_6, 2_5_6)   )

       if str(a   ).startswith("""mps"""   ):
        snake_case__  :    Dict							       =torch.manual_seed(a   )
       else:
        snake_case__  :    Optional[Any]							       =torch.Generator(device=a   ).manual_seed(a   )
       snake_case__  :    str							       ={
           """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 						lowercase__					(     self   ):
       snake_case__  :    List[str]							       ="""cpu"""

       snake_case__  :    Tuple							       =self.get_dummy_components()

       snake_case__  :    Optional[int]							       =self.pipeline_class(**a   )
       snake_case__  :    Union[str, Any]							       =pipe.to(a   )

       pipe.set_progress_bar_config(disable=a   )

       snake_case__  :    List[Any]							       =pipe(**self.get_dummy_inputs(a   )   )
       snake_case__  :    List[str]							       =output.images

       snake_case__  :    int							       =pipe(
           **self.get_dummy_inputs(a   )				,			return_dict=a				,			)[0]

       snake_case__  :    str							       =image[0, -3:, -3:, -1]
       snake_case__  :    Tuple							       =image_from_tuple[0, -3:, -3:, -1]

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

       snake_case__  :    List[str]							       =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  _lowercase   (   unittest.TestCase						):
      def 						lowercase__					(     self   ):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()



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

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

       snake_case__  :    Dict							       =KandinskyPriorPipeline.from_pretrained(
           """kandinsky-community/kandinsky-2-1-prior"""				,			torch_dtype=torch.floataa   )
       pipe_prior.to(a   )

       snake_case__  :    Tuple							       =KandinskyImgaImgPipeline.from_pretrained(
           """kandinsky-community/kandinsky-2-1"""				,			torch_dtype=torch.floataa   )
       snake_case__  :    str							       =pipeline.to(a   )

       pipeline.set_progress_bar_config(disable=a   )

       snake_case__  :    Optional[Any]							       =torch.Generator(device="""cpu"""   ).manual_seed(0   )
       snake_case__       ,   snake_case__  :    Any							       =pipe_prior(
           a				,			generator=a				,			num_inference_steps=5				,			negative_prompt=""""""				,			).to_tuple()

       snake_case__  :    str							       =pipeline(
           a				,			image=a				,			image_embeds=a				,			negative_image_embeds=a				,			generator=a				,			num_inference_steps=1_0_0				,			height=7_6_8				,			width=7_6_8				,			strength=0.2				,			output_type="""np"""				,			)

       snake_case__  :    Optional[int]							       =output.images[0]

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

       assert_mean_pixel_difference(a				,			a   )













from collections import OrderedDict
from typing import Mapping

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


__lowerCamelCase	:		Dict										=   logging.get_logger(__name__)

__lowerCamelCase	:		Union[str, Any]										=   {
    """xlm-mlm-en-2048""": """https://huggingface.co/xlm-mlm-en-2048/resolve/main/config.json""",
    """xlm-mlm-ende-1024""": """https://huggingface.co/xlm-mlm-ende-1024/resolve/main/config.json""",
    """xlm-mlm-enfr-1024""": """https://huggingface.co/xlm-mlm-enfr-1024/resolve/main/config.json""",
    """xlm-mlm-enro-1024""": """https://huggingface.co/xlm-mlm-enro-1024/resolve/main/config.json""",
    """xlm-mlm-tlm-xnli15-1024""": """https://huggingface.co/xlm-mlm-tlm-xnli15-1024/resolve/main/config.json""",
    """xlm-mlm-xnli15-1024""": """https://huggingface.co/xlm-mlm-xnli15-1024/resolve/main/config.json""",
    """xlm-clm-enfr-1024""": """https://huggingface.co/xlm-clm-enfr-1024/resolve/main/config.json""",
    """xlm-clm-ende-1024""": """https://huggingface.co/xlm-clm-ende-1024/resolve/main/config.json""",
    """xlm-mlm-17-1280""": """https://huggingface.co/xlm-mlm-17-1280/resolve/main/config.json""",
    """xlm-mlm-100-1280""": """https://huggingface.co/xlm-mlm-100-1280/resolve/main/config.json""",
}







class  _lowercase   (   _A						):
      _a							:						int										=							'xlm'
      _a							:						List[Any]										=							{
          'hidden_size': 'emb_dim',
          'num_attention_heads': 'n_heads',
          'num_hidden_layers': 'n_layers',
          'n_words': 'vocab_size',  # For backward compatibility
      }
      def __init__(     self				,			a=3_0_1_4_5				,			a=2_0_4_8				,			a=1_2				,			a=1_6				,			a=0.1				,			a=0.1				,			a=True				,			a=False				,			a=False				,			a=False				,			a=1				,			a=True				,			a=5_1_2				,			a=2_0_4_8**-0.5				,			a=1e-12				,			a=0.02				,			a=0				,			a=1				,			a=2				,			a=3				,			a=5				,			a=True				,			a="first"				,			a=True				,			a=None				,			a=True				,			a=0.1				,			a=5				,			a=5				,			a=0				,			a=0				,			a=2				,			a=0				,			**a				,			):
       snake_case__  :    str							       =vocab_size
       snake_case__  :    List[Any]							       =emb_dim
       snake_case__  :    List[str]							       =n_layers
       snake_case__  :    Any							       =n_heads
       snake_case__  :    Any							       =dropout
       snake_case__  :    Optional[int]							       =attention_dropout
       snake_case__  :    Dict							       =gelu_activation
       snake_case__  :    Tuple							       =sinusoidal_embeddings
       snake_case__  :    Optional[int]							       =causal
       snake_case__  :    Optional[int]							       =asm
       snake_case__  :    Optional[int]							       =n_langs
       snake_case__  :    Union[str, Any]							       =use_lang_emb
       snake_case__  :    Optional[Any]							       =layer_norm_eps
       snake_case__  :    Union[str, Any]							       =bos_index
       snake_case__  :    List[Any]							       =eos_index
       snake_case__  :    Optional[Any]							       =pad_index
       snake_case__  :    Union[str, Any]							       =unk_index
       snake_case__  :    Dict							       =mask_index
       snake_case__  :    str							       =is_encoder
       snake_case__  :    Optional[Any]							       =max_position_embeddings
       snake_case__  :    Any							       =embed_init_std
       snake_case__  :    Optional[Any]							       =init_std
       snake_case__  :    str							       =summary_type
       snake_case__  :    List[Any]							       =summary_use_proj
       snake_case__  :    str							       =summary_activation
       snake_case__  :    Optional[Any]							       =summary_proj_to_labels
       snake_case__  :    Optional[int]							       =summary_first_dropout
       snake_case__  :    Optional[Any]							       =start_n_top
       snake_case__  :    List[Any]							       =end_n_top
       snake_case__  :    Union[str, Any]							       =mask_token_id
       snake_case__  :    Any							       =lang_id

       if "n_words" in kwargs:
        snake_case__  :    int							       =kwargs["""n_words"""]

       super().__init__(pad_token_id=a				,			bos_token_id=a				,			**a   )







class  _lowercase   (   _A						):
      @property
      def 						lowercase__					(     self   ):
       if self.task == "multiple-choice":
        snake_case__  :    Any							       ={0: """batch""", 1: """choice""", 2: """sequence"""}
       else:
        snake_case__  :    Union[str, Any]							       ={0: """batch""", 1: """sequence"""}
       return OrderedDict(
           [
               ("""input_ids""", dynamic_axis),
               ("""attention_mask""", dynamic_axis),
               ("""token_type_ids""", dynamic_axis),
           ]   )





