code: 
  
  
  
  
  
  import inspect
  from typing import List, Optional, Union
  
  import torch
  from torch import nn
  from torch.nn import functional as F
  from torchvision import transforms
  from transformers import CLIPImageProcessor, CLIPModel, CLIPTextModel, CLIPTokenizer
  
  from diffusers import (
      AutoencoderKL,
      DDIMScheduler,
      DiffusionPipeline,
      DPMSolverMultistepScheduler,
      LMSDiscreteScheduler,
      PNDMScheduler,
      UNetaDConditionModel,
  )
  from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput
  
  
  
  
  class   _UpperCAmelCase	(   nn.Module					):
  
  
  
  
  					def __init__(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase=1.0       )						->  Tuple:
  							super().__init__()
  
  							UpperCamelCase__         =	cut_size
  							UpperCamelCase__         =	cut_power
  
  
  
  
  					def 	snake_case_		(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase       )						->  Optional[int]:
  							UpperCamelCase__  ,			UpperCamelCase__         =	pixel_values.shape[2:4]
  							UpperCamelCase__         =	min(__lowerCAmelCase    ,      __lowerCAmelCase       )
  							UpperCamelCase__         =	min(__lowerCAmelCase    ,      __lowerCAmelCase    ,      self.cut_size       )
  							UpperCamelCase__         =	[]
  							for _ in range(__lowerCAmelCase       ):
  									UpperCamelCase__         =	int(torch.rand([]       ) ** self.cut_power * (max_size - min_size) + min_size       )
  									UpperCamelCase__         =	torch.randint(0    ,      sideX - size + 1    ,      ()       )
  									UpperCamelCase__         =	torch.randint(0    ,      sideY - size + 1    ,      ()       )
  									UpperCamelCase__         =	pixel_values[:, :, offsety : offsety + size, offsetx : offsetx + size]
  									cutouts.append(F.adaptive_avg_poolad(__lowerCAmelCase    ,      self.cut_size       )       )
  							return torch.cat(__lowerCAmelCase       )
  
  
  
  def     a						(      _lowercase   :     Tuple ,						_lowercase   :     int       )     ->Tuple:
  		UpperCamelCase__         =	F.normalize(_lowercase ,						dim=-1       )
  		UpperCamelCase__         =	F.normalize(_lowercase ,						dim=-1       )
  		return (x - y).norm(dim=-1       ).div(2       ).arcsin().pow(2       ).mul(2       )
  
  def     a						(      _lowercase   :     List[Any] ,						_lowercase   :     str       )     ->Any:
  		for param in model.parameters():
  				UpperCamelCase__         =	value
  
  
  
  
  class   _UpperCAmelCase	(   _lowercase					):
  
  
  
  
  					def __init__(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      )						->  Optional[Any]:
  							super().__init__()
  							self.register_modules(
  							    vae=__lowerCAmelCase    ,      text_encoder=__lowerCAmelCase    ,      clip_model=__lowerCAmelCase    ,      tokenizer=__lowerCAmelCase    ,      unet=__lowerCAmelCase    ,      scheduler=__lowerCAmelCase    ,      feature_extractor=__lowerCAmelCase    ,      )
  
  							UpperCamelCase__         =	transforms.Normalize(mean=feature_extractor.image_mean    ,      std=feature_extractor.image_std       )
  							UpperCamelCase__         =	(
  							    feature_extractor.size
  							    if isinstance(feature_extractor.size    ,      __lowerCAmelCase       )
  							    else feature_extractor.size["""shortest_edge"""]
  							)
  							UpperCamelCase__         =	MakeCutouts(self.cut_out_size       )
  
  							set_requires_grad(self.text_encoder    ,      __lowerCAmelCase       )
  							set_requires_grad(self.clip_model    ,      __lowerCAmelCase       )
  
  
  
  
  					def 	snake_case_		(						self    ,      __lowerCAmelCase = "auto"       )						->  Optional[int]:
  							if slice_size == "auto":
  									# half the attention head size is usually a good trade-off between
  									# speed and memory
  									UpperCamelCase__         =	self.unet.config.attention_head_dim // 2
  							self.unet.set_attention_slice(__lowerCAmelCase       )
  
  
  
  
  					def 	snake_case_		(						self       )						->  str:
  							self.enable_attention_slicing(__lowerCAmelCase       )
  
  
  
  
  					def 	snake_case_		(						self       )						->  Any:
  							set_requires_grad(self.vae    ,      __lowerCAmelCase       )
  
  
  
  
  					def 	snake_case_		(						self       )						->  Optional[int]:
  							set_requires_grad(self.vae    ,      __lowerCAmelCase       )
  
  
  
  
  					def 	snake_case_		(						self       )						->  Union[str, Any]:
  							set_requires_grad(self.unet    ,      __lowerCAmelCase       )
  
  
  
  
  					def 	snake_case_		(						self       )						->  List[str]:
  							set_requires_grad(self.unet    ,      __lowerCAmelCase       )
  
  
  
  
  					@torch.enable_grad()
  					def 	snake_case_		(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase=True    ,      )						->  List[str]:
  							UpperCamelCase__         =	latents.detach().requires_grad_()
  
  							UpperCamelCase__         =	self.scheduler.scale_model_input(__lowerCAmelCase    ,      __lowerCAmelCase       )
  
  							# predict the noise residual
  							UpperCamelCase__         =	self.unet(__lowerCAmelCase    ,      __lowerCAmelCase    ,      encoder_hidden_states=__lowerCAmelCase       ).sample
  
  							if isinstance(self.scheduler    ,      (PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)       ):
  									UpperCamelCase__         =	self.scheduler.alphas_cumprod[timestep]
  									UpperCamelCase__         =	1 - alpha_prod_t
  									# compute predicted original sample from predicted noise also called
  									# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
  									UpperCamelCase__         =	(latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5
  
  									UpperCamelCase__         =	torch.sqrt(__lowerCAmelCase       )
  									UpperCamelCase__         =	pred_original_sample * (fac) + latents * (1 - fac)
  							elif isinstance(self.scheduler    ,      __lowerCAmelCase       ):
  									UpperCamelCase__         =	self.scheduler.sigmas[index]
  									UpperCamelCase__         =	latents - sigma * noise_pred
  							else:
  									raise ValueError(F"""scheduler type {type(self.scheduler       )} not supported"""       )
  
  							UpperCamelCase__         =	1 / self.vae.config.scaling_factor * sample
  							UpperCamelCase__         =	self.vae.decode(__lowerCAmelCase       ).sample
  							UpperCamelCase__         =	(image / 2 + 0.5).clamp(0    ,      1       )
  
  							if use_cutouts:
  									UpperCamelCase__         =	self.make_cutouts(__lowerCAmelCase    ,      __lowerCAmelCase       )
  							else:
  									UpperCamelCase__         =	transforms.Resize(self.cut_out_size       )(__lowerCAmelCase       )
  							UpperCamelCase__         =	self.normalize(__lowerCAmelCase       ).to(latents.dtype       )
  
  							UpperCamelCase__         =	self.clip_model.get_image_features(__lowerCAmelCase       )
  							UpperCamelCase__         =	image_embeddings_clip / image_embeddings_clip.norm(p=2    ,      dim=-1    ,      keepdim=__lowerCAmelCase       )
  
  							if use_cutouts:
  									UpperCamelCase__         =	spherical_dist_loss(__lowerCAmelCase    ,      __lowerCAmelCase       )
  									UpperCamelCase__         =	dists.view([num_cutouts, sample.shape[0], -1]       )
  									UpperCamelCase__         =	dists.sum(2       ).mean(0       ).sum() * clip_guidance_scale
  							else:
  									UpperCamelCase__         =	spherical_dist_loss(__lowerCAmelCase    ,      __lowerCAmelCase       ).mean() * clip_guidance_scale
  
  							UpperCamelCase__         =	-torch.autograd.grad(__lowerCAmelCase    ,      __lowerCAmelCase       )[0]
  
  							if isinstance(self.scheduler    ,      __lowerCAmelCase       ):
  									UpperCamelCase__         =	latents.detach() + grads * (sigma**2)
  									UpperCamelCase__         =	noise_pred_original
  							else:
  									UpperCamelCase__         =	noise_pred_original - torch.sqrt(__lowerCAmelCase       ) * grads
  							return noise_pred, latents
  
  
  
  
  					@torch.no_grad()
  					def __call__(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase = 512    ,      __lowerCAmelCase = 512    ,      __lowerCAmelCase = 50    ,      __lowerCAmelCase = 7.5    ,      __lowerCAmelCase = 1    ,      __lowerCAmelCase = 0.0    ,      __lowerCAmelCase = 100    ,      __lowerCAmelCase = None    ,      __lowerCAmelCase = 4    ,      __lowerCAmelCase = True    ,      __lowerCAmelCase = None    ,      __lowerCAmelCase = None    ,      __lowerCAmelCase = "pil"    ,      __lowerCAmelCase = True    ,      )						->  List[str]:
  							if isinstance(__lowerCAmelCase    ,      __lowerCAmelCase       ):
  									UpperCamelCase__         =	1
  							elif isinstance(__lowerCAmelCase    ,      __lowerCAmelCase       ):
  									UpperCamelCase__         =	len(__lowerCAmelCase       )
  							else:
  									raise ValueError(F"""`prompt` has to be of type `str` or `list` but is {type(__lowerCAmelCase       )}"""       )
  
  							if height % 8 != 0 or width % 8 != 0:
  									raise ValueError(F"""`height` and `width` have to be divisible by 8 but are {height} and {width}."""       )
  
  							# get prompt text embeddings
  							UpperCamelCase__         =	self.tokenizer(
  							    __lowerCAmelCase    ,      padding="""max_length"""    ,      max_length=self.tokenizer.model_max_length    ,      truncation=__lowerCAmelCase    ,      return_tensors="""pt"""    ,      )
  							UpperCamelCase__         =	self.text_encoder(text_input.input_ids.to(self.device       )       )[0]
  							# duplicate text embeddings for each generation per prompt
  							UpperCamelCase__         =	text_embeddings.repeat_interleave(__lowerCAmelCase    ,      dim=0       )
  
  							if clip_guidance_scale > 0:
  									if clip_prompt is not None:
  											UpperCamelCase__         =	self.tokenizer(
  											    __lowerCAmelCase    ,      padding="""max_length"""    ,      max_length=self.tokenizer.model_max_length    ,      truncation=__lowerCAmelCase    ,      return_tensors="""pt"""    ,      ).input_ids.to(self.device       )
  									else:
  											UpperCamelCase__         =	text_input.input_ids.to(self.device       )
  									UpperCamelCase__         =	self.clip_model.get_text_features(__lowerCAmelCase       )
  									UpperCamelCase__         =	text_embeddings_clip / text_embeddings_clip.norm(p=2    ,      dim=-1    ,      keepdim=__lowerCAmelCase       )
  									# duplicate text embeddings clip for each generation per prompt
  									UpperCamelCase__         =	text_embeddings_clip.repeat_interleave(__lowerCAmelCase    ,      dim=0       )
  
  							# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
  							# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
  							# corresponds to doing no classifier free guidance.
  							UpperCamelCase__         =	guidance_scale > 1.0
  							# get unconditional embeddings for classifier free guidance
  							if do_classifier_free_guidance:
  									UpperCamelCase__         =	text_input.input_ids.shape[-1]
  									UpperCamelCase__         =	self.tokenizer([""""""]    ,      padding="""max_length"""    ,      max_length=__lowerCAmelCase    ,      return_tensors="""pt"""       )
  									UpperCamelCase__         =	self.text_encoder(uncond_input.input_ids.to(self.device       )       )[0]
  									# duplicate unconditional embeddings for each generation per prompt
  									UpperCamelCase__         =	uncond_embeddings.repeat_interleave(__lowerCAmelCase    ,      dim=0       )
  
  									# For classifier free guidance, we need to do two forward passes.
  									# Here we concatenate the unconditional and text embeddings into a single batch
  									# to avoid doing two forward passes
  									UpperCamelCase__         =	torch.cat([uncond_embeddings, text_embeddings]       )
  
  							# get the initial random noise unless the user supplied it
  
  							# Unlike in other pipelines, latents need to be generated in the target device
  							# for 1-to-1 results reproducibility with the CompVis implementation.
  							# However this currently doesn't work in `mps`.
  							UpperCamelCase__         =	(batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
  							UpperCamelCase__         =	text_embeddings.dtype
  							if latents is None:
  									if self.device.type == "mps":
  											# randn does not work reproducibly on mps
  											UpperCamelCase__         =	torch.randn(__lowerCAmelCase    ,      generator=__lowerCAmelCase    ,      device="""cpu"""    ,      dtype=__lowerCAmelCase       ).to(
  											    self.device       )
  									else:
  											UpperCamelCase__         =	torch.randn(__lowerCAmelCase    ,      generator=__lowerCAmelCase    ,      device=self.device    ,      dtype=__lowerCAmelCase       )
  							else:
  									if latents.shape != latents_shape:
  											raise ValueError(F"""Unexpected latents shape, got {latents.shape}, expected {latents_shape}"""       )
  									UpperCamelCase__         =	latents.to(self.device       )
  
  							# set timesteps
  							UpperCamelCase__         =	"""offset""" in set(inspect.signature(self.scheduler.set_timesteps       ).parameters.keys()       )
  							UpperCamelCase__         =	{}
  							if accepts_offset:
  									UpperCamelCase__         =	1
  
  							self.scheduler.set_timesteps(__lowerCAmelCase    ,      **__lowerCAmelCase       )
  
  							# Some schedulers like PNDM have timesteps as arrays
  							# It's more optimized to move all timesteps to correct device beforehand
  							UpperCamelCase__         =	self.scheduler.timesteps.to(self.device       )
  
  							# scale the initial noise by the standard deviation required by the scheduler
  							UpperCamelCase__         =	latents * self.scheduler.init_noise_sigma
  
  							# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
  							# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
  							# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
  							# and should be between [0, 1]
  							UpperCamelCase__         =	"""eta""" in set(inspect.signature(self.scheduler.step       ).parameters.keys()       )
  							UpperCamelCase__         =	{}
  							if accepts_eta:
  									UpperCamelCase__         =	eta
  
  							# check if the scheduler accepts generator
  							UpperCamelCase__         =	"""generator""" in set(inspect.signature(self.scheduler.step       ).parameters.keys()       )
  							if accepts_generator:
  									UpperCamelCase__         =	generator
  
  							for i, t in enumerate(self.progress_bar(__lowerCAmelCase       )       ):
  									# expand the latents if we are doing classifier free guidance
  									UpperCamelCase__         =	torch.cat([latents] * 2       ) if do_classifier_free_guidance else latents
  									UpperCamelCase__         =	self.scheduler.scale_model_input(__lowerCAmelCase    ,      __lowerCAmelCase       )
  
  									# predict the noise residual
  									UpperCamelCase__         =	self.unet(__lowerCAmelCase    ,      __lowerCAmelCase    ,      encoder_hidden_states=__lowerCAmelCase       ).sample
  
  									# perform classifier free guidance
  									if do_classifier_free_guidance:
  											UpperCamelCase__  ,			UpperCamelCase__         =	noise_pred.chunk(2       )
  											UpperCamelCase__         =	noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
  
  									# perform clip guidance
  									if clip_guidance_scale > 0:
  											UpperCamelCase__         =	(
  											    text_embeddings.chunk(2       )[1] if do_classifier_free_guidance else text_embeddings
  											)
  											UpperCamelCase__  ,			UpperCamelCase__         =	self.cond_fn(
  											    __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      )
  
  									# compute the previous noisy sample x_t -> x_t-1
  									UpperCamelCase__         =	self.scheduler.step(__lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      **__lowerCAmelCase       ).prev_sample
  
  							# scale and decode the image latents with vae
  							UpperCamelCase__         =	1 / self.vae.config.scaling_factor * latents
  							UpperCamelCase__         =	self.vae.decode(__lowerCAmelCase       ).sample
  
  							UpperCamelCase__         =	(image / 2 + 0.5).clamp(0    ,      1       )
  							UpperCamelCase__         =	image.cpu().permute(0    ,      2    ,      3    ,      1       ).numpy()
  
  							if output_type == "pil":
  									UpperCamelCase__         =	self.numpy_to_pil(__lowerCAmelCase       )
  
  							if not return_dict:
  									return (image, None)
  
  							return StableDiffusionPipelineOutput(images=__lowerCAmelCase    ,      nsfw_content_detected=__lowerCAmelCase       )
  
  
  
  
  
  
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)
  
  
  
  
  
  