



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 			__lowerCAmelCase  (						lowerCAmelCase__	,						unittest.TestCase ):
 lowerCAmelCase__							       =	KandinskyImgaImgPipeline
 lowerCAmelCase__							       =	["""prompt""", """image_embeds""", """negative_image_embeds""", """image"""]
 lowerCAmelCase__							       =	[
     """prompt""",
     """negative_prompt""",
     """image_embeds""",
     """negative_image_embeds""",
     """image""",
 ]
 lowerCAmelCase__							       =	[
     """generator""",
     """height""",
     """width""",
     """strength""",
     """guidance_scale""",
     """negative_prompt""",
     """num_inference_steps""",
     """return_dict""",
     """guidance_scale""",
     """num_images_per_prompt""",
     """output_type""",
     """return_dict""",
 ]
 lowerCAmelCase__							       =	False




 @property
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    return 32




 @property
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    return 32




 @property
 def 							lowerCamelCase   (    self					):



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




 @property
 def 							lowerCamelCase   (    self					):



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




 @property
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    return 100




 @property
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     XLMRobertaTokenizerFast.from_pretrained('''YiYiXu/tiny-random-mclip-base'''					)
    return tokenizer




 @property
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    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(__UpperCAmelCase					)
    __lowerCamelCase    				=     text_encoder.eval()

    return text_encoder




 @property
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    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(**__UpperCAmelCase					)
    return model




 @property
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    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 							lowerCamelCase   (    self					):



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




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __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(**__UpperCAmelCase					)

    __lowerCamelCase    				=     {
        '''text_encoder''': text_encoder,
        '''tokenizer''': tokenizer,
        '''unet''': unet,
        '''scheduler''': scheduler,
        '''movq''': movq,
    }

    return components




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase=0					):



    '''simple docstring'''
    __lowerCamelCase    				=     floats_tensor((1, self.cross_attention_dim)						,					rng=random.Random(__UpperCAmelCase					)					).to(__UpperCAmelCase					)
    __lowerCamelCase    				=     floats_tensor((1, self.cross_attention_dim)						,					rng=random.Random(seed + 1					)					).to(__UpperCAmelCase					)
    # create init_image
    __lowerCamelCase    				=     floats_tensor((1, 3, 64, 64)						,					rng=random.Random(__UpperCAmelCase					)					).to(__UpperCAmelCase					)
    __lowerCamelCase    				=     image.cpu().permute(0						,					2						,					3						,					1					)[0]
    __lowerCamelCase    				=     Image.fromarray(np.uinta(__UpperCAmelCase					)					).convert('''RGB'''					).resize((256, 256)					)

    if str(__UpperCAmelCase					).startswith('''mps'''					):
       __lowerCamelCase    				=     torch.manual_seed(__UpperCAmelCase					)
    else:
       __lowerCamelCase    				=     torch.Generator(device=__UpperCAmelCase					).manual_seed(__UpperCAmelCase					)
    __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 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     '''cpu'''

    __lowerCamelCase    				=     self.get_dummy_components()

    __lowerCamelCase    				=     self.pipeline_class(**__UpperCAmelCase					)
    __lowerCamelCase    				=     pipe.to(__UpperCAmelCase					)

    pipe.set_progress_bar_config(disable=__UpperCAmelCase					)

    __lowerCamelCase    				=     pipe(**self.get_dummy_inputs(__UpperCAmelCase					)					)
    __lowerCamelCase    				=     output.images

    __lowerCamelCase    				=     pipe(
        **self.get_dummy_inputs(__UpperCAmelCase					)						,					return_dict=__UpperCAmelCase						,					)[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 			__lowerCAmelCase  (						unittest.TestCase ):




 def 							lowerCamelCase   (    self					):



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




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __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(__UpperCAmelCase					)

    __lowerCamelCase    				=     KandinskyImgaImgPipeline.from_pretrained(
        '''kandinsky-community/kandinsky-2-1'''						,					torch_dtype=torch.floataa					)
    __lowerCamelCase    				=     pipeline.to(__UpperCAmelCase					)

    pipeline.set_progress_bar_config(disable=__UpperCAmelCase					)

    __lowerCamelCase    				=     torch.Generator(device='''cpu'''					).manual_seed(0					)
    __lowerCamelCase				,__lowerCamelCase    				=     pipe_prior(
        __UpperCAmelCase						,					generator=__UpperCAmelCase						,					num_inference_steps=5						,					negative_prompt=''''''						,					).to_tuple()

    __lowerCamelCase    				=     pipeline(
        __UpperCAmelCase						,					image=__UpperCAmelCase						,					image_embeds=__UpperCAmelCase						,					negative_image_embeds=__UpperCAmelCase						,					generator=__UpperCAmelCase						,					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(__UpperCAmelCase						,					__UpperCAmelCase					)











from typing import Mapping

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


a_								=       logging.get_logger(__name__)

a_								=       {
    """t5-small""": """https://huggingface.co/t5-small/resolve/main/config.json""",
    """t5-base""": """https://huggingface.co/t5-base/resolve/main/config.json""",
    """t5-large""": """https://huggingface.co/t5-large/resolve/main/config.json""",
    """t5-3b""": """https://huggingface.co/t5-3b/resolve/main/config.json""",
    """t5-11b""": """https://huggingface.co/t5-11b/resolve/main/config.json""",
}



class 			__lowerCAmelCase  (						lowerCAmelCase__ ):
 lowerCAmelCase__							       =	"""t5"""
 lowerCAmelCase__							       =	["""past_key_values"""]
 lowerCAmelCase__							       =	{"""hidden_size""": """d_model""", """num_attention_heads""": """num_heads""", """num_hidden_layers""": """num_layers"""}




 def __init__(    self						,					__UpperCAmelCase=32128						,					__UpperCAmelCase=512						,					__UpperCAmelCase=64						,					__UpperCAmelCase=2048						,					__UpperCAmelCase=6						,					__UpperCAmelCase=None						,					__UpperCAmelCase=8						,					__UpperCAmelCase=32						,					__UpperCAmelCase=128						,					__UpperCAmelCase=0.1						,					__UpperCAmelCase=1E-6						,					__UpperCAmelCase=1.0						,					__UpperCAmelCase="relu"						,					__UpperCAmelCase=True						,					__UpperCAmelCase=True						,					__UpperCAmelCase=0						,					__UpperCAmelCase=1						,					**__UpperCAmelCase						,					):



    '''simple docstring'''
    __lowerCamelCase    				=     vocab_size
    __lowerCamelCase    				=     d_model
    __lowerCamelCase    				=     d_kv
    __lowerCamelCase    				=     d_ff
    __lowerCamelCase    				=     num_layers
    __lowerCamelCase    				=     (
        num_decoder_layers if num_decoder_layers is not None else self.num_layers
    )  # default = symmetry
    __lowerCamelCase    				=     num_heads
    __lowerCamelCase    				=     relative_attention_num_buckets
    __lowerCamelCase    				=     relative_attention_max_distance
    __lowerCamelCase    				=     dropout_rate
    __lowerCamelCase    				=     layer_norm_epsilon
    __lowerCamelCase    				=     initializer_factor
    __lowerCamelCase    				=     feed_forward_proj
    __lowerCamelCase    				=     use_cache

    __lowerCamelCase    				=     self.feed_forward_proj.split('''-'''					)
    __lowerCamelCase    				=     act_info[-1]
    __lowerCamelCase    				=     act_info[0] == '''gated'''

    if len(__UpperCAmelCase					) > 1 and act_info[0] != "gated" or len(__UpperCAmelCase					) > 2:
       raise ValueError(
           F"""`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."""
           '''Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. '''
           '''\'gated-gelu\' or \'relu\''''					)

    # for backwards compatibility
    if feed_forward_proj == "gated-gelu":
       __lowerCamelCase    				=     '''gelu_new'''

    super().__init__(
        pad_token_id=__UpperCAmelCase						,					eos_token_id=__UpperCAmelCase						,					is_encoder_decoder=__UpperCAmelCase						,					**__UpperCAmelCase						,					)



class 			__lowerCAmelCase  (						lowerCAmelCase__ ):




 @property
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     {
        '''input_ids''': {0: '''batch''', 1: '''encoder_sequence'''},
        '''attention_mask''': {0: '''batch''', 1: '''encoder_sequence'''},
    }
    if self.use_past:
       __lowerCamelCase    				=     '''past_encoder_sequence + sequence'''
       __lowerCamelCase    				=     {0: '''batch'''}
       __lowerCamelCase    				=     {0: '''batch''', 1: '''past_decoder_sequence + sequence'''}
    else:
       __lowerCamelCase    				=     {0: '''batch''', 1: '''decoder_sequence'''}
       __lowerCamelCase    				=     {0: '''batch''', 1: '''decoder_sequence'''}

    if self.use_past:
       self.fill_with_past_key_values_(__UpperCAmelCase						,					direction='''inputs'''					)

    return common_inputs




 @property
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    return 13






