






"""simple docstring"""





import unittest
from pathlib import Path
from shutil import copyfile

from transformers import SPIECE_UNDERLINE, is_sentencepiece_available
from transformers.models.speech_to_text import SpeechaTextTokenizer
from transformers.models.speech_to_text.tokenization_speech_to_text import VOCAB_FILES_NAMES, save_json
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin


UpperCamelCase_	:				str										= get_tests_dir("""fixtures/test_sentencepiece.model""")

if is_sentencepiece_available():
    import sentencepiece as sp


UpperCamelCase_	:				Union[str, Any]										= 5
UpperCamelCase_	:				Optional[int]										= 10


@require_sentencepiece
@require_tokenizers
class      _lowercase (     lowerCAmelCase		,  unittest.TestCase						):
      _a   :     int								=   SpeechaTextTokenizer
      _a   :     List[str]								=   False
      _a   :     Dict								=   True


      def _UpperCamelCase						(    self				:	Optional[int]						):





            """simple docstring"""







            super().setUp()

            __snake_case      :    str  			=sp.SentencePieceProcessor()
            spm_model.Load(a						)
            __snake_case      :    Optional[int]  			=['''<s>''', '''<pad>''', '''</s>''', '''<unk>''']

            vocab += [spm_model.IdToPiece(id_						) for id_ in range(len(a						)						)]
            __snake_case      :    Union[str, Any]  			=dict(zip(a							,   range(len(a						)						)						)						)

            __snake_case      :    List[Any]  			=Path(self.tmpdirname						)
            save_json(a							,   save_dir / VOCAB_FILES_NAMES['''vocab_file''']						)
            if not (save_dir / VOCAB_FILES_NAMES["spm_file"]).exists():
                  copyfile(a							,   save_dir / VOCAB_FILES_NAMES['''spm_file''']						)

            __snake_case      :    Union[str, Any]  			=SpeechaTextTokenizer.from_pretrained(self.tmpdirname						)
            tokenizer.save_pretrained(self.tmpdirname						)


      def _UpperCamelCase						(    self				:	List[Any]						):





            """simple docstring"""







            __snake_case      :    Tuple  			='''<pad>'''
            __snake_case      :    Optional[int]  			=1

            self.assertEqual(self.get_tokenizer()._convert_token_to_id(a						)							,   a						)
            self.assertEqual(self.get_tokenizer()._convert_id_to_token(a						)							,   a						)


      def _UpperCamelCase						(    self				:	Union[str, Any]						):





            """simple docstring"""







            __snake_case      :    Optional[int]  			=list(self.get_tokenizer().get_vocab().keys()						)

            self.assertEqual(vocab_keys[0]							,   '''<s>'''						)
            self.assertEqual(vocab_keys[1]							,   '''<pad>'''						)
            self.assertEqual(vocab_keys[-1]							,   '''j'''						)
            self.assertEqual(len(a						)							,   1_0_0_1						)


      def _UpperCamelCase						(    self				:	Optional[int]						):





            """simple docstring"""







            self.assertEqual(self.get_tokenizer().vocab_size							,   1_0_0_1						)


      def _UpperCamelCase						(    self				:	Any						):





            """simple docstring"""







            __snake_case      :    List[Any]  			=SpeechaTextTokenizer.from_pretrained(self.tmpdirname						)

            __snake_case      :    int  			=tokenizer.tokenize('''This is a test'''						)
            self.assertListEqual(a							,   ['''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est''']						)

            self.assertListEqual(
                tokenizer.convert_tokens_to_ids(a						)							,   [2_8_9, 5_0, 1_4, 1_7_4, 3_8_6]							,   )

            __snake_case      :    int  			=tokenizer.tokenize('''I was born in 92000, and this is falsé.'''						)
            self.assertListEqual(
                a							,   [SPIECE_UNDERLINE + '''I''', SPIECE_UNDERLINE + '''was''', SPIECE_UNDERLINE + '''b''', '''or''', '''n''', SPIECE_UNDERLINE + '''in''', SPIECE_UNDERLINE + '''''', '''9''', '''2''', '''0''', '''0''', '''0''', ''',''', SPIECE_UNDERLINE + '''and''', SPIECE_UNDERLINE + '''this''', SPIECE_UNDERLINE + '''is''', SPIECE_UNDERLINE + '''f''', '''al''', '''s''', '''é''', '''.''']							,   )
            __snake_case      :    int  			=tokenizer.convert_tokens_to_ids(a						)
            self.assertListEqual(a							,   [1_2, 2_5, 8_8, 5_9, 2_8, 2_3, 1_1, 4, 6_0_6, 3_5_1, 3_5_1, 3_5_1, 7, 1_6, 7_0, 5_0, 7_6, 8_4, 1_0, 4, 8]						)

            __snake_case      :    Any  			=tokenizer.convert_ids_to_tokens(a						)
            self.assertListEqual(
                a							,   [SPIECE_UNDERLINE + '''I''', SPIECE_UNDERLINE + '''was''', SPIECE_UNDERLINE + '''b''', '''or''', '''n''', SPIECE_UNDERLINE + '''in''', SPIECE_UNDERLINE + '''''', '''<unk>''', '''2''', '''0''', '''0''', '''0''', ''',''', SPIECE_UNDERLINE + '''and''', SPIECE_UNDERLINE + '''this''', SPIECE_UNDERLINE + '''is''', SPIECE_UNDERLINE + '''f''', '''al''', '''s''', '''<unk>''', '''.''']							,   )







      @slow
      def _UpperCamelCase						(    self				:	Any						):





            """simple docstring"""







            __snake_case      :    Tuple  			={'''input_ids''': [[3_7_9_1, 7_9_7, 3_1, 1_1, 6_4, 7_9_7, 3_1, 2_4_2_9, 4_3_3, 1_2, 1_1_7_6, 1_2, 2_0, 7_8_6, 9_1_5, 1_4_2, 2_4_1_3, 2_4_0, 3_7, 3_2_3_8, 7_9_7, 3_1, 1_1, 3_5, 9_3, 9_1_5, 1_4_2, 2_4_1_3, 2_4_0, 3_7, 5_5_4_0, 5_6_7, 1_2_7_6, 9_3, 3_7, 6_1_0, 4_0, 6_2, 4_5_5, 6_5_7, 1_0_4_2, 1_2_3, 7_8_0, 1_7_7, 3_7, 3_0_9, 2_4_1, 1_2_9_8, 5_1_4, 2_0, 2_9_2, 2_7_3_7, 1_1_4, 2_4_6_9, 2_4_1, 8_5, 6_4, 3_0_2, 5_4_8, 5_2_8, 4_2_3, 4, 5_0_9, 4_0_6, 4_2_3, 3_7, 6_0_1, 4, 7_7_7, 3_0_2, 5_4_8, 5_2_8, 4_2_3, 2_8_4, 4, 3_3_8_8, 5_1_1, 4_5_9, 4, 3_5_5_5, 4_0, 3_2_1, 3_0_2, 7_0_5, 4, 3_3_8_8, 5_1_1, 5_8_3, 3_2_6, 5, 5, 5, 6_2, 3_3_1_0, 5_6_0, 1_7_7, 2_6_8_0, 2_1_7, 1_5_0_8, 3_2, 3_1, 8_5_3, 4_1_8, 6_4, 5_8_3, 5_1_1, 1_6_0_5, 6_2, 3_5, 9_3, 5_6_0, 1_7_7, 2_6_8_0, 2_1_7, 1_5_0_8, 1_5_2_1, 6_4, 5_8_3, 5_1_1, 5_1_9, 6_2, 2_0, 1_5_1_5, 7_6_4, 2_0, 1_4_9, 2_6_1, 5_6_2_5, 7_9_7_2, 2_0, 5_5_4_0, 5_6_7, 1_2_7_6, 9_3, 3_9_2_5, 1_6_7_5, 1_1, 1_5, 8_0_2, 7_9_7_2, 5_7_6, 2_1_7, 1_5_0_8, 1_1, 3_5, 9_3, 1_2_5_3, 2_4_4_1, 1_5, 2_8_9, 6_5_2, 3_1, 4_1_6, 3_2_1, 3_8_4_2, 1_1_5, 4_0, 9_1_1, 8, 4_7_6, 6_1_9, 4, 3_8_0, 1_4_2, 4_2_3, 3_3_5, 2_4_0, 3_5, 9_3, 2_6_4, 8, 1_1, 3_3_5, 5_6_9, 4_2_0, 1_6_3, 5, 2], [2_6_0, 5_4_8, 5_2_8, 4_2_3, 2_0, 4_5_1, 2_0, 2_6_8_1, 1_1_5_3, 3_4_3_4, 2_0, 5_5_4_0, 3_7, 5_6_7, 1_2_6, 1_2_5_3, 2_4_4_1, 3_3_7_6, 4_4_9, 2_1_0, 4_3_1, 1_5_6_3, 1_7_7, 7_6_7, 5_5_4_0, 1_1, 1_2_0_3, 4_7_2, 1_1, 2_9_5_3, 6_8_5, 2_8_5, 3_6_4, 7_0_6, 1_1_5_3, 2_0, 6_7_9_9, 2_0, 2_8_6_9, 2_0, 4_4_6_4, 1_2_6, 4_0, 2_4_2_9, 2_0, 1_0_4_0, 8_6_6, 2_6_6_4, 4_1_8, 2_0, 3_1_8, 2_0, 1_7_2_6, 1_8_6, 2_0, 2_6_5, 5_2_2, 3_5, 9_3, 2_1_9_1, 4_6_3_4, 2_0, 1_0_4_0, 1_2, 6_7_9_9, 1_5, 2_2_8, 2_3_5_6, 1_4_2, 3_1, 1_1, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [2_5_7_5, 2_6_6_6, 6_8_4, 1_5_8_2, 1_1_7_6, 1_2, 6_2_7, 1_4_9, 6_1_9, 2_0, 4_9_0_2, 5_6_3, 1_1, 2_0, 1_4_9, 2_6_1, 3_4_2_0, 2_3_5_6, 1_7_4, 1_4_2, 4_7_1_4, 1_3_1, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], '''attention_mask''': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
            # fmt: on

            self.tokenizer_integration_test_util(
                expected_encoding=a							,   model_name='''facebook/s2t-small-mustc-en-de-st'''							,   revision='''a14f04cf0776c02f62a8cb800cf7909e15ea23ad'''							,   )



@require_sentencepiece
class      _lowercase (     unittest.TestCase						):
      _a   :     Dict								=   '''valhalla/s2t_mustc_multilinguial_medium'''

      _a   :     Any								=   '''C\'est trop cool'''
      _a   :     int								=   '''Esto es genial'''


      @classmethod
      def _UpperCamelCase						(    cls				:	Union[str, Any]						):





            """simple docstring"""







            __snake_case      :    SpeechaTextTokenizer  			=SpeechaTextTokenizer.from_pretrained(cls.checkpoint_name						)
            return cls


      def _UpperCamelCase						(    self				:	List[str]						):





            """simple docstring"""







            self.assertEqual(self.tokenizer.lang_code_to_id['''pt''']							,   4						)
            self.assertEqual(self.tokenizer.lang_code_to_id['''ru''']							,   6						)
            self.assertEqual(self.tokenizer.lang_code_to_id['''it''']							,   9						)
            self.assertEqual(self.tokenizer.lang_code_to_id['''de''']							,   1_1						)


      def _UpperCamelCase						(    self				:	str						):





            """simple docstring"""







            self.assertEqual(self.tokenizer.vocab_size							,   1_0_0_0_0						)


      def _UpperCamelCase						(    self				:	List[str]						):





            """simple docstring"""







            self.assertIn(a							,   self.tokenizer.all_special_ids						)
            __snake_case      :    int  			=[ES_CODE, 4, 1_6_0_1, 4_7, 7_6_4_7, 2]
            __snake_case      :    Any  			=self.tokenizer.decode(a							,   skip_special_tokens=a						)
            __snake_case      :    Any  			=self.tokenizer.decode(generated_ids[1:]							,   skip_special_tokens=a						)
            self.assertEqual(a							,   a						)
            self.assertNotIn(self.tokenizer.eos_token							,   a						)


      def _UpperCamelCase						(    self				:	Optional[int]						):





            """simple docstring"""







            __snake_case      :    Optional[int]  			='''fr'''
            __snake_case      :    str  			=self.tokenizer(self.french_text						).input_ids
            self.assertEqual(encoded[0]							,   a						)
            self.assertEqual(encoded[-1]							,   self.tokenizer.eos_token_id						)







      def _UpperCamelCase						(    self				:	int						):





            """simple docstring"""







            __snake_case      :    Optional[Any]  			='''fr'''
            self.assertListEqual(self.tokenizer.prefix_tokens							,   [FR_CODE]						)

            __snake_case      :    List[Any]  			='''es'''
            self.assertListEqual(self.tokenizer.prefix_tokens							,   [ES_CODE]						)










"""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					(    a :				Any       ,		a :				List[Any] )					->					Dict:
      __snake_case      :    Any  			=checkpoint

      __snake_case      :    Dict  			={}

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

      __snake_case      :    Tuple  			=vae_state_dict['''decoder.conv_in.weight''']
      __snake_case      :    str  			=vae_state_dict['''decoder.conv_in.bias''']
      __snake_case      :    List[str]  			=vae_state_dict['''decoder.conv_out.weight''']
      __snake_case      :    Tuple  			=vae_state_dict['''decoder.conv_out.bias''']
      __snake_case      :    Union[str, Any]  			=vae_state_dict['''decoder.norm_out.weight''']
      __snake_case      :    List[str]  			=vae_state_dict['''decoder.norm_out.bias''']

      __snake_case      :    Tuple  			=vae_state_dict['''quant_conv.weight''']
      __snake_case      :    List[str]  			=vae_state_dict['''quant_conv.bias''']
      __snake_case      :    Optional[int]  			=vae_state_dict['''post_quant_conv.weight''']
      __snake_case      :    Tuple  			=vae_state_dict['''post_quant_conv.bias''']

      # Retrieves the keys for the encoder down blocks only
      __snake_case      :    Union[str, Any]  			=len({'''.'''.join(layer.split('''.''' )[:3] ) for layer in vae_state_dict if '''encoder.down''' in layer} )
      __snake_case      :    Union[str, Any]  			={
          layer_id: [key for key in vae_state_dict if f'''down.{layer_id}''' in key] for layer_id in range(a )
      }

      # Retrieves the keys for the decoder up blocks only
      __snake_case      :    List[Any]  			=len({'''.'''.join(layer.split('''.''' )[:3] ) for layer in vae_state_dict if '''decoder.up''' in layer} )
      __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(a )
      }

      for i in range(a ):
            __snake_case      :    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:
                  __snake_case      :    Optional[Any]  			=vae_state_dict.pop(
                      f'''encoder.down.{i}.downsample.conv.weight''' )
                  __snake_case      :    Optional[int]  			=vae_state_dict.pop(
                      f'''encoder.down.{i}.downsample.conv.bias''' )

            __snake_case      :    Optional[Any]  			=renew_vae_resnet_paths(a )
            __snake_case      :    Union[str, Any]  			={'''old''': f'''down.{i}.block''', '''new''': f'''down_blocks.{i}.resnets'''}
            assign_to_checkpoint(a       ,		a       ,		a       ,		additional_replacements=[meta_path]       ,		config=a )

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

            __snake_case      :    List[str]  			=renew_vae_resnet_paths(a )
            __snake_case      :    List[str]  			={'''old''': f'''mid.block_{i}''', '''new''': f'''mid_block.resnets.{i - 1}'''}
            assign_to_checkpoint(a       ,		a       ,		a       ,		additional_replacements=[meta_path]       ,		config=a )

      __snake_case      :    int  			=[key for key in vae_state_dict if '''encoder.mid.attn''' in key]
      __snake_case      :    List[Any]  			=renew_vae_attention_paths(a )
      __snake_case      :    Union[str, Any]  			={'''old''': '''mid.attn_1''', '''new''': '''mid_block.attentions.0'''}
      assign_to_checkpoint(a       ,		a       ,		a       ,		additional_replacements=[meta_path]       ,		config=a )
      conv_attn_to_linear(a )

      for i in range(a ):
            __snake_case      :    List[Any]  			=num_up_blocks - 1 - i
            __snake_case      :    Dict  			=[
                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:
                  __snake_case      :    int  			=vae_state_dict[
                      f'''decoder.up.{block_id}.upsample.conv.weight'''
                  ]
                  __snake_case      :    Optional[Any]  			=vae_state_dict[
                      f'''decoder.up.{block_id}.upsample.conv.bias'''
                  ]

            __snake_case      :    List[str]  			=renew_vae_resnet_paths(a )
            __snake_case      :    int  			={'''old''': f'''up.{block_id}.block''', '''new''': f'''up_blocks.{i}.resnets'''}
            assign_to_checkpoint(a       ,		a       ,		a       ,		additional_replacements=[meta_path]       ,		config=a )

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

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

      __snake_case      :    Optional[Any]  			=[key for key in vae_state_dict if '''decoder.mid.attn''' in key]
      __snake_case      :    Dict  			=renew_vae_attention_paths(a )
      __snake_case      :    List[str]  			={'''old''': '''mid.attn_1''', '''new''': '''mid_block.attentions.0'''}
      assign_to_checkpoint(a       ,		a       ,		a       ,		additional_replacements=[meta_path]       ,		config=a )
      conv_attn_to_linear(a )
      return new_checkpoint



def     __lowercase					(    a :				str       ,		a :				str       ,		)					->					Optional[int]:
      # Only support V1
      __snake_case      :    Union[str, Any]  			=requests.get(
          ''' https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml''' )
      __snake_case      :    Union[str, Any]  			=io.BytesIO(r.content )

      __snake_case      :    Any  			=OmegaConf.load(a )
      __snake_case      :    Union[str, Any]  			=512
      __snake_case      :    List[str]  			='''cuda''' if torch.cuda.is_available() else '''cpu'''
      if checkpoint_path.endswith('''safetensors''' ):
            from safetensors import safe_open

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

      # Convert the VAE model.
      __snake_case      :    Dict  			=create_vae_diffusers_config(a       ,		image_size=a )
      __snake_case      :    Union[str, Any]  			=custom_convert_ldm_vae_checkpoint(a       ,		a )

      __snake_case      :    str  			=AutoencoderKL(**a )
      vae.load_state_dict(a )
      vae.save_pretrained(a )


if __name__ == "__main__":
    UpperCamelCase_	:				Optional[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_	:				Optional[Any]										= parser.parse_args()

    vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)


