



import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, 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__       =					KandinskyPipeline
 lowerCAmelCase__       =					[
     """prompt""",
     """image_embeds""",
     """negative_image_embeds""",
 ]
 lowerCAmelCase__       =					["""prompt""", """negative_prompt""", """image_embeds""", """negative_image_embeds"""]
 lowerCAmelCase__       =					[
     """generator""",
     """height""",
     """width""",
     """latents""",
     """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  							=      DDIMScheduler(
        num_train_timesteps=1000		,			beta_schedule='''linear'''		,			beta_start=0.00_085		,			beta_end=0.012		,			clip_sample=__UpperCAmelCase		,			set_alpha_to_one=__UpperCAmelCase		,			steps_offset=1		,			prediction_type='''epsilon'''		,			thresholding=__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			)
    if str(__UpperCAmelCase			).startswith('''mps'''			):
       __lowerCamelCase  							=      torch.manual_seed(__UpperCAmelCase			)
    else:
       __lowerCamelCase  							=      torch.Generator(device=__UpperCAmelCase			).manual_seed(__UpperCAmelCase			)
    __lowerCamelCase  							=      {
        '''prompt''': '''horse''',
        '''image_embeds''': image_embeds,
        '''negative_image_embeds''': negative_image_embeds,
        '''generator''': generator,
        '''height''': 64,
        '''width''': 64,
        '''guidance_scale''': 4.0,
        '''num_inference_steps''': 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.328_663, 1.0, 0.23_216_873, 1.0, 0.92_717_564, 0.4_639_046, 0.96_894_777, 0.31_713_378, 0.6_293_953]			)

    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_text2img_cat_fp16.npy'''			)

    __lowerCamelCase  							=      KandinskyPriorPipeline.from_pretrained(
        '''kandinsky-community/kandinsky-2-1-prior'''		,			torch_dtype=torch.floataa			)
    pipe_prior.to(__UpperCAmelCase			)

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

    __lowerCamelCase  							=      '''red cat, 4k photo'''

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

    __lowerCamelCase  							=      torch.Generator(device='''cuda'''			).manual_seed(0			)
    __lowerCamelCase  							=      pipeline(
        __UpperCAmelCase		,			image_embeds=__UpperCAmelCase		,			negative_image_embeds=__UpperCAmelCase		,			generator=__UpperCAmelCase		,			num_inference_steps=100		,			output_type='''np'''		,			)

    __lowerCamelCase  							=      output.images[0]

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

    assert_mean_pixel_difference(__UpperCAmelCase		,			__UpperCAmelCase			)











from typing import Optional, Tuple

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_openai import OpenAIGPTTokenizer


a_								=   logging.get_logger(__name__)

a_								=   {"""vocab_file""": """vocab.json""", """merges_file""": """merges.txt""", """tokenizer_file""": """tokenizer.json"""}

a_								=   {
    """vocab_file""": {"""openai-gpt""": """https://huggingface.co/openai-gpt/resolve/main/vocab.json"""},
    """merges_file""": {"""openai-gpt""": """https://huggingface.co/openai-gpt/resolve/main/merges.txt"""},
    """tokenizer_file""": {"""openai-gpt""": """https://huggingface.co/openai-gpt/resolve/main/tokenizer.json"""},
}

a_								=   {
    """openai-gpt""": 512,
}



class 	__lowerCAmelCase       (		lowerCAmelCase__				):
 lowerCAmelCase__       =					VOCAB_FILES_NAMES
 lowerCAmelCase__       =					PRETRAINED_VOCAB_FILES_MAP
 lowerCAmelCase__       =					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 lowerCAmelCase__       =					["""input_ids""", """attention_mask"""]
 lowerCAmelCase__       =					OpenAIGPTTokenizer




 def __init__(  self		,			__UpperCAmelCase=None		,			__UpperCAmelCase=None		,			__UpperCAmelCase=None		,			__UpperCAmelCase="<unk>"		,			**__UpperCAmelCase			):



    '''simple docstring'''
    super().__init__(__UpperCAmelCase		,			__UpperCAmelCase		,			tokenizer_file=__UpperCAmelCase		,			unk_token=__UpperCAmelCase		,			**__UpperCAmelCase			)




 @property
 def     lowerCamelCase (  self			):



    '''simple docstring'''
    return True




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



    '''simple docstring'''
    __lowerCamelCase  							=      self._tokenizer.model.save(__UpperCAmelCase		,			name=__UpperCAmelCase			)
    return tuple(__UpperCAmelCase			)






