


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__:  Optional[int]				,       UpperCamelCase__:  Dict    )			-> Tuple:





	"""simple docstring"""







	A             =					checkpoint

	A             =					{}

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

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

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

	# Retrieves the keys for the encoder down blocks only
	A             =					len({""".""".join(layer.split("""."""    )[:3]    ) for layer in vae_state_dict if """encoder.down""" in layer}    )
	A             =					{
	    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
	A             =					len({""".""".join(layer.split("""."""    )[:3]    ) for layer in vae_state_dict if """decoder.up""" in layer}    )
	A             =					{
	    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__    ):
		A             =					[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:
			A             =					vae_state_dict.pop(
			    f'encoder.down.{i}.downsample.conv.weight'    )
			A             =					vae_state_dict.pop(
			    f'encoder.down.{i}.downsample.conv.bias'    )

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

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

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

	A             =					[key for key in vae_state_dict if """encoder.mid.attn""" in key]
	A             =					renew_vae_attention_paths(UpperCamelCase__    )
	A             =					{"""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__    ):
		A             =					num_up_blocks - 1 - i
		A             =					[
		    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:
			A             =					vae_state_dict[
			    f'decoder.up.{block_id}.upsample.conv.weight'
			]
			A             =					vae_state_dict[
			    f'decoder.up.{block_id}.upsample.conv.bias'
			]

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

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

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

	A             =					[key for key in vae_state_dict if """decoder.mid.attn""" in key]
	A             =					renew_vae_attention_paths(UpperCamelCase__    )
	A             =					{"""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				,       )			-> Dict:





	"""simple docstring"""







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

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

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

	# Convert the VAE model.
	A             =					create_vae_diffusers_config(UpperCamelCase__				,       image_size=UpperCamelCase__    )
	A             =					custom_convert_ldm_vae_checkpoint(UpperCamelCase__				,       UpperCamelCase__    )

	A             =					AutoencoderKL(**UpperCamelCase__    )
	vae.load_state_dict(UpperCamelCase__    )
	vae.save_pretrained(UpperCamelCase__    )


if __name__ == "__main__":
				_lowercase							:   Optional[int]						=						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.")

				_lowercase							:   str						=						parser.parse_args()

				vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)







from transformers import HfArgumentParser, TensorFlowBenchmark, TensorFlowBenchmarkArguments





def 	_lowerCAmelCase					( )			-> str:





	"""simple docstring"""







	A             =					HfArgumentParser(UpperCamelCase__    )
	A             =					parser.parse_args_into_dataclasses()[0]
	A             =					TensorFlowBenchmark(args=UpperCamelCase__    )
	try:
		A             =					parser.parse_args_into_dataclasses()[0]
	except ValueError as e:
		A             =					"""Arg --no_{0} is no longer used, please use --no-{0} instead."""
		A             =					""" """.join(str(UpperCamelCase__    ).split(""" """    )[:-1]    )
		A             =					""""""
		A             =					eval(str(UpperCamelCase__    ).split(""" """    )[-1]    )
		A             =					[]
		for arg in depreciated_args:
			# arg[2:] removes '--'
			if arg[2:] in TensorFlowBenchmark.deprecated_args:
				# arg[5:] removes '--no_'
				full_error_msg += arg_error_msg.format(arg[5:]    )
			else:
				wrong_args.append(UpperCamelCase__    )
		if len(UpperCamelCase__    ) > 0:
			A             =					full_error_msg + begin_error_msg + str(UpperCamelCase__    )
		raise ValueError(UpperCamelCase__    )
	benchmark.run()


if __name__ == "__main__":
				main()



