'''simple docstring'''




from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import DeiTConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
						import tensorflow as tf

						from transformers import (
						    TFDeiTForImageClassification,
						    TFDeiTForImageClassificationWithTeacher,
						    TFDeiTForMaskedImageModeling,
						    TFDeiTModel,
						)
						from transformers.models.deit.modeling_tf_deit import TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
						from PIL import Image

						from transformers import DeiTImageProcessor





class     __magic_name__						:







			def __init__(	self				:  Any				,			lowercase_				:  str				,			lowercase_				:  List[str]=13				,			lowercase_				:  Optional[int]=30				,			lowercase_				:  Tuple=2				,			lowercase_				:  str=3				,			lowercase_				:  Tuple=True				,			lowercase_				:  List[Any]=True				,			lowercase_				:  Optional[int]=32				,			lowercase_				:  Optional[int]=2				,			lowercase_				:  int=4				,			lowercase_				:  Optional[Any]=37				,			lowercase_				:  Optional[Any]="gelu"				,			lowercase_				:  Optional[Any]=0.1				,			lowercase_				:  int=0.1				,			lowercase_				:  int=10				,			lowercase_				:  Optional[int]=0.02				,			lowercase_				:  Dict=3				,			lowercase_				:  Optional[int]=None				,			lowercase_				:  List[str]=2				,			):
										lowercase_					:			Optional[int]										=       parent
										lowercase_					:			Optional[Any]										=       batch_size
										lowercase_					:			List[str]										=       image_size
										lowercase_					:			Any										=       patch_size
										lowercase_					:			List[str]										=       num_channels
										lowercase_					:			Dict										=       is_training
										lowercase_					:			Union[str, Any]										=       use_labels
										lowercase_					:			List[str]										=       hidden_size
										lowercase_					:			Tuple										=       num_hidden_layers
										lowercase_					:			Union[str, Any]										=       num_attention_heads
										lowercase_					:			int										=       intermediate_size
										lowercase_					:			Tuple										=       hidden_act
										lowercase_					:			List[str]										=       hidden_dropout_prob
										lowercase_					:			Any										=       attention_probs_dropout_prob
										lowercase_					:			Tuple										=       type_sequence_label_size
										lowercase_					:			str										=       initializer_range
										lowercase_					:			str										=       scope
										lowercase_					:			Optional[Any]										=       encoder_stride

										# in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
										lowercase_					:			Tuple										=       (image_size // patch_size) ** 2
										lowercase_					:			Optional[Any]										=       num_patches + 2







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[Any]    ):
										lowercase_					:			Any										=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

										lowercase_					:			Optional[int]										=       None
										if self.use_labels:
																	lowercase_					:			Any										=       ids_tensor([self.batch_size]				,			self.type_sequence_label_size    )

										lowercase_					:			Union[str, Any]										=       self.get_config()

										return config, pixel_values, labels







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[str]    ):
										return DeiTConfig(
										    image_size=self.image_size				,			patch_size=self.patch_size				,			num_channels=self.num_channels				,			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				,			is_decoder=__a				,			initializer_range=self.initializer_range				,			encoder_stride=self.encoder_stride				,			)







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Union[str, Any]				,			lowercase_				:  str				,			lowercase_				:  Any				,			lowercase_				:  Dict    ):
										lowercase_					:			Tuple										=       TFDeiTModel(config=__a    )
										lowercase_					:			str										=       model(__a    )
										self.parent.assertEqual(result.last_hidden_state.shape				,			(self.batch_size, self.seq_length, self.hidden_size)    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  str				,			lowercase_				:  Optional[int]				,			lowercase_				:  List[str]				,			lowercase_				:  Tuple    ):
										lowercase_					:			Union[str, Any]										=       TFDeiTForMaskedImageModeling(config=__a    )
										lowercase_					:			Optional[Any]										=       model(__a    )
										self.parent.assertEqual(
										    result.reconstruction.shape				,			(self.batch_size, self.num_channels, self.image_size, self.image_size)    )

										# test greyscale images
										lowercase_					:			Optional[int]										=       1
										lowercase_					:			Union[str, Any]										=       TFDeiTForMaskedImageModeling(__a    )

										lowercase_					:			Any										=       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
										lowercase_					:			str										=       model(__a    )
										self.parent.assertEqual(result.reconstruction.shape				,			(self.batch_size, 1, self.image_size, self.image_size)    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]				,			lowercase_				:  List[str]				,			lowercase_				:  int				,			lowercase_				:  int    ):
										lowercase_					:			List[Any]										=       self.type_sequence_label_size
										lowercase_					:			Tuple										=       TFDeiTForImageClassification(__a    )
										lowercase_					:			Optional[int]										=       model(__a				,			labels=__a    )
										self.parent.assertEqual(result.logits.shape				,			(self.batch_size, self.type_sequence_label_size)    )

										# test greyscale images
										lowercase_					:			int										=       1
										lowercase_					:			Dict										=       TFDeiTForImageClassification(__a    )

										lowercase_					:			int										=       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
										lowercase_					:			Dict										=       model(__a				,			labels=__a    )
										self.parent.assertEqual(result.logits.shape				,			(self.batch_size, self.type_sequence_label_size)    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Dict    ):
										lowercase_					:			List[Any]										=       self.prepare_config_and_inputs()
										lowercase_						,						lowercase_						,						lowercase_					:			Union[str, Any]										=       config_and_inputs
										lowercase_					:			Union[str, Any]										=       {"""pixel_values""": pixel_values}
										return config, inputs_dict





@require_tf
class     __magic_name__						(    UpperCamelCase__,     UpperCamelCase__,     unittest.TestCase):
			UpperCamelCase__				=       (
			    (
			        TFDeiTModel,
			        TFDeiTForImageClassification,
			        TFDeiTForImageClassificationWithTeacher,
			        TFDeiTForMaskedImageModeling,
			    )
			    if is_tf_available()
			    else ()
			)
			UpperCamelCase__				=       (
			    {
			        """feature-extraction""": TFDeiTModel,
			        """image-classification""": (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
			    }
			    if is_tf_available()
			    else {}
			)

			UpperCamelCase__				=       False
			UpperCamelCase__				=       False
			UpperCamelCase__				=       False
			UpperCamelCase__				=       False







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]    ):
										lowercase_					:			Tuple										=       TFDeiTModelTester(self    )
										lowercase_					:			List[str]										=       ConfigTester(self				,			config_class=__a				,			has_text_modality=__a				,			hidden_size=37    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Union[str, Any]    ):
										self.config_tester.run_common_tests()







			@unittest.skip(reason="""DeiT does not use inputs_embeds"""    )
			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[str]    ):
										pass







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]    ):
										lowercase_						,						lowercase_					:			Optional[Any]										=       self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																	lowercase_					:			int										=       model_class(__a    )
																	self.assertIsInstance(model.get_input_embeddings()				,			(tf.keras.layers.Layer)    )
																	lowercase_					:			Optional[Any]										=       model.get_output_embeddings()
																	self.assertTrue(x is None or isinstance(__a				,			tf.keras.layers.Dense    )    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[Any]    ):
										lowercase_						,						lowercase_					:			Tuple										=       self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																	lowercase_					:			List[str]										=       model_class(__a    )
																	lowercase_					:			List[str]										=       inspect.signature(model.call    )
																	# signature.parameters is an OrderedDict => so arg_names order is deterministic
																	lowercase_					:			List[str]										=       [*signature.parameters.keys()]

																	lowercase_					:			str										=       ["""pixel_values"""]
																	self.assertListEqual(arg_names[:1]				,			__a    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Dict    ):
										lowercase_					:			Optional[int]										=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*__a    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Any    ):
										lowercase_					:			Union[str, Any]										=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_image_modeling(*__a    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Any    ):
										lowercase_					:			List[str]										=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*__a    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[Any]				,			lowercase_				:  Tuple				,			lowercase_				:  Tuple				,			lowercase_				:  Union[str, Any]=False    ):
										lowercase_					:			Optional[Any]										=       super()._prepare_for_class(__a				,			__a				,			return_labels=__a    )

										if return_labels:
																	if "labels" in inputs_dict and "labels" not in inspect.signature(model_class.call    ).parameters:
																								del inputs_dict["labels"]

										return inputs_dict







			@slow
			def 	SCREAMING_SNAKE_CASE_     (	self				:  int    ):
										for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	lowercase_					:			str										=       TFDeiTModel.from_pretrained(__a    )
																	self.assertIsNotNone(__a    )




def 							lowerCamelCase	(      )			->      Dict:
							lowercase_					:			List[str]										=       Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png""" )
							return image





@require_tf
@require_vision
class     __magic_name__						(    unittest.TestCase):







			@cached_property
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Union[str, Any]    ):
										return (
										    DeiTImageProcessor.from_pretrained("""facebook/deit-base-distilled-patch16-224"""    )
										    if is_vision_available()
										    else None
										)







			@slow
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[Any]    ):
										lowercase_					:			str										=       TFDeiTForImageClassificationWithTeacher.from_pretrained("""facebook/deit-base-distilled-patch16-224"""    )

										lowercase_					:			Optional[Any]										=       self.default_image_processor
										lowercase_					:			List[str]										=       prepare_img()
										lowercase_					:			Any										=       image_processor(images=__a				,			return_tensors="""tf"""    )

										# forward pass
										lowercase_					:			int										=       model(**__a    )

										# verify the logits
										lowercase_					:			int										=       tf.TensorShape((1, 1000)    )
										self.assertEqual(outputs.logits.shape				,			__a    )

										lowercase_					:			int										=       tf.constant([-1.02_66, 0.19_12, -1.28_61]    )

										self.assertTrue(np.allclose(outputs.logits[0, :3]				,			__a				,			atol=1E-4    )    )






'''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 	lowerCAmelCase_				(					snake_case_       : List[Any]					,    snake_case_       : int			)    ->      Optional[int]:


     '''simple docstring'''




     UpperCAmelCase_           =   checkpoint

     UpperCAmelCase_           =   {}

     UpperCAmelCase_           =   vae_state_dict["encoder.conv_in.weight"]
     UpperCAmelCase_           =   vae_state_dict["encoder.conv_in.bias"]
     UpperCAmelCase_           =   vae_state_dict["encoder.conv_out.weight"]
     UpperCAmelCase_           =   vae_state_dict["encoder.conv_out.bias"]
     UpperCAmelCase_           =   vae_state_dict["encoder.norm_out.weight"]
     UpperCAmelCase_           =   vae_state_dict["encoder.norm_out.bias"]

     UpperCAmelCase_           =   vae_state_dict["decoder.conv_in.weight"]
     UpperCAmelCase_           =   vae_state_dict["decoder.conv_in.bias"]
     UpperCAmelCase_           =   vae_state_dict["decoder.conv_out.weight"]
     UpperCAmelCase_           =   vae_state_dict["decoder.conv_out.bias"]
     UpperCAmelCase_           =   vae_state_dict["decoder.norm_out.weight"]
     UpperCAmelCase_           =   vae_state_dict["decoder.norm_out.bias"]

     UpperCAmelCase_           =   vae_state_dict["quant_conv.weight"]
     UpperCAmelCase_           =   vae_state_dict["quant_conv.bias"]
     UpperCAmelCase_           =   vae_state_dict["post_quant_conv.weight"]
     UpperCAmelCase_           =   vae_state_dict["post_quant_conv.bias"]

     # Retrieves the keys for the encoder down blocks only
     UpperCAmelCase_           =   len({".".join(layer.split("."			)[:3]			) for layer in vae_state_dict if "encoder.down" in layer}			)
     UpperCAmelCase_           =   {
         layer_id: [key for key in vae_state_dict if f"""down.{layer_id}""" in key] for layer_id in range(snake_case_			)
     }

     # Retrieves the keys for the decoder up blocks only
     UpperCAmelCase_           =   len({".".join(layer.split("."			)[:3]			) for layer in vae_state_dict if "decoder.up" in layer}			)
     UpperCAmelCase_           =   {
         layer_id: [key for key in vae_state_dict if f"""up.{layer_id}""" in key] for layer_id in range(snake_case_			)
     }

     for i in range(snake_case_			):
          UpperCAmelCase_           =   [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:
               UpperCAmelCase_           =   vae_state_dict.pop(
                   f"""encoder.down.{i}.downsample.conv.weight"""			)
               UpperCAmelCase_           =   vae_state_dict.pop(
                   f"""encoder.down.{i}.downsample.conv.bias"""			)

          UpperCAmelCase_           =   renew_vae_resnet_paths(snake_case_			)
          UpperCAmelCase_           =   {"old": f"""down.{i}.block""", "new": f"""down_blocks.{i}.resnets"""}
          assign_to_checkpoint(snake_case_					,    snake_case_					,    snake_case_					,    additional_replacements=[meta_path]					,    config=snake_case_			)

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

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

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

     for i in range(snake_case_			):
          UpperCAmelCase_           =   num_up_blocks - 1 - i
          UpperCAmelCase_           =   [
              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:
               UpperCAmelCase_           =   vae_state_dict[
                   f"""decoder.up.{block_id}.upsample.conv.weight"""
               ]
               UpperCAmelCase_           =   vae_state_dict[
                   f"""decoder.up.{block_id}.upsample.conv.bias"""
               ]

          UpperCAmelCase_           =   renew_vae_resnet_paths(snake_case_			)
          UpperCAmelCase_           =   {"old": f"""up.{block_id}.block""", "new": f"""up_blocks.{i}.resnets"""}
          assign_to_checkpoint(snake_case_					,    snake_case_					,    snake_case_					,    additional_replacements=[meta_path]					,    config=snake_case_			)

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

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

     UpperCAmelCase_           =   [key for key in vae_state_dict if "decoder.mid.attn" in key]
     UpperCAmelCase_           =   renew_vae_attention_paths(snake_case_			)
     UpperCAmelCase_           =   {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
     assign_to_checkpoint(snake_case_					,    snake_case_					,    snake_case_					,    additional_replacements=[meta_path]					,    config=snake_case_			)
     conv_attn_to_linear(snake_case_			)
     return new_checkpoint





def 	lowerCAmelCase_				(					snake_case_       : str					,    snake_case_       : str					,    )    ->      Dict:


     '''simple docstring'''




     UpperCAmelCase_           =   requests.get(
         " https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml"			)
     UpperCAmelCase_           =   io.BytesIO(r.content			)

     UpperCAmelCase_           =   OmegaConf.load(snake_case_			)
     UpperCAmelCase_           =   5_12
     UpperCAmelCase_           =   "cuda" if torch.cuda.is_available() else "cpu"
     if checkpoint_path.endswith("safetensors"			):
          from safetensors import safe_open

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

     # Convert the VAE model.
     UpperCAmelCase_           =   create_vae_diffusers_config(snake_case_					,    image_size=snake_case_			)
     UpperCAmelCase_           =   custom_convert_ldm_vae_checkpoint(snake_case_					,    snake_case_			)

     UpperCAmelCase_           =   AutoencoderKL(**snake_case_			)
     vae.load_state_dict(snake_case_			)
     vae.save_pretrained(snake_case_			)


if __name__ == "__main__":
    SCREAMING_SNAKE_CASE_:     Optional[int]  =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.')

    SCREAMING_SNAKE_CASE_:     str  =parser.parse_args()

    vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)




