






import copy
import os
from typing import TYPE_CHECKING, List, Union


if TYPE_CHECKING:
       pass

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__lowerCamelCase	:		Dict										=   logging.get_logger(__name__)

__lowerCamelCase	:		List[Any]										=   {
    """kakaobrain/align-base""": """https://huggingface.co/kakaobrain/align-base/resolve/main/config.json""",
}







class  _lowercase   (   _A						):
      _a							:						Optional[int]										=							'align_text_model'
      def __init__(     self				,			a=3_0_5_2_2				,			a=7_6_8				,			a=1_2				,			a=1_2				,			a=3_0_7_2				,			a="gelu"				,			a=0.1				,			a=0.1				,			a=5_1_2				,			a=2				,			a=0.02				,			a=1e-12				,			a=0				,			a="absolute"				,			a=True				,			**a				,			):
       super().__init__(**a   )

       snake_case__  :    Optional[int]							       =vocab_size
       snake_case__  :    str							       =hidden_size
       snake_case__  :    Dict							       =num_hidden_layers
       snake_case__  :    Tuple							       =num_attention_heads
       snake_case__  :    Optional[Any]							       =hidden_act
       snake_case__  :    List[Any]							       =intermediate_size
       snake_case__  :    List[Any]							       =hidden_dropout_prob
       snake_case__  :    int							       =attention_probs_dropout_prob
       snake_case__  :    int							       =max_position_embeddings
       snake_case__  :    Optional[int]							       =type_vocab_size
       snake_case__  :    Optional[Any]							       =initializer_range
       snake_case__  :    List[Any]							       =layer_norm_eps
       snake_case__  :    int							       =position_embedding_type
       snake_case__  :    List[str]							       =use_cache
       snake_case__  :    Any							       =pad_token_id



      @classmethod
      def 						lowercase__					(     cls				,			a				,			**a   ):
       cls._set_token_in_kwargs(a   )

       snake_case__       ,   snake_case__  :    int							       =cls.get_config_dict(a				,			**a   )

       # get the text config dict if we are loading from AlignConfig
       if config_dict.get("""model_type"""   ) == "align":
        snake_case__  :    Union[str, Any]							       =config_dict["""text_config"""]

       if "model_type" in config_dict and hasattr(cls				,			"""model_type"""   ) and config_dict["model_type"] != cls.model_type:
        logger.warning(
            F"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
            F"{cls.model_type}. This is not supported for all configurations of models and can yield errors."   )

       return cls.from_dict(a				,			**a   )







class  _lowercase   (   _A						):
      _a							:						List[Any]										=							'align_vision_model'
      def __init__(     self				,			a = 3				,			a = 6_0_0				,			a = 2.0				,			a = 3.1				,			a = 8				,			a = [3, 3, 5, 3, 5, 5, 3]				,			a = [3_2, 1_6, 2_4, 4_0, 8_0, 1_1_2, 1_9_2]				,			a = [1_6, 2_4, 4_0, 8_0, 1_1_2, 1_9_2, 3_2_0]				,			a = []				,			a = [1, 2, 2, 2, 1, 2, 1]				,			a = [1, 2, 2, 3, 3, 4, 1]				,			a = [1, 6, 6, 6, 6, 6, 6]				,			a = 0.25				,			a = "swish"				,			a = 2_5_6_0				,			a = "mean"				,			a = 0.02				,			a = 0.001				,			a = 0.99				,			a = 0.2				,			**a				,			):
       super().__init__(**a   )

       snake_case__  :    Dict							       =num_channels
       snake_case__  :    Tuple							       =image_size
       snake_case__  :    Optional[Any]							       =width_coefficient
       snake_case__  :    List[Any]							       =depth_coefficient
       snake_case__  :    str							       =depth_divisor
       snake_case__  :    List[Any]							       =kernel_sizes
       snake_case__  :    Optional[int]							       =in_channels
       snake_case__  :    Dict							       =out_channels
       snake_case__  :    Optional[Any]							       =depthwise_padding
       snake_case__  :    int							       =strides
       snake_case__  :    Union[str, Any]							       =num_block_repeats
       snake_case__  :    Any							       =expand_ratios
       snake_case__  :    int							       =squeeze_expansion_ratio
       snake_case__  :    List[Any]							       =hidden_act
       snake_case__  :    str							       =hidden_dim
       snake_case__  :    List[str]							       =pooling_type
       snake_case__  :    int							       =initializer_range
       snake_case__  :    List[str]							       =batch_norm_eps
       snake_case__  :    str							       =batch_norm_momentum
       snake_case__  :    Union[str, Any]							       =drop_connect_rate
       snake_case__  :    Optional[Any]							       =sum(a   ) * 4



      @classmethod
      def 						lowercase__					(     cls				,			a				,			**a   ):
       cls._set_token_in_kwargs(a   )

       snake_case__       ,   snake_case__  :    Union[str, Any]							       =cls.get_config_dict(a				,			**a   )

       # get the vision config dict if we are loading from AlignConfig
       if config_dict.get("""model_type"""   ) == "align":
        snake_case__  :    Optional[int]							       =config_dict["""vision_config"""]

       if "model_type" in config_dict and hasattr(cls				,			"""model_type"""   ) and config_dict["model_type"] != cls.model_type:
        logger.warning(
            F"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
            F"{cls.model_type}. This is not supported for all configurations of models and can yield errors."   )

       return cls.from_dict(a				,			**a   )







class  _lowercase   (   _A						):
      _a							:						str										=							'align'
      _a							:						int										=							True
      def __init__(     self				,			a=None				,			a=None				,			a=6_4_0				,			a=1.0				,			a=0.02				,			**a				,			):
       super().__init__(**a   )

       if text_config is None:
        snake_case__  :    str							       ={}
        logger.info("""text_config is None. Initializing the AlignTextConfig with default values."""   )

       if vision_config is None:
        snake_case__  :    int							       ={}
        logger.info("""vision_config is None. Initializing the AlignVisionConfig with default values."""   )

       snake_case__  :    Dict							       =AlignTextConfig(**a   )
       snake_case__  :    List[Any]							       =AlignVisionConfig(**a   )

       snake_case__  :    List[str]							       =projection_dim
       snake_case__  :    List[Any]							       =temperature_init_value
       snake_case__  :    Any							       =initializer_range
      @classmethod
      def 						lowercase__					(     cls				,			a				,			a				,			**a   ):
       return cls(text_config=text_config.to_dict()				,			vision_config=vision_config.to_dict()				,			**a   )



      def 						lowercase__					(     self   ):
       snake_case__  :    int							       =copy.deepcopy(self.__dict__   )
       snake_case__  :    Optional[int]							       =self.text_config.to_dict()
       snake_case__  :    Any							       =self.vision_config.to_dict()
       snake_case__  :    Optional[int]							       =self.__class__.model_type
       return output













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   )





