



'''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 							UpperCAmelCase_  (						__lowerCamelCase : Any     ,__lowerCamelCase : Optional[int]			):
   lowercase_      :Optional[int]         =			checkpoint

   lowercase_      :Dict         =			{}

   lowercase_      :int         =			vae_state_dict["encoder.conv_in.weight"]
   lowercase_      :Optional[int]         =			vae_state_dict["encoder.conv_in.bias"]
   lowercase_      :Optional[Any]         =			vae_state_dict["encoder.conv_out.weight"]
   lowercase_      :List[str]         =			vae_state_dict["encoder.conv_out.bias"]
   lowercase_      :Dict         =			vae_state_dict["encoder.norm_out.weight"]
   lowercase_      :Tuple         =			vae_state_dict["encoder.norm_out.bias"]

   lowercase_      :str         =			vae_state_dict["decoder.conv_in.weight"]
   lowercase_      :List[Any]         =			vae_state_dict["decoder.conv_in.bias"]
   lowercase_      :List[str]         =			vae_state_dict["decoder.conv_out.weight"]
   lowercase_      :Dict         =			vae_state_dict["decoder.conv_out.bias"]
   lowercase_      :Tuple         =			vae_state_dict["decoder.norm_out.weight"]
   lowercase_      :Optional[Any]         =			vae_state_dict["decoder.norm_out.bias"]

   lowercase_      :str         =			vae_state_dict["quant_conv.weight"]
   lowercase_      :List[str]         =			vae_state_dict["quant_conv.bias"]
   lowercase_      :Dict         =			vae_state_dict["post_quant_conv.weight"]
   lowercase_      :List[str]         =			vae_state_dict["post_quant_conv.bias"]

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

      lowercase_      :Tuple         =			renew_vae_resnet_paths(__lowerCamelCase			)
      lowercase_      :Union[str, Any]         =			{"old": F'down.{i}.block', "new": F'down_blocks.{i}.resnets'}
      assign_to_checkpoint(__lowerCamelCase     ,__lowerCamelCase     ,__lowerCamelCase     ,additional_replacements=[meta_path]     ,config=__lowerCamelCase			)

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

      lowercase_      :List[Any]         =			renew_vae_resnet_paths(__lowerCamelCase			)
      lowercase_      :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			)

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

      lowercase_      :List[Any]         =			renew_vae_resnet_paths(__lowerCamelCase			)
      lowercase_      :List[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			)

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

      lowercase_      :Tuple         =			renew_vae_resnet_paths(__lowerCamelCase			)
      lowercase_      :Optional[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			)

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

   lowercase_      :Union[str, Any]         =			OmegaConf.load(__lowerCamelCase			)
   lowercase_      :Optional[Any]         =			5_12
   lowercase_      :List[str]         =			"cuda" if torch.cuda.is_available() else "cpu"
   if checkpoint_path.endswith("safetensors"			):
      from safetensors import safe_open

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

   # Convert the VAE model.
   lowercase_      :List[str]         =			create_vae_diffusers_config(__lowerCamelCase     ,image_size=__lowerCamelCase			)
   lowercase_      :int         =			custom_convert_ldm_vae_checkpoint(__lowerCamelCase     ,__lowerCamelCase			)

   lowercase_      :Tuple         =			AutoencoderKL(**__lowerCamelCase			)
   vae.load_state_dict(__lowerCamelCase			)
   vae.save_pretrained(__lowerCamelCase			)


if __name__ == "__main__":
 lowerCAmelCase :					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.''')

 lowerCAmelCase :					Dict 							=parser.parse_args()

 vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)









'''simple docstring'''


import unittest

from transformers import DebertaConfig, is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import (
     DebertaForMaskedLM,
     DebertaForQuestionAnswering,
     DebertaForSequenceClassification,
     DebertaForTokenClassification,
     DebertaModel,
 )
 from transformers.models.deberta.modeling_deberta import DEBERTA_PRETRAINED_MODEL_ARCHIVE_LIST
class 				a_       (    _lowerCAmelCase							):





     def __init__(			self		: Any     ,				lowercase		: int     ,				lowercase		: Union[str, Any]=13     ,				lowercase		: List[str]=7     ,				lowercase		: List[str]=True     ,				lowercase		: int=True     ,				lowercase		: Tuple=True     ,				lowercase		: int=True     ,				lowercase		: List[Any]=99     ,				lowercase		: Optional[int]=32     ,				lowercase		: Dict=5     ,				lowercase		: Optional[int]=4     ,				lowercase		: List[str]=37     ,				lowercase		: Tuple="gelu"     ,				lowercase		: List[Any]=0.1     ,				lowercase		: Tuple=0.1     ,				lowercase		: List[str]=512     ,				lowercase		: str=16     ,				lowercase		: Tuple=2     ,				lowercase		: List[Any]=0.02     ,				lowercase		: Dict=False     ,				lowercase		: Dict=True     ,				lowercase		: int="None"     ,				lowercase		: Optional[Any]=3     ,				lowercase		: Dict=4     ,				lowercase		: List[Any]=None     ,				):





        """simple docstring"""




        lowercase_      :int         =			parent
        lowercase_      :str         =			batch_size
        lowercase_      :Tuple         =			seq_length
        lowercase_      :Union[str, Any]         =			is_training
        lowercase_      :Dict         =			use_input_mask
        lowercase_      :Any         =			use_token_type_ids
        lowercase_      :Tuple         =			use_labels
        lowercase_      :Dict         =			vocab_size
        lowercase_      :Tuple         =			hidden_size
        lowercase_      :Union[str, Any]         =			num_hidden_layers
        lowercase_      :int         =			num_attention_heads
        lowercase_      :List[Any]         =			intermediate_size
        lowercase_      :Tuple         =			hidden_act
        lowercase_      :str         =			hidden_dropout_prob
        lowercase_      :Any         =			attention_probs_dropout_prob
        lowercase_      :List[Any]         =			max_position_embeddings
        lowercase_      :Union[str, Any]         =			type_vocab_size
        lowercase_      :Union[str, Any]         =			type_sequence_label_size
        lowercase_      :Any         =			initializer_range
        lowercase_      :List[Any]         =			num_labels
        lowercase_      :str         =			num_choices
        lowercase_      :Optional[Any]         =			relative_attention
        lowercase_      :Tuple         =			position_biased_input
        lowercase_      :Union[str, Any]         =			pos_att_type
        lowercase_      :Tuple         =			scope





     def 							lowercase__				(			self		: Dict						):





        """simple docstring"""




        lowercase_      :List[str]         =			ids_tensor([self.batch_size, self.seq_length]     ,				self.vocab_size						)

        lowercase_      :Union[str, Any]         =			None
        if self.use_input_mask:
           lowercase_      :int         =			ids_tensor([self.batch_size, self.seq_length]     ,				vocab_size=2						)

        lowercase_      :List[Any]         =			None
        if self.use_token_type_ids:
           lowercase_      :Any         =			ids_tensor([self.batch_size, self.seq_length]     ,				self.type_vocab_size						)

        lowercase_      :str         =			None
        lowercase_      :Union[str, Any]         =			None
        lowercase_      :List[str]         =			None
        if self.use_labels:
           lowercase_      :List[Any]         =			ids_tensor([self.batch_size]     ,				self.type_sequence_label_size						)
           lowercase_      :Any         =			ids_tensor([self.batch_size, self.seq_length]     ,				self.num_labels						)
           lowercase_      :str         =			ids_tensor([self.batch_size]     ,				self.num_choices						)

        lowercase_      :Optional[Any]         =			self.get_config()

        return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels





     def 							lowercase__				(			self		: Optional[Any]						):





        """simple docstring"""




        return DebertaConfig(
            vocab_size=self.vocab_size     ,				hidden_size=self.hidden_size     ,				num_hidden_layers=self.num_hidden_layers     ,				num_attention_heads=self.num_attention_heads     ,				intermediate_size=self.intermediate_size     ,				hidden_act=self.hidden_act     ,				hidden_dropout_prob=self.hidden_dropout_prob     ,				attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,				max_position_embeddings=self.max_position_embeddings     ,				type_vocab_size=self.type_vocab_size     ,				initializer_range=self.initializer_range     ,				relative_attention=self.relative_attention     ,				position_biased_input=self.position_biased_input     ,				pos_att_type=self.pos_att_type     ,				)





     def 							lowercase__				(			self		: Optional[int]						):





        """simple docstring"""




        lowercase_      :Dict         =			self.get_config()
        lowercase_      :Optional[Any]         =			300
        return config





     def 							lowercase__				(			self		: Optional[Any]     ,				lowercase		: Dict						):





        """simple docstring"""




        self.parent.assertListEqual(list(result.loss.size()						)     ,				[]						)





     def 							lowercase__				(			self		: Optional[Any]     ,				lowercase		: Union[str, Any]     ,				lowercase		: List[Any]     ,				lowercase		: List[Any]     ,				lowercase		: Tuple     ,				lowercase		: str     ,				lowercase		: Optional[Any]     ,				lowercase		: Optional[int]						):





        """simple docstring"""




        lowercase_      :str         =			DebertaModel(config=lowercase						)
        model.to(lowercase						)
        model.eval()
        lowercase_      :Optional[int]         =			model(lowercase     ,				attention_mask=lowercase     ,				token_type_ids=lowercase						)[0]
        lowercase_      :Union[str, Any]         =			model(lowercase     ,				token_type_ids=lowercase						)[0]
        lowercase_      :Dict         =			model(lowercase						)[0]

        self.parent.assertListEqual(list(sequence_output.size()						)     ,				[self.batch_size, self.seq_length, self.hidden_size]						)





     def 							lowercase__				(			self		: Dict     ,				lowercase		: Union[str, Any]     ,				lowercase		: Optional[int]     ,				lowercase		: Tuple     ,				lowercase		: Dict     ,				lowercase		: Dict     ,				lowercase		: str     ,				lowercase		: Tuple						):





        """simple docstring"""




        lowercase_      :Dict         =			DebertaForMaskedLM(config=lowercase						)
        model.to(lowercase						)
        model.eval()
        lowercase_      :Optional[Any]         =			model(lowercase     ,				attention_mask=lowercase     ,				token_type_ids=lowercase     ,				labels=lowercase						)

        self.parent.assertEqual(result.logits.shape     ,				(self.batch_size, self.seq_length, self.vocab_size)						)





     def 							lowercase__				(			self		: Optional[int]     ,				lowercase		: Union[str, Any]     ,				lowercase		: Union[str, Any]     ,				lowercase		: Union[str, Any]     ,				lowercase		: Dict     ,				lowercase		: List[Any]     ,				lowercase		: int     ,				lowercase		: Dict						):





        """simple docstring"""




        lowercase_      :Dict         =			self.num_labels
        lowercase_      :int         =			DebertaForSequenceClassification(lowercase						)
        model.to(lowercase						)
        model.eval()
        lowercase_      :Union[str, Any]         =			model(lowercase     ,				attention_mask=lowercase     ,				token_type_ids=lowercase     ,				labels=lowercase						)
        self.parent.assertListEqual(list(result.logits.size()						)     ,				[self.batch_size, self.num_labels]						)
        self.check_loss_output(lowercase						)





     def 							lowercase__				(			self		: List[Any]     ,				lowercase		: List[str]     ,				lowercase		: Tuple     ,				lowercase		: Optional[Any]     ,				lowercase		: Optional[int]     ,				lowercase		: Union[str, Any]     ,				lowercase		: Dict     ,				lowercase		: int						):





        """simple docstring"""




        lowercase_      :List[str]         =			self.num_labels
        lowercase_      :Optional[int]         =			DebertaForTokenClassification(config=lowercase						)
        model.to(lowercase						)
        model.eval()
        lowercase_      :Dict         =			model(lowercase     ,				attention_mask=lowercase     ,				token_type_ids=lowercase     ,				labels=lowercase						)
        self.parent.assertEqual(result.logits.shape     ,				(self.batch_size, self.seq_length, self.num_labels)						)





     def 							lowercase__				(			self		: Union[str, Any]     ,				lowercase		: Tuple     ,				lowercase		: Any     ,				lowercase		: List[Any]     ,				lowercase		: List[Any]     ,				lowercase		: Union[str, Any]     ,				lowercase		: Optional[int]     ,				lowercase		: List[Any]						):





        """simple docstring"""




        lowercase_      :Any         =			DebertaForQuestionAnswering(config=lowercase						)
        model.to(lowercase						)
        model.eval()
        lowercase_      :List[Any]         =			model(
            lowercase     ,				attention_mask=lowercase     ,				token_type_ids=lowercase     ,				start_positions=lowercase     ,				end_positions=lowercase     ,				)
        self.parent.assertEqual(result.start_logits.shape     ,				(self.batch_size, self.seq_length)						)
        self.parent.assertEqual(result.end_logits.shape     ,				(self.batch_size, self.seq_length)						)







     def 							lowercase__				(			self		: Any						):





        """simple docstring"""




        lowercase_      :Optional[int]         =			self.prepare_config_and_inputs()
        (
            (
            lowercase_
        )      ,							(
            lowercase_
        )      ,							(
            lowercase_
        )      ,							(
            lowercase_
        )      ,							(
            lowercase_
        )      ,							(
            lowercase_
        )      ,							(
            lowercase_
        )      ,							
        )      :List[str]         =			config_and_inputs
        lowercase_      :Tuple         =			{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
        return config, inputs_dict







@require_torch
class 				a_       (    _lowerCAmelCase     ,						_lowerCAmelCase     ,						unittest.TestCase							):
     __A								=      (
         (
             DebertaModel,
             DebertaForMaskedLM,
             DebertaForSequenceClassification,
             DebertaForTokenClassification,
             DebertaForQuestionAnswering,
         )
         if is_torch_available()
         else ()
     )
     __A								=      (
         {
             "feature-extraction": DebertaModel,
             "fill-mask": DebertaForMaskedLM,
             "question-answering": DebertaForQuestionAnswering,
             "text-classification": DebertaForSequenceClassification,
             "token-classification": DebertaForTokenClassification,
             "zero-shot": DebertaForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )

     __A								=      True
     __A								=      False
     __A								=      False
     __A								=      False
     __A								=      False





     def 							lowercase__				(			self		: Optional[int]						):





        """simple docstring"""




        lowercase_      :List[Any]         =			DebertaModelTester(self						)
        lowercase_      :str         =			ConfigTester(self     ,				config_class=lowercase     ,				hidden_size=37						)





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





        """simple docstring"""




        self.config_tester.run_common_tests()





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





        """simple docstring"""




        lowercase_      :str         =			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_deberta_model(*lowercase						)





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





        """simple docstring"""




        lowercase_      :Optional[Any]         =			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_deberta_for_sequence_classification(*lowercase						)





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





        """simple docstring"""




        lowercase_      :Optional[int]         =			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_deberta_for_masked_lm(*lowercase						)





     def 							lowercase__				(			self		: int						):





        """simple docstring"""




        lowercase_      :Union[str, Any]         =			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_deberta_for_question_answering(*lowercase						)





     def 							lowercase__				(			self		: Tuple						):





        """simple docstring"""




        lowercase_      :Union[str, Any]         =			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_deberta_for_token_classification(*lowercase						)







     @slow
     def 							lowercase__				(			self		: Optional[int]						):





        """simple docstring"""




        for model_name in DEBERTA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           lowercase_      :Tuple         =			DebertaModel.from_pretrained(lowercase						)
           self.assertIsNotNone(lowercase						)







@require_torch
@require_sentencepiece
@require_tokenizers
class 				a_       (    unittest.TestCase							):





     @unittest.skip(reason="Model not available yet"						)
     def 							lowercase__				(			self		: Dict						):





        """simple docstring"""




        pass







     @slow
     def 							lowercase__				(			self		: List[Any]						):





        """simple docstring"""




        lowercase_      :Optional[Any]         =			DebertaModel.from_pretrained("microsoft/deberta-base"						)

        lowercase_      :Dict         =			torch.tensor([[0, 31_414, 232, 328, 740, 1_140, 12_695, 69, 46_078, 1_588, 2]]						)
        lowercase_      :Optional[Any]         =			torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]						)
        with torch.no_grad():
           lowercase_      :Optional[int]         =			model(lowercase     ,				attention_mask=lowercase						)[0]
        # compare the actual values for a slice.
        lowercase_      :List[Any]         =			torch.tensor(
            [[[-0.59_86, -0.80_55, -0.84_62], [1.44_84, -0.93_48, -0.80_59], [0.31_23, 0.00_32, -1.41_31]]]						)
        self.assertTrue(torch.allclose(output[:, 1:4, 1:4]     ,				lowercase     ,				atol=1e-4						)     ,				F'{output[:, 1:4, 1:4]}'						)


