




'''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


UpperCAmelCase_    :					Optional[int]					      =						logging.get_logger(__name__)





class       lowerCAmelCase				(	nn.Module):
       __lowercase					: HybridCLIPConfig
       __lowercase					: jnp.dtype	    =     jnp.floataa


       def  lowerCAmelCase (				self					) ->					Optional[Any]:



            '''simple docstring'''




            __snake_case         =					self.config.text_config
            __snake_case         =					self.config.vision_config

            __snake_case         =					self.config.projection_dim
            __snake_case         =					text_config.hidden_size
            __snake_case         =					vision_config.hidden_size

            __snake_case         =					FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
            __snake_case         =					FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__      ,					__SCREAMING_SNAKE_CASE					).module_class

            __snake_case         =					text_module(__SCREAMING_SNAKE_CASE      ,					dtype=self.dtype					)
            __snake_case         =					vision_module(__SCREAMING_SNAKE_CASE      ,					dtype=self.dtype					)

            __snake_case         =					nn.Dense(
                self.projection_dim      ,					dtype=self.dtype      ,					kernel_init=jax.nn.initializers.normal(0.02					)      ,					use_bias=__SCREAMING_SNAKE_CASE      ,					)
            __snake_case         =					nn.Dense(
                self.projection_dim      ,					dtype=self.dtype      ,					kernel_init=jax.nn.initializers.normal(0.02					)      ,					use_bias=__SCREAMING_SNAKE_CASE      ,					)
            __snake_case         =					self.param('''logit_scale'''      ,					jax.nn.initializers.ones      ,					[]					)




       def __call__(				self      ,					__SCREAMING_SNAKE_CASE=None      ,					__SCREAMING_SNAKE_CASE=None      ,					__SCREAMING_SNAKE_CASE=None      ,					__SCREAMING_SNAKE_CASE=None      ,					__SCREAMING_SNAKE_CASE=None      ,					__SCREAMING_SNAKE_CASE = True      ,					__SCREAMING_SNAKE_CASE=None      ,					__SCREAMING_SNAKE_CASE=None      ,					__SCREAMING_SNAKE_CASE=None      ,					) ->					Optional[int]:



            '''simple docstring'''




            __snake_case         =					return_dict if return_dict is not None else self.config.return_dict

            __snake_case         =					self.vision_model(
                pixel_values=__SCREAMING_SNAKE_CASE      ,					deterministic=__SCREAMING_SNAKE_CASE      ,					output_attentions=__SCREAMING_SNAKE_CASE      ,					output_hidden_states=__SCREAMING_SNAKE_CASE      ,					return_dict=__SCREAMING_SNAKE_CASE      ,					)

            __snake_case         =					self.text_model(
                input_ids=__SCREAMING_SNAKE_CASE      ,					attention_mask=__SCREAMING_SNAKE_CASE      ,					token_type_ids=__SCREAMING_SNAKE_CASE      ,					position_ids=__SCREAMING_SNAKE_CASE      ,					deterministic=__SCREAMING_SNAKE_CASE      ,					output_attentions=__SCREAMING_SNAKE_CASE      ,					output_hidden_states=__SCREAMING_SNAKE_CASE      ,					return_dict=__SCREAMING_SNAKE_CASE      ,					)

            __snake_case         =					vision_outputs[1]
            __snake_case         =					self.visual_projection(__SCREAMING_SNAKE_CASE					)

            __snake_case         =					text_outputs[1]
            __snake_case         =					self.text_projection(__SCREAMING_SNAKE_CASE					)

            # normalized features
            __snake_case         =					image_embeds / jnp.linalg.norm(__SCREAMING_SNAKE_CASE      ,					axis=-1      ,					keepdims=__SCREAMING_SNAKE_CASE					)
            __snake_case         =					text_embeds / jnp.linalg.norm(__SCREAMING_SNAKE_CASE      ,					axis=-1      ,					keepdims=__SCREAMING_SNAKE_CASE					)

            # cosine similarity as logits
            __snake_case         =					jnp.exp(self.logit_scale					)
            __snake_case         =					jnp.matmul(__SCREAMING_SNAKE_CASE      ,					image_embeds.T					) * logit_scale
            __snake_case         =					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=__SCREAMING_SNAKE_CASE      ,					logits_per_text=__SCREAMING_SNAKE_CASE      ,					text_embeds=__SCREAMING_SNAKE_CASE      ,					image_embeds=__SCREAMING_SNAKE_CASE      ,					text_model_output=__SCREAMING_SNAKE_CASE      ,					vision_model_output=__SCREAMING_SNAKE_CASE      ,					)





class       lowerCAmelCase				(	__lowerCAmelCase):
       __lowercase					: Dict	    =     HybridCLIPConfig
       __lowercase					: str	    =     FlaxHybridCLIPModule


       def __init__(				self      ,					__SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE = None      ,					__SCREAMING_SNAKE_CASE = 0      ,					__SCREAMING_SNAKE_CASE = jnp.floataa      ,					**__SCREAMING_SNAKE_CASE      ,					) ->					Dict:



            '''simple docstring'''




            if input_shape is None:
                 __snake_case         =					((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))

            __snake_case         =					self.module_class(config=__SCREAMING_SNAKE_CASE      ,					dtype=__SCREAMING_SNAKE_CASE      ,					**__SCREAMING_SNAKE_CASE					)
            super().__init__(__SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE      ,					input_shape=__SCREAMING_SNAKE_CASE      ,					seed=__SCREAMING_SNAKE_CASE      ,					dtype=__SCREAMING_SNAKE_CASE					)


       def  lowerCAmelCase (				self      ,					__SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE = None					) ->					FrozenDict:



            '''simple docstring'''




            __snake_case         =					jnp.zeros(input_shape[0]      ,					dtype='''i4'''					)
            __snake_case         =					jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__SCREAMING_SNAKE_CASE					).shape[-1]					)      ,					input_shape[0]					)
            __snake_case         =					jnp.ones_like(__SCREAMING_SNAKE_CASE					)
            __snake_case         =					jnp.ones_like(__SCREAMING_SNAKE_CASE					)

            __snake_case         =					jax.random.normal(__SCREAMING_SNAKE_CASE      ,					input_shape[1]					)

            __snake_case   ,      __snake_case         =					jax.random.split(__SCREAMING_SNAKE_CASE					)
            __snake_case         =					{'''params''': params_rng, '''dropout''': dropout_rng}

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


       def __call__(				self      ,					__SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE=None      ,					__SCREAMING_SNAKE_CASE=None      ,					__SCREAMING_SNAKE_CASE=None      ,					__SCREAMING_SNAKE_CASE = None      ,					__SCREAMING_SNAKE_CASE = None      ,					__SCREAMING_SNAKE_CASE = False      ,					__SCREAMING_SNAKE_CASE = None      ,					__SCREAMING_SNAKE_CASE = None      ,					__SCREAMING_SNAKE_CASE = None      ,					) ->					Optional[int]:



            '''simple docstring'''




            __snake_case         =					output_attentions if output_attentions is not None else self.config.output_attentions
            __snake_case         =					(
                output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
            )
            __snake_case         =					return_dict if return_dict is not None else self.config.return_dict

            if position_ids is None:
                 __snake_case         =					jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__SCREAMING_SNAKE_CASE					).shape[-1]					)      ,					input_ids.shape					)

            if token_type_ids is None:
                 __snake_case         =					jnp.zeros_like(__SCREAMING_SNAKE_CASE					)

            if attention_mask is None:
                 __snake_case         =					jnp.ones_like(__SCREAMING_SNAKE_CASE					)

            # Handle any PRNG if needed
            __snake_case         =					{}
            if dropout_rng is not None:
                 __snake_case         =					dropout_rng

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


       def  lowerCAmelCase (				self      ,					__SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE=None      ,					__SCREAMING_SNAKE_CASE=None      ,					__SCREAMING_SNAKE_CASE=None      ,					__SCREAMING_SNAKE_CASE = None      ,					__SCREAMING_SNAKE_CASE = None      ,					__SCREAMING_SNAKE_CASE=False      ,					) ->					Any:



            '''simple docstring'''




            if position_ids is None:
                 __snake_case         =					jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__SCREAMING_SNAKE_CASE					).shape[-1]					)      ,					input_ids.shape					)

            if token_type_ids is None:
                 __snake_case         =					jnp.zeros_like(__SCREAMING_SNAKE_CASE					)

            if attention_mask is None:
                 __snake_case         =					jnp.ones_like(__SCREAMING_SNAKE_CASE					)

            # Handle any PRNG if needed
            __snake_case         =					{}
            if dropout_rng is not None:
                 __snake_case         =					dropout_rng

            def _get_features(__SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE					):
                 __snake_case         =					module.text_model(
                     input_ids=__SCREAMING_SNAKE_CASE      ,					attention_mask=__SCREAMING_SNAKE_CASE      ,					position_ids=__SCREAMING_SNAKE_CASE      ,					token_type_ids=__SCREAMING_SNAKE_CASE      ,					deterministic=__SCREAMING_SNAKE_CASE      ,					)
                 __snake_case         =					text_outputs[1]
                 __snake_case         =					module.text_projection(__SCREAMING_SNAKE_CASE					)
                 return text_features

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


       def  lowerCAmelCase (				self      ,					__SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE = None      ,					__SCREAMING_SNAKE_CASE = None      ,					__SCREAMING_SNAKE_CASE=False					) ->					Tuple:



            '''simple docstring'''




            __snake_case         =					{}
            if dropout_rng is not None:
                 __snake_case         =					dropout_rng

            def _get_features(__SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE					):
                 __snake_case         =					module.vision_model(pixel_values=__SCREAMING_SNAKE_CASE      ,					deterministic=__SCREAMING_SNAKE_CASE					)
                 __snake_case         =					vision_outputs[1]  # pooled_output
                 __snake_case         =					module.visual_projection(__SCREAMING_SNAKE_CASE					)
                 return image_features

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




       @classmethod
       def  lowerCAmelCase (				cls      ,					__SCREAMING_SNAKE_CASE = None      ,					__SCREAMING_SNAKE_CASE = None      ,					*__SCREAMING_SNAKE_CASE      ,					**__SCREAMING_SNAKE_CASE      ,					) ->					FlaxPreTrainedModel:



            '''simple docstring'''




            __snake_case         =					{
                argument[len('''text_'''					) :]: value for argument, value in kwargs.items() if argument.startswith('''text_'''					)
            }

            __snake_case         =					{
                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
            __snake_case         =					kwargs_text.pop('''model'''      ,					__SCREAMING_SNAKE_CASE					)
            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

                      __snake_case         =					AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE					)
                      __snake_case         =					text_config

                 __snake_case         =					FlaxAutoModel.from_pretrained(__SCREAMING_SNAKE_CASE      ,					*__SCREAMING_SNAKE_CASE      ,					**__SCREAMING_SNAKE_CASE					)

            __snake_case         =					kwargs_vision.pop('''model'''      ,					__SCREAMING_SNAKE_CASE					)
            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

                      __snake_case         =					AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE					)
                      __snake_case         =					vision_config

                 __snake_case         =					FlaxAutoModel.from_pretrained(__SCREAMING_SNAKE_CASE      ,					*__SCREAMING_SNAKE_CASE      ,					**__SCREAMING_SNAKE_CASE					)

            # instantiate config with corresponding kwargs
            __snake_case         =					kwargs.pop('''dtype'''      ,					jnp.floataa					)
            __snake_case         =					HybridCLIPConfig.from_text_vision_configs(text_model.config      ,					vision_model.config      ,					**__SCREAMING_SNAKE_CASE					)

            # init model
            __snake_case         =					cls(__SCREAMING_SNAKE_CASE      ,					*__SCREAMING_SNAKE_CASE      ,					dtype=__SCREAMING_SNAKE_CASE      ,					**__SCREAMING_SNAKE_CASE					)

            if vision_config.model_type == "clip":
                 __snake_case         =					vision_model.params['''vision_model''']
                 __snake_case         =					vision_model.params['''visual_projection''']['''kernel''']
            else:
                 __snake_case         =					vision_model.params

            __snake_case         =					text_model.params

            return model












'''simple docstring'''



import inspect

from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
     import jax





@require_flax
class       lowerCAmelCase				:


       def  lowerCAmelCase (				self					) ->					int:



            '''simple docstring'''




            __snake_case   ,      __snake_case         =					self.prepare_init_args_and_inputs_for_common()

            __snake_case         =					self.model_class(**__SCREAMING_SNAKE_CASE					)
            __snake_case         =					model.init(inputs_dict['''prng_key''']      ,					inputs_dict['''sample''']					)
            jax.lax.stop_gradient(__SCREAMING_SNAKE_CASE					)

            __snake_case         =					model.apply(__SCREAMING_SNAKE_CASE      ,					inputs_dict['''sample''']					)

            if isinstance(__SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE					):
                 __snake_case         =					output.sample

            self.assertIsNotNone(__SCREAMING_SNAKE_CASE					)
            __snake_case         =					inputs_dict['''sample'''].shape
            self.assertEqual(output.shape      ,					__SCREAMING_SNAKE_CASE      ,					'''Input and output shapes do not match'''					)


       def  lowerCAmelCase (				self					) ->					Tuple:



            '''simple docstring'''




            __snake_case   ,      __snake_case         =					self.prepare_init_args_and_inputs_for_common()

            __snake_case         =					16
            __snake_case         =					(16, 32)

            __snake_case         =					self.model_class(**__SCREAMING_SNAKE_CASE					)
            __snake_case         =					model.init(inputs_dict['''prng_key''']      ,					inputs_dict['''sample''']					)
            jax.lax.stop_gradient(__SCREAMING_SNAKE_CASE					)

            __snake_case         =					model.apply(__SCREAMING_SNAKE_CASE      ,					inputs_dict['''sample''']					)

            if isinstance(__SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE					):
                 __snake_case         =					output.sample

            self.assertIsNotNone(__SCREAMING_SNAKE_CASE					)
            __snake_case         =					inputs_dict['''sample'''].shape
            self.assertEqual(output.shape      ,					__SCREAMING_SNAKE_CASE      ,					'''Input and output shapes do not match'''					)




       def  lowerCAmelCase (				self					) ->					str:



            '''simple docstring'''




            __snake_case         =					'''kwargs''' in inspect.signature(self.model_class.__init__					).parameters
            __snake_case         =					len(self.model_class._deprecated_kwargs					) > 0

            if has_kwarg_in_model_class and not has_deprecated_kwarg:
                 raise ValueError(
                     F'''{self.model_class} has `**kwargs` in its __init__ method but has not defined any deprecated kwargs'''
                     ''' under the `_deprecated_kwargs` class attribute. Make sure to either remove `**kwargs` if there are'''
                     ''' no deprecated arguments or add the deprecated argument with `_deprecated_kwargs ='''
                     ''' [<deprecated_argument>]`'''					)

            if not has_kwarg_in_model_class and has_deprecated_kwarg:
                 raise ValueError(
                     F'''{self.model_class} doesn\'t have `**kwargs` in its __init__ method but has defined deprecated kwargs'''
                     ''' under the `_deprecated_kwargs` class attribute. Make sure to either add the `**kwargs` argument to'''
                     F''' {self.model_class}.__init__ if there are deprecated arguments or remove the deprecated argument'''
                     ''' from `_deprecated_kwargs = [<deprecated_argument>]`'''					)






