





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



 @property
 def 							UpperCamelCase     (				self					):
        return 32



 @property
 def 							UpperCamelCase     (				self					):
        return 32



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



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



 @property
 def 							UpperCamelCase     (				self					):
        return 100



 @property
 def 							UpperCamelCase     (				self					):
        A__		   =		XLMRobertaTokenizerFast.from_pretrained('''YiYiXu/tiny-random-mclip-base'''					)
        return tokenizer



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

        A__		   =		MultilingualCLIP(__lowerCamelCase					)
        A__		   =		text_encoder.eval()

        return text_encoder



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

        A__		   =		{
            '''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__		   =		UNetaDConditionModel(**__lowerCamelCase					)
        return model



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



 def 							UpperCamelCase     (				self					):
        A__		   =		self.dummy_text_encoder
        A__		   =		self.dummy_tokenizer
        A__		   =		self.dummy_unet
        A__		   =		self.dummy_movq

        A__		   =		{
            '''num_train_timesteps''': 1000,
            '''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,
        }

        A__		   =		DDIMScheduler(**__lowerCamelCase					)

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

        return components



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

        if str(__lowerCamelCase					).startswith('''mps'''					):
               A__		   =		torch.manual_seed(__lowerCamelCase					)
        else:
               A__		   =		torch.Generator(device=__lowerCamelCase					).manual_seed(__lowerCamelCase					)
        A__		   =		{
            '''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 							UpperCamelCase     (				self					):
        A__		   =		'''cpu'''

        A__		   =		self.get_dummy_components()

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

        pipe.set_progress_bar_config(disable=__lowerCamelCase					)

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

        A__		   =		pipe(
            **self.get_dummy_inputs(__lowerCamelCase					),return_dict=__lowerCamelCase,)[0]

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

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

        A__		   =		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        SCREAMING_SNAKE_CASE__					(		unittest.TestCase ):



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




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

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

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

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

        pipeline.set_progress_bar_config(disable=__lowerCamelCase					)

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

        A__		   =		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''',)

        A__		   =		output.images[0]

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

        assert_mean_pixel_difference(__lowerCamelCase,__lowerCamelCase					)












from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


a__: str							   =						{
    'configuration_mega': ['MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP', 'MegaConfig', 'MegaOnnxConfig'],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       a__: Union[str, Any]							   =						[
           'MEGA_PRETRAINED_MODEL_ARCHIVE_LIST',
           'MegaForCausalLM',
           'MegaForMaskedLM',
           'MegaForMultipleChoice',
           'MegaForQuestionAnswering',
           'MegaForSequenceClassification',
           'MegaForTokenClassification',
           'MegaModel',
           'MegaPreTrainedModel',
       ]

if TYPE_CHECKING:
       from .configuration_mega import MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP, MegaConfig, MegaOnnxConfig

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_mega import (
                  MEGA_PRETRAINED_MODEL_ARCHIVE_LIST,
                  MegaForCausalLM,
                  MegaForMaskedLM,
                  MegaForMultipleChoice,
                  MegaForQuestionAnswering,
                  MegaForSequenceClassification,
                  MegaForTokenClassification,
                  MegaModel,
                  MegaPreTrainedModel,
              )

else:
       import sys

       a__: Union[str, Any]							   =						_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)





