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()
_SCREAMING_SNAKE_CASE:						Union[str, Any]    	=       logging.get_logger(__name__)






def    UpperCAmelCase__(						lowerCamelCase_  ):



			"""simple docstring"""





			_snake_case              =      tf.image.resize(lowerCamelCase_							,	(346, 346)  )
			_snake_case              =      tf.image.crop_to_bounding_box(lowerCamelCase_							,	(346 - 289) // 2							,	(346 - 289) // 2							,	289							,	289  )
			return image






def    UpperCAmelCase__(						):



			"""simple docstring"""





			_snake_case              =      EfficientNetConfig.from_pretrained("google/efficientnet-b7"  )
			_snake_case              =      289
			_snake_case              =      640
			_snake_case              =      {"0": "LABEL_0", "1": "LABEL_1"}
			_snake_case              =      {"LABEL_0": 0, "LABEL_1": 1}
			_snake_case              =      []

			_snake_case              =      BertConfig()
			_snake_case              =      AlignConfig.from_text_vision_configs(
			    text_config=lowerCamelCase_							,	vision_config=lowerCamelCase_							,	projection_dim=640  )
			return config






def    UpperCAmelCase__(						):



			"""simple docstring"""





			_snake_case              =      "http://images.cocodataset.org/val2017/000000039769.jpg"
			_snake_case              =      Image.open(requests.get(lowerCamelCase_							,	stream=lowerCamelCase_  ).raw  )
			return im






def    UpperCAmelCase__(						):



			"""simple docstring"""





			_snake_case              =      EfficientNetImageProcessor(
			    do_center_crop=lowerCamelCase_							,	rescale_factor=1 / 1_2_7.5							,	rescale_offset=lowerCamelCase_							,	do_normalize=lowerCamelCase_							,	include_top=lowerCamelCase_							,	resample=Image.BILINEAR							,	)
			_snake_case              =      BertTokenizer.from_pretrained("bert-base-uncased"  )
			_snake_case              =      64
			_snake_case              =      AlignProcessor(image_processor=lowerCamelCase_							,	tokenizer=lowerCamelCase_  )
			return processor






def    UpperCAmelCase__(						lowerCamelCase_  ):



			"""simple docstring"""





			_snake_case              =      [v.split("_"  )[0].split("block"  )[1] for v in original_param_names if v.startswith("block"  )]
			_snake_case              =      list(set(lowerCamelCase_  )  )
			_snake_case              =      sorted(lowerCamelCase_  )
			_snake_case              =      len(lowerCamelCase_  )
			_snake_case              =      {b: str(lowerCamelCase_  ) for b, i in zip(lowerCamelCase_							,	range(lowerCamelCase_  )  )}

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

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

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

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

    # BERT text encoder
			_snake_case              =      []
			_snake_case              =      "tf_bert_model/bert"
			_snake_case              =      "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:
									_snake_case              =      item[1]
			return key_mapping






def    UpperCAmelCase__(						lowerCamelCase_							,	lowerCamelCase_							,	lowerCamelCase_  ):



			"""simple docstring"""





			list(hf_params.keys()  )

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

						_snake_case              =      key_mapping[key]
						if "_conv" in key and "kernel" in key:
									_snake_case              =      torch.from_numpy(lowerCamelCase_  ).permute(3							,	2							,	0							,	1  )
						elif "embeddings" in key:
									_snake_case              =      torch.from_numpy(lowerCamelCase_  )
						elif "depthwise_kernel" in key:
									_snake_case              =      torch.from_numpy(lowerCamelCase_  ).permute(2							,	3							,	0							,	1  )
						elif "kernel" in key:
									_snake_case              =      torch.from_numpy(np.transpose(lowerCamelCase_  )  )
						elif "temperature" in key:
									_snake_case              =      value
						elif "bn/gamma" or "bn/beta" in key:
									_snake_case              =      torch.from_numpy(np.transpose(lowerCamelCase_  )  ).squeeze()
						else:
									_snake_case              =      torch.from_numpy(lowerCamelCase_  )

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






@torch.no_grad()
def    UpperCAmelCase__(						lowerCamelCase_							,	lowerCamelCase_							,	lowerCamelCase_							,	lowerCamelCase_  ):



			"""simple docstring"""





			_snake_case              =      64
			_snake_case              =      Tokenizer(lowerCamelCase_  )
			_snake_case              =      align.Align("efficientnet-b7"							,	"bert-base"							,	640							,	lowerCamelCase_							,	tok.get_vocab_size()  )
			original_model.compile()
			original_model.load_weights(lowerCamelCase_  )

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

			# Load HuggingFace model
			_snake_case              =      get_align_config()
			_snake_case              =      AlignModel(lowerCamelCase_  ).eval()
			_snake_case              =      hf_model.state_dict()

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

			# Initialize processor
			_snake_case              =      get_processor()
			_snake_case              =      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():
						_snake_case              =      hf_model(**lowerCamelCase_  )

			_snake_case              =      outputs.image_embeds.detach().numpy()
			_snake_case              =      outputs.text_embeds.detach().numpy()

			# Original model inference
			_snake_case              =      False
			_snake_case              =      EfficientNetImageProcessor(
			    do_center_crop=lowerCamelCase_							,	do_rescale=lowerCamelCase_							,	do_normalize=lowerCamelCase_							,	include_top=lowerCamelCase_							,	resample=Image.BILINEAR							,	)
			_snake_case              =      tf_image_processor(images=prepare_img()							,	return_tensors="tf"							,	data_format="channels_last"  )["pixel_values"]
			_snake_case              =      tok(tf.constant(["A picture of a cat"]  )  )

			_snake_case              =      original_model.image_encoder(lowerCamelCase_							,	training=lowerCamelCase_  )
			_snake_case              =      original_model.text_encoder(lowerCamelCase_							,	training=lowerCamelCase_  )

			_snake_case              =      tf.nn.la_normalize(lowerCamelCase_							,	axis=-1  )
			_snake_case              =      tf.nn.la_normalize(lowerCamelCase_							,	axis=-1  )

			# Check whether original and HF model outputs match  -> np.allclose
			if not np.allclose(lowerCamelCase_							,	lowerCamelCase_							,	atol=1E-3  ):
						raise ValueError("The predicted image features are not the same."  )
			if not np.allclose(lowerCamelCase_							,	lowerCamelCase_							,	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(lowerCamelCase_  ):
									os.mkdir(lowerCamelCase_  )
						# Save converted model and image processor
						hf_model.save_pretrained(lowerCamelCase_  )
						processor.save_pretrained(lowerCamelCase_  )

			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__":
			_SCREAMING_SNAKE_CASE:						int    	=       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''')

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


import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp






def    UpperCAmelCase__(						lowerCamelCase_							,	lowerCamelCase_							,	lowerCamelCase_							,	lowerCamelCase_							,	lowerCamelCase_ = 4096  ):



			"""simple docstring"""





			_snake_case       , _snake_case       , _snake_case              =      key.shape[-3:]
			_snake_case              =      value.shape[-1]
			_snake_case              =      min(lowerCamelCase_							,	lowerCamelCase_  )
			_snake_case              =      query / jnp.sqrt(lowerCamelCase_  )

			@functools.partial(jax.checkpoint							,	prevent_cse=lowerCamelCase_  )
			def summarize_chunk(lowerCamelCase_							,	lowerCamelCase_							,	lowerCamelCase_  ):
						_snake_case              =      jnp.einsum("...qhd,...khd->...qhk"							,	lowerCamelCase_							,	lowerCamelCase_							,	precision=lowerCamelCase_  )

						_snake_case              =      jnp.max(lowerCamelCase_							,	axis=-1							,	keepdims=lowerCamelCase_  )
						_snake_case              =      jax.lax.stop_gradient(lowerCamelCase_  )
						_snake_case              =      jnp.exp(attn_weights - max_score  )

						_snake_case              =      jnp.einsum("...vhf,...qhv->...qhf"							,	lowerCamelCase_							,	lowerCamelCase_							,	precision=lowerCamelCase_  )
						_snake_case              =      jnp.einsum("...qhk->...qh"							,	lowerCamelCase_  )

						return (exp_values, exp_weights.sum(axis=-1  ), max_score)

			def chunk_scanner(lowerCamelCase_  ):
						# julienne key array
						_snake_case              =      jax.lax.dynamic_slice(
						    operand=lowerCamelCase_							,	start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0]							,	slice_sizes=list(key.shape[:-3]  ) + [key_chunk_size, num_heads, k_features]							,	)

						# julienne value array
						_snake_case              =      jax.lax.dynamic_slice(
						    operand=lowerCamelCase_							,	start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0]							,	slice_sizes=list(value.shape[:-3]  ) + [key_chunk_size, num_heads, v_features]							,	)

						return summarize_chunk(lowerCamelCase_							,	lowerCamelCase_							,	lowerCamelCase_  )

			_snake_case       , _snake_case       , _snake_case              =      jax.lax.map(f=lowerCamelCase_							,	xs=jnp.arange(0							,	lowerCamelCase_							,	lowerCamelCase_  )  )

			_snake_case              =      jnp.max(lowerCamelCase_							,	axis=0							,	keepdims=lowerCamelCase_  )
			_snake_case              =      jnp.exp(chunk_max - global_max  )

			chunk_values *= jnp.expand_dims(lowerCamelCase_							,	axis=-1  )
			chunk_weights *= max_diffs

			_snake_case              =      chunk_values.sum(axis=0  )
			_snake_case              =      jnp.expand_dims(lowerCamelCase_							,	-1  ).sum(axis=0  )

			return all_values / all_weights






def    UpperCAmelCase__(						lowerCamelCase_							,	lowerCamelCase_							,	lowerCamelCase_							,	lowerCamelCase_=jax.lax.Precision.HIGHEST							,	lowerCamelCase_ = 1024							,	lowerCamelCase_ = 4096  ):



			"""simple docstring"""





			_snake_case       , _snake_case       , _snake_case              =      query.shape[-3:]

			def chunk_scanner(lowerCamelCase_							,	lowerCamelCase_  ):
						# julienne query array
						_snake_case              =      jax.lax.dynamic_slice(
						    operand=lowerCamelCase_							,	start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]							,	slice_sizes=list(query.shape[:-3]  ) + [min(lowerCamelCase_							,	lowerCamelCase_  ), num_heads, q_features]							,	)

						return (
						    chunk_idx + query_chunk_size,  # unused ignore it
						    _query_chunk_attention(
						        query=lowerCamelCase_							,	key=lowerCamelCase_							,	value=lowerCamelCase_							,	precision=lowerCamelCase_							,	key_chunk_size=lowerCamelCase_  ),
						)

			_snake_case       , _snake_case              =      jax.lax.scan(
			    f=lowerCamelCase_							,	init=0							,	xs=lowerCamelCase_							,	length=math.ceil(num_q / query_chunk_size  )  )

			return jnp.concatenate(lowerCamelCase_							,	axis=-3  )  # fuse the chunked result back

class 	_SCREAMING_SNAKE_CASE							(     nn.Module							):






		"""simple docstring"""







		A_						 =      42
		A_						 =      8
		A_						 =      64
		A_						 =      0.0
		A_						 =      False
		A_						 =      jnp.floataa



		def 				_a				(			self    :				Union[str, Any]   ):






					"""simple docstring"""
					_snake_case              =      self.dim_head * self.heads
					_snake_case              =      self.dim_head**-0.5

					# Weights were exported with old names {to_q, to_k, to_v, to_out}
					_snake_case              =      nn.Dense(__SCREAMING_SNAKE_CASE  ,					use_bias=__SCREAMING_SNAKE_CASE  ,					dtype=self.dtype  ,					name="to_q"   )
					_snake_case              =      nn.Dense(__SCREAMING_SNAKE_CASE  ,					use_bias=__SCREAMING_SNAKE_CASE  ,					dtype=self.dtype  ,					name="to_k"   )
					_snake_case              =      nn.Dense(__SCREAMING_SNAKE_CASE  ,					use_bias=__SCREAMING_SNAKE_CASE  ,					dtype=self.dtype  ,					name="to_v"   )

					_snake_case              =      nn.Dense(self.query_dim  ,					dtype=self.dtype  ,					name="to_out_0"   )
					_snake_case              =      nn.Dropout(rate=self.dropout   )



		def 				_a				(			self    :				Optional[int]  ,					__SCREAMING_SNAKE_CASE    :				Any   ):






					"""simple docstring"""
					_snake_case       , _snake_case       , _snake_case              =      tensor.shape
					_snake_case              =      self.heads
					_snake_case              =      tensor.reshape(__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					dim // head_size   )
					_snake_case              =      jnp.transpose(__SCREAMING_SNAKE_CASE  ,					(0, 2, 1, 3)   )
					_snake_case              =      tensor.reshape(batch_size * head_size  ,					__SCREAMING_SNAKE_CASE  ,					dim // head_size   )
					return tensor



		def 				_a				(			self    :				str  ,					__SCREAMING_SNAKE_CASE    :				Tuple   ):






					"""simple docstring"""
					_snake_case       , _snake_case       , _snake_case              =      tensor.shape
					_snake_case              =      self.heads
					_snake_case              =      tensor.reshape(batch_size // head_size  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE   )
					_snake_case              =      jnp.transpose(__SCREAMING_SNAKE_CASE  ,					(0, 2, 1, 3)   )
					_snake_case              =      tensor.reshape(batch_size // head_size  ,					__SCREAMING_SNAKE_CASE  ,					dim * head_size   )
					return tensor







		def __call__(			self    :				int  ,					__SCREAMING_SNAKE_CASE    :				int  ,					__SCREAMING_SNAKE_CASE    :				Tuple=None  ,					__SCREAMING_SNAKE_CASE    :				List[Any]=True   ):






					"""simple docstring"""
					_snake_case              =      hidden_states if context is None else context

					_snake_case              =      self.query(__SCREAMING_SNAKE_CASE   )
					_snake_case              =      self.key(__SCREAMING_SNAKE_CASE   )
					_snake_case              =      self.value(__SCREAMING_SNAKE_CASE   )

					_snake_case              =      self.reshape_heads_to_batch_dim(__SCREAMING_SNAKE_CASE   )
					_snake_case              =      self.reshape_heads_to_batch_dim(__SCREAMING_SNAKE_CASE   )
					_snake_case              =      self.reshape_heads_to_batch_dim(__SCREAMING_SNAKE_CASE   )

					if self.use_memory_efficient_attention:
								_snake_case              =      query_states.transpose(1  ,					0  ,					2   )
								_snake_case              =      key_states.transpose(1  ,					0  ,					2   )
								_snake_case              =      value_states.transpose(1  ,					0  ,					2   )

								# this if statement create a chunk size for each layer of the unet
								# the chunk size is equal to the query_length dimension of the deepest layer of the unet

								_snake_case              =      query_states.shape[-3]
								if flatten_latent_dim % 64 == 0:
											_snake_case              =      int(flatten_latent_dim / 64   )
								elif flatten_latent_dim % 16 == 0:
											_snake_case              =      int(flatten_latent_dim / 16   )
								elif flatten_latent_dim % 4 == 0:
											_snake_case              =      int(flatten_latent_dim / 4   )
								else:
											_snake_case              =      int(__SCREAMING_SNAKE_CASE   )

								_snake_case              =      jax_memory_efficient_attention(
								    __SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					query_chunk_size=__SCREAMING_SNAKE_CASE  ,					key_chunk_size=4_096 * 4   )

								_snake_case              =      hidden_states.transpose(1  ,					0  ,					2   )
					else:
								# compute attentions
								_snake_case              =      jnp.einsum("b i d, b j d->b i j"  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE   )
								_snake_case              =      attention_scores * self.scale
								_snake_case              =      nn.softmax(__SCREAMING_SNAKE_CASE  ,					axis=2   )

								# attend to values
								_snake_case              =      jnp.einsum("b i j, b j d -> b i d"  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE   )

					_snake_case              =      self.reshape_batch_dim_to_heads(__SCREAMING_SNAKE_CASE   )
					_snake_case              =      self.proj_attn(__SCREAMING_SNAKE_CASE   )
					return self.dropout_layer(__SCREAMING_SNAKE_CASE  ,					deterministic=__SCREAMING_SNAKE_CASE   )

class 	_SCREAMING_SNAKE_CASE							(     nn.Module							):






		"""simple docstring"""







		A_						 =      42
		A_						 =      42
		A_						 =      42
		A_						 =      0.0
		A_						 =      False
		A_						 =      jnp.floataa
		A_						 =      False



		def 				_a				(			self    :				Any   ):






					"""simple docstring"""
					# self attention (or cross_attention if only_cross_attention is True)
					_snake_case              =      FlaxAttention(
					    self.dim  ,					self.n_heads  ,					self.d_head  ,					self.dropout  ,					self.use_memory_efficient_attention  ,					dtype=self.dtype   )
					# cross attention
					_snake_case              =      FlaxAttention(
					    self.dim  ,					self.n_heads  ,					self.d_head  ,					self.dropout  ,					self.use_memory_efficient_attention  ,					dtype=self.dtype   )
					_snake_case              =      FlaxFeedForward(dim=self.dim  ,					dropout=self.dropout  ,					dtype=self.dtype   )
					_snake_case              =      nn.LayerNorm(epsilon=1e-5  ,					dtype=self.dtype   )
					_snake_case              =      nn.LayerNorm(epsilon=1e-5  ,					dtype=self.dtype   )
					_snake_case              =      nn.LayerNorm(epsilon=1e-5  ,					dtype=self.dtype   )
					_snake_case              =      nn.Dropout(rate=self.dropout   )







		def __call__(			self    :				List[Any]  ,					__SCREAMING_SNAKE_CASE    :				Any  ,					__SCREAMING_SNAKE_CASE    :				List[str]  ,					__SCREAMING_SNAKE_CASE    :				Any=True   ):






					"""simple docstring"""
					# self attention
					_snake_case              =      hidden_states
					if self.only_cross_attention:
								_snake_case              =      self.attna(self.norma(__SCREAMING_SNAKE_CASE   )  ,					__SCREAMING_SNAKE_CASE  ,					deterministic=__SCREAMING_SNAKE_CASE   )
					else:
								_snake_case              =      self.attna(self.norma(__SCREAMING_SNAKE_CASE   )  ,					deterministic=__SCREAMING_SNAKE_CASE   )
					_snake_case              =      hidden_states + residual

					# cross attention
					_snake_case              =      hidden_states
					_snake_case              =      self.attna(self.norma(__SCREAMING_SNAKE_CASE   )  ,					__SCREAMING_SNAKE_CASE  ,					deterministic=__SCREAMING_SNAKE_CASE   )
					_snake_case              =      hidden_states + residual

					# feed forward
					_snake_case              =      hidden_states
					_snake_case              =      self.ff(self.norma(__SCREAMING_SNAKE_CASE   )  ,					deterministic=__SCREAMING_SNAKE_CASE   )
					_snake_case              =      hidden_states + residual

					return self.dropout_layer(__SCREAMING_SNAKE_CASE  ,					deterministic=__SCREAMING_SNAKE_CASE   )

class 	_SCREAMING_SNAKE_CASE							(     nn.Module							):






		"""simple docstring"""







		A_						 =      42
		A_						 =      42
		A_						 =      42
		A_						 =      1
		A_						 =      0.0
		A_						 =      False
		A_						 =      False
		A_						 =      jnp.floataa
		A_						 =      False



		def 				_a				(			self    :				Union[str, Any]   ):






					"""simple docstring"""
					_snake_case              =      nn.GroupNorm(num_groups=32  ,					epsilon=1e-5   )

					_snake_case              =      self.n_heads * self.d_head
					if self.use_linear_projection:
								_snake_case              =      nn.Dense(__SCREAMING_SNAKE_CASE  ,					dtype=self.dtype   )
					else:
								_snake_case              =      nn.Conv(
								    __SCREAMING_SNAKE_CASE  ,					kernel_size=(1, 1)  ,					strides=(1, 1)  ,					padding="VALID"  ,					dtype=self.dtype  ,					)

					_snake_case              =      [
					    FlaxBasicTransformerBlock(
					        __SCREAMING_SNAKE_CASE  ,					self.n_heads  ,					self.d_head  ,					dropout=self.dropout  ,					only_cross_attention=self.only_cross_attention  ,					dtype=self.dtype  ,					use_memory_efficient_attention=self.use_memory_efficient_attention  ,					)
					    for _ in range(self.depth   )
					]

					if self.use_linear_projection:
								_snake_case              =      nn.Dense(__SCREAMING_SNAKE_CASE  ,					dtype=self.dtype   )
					else:
								_snake_case              =      nn.Conv(
								    __SCREAMING_SNAKE_CASE  ,					kernel_size=(1, 1)  ,					strides=(1, 1)  ,					padding="VALID"  ,					dtype=self.dtype  ,					)

					_snake_case              =      nn.Dropout(rate=self.dropout   )







		def __call__(			self    :				Optional[int]  ,					__SCREAMING_SNAKE_CASE    :				List[str]  ,					__SCREAMING_SNAKE_CASE    :				Union[str, Any]  ,					__SCREAMING_SNAKE_CASE    :				Union[str, Any]=True   ):






					"""simple docstring"""
					_snake_case       , _snake_case       , _snake_case       , _snake_case              =      hidden_states.shape
					_snake_case              =      hidden_states
					_snake_case              =      self.norm(__SCREAMING_SNAKE_CASE   )
					if self.use_linear_projection:
								_snake_case              =      hidden_states.reshape(__SCREAMING_SNAKE_CASE  ,					height * width  ,					__SCREAMING_SNAKE_CASE   )
								_snake_case              =      self.proj_in(__SCREAMING_SNAKE_CASE   )
					else:
								_snake_case              =      self.proj_in(__SCREAMING_SNAKE_CASE   )
								_snake_case              =      hidden_states.reshape(__SCREAMING_SNAKE_CASE  ,					height * width  ,					__SCREAMING_SNAKE_CASE   )

					for transformer_block in self.transformer_blocks:
								_snake_case              =      transformer_block(__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					deterministic=__SCREAMING_SNAKE_CASE   )

					if self.use_linear_projection:
								_snake_case              =      self.proj_out(__SCREAMING_SNAKE_CASE   )
								_snake_case              =      hidden_states.reshape(__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE   )
					else:
								_snake_case              =      hidden_states.reshape(__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE   )
								_snake_case              =      self.proj_out(__SCREAMING_SNAKE_CASE   )

					_snake_case              =      hidden_states + residual
					return self.dropout_layer(__SCREAMING_SNAKE_CASE  ,					deterministic=__SCREAMING_SNAKE_CASE   )

class 	_SCREAMING_SNAKE_CASE							(     nn.Module							):






		"""simple docstring"""







		A_						 =      42
		A_						 =      0.0
		A_						 =      jnp.floataa



		def 				_a				(			self    :				Optional[int]   ):






					"""simple docstring"""
					# The second linear layer needs to be called
					# net_2 for now to match the index of the Sequential layer
					_snake_case              =      FlaxGEGLU(self.dim  ,					self.dropout  ,					self.dtype   )
					_snake_case              =      nn.Dense(self.dim  ,					dtype=self.dtype   )







		def __call__(			self    :				Tuple  ,					__SCREAMING_SNAKE_CASE    :				List[str]  ,					__SCREAMING_SNAKE_CASE    :				List[Any]=True   ):






					"""simple docstring"""
					_snake_case              =      self.net_a(__SCREAMING_SNAKE_CASE  ,					deterministic=__SCREAMING_SNAKE_CASE   )
					_snake_case              =      self.net_a(__SCREAMING_SNAKE_CASE   )
					return hidden_states

class 	_SCREAMING_SNAKE_CASE							(     nn.Module							):






		"""simple docstring"""







		A_						 =      42
		A_						 =      0.0
		A_						 =      jnp.floataa



		def 				_a				(			self    :				Optional[Any]   ):






					"""simple docstring"""
					_snake_case              =      self.dim * 4
					_snake_case              =      nn.Dense(inner_dim * 2  ,					dtype=self.dtype   )
					_snake_case              =      nn.Dropout(rate=self.dropout   )







		def __call__(			self    :				Optional[int]  ,					__SCREAMING_SNAKE_CASE    :				Union[str, Any]  ,					__SCREAMING_SNAKE_CASE    :				List[str]=True   ):






					"""simple docstring"""
					_snake_case              =      self.proj(__SCREAMING_SNAKE_CASE   )
					_snake_case       , _snake_case              =      jnp.split(__SCREAMING_SNAKE_CASE  ,					2  ,					axis=2   )
					return self.dropout_layer(hidden_linear * nn.gelu(__SCREAMING_SNAKE_CASE   )  ,					deterministic=__SCREAMING_SNAKE_CASE   )

