




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			(							lowercase	:       Dict	,							lowercase	:       Optional[int]     ):


      '''simple docstring'''



      lowerCamelCase_		     =    checkpoint

      lowerCamelCase_		     =    {}

      lowerCamelCase_		     =    vae_state_dict['encoder.conv_in.weight']
      lowerCamelCase_		     =    vae_state_dict['encoder.conv_in.bias']
      lowerCamelCase_		     =    vae_state_dict['encoder.conv_out.weight']
      lowerCamelCase_		     =    vae_state_dict['encoder.conv_out.bias']
      lowerCamelCase_		     =    vae_state_dict['encoder.norm_out.weight']
      lowerCamelCase_		     =    vae_state_dict['encoder.norm_out.bias']

      lowerCamelCase_		     =    vae_state_dict['decoder.conv_in.weight']
      lowerCamelCase_		     =    vae_state_dict['decoder.conv_in.bias']
      lowerCamelCase_		     =    vae_state_dict['decoder.conv_out.weight']
      lowerCamelCase_		     =    vae_state_dict['decoder.conv_out.bias']
      lowerCamelCase_		     =    vae_state_dict['decoder.norm_out.weight']
      lowerCamelCase_		     =    vae_state_dict['decoder.norm_out.bias']

      lowerCamelCase_		     =    vae_state_dict['quant_conv.weight']
      lowerCamelCase_		     =    vae_state_dict['quant_conv.bias']
      lowerCamelCase_		     =    vae_state_dict['post_quant_conv.weight']
      lowerCamelCase_		     =    vae_state_dict['post_quant_conv.bias']

      # Retrieves the keys for the encoder down blocks only
      lowerCamelCase_		     =    len({'.'.join(layer.split('.'     )[:3]     ) for layer in vae_state_dict if 'encoder.down' in layer}     )
      lowerCamelCase_		     =    {
          layer_id: [key for key in vae_state_dict if f"""down.{layer_id}""" in key] for layer_id in range(lowercase     )
      }

      # Retrieves the keys for the decoder up blocks only
      lowerCamelCase_		     =    len({'.'.join(layer.split('.'     )[:3]     ) for layer in vae_state_dict if 'decoder.up' in layer}     )
      lowerCamelCase_		     =    {
          layer_id: [key for key in vae_state_dict if f"""up.{layer_id}""" in key] for layer_id in range(lowercase     )
      }

      for i in range(lowercase     ):
            lowerCamelCase_		     =    [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:
                  lowerCamelCase_		     =    vae_state_dict.pop(
                      f"""encoder.down.{i}.downsample.conv.weight"""     )
                  lowerCamelCase_		     =    vae_state_dict.pop(
                      f"""encoder.down.{i}.downsample.conv.bias"""     )

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

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

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

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

      for i in range(lowercase     ):
            lowerCamelCase_		     =    num_up_blocks - 1 - i
            lowerCamelCase_		     =    [
                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:
                  lowerCamelCase_		     =    vae_state_dict[
                      f"""decoder.up.{block_id}.upsample.conv.weight"""
                  ]
                  lowerCamelCase_		     =    vae_state_dict[
                      f"""decoder.up.{block_id}.upsample.conv.bias"""
                  ]

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

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

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

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

def 			_SCREAMING_SNAKE_CASE			(							lowercase	:       str	,							lowercase	:       str	,							):


      '''simple docstring'''



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

      lowerCamelCase_		     =    OmegaConf.load(lowercase     )
      lowerCamelCase_		     =    5_12
      lowerCamelCase_		     =    'cuda' if torch.cuda.is_available() else 'cpu'
      if checkpoint_path.endswith('safetensors'     ):
            from safetensors import safe_open

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

      # Convert the VAE model.
      lowerCamelCase_		     =    create_vae_diffusers_config(lowercase	,							image_size=lowercase     )
      lowerCamelCase_		     =    custom_convert_ldm_vae_checkpoint(lowercase	,							lowercase     )

      lowerCamelCase_		     =    AutoencoderKL(**lowercase     )
      vae.load_state_dict(lowercase     )
      vae.save_pretrained(lowercase     )


if __name__ == "__main__":
 lowerCamelCase				:		str     =   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				:		str     =   parser.parse_args()

 vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)









def 			_SCREAMING_SNAKE_CASE			(							lowercase	:       int = 10_00     ):


      '''simple docstring'''



      lowerCamelCase_		     =    2**power
      lowerCamelCase_		     =    0
      while n:
            lowerCamelCase_     ,				lowerCamelCase_		     =    r + n % 10, n // 10
      return r


if __name__ == "__main__":
 print(solution(int(str(input()).strip())))



