




'''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


UpperCAmelCase_ 		=      logging.get_logger(__name__)

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





class    __lowercase						(  __magic_name__ ):
   _a							=      """align_text_model"""







   def __init__(	self      ,							UpperCamelCase=3_0522      ,							UpperCamelCase=768      ,							UpperCamelCase=12      ,							UpperCamelCase=12      ,							UpperCamelCase=3072      ,							UpperCamelCase="gelu"      ,							UpperCamelCase=0.1      ,							UpperCamelCase=0.1      ,							UpperCamelCase=512      ,							UpperCamelCase=2      ,							UpperCamelCase=0.02      ,							UpperCamelCase=1e-12      ,							UpperCamelCase=0      ,							UpperCamelCase="absolute"      ,							UpperCamelCase=True      ,							**UpperCamelCase      ,							)				->    Tuple:
     super().__init__(**UpperCamelCase      )

     __a			=					vocab_size
     __a			=					hidden_size
     __a			=					num_hidden_layers
     __a			=					num_attention_heads
     __a			=					hidden_act
     __a			=					intermediate_size
     __a			=					hidden_dropout_prob
     __a			=					attention_probs_dropout_prob
     __a			=					max_position_embeddings
     __a			=					type_vocab_size
     __a			=					initializer_range
     __a			=					layer_norm_eps
     __a			=					position_embedding_type
     __a			=					use_cache
     __a			=					pad_token_id







   @classmethod
   def 		UpperCamelCase__			(	cls      ,							UpperCamelCase      ,							**UpperCamelCase      )				->    "PretrainedConfig":
     cls._set_token_in_kwargs(UpperCamelCase      )

     __a	,       __a			=					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			=					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    __lowercase						(  __magic_name__ ):
   _a							=      """align_vision_model"""







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

     __a			=					num_channels
     __a			=					image_size
     __a			=					width_coefficient
     __a			=					depth_coefficient
     __a			=					depth_divisor
     __a			=					kernel_sizes
     __a			=					in_channels
     __a			=					out_channels
     __a			=					depthwise_padding
     __a			=					strides
     __a			=					num_block_repeats
     __a			=					expand_ratios
     __a			=					squeeze_expansion_ratio
     __a			=					hidden_act
     __a			=					hidden_dim
     __a			=					pooling_type
     __a			=					initializer_range
     __a			=					batch_norm_eps
     __a			=					batch_norm_momentum
     __a			=					drop_connect_rate
     __a			=					sum(UpperCamelCase      ) * 4







   @classmethod
   def 		UpperCamelCase__			(	cls      ,							UpperCamelCase      ,							**UpperCamelCase      )				->    "PretrainedConfig":
     cls._set_token_in_kwargs(UpperCamelCase      )

     __a	,       __a			=					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			=					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    __lowercase						(  __magic_name__ ):
   _a							=      """align"""
   _a							=      True







   def __init__(	self      ,							UpperCamelCase=None      ,							UpperCamelCase=None      ,							UpperCamelCase=640      ,							UpperCamelCase=1.0      ,							UpperCamelCase=0.02      ,							**UpperCamelCase      ,							)				->    int:
     super().__init__(**UpperCamelCase      )

     if text_config is None:
       __a			=					{}
       logger.info('text_config is None. Initializing the AlignTextConfig with default values.'      )

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

     __a			=					AlignTextConfig(**UpperCamelCase      )
     __a			=					AlignVisionConfig(**UpperCamelCase      )

     __a			=					projection_dim
     __a			=					temperature_init_value
     __a			=					initializer_range







   @classmethod
   def 		UpperCamelCase__			(	cls      ,							UpperCamelCase      ,							UpperCamelCase      ,							**UpperCamelCase      )				->    Union[str, Any]:
     return cls(text_config=text_config.to_dict()      ,							vision_config=vision_config.to_dict()      ,							**UpperCamelCase      )







   def 		UpperCamelCase__			(	self      )				->    Optional[Any]:
     __a			=					copy.deepcopy(self.__dict__      )
     __a			=					self.text_config.to_dict()
     __a			=					self.vision_config.to_dict()
     __a			=					self.__class__.model_type
     return output









'''simple docstring'''


import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np





class    __lowercase						:







   def 		UpperCamelCase__			(	self      )				->    List[Any]:
     torch.manual_seed(0      )
     __a			=					TaEncoderModel.from_pretrained('hf-internal-testing/tiny-random-t5'      )

     torch.manual_seed(0      )
     __a			=					AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-t5'      )

     torch.manual_seed(0      )
     __a			=					UNetaDConditionModel(
         sample_size=32      ,							layers_per_block=1      ,							block_out_channels=[32, 64]      ,							down_block_types=[
             'ResnetDownsampleBlock2D',
             'SimpleCrossAttnDownBlock2D',
         ]      ,							mid_block_type='UNetMidBlock2DSimpleCrossAttn'      ,							up_block_types=['SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D']      ,							in_channels=3      ,							out_channels=6      ,							cross_attention_dim=32      ,							encoder_hid_dim=32      ,							attention_head_dim=8      ,							addition_embed_type='text'      ,							addition_embed_type_num_heads=2      ,							cross_attention_norm='group_norm'      ,							resnet_time_scale_shift='scale_shift'      ,							act_fn='gelu'      ,							)
     unet.set_attn_processor(AttnAddedKVProcessor()      )  # For reproducibility tests

     torch.manual_seed(0      )
     __a			=					DDPMScheduler(
         num_train_timesteps=1000      ,							beta_schedule='squaredcos_cap_v2'      ,							beta_start=0.0_001      ,							beta_end=0.02      ,							thresholding=UpperCamelCase      ,							dynamic_thresholding_ratio=0.95      ,							sample_max_value=1.0      ,							prediction_type='epsilon'      ,							variance_type='learned_range'      ,							)

     torch.manual_seed(0      )
     __a			=					IFWatermarker()

     return {
         "text_encoder": text_encoder,
         "tokenizer": tokenizer,
         "unet": unet,
         "scheduler": scheduler,
         "watermarker": watermarker,
         "safety_checker": None,
         "feature_extractor": None,
     }







   def 		UpperCamelCase__			(	self      )				->    Tuple:
     torch.manual_seed(0      )
     __a			=					TaEncoderModel.from_pretrained('hf-internal-testing/tiny-random-t5'      )

     torch.manual_seed(0      )
     __a			=					AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-t5'      )

     torch.manual_seed(0      )
     __a			=					UNetaDConditionModel(
         sample_size=32      ,							layers_per_block=[1, 2]      ,							block_out_channels=[32, 64]      ,							down_block_types=[
             'ResnetDownsampleBlock2D',
             'SimpleCrossAttnDownBlock2D',
         ]      ,							mid_block_type='UNetMidBlock2DSimpleCrossAttn'      ,							up_block_types=['SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D']      ,							in_channels=6      ,							out_channels=6      ,							cross_attention_dim=32      ,							encoder_hid_dim=32      ,							attention_head_dim=8      ,							addition_embed_type='text'      ,							addition_embed_type_num_heads=2      ,							cross_attention_norm='group_norm'      ,							resnet_time_scale_shift='scale_shift'      ,							act_fn='gelu'      ,							class_embed_type='timestep'      ,							mid_block_scale_factor=1.414      ,							time_embedding_act_fn='gelu'      ,							time_embedding_dim=32      ,							)
     unet.set_attn_processor(AttnAddedKVProcessor()      )  # For reproducibility tests

     torch.manual_seed(0      )
     __a			=					DDPMScheduler(
         num_train_timesteps=1000      ,							beta_schedule='squaredcos_cap_v2'      ,							beta_start=0.0_001      ,							beta_end=0.02      ,							thresholding=UpperCamelCase      ,							dynamic_thresholding_ratio=0.95      ,							sample_max_value=1.0      ,							prediction_type='epsilon'      ,							variance_type='learned_range'      ,							)

     torch.manual_seed(0      )
     __a			=					DDPMScheduler(
         num_train_timesteps=1000      ,							beta_schedule='squaredcos_cap_v2'      ,							beta_start=0.0_001      ,							beta_end=0.02      ,							)

     torch.manual_seed(0      )
     __a			=					IFWatermarker()

     return {
         "text_encoder": text_encoder,
         "tokenizer": tokenizer,
         "unet": unet,
         "scheduler": scheduler,
         "image_noising_scheduler": image_noising_scheduler,
         "watermarker": watermarker,
         "safety_checker": None,
         "feature_extractor": None,
     }







   def 		UpperCamelCase__			(	self      )				->    Tuple:
     __a			=					self.get_dummy_components()
     __a			=					self.pipeline_class(**UpperCamelCase      )
     pipe.to(UpperCamelCase      )
     pipe.set_progress_bar_config(disable=UpperCamelCase      )

     __a			=					self.get_dummy_inputs(UpperCamelCase      )

     __a			=					inputs['prompt']
     __a			=					inputs['generator']
     __a			=					inputs['num_inference_steps']
     __a			=					inputs['output_type']

     if "image" in inputs:
       __a			=					inputs['image']
     else:
       __a			=					None

     if "mask_image" in inputs:
       __a			=					inputs['mask_image']
     else:
       __a			=					None

     if "original_image" in inputs:
       __a			=					inputs['original_image']
     else:
       __a			=					None

     __a	,       __a			=					pipe.encode_prompt(UpperCamelCase      )

     # inputs with prompt converted to embeddings
     __a			=					{
         'prompt_embeds': prompt_embeds,
         'negative_prompt_embeds': negative_prompt_embeds,
         'generator': generator,
         'num_inference_steps': num_inference_steps,
         'output_type': output_type,
     }

     if image is not None:
       __a			=					image

     if mask_image is not None:
       __a			=					mask_image

     if original_image is not None:
       __a			=					original_image

     # set all optional components to None
     for optional_component in pipe._optional_components:
       setattr(UpperCamelCase      ,							UpperCamelCase      ,							UpperCamelCase      )

     __a			=					pipe(**UpperCamelCase      )[0]

     with tempfile.TemporaryDirectory() as tmpdir:
       pipe.save_pretrained(UpperCamelCase      )
       __a			=					self.pipeline_class.from_pretrained(UpperCamelCase      )
       pipe_loaded.to(UpperCamelCase      )
       pipe_loaded.set_progress_bar_config(disable=UpperCamelCase      )

     pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()      )  # For reproducibility tests

     for optional_component in pipe._optional_components:
       self.assertTrue(
           getattr(UpperCamelCase      ,							UpperCamelCase      ) is None      ,							f"`{optional_component}` did not stay set to None after loading."      ,							)

     __a			=					self.get_dummy_inputs(UpperCamelCase      )

     __a			=					inputs['generator']
     __a			=					inputs['num_inference_steps']
     __a			=					inputs['output_type']

     # inputs with prompt converted to embeddings
     __a			=					{
         'prompt_embeds': prompt_embeds,
         'negative_prompt_embeds': negative_prompt_embeds,
         'generator': generator,
         'num_inference_steps': num_inference_steps,
         'output_type': output_type,
     }

     if image is not None:
       __a			=					image

     if mask_image is not None:
       __a			=					mask_image

     if original_image is not None:
       __a			=					original_image

     __a			=					pipe_loaded(**UpperCamelCase      )[0]

     __a			=					np.abs(to_np(UpperCamelCase      ) - to_np(UpperCamelCase      )      ).max()
     self.assertLess(UpperCamelCase      ,							1e-4      )







   def 		UpperCamelCase__			(	self      )				->    str:
     __a			=					self.get_dummy_components()
     __a			=					self.pipeline_class(**UpperCamelCase      )
     pipe.to(UpperCamelCase      )
     pipe.set_progress_bar_config(disable=UpperCamelCase      )

     __a			=					self.get_dummy_inputs(UpperCamelCase      )
     __a			=					pipe(**UpperCamelCase      )[0]

     with tempfile.TemporaryDirectory() as tmpdir:
       pipe.save_pretrained(UpperCamelCase      )
       __a			=					self.pipeline_class.from_pretrained(UpperCamelCase      )
       pipe_loaded.to(UpperCamelCase      )
       pipe_loaded.set_progress_bar_config(disable=UpperCamelCase      )

     pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()      )  # For reproducibility tests

     __a			=					self.get_dummy_inputs(UpperCamelCase      )
     __a			=					pipe_loaded(**UpperCamelCase      )[0]

     __a			=					np.abs(to_np(UpperCamelCase      ) - to_np(UpperCamelCase      )      ).max()
     self.assertLess(UpperCamelCase      ,							1e-4      )



