

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 						_snake_case(   SCREAMING_SNAKE_CASE__      :     Any   ,      SCREAMING_SNAKE_CASE__      :     Union[str, Any]				)							->				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(lowerCamelCase_				)
    }

    # 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(lowerCamelCase_				)
    }

    for i in range(lowerCamelCase_				):
        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(lowerCamelCase_				)
        A__   				=       {'''old''': f'down.{i}.block', '''new''': f'down_blocks.{i}.resnets'}
        assign_to_checkpoint(lowerCamelCase_   ,      lowerCamelCase_   ,      lowerCamelCase_   ,      additional_replacements=[meta_path]   ,      config=lowerCamelCase_				)

    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(lowerCamelCase_				)
        A__   				=       {'''old''': f'mid.block_{i}', '''new''': f'mid_block.resnets.{i - 1}'}
        assign_to_checkpoint(lowerCamelCase_   ,      lowerCamelCase_   ,      lowerCamelCase_   ,      additional_replacements=[meta_path]   ,      config=lowerCamelCase_				)

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

    for i in range(lowerCamelCase_				):
        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(lowerCamelCase_				)
        A__   				=       {'''old''': f'up.{block_id}.block', '''new''': f'up_blocks.{i}.resnets'}
        assign_to_checkpoint(lowerCamelCase_   ,      lowerCamelCase_   ,      lowerCamelCase_   ,      additional_replacements=[meta_path]   ,      config=lowerCamelCase_				)

    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(lowerCamelCase_				)
        A__   				=       {'''old''': f'mid.block_{i}', '''new''': f'mid_block.resnets.{i - 1}'}
        assign_to_checkpoint(lowerCamelCase_   ,      lowerCamelCase_   ,      lowerCamelCase_   ,      additional_replacements=[meta_path]   ,      config=lowerCamelCase_				)

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





    '''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(lowerCamelCase_				)
    A__   				=       512
    A__   				=       '''cuda''' if torch.cuda.is_available() else '''cpu'''
    if checkpoint_path.endswith('safetensors'				):
        from safetensors import safe_open

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

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

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


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

     vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)





from ...utils import (
    OptionalDependencyNotAvailable,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)


try:
     if not (is_transformers_available() and is_torch_available()):
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     from ...utils.dummy_torch_and_transformers_objects import (
         ImageTextPipelineOutput,
         UniDiffuserPipeline,
     )
else:
     from .modeling_text_decoder import UniDiffuserTextDecoder
     from .modeling_uvit import UniDiffuserModel, UTransformeraDModel
     from .pipeline_unidiffuser import ImageTextPipelineOutput, UniDiffuserPipeline

