

'''simple docstring'''
lowerCamelCase			:				List[Any] 			=					{
    "Pillow": "Pillow<10.0.0",
    "accelerate": "accelerate>=0.20.3",
    "av": "av==9.2.0",
    "beautifulsoup4": "beautifulsoup4",
    "black": "black~=23.1",
    "codecarbon": "codecarbon==1.2.0",
    "cookiecutter": "cookiecutter==1.7.3",
    "dataclasses": "dataclasses",
    "datasets": "datasets!=2.5.0",
    "decord": "decord==0.6.0",
    "deepspeed": "deepspeed>=0.9.3",
    "diffusers": "diffusers",
    "dill": "dill<0.3.5",
    "evaluate": "evaluate>=0.2.0",
    "fairscale": "fairscale>0.3",
    "faiss-cpu": "faiss-cpu",
    "fastapi": "fastapi",
    "filelock": "filelock",
    "flax": "flax>=0.4.1,<=0.7.0",
    "ftfy": "ftfy",
    "fugashi": "fugashi>=1.0",
    "GitPython": "GitPython<3.1.19",
    "hf-doc-builder": "hf-doc-builder>=0.3.0",
    "huggingface-hub": "huggingface-hub>=0.14.1,<1.0",
    "importlib_metadata": "importlib_metadata",
    "ipadic": "ipadic>=1.0.0,<2.0",
    "isort": "isort>=5.5.4",
    "jax": "jax>=0.2.8,!=0.3.2,<=0.4.13",
    "jaxlib": "jaxlib>=0.1.65,<=0.4.13",
    "jieba": "jieba",
    "kenlm": "kenlm",
    "keras-nlp": "keras-nlp>=0.3.1",
    "librosa": "librosa",
    "nltk": "nltk",
    "natten": "natten>=0.14.6",
    "numpy": "numpy>=1.17",
    "onnxconverter-common": "onnxconverter-common",
    "onnxruntime-tools": "onnxruntime-tools>=1.4.2",
    "onnxruntime": "onnxruntime>=1.4.0",
    "opencv-python": "opencv-python",
    "optuna": "optuna",
    "optax": "optax>=0.0.8,<=0.1.4",
    "packaging": "packaging>=20.0",
    "parameterized": "parameterized",
    "phonemizer": "phonemizer",
    "protobuf": "protobuf",
    "psutil": "psutil",
    "pyyaml": "pyyaml>=5.1",
    "pydantic": "pydantic<2",
    "pytest": "pytest>=7.2.0",
    "pytest-timeout": "pytest-timeout",
    "pytest-xdist": "pytest-xdist",
    "python": "python>=3.8.0",
    "ray[tune]": "ray[tune]",
    "regex": "regex!=2019.12.17",
    "requests": "requests",
    "rhoknp": "rhoknp>=1.1.0,<1.3.1",
    "rjieba": "rjieba",
    "rouge-score": "rouge-score!=0.0.7,!=0.0.8,!=0.1,!=0.1.1",
    "ruff": "ruff>=0.0.241,<=0.0.259",
    "sacrebleu": "sacrebleu>=1.4.12,<2.0.0",
    "sacremoses": "sacremoses",
    "safetensors": "safetensors>=0.3.1",
    "sagemaker": "sagemaker>=2.31.0",
    "scikit-learn": "scikit-learn",
    "sentencepiece": "sentencepiece>=0.1.91,!=0.1.92",
    "sigopt": "sigopt",
    "starlette": "starlette",
    "sudachipy": "sudachipy>=0.6.6",
    "sudachidict_core": "sudachidict_core>=20220729",
    "tensorflow-cpu": "tensorflow-cpu>=2.6,<2.14",
    "tensorflow": "tensorflow>=2.6,<2.14",
    "tensorflow-text": "tensorflow-text<2.14",
    "tf2onnx": "tf2onnx",
    "timeout-decorator": "timeout-decorator",
    "timm": "timm",
    "tokenizers": "tokenizers>=0.11.1,!=0.11.3,<0.14",
    "torch": "torch>=1.9,!=1.12.0",
    "torchaudio": "torchaudio",
    "torchvision": "torchvision",
    "pyctcdecode": "pyctcdecode>=0.4.0",
    "tqdm": "tqdm>=4.27",
    "unidic": "unidic>=1.0.2",
    "unidic_lite": "unidic_lite>=1.0.7",
    "urllib3": "urllib3<2.0.0",
    "uvicorn": "uvicorn",
}



'''simple docstring'''
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 	_lowerCAmelCase (						_UpperCamelCase						:					Tuple							,			_UpperCamelCase						:					List[Any]    )				->      Dict:
     """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 	_lowerCAmelCase (						_UpperCamelCase						:					str							,			_UpperCamelCase						:					str							,			)				->      List[Any]:
     """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			:				Union[str, 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			:				List[str] 			=					parser.parse_args()

     vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)
