import warnings
from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


__lowerCamelCase      :			Any							       =						logging.get_logger(__name__)

__lowerCamelCase      :			Tuple							       =						{
    "facebook/bart-large": "https://huggingface.co/facebook/bart-large/resolve/main/config.json",
    # See all BART models at https://huggingface.co/models?filter=bart
}







class __snake_case    (						SCREAMING_SNAKE_CASE__		):
   lowerCAmelCase_														=							'''bart'''
   lowerCAmelCase_														=							['''past_key_values''']
   lowerCAmelCase_														=							{'''num_attention_heads''': '''encoder_attention_heads''', '''hidden_size''': '''d_model'''}




   def __init__(   self      :    Dict	,     _lowercase      :    Optional[Any]=5_02_65	,     _lowercase      :    Any=10_24	,     _lowercase      :    Union[str, Any]=12	,     _lowercase      :    Union[str, Any]=40_96	,     _lowercase      :    List[str]=16	,     _lowercase      :    Optional[Any]=12	,     _lowercase      :    Optional[Any]=40_96	,     _lowercase      :    Any=16	,     _lowercase      :    str=0.0	,     _lowercase      :    List[Any]=0.0	,     _lowercase      :    str="gelu"	,     _lowercase      :    Any=10_24	,     _lowercase      :    str=0.1	,     _lowercase      :    Tuple=0.0	,     _lowercase      :    int=0.0	,     _lowercase      :    str=0.02	,     _lowercase      :    Tuple=0.0	,     _lowercase      :    Any=False	,     _lowercase      :    List[Any]=True	,     _lowercase      :    Tuple=3	,     _lowercase      :    Optional[Any]=1	,     _lowercase      :    List[str]=0	,     _lowercase      :    Tuple=2	,     _lowercase      :    Optional[Any]=True	,     _lowercase      :    Any=2	,     _lowercase      :    List[str]=2	,     **_lowercase      :    Tuple	,     ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				vocab_size
     SCREAMING_SNAKE_CASE__      							=				max_position_embeddings
     SCREAMING_SNAKE_CASE__      							=				d_model
     SCREAMING_SNAKE_CASE__      							=				encoder_ffn_dim
     SCREAMING_SNAKE_CASE__      							=				encoder_layers
     SCREAMING_SNAKE_CASE__      							=				encoder_attention_heads
     SCREAMING_SNAKE_CASE__      							=				decoder_ffn_dim
     SCREAMING_SNAKE_CASE__      							=				decoder_layers
     SCREAMING_SNAKE_CASE__      							=				decoder_attention_heads
     SCREAMING_SNAKE_CASE__      							=				dropout
     SCREAMING_SNAKE_CASE__      							=				attention_dropout
     SCREAMING_SNAKE_CASE__      							=				activation_dropout
     SCREAMING_SNAKE_CASE__      							=				activation_function
     SCREAMING_SNAKE_CASE__      							=				init_std
     SCREAMING_SNAKE_CASE__      							=				encoder_layerdrop
     SCREAMING_SNAKE_CASE__      							=				decoder_layerdrop
     SCREAMING_SNAKE_CASE__      							=				classifier_dropout
     SCREAMING_SNAKE_CASE__      							=				use_cache
     SCREAMING_SNAKE_CASE__      							=				encoder_layers
     SCREAMING_SNAKE_CASE__      							=				scale_embedding  # scale factor will be sqrt(d_model) if True

     super().__init__(
         num_labels=A__	,     pad_token_id=A__	,     bos_token_id=A__	,     eos_token_id=A__	,     is_encoder_decoder=A__	,     decoder_start_token_id=A__	,     forced_eos_token_id=A__	,     **A__	,     )

     # ensure backward compatibility for BART CNN models
     if self.forced_bos_token_id is None and kwargs.get("""force_bos_token_to_be_generated"""	,     A__				):
       SCREAMING_SNAKE_CASE__      							=				self.bos_token_id
       warnings.warn(
           f"""Please make sure the config includes `forced_bos_token_id={self.bos_token_id}` in future versions. """
           """The config can simply be saved and uploaded again to be fixed."""				)







class __snake_case    (						SCREAMING_SNAKE_CASE__		):




   @property
   def      __a						(   self      :    Any				):



     """simple docstring"""






     if self.task in ["default", "seq2seq-lm"]:
       SCREAMING_SNAKE_CASE__      							=				OrderedDict(
           [
               ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
               ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
           ]				)

       if self.use_past:
         SCREAMING_SNAKE_CASE__      							=				{0: """batch"""}
         SCREAMING_SNAKE_CASE__      							=				{0: """batch""", 1: """past_decoder_sequence + sequence"""}
       else:
         SCREAMING_SNAKE_CASE__      							=				{0: """batch""", 1: """decoder_sequence"""}
         SCREAMING_SNAKE_CASE__      							=				{0: """batch""", 1: """decoder_sequence"""}

       if self.use_past:
         self.fill_with_past_key_values_(A__	,     direction="""inputs"""				)
     elif self.task == "causal-lm":
       # TODO: figure this case out.
       SCREAMING_SNAKE_CASE__      							=				OrderedDict(
           [
               ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
               ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
           ]				)
       if self.use_past:
         SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__      							=				self.num_layers
         for i in range(A__				):
           SCREAMING_SNAKE_CASE__      							=				{0: """batch""", 2: """past_sequence + sequence"""}
           SCREAMING_SNAKE_CASE__      							=				{0: """batch""", 2: """past_sequence + sequence"""}
     else:
       SCREAMING_SNAKE_CASE__      							=				OrderedDict(
           [
               ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
               ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
               ("""decoder_input_ids""", {0: """batch""", 1: """decoder_sequence"""}),
               ("""decoder_attention_mask""", {0: """batch""", 1: """decoder_sequence"""}),
           ]				)

     return common_inputs




   @property
   def      __a						(   self      :    int				):



     """simple docstring"""






     if self.task in ["default", "seq2seq-lm"]:
       SCREAMING_SNAKE_CASE__      							=				super().outputs
     else:
       SCREAMING_SNAKE_CASE__      							=				super(A__	,     self				).outputs
       if self.use_past:
         SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__      							=				self.num_layers
         for i in range(A__				):
           SCREAMING_SNAKE_CASE__      							=				{0: """batch""", 2: """past_sequence + sequence"""}
           SCREAMING_SNAKE_CASE__      							=				{0: """batch""", 2: """past_sequence + sequence"""}
     return common_outputs




   def      __a						(   self      :    Dict	,     _lowercase      :    int	,     _lowercase      :    int = -1	,     _lowercase      :    Tuple = -1	,     _lowercase      :    List[str] = False	,     _lowercase      :    Union[str, Any] = None	,     ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
         A__	,     A__	,     A__	,     A__	,     A__				)

     # Generate decoder inputs
     SCREAMING_SNAKE_CASE__      							=				seq_length if not self.use_past else 1
     SCREAMING_SNAKE_CASE__      							=				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
         A__	,     A__	,     A__	,     A__	,     A__				)
     SCREAMING_SNAKE_CASE__      							=				{f"""decoder_{name}""": tensor for name, tensor in decoder_inputs.items()}
     SCREAMING_SNAKE_CASE__      							=				dict(**A__	,     **A__				)

     if self.use_past:
       if not is_torch_available():
         raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""				)
       else:
         import torch
       SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__      							=				common_inputs["""input_ids"""].shape
       SCREAMING_SNAKE_CASE__      							=				common_inputs["""decoder_input_ids"""].shape[1]
       SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__      							=				self.num_attention_heads
       SCREAMING_SNAKE_CASE__      							=				(
           batch,
           num_encoder_attention_heads,
           encoder_seq_length,
           self._config.hidden_size // num_encoder_attention_heads,
       )
       SCREAMING_SNAKE_CASE__      							=				decoder_seq_length + 3
       SCREAMING_SNAKE_CASE__      							=				(
           batch,
           num_decoder_attention_heads,
           decoder_past_length,
           self._config.hidden_size // num_decoder_attention_heads,
       )

       SCREAMING_SNAKE_CASE__      							=				torch.cat(
           [common_inputs["""decoder_attention_mask"""], torch.ones(A__	,     A__				)]	,     dim=1				)

       SCREAMING_SNAKE_CASE__      							=				[]
       # If the number of encoder and decoder layers are present in the model configuration, both are considered
       SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__      							=				self.num_layers
       SCREAMING_SNAKE_CASE__      							=				min(A__	,     A__				)
       SCREAMING_SNAKE_CASE__      							=				max(A__	,     A__				) - min_num_layers
       SCREAMING_SNAKE_CASE__      							=				"""encoder""" if num_encoder_layers > num_decoder_layers else """decoder"""

       for _ in range(A__				):
         common_inputs["past_key_values"].append(
             (
                 torch.zeros(A__				),
                 torch.zeros(A__				),
                 torch.zeros(A__				),
                 torch.zeros(A__				),
             )				)
       # TODO: test this.
       SCREAMING_SNAKE_CASE__      							=				encoder_shape if remaining_side_name == """encoder""" else decoder_shape
       for _ in range(A__	,     A__				):
         common_inputs["past_key_values"].append((torch.zeros(A__				), torch.zeros(A__				))				)
     return common_inputs




   def      __a						(   self      :    List[Any]	,     _lowercase      :    Tuple	,     _lowercase      :    List[Any] = -1	,     _lowercase      :    int = -1	,     _lowercase      :    str = False	,     _lowercase      :    Any = None	,     ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
         A__	,     A__	,     A__	,     A__	,     A__				)

     if self.use_past:
       if not is_torch_available():
         raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""				)
       else:
         import torch
       SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__      							=				common_inputs["""input_ids"""].shape
       # Not using the same length for past_key_values
       SCREAMING_SNAKE_CASE__      							=				seqlen + 2
       SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__      							=				self.num_layers
       SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__      							=				self.num_attention_heads
       SCREAMING_SNAKE_CASE__      							=				(
           batch,
           num_encoder_attention_heads,
           past_key_values_length,
           self._config.hidden_size // num_encoder_attention_heads,
       )

       SCREAMING_SNAKE_CASE__      							=				common_inputs["""attention_mask"""].dtype
       SCREAMING_SNAKE_CASE__      							=				torch.cat(
           [common_inputs["""attention_mask"""], torch.ones(A__	,     A__	,     dtype=A__				)]	,     dim=1				)
       SCREAMING_SNAKE_CASE__      							=				[
           (torch.zeros(A__				), torch.zeros(A__				)) for _ in range(A__				)
       ]
     return common_inputs




   def      __a						(   self      :    List[Any]	,     _lowercase      :    str	,     _lowercase      :    Dict = -1	,     _lowercase      :    Union[str, Any] = -1	,     _lowercase      :    Union[str, Any] = False	,     _lowercase      :    Tuple = None	,     ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				compute_effective_axis_dimension(
         A__	,     fixed_dimension=OnnxConfig.default_fixed_batch	,     num_token_to_add=0				)

     # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
     SCREAMING_SNAKE_CASE__      							=				tokenizer.num_special_tokens_to_add(A__				)
     SCREAMING_SNAKE_CASE__      							=				compute_effective_axis_dimension(
         A__	,     fixed_dimension=OnnxConfig.default_fixed_sequence	,     num_token_to_add=A__				)

     # Generate dummy inputs according to compute batch and sequence
     SCREAMING_SNAKE_CASE__      							=				[""" """.join([tokenizer.unk_token]				) * seq_length] * batch_size
     SCREAMING_SNAKE_CASE__      							=				dict(tokenizer(A__	,     return_tensors=A__				)				)
     return common_inputs




   def      __a						(   self      :    Union[str, Any]	,     _lowercase      :    Optional[int]	,     _lowercase      :    List[str] = -1	,     _lowercase      :    Optional[Any] = -1	,     _lowercase      :    Tuple = False	,     _lowercase      :    List[Any] = None	,     ):



     """simple docstring"""






     if self.task in ["default", "seq2seq-lm"]:
       SCREAMING_SNAKE_CASE__      							=				self._generate_dummy_inputs_for_default_and_seqaseq_lm(
           A__	,     batch_size=A__	,     seq_length=A__	,     is_pair=A__	,     framework=A__				)

     elif self.task == "causal-lm":
       SCREAMING_SNAKE_CASE__      							=				self._generate_dummy_inputs_for_causal_lm(
           A__	,     batch_size=A__	,     seq_length=A__	,     is_pair=A__	,     framework=A__				)
     else:
       SCREAMING_SNAKE_CASE__      							=				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
           A__	,     batch_size=A__	,     seq_length=A__	,     is_pair=A__	,     framework=A__				)

     return common_inputs




   def      __a						(   self      :    Optional[int]	,     _lowercase      :    Any	,     _lowercase      :    Any	,     _lowercase      :    Optional[int]	,     _lowercase      :    Union[str, Any]				):



     """simple docstring"""






     if self.task in ["default", "seq2seq-lm"]:
       SCREAMING_SNAKE_CASE__      							=				super()._flatten_past_key_values_(A__	,     A__	,     A__	,     A__				)
     else:
       SCREAMING_SNAKE_CASE__      							=				super(A__	,     self				)._flatten_past_key_values_(
           A__	,     A__	,     A__	,     A__				)

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)
