
import argparse
import os

import align
import numpy as np
import requests
import tensorflow as tf
import torch
from PIL import Image
from tokenizer import Tokenizer

from transformers import (
    AlignConfig,
    AlignModel,
    AlignProcessor,
    BertConfig,
    BertTokenizer,
    EfficientNetConfig,
    EfficientNetImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
__snake_case					:							List[Any]           =							logging.get_logger(__name__)




def 		A						(  SCREAMING_SNAKE_CASE							):




			"""simple docstring"""


			UpperCAmelCase__		:Optional[Any] 		=    tf.image.resize(SCREAMING_SNAKE_CASE     ,						(346, 346)							)
			UpperCAmelCase__		:int 		=    tf.image.crop_to_bounding_box(SCREAMING_SNAKE_CASE     ,						(346 - 289) // 2     ,						(346 - 289) // 2     ,						289     ,						289							)
			return image




def 		A						(  ):




			"""simple docstring"""


			UpperCAmelCase__		:Any 		=    EfficientNetConfig.from_pretrained('google/efficientnet-b7'							)
			UpperCAmelCase__		:Tuple 		=    289
			UpperCAmelCase__		:int 		=    640
			UpperCAmelCase__		:str 		=    {'0': 'LABEL_0', '1': 'LABEL_1'}
			UpperCAmelCase__		:Tuple 		=    {'LABEL_0': 0, 'LABEL_1': 1}
			UpperCAmelCase__		:Dict 		=    []

			UpperCAmelCase__		:List[str] 		=    BertConfig()
			UpperCAmelCase__		:List[str] 		=    AlignConfig.from_text_vision_configs(
			    text_config=SCREAMING_SNAKE_CASE     ,						vision_config=SCREAMING_SNAKE_CASE     ,						projection_dim=640							)
			return config




def 		A						(  ):




			"""simple docstring"""


			UpperCAmelCase__		:Tuple 		=    'http://images.cocodataset.org/val2017/000000039769.jpg'
			UpperCAmelCase__		:str 		=    Image.open(requests.get(SCREAMING_SNAKE_CASE     ,						stream=SCREAMING_SNAKE_CASE							).raw							)
			return im




def 		A						(  ):




			"""simple docstring"""


			UpperCAmelCase__		:str 		=    EfficientNetImageProcessor(
			    do_center_crop=SCREAMING_SNAKE_CASE     ,						rescale_factor=1 / 1_27.5     ,						rescale_offset=SCREAMING_SNAKE_CASE     ,						do_normalize=SCREAMING_SNAKE_CASE     ,						include_top=SCREAMING_SNAKE_CASE     ,						resample=Image.BILINEAR     ,						)
			UpperCAmelCase__		:List[str] 		=    BertTokenizer.from_pretrained('bert-base-uncased'							)
			UpperCAmelCase__		:List[str] 		=    64
			UpperCAmelCase__		:Tuple 		=    AlignProcessor(image_processor=SCREAMING_SNAKE_CASE     ,						tokenizer=SCREAMING_SNAKE_CASE							)
			return processor




def 		A						(  SCREAMING_SNAKE_CASE							):




			"""simple docstring"""


			UpperCAmelCase__		:Dict 		=    [v.split('_'							)[0].split('block'							)[1] for v in original_param_names if v.startswith('block'							)]
			UpperCAmelCase__		:Optional[int] 		=    list(set(SCREAMING_SNAKE_CASE							)							)
			UpperCAmelCase__		:Optional[int] 		=    sorted(SCREAMING_SNAKE_CASE							)
			UpperCAmelCase__		:List[str] 		=    len(SCREAMING_SNAKE_CASE							)
			UpperCAmelCase__		:Dict 		=    {b: str(SCREAMING_SNAKE_CASE							) for b, i in zip(SCREAMING_SNAKE_CASE     ,						range(SCREAMING_SNAKE_CASE							)							)}

			UpperCAmelCase__		:Dict 		=    []
			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:
						UpperCAmelCase__		:int 		=    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""")							)

			UpperCAmelCase__		:Dict 		=    {}
			for item in rename_keys:
						if item[0] in original_param_names:
									UpperCAmelCase__		:Optional[int] 		=    'vision_model.' + item[1]

    # BERT text encoder
			UpperCAmelCase__		:Any 		=    []
			UpperCAmelCase__		:List[str] 		=    'tf_bert_model/bert'
			UpperCAmelCase__		:Union[str, Any] 		=    'text_model'
			for i in range(12							):
						rename_keys.append(
						    (
						        f"""{old}/encoder/layer_._{i}/attention/self/query/kernel:0""",
						        f"""{new}.encoder.layer.{i}.attention.self.query.weight""",
						    )							)
						rename_keys.append(
						    (
						        f"""{old}/encoder/layer_._{i}/attention/self/query/bias:0""",
						        f"""{new}.encoder.layer.{i}.attention.self.query.bias""",
						    )							)
						rename_keys.append(
						    (
						        f"""{old}/encoder/layer_._{i}/attention/self/key/kernel:0""",
						        f"""{new}.encoder.layer.{i}.attention.self.key.weight""",
						    )							)
						rename_keys.append(
						    (
						        f"""{old}/encoder/layer_._{i}/attention/self/key/bias:0""",
						        f"""{new}.encoder.layer.{i}.attention.self.key.bias""",
						    )							)
						rename_keys.append(
						    (
						        f"""{old}/encoder/layer_._{i}/attention/self/value/kernel:0""",
						        f"""{new}.encoder.layer.{i}.attention.self.value.weight""",
						    )							)
						rename_keys.append(
						    (
						        f"""{old}/encoder/layer_._{i}/attention/self/value/bias:0""",
						        f"""{new}.encoder.layer.{i}.attention.self.value.bias""",
						    )							)
						rename_keys.append(
						    (
						        f"""{old}/encoder/layer_._{i}/attention/output/dense/kernel:0""",
						        f"""{new}.encoder.layer.{i}.attention.output.dense.weight""",
						    )							)
						rename_keys.append(
						    (
						        f"""{old}/encoder/layer_._{i}/attention/output/dense/bias:0""",
						        f"""{new}.encoder.layer.{i}.attention.output.dense.bias""",
						    )							)
						rename_keys.append(
						    (
						        f"""{old}/encoder/layer_._{i}/attention/output/LayerNorm/gamma:0""",
						        f"""{new}.encoder.layer.{i}.attention.output.LayerNorm.weight""",
						    )							)
						rename_keys.append(
						    (
						        f"""{old}/encoder/layer_._{i}/attention/output/LayerNorm/beta:0""",
						        f"""{new}.encoder.layer.{i}.attention.output.LayerNorm.bias""",
						    )							)
						rename_keys.append(
						    (
						        f"""{old}/encoder/layer_._{i}/intermediate/dense/kernel:0""",
						        f"""{new}.encoder.layer.{i}.intermediate.dense.weight""",
						    )							)
						rename_keys.append(
						    (
						        f"""{old}/encoder/layer_._{i}/intermediate/dense/bias:0""",
						        f"""{new}.encoder.layer.{i}.intermediate.dense.bias""",
						    )							)
						rename_keys.append(
						    (f"""{old}/encoder/layer_._{i}/output/dense/kernel:0""", f"""{new}.encoder.layer.{i}.output.dense.weight""")							)
						rename_keys.append(
						    (f"""{old}/encoder/layer_._{i}/output/dense/bias:0""", f"""{new}.encoder.layer.{i}.output.dense.bias""")							)
						rename_keys.append(
						    (f"""{old}/encoder/layer_._{i}/output/LayerNorm/gamma:0""", f"""{new}.encoder.layer.{i}.output.LayerNorm.weight""")							)
						rename_keys.append(
						    (f"""{old}/encoder/layer_._{i}/output/LayerNorm/beta:0""", f"""{new}.encoder.layer.{i}.output.LayerNorm.bias""")							)

			rename_keys.append((f"""{old}/embeddings/word_embeddings/weight:0""", f"""{new}.embeddings.word_embeddings.weight""")							)
			rename_keys.append(
			    (f"""{old}/embeddings/position_embeddings/embeddings:0""", f"""{new}.embeddings.position_embeddings.weight""")							)
			rename_keys.append(
			    (f"""{old}/embeddings/token_type_embeddings/embeddings:0""", f"""{new}.embeddings.token_type_embeddings.weight""")							)
			rename_keys.append((f"""{old}/embeddings/LayerNorm/gamma:0""", f"""{new}.embeddings.LayerNorm.weight""")							)
			rename_keys.append((f"""{old}/embeddings/LayerNorm/beta:0""", f"""{new}.embeddings.LayerNorm.bias""")							)

			rename_keys.append((f"""{old}/pooler/dense/kernel:0""", f"""{new}.pooler.dense.weight""")							)
			rename_keys.append((f"""{old}/pooler/dense/bias:0""", f"""{new}.pooler.dense.bias""")							)
			rename_keys.append(('dense/kernel:0', 'text_projection.weight')							)
			rename_keys.append(('dense/bias:0', 'text_projection.bias')							)
			rename_keys.append(('dense/bias:0', 'text_projection.bias')							)
			rename_keys.append(('temperature:0', 'temperature')							)

			for item in rename_keys:
						if item[0] in original_param_names:
									UpperCAmelCase__		:Any 		=    item[1]
			return key_mapping




def 		A						(  SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE							):




			"""simple docstring"""


			list(hf_params.keys()							)

			for key, value in tf_params.items():
						if key not in key_mapping:
									continue

						UpperCAmelCase__		:str 		=    key_mapping[key]
						if "_conv" in key and "kernel" in key:
									UpperCAmelCase__		:Optional[Any] 		=    torch.from_numpy(SCREAMING_SNAKE_CASE							).permute(3     ,						2     ,						0     ,						1							)
						elif "embeddings" in key:
									UpperCAmelCase__		:int 		=    torch.from_numpy(SCREAMING_SNAKE_CASE							)
						elif "depthwise_kernel" in key:
									UpperCAmelCase__		:Union[str, Any] 		=    torch.from_numpy(SCREAMING_SNAKE_CASE							).permute(2     ,						3     ,						0     ,						1							)
						elif "kernel" in key:
									UpperCAmelCase__		:Dict 		=    torch.from_numpy(np.transpose(SCREAMING_SNAKE_CASE							)							)
						elif "temperature" in key:
									UpperCAmelCase__		:Optional[int] 		=    value
						elif "bn/gamma" or "bn/beta" in key:
									UpperCAmelCase__		:Tuple 		=    torch.from_numpy(np.transpose(SCREAMING_SNAKE_CASE							)							).squeeze()
						else:
									UpperCAmelCase__		:List[str] 		=    torch.from_numpy(SCREAMING_SNAKE_CASE							)

						# Replace HF parameters with original TF model parameters
						hf_params[hf_key].copy_(SCREAMING_SNAKE_CASE							)




@torch.no_grad()
def 		A						(  SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE							):




			"""simple docstring"""


			UpperCAmelCase__		:Dict 		=    64
			UpperCAmelCase__		:Optional[Any] 		=    Tokenizer(SCREAMING_SNAKE_CASE							)
			UpperCAmelCase__		:Optional[int] 		=    align.Align('efficientnet-b7'     ,						'bert-base'     ,						640     ,						SCREAMING_SNAKE_CASE     ,						tok.get_vocab_size()							)
			original_model.compile()
			original_model.load_weights(SCREAMING_SNAKE_CASE							)

			UpperCAmelCase__		:Tuple 		=    original_model.trainable_variables
			UpperCAmelCase__		:List[Any] 		=    original_model.non_trainable_variables
			UpperCAmelCase__		:List[str] 		=    {param.name: param.numpy() for param in tf_params}
			for param in tf_non_train_params:
						UpperCAmelCase__		:int 		=    param.numpy()
			UpperCAmelCase__		:List[str] 		=    list(tf_params.keys()							)

			# Load HuggingFace model
			UpperCAmelCase__		:List[str] 		=    get_align_config()
			UpperCAmelCase__		:str 		=    AlignModel(SCREAMING_SNAKE_CASE							).eval()
			UpperCAmelCase__		:Any 		=    hf_model.state_dict()

			# Create src-to-dst parameter name mapping dictionary
			print('Converting parameters...'							)
			UpperCAmelCase__		:List[str] 		=    rename_keys(SCREAMING_SNAKE_CASE							)
			replace_params(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE							)

			# Initialize processor
			UpperCAmelCase__		:int 		=    get_processor()
			UpperCAmelCase__		:int 		=    processor(
			    images=prepare_img()     ,						text='A picture of a cat'     ,						padding='max_length'     ,						max_length=64     ,						return_tensors='pt'							)

			# HF model inference
			hf_model.eval()
			with torch.no_grad():
						UpperCAmelCase__		:List[str] 		=    hf_model(**SCREAMING_SNAKE_CASE							)

			UpperCAmelCase__		:str 		=    outputs.image_embeds.detach().numpy()
			UpperCAmelCase__		:Tuple 		=    outputs.text_embeds.detach().numpy()

			# Original model inference
			UpperCAmelCase__		:Union[str, Any] 		=    False
			UpperCAmelCase__		:Dict 		=    EfficientNetImageProcessor(
			    do_center_crop=SCREAMING_SNAKE_CASE     ,						do_rescale=SCREAMING_SNAKE_CASE     ,						do_normalize=SCREAMING_SNAKE_CASE     ,						include_top=SCREAMING_SNAKE_CASE     ,						resample=Image.BILINEAR     ,						)
			UpperCAmelCase__		:int 		=    tf_image_processor(images=prepare_img()     ,						return_tensors='tf'     ,						data_format='channels_last'							)['pixel_values']
			UpperCAmelCase__		:Optional[int] 		=    tok(tf.constant(['A picture of a cat']							)							)

			UpperCAmelCase__		:Optional[int] 		=    original_model.image_encoder(SCREAMING_SNAKE_CASE     ,						training=SCREAMING_SNAKE_CASE							)
			UpperCAmelCase__		:Union[str, Any] 		=    original_model.text_encoder(SCREAMING_SNAKE_CASE     ,						training=SCREAMING_SNAKE_CASE							)

			UpperCAmelCase__		:str 		=    tf.nn.la_normalize(SCREAMING_SNAKE_CASE     ,						axis=-1							)
			UpperCAmelCase__		:Union[str, Any] 		=    tf.nn.la_normalize(SCREAMING_SNAKE_CASE     ,						axis=-1							)

			# Check whether original and HF model outputs match  -> np.allclose
			if not np.allclose(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE     ,						atol=1E-3							):
						raise ValueError('The predicted image features are not the same.'							)
			if not np.allclose(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE     ,						atol=1E-3							):
						raise ValueError('The predicted text features are not the same.'							)
			print('Model outputs match!'							)

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

			if push_to_hub:
						# Push model and image processor to hub
						print('Pushing converted ALIGN to the hub...'							)
						processor.push_to_hub('align-base'							)
						hf_model.push_to_hub('align-base'							)


if __name__ == "__main__":
						__snake_case					:							Tuple           =							argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '--checkpoint_path',
						    default='./weights/model-weights',
						    type=str,
						    help='Path to the pretrained TF ALIGN checkpoint.',
						)
						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')

						__snake_case					:							List[str]           =							parser.parse_args()
						convert_align_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)







import bisect
import copy
from collections import defaultdict

import numpy as np
from torch.utils.data import BatchSampler, Sampler

from utils import logger




def 		A						(  SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE							):




			"""simple docstring"""


			UpperCAmelCase__		:int 		=    copy.deepcopy(SCREAMING_SNAKE_CASE							)
			UpperCAmelCase__		:List[Any] 		=    sorted(SCREAMING_SNAKE_CASE							)
			UpperCAmelCase__		:Dict 		=    [bisect.bisect_right(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE							) for y in x]
			return quantized




def 		A						(  SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE=0							):




			"""simple docstring"""


			UpperCAmelCase__		:Tuple 		=    np.arange(start=3     ,						stop=SCREAMING_SNAKE_CASE     ,						step=4							).tolist() if k > 0 else [10]
			UpperCAmelCase__		:List[Any] 		=    _quantize(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE							)
			# count number of elements per group
			UpperCAmelCase__		:Dict 		=    np.unique(SCREAMING_SNAKE_CASE     ,						return_counts=SCREAMING_SNAKE_CASE							)[1]
			UpperCAmelCase__		:Any 		=    [0] + bins + [np.inf]
			logger.info('Using {} as bins for aspect lengths quantization'.format(SCREAMING_SNAKE_CASE							)							)
			logger.info('Count of instances per bin: {}'.format(SCREAMING_SNAKE_CASE							)							)
			return groups




class 	UpperCamelCase__       (  UpperCAmelCase__):

						'''simple docstring'''



						def __init__(  self  , A  , A  , A	)  ->Union[str, Any]:
									if not isinstance(A  , A	):
												raise ValueError(
												    'sampler should be an instance of torch.utils.data.Sampler, but got sampler={}'.format(A	)	)
									UpperCAmelCase__		:Tuple 		=    sampler
									UpperCAmelCase__		:Tuple 		=    group_ids
									UpperCAmelCase__		:str 		=    batch_size


						def __iter__(  self	)  ->Dict:
									UpperCAmelCase__		:List[str] 		=    defaultdict(A	)
									UpperCAmelCase__		:List[str] 		=    defaultdict(A	)

									UpperCAmelCase__		:Any 		=    0
									for idx in self.sampler:
												UpperCAmelCase__		:List[Any] 		=    self.group_ids[idx]
												buffer_per_group[group_id].append(A	)
												samples_per_group[group_id].append(A	)
												if len(buffer_per_group[group_id]	) == self.batch_size:
															yield buffer_per_group[group_id]  # TODO
															num_batches += 1
															del buffer_per_group[group_id]
												assert len(buffer_per_group[group_id]	) < self.batch_size

									# now we have run out of elements that satisfy
									# the group criteria, let's return the remaining
									# elements so that the size of the sampler is
									# deterministic
									UpperCAmelCase__		:Union[str, Any] 		=    len(self	)
									UpperCAmelCase__		:Tuple 		=    expected_num_batches - num_batches
									if num_remaining > 0:
												# for the remaining batches, group the batches by similar lengths
												UpperCAmelCase__		:str 		=    []
												for group_id, idxs in sorted(buffer_per_group.items()  , key=lambda A	: x[0]	):
															batch_idx.extend(A	)
															if len(A	) >= self.batch_size:
																		yield batch_idx[: self.batch_size]
																		UpperCAmelCase__		:int 		=    batch_idx[self.batch_size :]
																		num_remaining -= 1
												if len(A	) > 0:
															yield batch_idx
															num_remaining -= 1
									assert num_remaining == 0


						def __len__(  self	)  ->Union[str, Any]:
									return (len(self.sampler	) + self.batch_size - 1) // self.batch_size





