import colorsys

from PIL import Image  # type: ignore



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	float			,					__UpperCamelCase   :	float			,					__UpperCamelCase   :	int     )							-> float:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__      							=				x
  SCREAMING_SNAKE_CASE__      							=				y
  for step in range(__UpperCamelCase     ):  # noqa: B007
    SCREAMING_SNAKE_CASE__      							=				a * a - b * b + x
    SCREAMING_SNAKE_CASE__      							=				2 * a * b + y
    SCREAMING_SNAKE_CASE__      							=				a_new

    # divergence happens for all complex number with an absolute value
    # greater than 4
    if a * a + b * b > 4:
      break
  return step / (max_step - 1)



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	float     )							-> tuple:







  """simple docstring"""





  if distance == 1:
    return (0, 0, 0)
  else:
    return (2_55, 2_55, 2_55)



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	float     )							-> tuple:







  """simple docstring"""





  if distance == 1:
    return (0, 0, 0)
  else:
    return tuple(round(i * 2_55     ) for i in colorsys.hsv_to_rgb(__UpperCamelCase			,					1			,					1     )     )



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	int = 8_00			,					__UpperCamelCase   :	int = 6_00			,					__UpperCamelCase   :	float = -0.6			,					__UpperCamelCase   :	float = 0			,					__UpperCamelCase   :	float = 3.2			,					__UpperCamelCase   :	int = 50			,					__UpperCamelCase   :	bool = True			,					)							-> Image.Image:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__      							=				Image.new("""RGB"""			,					(image_width, image_height)     )
  SCREAMING_SNAKE_CASE__      							=				img.load()

  # loop through the image-coordinates
  for image_x in range(__UpperCamelCase     ):
    for image_y in range(__UpperCamelCase     ):
      # determine the figure-coordinates based on the image-coordinates
      SCREAMING_SNAKE_CASE__      							=				figure_width / image_width * image_height
      SCREAMING_SNAKE_CASE__      							=				figure_center_x + (image_x / image_width - 0.5) * figure_width
      SCREAMING_SNAKE_CASE__      							=				figure_center_y + (image_y / image_height - 0.5) * figure_height

      SCREAMING_SNAKE_CASE__      							=				get_distance(__UpperCamelCase			,					__UpperCamelCase			,					__UpperCamelCase     )

      # color the corresponding pixel based on the selected coloring-function
      if use_distance_color_coding:
        SCREAMING_SNAKE_CASE__      							=				get_color_coded_rgb(__UpperCamelCase     )
      else:
        SCREAMING_SNAKE_CASE__      							=				get_black_and_white_rgb(__UpperCamelCase     )

  return img


if __name__ == "__main__":
   import doctest

   doctest.testmod()

   # colored version, full figure
   __lowerCamelCase      :			int							       =						get_image()

   # uncomment for colored version, different section, zoomed in
   # img = get_image(figure_center_x = -0.6, figure_center_y = -0.4,
   # figure_width = 0.8)

   # uncomment for black and white version, full figure
   # img = get_image(use_distance_color_coding = False)

   # uncomment to save the image
   # img.save("mandelbrot.png")

   img.show()

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  ( __UpperCamelCase   :	Optional[int]			,					__UpperCamelCase   :	str     )							-> Tuple:







  """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      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	str			,					__UpperCamelCase   :	str			,					)							-> Optional[int]:







  """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      :			List[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      :			Optional[int]							       =		parser.parse_args()

   vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)
