


'''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 			_lowercase       (				UpperCamelCase__       :					Optional[Any],				UpperCamelCase__       :					Dict			):
 __A  :		Optional[Any]       =			checkpoint

 __A  :		Any       =			{}

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

 __A  :		Optional[int]       =			vae_state_dict['decoder.conv_in.weight']
 __A  :		Dict       =			vae_state_dict['decoder.conv_in.bias']
 __A  :		int       =			vae_state_dict['decoder.conv_out.weight']
 __A  :		str       =			vae_state_dict['decoder.conv_out.bias']
 __A  :		Union[str, Any]       =			vae_state_dict['decoder.norm_out.weight']
 __A  :		str       =			vae_state_dict['decoder.norm_out.bias']

 __A  :		str       =			vae_state_dict['quant_conv.weight']
 __A  :		str       =			vae_state_dict['quant_conv.bias']
 __A  :		List[str]       =			vae_state_dict['post_quant_conv.weight']
 __A  :		List[str]       =			vae_state_dict['post_quant_conv.bias']

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

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

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

  __A  :		Union[str, Any]       =			renew_vae_resnet_paths(UpperCamelCase__			)
  __A  :		Tuple       =			{'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  :		List[str]       =			[key for key in vae_state_dict if 'encoder.mid.attn' in key]
 __A  :		List[str]       =			renew_vae_attention_paths(UpperCamelCase__			)
 __A  :		List[str]       =			{'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  :		Dict       =			num_up_blocks - 1 - i
  __A  :		Optional[Any]       =			[
      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  :		Union[str, Any]       =			vae_state_dict[
       f"""decoder.up.{block_id}.upsample.conv.weight"""
   ]
   __A  :		Optional[Any]       =			vae_state_dict[
       f"""decoder.up.{block_id}.upsample.conv.bias"""
   ]

  __A  :		List[str]       =			renew_vae_resnet_paths(UpperCamelCase__			)
  __A  :		List[str]       =			{'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  :		Optional[int]       =			[key for key in vae_state_dict if 'decoder.mid.block' in key]
 __A  :		Any       =			2
 for i in range(1,				num_mid_res_blocks + 1			):
  __A  :		Optional[Any]       =			[key for key in mid_resnets if f"""decoder.mid.block_{i}""" in key]

  __A  :		List[str]       =			renew_vae_resnet_paths(UpperCamelCase__			)
  __A  :		List[str]       =			{'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  :		int       =			[key for key in vae_state_dict if 'decoder.mid.attn' in key]
 __A  :		Dict       =			renew_vae_attention_paths(UpperCamelCase__			)
 __A  :		Tuple       =			{'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 			_lowercase       (				UpperCamelCase__       :					str,				UpperCamelCase__       :					str,				):
 # Only support V1
 __A  :		Tuple       =			requests.get(
     ' https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml'			)
 __A  :		Optional[int]       =			io.BytesIO(r.content			)

 __A  :		str       =			OmegaConf.load(UpperCamelCase__			)
 __A  :		Optional[Any]       =			512
 __A  :		str       =			'cuda' if torch.cuda.is_available() else 'cpu'
 if checkpoint_path.endswith('safetensors'			):
  from safetensors import safe_open

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

 # Convert the VAE model.
 __A  :		int       =			create_vae_diffusers_config(UpperCamelCase__,				image_size=UpperCamelCase__			)
 __A  :		Tuple       =			custom_convert_ldm_vae_checkpoint(UpperCamelCase__,				UpperCamelCase__			)

 __A  :		int       =			AutoencoderKL(**UpperCamelCase__			)
 vae.load_state_dict(UpperCamelCase__			)
 vae.save_pretrained(UpperCamelCase__			)


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

 UpperCAmelCase_ :     List[Any]								=      parser.parse_args()

 vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)










'''simple docstring'''
from collections import defaultdict
from math import ceil, sqrt


def 			_lowercase       (				UpperCamelCase__       :					int = 1000000,				UpperCamelCase__       :					int = 10			):
 __A  :		defaultdict       =			defaultdict(UpperCamelCase__			)

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

  hole_width_lower_bound += (outer_width - hole_width_lower_bound) % 2

  for hole_width in range(UpperCamelCase__,				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() = }''')






