

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       A_     (							a    ,      a     ):

      """simple docstring"""



      SCREAMING_SNAKE_CASE_ :     Optional[Any]							=   checkpoint

      SCREAMING_SNAKE_CASE_ :     List[Any]							=   {}

      SCREAMING_SNAKE_CASE_ :     int							=   vae_state_dict['encoder.conv_in.weight']
      SCREAMING_SNAKE_CASE_ :     Optional[int]							=   vae_state_dict['encoder.conv_in.bias']
      SCREAMING_SNAKE_CASE_ :     str							=   vae_state_dict['encoder.conv_out.weight']
      SCREAMING_SNAKE_CASE_ :     int							=   vae_state_dict['encoder.conv_out.bias']
      SCREAMING_SNAKE_CASE_ :     Tuple							=   vae_state_dict['encoder.norm_out.weight']
      SCREAMING_SNAKE_CASE_ :     str							=   vae_state_dict['encoder.norm_out.bias']

      SCREAMING_SNAKE_CASE_ :     Any							=   vae_state_dict['decoder.conv_in.weight']
      SCREAMING_SNAKE_CASE_ :     Optional[int]							=   vae_state_dict['decoder.conv_in.bias']
      SCREAMING_SNAKE_CASE_ :     str							=   vae_state_dict['decoder.conv_out.weight']
      SCREAMING_SNAKE_CASE_ :     List[Any]							=   vae_state_dict['decoder.conv_out.bias']
      SCREAMING_SNAKE_CASE_ :     str							=   vae_state_dict['decoder.norm_out.weight']
      SCREAMING_SNAKE_CASE_ :     str							=   vae_state_dict['decoder.norm_out.bias']

      SCREAMING_SNAKE_CASE_ :     int							=   vae_state_dict['quant_conv.weight']
      SCREAMING_SNAKE_CASE_ :     Union[str, Any]							=   vae_state_dict['quant_conv.bias']
      SCREAMING_SNAKE_CASE_ :     Union[str, Any]							=   vae_state_dict['post_quant_conv.weight']
      SCREAMING_SNAKE_CASE_ :     Union[str, Any]							=   vae_state_dict['post_quant_conv.bias']

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

            SCREAMING_SNAKE_CASE_ :     Optional[Any]							=   renew_vae_resnet_paths(lowerCAmelCase__     )
            SCREAMING_SNAKE_CASE_ :     List[str]							=   {'old': f"down.{i}.block", 'new': f"down_blocks.{i}.resnets"}
            assign_to_checkpoint(lowerCAmelCase__    ,      lowerCAmelCase__    ,      lowerCAmelCase__    ,      additional_replacements=[meta_path]    ,      config=lowerCAmelCase__     )

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

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

      SCREAMING_SNAKE_CASE_ :     Optional[int]							=   [key for key in vae_state_dict if 'encoder.mid.attn' in key]
      SCREAMING_SNAKE_CASE_ :     Any							=   renew_vae_attention_paths(lowerCAmelCase__     )
      SCREAMING_SNAKE_CASE_ :     Tuple							=   {'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__     ):
            SCREAMING_SNAKE_CASE_ :     int							=   num_up_blocks - 1 - i
            SCREAMING_SNAKE_CASE_ :     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:
                  SCREAMING_SNAKE_CASE_ :     int							=   vae_state_dict[
                      f"decoder.up.{block_id}.upsample.conv.weight"
                  ]
                  SCREAMING_SNAKE_CASE_ :     List[Any]							=   vae_state_dict[
                      f"decoder.up.{block_id}.upsample.conv.bias"
                  ]

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

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

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

      SCREAMING_SNAKE_CASE_ :     Tuple							=   [key for key in vae_state_dict if 'decoder.mid.attn' in key]
      SCREAMING_SNAKE_CASE_ :     Tuple							=   renew_vae_attention_paths(lowerCAmelCase__     )
      SCREAMING_SNAKE_CASE_ :     Any							=   {'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       A_     (							a    ,      a    ,      ):

      """simple docstring"""



      SCREAMING_SNAKE_CASE_ :     str							=   requests.get(
          ' https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml'     )
      SCREAMING_SNAKE_CASE_ :     Union[str, Any]							=   io.BytesIO(r.content     )

      SCREAMING_SNAKE_CASE_ :     str							=   OmegaConf.load(lowerCAmelCase__     )
      SCREAMING_SNAKE_CASE_ :     str							=   5_1_2
      SCREAMING_SNAKE_CASE_ :     int							=   'cuda' if torch.cuda.is_available() else 'cpu'
      if checkpoint_path.endswith('safetensors'     ):
            from safetensors import safe_open

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

      # Convert the VAE model.
      SCREAMING_SNAKE_CASE_ :     List[str]							=   create_vae_diffusers_config(lowerCAmelCase__    ,      image_size=lowerCAmelCase__     )
      SCREAMING_SNAKE_CASE_ :     str							=   custom_convert_ldm_vae_checkpoint(lowerCAmelCase__    ,      lowerCAmelCase__     )

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


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   :      int     =     parser.parse_args()

    vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)











"""simple docstring"""







import json
import logging
import os
import sys
from time import time
from unittest.mock import patch

from transformers.testing_utils import TestCasePlus, require_torch_tpu


logging.basicConfig(level=logging.DEBUG)

lowercase__		:				List[str]          =	logging.getLogger()

def 				UpperCamelCase_    (    lowerCAmelCase__      :				str						)						->   Dict:

  """simple docstring"""




  lowerCAmelCase_   :				Dict								= {}
  lowerCAmelCase_   :				Union[str, Any]								= os.path.join(lowerCAmelCase__    ,   'all_results.json'						)
  if os.path.exists(lowerCAmelCase__						):
    with open(lowerCAmelCase__    ,   'r'						) as f:
      lowerCAmelCase_   :				Any								= json.load(lowerCAmelCase__						)
  else:
    raise ValueError(f"can't find {path}"						)
  return results


lowercase__		:				Optional[int]          =	logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)







@require_torch_tpu
class       UpperCamelCase__   (	lowercase_  ):




      """simple docstring"""




      def      SCREAMING_SNAKE_CASE__ (      self				:      Any       ):
        import xla_spawn

        lowerCAmelCase_   :				Union[str, Any]								= self.get_auto_remove_tmp_dir()
        lowerCAmelCase_   :				Optional[Any]								= F"\n            ./examples/pytorch/text-classification/run_glue.py\n            --num_cores=8\n            ./examples/pytorch/text-classification/run_glue.py\n            --model_name_or_path distilbert-base-uncased\n            --output_dir {tmp_dir}\n            --overwrite_output_dir\n            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv\n            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv\n            --do_train\n            --do_eval\n            --debug tpu_metrics_debug\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n            --learning_rate=1e-4\n            --max_steps=10\n            --warmup_steps=2\n            --seed=42\n            --max_seq_length=128\n            ".split()

        with patch.object(SCREAMING_SNAKE_CASE_    ,			'argv'    ,			SCREAMING_SNAKE_CASE_       ):
          lowerCAmelCase_   :				Optional[Any]								= time()
          xla_spawn.main()
          lowerCAmelCase_   :				Any								= time()

          lowerCAmelCase_   :				Dict								= get_results(SCREAMING_SNAKE_CASE_       )
          self.assertGreaterEqual(result['eval_accuracy']    ,			0.75       )

          # Assert that the script takes less than 500 seconds to make sure it doesn't hang.
          self.assertLess(end - start    ,			5_0_0       )





      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]       ):
        import xla_spawn

        lowerCAmelCase_   :				Dict								= '\n            ./tests/test_trainer_tpu.py\n            --num_cores=8\n            ./tests/test_trainer_tpu.py\n            '.split()
        with patch.object(SCREAMING_SNAKE_CASE_    ,			'argv'    ,			SCREAMING_SNAKE_CASE_       ):
          xla_spawn.main()





