

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    SCREAMING_SNAKE_CASE__       (  lowerCamelCase__	, lowerCamelCase__			)  ->     List[str]:
	__lowerCamelCase				:				str					       = checkpoint

	__lowerCamelCase				:				Optional[Any]					       = {}

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

	__lowerCamelCase				:				Optional[Any]					       = vae_state_dict['decoder.conv_in.weight']
	__lowerCamelCase				:				Optional[int]					       = vae_state_dict['decoder.conv_in.bias']
	__lowerCamelCase				:				Union[str, Any]					       = vae_state_dict['decoder.conv_out.weight']
	__lowerCamelCase				:				Any					       = vae_state_dict['decoder.conv_out.bias']
	__lowerCamelCase				:				int					       = vae_state_dict['decoder.norm_out.weight']
	__lowerCamelCase				:				Optional[int]					       = vae_state_dict['decoder.norm_out.bias']

	__lowerCamelCase				:				Dict					       = vae_state_dict['quant_conv.weight']
	__lowerCamelCase				:				Optional[Any]					       = vae_state_dict['quant_conv.bias']
	__lowerCamelCase				:				Optional[int]					       = vae_state_dict['post_quant_conv.weight']
	__lowerCamelCase				:				str					       = vae_state_dict['post_quant_conv.bias']

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

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

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

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

		__lowerCamelCase				:				str					       = renew_vae_resnet_paths(lowerCamelCase__			)
		__lowerCamelCase				:				Optional[int]					       = {'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				:				List[str]					       = [key for key in vae_state_dict if 'decoder.mid.block' in key]
	__lowerCamelCase				:				List[Any]					       = 2
	for i in range(1	, num_mid_res_blocks + 1			):
		__lowerCamelCase				:				Optional[Any]					       = [key for key in mid_resnets if F"decoder.mid.block_{i}" in key]

		__lowerCamelCase				:				Optional[Any]					       = renew_vae_resnet_paths(lowerCamelCase__			)
		__lowerCamelCase				:				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__			)

	__lowerCamelCase				:				Union[str, Any]					       = [key for key in vae_state_dict if 'decoder.mid.attn' in key]
	__lowerCamelCase				:				Optional[int]					       = renew_vae_attention_paths(lowerCamelCase__			)
	__lowerCamelCase				:				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__			)
	return new_checkpoint






def    SCREAMING_SNAKE_CASE__       (  lowerCamelCase__	, lowerCamelCase__	, )  ->     str:
	# Only support V1
	__lowerCamelCase				:				Tuple					       = requests.get(
	    ' https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml'			)
	__lowerCamelCase				:				List[Any]					       = io.BytesIO(r.content			)

	__lowerCamelCase				:				int					       = OmegaConf.load(lowerCamelCase__			)
	__lowerCamelCase				:				int					       = 5_1_2
	__lowerCamelCase				:				Optional[Any]					       = 'cuda' if torch.cuda.is_available() else 'cpu'
	if checkpoint_path.endswith('safetensors'			):
		from safetensors import safe_open

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

	# Convert the VAE model.
	__lowerCamelCase				:				Dict					       = create_vae_diffusers_config(lowerCamelCase__	, image_size=lowerCamelCase__			)
	__lowerCamelCase				:				Tuple					       = custom_convert_ldm_vae_checkpoint(lowerCamelCase__	, lowerCamelCase__			)

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


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

					a		=parser.parse_args()

					vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)





import unittest

from transformers import MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING, is_vision_available, pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
					from PIL import Image
else:





					class 					lowercase__						:


											@staticmethod
											def        __A       (       *UpperCamelCase__	:				Union[str, Any]    ,		**UpperCamelCase__	:				Optional[Any]						):





																	'''simple docstring'''





																	pass





@is_pipeline_test
@require_vision
@require_torch
class 					lowercase__						(   unittest.TestCase):
						UpperCamelCase_     =		MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING


						def        __A       (       self	:				Union[str, Any]    ,		UpperCamelCase__	:				Tuple    ,		UpperCamelCase__	:				List[Any]    ,		UpperCamelCase__	:				int						):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:    List[Any]							=							pipeline(
												    '''zero-shot-object-detection'''    ,		model='''hf-internal-testing/tiny-random-owlvit-object-detection'''						)

												SCREAMING_SNAKE_CASE							:    List[str]							=							[
												    {
												        '''image''': '''./tests/fixtures/tests_samples/COCO/000000039769.png''',
												        '''candidate_labels''': ['''cat''', '''remote''', '''couch'''],
												    }
												]
												return object_detector, examples


						def        __A       (       self	:				Union[str, Any]    ,		UpperCamelCase__	:				Any    ,		UpperCamelCase__	:				Optional[Any]						):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:    Optional[int]							=							object_detector(examples[0]    ,		threshold=0.0						)

												SCREAMING_SNAKE_CASE							:    Tuple							=							len(UpperCamelCase__						)
												self.assertGreater(UpperCamelCase__    ,		0						)
												self.assertEqual(
												    UpperCamelCase__    ,		[
												        {
												            '''score''': ANY(UpperCamelCase__						),
												            '''label''': ANY(UpperCamelCase__						),
												            '''box''': {'''xmin''': ANY(UpperCamelCase__						), '''ymin''': ANY(UpperCamelCase__						), '''xmax''': ANY(UpperCamelCase__						), '''ymax''': ANY(UpperCamelCase__						)},
												        }
												        for i in range(UpperCamelCase__						)
												    ]    ,		)


						@require_tf
						@unittest.skip('''Zero Shot Object Detection not implemented in TF'''						)
						def        __A       (       self	:				Optional[Any]						):





												'''simple docstring'''





												pass


						@require_torch
						def        __A       (       self	:				Dict						):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:    Tuple							=							pipeline(
												    '''zero-shot-object-detection'''    ,		model='''hf-internal-testing/tiny-random-owlvit-object-detection'''						)

												SCREAMING_SNAKE_CASE							:    str							=							object_detector(
												    '''./tests/fixtures/tests_samples/COCO/000000039769.png'''    ,		candidate_labels=['''cat''', '''remote''', '''couch''']    ,		threshold=0.64    ,		)

												self.assertEqual(
												    nested_simplify(UpperCamelCase__    ,		decimals=4						)    ,		[
												        {'''score''': 0.7235, '''label''': '''cat''', '''box''': {'''xmin''': 204, '''ymin''': 167, '''xmax''': 232, '''ymax''': 190}},
												        {'''score''': 0.7218, '''label''': '''remote''', '''box''': {'''xmin''': 204, '''ymin''': 167, '''xmax''': 232, '''ymax''': 190}},
												        {'''score''': 0.7184, '''label''': '''couch''', '''box''': {'''xmin''': 204, '''ymin''': 167, '''xmax''': 232, '''ymax''': 190}},
												        {'''score''': 0.6748, '''label''': '''remote''', '''box''': {'''xmin''': 571, '''ymin''': 83, '''xmax''': 598, '''ymax''': 103}},
												        {'''score''': 0.6656, '''label''': '''cat''', '''box''': {'''xmin''': 571, '''ymin''': 83, '''xmax''': 598, '''ymax''': 103}},
												        {'''score''': 0.6614, '''label''': '''couch''', '''box''': {'''xmin''': 571, '''ymin''': 83, '''xmax''': 598, '''ymax''': 103}},
												        {'''score''': 0.6456, '''label''': '''remote''', '''box''': {'''xmin''': 494, '''ymin''': 105, '''xmax''': 521, '''ymax''': 127}},
												        {'''score''': 0.642, '''label''': '''remote''', '''box''': {'''xmin''': 67, '''ymin''': 274, '''xmax''': 93, '''ymax''': 297}},
												        {'''score''': 0.6419, '''label''': '''cat''', '''box''': {'''xmin''': 494, '''ymin''': 105, '''xmax''': 521, '''ymax''': 127}},
												    ]    ,		)

												SCREAMING_SNAKE_CASE							:    str							=							object_detector(
												    [
												        {
												            '''image''': '''./tests/fixtures/tests_samples/COCO/000000039769.png''',
												            '''candidate_labels''': ['''cat''', '''remote''', '''couch'''],
												        }
												    ]    ,		threshold=0.64    ,		)

												self.assertEqual(
												    nested_simplify(UpperCamelCase__    ,		decimals=4						)    ,		[
												        [
												            {'''score''': 0.7235, '''label''': '''cat''', '''box''': {'''xmin''': 204, '''ymin''': 167, '''xmax''': 232, '''ymax''': 190}},
												            {'''score''': 0.7218, '''label''': '''remote''', '''box''': {'''xmin''': 204, '''ymin''': 167, '''xmax''': 232, '''ymax''': 190}},
												            {'''score''': 0.7184, '''label''': '''couch''', '''box''': {'''xmin''': 204, '''ymin''': 167, '''xmax''': 232, '''ymax''': 190}},
												            {'''score''': 0.6748, '''label''': '''remote''', '''box''': {'''xmin''': 571, '''ymin''': 83, '''xmax''': 598, '''ymax''': 103}},
												            {'''score''': 0.6656, '''label''': '''cat''', '''box''': {'''xmin''': 571, '''ymin''': 83, '''xmax''': 598, '''ymax''': 103}},
												            {'''score''': 0.6614, '''label''': '''couch''', '''box''': {'''xmin''': 571, '''ymin''': 83, '''xmax''': 598, '''ymax''': 103}},
												            {'''score''': 0.6456, '''label''': '''remote''', '''box''': {'''xmin''': 494, '''ymin''': 105, '''xmax''': 521, '''ymax''': 127}},
												            {'''score''': 0.642, '''label''': '''remote''', '''box''': {'''xmin''': 67, '''ymin''': 274, '''xmax''': 93, '''ymax''': 297}},
												            {'''score''': 0.6419, '''label''': '''cat''', '''box''': {'''xmin''': 494, '''ymin''': 105, '''xmax''': 521, '''ymax''': 127}},
												        ]
												    ]    ,		)


						@require_torch
						@slow
						def        __A       (       self	:				Tuple						):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:    int							=							pipeline('''zero-shot-object-detection'''						)

												SCREAMING_SNAKE_CASE							:    Optional[int]							=							object_detector(
												    '''http://images.cocodataset.org/val2017/000000039769.jpg'''    ,		candidate_labels=['''cat''', '''remote''', '''couch''']    ,		)
												self.assertEqual(
												    nested_simplify(UpperCamelCase__    ,		decimals=4						)    ,		[
												        {'''score''': 0.2868, '''label''': '''cat''', '''box''': {'''xmin''': 324, '''ymin''': 20, '''xmax''': 640, '''ymax''': 373}},
												        {'''score''': 0.277, '''label''': '''remote''', '''box''': {'''xmin''': 40, '''ymin''': 72, '''xmax''': 177, '''ymax''': 115}},
												        {'''score''': 0.2537, '''label''': '''cat''', '''box''': {'''xmin''': 1, '''ymin''': 55, '''xmax''': 315, '''ymax''': 472}},
												        {'''score''': 0.1474, '''label''': '''remote''', '''box''': {'''xmin''': 335, '''ymin''': 74, '''xmax''': 371, '''ymax''': 187}},
												        {'''score''': 0.1208, '''label''': '''couch''', '''box''': {'''xmin''': 4, '''ymin''': 0, '''xmax''': 642, '''ymax''': 476}},
												    ]    ,		)

												SCREAMING_SNAKE_CASE							:    int							=							object_detector(
												    [
												        {
												            '''image''': '''http://images.cocodataset.org/val2017/000000039769.jpg''',
												            '''candidate_labels''': ['''cat''', '''remote''', '''couch'''],
												        },
												        {
												            '''image''': '''http://images.cocodataset.org/val2017/000000039769.jpg''',
												            '''candidate_labels''': ['''cat''', '''remote''', '''couch'''],
												        },
												    ]    ,		)
												self.assertEqual(
												    nested_simplify(UpperCamelCase__    ,		decimals=4						)    ,		[
												        [
												            {'''score''': 0.2868, '''label''': '''cat''', '''box''': {'''xmin''': 324, '''ymin''': 20, '''xmax''': 640, '''ymax''': 373}},
												            {'''score''': 0.277, '''label''': '''remote''', '''box''': {'''xmin''': 40, '''ymin''': 72, '''xmax''': 177, '''ymax''': 115}},
												            {'''score''': 0.2537, '''label''': '''cat''', '''box''': {'''xmin''': 1, '''ymin''': 55, '''xmax''': 315, '''ymax''': 472}},
												            {'''score''': 0.1474, '''label''': '''remote''', '''box''': {'''xmin''': 335, '''ymin''': 74, '''xmax''': 371, '''ymax''': 187}},
												            {'''score''': 0.1208, '''label''': '''couch''', '''box''': {'''xmin''': 4, '''ymin''': 0, '''xmax''': 642, '''ymax''': 476}},
												        ],
												        [
												            {'''score''': 0.2868, '''label''': '''cat''', '''box''': {'''xmin''': 324, '''ymin''': 20, '''xmax''': 640, '''ymax''': 373}},
												            {'''score''': 0.277, '''label''': '''remote''', '''box''': {'''xmin''': 40, '''ymin''': 72, '''xmax''': 177, '''ymax''': 115}},
												            {'''score''': 0.2537, '''label''': '''cat''', '''box''': {'''xmin''': 1, '''ymin''': 55, '''xmax''': 315, '''ymax''': 472}},
												            {'''score''': 0.1474, '''label''': '''remote''', '''box''': {'''xmin''': 335, '''ymin''': 74, '''xmax''': 371, '''ymax''': 187}},
												            {'''score''': 0.1208, '''label''': '''couch''', '''box''': {'''xmin''': 4, '''ymin''': 0, '''xmax''': 642, '''ymax''': 476}},
												        ],
												    ]    ,		)


						@require_tf
						@unittest.skip('''Zero Shot Object Detection not implemented in TF'''						)
						def        __A       (       self	:				str						):





												'''simple docstring'''





												pass


						@require_torch
						@slow
						def        __A       (       self	:				List[str]						):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:    Tuple							=							0.2
												SCREAMING_SNAKE_CASE							:    Optional[int]							=							pipeline('''zero-shot-object-detection'''						)

												SCREAMING_SNAKE_CASE							:    Dict							=							object_detector(
												    '''http://images.cocodataset.org/val2017/000000039769.jpg'''    ,		candidate_labels=['''cat''', '''remote''', '''couch''']    ,		threshold=UpperCamelCase__    ,		)
												self.assertEqual(
												    nested_simplify(UpperCamelCase__    ,		decimals=4						)    ,		[
												        {'''score''': 0.2868, '''label''': '''cat''', '''box''': {'''xmin''': 324, '''ymin''': 20, '''xmax''': 640, '''ymax''': 373}},
												        {'''score''': 0.277, '''label''': '''remote''', '''box''': {'''xmin''': 40, '''ymin''': 72, '''xmax''': 177, '''ymax''': 115}},
												        {'''score''': 0.2537, '''label''': '''cat''', '''box''': {'''xmin''': 1, '''ymin''': 55, '''xmax''': 315, '''ymax''': 472}},
												    ]    ,		)



						@require_torch
						@slow
						def        __A       (       self	:				List[Any]						):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:    List[str]							=							2
												SCREAMING_SNAKE_CASE							:    Optional[Any]							=							pipeline('''zero-shot-object-detection'''						)

												SCREAMING_SNAKE_CASE							:    List[str]							=							object_detector(
												    '''http://images.cocodataset.org/val2017/000000039769.jpg'''    ,		candidate_labels=['''cat''', '''remote''', '''couch''']    ,		top_k=UpperCamelCase__    ,		)
												self.assertEqual(
												    nested_simplify(UpperCamelCase__    ,		decimals=4						)    ,		[
												        {'''score''': 0.2868, '''label''': '''cat''', '''box''': {'''xmin''': 324, '''ymin''': 20, '''xmax''': 640, '''ymax''': 373}},
												        {'''score''': 0.277, '''label''': '''remote''', '''box''': {'''xmin''': 40, '''ymin''': 72, '''xmax''': 177, '''ymax''': 115}},
												    ]    ,		)





