from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_DEFAULT_MEAN,
    IMAGENET_DEFAULT_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
   import PIL


__lowerCamelCase      :			Any							       =						logging.get_logger(__name__)







class __snake_case    (						SCREAMING_SNAKE_CASE__		):
   lowerCAmelCase_														=							['pixel_values']




   def __init__(   self      :    Optional[Any]	,     _lowercase      :    bool = True	,     _lowercase      :    Dict[str, int] = None	,     _lowercase      :    int = 0.9	,     _lowercase      :    PILImageResampling = PILImageResampling.BICUBIC	,     _lowercase      :    bool = True	,     _lowercase      :    Dict[str, int] = None	,     _lowercase      :    Union[int, float] = 1 / 2_55	,     _lowercase      :    bool = True	,     _lowercase      :    bool = True	,     _lowercase      :    Optional[Union[float, List[float]]] = None	,     _lowercase      :    Optional[Union[float, List[float]]] = None	,     **_lowercase      :    Any	,     ):



     """simple docstring"""






     super().__init__(**a_				)
     SCREAMING_SNAKE_CASE__      							=				size if size is not None else {"""shortest_edge""": 2_24}
     SCREAMING_SNAKE_CASE__      							=				get_size_dict(a_	,     default_to_square=a_				)
     SCREAMING_SNAKE_CASE__      							=				crop_size if crop_size is not None else {"""height""": 2_24, """width""": 2_24}
     SCREAMING_SNAKE_CASE__      							=				get_size_dict(a_	,     param_name="""crop_size"""				)

     SCREAMING_SNAKE_CASE__      							=				do_resize
     SCREAMING_SNAKE_CASE__      							=				size
     SCREAMING_SNAKE_CASE__      							=				crop_pct
     SCREAMING_SNAKE_CASE__      							=				resample
     SCREAMING_SNAKE_CASE__      							=				do_center_crop
     SCREAMING_SNAKE_CASE__      							=				crop_size
     SCREAMING_SNAKE_CASE__      							=				do_rescale
     SCREAMING_SNAKE_CASE__      							=				rescale_factor
     SCREAMING_SNAKE_CASE__      							=				do_normalize
     SCREAMING_SNAKE_CASE__      							=				image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
     SCREAMING_SNAKE_CASE__      							=				image_std if image_std is not None else IMAGENET_DEFAULT_STD




   def      __a						(   self      :    Any	,     _lowercase      :    np.ndarray	,     _lowercase      :    Dict[str, int]	,     _lowercase      :    Optional[float] = None	,     _lowercase      :    PILImageResampling = PILImageResampling.BICUBIC	,     _lowercase      :    Optional[Union[str, ChannelDimension]] = None	,     **_lowercase      :    Dict	,     ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				get_size_dict(a_	,     default_to_square=a_				)
     if "shortest_edge" not in size and ("height" not in size or "width" not in size):
       raise ValueError(f"""size must contain \'height\' and \'width\' or \'shortest_edge\' as keys. Got {size.keys()}"""				)
     if crop_pct is not None:
       if "shortest_edge" in size:
         SCREAMING_SNAKE_CASE__      							=				int(size["""shortest_edge"""] / crop_pct				)
       elif "height" in size and "width" in size:
         if size["height"] == size["width"]:
           SCREAMING_SNAKE_CASE__      							=				int(size["""height"""] / crop_pct				)
         else:
           SCREAMING_SNAKE_CASE__      							=				(int(size["""height"""] / crop_pct				), int(size["""width"""] / crop_pct				))
       else:
         raise ValueError("""Invalid size for resize: {}""".format(a_				)				)

       SCREAMING_SNAKE_CASE__      							=				get_resize_output_image_size(a_	,     size=a_	,     default_to_square=a_				)
     else:
       if "shortest_edge" in size:
         SCREAMING_SNAKE_CASE__      							=				get_resize_output_image_size(a_	,     size=size["""shortest_edge"""]	,     default_to_square=a_				)
       elif "height" in size and "width" in size:
         SCREAMING_SNAKE_CASE__      							=				(size["""height"""], size["""width"""])
       else:
         raise ValueError("""Invalid size for resize: {}""".format(a_				)				)

     return resize(a_	,     size=a_	,     resample=a_	,     data_format=a_	,     **a_				)




   def      __a						(   self      :    Any	,     _lowercase      :    np.ndarray	,     _lowercase      :    Dict[str, int]	,     _lowercase      :    Optional[Union[str, ChannelDimension]] = None	,     **_lowercase      :    List[str]	,     ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				get_size_dict(a_				)
     if "height" not in size or "width" not in size:
       raise ValueError(f"""size must contain \'height\' and \'width\' as keys. Got {size.keys()}"""				)
     return center_crop(a_	,     size=(size["""height"""], size["""width"""])	,     data_format=a_	,     **a_				)




   def      __a						(   self      :    List[Any]	,     _lowercase      :    np.ndarray	,     _lowercase      :    Union[int, float]	,     _lowercase      :    Optional[Union[str, ChannelDimension]] = None	,     **_lowercase      :    int	,     ):



     """simple docstring"""






     return rescale(a_	,     scale=a_	,     data_format=a_	,     **a_				)




   def      __a						(   self      :    Dict	,     _lowercase      :    np.ndarray	,     _lowercase      :    Union[float, List[float]]	,     _lowercase      :    Union[float, List[float]]	,     _lowercase      :    Optional[Union[str, ChannelDimension]] = None	,     **_lowercase      :    Tuple	,     ):



     """simple docstring"""






     return normalize(a_	,     mean=a_	,     std=a_	,     data_format=a_	,     **a_				)




   def      __a						(   self      :    str	,     _lowercase      :    ImageInput	,     _lowercase      :    bool = None	,     _lowercase      :    Dict[str, int] = None	,     _lowercase      :    int = None	,     _lowercase      :    PILImageResampling = None	,     _lowercase      :    bool = None	,     _lowercase      :    Dict[str, int] = None	,     _lowercase      :    bool = None	,     _lowercase      :    float = None	,     _lowercase      :    bool = None	,     _lowercase      :    Optional[Union[float, List[float]]] = None	,     _lowercase      :    Optional[Union[float, List[float]]] = None	,     _lowercase      :    Optional[Union[str, TensorType]] = None	,     _lowercase      :    ChannelDimension = ChannelDimension.FIRST	,     **_lowercase      :    str	,     ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				do_resize if do_resize is not None else self.do_resize
     SCREAMING_SNAKE_CASE__      							=				crop_pct if crop_pct is not None else self.crop_pct
     SCREAMING_SNAKE_CASE__      							=				resample if resample is not None else self.resample
     SCREAMING_SNAKE_CASE__      							=				do_center_crop if do_center_crop is not None else self.do_center_crop
     SCREAMING_SNAKE_CASE__      							=				do_rescale if do_rescale is not None else self.do_rescale
     SCREAMING_SNAKE_CASE__      							=				rescale_factor if rescale_factor is not None else self.rescale_factor
     SCREAMING_SNAKE_CASE__      							=				do_normalize if do_normalize is not None else self.do_normalize
     SCREAMING_SNAKE_CASE__      							=				image_mean if image_mean is not None else self.image_mean
     SCREAMING_SNAKE_CASE__      							=				image_std if image_std is not None else self.image_std

     SCREAMING_SNAKE_CASE__      							=				size if size is not None else self.size
     SCREAMING_SNAKE_CASE__      							=				get_size_dict(a_	,     default_to_square=a_				)
     SCREAMING_SNAKE_CASE__      							=				crop_size if crop_size is not None else self.crop_size
     SCREAMING_SNAKE_CASE__      							=				get_size_dict(a_	,     param_name="""crop_size"""				)

     SCREAMING_SNAKE_CASE__      							=				make_list_of_images(a_				)

     if not valid_images(a_				):
       raise ValueError(
           """Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, """
           """torch.Tensor, tf.Tensor or jax.ndarray."""				)

     if do_resize and size is None or resample is None:
       raise ValueError("""Size and resample must be specified if do_resize is True."""				)

     if do_center_crop and crop_pct is None:
       raise ValueError("""Crop_pct must be specified if do_center_crop is True."""				)

     if do_rescale and rescale_factor is None:
       raise ValueError("""Rescale factor must be specified if do_rescale is True."""				)

     if do_normalize and (image_mean is None or image_std is None):
       raise ValueError("""Image mean and std must be specified if do_normalize is True."""				)

     # All transformations expect numpy arrays.
     SCREAMING_SNAKE_CASE__      							=				[to_numpy_array(a_				) for image in images]

     if do_resize:
       SCREAMING_SNAKE_CASE__      							=				[self.resize(image=a_	,     size=a_	,     crop_pct=a_	,     resample=a_				) for image in images]

     if do_center_crop:
       SCREAMING_SNAKE_CASE__      							=				[self.center_crop(image=a_	,     size=a_				) for image in images]

     if do_rescale:
       SCREAMING_SNAKE_CASE__      							=				[self.rescale(image=a_	,     scale=a_				) for image in images]

     if do_normalize:
       SCREAMING_SNAKE_CASE__      							=				[self.normalize(image=a_	,     mean=a_	,     std=a_				) for image in images]

     SCREAMING_SNAKE_CASE__      							=				[to_channel_dimension_format(a_	,     a_				) for image in images]

     SCREAMING_SNAKE_CASE__      							=				{"""pixel_values""": images}
     return BatchFeature(data=a_	,     tensor_type=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)
