






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 					lowercase__				(					__snake_case   :							Any					,	__snake_case   :							Optional[int]  ):


					'''simple docstring'''


					UpperCAmelCase_   :				Optional[int]						=						checkpoint

					UpperCAmelCase_   :				Optional[int]						=						{}

					UpperCAmelCase_   :				Any						=						vae_state_dict['encoder.conv_in.weight']
					UpperCAmelCase_   :				Any						=						vae_state_dict['encoder.conv_in.bias']
					UpperCAmelCase_   :				str						=						vae_state_dict['encoder.conv_out.weight']
					UpperCAmelCase_   :				Union[str, Any]						=						vae_state_dict['encoder.conv_out.bias']
					UpperCAmelCase_   :				int						=						vae_state_dict['encoder.norm_out.weight']
					UpperCAmelCase_   :				List[str]						=						vae_state_dict['encoder.norm_out.bias']

					UpperCAmelCase_   :				int						=						vae_state_dict['decoder.conv_in.weight']
					UpperCAmelCase_   :				Tuple						=						vae_state_dict['decoder.conv_in.bias']
					UpperCAmelCase_   :				List[Any]						=						vae_state_dict['decoder.conv_out.weight']
					UpperCAmelCase_   :				List[str]						=						vae_state_dict['decoder.conv_out.bias']
					UpperCAmelCase_   :				Dict						=						vae_state_dict['decoder.norm_out.weight']
					UpperCAmelCase_   :				int						=						vae_state_dict['decoder.norm_out.bias']

					UpperCAmelCase_   :				Union[str, Any]						=						vae_state_dict['quant_conv.weight']
					UpperCAmelCase_   :				Dict						=						vae_state_dict['quant_conv.bias']
					UpperCAmelCase_   :				Union[str, Any]						=						vae_state_dict['post_quant_conv.weight']
					UpperCAmelCase_   :				Any						=						vae_state_dict['post_quant_conv.bias']

					# Retrieves the keys for the encoder down blocks only
					UpperCAmelCase_   :				Dict						=						len({'.'.join(layer.split('.'  )[:3]  ) for layer in vae_state_dict if 'encoder.down' in layer}  )
					UpperCAmelCase_   :				List[Any]						=						{
					    layer_id: [key for key in vae_state_dict if F"down.{layer_id}" in key] for layer_id in range(__snake_case  )
					}

					# Retrieves the keys for the decoder up blocks only
					UpperCAmelCase_   :				Union[str, Any]						=						len({'.'.join(layer.split('.'  )[:3]  ) for layer in vae_state_dict if 'decoder.up' in layer}  )
					UpperCAmelCase_   :				Any						=						{
					    layer_id: [key for key in vae_state_dict if F"up.{layer_id}" in key] for layer_id in range(__snake_case  )
					}

					for i in range(__snake_case  ):
										UpperCAmelCase_   :				List[str]						=						[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:
															UpperCAmelCase_   :				Optional[int]						=						vae_state_dict.pop(
															    F"encoder.down.{i}.downsample.conv.weight"  )
															UpperCAmelCase_   :				int						=						vae_state_dict.pop(
															    F"encoder.down.{i}.downsample.conv.bias"  )

										UpperCAmelCase_   :				Dict						=						renew_vae_resnet_paths(__snake_case  )
										UpperCAmelCase_   :				Union[str, Any]						=						{'old': F"down.{i}.block", 'new': F"down_blocks.{i}.resnets"}
										assign_to_checkpoint(__snake_case					,	__snake_case					,	__snake_case					,	additional_replacements=[meta_path]					,	config=__snake_case  )

					UpperCAmelCase_   :				List[str]						=						[key for key in vae_state_dict if 'encoder.mid.block' in key]
					UpperCAmelCase_   :				Any						=						2
					for i in range(1					,	num_mid_res_blocks + 1  ):
										UpperCAmelCase_   :				Union[str, Any]						=						[key for key in mid_resnets if F"encoder.mid.block_{i}" in key]

										UpperCAmelCase_   :				List[str]						=						renew_vae_resnet_paths(__snake_case  )
										UpperCAmelCase_   :				str						=						{'old': F"mid.block_{i}", 'new': F"mid_block.resnets.{i - 1}"}
										assign_to_checkpoint(__snake_case					,	__snake_case					,	__snake_case					,	additional_replacements=[meta_path]					,	config=__snake_case  )

					UpperCAmelCase_   :				int						=						[key for key in vae_state_dict if 'encoder.mid.attn' in key]
					UpperCAmelCase_   :				str						=						renew_vae_attention_paths(__snake_case  )
					UpperCAmelCase_   :				Dict						=						{'old': 'mid.attn_1', 'new': 'mid_block.attentions.0'}
					assign_to_checkpoint(__snake_case					,	__snake_case					,	__snake_case					,	additional_replacements=[meta_path]					,	config=__snake_case  )
					conv_attn_to_linear(__snake_case  )

					for i in range(__snake_case  ):
										UpperCAmelCase_   :				Dict						=						num_up_blocks - 1 - i
										UpperCAmelCase_   :				str						=						[
										    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:
															UpperCAmelCase_   :				Optional[int]						=						vae_state_dict[
															    F"decoder.up.{block_id}.upsample.conv.weight"
															]
															UpperCAmelCase_   :				int						=						vae_state_dict[
															    F"decoder.up.{block_id}.upsample.conv.bias"
															]

										UpperCAmelCase_   :				List[str]						=						renew_vae_resnet_paths(__snake_case  )
										UpperCAmelCase_   :				int						=						{'old': F"up.{block_id}.block", 'new': F"up_blocks.{i}.resnets"}
										assign_to_checkpoint(__snake_case					,	__snake_case					,	__snake_case					,	additional_replacements=[meta_path]					,	config=__snake_case  )

					UpperCAmelCase_   :				Optional[int]						=						[key for key in vae_state_dict if 'decoder.mid.block' in key]
					UpperCAmelCase_   :				Tuple						=						2
					for i in range(1					,	num_mid_res_blocks + 1  ):
										UpperCAmelCase_   :				List[Any]						=						[key for key in mid_resnets if F"decoder.mid.block_{i}" in key]

										UpperCAmelCase_   :				Tuple						=						renew_vae_resnet_paths(__snake_case  )
										UpperCAmelCase_   :				List[str]						=						{'old': F"mid.block_{i}", 'new': F"mid_block.resnets.{i - 1}"}
										assign_to_checkpoint(__snake_case					,	__snake_case					,	__snake_case					,	additional_replacements=[meta_path]					,	config=__snake_case  )

					UpperCAmelCase_   :				Optional[int]						=						[key for key in vae_state_dict if 'decoder.mid.attn' in key]
					UpperCAmelCase_   :				Dict						=						renew_vae_attention_paths(__snake_case  )
					UpperCAmelCase_   :				str						=						{'old': 'mid.attn_1', 'new': 'mid_block.attentions.0'}
					assign_to_checkpoint(__snake_case					,	__snake_case					,	__snake_case					,	additional_replacements=[meta_path]					,	config=__snake_case  )
					conv_attn_to_linear(__snake_case  )
					return new_checkpoint




def 					lowercase__				(					__snake_case   :							str					,	__snake_case   :							str					,	):


					'''simple docstring'''


					UpperCAmelCase_   :				Optional[int]						=						requests.get(
					    ' https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml'  )
					UpperCAmelCase_   :				Tuple						=						io.BytesIO(r.content  )

					UpperCAmelCase_   :				Dict						=						OmegaConf.load(__snake_case  )
					UpperCAmelCase_   :				int						=						512
					UpperCAmelCase_   :				Optional[int]						=						'cuda' if torch.cuda.is_available() else 'cpu'
					if checkpoint_path.endswith('safetensors'  ):
										from safetensors import safe_open

										UpperCAmelCase_   :				Optional[Any]						=						{}
										with safe_open(__snake_case					,	framework='pt'					,	device='cpu'  ) as f:
															for key in f.keys():
																				UpperCAmelCase_   :				str						=						f.get_tensor(__snake_case  )
					else:
										UpperCAmelCase_   :				Optional[int]						=						torch.load(__snake_case					,	map_location=__snake_case  )['state_dict']

					# Convert the VAE model.
					UpperCAmelCase_   :				str						=						create_vae_diffusers_config(__snake_case					,	image_size=__snake_case  )
					UpperCAmelCase_   :				Any						=						custom_convert_ldm_vae_checkpoint(__snake_case					,	__snake_case  )

					UpperCAmelCase_   :				Union[str, Any]						=						AutoencoderKL(**__snake_case  )
					vae.load_state_dict(__snake_case  )
					vae.save_pretrained(__snake_case  )


if __name__ == "__main__":
		__UpperCAmelCase     							=						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.')

		__UpperCAmelCase     							=						parser.parse_args()

		vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)














from collections import defaultdict
from math import ceil, sqrt




def 					lowercase__				(					__snake_case   :							int = 1_000_000					,	__snake_case   :							int = 10  ):


					'''simple docstring'''


					UpperCAmelCase_   :				defaultdict						=						defaultdict(__snake_case  )

					for outer_width in range(3					,	(t_limit // 4) + 2  ):
										if outer_width * outer_width > t_limit:
															UpperCAmelCase_   :				Union[str, Any]						=						max(
															    ceil(sqrt(outer_width * outer_width - t_limit  )  )					,	1  )
										else:
															UpperCAmelCase_   :				int						=						1

										hole_width_lower_bound += (outer_width - hole_width_lower_bound) % 2

										for hole_width in range(__snake_case					,	outer_width - 1					,	2  ):
															count[outer_width * outer_width - hole_width * hole_width] += 1

					return sum(1 for n in count.values() if 1 <= n <= 10  )


if __name__ == "__main__":
		print(F'{solution() = }')






