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      :			Any							       =		logging.get_logger(__name__)

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







class __snake_case    (						lowerCamelCase_		):
   lowerCAmelCase_														=							"align_text_model"




   def __init__(   self      :    Optional[Any]	,     _lowercase      :    Optional[Any]=3_05_22	,     _lowercase      :    Dict=7_68	,     _lowercase      :    Tuple=12	,     _lowercase      :    Optional[Any]=12	,     _lowercase      :    Union[str, Any]=30_72	,     _lowercase      :    List[str]="gelu"	,     _lowercase      :    Dict=0.1	,     _lowercase      :    List[str]=0.1	,     _lowercase      :    Any=5_12	,     _lowercase      :    List[Any]=2	,     _lowercase      :    Tuple=0.02	,     _lowercase      :    List[Any]=1E-12	,     _lowercase      :    Union[str, Any]=0	,     _lowercase      :    Union[str, Any]="absolute"	,     _lowercase      :    Optional[int]=True	,     **_lowercase      :    Optional[Any]	,     ):



     """simple docstring"""






     super().__init__(**_lowercase				)

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




   @classmethod
   def      __a						(   cls      :    str	,     _lowercase      :    Union[str, os.PathLike]	,     **_lowercase      :    Optional[int]				):



     """simple docstring"""






     cls._set_token_in_kwargs(_lowercase				)

     SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__      							=				cls.get_config_dict(_lowercase	,     **_lowercase				)

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







class __snake_case    (						lowerCamelCase_		):
   lowerCAmelCase_														=							"align_vision_model"




   def __init__(   self      :    List[Any]	,     _lowercase      :    int = 3	,     _lowercase      :    int = 6_00	,     _lowercase      :    float = 2.0	,     _lowercase      :    float = 3.1	,     _lowercase      :    int = 8	,     _lowercase      :    List[int] = [3, 3, 5, 3, 5, 5, 3]	,     _lowercase      :    List[int] = [32, 16, 24, 40, 80, 1_12, 1_92]	,     _lowercase      :    List[int] = [16, 24, 40, 80, 1_12, 1_92, 3_20]	,     _lowercase      :    List[int] = []	,     _lowercase      :    List[int] = [1, 2, 2, 2, 1, 2, 1]	,     _lowercase      :    List[int] = [1, 2, 2, 3, 3, 4, 1]	,     _lowercase      :    List[int] = [1, 6, 6, 6, 6, 6, 6]	,     _lowercase      :    float = 0.25	,     _lowercase      :    str = "swish"	,     _lowercase      :    int = 25_60	,     _lowercase      :    str = "mean"	,     _lowercase      :    float = 0.02	,     _lowercase      :    float = 0.0_01	,     _lowercase      :    float = 0.99	,     _lowercase      :    float = 0.2	,     **_lowercase      :    List[str]	,     ):



     """simple docstring"""






     super().__init__(**_lowercase				)

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




   @classmethod
   def      __a						(   cls      :    Any	,     _lowercase      :    Union[str, os.PathLike]	,     **_lowercase      :    List[Any]				):



     """simple docstring"""






     cls._set_token_in_kwargs(_lowercase				)

     SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__      							=				cls.get_config_dict(_lowercase	,     **_lowercase				)

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







class __snake_case    (						lowerCamelCase_		):
   lowerCAmelCase_														=							"align"
   lowerCAmelCase_														=							True




   def __init__(   self      :    Optional[Any]	,     _lowercase      :    Optional[Any]=None	,     _lowercase      :    str=None	,     _lowercase      :    Union[str, Any]=6_40	,     _lowercase      :    Union[str, Any]=1.0	,     _lowercase      :    Any=0.02	,     **_lowercase      :    int	,     ):



     """simple docstring"""






     super().__init__(**_lowercase				)

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

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

     SCREAMING_SNAKE_CASE__      							=				AlignTextConfig(**_lowercase				)
     SCREAMING_SNAKE_CASE__      							=				AlignVisionConfig(**_lowercase				)

     SCREAMING_SNAKE_CASE__      							=				projection_dim
     SCREAMING_SNAKE_CASE__      							=				temperature_init_value
     SCREAMING_SNAKE_CASE__      							=				initializer_range




   @classmethod
   def      __a						(   cls      :    Optional[int]	,     _lowercase      :    AlignTextConfig	,     _lowercase      :    AlignVisionConfig	,     **_lowercase      :    str				):



     """simple docstring"""






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




   def      __a						(   self      :    Tuple				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				copy.deepcopy(self.__dict__				)
     SCREAMING_SNAKE_CASE__      							=				self.text_config.to_dict()
     SCREAMING_SNAKE_CASE__      							=				self.vision_config.to_dict()
     SCREAMING_SNAKE_CASE__      							=				self.__class__.model_type
     return output

import argparse
import io

import requests
import torch
from omegaconf import OmegaConf

from diffusers import AutoencoderKL
from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
    assign_to_checkpoint,
    conv_attn_to_linear,
    create_vae_diffusers_config,
    renew_vae_attention_paths,
    renew_vae_resnet_paths,
)



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	Optional[int]			,					__UpperCamelCase   :	str     )							-> Tuple:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__      							=				checkpoint

  SCREAMING_SNAKE_CASE__      							=				{}

  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""encoder.conv_in.weight"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""encoder.conv_in.bias"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""encoder.conv_out.weight"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""encoder.conv_out.bias"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""encoder.norm_out.weight"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""encoder.norm_out.bias"""]

  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""decoder.conv_in.weight"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""decoder.conv_in.bias"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""decoder.conv_out.weight"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""decoder.conv_out.bias"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""decoder.norm_out.weight"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""decoder.norm_out.bias"""]

  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""quant_conv.weight"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""quant_conv.bias"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""post_quant_conv.weight"""]
  SCREAMING_SNAKE_CASE__      							=				vae_state_dict["""post_quant_conv.bias"""]

  # Retrieves the keys for the encoder down blocks only
  SCREAMING_SNAKE_CASE__      							=				len({""".""".join(layer.split("""."""     )[:3]     ) for layer in vae_state_dict if """encoder.down""" in layer}     )
  SCREAMING_SNAKE_CASE__      							=				{
      layer_id: [key for key in vae_state_dict if f"""down.{layer_id}""" in key] for layer_id in range(__UpperCamelCase     )
  }

  # Retrieves the keys for the decoder up blocks only
  SCREAMING_SNAKE_CASE__      							=				len({""".""".join(layer.split("""."""     )[:3]     ) for layer in vae_state_dict if """decoder.up""" in layer}     )
  SCREAMING_SNAKE_CASE__      							=				{
      layer_id: [key for key in vae_state_dict if f"""up.{layer_id}""" in key] for layer_id in range(__UpperCamelCase     )
  }

  for i in range(__UpperCamelCase     ):
    SCREAMING_SNAKE_CASE__      							=				[key for key in down_blocks[i] if f"""down.{i}""" in key and f"""down.{i}.downsample""" not in key]

    if f"""encoder.down.{i}.downsample.conv.weight""" in vae_state_dict:
      SCREAMING_SNAKE_CASE__      							=				vae_state_dict.pop(
          f"""encoder.down.{i}.downsample.conv.weight"""     )
      SCREAMING_SNAKE_CASE__      							=				vae_state_dict.pop(
          f"""encoder.down.{i}.downsample.conv.bias"""     )

    SCREAMING_SNAKE_CASE__      							=				renew_vae_resnet_paths(__UpperCamelCase     )
    SCREAMING_SNAKE_CASE__      							=				{"""old""": f"""down.{i}.block""", """new""": f"""down_blocks.{i}.resnets"""}
    assign_to_checkpoint(__UpperCamelCase			,					__UpperCamelCase			,					__UpperCamelCase			,					additional_replacements=[meta_path]			,					config=__UpperCamelCase     )

  SCREAMING_SNAKE_CASE__      							=				[key for key in vae_state_dict if """encoder.mid.block""" in key]
  SCREAMING_SNAKE_CASE__      							=				2
  for i in range(1			,					num_mid_res_blocks + 1     ):
    SCREAMING_SNAKE_CASE__      							=				[key for key in mid_resnets if f"""encoder.mid.block_{i}""" in key]

    SCREAMING_SNAKE_CASE__      							=				renew_vae_resnet_paths(__UpperCamelCase     )
    SCREAMING_SNAKE_CASE__      							=				{"""old""": f"""mid.block_{i}""", """new""": f"""mid_block.resnets.{i - 1}"""}
    assign_to_checkpoint(__UpperCamelCase			,					__UpperCamelCase			,					__UpperCamelCase			,					additional_replacements=[meta_path]			,					config=__UpperCamelCase     )

  SCREAMING_SNAKE_CASE__      							=				[key for key in vae_state_dict if """encoder.mid.attn""" in key]
  SCREAMING_SNAKE_CASE__      							=				renew_vae_attention_paths(__UpperCamelCase     )
  SCREAMING_SNAKE_CASE__      							=				{"""old""": """mid.attn_1""", """new""": """mid_block.attentions.0"""}
  assign_to_checkpoint(__UpperCamelCase			,					__UpperCamelCase			,					__UpperCamelCase			,					additional_replacements=[meta_path]			,					config=__UpperCamelCase     )
  conv_attn_to_linear(__UpperCamelCase     )

  for i in range(__UpperCamelCase     ):
    SCREAMING_SNAKE_CASE__      							=				num_up_blocks - 1 - i
    SCREAMING_SNAKE_CASE__      							=				[
        key for key in up_blocks[block_id] if f"""up.{block_id}""" in key and f"""up.{block_id}.upsample""" not in key
    ]

    if f"""decoder.up.{block_id}.upsample.conv.weight""" in vae_state_dict:
      SCREAMING_SNAKE_CASE__      							=				vae_state_dict[
          f"""decoder.up.{block_id}.upsample.conv.weight"""
      ]
      SCREAMING_SNAKE_CASE__      							=				vae_state_dict[
          f"""decoder.up.{block_id}.upsample.conv.bias"""
      ]

    SCREAMING_SNAKE_CASE__      							=				renew_vae_resnet_paths(__UpperCamelCase     )
    SCREAMING_SNAKE_CASE__      							=				{"""old""": f"""up.{block_id}.block""", """new""": f"""up_blocks.{i}.resnets"""}
    assign_to_checkpoint(__UpperCamelCase			,					__UpperCamelCase			,					__UpperCamelCase			,					additional_replacements=[meta_path]			,					config=__UpperCamelCase     )

  SCREAMING_SNAKE_CASE__      							=				[key for key in vae_state_dict if """decoder.mid.block""" in key]
  SCREAMING_SNAKE_CASE__      							=				2
  for i in range(1			,					num_mid_res_blocks + 1     ):
    SCREAMING_SNAKE_CASE__      							=				[key for key in mid_resnets if f"""decoder.mid.block_{i}""" in key]

    SCREAMING_SNAKE_CASE__      							=				renew_vae_resnet_paths(__UpperCamelCase     )
    SCREAMING_SNAKE_CASE__      							=				{"""old""": f"""mid.block_{i}""", """new""": f"""mid_block.resnets.{i - 1}"""}
    assign_to_checkpoint(__UpperCamelCase			,					__UpperCamelCase			,					__UpperCamelCase			,					additional_replacements=[meta_path]			,					config=__UpperCamelCase     )

  SCREAMING_SNAKE_CASE__      							=				[key for key in vae_state_dict if """decoder.mid.attn""" in key]
  SCREAMING_SNAKE_CASE__      							=				renew_vae_attention_paths(__UpperCamelCase     )
  SCREAMING_SNAKE_CASE__      							=				{"""old""": """mid.attn_1""", """new""": """mid_block.attentions.0"""}
  assign_to_checkpoint(__UpperCamelCase			,					__UpperCamelCase			,					__UpperCamelCase			,					additional_replacements=[meta_path]			,					config=__UpperCamelCase     )
  conv_attn_to_linear(__UpperCamelCase     )
  return new_checkpoint



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	str			,					__UpperCamelCase   :	str			,					)							-> Optional[int]:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__      							=				requests.get(
      """ https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml"""     )
  SCREAMING_SNAKE_CASE__      							=				io.BytesIO(r.content     )

  SCREAMING_SNAKE_CASE__      							=				OmegaConf.load(__UpperCamelCase     )
  SCREAMING_SNAKE_CASE__      							=				5_12
  SCREAMING_SNAKE_CASE__      							=				"""cuda""" if torch.cuda.is_available() else """cpu"""
  if checkpoint_path.endswith("""safetensors"""     ):
    from safetensors import safe_open

    SCREAMING_SNAKE_CASE__      							=				{}
    with safe_open(__UpperCamelCase			,					framework="""pt"""			,					device="""cpu"""     ) as f:
      for key in f.keys():
        SCREAMING_SNAKE_CASE__      							=				f.get_tensor(__UpperCamelCase     )
  else:
    SCREAMING_SNAKE_CASE__      							=				torch.load(__UpperCamelCase			,					map_location=__UpperCamelCase     )["""state_dict"""]

  # Convert the VAE model.
  SCREAMING_SNAKE_CASE__      							=				create_vae_diffusers_config(__UpperCamelCase			,					image_size=__UpperCamelCase     )
  SCREAMING_SNAKE_CASE__      							=				custom_convert_ldm_vae_checkpoint(__UpperCamelCase			,					__UpperCamelCase     )

  SCREAMING_SNAKE_CASE__      							=				AutoencoderKL(**__UpperCamelCase     )
  vae.load_state_dict(__UpperCamelCase     )
  vae.save_pretrained(__UpperCamelCase     )


if __name__ == "__main__":
   __lowerCamelCase      :			List[Any]							       =		argparse.ArgumentParser()

   parser.add_argument('''--vae_pt_path''', default=None, type=str, required=True, help='''Path to the VAE.pt to convert.''')
   parser.add_argument('''--dump_path''', default=None, type=str, required=True, help='''Path to the VAE.pt to convert.''')

   __lowerCamelCase      :			Optional[int]							       =		parser.parse_args()

   vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)
