





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





					"""simple docstring"""


					A__				 =    KandinskyImgaImgPipeline
					A__				 =    ["prompt", "image_embeds", "negative_image_embeds", "image"]
					A__				 =    [
					    "prompt",
					    "negative_prompt",
					    "image_embeds",
					    "negative_image_embeds",
					    "image",
					]
					A__				 =    [
					    "generator",
					    "height",
					    "width",
					    "strength",
					    "guidance_scale",
					    "negative_prompt",
					    "num_inference_steps",
					    "return_dict",
					    "guidance_scale",
					    "num_images_per_prompt",
					    "output_type",
					    "return_dict",
					]
					A__				 =    False





					@property
					def 					__magic_name__ ( self	):
										return 32





					@property
					def 					__magic_name__ ( self	):
										return 32





					@property
					def 					__magic_name__ ( self	):
										return self.time_input_dim





					@property
					def 					__magic_name__ ( self	):
										return self.time_input_dim * 4





					@property
					def 					__magic_name__ ( self	):
										return 100





					@property
					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       XLMRobertaTokenizerFast.from_pretrained("YiYiXu/tiny-random-mclip-base"	)
										return tokenizer





					@property
					def 					__magic_name__ ( self	):
										torch.manual_seed(0	)
										lowerCamelCase__       =       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	,  )

										lowerCamelCase__       =       MultilingualCLIP(_lowerCAmelCase	)
										lowerCamelCase__       =       text_encoder.eval()

										return text_encoder





					@property
					def 					__magic_name__ ( self	):
										torch.manual_seed(0	)

										lowerCamelCase__       =       {
										    "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,
										}

										lowerCamelCase__       =       UNetaDConditionModel(**_lowerCAmelCase	)
										return model





					@property
					def 					__magic_name__ ( self	):
										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 					__magic_name__ ( self	):
										torch.manual_seed(0	)
										lowerCamelCase__       =       VQModel(**self.dummy_movq_kwargs	)
										return model





					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       self.dummy_text_encoder
										lowerCamelCase__       =       self.dummy_tokenizer
										lowerCamelCase__       =       self.dummy_unet
										lowerCamelCase__       =       self.dummy_movq

										lowerCamelCase__       =       {
										    "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,
										}

										lowerCamelCase__       =       DDIMScheduler(**_lowerCAmelCase	)

										lowerCamelCase__       =       {
										    "text_encoder": text_encoder,
										    "tokenizer": tokenizer,
										    "unet": unet,
										    "scheduler": scheduler,
										    "movq": movq,
										}

										return components





					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase=0	):
										lowerCamelCase__       =       floats_tensor((1, self.cross_attention_dim)	,  rng=random.Random(_lowerCAmelCase	)	).to(_lowerCAmelCase	)
										lowerCamelCase__       =       floats_tensor((1, self.cross_attention_dim)	,  rng=random.Random(seed + 1	)	).to(_lowerCAmelCase	)
										# create init_image
										lowerCamelCase__       =       floats_tensor((1, 3, 64, 64)	,  rng=random.Random(_lowerCAmelCase	)	).to(_lowerCAmelCase	)
										lowerCamelCase__       =       image.cpu().permute(0	,  2	,  3	,  1	)[0]
										lowerCamelCase__       =       Image.fromarray(np.uinta(_lowerCAmelCase	)	).convert("RGB"	).resize((256, 256)	)

										if str(_lowerCAmelCase	).startswith("mps"	):
															lowerCamelCase__       =       torch.manual_seed(_lowerCAmelCase	)
										else:
															lowerCamelCase__       =       torch.Generator(device=_lowerCAmelCase	).manual_seed(_lowerCAmelCase	)
										lowerCamelCase__       =       {
										    "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 					__magic_name__ ( self	):
										lowerCamelCase__       =       "cpu"

										lowerCamelCase__       =       self.get_dummy_components()

										lowerCamelCase__       =       self.pipeline_class(**_lowerCAmelCase	)
										lowerCamelCase__       =       pipe.to(_lowerCAmelCase	)

										pipe.set_progress_bar_config(disable=_lowerCAmelCase	)

										lowerCamelCase__       =       pipe(**self.get_dummy_inputs(_lowerCAmelCase	)	)
										lowerCamelCase__       =       output.images

										lowerCamelCase__       =       pipe(
										    **self.get_dummy_inputs(_lowerCAmelCase	)	,  return_dict=_lowerCAmelCase	,  )[0]

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

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

										lowerCamelCase__       =       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 SCREAMING_SNAKE_CASE_ (					unittest.TestCase   ):





					"""simple docstring"""


					def 					__magic_name__ ( self	):
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()






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

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

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

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

										pipeline.set_progress_bar_config(disable=_lowerCAmelCase	)

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

										lowerCamelCase__       =       pipeline(
										    _lowerCAmelCase	,  image=_lowerCAmelCase	,  image_embeds=_lowerCAmelCase	,  negative_image_embeds=_lowerCAmelCase	,  generator=_lowerCAmelCase	,  num_inference_steps=100	,  height=768	,  width=768	,  strength=0.2	,  output_type="np"	,  )

										lowerCamelCase__       =       output.images[0]

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

										assert_mean_pixel_difference(_lowerCAmelCase	,  _lowerCAmelCase	)








from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig


A_												=     {
    "albert-base-v1": "https://huggingface.co/albert-base-v1/resolve/main/config.json",
    "albert-large-v1": "https://huggingface.co/albert-large-v1/resolve/main/config.json",
    "albert-xlarge-v1": "https://huggingface.co/albert-xlarge-v1/resolve/main/config.json",
    "albert-xxlarge-v1": "https://huggingface.co/albert-xxlarge-v1/resolve/main/config.json",
    "albert-base-v2": "https://huggingface.co/albert-base-v2/resolve/main/config.json",
    "albert-large-v2": "https://huggingface.co/albert-large-v2/resolve/main/config.json",
    "albert-xlarge-v2": "https://huggingface.co/albert-xlarge-v2/resolve/main/config.json",
    "albert-xxlarge-v2": "https://huggingface.co/albert-xxlarge-v2/resolve/main/config.json",
}

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					A__				 =    "albert"





					def __init__( self	,  _lowerCAmelCase=3_0000	,  _lowerCAmelCase=128	,  _lowerCAmelCase=4096	,  _lowerCAmelCase=12	,  _lowerCAmelCase=1	,  _lowerCAmelCase=64	,  _lowerCAmelCase=1_6384	,  _lowerCAmelCase=1	,  _lowerCAmelCase="gelu_new"	,  _lowerCAmelCase=0	,  _lowerCAmelCase=0	,  _lowerCAmelCase=512	,  _lowerCAmelCase=2	,  _lowerCAmelCase=0.02	,  _lowerCAmelCase=1E-12	,  _lowerCAmelCase=0.1	,  _lowerCAmelCase="absolute"	,  _lowerCAmelCase=0	,  _lowerCAmelCase=2	,  _lowerCAmelCase=3	,  **_lowerCAmelCase	,  ):
										super().__init__(pad_token_id=_lowerCAmelCase	,  bos_token_id=_lowerCAmelCase	,  eos_token_id=_lowerCAmelCase	,  **_lowerCAmelCase	)

										lowerCamelCase__       =       vocab_size
										lowerCamelCase__       =       embedding_size
										lowerCamelCase__       =       hidden_size
										lowerCamelCase__       =       num_hidden_layers
										lowerCamelCase__       =       num_hidden_groups
										lowerCamelCase__       =       num_attention_heads
										lowerCamelCase__       =       inner_group_num
										lowerCamelCase__       =       hidden_act
										lowerCamelCase__       =       intermediate_size
										lowerCamelCase__       =       hidden_dropout_prob
										lowerCamelCase__       =       attention_probs_dropout_prob
										lowerCamelCase__       =       max_position_embeddings
										lowerCamelCase__       =       type_vocab_size
										lowerCamelCase__       =       initializer_range
										lowerCamelCase__       =       layer_norm_eps
										lowerCamelCase__       =       classifier_dropout_prob
										lowerCamelCase__       =       position_embedding_type

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					@property
					def 					__magic_name__ ( self	):
										if self.task == "multiple-choice":
															lowerCamelCase__       =       {0: "batch", 1: "choice", 2: "sequence"}
										else:
															lowerCamelCase__       =       {0: "batch", 1: "sequence"}
										return OrderedDict(
										    [
										        ("input_ids", dynamic_axis),
										        ("attention_mask", dynamic_axis),
										        ("token_type_ids", dynamic_axis),
										    ]	)

