



"""simple docstring"""

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


if TYPE_CHECKING:
	pass

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


lowercase__			   =  logging.get_logger(__name__)

lowercase__			   =  {
    "kakaobrain/align-base": "https://huggingface.co/kakaobrain/align-base/resolve/main/config.json",
}




class SCREAMING_SNAKE_CASE__					(		_UpperCamelCase      ):
						_lowerCAmelCase							=					"align_text_model"



						def __init__(self						,     _lowercase=30522						,     _lowercase=768						,     _lowercase=12						,     _lowercase=12						,     _lowercase=3072						,     _lowercase="gelu"						,     _lowercase=0.1						,     _lowercase=0.1						,     _lowercase=512						,     _lowercase=2						,     _lowercase=0.02						,     _lowercase=1e-12						,     _lowercase=0						,     _lowercase="absolute"						,     _lowercase=True						,     **_lowercase						,     ):





													'''simple docstring'''

													super().__init__(**_UpperCAmelCase							)

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





						@classmethod
						def     lowerCAmelCase__(cls						,     _lowercase						,     **_lowercase							):





													'''simple docstring'''

													cls._set_token_in_kwargs(_UpperCAmelCase							)

													__a		:		Any									=       cls.get_config_dict(_UpperCAmelCase						,     **_UpperCAmelCase							)

													# get the text config dict if we are loading from AlignConfig
													if config_dict.get("""model_type"""							) == "align":
																				__a		:		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(_UpperCAmelCase						,     **_UpperCAmelCase							)







class SCREAMING_SNAKE_CASE__					(		_UpperCamelCase      ):
						_lowerCAmelCase							=					"align_vision_model"



						def __init__(self						,     _lowercase = 3						,     _lowercase = 600						,     _lowercase = 2.0						,     _lowercase = 3.1						,     _lowercase = 8						,     _lowercase = [3, 3, 5, 3, 5, 5, 3]						,     _lowercase = [32, 16, 24, 40, 80, 112, 192]						,     _lowercase = [16, 24, 40, 80, 112, 192, 320]						,     _lowercase = []						,     _lowercase = [1, 2, 2, 2, 1, 2, 1]						,     _lowercase = [1, 2, 2, 3, 3, 4, 1]						,     _lowercase = [1, 6, 6, 6, 6, 6, 6]						,     _lowercase = 0.25						,     _lowercase = "swish"						,     _lowercase = 2560						,     _lowercase = "mean"						,     _lowercase = 0.02						,     _lowercase = 0.001						,     _lowercase = 0.99						,     _lowercase = 0.2						,     **_lowercase						,     ):





													'''simple docstring'''

													super().__init__(**_UpperCAmelCase							)

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





						@classmethod
						def     lowerCAmelCase__(cls						,     _lowercase						,     **_lowercase							):





													'''simple docstring'''

													cls._set_token_in_kwargs(_UpperCAmelCase							)

													__a		:		int									=       cls.get_config_dict(_UpperCAmelCase						,     **_UpperCAmelCase							)

													# get the vision config dict if we are loading from AlignConfig
													if config_dict.get("""model_type"""							) == "align":
																				__a		:		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(_UpperCAmelCase						,     **_UpperCAmelCase							)







class SCREAMING_SNAKE_CASE__					(		_UpperCamelCase      ):
						_lowerCAmelCase							=					"align"
						_lowerCAmelCase							=					True



						def __init__(self						,     _lowercase=None						,     _lowercase=None						,     _lowercase=640						,     _lowercase=1.0						,     _lowercase=0.02						,     **_lowercase						,     ):





													'''simple docstring'''

													super().__init__(**_UpperCAmelCase							)

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

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

													__a		:		Union[str, Any]									=       AlignTextConfig(**_UpperCAmelCase							)
													__a		:		Any									=       AlignVisionConfig(**_UpperCAmelCase							)

													__a		:		int									=       projection_dim
													__a		:		Union[str, Any]									=       temperature_init_value
													__a		:		Dict									=       initializer_range



						@classmethod
						def     lowerCAmelCase__(cls						,     _lowercase						,     _lowercase						,     **_lowercase							):





													'''simple docstring'''

													return cls(text_config=text_config.to_dict()						,     vision_config=vision_config.to_dict()						,     **_UpperCAmelCase							)





						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Any									=       copy.deepcopy(self.__dict__							)
													__a		:		str									=       self.text_config.to_dict()
													__a		:		List[str]									=       self.vision_config.to_dict()
													__a		:		List[Any]									=       self.__class__.model_type
													return output






"""simple docstring"""

import unittest

import numpy as np
import torch
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyVaaPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class SCREAMING_SNAKE_CASE__					(		__snake_case     ,					unittest.TestCase      ):
						_lowerCAmelCase							=					KandinskyVaaPriorPipeline
						_lowerCAmelCase							=					["prompt"]
						_lowerCAmelCase							=					["prompt", "negative_prompt"]
						_lowerCAmelCase							=					[
						    "num_images_per_prompt",
						    "generator",
						    "num_inference_steps",
						    "latents",
						    "negative_prompt",
						    "guidance_scale",
						    "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'''

													__a		:		List[str]									=       CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""							)
													return tokenizer



						@property
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													torch.manual_seed(0							)
													__a		:		str									=       CLIPTextConfig(
													    bos_token_id=0						,     eos_token_id=2						,     hidden_size=self.text_embedder_hidden_size						,     projection_dim=self.text_embedder_hidden_size						,     intermediate_size=37						,     layer_norm_eps=1e-05						,     num_attention_heads=4						,     num_hidden_layers=5						,     pad_token_id=1						,     vocab_size=1000						,     )
													return CLIPTextModelWithProjection(_lowercase							)



						@property
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													torch.manual_seed(0							)

													__a		:		Dict									=       {
													    """num_attention_heads""": 2,
													    """attention_head_dim""": 12,
													    """embedding_dim""": self.text_embedder_hidden_size,
													    """num_layers""": 1,
													}

													__a		:		Tuple									=       PriorTransformer(**_lowercase							)
													# clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
													__a		:		int									=       nn.Parameter(torch.ones(model.clip_std.shape							)							)
													return model



						@property
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													torch.manual_seed(0							)
													__a		:		List[str]									=       CLIPVisionConfig(
													    hidden_size=self.text_embedder_hidden_size						,     image_size=224						,     projection_dim=self.text_embedder_hidden_size						,     intermediate_size=37						,     num_attention_heads=4						,     num_channels=3						,     num_hidden_layers=5						,     patch_size=14						,     )

													__a		:		Optional[Any]									=       CLIPVisionModelWithProjection(_lowercase							)
													return model



						@property
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Optional[Any]									=       CLIPImageProcessor(
													    crop_size=224						,     do_center_crop=_lowercase						,     do_normalize=_lowercase						,     do_resize=_lowercase						,     image_mean=[0.4814_5466, 0.457_8275, 0.4082_1073]						,     image_std=[0.2686_2954, 0.2613_0258, 0.2757_7711]						,     resample=3						,     size=224						,     )

													return image_processor



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Union[str, Any]									=       self.dummy_prior
													__a		:		int									=       self.dummy_image_encoder
													__a		:		Any									=       self.dummy_text_encoder
													__a		:		int									=       self.dummy_tokenizer
													__a		:		Optional[Any]									=       self.dummy_image_processor

													__a		:		List[Any]									=       UnCLIPScheduler(
													    variance_type="""fixed_small_log"""						,     prediction_type="""sample"""						,     num_train_timesteps=1000						,     clip_sample=_lowercase						,     clip_sample_range=10.0						,     )

													__a		:		List[Any]									=       {
													    """prior""": prior,
													    """image_encoder""": image_encoder,
													    """text_encoder""": text_encoder,
													    """tokenizer""": tokenizer,
													    """scheduler""": scheduler,
													    """image_processor""": image_processor,
													}

													return components



						def     lowerCAmelCase__(self						,     _lowercase						,     _lowercase=0							):





													'''simple docstring'''

													if str(_lowercase							).startswith("""mps"""							):
																				__a		:		Dict									=       torch.manual_seed(_lowercase							)
													else:
																				__a		:		Union[str, Any]									=       torch.Generator(device=_lowercase							).manual_seed(_lowercase							)
													__a		:		Union[str, Any]									=       {
													    """prompt""": """horse""",
													    """generator""": generator,
													    """guidance_scale""": 4.0,
													    """num_inference_steps""": 2,
													    """output_type""": """np""",
													}
													return inputs



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Union[str, Any]									=       """cpu"""

													__a		:		Union[str, Any]									=       self.get_dummy_components()

													__a		:		Dict									=       self.pipeline_class(**_lowercase							)
													__a		:		Tuple									=       pipe.to(_lowercase							)

													pipe.set_progress_bar_config(disable=_lowercase							)

													__a		:		Optional[int]									=       pipe(**self.get_dummy_inputs(_lowercase							)							)
													__a		:		str									=       output.image_embeds

													__a		:		Any									=       pipe(
													    **self.get_dummy_inputs(_lowercase							)						,     return_dict=_lowercase						,     )[0]

													__a		:		List[Any]									=       image[0, -10:]
													__a		:		List[Any]									=       image_from_tuple[0, -10:]

													assert image.shape == (1, 32)

													__a		:		Optional[Any]									=       np.array(
													    [-0.0532, 1.7120, 0.3656, -1.0852, -0.8946, -1.1756, 0.4348, 0.2482, 0.5146, -0.1156]							)

													assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-2
													assert np.abs(image_from_tuple_slice.flatten() - expected_slice							).max() < 1e-2



						@skip_mps
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Any									=       torch_device == """cpu"""
													__a		:		Any									=       True
													__a		:		Any									=       False

													self._test_inference_batch_single_identical(
													    test_max_difference=_lowercase						,     relax_max_difference=_lowercase						,     test_mean_pixel_difference=_lowercase						,     )





						@skip_mps
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Optional[int]									=       torch_device == """cpu"""
													__a		:		Union[str, Any]									=       False

													self._test_attention_slicing_forward_pass(
													    test_max_difference=_lowercase						,     test_mean_pixel_difference=_lowercase						,     )

