






"""simple docstring"""







import argparse
import json
import os

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ConvNextImageProcessor, ConvNextVaConfig, ConvNextVaForImageClassification
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
a :				Union[str, Any]      					=		logging.get_logger(__name__)





def 						lowercase__(A						)				->Union[str, Any]:




						"""simple docstring"""



						lowercase__				:      Tuple=       ConvNextVaConfig()

						if "atto" in checkpoint_url:
												lowercase__				:      Optional[int]=       [2, 2, 6, 2]
												lowercase__				:      Union[str, Any]=       [40, 80, 160, 320]
						if "femto" in checkpoint_url:
												lowercase__				:      List[str]=       [2, 2, 6, 2]
												lowercase__				:      str=       [48, 96, 192, 384]
						if "pico" in checkpoint_url:
												lowercase__				:      List[Any]=       [2, 2, 6, 2]
												lowercase__				:      Optional[int]=       [64, 128, 256, 512]
						if "nano" in checkpoint_url:
												lowercase__				:      List[str]=       [2, 2, 8, 2]
												lowercase__				:      Union[str, Any]=       [80, 160, 320, 640]
						if "tiny" in checkpoint_url:
												lowercase__				:      Tuple=       [3, 3, 9, 3]
												lowercase__				:      Tuple=       [96, 192, 384, 768]
						if "base" in checkpoint_url:
												lowercase__				:      Union[str, Any]=       [3, 3, 27, 3]
												lowercase__				:      Dict=       [128, 256, 512, 1_024]
						if "large" in checkpoint_url:
												lowercase__				:      List[Any]=       [3, 3, 27, 3]
												lowercase__				:      int=       [192, 384, 768, 1_536]
						if "huge" in checkpoint_url:
												lowercase__				:      Optional[Any]=       [3, 3, 27, 3]
												lowercase__				:      Any=       [352, 704, 1_408, 2_816]

						lowercase__				:      Optional[Any]=       1_000
						lowercase__				:      str=       "imagenet-1k-id2label.json"
						lowercase__				:      Optional[int]=       (1, 1_000)

						lowercase__				:      int=       "huggingface/label-files"
						lowercase__				:      Optional[int]=       num_labels
						lowercase__				:      Any=       json.load(open(hf_hub_download(A					,		A					,		repo_type="dataset"						)					,		"r"						)						)
						lowercase__				:      str=       {int(A						): v for k, v in idalabel.items()}

						lowercase__				:      Union[str, Any]=       idalabel
						lowercase__				:      Union[str, Any]=       {v: k for k, v in idalabel.items()}
						lowercase__				:      Optional[int]=       hidden_sizes
						lowercase__				:      Dict=       depths

						return config, expected_shape





def 						lowercase__(A						)				->Optional[int]:




						"""simple docstring"""



						if "downsample_layers.0.0" in name:
												lowercase__				:      Tuple=       name.replace("downsample_layers.0.0"					,		"embeddings.patch_embeddings"						)
						if "downsample_layers.0.1" in name:
												lowercase__				:      Dict=       name.replace("downsample_layers.0.1"					,		"embeddings.norm"						)  # we rename to layernorm later on
						if "downsample_layers.1.0" in name:
												lowercase__				:      Dict=       name.replace("downsample_layers.1.0"					,		"stages.1.downsampling_layer.0"						)
						if "downsample_layers.1.1" in name:
												lowercase__				:      Union[str, Any]=       name.replace("downsample_layers.1.1"					,		"stages.1.downsampling_layer.1"						)
						if "downsample_layers.2.0" in name:
												lowercase__				:      int=       name.replace("downsample_layers.2.0"					,		"stages.2.downsampling_layer.0"						)
						if "downsample_layers.2.1" in name:
												lowercase__				:      Dict=       name.replace("downsample_layers.2.1"					,		"stages.2.downsampling_layer.1"						)
						if "downsample_layers.3.0" in name:
												lowercase__				:      int=       name.replace("downsample_layers.3.0"					,		"stages.3.downsampling_layer.0"						)
						if "downsample_layers.3.1" in name:
												lowercase__				:      Tuple=       name.replace("downsample_layers.3.1"					,		"stages.3.downsampling_layer.1"						)
						if "stages" in name and "downsampling_layer" not in name:
												# stages.0.0. for instance should be renamed to stages.0.layers.0.
												lowercase__				:      Union[str, Any]=       name[: len("stages.0"						)] + ".layers" + name[len("stages.0"						) :]
						if "gamma" in name:
												lowercase__				:      List[Any]=       name.replace("gamma"					,		"weight"						)
						if "beta" in name:
												lowercase__				:      Tuple=       name.replace("beta"					,		"bias"						)
						if "stages" in name:
												lowercase__				:      Optional[Any]=       name.replace("stages"					,		"encoder.stages"						)
						if "norm" in name:
												lowercase__				:      List[Any]=       name.replace("norm"					,		"layernorm"						)
						if "head" in name:
												lowercase__				:      int=       name.replace("head"					,		"classifier"						)

						return name





def 						lowercase__()				->Dict:




						"""simple docstring"""



						lowercase__				:      Optional[Any]=       "http://images.cocodataset.org/val2017/000000039769.jpg"
						lowercase__				:      Optional[int]=       Image.open(requests.get(A					,		stream=A						).raw						)
						return im





def 						lowercase__(A						)				->str:




						"""simple docstring"""



						if "224" in checkpoint_url:
												lowercase__				:      Any=       224
												lowercase__				:      Optional[int]=       224 / 256
						elif "384" in checkpoint_url:
												lowercase__				:      Optional[int]=       384
												lowercase__				:      Dict=       None
						else:
												lowercase__				:      Optional[Any]=       512
												lowercase__				:      Dict=       None

						return ConvNextImageProcessor(
						    size=A					,		crop_pct=A					,		image_mean=[0.485, 0.456, 0.406]					,		image_std=[0.229, 0.224, 0.225]					,		resample=PILImageResampling.BICUBIC					,		)





@torch.no_grad()
def 						lowercase__(A					,		A					,		A					,		A						)				->Union[str, Any]:




						"""simple docstring"""



						print("Downloading original model from checkpoint..."						)
						# define ConvNeXTV2 configuration based on URL
						lowercase__,	lowercase__				:      Dict=       get_convnextva_config(A						)
						# load original state_dict from URL
						lowercase__				:      Tuple=       torch.hub.load_state_dict_from_url(A						)["model"]

						print("Converting model parameters..."						)
						# rename keys
						for key in state_dict.copy().keys():
												lowercase__				:      Optional[int]=       state_dict.pop(A						)
												lowercase__				:      List[str]=       val
						# add prefix to all keys expect classifier head
						for key in state_dict.copy().keys():
												lowercase__				:      Optional[Any]=       state_dict.pop(A						)
												if not key.startswith("classifier"						):
																		lowercase__				:      List[str]=       "convnextv2." + key
												lowercase__				:      Optional[Any]=       val

						# load HuggingFace model
						lowercase__				:      Union[str, Any]=       ConvNextVaForImageClassification(A						)
						model.load_state_dict(A						)
						model.eval()

						# Check outputs on an image, prepared by ConvNextImageProcessor
						lowercase__				:      str=       convert_preprocessor(A						)
						lowercase__				:      Optional[int]=       preprocessor(images=prepare_img()					,		return_tensors="pt"						)
						lowercase__				:      Optional[int]=       model(**A						).logits

						# note: the logits below were obtained without center cropping
						if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt":
												lowercase__				:      Union[str, Any]=       torch.tensor([-0.3_930, 0.1_747, -0.5_246, 0.4_177, 0.4_295]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
												lowercase__				:      Optional[int]=       torch.tensor([-0.1_727, -0.5_341, -0.7_818, -0.4_745, -0.6_566]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
												lowercase__				:      str=       torch.tensor([-0.0_333, 0.1_563, -0.9_137, 0.1_054, 0.0_381]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
												lowercase__				:      Optional[int]=       torch.tensor([-0.1_744, -0.1_555, -0.0_713, 0.0_950, -0.1_431]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
												lowercase__				:      Optional[int]=       torch.tensor([0.9_996, 0.1_966, -0.4_386, -0.3_472, 0.6_661]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
												lowercase__				:      List[Any]=       torch.tensor([-0.2_553, -0.6_708, -0.1_359, 0.2_518, -0.2_488]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
												lowercase__				:      Tuple=       torch.tensor([-0.0_673, -0.5_627, -0.3_753, -0.2_722, 0.0_178]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
												lowercase__				:      Union[str, Any]=       torch.tensor([-0.6_377, -0.7_458, -0.2_150, 0.1_184, -0.0_597]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
												lowercase__				:      Tuple=       torch.tensor([1.0_799, 0.2_322, -0.8_860, 1.0_219, 0.6_231]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
												lowercase__				:      Optional[int]=       torch.tensor([0.3_766, 0.4_917, -1.1_426, 0.9_942, 0.6_024]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
												lowercase__				:      Tuple=       torch.tensor([0.4_220, -0.6_919, -0.4_317, -0.2_881, -0.6_609]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
												lowercase__				:      Tuple=       torch.tensor([0.1_082, -0.8_286, -0.5_095, 0.4_681, -0.8_085]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
												lowercase__				:      str=       torch.tensor([-0.2_419, -0.6_221, 0.2_176, -0.0_980, -0.7_527]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
												lowercase__				:      Union[str, Any]=       torch.tensor([0.0_391, -0.4_371, 0.3_786, 0.1_251, -0.2_784]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
												lowercase__				:      Optional[int]=       torch.tensor([-0.0_504, 0.5_636, -0.1_729, -0.6_507, -0.3_949]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
												lowercase__				:      int=       torch.tensor([0.3_560, 0.9_486, 0.3_149, -0.2_667, -0.5_138]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
												lowercase__				:      str=       torch.tensor([-0.2_469, -0.4_550, -0.5_853, -0.0_810, 0.0_309]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
												lowercase__				:      Optional[Any]=       torch.tensor([-0.3_090, 0.0_802, -0.0_682, -0.1_979, -0.2_826]						)
						else:
												raise ValueError(f'''Unknown URL: {checkpoint_url}'''						)

						assert torch.allclose(logits[0, :5]					,		A					,		atol=1e-3						)
						assert logits.shape == expected_shape
						print("Model outputs match the original results!"						)

						if save_model:
												print("Saving model to local..."						)
												# Create folder to save model
												if not os.path.isdir(A						):
																		os.mkdir(A						)

												model.save_pretrained(A						)
												preprocessor.save_pretrained(A						)

						lowercase__				:      int=       "convnextv2"
						if "atto" in checkpoint_url:
												model_name += "-atto"
						if "femto" in checkpoint_url:
												model_name += "-femto"
						if "pico" in checkpoint_url:
												model_name += "-pico"
						if "nano" in checkpoint_url:
												model_name += "-nano"
						elif "tiny" in checkpoint_url:
												model_name += "-tiny"
						elif "base" in checkpoint_url:
												model_name += "-base"
						elif "large" in checkpoint_url:
												model_name += "-large"
						elif "huge" in checkpoint_url:
												model_name += "-huge"
						if "22k" in checkpoint_url and "1k" not in checkpoint_url:
												model_name += "-22k"
						elif "22k" in checkpoint_url and "1k" in checkpoint_url:
												model_name += "-22k-1k"
						elif "1k" in checkpoint_url:
												model_name += "-1k"
						if "224" in checkpoint_url:
												model_name += "-224"
						elif "384" in checkpoint_url:
												model_name += "-384"
						elif "512" in checkpoint_url:
												model_name += "-512"

						if push_to_hub:
												print(f'''Pushing {model_name} to the hub...'''						)
												model.push_to_hub(A						)
												preprocessor.push_to_hub(A						)


if __name__ == "__main__":
			a :				str      					=		argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    """--checkpoint_url""",
			    default="""https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt""",
			    type=str,
			    help="""URL of the original ConvNeXTV2 checkpoint you'd like to convert.""",
			)
			parser.add_argument(
			    """--pytorch_dump_folder_path""",
			    default="""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 preprocessor to the hub""")

			a :				Any      					=		parser.parse_args()
			convert_convnextva_checkpoint(
			    args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
			)








'''simple docstring'''

from typing import Optional, Tuple

import flax.linen as nn
import jax
import jax.numpy as jnp
from configuration_hybrid_clip import HybridCLIPConfig
from flax.core.frozen_dict import FrozenDict

from transformers import FLAX_MODEL_MAPPING, FlaxCLIPVisionModel
from transformers.modeling_flax_utils import FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPOutput
from transformers.utils import logging


lowerCAmelCase_			:			List[str]			   =     logging.get_logger(__name__)







class 	__lowerCAmelCase    (  nn.Module				):
     snake_case   :    HybridCLIPConfig
     snake_case   :    jnp.dtype							      =      jnp.floataa



     def     snake_case_ (self	):
          _UpperCAmelCase					:  int         =    self.config.text_config
          _UpperCAmelCase					:  Dict         =    self.config.vision_config

          _UpperCAmelCase					:  str         =    self.config.projection_dim
          _UpperCAmelCase					:  int         =    text_config.hidden_size
          _UpperCAmelCase					:  int         =    vision_config.hidden_size

          _UpperCAmelCase					:  Optional[int]         =    FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
          _UpperCAmelCase					:  List[Any]         =    FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__						,	lowerCAmelCase__	).module_class

          _UpperCAmelCase					:  List[Any]         =    text_module(lowerCAmelCase__						,	dtype=self.dtype	)
          _UpperCAmelCase					:  Union[str, Any]         =    vision_module(lowerCAmelCase__						,	dtype=self.dtype	)

          _UpperCAmelCase					:  Optional[int]         =    nn.Dense(
              self.projection_dim						,	dtype=self.dtype						,	kernel_init=jax.nn.initializers.normal(0.0_2	)						,	use_bias=lowerCAmelCase__						,	)
          _UpperCAmelCase					:  List[str]         =    nn.Dense(
              self.projection_dim						,	dtype=self.dtype						,	kernel_init=jax.nn.initializers.normal(0.0_2	)						,	use_bias=lowerCAmelCase__						,	)
          _UpperCAmelCase					:  List[Any]         =    self.param("""logit_scale"""						,	jax.nn.initializers.ones						,	[]	)



     def __call__(self						,	lowerCAmelCase__=None						,	lowerCAmelCase__=None						,	lowerCAmelCase__=None						,	lowerCAmelCase__=None						,	lowerCAmelCase__=None						,	lowerCAmelCase__ = True						,	lowerCAmelCase__=None						,	lowerCAmelCase__=None						,	lowerCAmelCase__=None						,	):
          _UpperCAmelCase					:  int         =    return_dict if return_dict is not None else self.config.return_dict

          _UpperCAmelCase					:  Dict         =    self.vision_model(
              pixel_values=lowerCAmelCase__						,	deterministic=lowerCAmelCase__						,	output_attentions=lowerCAmelCase__						,	output_hidden_states=lowerCAmelCase__						,	return_dict=lowerCAmelCase__						,	)

          _UpperCAmelCase					:  int         =    self.text_model(
              input_ids=lowerCAmelCase__						,	attention_mask=lowerCAmelCase__						,	token_type_ids=lowerCAmelCase__						,	position_ids=lowerCAmelCase__						,	deterministic=lowerCAmelCase__						,	output_attentions=lowerCAmelCase__						,	output_hidden_states=lowerCAmelCase__						,	return_dict=lowerCAmelCase__						,	)

          _UpperCAmelCase					:  int         =    vision_outputs[1]
          _UpperCAmelCase					:  Dict         =    self.visual_projection(lowerCAmelCase__	)

          _UpperCAmelCase					:  List[str]         =    text_outputs[1]
          _UpperCAmelCase					:  List[Any]         =    self.text_projection(lowerCAmelCase__	)

          # normalized features
          _UpperCAmelCase					:  List[Any]         =    image_embeds / jnp.linalg.norm(lowerCAmelCase__						,	axis=-1						,	keepdims=lowerCAmelCase__	)
          _UpperCAmelCase					:  Optional[Any]         =    text_embeds / jnp.linalg.norm(lowerCAmelCase__						,	axis=-1						,	keepdims=lowerCAmelCase__	)

          # cosine similarity as logits
          _UpperCAmelCase					:  str         =    jnp.exp(self.logit_scale	)
          _UpperCAmelCase					:  str         =    jnp.matmul(lowerCAmelCase__						,	image_embeds.T	) * logit_scale
          _UpperCAmelCase					:  Any         =    logits_per_text.T

          if not return_dict:
               return (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)

          return FlaxCLIPOutput(
              logits_per_image=lowerCAmelCase__						,	logits_per_text=lowerCAmelCase__						,	text_embeds=lowerCAmelCase__						,	image_embeds=lowerCAmelCase__						,	text_model_output=lowerCAmelCase__						,	vision_model_output=lowerCAmelCase__						,	)







class 	__lowerCAmelCase    (  __a				):
     snake_case   :    List[Any]							      =      HybridCLIPConfig
     snake_case   :    int							      =      FlaxHybridCLIPModule



     def __init__(self						,	lowerCAmelCase__						,	lowerCAmelCase__ = None						,	lowerCAmelCase__ = 0						,	lowerCAmelCase__ = jnp.floataa						,	**lowerCAmelCase__						,	):
          if input_shape is None:
               _UpperCAmelCase					:  Optional[Any]         =    ((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))

          _UpperCAmelCase					:  Union[str, Any]         =    self.module_class(config=lowerCAmelCase__						,	dtype=lowerCAmelCase__						,	**lowerCAmelCase__	)
          super().__init__(lowerCAmelCase__						,	lowerCAmelCase__						,	input_shape=lowerCAmelCase__						,	seed=lowerCAmelCase__						,	dtype=lowerCAmelCase__	)



     def     snake_case_ (self						,	lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__ = None	):
          # init input tensor
          _UpperCAmelCase					:  List[Any]         =    jnp.zeros(input_shape[0]						,	dtype="""i4"""	)
          _UpperCAmelCase					:  str         =    jnp.broadcast_to(jnp.arange(jnp.atleast_ad(lowerCAmelCase__	).shape[-1]	)						,	input_shape[0]	)
          _UpperCAmelCase					:  List[Any]         =    jnp.ones_like(lowerCAmelCase__	)
          _UpperCAmelCase					:  List[str]         =    jnp.ones_like(lowerCAmelCase__	)

          _UpperCAmelCase					:  Any         =    jax.random.normal(lowerCAmelCase__						,	input_shape[1]	)

          _UpperCAmelCase  ,     _UpperCAmelCase					:  int         =    jax.random.split(lowerCAmelCase__	)
          _UpperCAmelCase					:  Any         =    {"""params""": params_rng, """dropout""": dropout_rng}

          return self.module.init(lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__	)["params"]



     def __call__(self						,	lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__=None						,	lowerCAmelCase__=None						,	lowerCAmelCase__=None						,	lowerCAmelCase__ = None						,	lowerCAmelCase__ = None						,	lowerCAmelCase__ = False						,	lowerCAmelCase__ = None						,	lowerCAmelCase__ = None						,	lowerCAmelCase__ = None						,	):
          _UpperCAmelCase					:  Any         =    output_attentions if output_attentions is not None else self.config.output_attentions
          _UpperCAmelCase					:  int         =    (
              output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
          )
          _UpperCAmelCase					:  int         =    return_dict if return_dict is not None else self.config.return_dict

          if position_ids is None:
               _UpperCAmelCase					:  List[Any]         =    jnp.broadcast_to(jnp.arange(jnp.atleast_ad(lowerCAmelCase__	).shape[-1]	)						,	input_ids.shape	)

          if token_type_ids is None:
               _UpperCAmelCase					:  List[str]         =    jnp.zeros_like(lowerCAmelCase__	)

          if attention_mask is None:
               _UpperCAmelCase					:  Any         =    jnp.ones_like(lowerCAmelCase__	)

          # Handle any PRNG if needed
          _UpperCAmelCase					:  List[str]         =    {}
          if dropout_rng is not None:
               _UpperCAmelCase					:  Union[str, Any]         =    dropout_rng

          return self.module.apply(
              {"""params""": params or self.params}						,	jnp.array(lowerCAmelCase__						,	dtype="""i4"""	)						,	jnp.array(lowerCAmelCase__						,	dtype=jnp.floataa	)						,	jnp.array(lowerCAmelCase__						,	dtype="""i4"""	)						,	jnp.array(lowerCAmelCase__						,	dtype="""i4"""	)						,	jnp.array(lowerCAmelCase__						,	dtype="""i4"""	)						,	not train						,	lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__						,	rngs=lowerCAmelCase__						,	)



     def     snake_case_ (self						,	lowerCAmelCase__						,	lowerCAmelCase__=None						,	lowerCAmelCase__=None						,	lowerCAmelCase__=None						,	lowerCAmelCase__ = None						,	lowerCAmelCase__ = None						,	lowerCAmelCase__=False						,	):
          if position_ids is None:
               _UpperCAmelCase					:  str         =    jnp.broadcast_to(jnp.arange(jnp.atleast_ad(lowerCAmelCase__	).shape[-1]	)						,	input_ids.shape	)

          if token_type_ids is None:
               _UpperCAmelCase					:  str         =    jnp.zeros_like(lowerCAmelCase__	)

          if attention_mask is None:
               _UpperCAmelCase					:  Optional[Any]         =    jnp.ones_like(lowerCAmelCase__	)

          # Handle any PRNG if needed
          _UpperCAmelCase					:  str         =    {}
          if dropout_rng is not None:
               _UpperCAmelCase					:  List[Any]         =    dropout_rng

          def _get_features(lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__	):
               _UpperCAmelCase					:  Union[str, Any]         =    module.text_model(
                   input_ids=lowerCAmelCase__						,	attention_mask=lowerCAmelCase__						,	position_ids=lowerCAmelCase__						,	token_type_ids=lowerCAmelCase__						,	deterministic=lowerCAmelCase__						,	)
               _UpperCAmelCase					:  Optional[int]         =    text_outputs[1]
               _UpperCAmelCase					:  List[Any]         =    module.text_projection(lowerCAmelCase__	)
               return text_features

          return self.module.apply(
              {"""params""": params or self.params}						,	jnp.array(lowerCAmelCase__						,	dtype="""i4"""	)						,	jnp.array(lowerCAmelCase__						,	dtype="""i4"""	)						,	jnp.array(lowerCAmelCase__						,	dtype="""i4"""	)						,	jnp.array(lowerCAmelCase__						,	dtype="""i4"""	)						,	not train						,	method=_get_features						,	rngs=lowerCAmelCase__						,	)



     def     snake_case_ (self						,	lowerCAmelCase__						,	lowerCAmelCase__ = None						,	lowerCAmelCase__ = None						,	lowerCAmelCase__=False	):
          _UpperCAmelCase					:  Any         =    {}
          if dropout_rng is not None:
               _UpperCAmelCase					:  int         =    dropout_rng

          def _get_features(lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__	):
               _UpperCAmelCase					:  Optional[Any]         =    module.vision_model(pixel_values=lowerCAmelCase__						,	deterministic=lowerCAmelCase__	)
               _UpperCAmelCase					:  int         =    vision_outputs[1]  # pooled_output
               _UpperCAmelCase					:  Optional[int]         =    module.visual_projection(lowerCAmelCase__	)
               return image_features

          return self.module.apply(
              {"""params""": params or self.params}						,	jnp.array(lowerCAmelCase__						,	dtype=jnp.floataa	)						,	not train						,	method=_get_features						,	rngs=lowerCAmelCase__						,	)



     @classmethod
     def     snake_case_ (cls						,	lowerCAmelCase__ = None						,	lowerCAmelCase__ = None						,	*lowerCAmelCase__						,	**lowerCAmelCase__						,	):
          _UpperCAmelCase					:  List[Any]         =    {
              argument[len("""text_"""	) :]: value for argument, value in kwargs.items() if argument.startswith("""text_"""	)
          }

          _UpperCAmelCase					:  Tuple         =    {
              argument[len("""vision_"""	) :]: value for argument, value in kwargs.items() if argument.startswith("""vision_"""	)
          }

          # remove text, vision kwargs from kwargs
          for key in kwargs_text.keys():
               del kwargs["text_" + key]
          for key in kwargs_vision.keys():
               del kwargs["vision_" + key]

          # Load and initialize the text and vision model
          _UpperCAmelCase					:  List[Any]         =    kwargs_text.pop("""model"""						,	lowerCAmelCase__	)
          if text_model is None:
               assert (
                   text_model_name_or_path is not None
               ), "If `model` is not defined as an argument, a `text_model_name_or_path` has to be defined"
               from transformers import FlaxAutoModel

               if "config" not in kwargs_text:
                    from transformers import AutoConfig

                    _UpperCAmelCase					:  Any         =    AutoConfig.from_pretrained(lowerCAmelCase__	)
                    _UpperCAmelCase					:  int         =    text_config

               _UpperCAmelCase					:  Any         =    FlaxAutoModel.from_pretrained(lowerCAmelCase__						,	*lowerCAmelCase__						,	**lowerCAmelCase__	)

          _UpperCAmelCase					:  List[Any]         =    kwargs_vision.pop("""model"""						,	lowerCAmelCase__	)
          if vision_model is None:
               assert (
                   vision_model_name_or_path is not None
               ), "If `model` is not defined as an argument, a `vision_model_name_or_path` has to be defined"
               from transformers import FlaxAutoModel

               if "config" not in kwargs_vision:
                    from transformers import AutoConfig

                    _UpperCAmelCase					:  Union[str, Any]         =    AutoConfig.from_pretrained(lowerCAmelCase__	)
                    _UpperCAmelCase					:  Any         =    vision_config

               _UpperCAmelCase					:  Tuple         =    FlaxAutoModel.from_pretrained(lowerCAmelCase__						,	*lowerCAmelCase__						,	**lowerCAmelCase__	)

          # instantiate config with corresponding kwargs
          _UpperCAmelCase					:  List[Any]         =    kwargs.pop("""dtype"""						,	jnp.floataa	)
          _UpperCAmelCase					:  List[Any]         =    HybridCLIPConfig.from_text_vision_configs(text_model.config						,	vision_model.config						,	**lowerCAmelCase__	)

          # init model
          _UpperCAmelCase					:  Optional[Any]         =    cls(lowerCAmelCase__						,	*lowerCAmelCase__						,	dtype=lowerCAmelCase__						,	**lowerCAmelCase__	)

          if vision_config.model_type == "clip":
               _UpperCAmelCase					:  str         =    vision_model.params["""vision_model"""]
               _UpperCAmelCase					:  int         =    vision_model.params["""visual_projection"""]["""kernel"""]
          else:
               _UpperCAmelCase					:  List[str]         =    vision_model.params

          _UpperCAmelCase					:  Any         =    text_model.params

          return model


