code: 
  
  
  
  
  
  import argparse
  
  import torch
  
  from diffusers import UNetaDConditionModel
  
  def     a						(      _lowercase   :     List[Any] ,						_lowercase   :     List[str] ,						_lowercase   :     Union[str, Any] ,						_lowercase   :     Any=None ,						_lowercase   :     Union[str, Any]=None ,						_lowercase   :     Optional[int]=None       )     ->List[str]:
  		assert isinstance(_lowercase ,						_lowercase       ), "Paths should be a list of dicts containing 'old' and 'new' keys."
  
  		# Splits the attention layers into three variables.
  		if attention_paths_to_split is not None:
  				for path, path_map in attention_paths_to_split.items():
  						UpperCamelCase__         =	old_checkpoint[path]
  						UpperCamelCase__         =	old_tensor.shape[0] // 3
  
  						UpperCamelCase__         =	(-1, channels) if len(old_tensor.shape       ) == 3 else (-1)
  
  						UpperCamelCase__         =	old_tensor.shape[0] // config["""num_head_channels"""] // 3
  
  						UpperCamelCase__         =	old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]       )
  						UpperCamelCase__  ,			UpperCamelCase__  ,			UpperCamelCase__         =	old_tensor.split(channels // num_heads ,						dim=1       )
  
  						UpperCamelCase__         =	query.reshape(_lowercase       )
  						UpperCamelCase__         =	key.reshape(_lowercase       )
  						UpperCamelCase__         =	value.reshape(_lowercase       )
  
  		for path in paths:
  				UpperCamelCase__         =	path["""new"""]
  
  				# These have already been assigned
  				if attention_paths_to_split is not None and new_path in attention_paths_to_split:
  						continue
  
  				if additional_replacements is not None:
  						for replacement in additional_replacements:
  								UpperCamelCase__         =	new_path.replace(replacement["""old"""] ,						replacement["""new"""]       )
  
          # proj_attn.weight has to be converted from conv 1D to linear
  				UpperCamelCase__         =	old_checkpoint[path["""old"""]]
  				UpperCamelCase__         =	["""proj_attn.weight"""]
  				UpperCamelCase__         =	["""proj_out.weight""", """proj_in.weight"""]
  				if any(k in new_path for k in names       ):
  						UpperCamelCase__         =	weight[:, :, 0]
  				elif any(k in new_path for k in names_a       ) and len(weight.shape       ) > 2 and ".attentions." not in new_path:
  						UpperCamelCase__         =	weight[:, :, 0]
  				else:
  						UpperCamelCase__         =	weight
  
  def     a						(      _lowercase   :     List[Any] ,						_lowercase   :     Dict=0       )     ->List[str]:
  		UpperCamelCase__         =	[]
  		for old_item in old_list:
  				UpperCamelCase__         =	old_item
  
  				#         new_item = new_item.replace('norm.weight', 'group_norm.weight')
  				#         new_item = new_item.replace('norm.bias', 'group_norm.bias')
  
  				#         new_item = new_item.replace('proj_out.weight', 'proj_attn.weight')
  				#         new_item = new_item.replace('proj_out.bias', 'proj_attn.bias')
  
  				#         new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
  
  				mapping.append({"""old""": old_item, """new""": new_item}       )
  
  		return mapping
  
  def     a						(      _lowercase   :     Any ,						_lowercase   :     Optional[Any]=1       )     ->Union[str, Any]:
  		if n_shave_prefix_segments >= 0:
  				return ".".join(path.split("""."""       )[n_shave_prefix_segments:]       )
  		else:
  				return ".".join(path.split("""."""       )[:n_shave_prefix_segments]       )
  
  def     a						(      _lowercase   :     Tuple ,						_lowercase   :     Optional[Any]=0       )     ->str:
  		UpperCamelCase__         =	[]
  		for old_item in old_list:
  				mapping.append({"""old""": old_item, """new""": old_item}       )
  
  		return mapping
  
  def     a						(      _lowercase   :     List[Any] ,						_lowercase   :     Union[str, Any]=0       )     ->Any:
  		UpperCamelCase__         =	[]
  		for old_item in old_list:
  				UpperCamelCase__         =	old_item.replace("""in_layers.0""" ,						"""norm1"""       )
  				UpperCamelCase__         =	new_item.replace("""in_layers.2""" ,						"""conv1"""       )
  
  				UpperCamelCase__         =	new_item.replace("""out_layers.0""" ,						"""norm2"""       )
  				UpperCamelCase__         =	new_item.replace("""out_layers.3""" ,						"""conv2"""       )
  
  				UpperCamelCase__         =	new_item.replace("""emb_layers.1""" ,						"""time_emb_proj"""       )
  				UpperCamelCase__         =	new_item.replace("""skip_connection""" ,						"""conv_shortcut"""       )
  
  				UpperCamelCase__         =	shave_segments(_lowercase ,						n_shave_prefix_segments=_lowercase       )
  
  				if "temopral_conv" not in old_item:
  						mapping.append({"""old""": old_item, """new""": new_item}       )
  
  		return mapping
  
  def     a						(      _lowercase   :     int ,						_lowercase   :     Optional[Any] ,						_lowercase   :     int=None ,						_lowercase   :     Any=False       )     ->Tuple:
  		UpperCamelCase__         =	{}
  		UpperCamelCase__         =	list(checkpoint.keys()       )
  
  		UpperCamelCase__         =	"""model.diffusion_model."""
  
  		# at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA
  		if sum(k.startswith("""model_ema"""       ) for k in keys       ) > 1_0_0 and extract_ema:
  				print(F"""Checkpoint {path} has both EMA and non-EMA weights."""       )
  				print(
  				    """In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA"""
  				    """ weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag."""       )
  				for key in keys:
  						if key.startswith("""model.diffusion_model"""       ):
  								UpperCamelCase__         =	"""model_ema.""" + """""".join(key.split("""."""       )[1:]       )
  								UpperCamelCase__         =	checkpoint.pop(_lowercase       )
  		else:
  				if sum(k.startswith("""model_ema"""       ) for k in keys       ) > 1_0_0:
  						print(
  						    """In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA"""
  						    """ weights (usually better for inference), please make sure to add the `--extract_ema` flag."""       )
  
  				for key in keys:
  						UpperCamelCase__         =	checkpoint.pop(_lowercase       )
  
  		UpperCamelCase__         =	{}
  
  		UpperCamelCase__         =	unet_state_dict["""time_embed.0.weight"""]
  		UpperCamelCase__         =	unet_state_dict["""time_embed.0.bias"""]
  		UpperCamelCase__         =	unet_state_dict["""time_embed.2.weight"""]
  		UpperCamelCase__         =	unet_state_dict["""time_embed.2.bias"""]
  
  		if config["class_embed_type"] is None:
  				# No parameters to port
  				...
  		elif config["class_embed_type"] == "timestep" or config["class_embed_type"] == "projection":
  				UpperCamelCase__         =	unet_state_dict["""label_emb.0.0.weight"""]
  				UpperCamelCase__         =	unet_state_dict["""label_emb.0.0.bias"""]
  				UpperCamelCase__         =	unet_state_dict["""label_emb.0.2.weight"""]
  				UpperCamelCase__         =	unet_state_dict["""label_emb.0.2.bias"""]
  		else:
  				raise NotImplementedError(F"""Not implemented `class_embed_type`: {config['class_embed_type']}"""       )
  
  		UpperCamelCase__         =	unet_state_dict["""input_blocks.0.0.weight"""]
  		UpperCamelCase__         =	unet_state_dict["""input_blocks.0.0.bias"""]
  
  		UpperCamelCase__         =	[v for v in unet_state_dict if v.startswith("""input_blocks.0.1"""       )]
  		UpperCamelCase__         =	renew_attention_paths(_lowercase       )
  		UpperCamelCase__         =	{"""old""": """input_blocks.0.1""", """new""": """transformer_in"""}
  		assign_to_checkpoint(_lowercase ,						_lowercase ,						_lowercase ,						additional_replacements=[meta_path] ,						config=_lowercase       )
  
  		UpperCamelCase__         =	unet_state_dict["""out.0.weight"""]
  		UpperCamelCase__         =	unet_state_dict["""out.0.bias"""]
  		UpperCamelCase__         =	unet_state_dict["""out.2.weight"""]
  		UpperCamelCase__         =	unet_state_dict["""out.2.bias"""]
  
  		# Retrieves the keys for the input blocks only
  		UpperCamelCase__         =	len({""".""".join(layer.split("""."""       )[:2]       ) for layer in unet_state_dict if """input_blocks""" in layer}       )
  		UpperCamelCase__         =	{
  		    layer_id: [key for key in unet_state_dict if F"""input_blocks.{layer_id}""" in key]
  		    for layer_id in range(_lowercase       )
  		}
  
  		# Retrieves the keys for the middle blocks only
  		UpperCamelCase__         =	len({""".""".join(layer.split("""."""       )[:2]       ) for layer in unet_state_dict if """middle_block""" in layer}       )
  		UpperCamelCase__         =	{
  		    layer_id: [key for key in unet_state_dict if F"""middle_block.{layer_id}""" in key]
  		    for layer_id in range(_lowercase       )
  		}
  
  		# Retrieves the keys for the output blocks only
  		UpperCamelCase__         =	len({""".""".join(layer.split("""."""       )[:2]       ) for layer in unet_state_dict if """output_blocks""" in layer}       )
  		UpperCamelCase__         =	{
  		    layer_id: [key for key in unet_state_dict if F"""output_blocks.{layer_id}""" in key]
  		    for layer_id in range(_lowercase       )
  		}
  
  		for i in range(1 ,						_lowercase       ):
  				UpperCamelCase__         =	(i - 1) // (config["""layers_per_block"""] + 1)
  				UpperCamelCase__         =	(i - 1) % (config["""layers_per_block"""] + 1)
  
  				UpperCamelCase__         =	[
  				    key for key in input_blocks[i] if F"""input_blocks.{i}.0""" in key and F"""input_blocks.{i}.0.op""" not in key
  				]
  				UpperCamelCase__         =	[key for key in input_blocks[i] if F"""input_blocks.{i}.1""" in key]
  				UpperCamelCase__         =	[key for key in input_blocks[i] if F"""input_blocks.{i}.2""" in key]
  
  				if F"""input_blocks.{i}.op.weight""" in unet_state_dict:
  						UpperCamelCase__         =	unet_state_dict.pop(
  						    F"""input_blocks.{i}.op.weight"""       )
  						UpperCamelCase__         =	unet_state_dict.pop(
  						    F"""input_blocks.{i}.op.bias"""       )
  
  				UpperCamelCase__         =	renew_resnet_paths(_lowercase       )
  				UpperCamelCase__         =	{"""old""": F"""input_blocks.{i}.0""", """new""": F"""down_blocks.{block_id}.resnets.{layer_in_block_id}"""}
  				assign_to_checkpoint(
  				    _lowercase ,						_lowercase ,						_lowercase ,						additional_replacements=[meta_path] ,						config=_lowercase       )
  
  				UpperCamelCase__         =	[key for key in resnets if """temopral_conv""" in key]
  				UpperCamelCase__         =	renew_temp_conv_paths(_lowercase       )
  				UpperCamelCase__         =	{
  				    """old""": F"""input_blocks.{i}.0.temopral_conv""",
  				    """new""": F"""down_blocks.{block_id}.temp_convs.{layer_in_block_id}""",
  				}
  				assign_to_checkpoint(
  				    _lowercase ,						_lowercase ,						_lowercase ,						additional_replacements=[meta_path] ,						config=_lowercase       )
  
  				if len(_lowercase       ):
  						UpperCamelCase__         =	renew_attention_paths(_lowercase       )
  						UpperCamelCase__         =	{"""old""": F"""input_blocks.{i}.1""", """new""": F"""down_blocks.{block_id}.attentions.{layer_in_block_id}"""}
  						assign_to_checkpoint(
  						    _lowercase ,						_lowercase ,						_lowercase ,						additional_replacements=[meta_path] ,						config=_lowercase       )
  
  				if len(_lowercase       ):
  						UpperCamelCase__         =	renew_attention_paths(_lowercase       )
  						UpperCamelCase__         =	{
  						    """old""": F"""input_blocks.{i}.2""",
  						    """new""": F"""down_blocks.{block_id}.temp_attentions.{layer_in_block_id}""",
  						}
  						assign_to_checkpoint(
  						    _lowercase ,						_lowercase ,						_lowercase ,						additional_replacements=[meta_path] ,						config=_lowercase       )
  
  		UpperCamelCase__         =	middle_blocks[0]
  		UpperCamelCase__         =	[key for key in resnet_a if """temopral_conv""" in key]
  		UpperCamelCase__         =	middle_blocks[1]
  		UpperCamelCase__         =	middle_blocks[2]
  		UpperCamelCase__         =	middle_blocks[3]
  		UpperCamelCase__         =	[key for key in resnet_a if """temopral_conv""" in key]
  
  		UpperCamelCase__         =	renew_resnet_paths(_lowercase       )
  		UpperCamelCase__         =	{"""old""": """middle_block.0""", """new""": """mid_block.resnets.0"""}
  		assign_to_checkpoint(
  		    _lowercase ,						_lowercase ,						_lowercase ,						config=_lowercase ,						additional_replacements=[meta_path]       )
  
  		UpperCamelCase__         =	renew_temp_conv_paths(_lowercase       )
  		UpperCamelCase__         =	{"""old""": """middle_block.0.temopral_conv""", """new""": """mid_block.temp_convs.0"""}
  		assign_to_checkpoint(
  		    _lowercase ,						_lowercase ,						_lowercase ,						config=_lowercase ,						additional_replacements=[meta_path]       )
  
  		UpperCamelCase__         =	renew_resnet_paths(_lowercase       )
  		UpperCamelCase__         =	{"""old""": """middle_block.3""", """new""": """mid_block.resnets.1"""}
  		assign_to_checkpoint(
  		    _lowercase ,						_lowercase ,						_lowercase ,						config=_lowercase ,						additional_replacements=[meta_path]       )
  
  		UpperCamelCase__         =	renew_temp_conv_paths(_lowercase       )
  		UpperCamelCase__         =	{"""old""": """middle_block.3.temopral_conv""", """new""": """mid_block.temp_convs.1"""}
  		assign_to_checkpoint(
  		    _lowercase ,						_lowercase ,						_lowercase ,						config=_lowercase ,						additional_replacements=[meta_path]       )
  
  		UpperCamelCase__         =	renew_attention_paths(_lowercase       )
  		UpperCamelCase__         =	{"""old""": """middle_block.1""", """new""": """mid_block.attentions.0"""}
  		assign_to_checkpoint(
  		    _lowercase ,						_lowercase ,						_lowercase ,						additional_replacements=[meta_path] ,						config=_lowercase       )
  
  		UpperCamelCase__         =	renew_attention_paths(_lowercase       )
  		UpperCamelCase__         =	{"""old""": """middle_block.2""", """new""": """mid_block.temp_attentions.0"""}
  		assign_to_checkpoint(
  		    _lowercase ,						_lowercase ,						_lowercase ,						additional_replacements=[meta_path] ,						config=_lowercase       )
  
  		for i in range(_lowercase       ):
  				UpperCamelCase__         =	i // (config["""layers_per_block"""] + 1)
  				UpperCamelCase__         =	i % (config["""layers_per_block"""] + 1)
  				UpperCamelCase__         =	[shave_segments(_lowercase ,						2       ) for name in output_blocks[i]]
  				UpperCamelCase__         =	{}
  
  				for layer in output_block_layers:
  						UpperCamelCase__  ,			UpperCamelCase__         =	layer.split("""."""       )[0], shave_segments(_lowercase ,						1       )
  						if layer_id in output_block_list:
  								output_block_list[layer_id].append(_lowercase       )
  						else:
  								UpperCamelCase__         =	[layer_name]
  
  				if len(_lowercase       ) > 1:
  						UpperCamelCase__         =	[key for key in output_blocks[i] if F"""output_blocks.{i}.0""" in key]
  						UpperCamelCase__         =	[key for key in output_blocks[i] if F"""output_blocks.{i}.1""" in key]
  						UpperCamelCase__         =	[key for key in output_blocks[i] if F"""output_blocks.{i}.2""" in key]
  
  						UpperCamelCase__         =	renew_resnet_paths(_lowercase       )
  						UpperCamelCase__         =	renew_resnet_paths(_lowercase       )
  
  						UpperCamelCase__         =	{"""old""": F"""output_blocks.{i}.0""", """new""": F"""up_blocks.{block_id}.resnets.{layer_in_block_id}"""}
  						assign_to_checkpoint(
  						    _lowercase ,						_lowercase ,						_lowercase ,						additional_replacements=[meta_path] ,						config=_lowercase       )
  
  						UpperCamelCase__         =	[key for key in resnets if """temopral_conv""" in key]
  						UpperCamelCase__         =	renew_temp_conv_paths(_lowercase       )
  						UpperCamelCase__         =	{
  						    """old""": F"""output_blocks.{i}.0.temopral_conv""",
  						    """new""": F"""up_blocks.{block_id}.temp_convs.{layer_in_block_id}""",
  						}
  						assign_to_checkpoint(
  						    _lowercase ,						_lowercase ,						_lowercase ,						additional_replacements=[meta_path] ,						config=_lowercase       )
  
  						UpperCamelCase__         =	{k: sorted(_lowercase       ) for k, v in output_block_list.items()}
  						if ["conv.bias", "conv.weight"] in output_block_list.values():
  								UpperCamelCase__         =	list(output_block_list.values()       ).index(["""conv.bias""", """conv.weight"""]       )
  								UpperCamelCase__         =	unet_state_dict[
  								    F"""output_blocks.{i}.{index}.conv.weight"""
  								]
  								UpperCamelCase__         =	unet_state_dict[
  								    F"""output_blocks.{i}.{index}.conv.bias"""
  								]
  
  								# Clear attentions as they have been attributed above.
  								if len(_lowercase       ) == 2:
  										UpperCamelCase__         =	[]
  
  						if len(_lowercase       ):
  								UpperCamelCase__         =	renew_attention_paths(_lowercase       )
  								UpperCamelCase__         =	{
  								    """old""": F"""output_blocks.{i}.1""",
  								    """new""": F"""up_blocks.{block_id}.attentions.{layer_in_block_id}""",
  								}
  								assign_to_checkpoint(
  								    _lowercase ,						_lowercase ,						_lowercase ,						additional_replacements=[meta_path] ,						config=_lowercase       )
  
  						if len(_lowercase       ):
  								UpperCamelCase__         =	renew_attention_paths(_lowercase       )
  								UpperCamelCase__         =	{
  								    """old""": F"""output_blocks.{i}.2""",
  								    """new""": F"""up_blocks.{block_id}.temp_attentions.{layer_in_block_id}""",
  								}
  								assign_to_checkpoint(
  								    _lowercase ,						_lowercase ,						_lowercase ,						additional_replacements=[meta_path] ,						config=_lowercase       )
  				else:
  						UpperCamelCase__         =	renew_resnet_paths(_lowercase ,						n_shave_prefix_segments=1       )
  						for path in resnet_0_paths:
  								UpperCamelCase__         =	""".""".join(["""output_blocks""", str(_lowercase       ), path["""old"""]]       )
  								UpperCamelCase__         =	""".""".join(["""up_blocks""", str(_lowercase       ), """resnets""", str(_lowercase       ), path["""new"""]]       )
  								UpperCamelCase__         =	unet_state_dict[old_path]
  
  						UpperCamelCase__         =	[l for l in output_block_layers if """temopral_conv""" in l]
  						for path in temopral_conv_paths:
  								UpperCamelCase__         =	path.split("""temopral_conv."""       )[-1]
  								UpperCamelCase__         =	""".""".join(["""output_blocks""", str(_lowercase       ), str(_lowercase       ), """temopral_conv""", pruned_path]       )
  								UpperCamelCase__         =	""".""".join(["""up_blocks""", str(_lowercase       ), """temp_convs""", str(_lowercase       ), pruned_path]       )
  								UpperCamelCase__         =	unet_state_dict[old_path]
  
  		return new_checkpoint
  
  
  if __name__ == "__main__":
  				__lowerCAmelCase     =			argparse.ArgumentParser()
  
  				parser.add_argument(
  				    "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert."
  				)
  				parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.")
  				__lowerCAmelCase     =			parser.parse_args()
  
  				__lowerCAmelCase     =			torch.load(args.checkpoint_path, map_location="cpu")
  				__lowerCAmelCase     =			UNetaDConditionModel()
  
  				__lowerCAmelCase     =			convert_ldm_unet_checkpoint(unet_checkpoint, unet.config)
  
  				__lowerCAmelCase     =			set(unet.state_dict().keys()) - set(converted_ckpt.keys())
  				__lowerCAmelCase     =			set(converted_ckpt.keys()) - set(unet.state_dict().keys())
  
  				assert len(diff_a) == len(diff_a) == 0, "Converted weights don't match"
  
  				# load state_dict
  				unet.load_state_dict(converted_ckpt)
  
  				unet.save_pretrained(args.dump_path)
  
  				# -- finish converting the unet --
  
  
  
  
  
  
style_context: 
  
  
  
  
  
  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()
  __lowerCAmelCase     =			logging.get_logger(__name__)
  
  def     a						(      _lowercase   :     Union[str, Any]       )     ->str:
  		UpperCamelCase__         =	tf.image.resize(_lowercase ,						(3_4_6, 3_4_6)       )
  		UpperCamelCase__         =	tf.image.crop_to_bounding_box(_lowercase ,						(3_4_6 - 2_8_9) // 2 ,						(3_4_6 - 2_8_9) // 2 ,						2_8_9 ,						2_8_9       )
  		return image
  
  def     a						(      )     ->Optional[int]:
  		UpperCamelCase__         =	EfficientNetConfig.from_pretrained("""google/efficientnet-b7"""       )
  		UpperCamelCase__         =	2_8_9
  		UpperCamelCase__         =	6_4_0
  		UpperCamelCase__         =	{"""0""": """LABEL_0""", """1""": """LABEL_1"""}
  		UpperCamelCase__         =	{"""LABEL_0""": 0, """LABEL_1""": 1}
  		UpperCamelCase__         =	[]
  
  		UpperCamelCase__         =	BertConfig()
  		UpperCamelCase__         =	AlignConfig.from_text_vision_configs(
  		    text_config=_lowercase ,						vision_config=_lowercase ,						projection_dim=6_4_0       )
  		return config
  
  def     a						(      )     ->Optional[int]:
  		UpperCamelCase__         =	"""http://images.cocodataset.org/val2017/000000039769.jpg"""
  		UpperCamelCase__         =	Image.open(requests.get(_lowercase ,						stream=_lowercase       ).raw       )
  		return im
  
  def     a						(      )     ->Optional[int]:
  		UpperCamelCase__         =	EfficientNetImageProcessor(
  		    do_center_crop=_lowercase ,						rescale_factor=1 / 127.5 ,						rescale_offset=_lowercase ,						do_normalize=_lowercase ,						include_top=_lowercase ,						resample=Image.BILINEAR ,						)
  		UpperCamelCase__         =	BertTokenizer.from_pretrained("""bert-base-uncased"""       )
  		UpperCamelCase__         =	6_4
  		UpperCamelCase__         =	AlignProcessor(image_processor=_lowercase ,						tokenizer=_lowercase       )
  		return processor
  
  def     a						(      _lowercase   :     str       )     ->int:
  		# EfficientNet image encoder
  		UpperCamelCase__         =	[v.split("""_"""       )[0].split("""block"""       )[1] for v in original_param_names if v.startswith("""block"""       )]
  		UpperCamelCase__         =	list(set(_lowercase       )       )
  		UpperCamelCase__         =	sorted(_lowercase       )
  		UpperCamelCase__         =	len(_lowercase       )
  		UpperCamelCase__         =	{b: str(_lowercase       ) for b, i in zip(_lowercase ,						range(_lowercase       )       )}
  
  		UpperCamelCase__         =	[]
  		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__         =	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__         =	{}
  		for item in rename_keys:
  				if item[0] in original_param_names:
  						UpperCamelCase__         =	"""vision_model.""" + item[1]
  
      # BERT text encoder
  		UpperCamelCase__         =	[]
  		UpperCamelCase__         =	"""tf_bert_model/bert"""
  		UpperCamelCase__         =	"""text_model"""
  		for i in range(1_2       ):
  				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__         =	item[1]
  		return key_mapping
  
  def     a						(      _lowercase   :     int ,						_lowercase   :     str ,						_lowercase   :     int       )     ->Optional[Any]:
  		list(hf_params.keys()       )
  
  		for key, value in tf_params.items():
  				if key not in key_mapping:
  						continue
  
  				UpperCamelCase__         =	key_mapping[key]
  				if "_conv" in key and "kernel" in key:
  						UpperCamelCase__         =	torch.from_numpy(_lowercase       ).permute(3 ,						2 ,						0 ,						1       )
  				elif "embeddings" in key:
  						UpperCamelCase__         =	torch.from_numpy(_lowercase       )
  				elif "depthwise_kernel" in key:
  						UpperCamelCase__         =	torch.from_numpy(_lowercase       ).permute(2 ,						3 ,						0 ,						1       )
  				elif "kernel" in key:
  						UpperCamelCase__         =	torch.from_numpy(np.transpose(_lowercase       )       )
  				elif "temperature" in key:
  						UpperCamelCase__         =	value
  				elif "bn/gamma" or "bn/beta" in key:
  						UpperCamelCase__         =	torch.from_numpy(np.transpose(_lowercase       )       ).squeeze()
  				else:
  						UpperCamelCase__         =	torch.from_numpy(_lowercase       )
  
  				# Replace HF parameters with original TF model parameters
  				hf_params[hf_key].copy_(_lowercase       )
  
  @torch.no_grad()
  def     a						(      _lowercase   :     List[str] ,						_lowercase   :     Any ,						_lowercase   :     Dict ,						_lowercase   :     Optional[Any]       )     ->str:
  		UpperCamelCase__         =	6_4
  		UpperCamelCase__         =	Tokenizer(_lowercase       )
  		UpperCamelCase__         =	align.Align("""efficientnet-b7""" ,						"""bert-base""" ,						6_4_0 ,						_lowercase ,						tok.get_vocab_size()       )
  		original_model.compile()
  		original_model.load_weights(_lowercase       )
  
  		UpperCamelCase__         =	original_model.trainable_variables
  		UpperCamelCase__         =	original_model.non_trainable_variables
  		UpperCamelCase__         =	{param.name: param.numpy() for param in tf_params}
  		for param in tf_non_train_params:
  				UpperCamelCase__         =	param.numpy()
  		UpperCamelCase__         =	list(tf_params.keys()       )
  
  		# Load HuggingFace model
  		UpperCamelCase__         =	get_align_config()
  		UpperCamelCase__         =	AlignModel(_lowercase       ).eval()
  		UpperCamelCase__         =	hf_model.state_dict()
  
  		# Create src-to-dst parameter name mapping dictionary
  		print("""Converting parameters..."""       )
  		UpperCamelCase__         =	rename_keys(_lowercase       )
  		replace_params(_lowercase ,						_lowercase ,						_lowercase       )
  
  		# Initialize processor
  		UpperCamelCase__         =	get_processor()
  		UpperCamelCase__         =	processor(
  		    images=prepare_img() ,						text="""A picture of a cat""" ,						padding="""max_length""" ,						max_length=6_4 ,						return_tensors="""pt"""       )
  
  		# HF model inference
  		hf_model.eval()
  		with torch.no_grad():
  				UpperCamelCase__         =	hf_model(**_lowercase       )
  
  		UpperCamelCase__         =	outputs.image_embeds.detach().numpy()
  		UpperCamelCase__         =	outputs.text_embeds.detach().numpy()
  
  		# Original model inference
  		UpperCamelCase__         =	False
  		UpperCamelCase__         =	EfficientNetImageProcessor(
  		    do_center_crop=_lowercase ,						do_rescale=_lowercase ,						do_normalize=_lowercase ,						include_top=_lowercase ,						resample=Image.BILINEAR ,						)
  		UpperCamelCase__         =	tf_image_processor(images=prepare_img() ,						return_tensors="""tf""" ,						data_format="""channels_last"""       )["""pixel_values"""]
  		UpperCamelCase__         =	tok(tf.constant(["""A picture of a cat"""]       )       )
  
  		UpperCamelCase__         =	original_model.image_encoder(_lowercase ,						training=_lowercase       )
  		UpperCamelCase__         =	original_model.text_encoder(_lowercase ,						training=_lowercase       )
  
  		UpperCamelCase__         =	tf.nn.la_normalize(_lowercase ,						axis=-1       )
  		UpperCamelCase__         =	tf.nn.la_normalize(_lowercase ,						axis=-1       )
  
  		# Check whether original and HF model outputs match  -> np.allclose
  		if not np.allclose(_lowercase ,						_lowercase ,						atol=1e-3       ):
  				raise ValueError("""The predicted image features are not the same."""       )
  		if not np.allclose(_lowercase ,						_lowercase ,						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(_lowercase       ):
  						os.mkdir(_lowercase       )
  				# Save converted model and image processor
  				hf_model.save_pretrained(_lowercase       )
  				processor.save_pretrained(_lowercase       )
  
  		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__":
  				__lowerCAmelCase     =			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")
  
  				__lowerCAmelCase     =			parser.parse_args()
  				convert_align_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)
  
  
  
  
  
  