
import argparse
import os.path as osp
import re

import torch
from safetensors.torch import load_file, save_file


# =================#
# UNet Conversion #
# =================#

_lowerCamelCase	:    Dict          =      [
    # (stable-diffusion, HF Diffusers)
    ('''time_embed.0.weight''', '''time_embedding.linear_1.weight'''),
    ('''time_embed.0.bias''', '''time_embedding.linear_1.bias'''),
    ('''time_embed.2.weight''', '''time_embedding.linear_2.weight'''),
    ('''time_embed.2.bias''', '''time_embedding.linear_2.bias'''),
    ('''input_blocks.0.0.weight''', '''conv_in.weight'''),
    ('''input_blocks.0.0.bias''', '''conv_in.bias'''),
    ('''out.0.weight''', '''conv_norm_out.weight'''),
    ('''out.0.bias''', '''conv_norm_out.bias'''),
    ('''out.2.weight''', '''conv_out.weight'''),
    ('''out.2.bias''', '''conv_out.bias'''),
]

_lowerCamelCase	:    Dict          =      [
    # (stable-diffusion, HF Diffusers)
    ('''in_layers.0''', '''norm1'''),
    ('''in_layers.2''', '''conv1'''),
    ('''out_layers.0''', '''norm2'''),
    ('''out_layers.3''', '''conv2'''),
    ('''emb_layers.1''', '''time_emb_proj'''),
    ('''skip_connection''', '''conv_shortcut'''),
]

_lowerCamelCase	:    List[str]          =      []
# hardcoded number of downblocks and resnets/attentions...
# would need smarter logic for other networks.
for i in range(4):
			# loop over downblocks/upblocks

			for j in range(2):
						# loop over resnets/attentions for downblocks
						_lowerCamelCase	:    Tuple          =      F'down_blocks.{i}.resnets.{j}.'
						_lowerCamelCase	:    List[str]          =      F'input_blocks.{3*i + j + 1}.0.'
						unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))

						if i < 3:
									# no attention layers in down_blocks.3
									_lowerCamelCase	:    Optional[Any]          =      F'down_blocks.{i}.attentions.{j}.'
									_lowerCamelCase	:    Optional[Any]          =      F'input_blocks.{3*i + j + 1}.1.'
									unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))

			for j in range(3):
						# loop over resnets/attentions for upblocks
						_lowerCamelCase	:    Tuple          =      F'up_blocks.{i}.resnets.{j}.'
						_lowerCamelCase	:    int          =      F'output_blocks.{3*i + j}.0.'
						unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))

						if i > 0:
									# no attention layers in up_blocks.0
									_lowerCamelCase	:    Any          =      F'up_blocks.{i}.attentions.{j}.'
									_lowerCamelCase	:    List[Any]          =      F'output_blocks.{3*i + j}.1.'
									unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))

			if i < 3:
						# no downsample in down_blocks.3
						_lowerCamelCase	:    Any          =      F'down_blocks.{i}.downsamplers.0.conv.'
						_lowerCamelCase	:    Any          =      F'input_blocks.{3*(i+1)}.0.op.'
						unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))

						# no upsample in up_blocks.3
						_lowerCamelCase	:    Optional[Any]          =      F'up_blocks.{i}.upsamplers.0.'
						_lowerCamelCase	:    List[Any]          =      F'output_blocks.{3*i + 2}.{1 if i == 0 else 2}.'
						unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))

_lowerCamelCase	:    Optional[Any]          =      '''mid_block.attentions.0.'''
_lowerCamelCase	:    Optional[int]          =      '''middle_block.1.'''
unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))

for j in range(2):
			_lowerCamelCase	:    Optional[int]          =      F'mid_block.resnets.{j}.'
			_lowerCamelCase	:    Tuple          =      F'middle_block.{2*j}.'
			unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))






def 	a_		(	__lowercase       :     List[str]			)			->  Any:
			# buyer beware: this is a *brittle* function,
			# and correct output requires that all of these pieces interact in
			# the exact order in which I have arranged them.
			_snake_case          =      {k: k for k in unet_state_dict.keys()}
			for sd_name, hf_name in unet_conversion_map:
						_snake_case          =      sd_name
			for k, v in mapping.items():
						if "resnets" in k:
									for sd_part, hf_part in unet_conversion_map_resnet:
												_snake_case          =      v.replace(__lowercase	, __lowercase			)
									_snake_case          =      v
			for k, v in mapping.items():
						for sd_part, hf_part in unet_conversion_map_layer:
									_snake_case          =      v.replace(__lowercase	, __lowercase			)
						_snake_case          =      v
			_snake_case          =      {v: unet_state_dict[k] for k, v in mapping.items()}
			return new_state_dict


# ================#
# VAE Conversion #
# ================#

_lowerCamelCase	:    int          =      [
    # (stable-diffusion, HF Diffusers)
    ('''nin_shortcut''', '''conv_shortcut'''),
    ('''norm_out''', '''conv_norm_out'''),
    ('''mid.attn_1.''', '''mid_block.attentions.0.'''),
]

for i in range(4):
			# down_blocks have two resnets
			for j in range(2):
						_lowerCamelCase	:    Dict          =      F'encoder.down_blocks.{i}.resnets.{j}.'
						_lowerCamelCase	:    List[Any]          =      F'encoder.down.{i}.block.{j}.'
						vae_conversion_map.append((sd_down_prefix, hf_down_prefix))

			if i < 3:
						_lowerCamelCase	:    Dict          =      F'down_blocks.{i}.downsamplers.0.'
						_lowerCamelCase	:    Union[str, Any]          =      F'down.{i}.downsample.'
						vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))

						_lowerCamelCase	:    Dict          =      F'up_blocks.{i}.upsamplers.0.'
						_lowerCamelCase	:    Union[str, Any]          =      F'up.{3-i}.upsample.'
						vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))

			# up_blocks have three resnets
			# also, up blocks in hf are numbered in reverse from sd
			for j in range(3):
						_lowerCamelCase	:    Dict          =      F'decoder.up_blocks.{i}.resnets.{j}.'
						_lowerCamelCase	:    Optional[int]          =      F'decoder.up.{3-i}.block.{j}.'
						vae_conversion_map.append((sd_up_prefix, hf_up_prefix))

# this part accounts for mid blocks in both the encoder and the decoder
for i in range(2):
			_lowerCamelCase	:    Union[str, Any]          =      F'mid_block.resnets.{i}.'
			_lowerCamelCase	:    Tuple          =      F'mid.block_{i+1}.'
			vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))


_lowerCamelCase	:    int          =      [
    # (stable-diffusion, HF Diffusers)
    ('''norm.''', '''group_norm.'''),
    ('''q.''', '''query.'''),
    ('''k.''', '''key.'''),
    ('''v.''', '''value.'''),
    ('''proj_out.''', '''proj_attn.'''),
]






def 	a_		(	__lowercase       :     str			)			->  str:
			# convert HF linear weights to SD conv2d weights
			return w.reshape(*w.shape	, 1	, 1			)






def 	a_		(	__lowercase       :     List[str]			)			->  List[Any]:
			_snake_case          =      {k: k for k in vae_state_dict.keys()}
			for k, v in mapping.items():
						for sd_part, hf_part in vae_conversion_map:
									_snake_case          =      v.replace(__lowercase	, __lowercase			)
						_snake_case          =      v
			for k, v in mapping.items():
						if "attentions" in k:
									for sd_part, hf_part in vae_conversion_map_attn:
												_snake_case          =      v.replace(__lowercase	, __lowercase			)
									_snake_case          =      v
			_snake_case          =      {v: vae_state_dict[k] for k, v in mapping.items()}
			_snake_case          =      ['q', 'k', 'v', 'proj_out']
			for k, v in new_state_dict.items():
						for weight_name in weights_to_convert:
									if f'''mid.attn_1.{weight_name}.weight''' in k:
												print(f'''Reshaping {k} for SD format'''			)
												_snake_case          =      reshape_weight_for_sd(__lowercase			)
			return new_state_dict


# =========================#
# Text Encoder Conversion #
# =========================#


_lowerCamelCase	:    Dict          =      [
    # (stable-diffusion, HF Diffusers)
    ('''resblocks.''', '''text_model.encoder.layers.'''),
    ('''ln_1''', '''layer_norm1'''),
    ('''ln_2''', '''layer_norm2'''),
    ('''.c_fc.''', '''.fc1.'''),
    ('''.c_proj.''', '''.fc2.'''),
    ('''.attn''', '''.self_attn'''),
    ('''ln_final.''', '''transformer.text_model.final_layer_norm.'''),
    ('''token_embedding.weight''', '''transformer.text_model.embeddings.token_embedding.weight'''),
    ('''positional_embedding''', '''transformer.text_model.embeddings.position_embedding.weight'''),
]
_lowerCamelCase	:    Tuple          =      {re.escape(x[1]): x[0] for x in textenc_conversion_lst}
_lowerCamelCase	:    Union[str, Any]          =      re.compile('''|'''.join(protected.keys()))

# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
_lowerCamelCase	:    int          =      {'''q''': 0, '''k''': 1, '''v''': 2}






def 	a_		(	__lowercase       :     Optional[Any]			)			->  int:
			_snake_case          =      {}
			_snake_case          =      {}
			_snake_case          =      {}
			for k, v in text_enc_dict.items():
						if (
						    k.endswith('.self_attn.q_proj.weight'			)
						    or k.endswith('.self_attn.k_proj.weight'			)
						    or k.endswith('.self_attn.v_proj.weight'			)
						):
									_snake_case          =      k[: -len('.q_proj.weight'			)]
									_snake_case          =      k[-len('q_proj.weight'			)]
									if k_pre not in capture_qkv_weight:
												_snake_case          =      [None, None, None]
									_snake_case          =      v
									continue

						if (
						    k.endswith('.self_attn.q_proj.bias'			)
						    or k.endswith('.self_attn.k_proj.bias'			)
						    or k.endswith('.self_attn.v_proj.bias'			)
						):
									_snake_case          =      k[: -len('.q_proj.bias'			)]
									_snake_case          =      k[-len('q_proj.bias'			)]
									if k_pre not in capture_qkv_bias:
												_snake_case          =      [None, None, None]
									_snake_case          =      v
									continue

						_snake_case          =      textenc_pattern.sub(lambda __lowercase			: protected[re.escape(m.group(0			)			)]	, __lowercase			)
						_snake_case          =      v

			for k_pre, tensors in capture_qkv_weight.items():
						if None in tensors:
									raise Exception('CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing'			)
						_snake_case          =      textenc_pattern.sub(lambda __lowercase			: protected[re.escape(m.group(0			)			)]	, __lowercase			)
						_snake_case          =      torch.cat(__lowercase			)

			for k_pre, tensors in capture_qkv_bias.items():
						if None in tensors:
									raise Exception('CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing'			)
						_snake_case          =      textenc_pattern.sub(lambda __lowercase			: protected[re.escape(m.group(0			)			)]	, __lowercase			)
						_snake_case          =      torch.cat(__lowercase			)

			return new_state_dict






def 	a_		(	__lowercase       :     Optional[int]			)			->  Union[str, Any]:
			return text_enc_dict


if __name__ == "__main__":
			_lowerCamelCase	:    Any          =      argparse.ArgumentParser()

			parser.add_argument('''--model_path''', default=None, type=str, required=True, help='''Path to the model to convert.''')
			parser.add_argument('''--checkpoint_path''', default=None, type=str, required=True, help='''Path to the output model.''')
			parser.add_argument('''--half''', action='''store_true''', help='''Save weights in half precision.''')
			parser.add_argument(
			    '''--use_safetensors''', action='''store_true''', help='''Save weights use safetensors, default is ckpt.'''
			)

			_lowerCamelCase	:    str          =      parser.parse_args()

			assert args.model_path is not None, "Must provide a model path!"

			assert args.checkpoint_path is not None, "Must provide a checkpoint path!"

			# Path for safetensors
			_lowerCamelCase	:    int          =      osp.join(args.model_path, '''unet''', '''diffusion_pytorch_model.safetensors''')
			_lowerCamelCase	:    int          =      osp.join(args.model_path, '''vae''', '''diffusion_pytorch_model.safetensors''')
			_lowerCamelCase	:    str          =      osp.join(args.model_path, '''text_encoder''', '''model.safetensors''')

			# Load models from safetensors if it exists, if it doesn't pytorch
			if osp.exists(unet_path):
						_lowerCamelCase	:    Any          =      load_file(unet_path, device='''cpu''')
			else:
						_lowerCamelCase	:    str          =      osp.join(args.model_path, '''unet''', '''diffusion_pytorch_model.bin''')
						_lowerCamelCase	:    Optional[Any]          =      torch.load(unet_path, map_location='''cpu''')

			if osp.exists(vae_path):
						_lowerCamelCase	:    Dict          =      load_file(vae_path, device='''cpu''')
			else:
						_lowerCamelCase	:    Optional[int]          =      osp.join(args.model_path, '''vae''', '''diffusion_pytorch_model.bin''')
						_lowerCamelCase	:    Optional[int]          =      torch.load(vae_path, map_location='''cpu''')

			if osp.exists(text_enc_path):
						_lowerCamelCase	:    Any          =      load_file(text_enc_path, device='''cpu''')
			else:
						_lowerCamelCase	:    str          =      osp.join(args.model_path, '''text_encoder''', '''pytorch_model.bin''')
						_lowerCamelCase	:    Union[str, Any]          =      torch.load(text_enc_path, map_location='''cpu''')

			# Convert the UNet model
			_lowerCamelCase	:    List[Any]          =      convert_unet_state_dict(unet_state_dict)
			_lowerCamelCase	:    List[Any]          =      {'''model.diffusion_model.''' + k: v for k, v in unet_state_dict.items()}

			# Convert the VAE model
			_lowerCamelCase	:    Any          =      convert_vae_state_dict(vae_state_dict)
			_lowerCamelCase	:    Dict          =      {'''first_stage_model.''' + k: v for k, v in vae_state_dict.items()}

			# Easiest way to identify v2.0 model seems to be that the text encoder (OpenCLIP) is deeper
			_lowerCamelCase	:    List[Any]          =      '''text_model.encoder.layers.22.layer_norm2.bias''' in text_enc_dict

			if is_vaa_model:
						# Need to add the tag 'transformer' in advance so we can knock it out from the final layer-norm
						_lowerCamelCase	:    Optional[int]          =      {'''transformer.''' + k: v for k, v in text_enc_dict.items()}
						_lowerCamelCase	:    Optional[Any]          =      convert_text_enc_state_dict_vaa(text_enc_dict)
						_lowerCamelCase	:    List[str]          =      {'''cond_stage_model.model.''' + k: v for k, v in text_enc_dict.items()}
			else:
						_lowerCamelCase	:    Any          =      convert_text_enc_state_dict(text_enc_dict)
						_lowerCamelCase	:    Tuple          =      {'''cond_stage_model.transformer.''' + k: v for k, v in text_enc_dict.items()}

			# Put together new checkpoint
			_lowerCamelCase	:    Dict          =      {**unet_state_dict, **vae_state_dict, **text_enc_dict}
			if args.half:
						_lowerCamelCase	:    Union[str, Any]          =      {k: v.half() for k, v in state_dict.items()}

			if args.use_safetensors:
						save_file(state_dict, args.checkpoint_path)
			else:
						_lowerCamelCase	:    Optional[int]          =      {'''state_dict''': state_dict}
						torch.save(state_dict, args.checkpoint_path)

import argparse
import json
import os

import numpy as np
import PIL
import requests
import tensorflow.keras.applications.efficientnet as efficientnet
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from tensorflow.keras.preprocessing import image

from transformers import (
    EfficientNetConfig,
    EfficientNetForImageClassification,
    EfficientNetImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
_lowerCamelCase	:    int          =      logging.get_logger(__name__)

_lowerCamelCase	:    Tuple          =      {
    '''b0''': efficientnet.EfficientNetBa,
    '''b1''': efficientnet.EfficientNetBa,
    '''b2''': efficientnet.EfficientNetBa,
    '''b3''': efficientnet.EfficientNetBa,
    '''b4''': efficientnet.EfficientNetBa,
    '''b5''': efficientnet.EfficientNetBa,
    '''b6''': efficientnet.EfficientNetBa,
    '''b7''': efficientnet.EfficientNetBa,
}

_lowerCamelCase	:    List[str]          =      {
    '''b0''': {
        '''hidden_dim''': 1_280,
        '''width_coef''': 1.0,
        '''depth_coef''': 1.0,
        '''image_size''': 224,
        '''dropout_rate''': 0.2,
        '''dw_padding''': [],
    },
    '''b1''': {
        '''hidden_dim''': 1_280,
        '''width_coef''': 1.0,
        '''depth_coef''': 1.1,
        '''image_size''': 240,
        '''dropout_rate''': 0.2,
        '''dw_padding''': [16],
    },
    '''b2''': {
        '''hidden_dim''': 1_408,
        '''width_coef''': 1.1,
        '''depth_coef''': 1.2,
        '''image_size''': 260,
        '''dropout_rate''': 0.3,
        '''dw_padding''': [5, 8, 16],
    },
    '''b3''': {
        '''hidden_dim''': 1_536,
        '''width_coef''': 1.2,
        '''depth_coef''': 1.4,
        '''image_size''': 300,
        '''dropout_rate''': 0.3,
        '''dw_padding''': [5, 18],
    },
    '''b4''': {
        '''hidden_dim''': 1_792,
        '''width_coef''': 1.4,
        '''depth_coef''': 1.8,
        '''image_size''': 380,
        '''dropout_rate''': 0.4,
        '''dw_padding''': [6],
    },
    '''b5''': {
        '''hidden_dim''': 2_048,
        '''width_coef''': 1.6,
        '''depth_coef''': 2.2,
        '''image_size''': 456,
        '''dropout_rate''': 0.4,
        '''dw_padding''': [13, 27],
    },
    '''b6''': {
        '''hidden_dim''': 2_304,
        '''width_coef''': 1.8,
        '''depth_coef''': 2.6,
        '''image_size''': 528,
        '''dropout_rate''': 0.5,
        '''dw_padding''': [31],
    },
    '''b7''': {
        '''hidden_dim''': 2_560,
        '''width_coef''': 2.0,
        '''depth_coef''': 3.1,
        '''image_size''': 600,
        '''dropout_rate''': 0.5,
        '''dw_padding''': [18],
    },
}






def 	a_		(	__lowercase       :     int			)			->  Optional[int]:
			_snake_case          =      EfficientNetConfig()
			_snake_case          =      CONFIG_MAP[model_name]['hidden_dim']
			_snake_case          =      CONFIG_MAP[model_name]['width_coef']
			_snake_case          =      CONFIG_MAP[model_name]['depth_coef']
			_snake_case          =      CONFIG_MAP[model_name]['image_size']
			_snake_case          =      CONFIG_MAP[model_name]['dropout_rate']
			_snake_case          =      CONFIG_MAP[model_name]['dw_padding']

			_snake_case          =      'huggingface/label-files'
			_snake_case          =      'imagenet-1k-id2label.json'
			_snake_case          =      1_000
			_snake_case          =      json.load(open(hf_hub_download(__lowercase	, __lowercase	, repo_type='dataset'			)	, 'r'			)			)
			_snake_case          =      {int(__lowercase			): v for k, v in idalabel.items()}

			_snake_case          =      idalabel
			_snake_case          =      {v: k for k, v in idalabel.items()}
			return config






def 	a_		(	)			->  Any:
			_snake_case          =      'http://images.cocodataset.org/val2017/000000039769.jpg'
			_snake_case          =      Image.open(requests.get(__lowercase	, stream=__lowercase			).raw			)
			return im






def 	a_		(	__lowercase       :     Union[str, Any]			)			->  Tuple:
			_snake_case          =      CONFIG_MAP[model_name]['image_size']
			_snake_case          =      EfficientNetImageProcessor(
			    size={'height': size, 'width': size}	, image_mean=[0.4_8_5, 0.4_5_6, 0.4_0_6]	, image_std=[0.4_7_8_5_3_9_4_4, 0.4_7_3_2_8_6_4, 0.4_7_4_3_4_1_6_3]	, do_center_crop=__lowercase	, )
			return preprocessor






def 	a_		(	__lowercase       :     str			)			->  List[Any]:
			_snake_case          =      [v.split('_'			)[0].split('block'			)[1] for v in original_param_names if v.startswith('block'			)]
			_snake_case          =      sorted(set(__lowercase			)			)
			_snake_case          =      len(__lowercase			)
			_snake_case          =      {b: str(__lowercase			) for b, i in zip(__lowercase	, range(__lowercase			)			)}

			_snake_case          =      []
			rename_keys.append(('stem_conv/kernel:0', 'embeddings.convolution.weight')			)
			rename_keys.append(('stem_bn/gamma:0', 'embeddings.batchnorm.weight')			)
			rename_keys.append(('stem_bn/beta:0', 'embeddings.batchnorm.bias')			)
			rename_keys.append(('stem_bn/moving_mean:0', 'embeddings.batchnorm.running_mean')			)
			rename_keys.append(('stem_bn/moving_variance:0', 'embeddings.batchnorm.running_var')			)

			for b in block_names:
						_snake_case          =      block_name_mapping[b]
						rename_keys.append((f'''block{b}_expand_conv/kernel:0''', f'''encoder.blocks.{hf_b}.expansion.expand_conv.weight''')			)
						rename_keys.append((f'''block{b}_expand_bn/gamma:0''', f'''encoder.blocks.{hf_b}.expansion.expand_bn.weight''')			)
						rename_keys.append((f'''block{b}_expand_bn/beta:0''', f'''encoder.blocks.{hf_b}.expansion.expand_bn.bias''')			)
						rename_keys.append(
						    (f'''block{b}_expand_bn/moving_mean:0''', f'''encoder.blocks.{hf_b}.expansion.expand_bn.running_mean''')			)
						rename_keys.append(
						    (f'''block{b}_expand_bn/moving_variance:0''', f'''encoder.blocks.{hf_b}.expansion.expand_bn.running_var''')			)
						rename_keys.append(
						    (f'''block{b}_dwconv/depthwise_kernel:0''', f'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_conv.weight''')			)
						rename_keys.append((f'''block{b}_bn/gamma:0''', f'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.weight''')			)
						rename_keys.append((f'''block{b}_bn/beta:0''', f'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.bias''')			)
						rename_keys.append(
						    (f'''block{b}_bn/moving_mean:0''', f'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_mean''')			)
						rename_keys.append(
						    (f'''block{b}_bn/moving_variance:0''', f'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_var''')			)

						rename_keys.append((f'''block{b}_se_reduce/kernel:0''', f'''encoder.blocks.{hf_b}.squeeze_excite.reduce.weight''')			)
						rename_keys.append((f'''block{b}_se_reduce/bias:0''', f'''encoder.blocks.{hf_b}.squeeze_excite.reduce.bias''')			)
						rename_keys.append((f'''block{b}_se_expand/kernel:0''', f'''encoder.blocks.{hf_b}.squeeze_excite.expand.weight''')			)
						rename_keys.append((f'''block{b}_se_expand/bias:0''', f'''encoder.blocks.{hf_b}.squeeze_excite.expand.bias''')			)
						rename_keys.append(
						    (f'''block{b}_project_conv/kernel:0''', f'''encoder.blocks.{hf_b}.projection.project_conv.weight''')			)
						rename_keys.append((f'''block{b}_project_bn/gamma:0''', f'''encoder.blocks.{hf_b}.projection.project_bn.weight''')			)
						rename_keys.append((f'''block{b}_project_bn/beta:0''', f'''encoder.blocks.{hf_b}.projection.project_bn.bias''')			)
						rename_keys.append(
						    (f'''block{b}_project_bn/moving_mean:0''', f'''encoder.blocks.{hf_b}.projection.project_bn.running_mean''')			)
						rename_keys.append(
						    (f'''block{b}_project_bn/moving_variance:0''', f'''encoder.blocks.{hf_b}.projection.project_bn.running_var''')			)

			rename_keys.append(('top_conv/kernel:0', 'encoder.top_conv.weight')			)
			rename_keys.append(('top_bn/gamma:0', 'encoder.top_bn.weight')			)
			rename_keys.append(('top_bn/beta:0', 'encoder.top_bn.bias')			)
			rename_keys.append(('top_bn/moving_mean:0', 'encoder.top_bn.running_mean')			)
			rename_keys.append(('top_bn/moving_variance:0', 'encoder.top_bn.running_var')			)

			_snake_case          =      {}
			for item in rename_keys:
						if item[0] in original_param_names:
									_snake_case          =      'efficientnet.' + item[1]

			_snake_case          =      'classifier.weight'
			_snake_case          =      'classifier.bias'
			return key_mapping






def 	a_		(	__lowercase       :     Any	, __lowercase       :     Any	, __lowercase       :     Any			)			->  Optional[Any]:
			for key, value in tf_params.items():
						if "normalization" in key:
									continue

						_snake_case          =      key_mapping[key]
						if "_conv" in key and "kernel" in key:
									_snake_case          =      torch.from_numpy(__lowercase			).permute(3	, 2	, 0	, 1			)
						elif "depthwise_kernel" in key:
									_snake_case          =      torch.from_numpy(__lowercase			).permute(2	, 3	, 0	, 1			)
						elif "kernel" in key:
									_snake_case          =      torch.from_numpy(np.transpose(__lowercase			)			)
						else:
									_snake_case          =      torch.from_numpy(__lowercase			)

						# Replace HF parameters with original TF model parameters
						assert hf_params[hf_key].shape == new_hf_value.shape
						hf_params[hf_key].copy_(__lowercase			)






@torch.no_grad()
def 	a_		(	__lowercase       :     List[Any]	, __lowercase       :     Any	, __lowercase       :     int	, __lowercase       :     str			)			->  Dict:
			_snake_case          =      model_classes[model_name](
			    include_top=__lowercase	, weights='imagenet'	, input_tensor=__lowercase	, input_shape=__lowercase	, pooling=__lowercase	, classes=1_000	, classifier_activation='softmax'	, )

			_snake_case          =      original_model.trainable_variables
			_snake_case          =      original_model.non_trainable_variables
			_snake_case          =      {param.name: param.numpy() for param in tf_params}
			for param in tf_non_train_params:
						_snake_case          =      param.numpy()
			_snake_case          =      list(tf_params.keys()			)

			# Load HuggingFace model
			_snake_case          =      get_efficientnet_config(__lowercase			)
			_snake_case          =      EfficientNetForImageClassification(__lowercase			).eval()
			_snake_case          =      hf_model.state_dict()

			# Create src-to-dst parameter name mapping dictionary
			print('Converting parameters...'			)
			_snake_case          =      rename_keys(__lowercase			)
			replace_params(__lowercase	, __lowercase	, __lowercase			)

			# Initialize preprocessor and preprocess input image
			_snake_case          =      convert_image_processor(__lowercase			)
			_snake_case          =      preprocessor(images=prepare_img()	, return_tensors='pt'			)

			# HF model inference
			hf_model.eval()
			with torch.no_grad():
						_snake_case          =      hf_model(**__lowercase			)
			_snake_case          =      outputs.logits.detach().numpy()

			# Original model inference
			_snake_case          =      False
			_snake_case          =      CONFIG_MAP[model_name]['image_size']
			_snake_case          =      prepare_img().resize((image_size, image_size)	, resample=PIL.Image.NEAREST			)
			_snake_case          =      image.img_to_array(__lowercase			)
			_snake_case          =      np.expand_dims(__lowercase	, axis=0			)
			_snake_case          =      original_model.predict(__lowercase			)

			# Check whether original and HF model outputs match  -> np.allclose
			assert np.allclose(__lowercase	, __lowercase	, atol=1E-3			), "The predicted logits are not the same."
			print('Model outputs match!'			)

			if save_model:
						# Create folder to save model
						if not os.path.isdir(__lowercase			):
									os.mkdir(__lowercase			)
						# Save converted model and image processor
						hf_model.save_pretrained(__lowercase			)
						preprocessor.save_pretrained(__lowercase			)

			if push_to_hub:
						# Push model and image processor to hub
						print(f'''Pushing converted {model_name} to the hub...'''			)
						_snake_case          =      f'''efficientnet-{model_name}'''
						preprocessor.push_to_hub(__lowercase			)
						hf_model.push_to_hub(__lowercase			)


if __name__ == "__main__":
			_lowerCamelCase	:    int          =      argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    '''--model_name''',
			    default='''b0''',
			    type=str,
			    help='''Version name of the EfficientNet model you want to convert, select from [b0, b1, b2, b3, b4, b5, b6, b7].''',
			)
			parser.add_argument(
			    '''--pytorch_dump_folder_path''',
			    default='''hf_model''',
			    type=str,
			    help='''Path to the output PyTorch model directory.''',
			)
			parser.add_argument('''--save_model''', action='''store_true''', help='''Save model to local''')
			parser.add_argument('''--push_to_hub''', action='''store_true''', help='''Push model and image processor to the hub''')

			_lowerCamelCase	:    List[str]          =      parser.parse_args()
			convert_efficientnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)