



'''simple docstring'''





import shutil
import tempfile
import unittest

from transformers import SPIECE_UNDERLINE, BatchEncoding, MBartaaTokenizer, MBartaaTokenizerFast, is_torch_available
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
)

from ...test_tokenization_common import TokenizerTesterMixin


_UpperCamelCase											=						get_tests_dir('fixtures/test_sentencepiece.model')

if is_torch_available():
   from transformers.models.mbart.modeling_mbart import shift_tokens_right

_UpperCamelCase											=						250004
_UpperCamelCase											=						250020






@require_sentencepiece
@require_tokenizers
class 					lowerCamelCase_    (				SCREAMING_SNAKE_CASE_   ,   unittest.TestCase     ):
  """simple docstring"""




  a_												=MBartaaTokenizer
  a_												=MBartaaTokenizerFast
  a_												=True
  a_												=True




  def 						_lowercase     ( self      :	str   )					->					Dict:
    super().setUp()

    # We have a SentencePiece fixture for testing
    __lowerCamelCase       :       Any           = MBartaaTokenizer(_a			,	src_lang='en_XX'			,	tgt_lang='ro_RO'			,	keep_accents=_a   )
    tokenizer.save_pretrained(self.tmpdirname   )




  def 						_lowercase     ( self      :	Any   )					->					Any:
    __lowerCamelCase       :       str           = '<s>'
    __lowerCamelCase       :       List[Any]           = 0

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




  def 						_lowercase     ( self      :	Tuple   )					->					Optional[Any]:
    __lowerCamelCase       :       Tuple           = list(self.get_tokenizer().get_vocab().keys()   )

    self.assertEqual(vocab_keys[0]			,	'<s>'   )
    self.assertEqual(vocab_keys[1]			,	'<pad>'   )
    self.assertEqual(vocab_keys[-1]			,	'<mask>'   )
    self.assertEqual(len(_a   )			,	1054   )




  def 						_lowercase     ( self      :	List[Any]   )					->					Tuple:
    self.assertEqual(self.get_tokenizer().vocab_size			,	1054   )




  def 						_lowercase     ( self      :	Any   )					->					Tuple:
    __lowerCamelCase       :       Dict           = MBartaaTokenizer(_a			,	src_lang='en_XX'			,	tgt_lang='ro_RO'			,	keep_accents=_a   )

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

    self.assertListEqual(
        tokenizer.convert_tokens_to_ids(_a   )			,	[value + tokenizer.fairseq_offset for value in [285, 46, 10, 170, 382]]			,	)

    __lowerCamelCase       :       Optional[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', 'é', '.']			,	)
    __lowerCamelCase       :       Optional[Any]           = tokenizer.convert_tokens_to_ids(_a   )
    self.assertListEqual(
        _a			,	[
            value + tokenizer.fairseq_offset
            for value in [8, 21, 84, 55, 24, 19, 7, 2, 602, 347, 347, 347, 3, 12, 66, 46, 72, 80, 6, 2, 4]
        ]			,	)

    __lowerCamelCase       :       Tuple           = 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 						_lowercase     ( self      :	Optional[int]   )					->					Dict:
    # fmt: off
    __lowerCamelCase       :       str           = {'input_ids': [[25_0004, 1_1062, 8_2772, 7, 15, 8_2772, 538, 5_1529, 237, 1_7198, 1290, 206, 9, 21_5175, 1314, 136, 1_7198, 1290, 206, 9, 5_6359, 42, 12_2009, 9, 1_6466, 16, 8_7344, 4537, 9, 4717, 7_8381, 6, 15_9958, 7, 15, 2_4480, 618, 4, 527, 2_2693, 5428, 4, 2777, 2_4480, 9874, 4, 4_3523, 594, 4, 803, 1_8392, 3_3189, 18, 4, 4_3523, 2_4447, 1_2399, 100, 2_4955, 8_3658, 9626, 14_4057, 15, 839, 2_2335, 16, 136, 2_4955, 8_3658, 8_3479, 15, 3_9102, 724, 16, 678, 645, 2789, 1328, 4589, 42, 12_2009, 11_5774, 23, 805, 1328, 4_6876, 7, 136, 5_3894, 1940, 4_2227, 4_1159, 1_7721, 823, 425, 4, 2_7512, 9_8722, 206, 136, 5531, 4970, 919, 1_7336, 5, 2], [25_0004, 2_0080, 618, 83, 8_2775, 47, 479, 9, 1517, 73, 5_3894, 333, 8_0581, 11_0117, 1_8811, 5256, 1295, 51, 15_2526, 297, 7986, 390, 12_4416, 538, 3_5431, 214, 98, 1_5044, 2_5737, 136, 7108, 4_3701, 23, 756, 13_5355, 7, 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], [25_0004, 581, 6_3773, 11_9455, 6, 14_7797, 8_8203, 7, 645, 70, 21, 3285, 1_0269, 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]], '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, 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, 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/mbart-large-50'			,	revision='d3913889c59cd5c9e456b269c376325eabad57e2'			,	)




  def 						_lowercase     ( self      :	Any   )					->					List[str]:
    if not self.test_slow_tokenizer:
      # as we don't have a slow version, we can't compare the outputs between slow and fast versions
      return

    __lowerCamelCase       :       Optional[Any]           = (self.rust_tokenizer_class, 'hf-internal-testing/tiny-random-mbart50', {})
    for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
      with self.subTest(f'{tokenizer.__class__.__name__} ({pretrained_name})'   ):
        __lowerCamelCase       :       int           = self.rust_tokenizer_class.from_pretrained(_a			,	**_a   )
        __lowerCamelCase       :       Any           = self.tokenizer_class.from_pretrained(_a			,	**_a   )

        __lowerCamelCase       :       Optional[int]           = tempfile.mkdtemp()

        __lowerCamelCase       :       Tuple           = tokenizer_r.save_pretrained(_a   )
        __lowerCamelCase       :       Tuple           = tokenizer_p.save_pretrained(_a   )

        # Checks it save with the same files + the tokenizer.json file for the fast one
        self.assertTrue(any('tokenizer.json' in f for f in tokenizer_r_files   )   )
        __lowerCamelCase       :       Tuple           = tuple(f for f in tokenizer_r_files if 'tokenizer.json' not in f   )
        self.assertSequenceEqual(_a			,	_a   )

        # Checks everything loads correctly in the same way
        __lowerCamelCase       :       List[Any]           = tokenizer_r.from_pretrained(_a   )
        __lowerCamelCase       :       Optional[Any]           = tokenizer_p.from_pretrained(_a   )

        # Check special tokens are set accordingly on Rust and Python
        for key in tokenizer_pp.special_tokens_map:
          self.assertTrue(hasattr(_a			,	_a   )   )
          # self.assertEqual(getattr(tokenizer_rp, key), getattr(tokenizer_pp, key))
          # self.assertEqual(getattr(tokenizer_rp, key + "_id"), getattr(tokenizer_pp, key + "_id"))

        shutil.rmtree(_a   )

        # Save tokenizer rust, legacy_format=True
        __lowerCamelCase       :       List[str]           = tempfile.mkdtemp()

        __lowerCamelCase       :       str           = tokenizer_r.save_pretrained(_a			,	legacy_format=_a   )
        __lowerCamelCase       :       List[Any]           = tokenizer_p.save_pretrained(_a   )

        # Checks it save with the same files
        self.assertSequenceEqual(_a			,	_a   )

        # Checks everything loads correctly in the same way
        __lowerCamelCase       :       str           = tokenizer_r.from_pretrained(_a   )
        __lowerCamelCase       :       Optional[int]           = tokenizer_p.from_pretrained(_a   )

        # Check special tokens are set accordingly on Rust and Python
        for key in tokenizer_pp.special_tokens_map:
          self.assertTrue(hasattr(_a			,	_a   )   )

        shutil.rmtree(_a   )

        # Save tokenizer rust, legacy_format=False
        __lowerCamelCase       :       Union[str, Any]           = tempfile.mkdtemp()

        __lowerCamelCase       :       Optional[Any]           = tokenizer_r.save_pretrained(_a			,	legacy_format=_a   )
        __lowerCamelCase       :       Any           = tokenizer_p.save_pretrained(_a   )

        # Checks it saved the tokenizer.json file
        self.assertTrue(any('tokenizer.json' in f for f in tokenizer_r_files   )   )

        # Checks everything loads correctly in the same way
        __lowerCamelCase       :       Tuple           = tokenizer_r.from_pretrained(_a   )
        __lowerCamelCase       :       Optional[int]           = tokenizer_p.from_pretrained(_a   )

        # Check special tokens are set accordingly on Rust and Python
        for key in tokenizer_pp.special_tokens_map:
          self.assertTrue(hasattr(_a			,	_a   )   )

        shutil.rmtree(_a   )






@require_torch
@require_sentencepiece
@require_tokenizers
class 					lowerCamelCase_    (				unittest.TestCase     ):
  """simple docstring"""




  a_												="""facebook/mbart-large-50-one-to-many-mmt"""
  a_												=[
      """ UN Chief Says There Is No Military Solution in Syria""",
      """ Secretary-General Ban Ki-moon says his response to Russia's stepped up military support for Syria is that \"there is no military solution\" to the nearly five-year conflict and more weapons will only worsen the violence and misery for millions of people.""",
  ]
  a_												=[
      """Şeful ONU declară că nu există o soluţie militară în Siria""",
      """Secretarul General Ban Ki-moon declară că răspunsul său la intensificarea sprijinului militar al Rusiei"""
      """ pentru Siria este că \"nu există o soluţie militară\" la conflictul de aproape cinci ani şi că noi arme nu vor"""
      """ face decât să înrăutăţească violenţele şi mizeria pentru milioane de oameni.""",
  ]
  a_												=[EN_CODE, 8274, 12_7873, 2_5916, 7, 8622, 2071, 438, 6_7485, 53, 18_7895, 23, 5_1712, 2]




  @classmethod
  def 						_lowercase     ( cls      :	Optional[int]   )					->					Tuple:
    __lowerCamelCase       :       MBartaaTokenizer           = MBartaaTokenizer.from_pretrained(
        cls.checkpoint_name			,	src_lang='en_XX'			,	tgt_lang='ro_RO'   )
    __lowerCamelCase       :       int           = 1
    return cls




  def 						_lowercase     ( self      :	Any   )					->					Dict:
    self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['ar_AR']			,	25_0001   )
    self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['en_EN']			,	25_0004   )
    self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['ro_RO']			,	25_0020   )
    self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['mr_IN']			,	25_0038   )




  def 						_lowercase     ( self      :	Tuple   )					->					Dict:
    __lowerCamelCase       :       List[Any]           = self.tokenizer.batch_encode_plus(self.src_text   ).input_ids[0]
    self.assertListEqual(self.expected_src_tokens			,	_a   )




  def 						_lowercase     ( self      :	str   )					->					List[str]:
    self.assertIn(_a			,	self.tokenizer.all_special_ids   )
    __lowerCamelCase       :       Union[str, Any]           = [RO_CODE, 884, 9019, 96, 9, 916, 8_6792, 36, 1_8743, 1_5596, 5, 2]
    __lowerCamelCase       :       Any           = self.tokenizer.decode(_a			,	skip_special_tokens=_a   )
    __lowerCamelCase       :       Union[str, Any]           = self.tokenizer.decode(generated_ids[1:]			,	skip_special_tokens=_a   )
    self.assertEqual(_a			,	_a   )
    self.assertNotIn(self.tokenizer.eos_token			,	_a   )




  def 						_lowercase     ( self      :	Union[str, Any]   )					->					Dict:
    __lowerCamelCase       :       Optional[Any]           = ['this is gunna be a long sentence ' * 20]
    assert isinstance(src_text[0]			,	_a   )
    __lowerCamelCase       :       Union[str, Any]           = 10
    __lowerCamelCase       :       Optional[int]           = self.tokenizer(_a			,	max_length=_a			,	truncation=_a   ).input_ids[0]
    self.assertEqual(ids[0]			,	_a   )
    self.assertEqual(ids[-1]			,	2   )
    self.assertEqual(len(_a   )			,	_a   )




  def 						_lowercase     ( self      :	Union[str, Any]   )					->					List[Any]:
    self.assertListEqual(self.tokenizer.convert_tokens_to_ids(['<mask>', 'ar_AR']   )			,	[25_0053, 25_0001]   )




  def 						_lowercase     ( self      :	List[str]   )					->					Tuple:
    __lowerCamelCase       :       Any           = tempfile.mkdtemp()
    __lowerCamelCase       :       Optional[Any]           = self.tokenizer.fairseq_tokens_to_ids
    self.tokenizer.save_pretrained(_a   )
    __lowerCamelCase       :       List[str]           = MBartaaTokenizer.from_pretrained(_a   )
    self.assertDictEqual(new_tok.fairseq_tokens_to_ids			,	_a   )




  @require_torch
  def 						_lowercase     ( self      :	Any   )					->					List[str]:
    __lowerCamelCase       :       Optional[int]           = self.tokenizer(self.src_text			,	text_target=self.tgt_text			,	padding=_a			,	return_tensors='pt'   )
    __lowerCamelCase       :       Optional[int]           = shift_tokens_right(batch['labels']			,	self.tokenizer.pad_token_id   )

    # fairseq batch: https://gist.github.com/sshleifer/cba08bc2109361a74ac3760a7e30e4f4
    assert batch.input_ids[1][0] == EN_CODE
    assert batch.input_ids[1][-1] == 2
    assert batch.labels[1][0] == RO_CODE
    assert batch.labels[1][-1] == 2
    assert batch.decoder_input_ids[1][:2].tolist() == [2, RO_CODE]




  @require_torch
  def 						_lowercase     ( self      :	int   )					->					int:
    __lowerCamelCase       :       List[Any]           = self.tokenizer(
        self.src_text			,	text_target=self.tgt_text			,	padding=_a			,	truncation=_a			,	max_length=len(self.expected_src_tokens   )			,	return_tensors='pt'			,	)
    __lowerCamelCase       :       Dict           = shift_tokens_right(batch['labels']			,	self.tokenizer.pad_token_id   )

    self.assertIsInstance(_a			,	_a   )

    self.assertEqual((2, 14)			,	batch.input_ids.shape   )
    self.assertEqual((2, 14)			,	batch.attention_mask.shape   )
    __lowerCamelCase       :       Any           = batch.input_ids.tolist()[0]
    self.assertListEqual(self.expected_src_tokens			,	_a   )
    self.assertEqual(2			,	batch.decoder_input_ids[0, 0]   )  # decoder_start_token_id
    # Test that special tokens are reset
    self.assertEqual(self.tokenizer.prefix_tokens			,	[EN_CODE]   )
    self.assertEqual(self.tokenizer.suffix_tokens			,	[self.tokenizer.eos_token_id]   )




  def 						_lowercase     ( self      :	Union[str, Any]   )					->					Union[str, Any]:
    __lowerCamelCase       :       Any           = self.tokenizer(self.src_text			,	padding=_a			,	truncation=_a			,	max_length=3			,	return_tensors='pt'   )
    __lowerCamelCase       :       Dict           = self.tokenizer(
        text_target=self.tgt_text			,	padding=_a			,	truncation=_a			,	max_length=10			,	return_tensors='pt'   )
    __lowerCamelCase       :       Dict           = targets['input_ids']
    __lowerCamelCase       :       Any           = shift_tokens_right(_a			,	self.tokenizer.pad_token_id   )

    self.assertEqual(batch.input_ids.shape[1]			,	3   )
    self.assertEqual(batch.decoder_input_ids.shape[1]			,	10   )




  @require_torch
  def 						_lowercase     ( self      :	List[Any]   )					->					List[str]:
    __lowerCamelCase       :       Any           = self.tokenizer._build_translation_inputs(
        'A test'			,	return_tensors='pt'			,	src_lang='en_XX'			,	tgt_lang='ar_AR'   )

    self.assertEqual(
        nested_simplify(_a   )			,	{
            # en_XX, A, test, EOS
            'input_ids': [[25_0004, 62, 3034, 2]],
            'attention_mask': [[1, 1, 1, 1]],
            # ar_AR
            'forced_bos_token_id': 25_0001,
        }			,	)











'''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        a_     (							_lowerCAmelCase       ,_lowerCAmelCase       )      ->      Dict:
  __lowerCamelCase       :       Any           = checkpoint

  __lowerCamelCase       :       List[str]           = {}

  __lowerCamelCase       :       List[str]           = vae_state_dict['encoder.conv_in.weight']
  __lowerCamelCase       :       Dict           = vae_state_dict['encoder.conv_in.bias']
  __lowerCamelCase       :       Optional[Any]           = vae_state_dict['encoder.conv_out.weight']
  __lowerCamelCase       :       Optional[int]           = vae_state_dict['encoder.conv_out.bias']
  __lowerCamelCase       :       Optional[int]           = vae_state_dict['encoder.norm_out.weight']
  __lowerCamelCase       :       Optional[int]           = vae_state_dict['encoder.norm_out.bias']

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

  __lowerCamelCase       :       Any           = vae_state_dict['quant_conv.weight']
  __lowerCamelCase       :       Union[str, Any]           = vae_state_dict['quant_conv.bias']
  __lowerCamelCase       :       Tuple           = vae_state_dict['post_quant_conv.weight']
  __lowerCamelCase       :       List[Any]           = vae_state_dict['post_quant_conv.bias']

  # Retrieves the keys for the encoder down blocks only
  __lowerCamelCase       :       str           = len({'.'.join(layer.split('.'       )[:3]       ) for layer in vae_state_dict if 'encoder.down' in layer}       )
  __lowerCamelCase       :       int           = {
      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
  __lowerCamelCase       :       Dict           = len({'.'.join(layer.split('.'       )[:3]       ) for layer in vae_state_dict if 'decoder.up' in layer}       )
  __lowerCamelCase       :       Dict           = {
      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       ):
    __lowerCamelCase       :       List[str]           = [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       :       Any           = vae_state_dict.pop(
          F'encoder.down.{i}.downsample.conv.weight'       )
      __lowerCamelCase       :       Any           = vae_state_dict.pop(
          F'encoder.down.{i}.downsample.conv.bias'       )

    __lowerCamelCase       :       Optional[int]           = renew_vae_resnet_paths(_lowerCAmelCase       )
    __lowerCamelCase       :       Any           = {'old': F'down.{i}.block', 'new': F'down_blocks.{i}.resnets'}
    assign_to_checkpoint(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,additional_replacements=[meta_path]       ,config=_lowerCAmelCase       )

  __lowerCamelCase       :       List[Any]           = [key for key in vae_state_dict if 'encoder.mid.block' in key]
  __lowerCamelCase       :       str           = 2
  for i in range(1       ,num_mid_res_blocks + 1       ):
    __lowerCamelCase       :       int           = [key for key in mid_resnets if F'encoder.mid.block_{i}' in key]

    __lowerCamelCase       :       Optional[Any]           = renew_vae_resnet_paths(_lowerCAmelCase       )
    __lowerCamelCase       :       Dict           = {'old': F'mid.block_{i}', 'new': F'mid_block.resnets.{i - 1}'}
    assign_to_checkpoint(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,additional_replacements=[meta_path]       ,config=_lowerCAmelCase       )

  __lowerCamelCase       :       Optional[Any]           = [key for key in vae_state_dict if 'encoder.mid.attn' in key]
  __lowerCamelCase       :       int           = renew_vae_attention_paths(_lowerCAmelCase       )
  __lowerCamelCase       :       Optional[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       )

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

    __lowerCamelCase       :       Tuple           = renew_vae_resnet_paths(_lowerCAmelCase       )
    __lowerCamelCase       :       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       )

  __lowerCamelCase       :       Optional[Any]           = [key for key in vae_state_dict if 'decoder.mid.block' in key]
  __lowerCamelCase       :       Union[str, Any]           = 2
  for i in range(1       ,num_mid_res_blocks + 1       ):
    __lowerCamelCase       :       List[str]           = [key for key in mid_resnets if F'decoder.mid.block_{i}' in key]

    __lowerCamelCase       :       str           = renew_vae_resnet_paths(_lowerCAmelCase       )
    __lowerCamelCase       :       Any           = {'old': F'mid.block_{i}', 'new': F'mid_block.resnets.{i - 1}'}
    assign_to_checkpoint(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       ,additional_replacements=[meta_path]       ,config=_lowerCAmelCase       )

  __lowerCamelCase       :       Any           = [key for key in vae_state_dict if 'decoder.mid.attn' in key]
  __lowerCamelCase       :       List[str]           = renew_vae_attention_paths(_lowerCAmelCase       )
  __lowerCamelCase       :       Optional[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_     (							_lowerCAmelCase       ,_lowerCAmelCase       ,)      ->      List[Any]:
  # Only support V1
  __lowerCamelCase       :       Any           = requests.get(
      ' https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml'       )
  __lowerCamelCase       :       str           = io.BytesIO(r.content       )

  __lowerCamelCase       :       Any           = OmegaConf.load(_lowerCAmelCase       )
  __lowerCamelCase       :       List[Any]           = 512
  __lowerCamelCase       :       Dict           = 'cuda' if torch.cuda.is_available() else 'cpu'
  if checkpoint_path.endswith('safetensors'       ):
    from safetensors import safe_open

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

  # Convert the VAE model.
  __lowerCamelCase       :       int           = create_vae_diffusers_config(_lowerCAmelCase       ,image_size=_lowerCAmelCase       )
  __lowerCamelCase       :       int           = custom_convert_ldm_vae_checkpoint(_lowerCAmelCase       ,_lowerCAmelCase       )

  __lowerCamelCase       :       Any           = AutoencoderKL(**_lowerCAmelCase       )
  vae.load_state_dict(_lowerCAmelCase       )
  vae.save_pretrained(_lowerCAmelCase       )


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

   vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)






