

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__			:					int							  =					logging.get_logger(__name__)






class       _UpperCAmelCase     (						nn.Module):
			__a	:    HybridCLIPConfig
			__a	:    jnp.dtype			   =							jnp.floataa
			def 							__snake_case	(						self       )      ->      Tuple:






				'''simple docstring'''







				_UpperCAmelCase      :							List[Any]       =   self.config.text_config
				_UpperCAmelCase      :							Optional[Any]       =   self.config.vision_config

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

				_UpperCAmelCase      :							str       =   FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
				_UpperCAmelCase      :							int       =   FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__ ,						_UpperCamelCase       ).module_class

				_UpperCAmelCase      :							int       =   text_module(_UpperCamelCase ,						dtype=self.dtype       )
				_UpperCAmelCase      :							str       =   vision_module(_UpperCamelCase ,						dtype=self.dtype       )

				_UpperCAmelCase      :							int       =   nn.Dense(
				    self.projection_dim ,						dtype=self.dtype ,						kernel_init=jax.nn.initializers.normal(0.02       ) ,						use_bias=_UpperCamelCase ,						)
				_UpperCAmelCase      :							Dict       =   nn.Dense(
				    self.projection_dim ,						dtype=self.dtype ,						kernel_init=jax.nn.initializers.normal(0.02       ) ,						use_bias=_UpperCamelCase ,						)
				_UpperCAmelCase      :							List[Any]       =   self.param("""logit_scale""" ,						jax.nn.initializers.ones ,						[]       )
			def __call__(						self ,						_A=None ,						_A=None ,						_A=None ,						_A=None ,						_A=None ,						_A = True ,						_A=None ,						_A=None ,						_A=None ,						)      ->      str:






				'''simple docstring'''







				_UpperCAmelCase      :							str       =   return_dict if return_dict is not None else self.config.return_dict

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

				_UpperCAmelCase      :							List[str]       =   self.text_model(
				    input_ids=_UpperCamelCase ,						attention_mask=_UpperCamelCase ,						token_type_ids=_UpperCamelCase ,						position_ids=_UpperCamelCase ,						deterministic=_UpperCamelCase ,						output_attentions=_UpperCamelCase ,						output_hidden_states=_UpperCamelCase ,						return_dict=_UpperCamelCase ,						)

				_UpperCAmelCase      :							Optional[int]       =   vision_outputs[1]
				_UpperCAmelCase      :							List[Any]       =   self.visual_projection(_UpperCamelCase       )

				_UpperCAmelCase      :							Union[str, Any]       =   text_outputs[1]
				_UpperCAmelCase      :							str       =   self.text_projection(_UpperCamelCase       )

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

				# cosine similarity as logits
				_UpperCAmelCase      :							Union[str, Any]       =   jnp.exp(self.logit_scale       )
				_UpperCAmelCase      :							int       =   jnp.matmul(_UpperCamelCase ,						image_embeds.T       ) * logit_scale
				_UpperCAmelCase      :							Union[str, 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=_UpperCamelCase ,						logits_per_text=_UpperCamelCase ,						text_embeds=_UpperCamelCase ,						image_embeds=_UpperCamelCase ,						text_model_output=_UpperCamelCase ,						vision_model_output=_UpperCamelCase ,						)






class       _UpperCAmelCase     (						__lowerCamelCase):
			__a	:    Tuple			   =							HybridCLIPConfig
			__a	:    Optional[int]			   =							FlaxHybridCLIPModule
			def __init__(						self ,						_A ,						_A = None ,						_A = 0 ,						_A = jnp.floataa ,						**_A ,						)      ->      Dict:






				'''simple docstring'''







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

				_UpperCAmelCase      :							Tuple       =   self.module_class(config=_UpperCamelCase ,						dtype=_UpperCamelCase ,						**_UpperCamelCase       )
				super().__init__(_UpperCamelCase ,						_UpperCamelCase ,						input_shape=_UpperCamelCase ,						seed=_UpperCamelCase ,						dtype=_UpperCamelCase       )
			def 							__snake_case	(						self ,						_A ,						_A ,						_A = None       )      ->      FrozenDict:






				'''simple docstring'''







				_UpperCAmelCase      :							int       =   jnp.zeros(input_shape[0] ,						dtype="""i4"""       )
				_UpperCAmelCase      :							str       =   jnp.broadcast_to(jnp.arange(jnp.atleast_ad(_UpperCamelCase       ).shape[-1]       ) ,						input_shape[0]       )
				_UpperCAmelCase      :							Optional[int]       =   jnp.ones_like(_UpperCamelCase       )
				_UpperCAmelCase      :							Union[str, Any]       =   jnp.ones_like(_UpperCamelCase       )

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

				_UpperCAmelCase  ,				_UpperCAmelCase      :							Union[str, Any]       =   jax.random.split(_UpperCamelCase       )
				_UpperCAmelCase      :							Optional[Any]       =   {"""params""": params_rng, """dropout""": dropout_rng}

				return self.module.init(_UpperCamelCase ,						_UpperCamelCase ,						_UpperCamelCase ,						_UpperCamelCase ,						_UpperCamelCase ,						_UpperCamelCase       )["params"]
			def __call__(						self ,						_A ,						_A ,						_A=None ,						_A=None ,						_A=None ,						_A = None ,						_A = None ,						_A = False ,						_A = None ,						_A = None ,						_A = None ,						)      ->      List[Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							Optional[int]       =   output_attentions if output_attentions is not None else self.config.output_attentions
				_UpperCAmelCase      :							List[str]       =   (
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)
				_UpperCAmelCase      :							Dict       =   return_dict if return_dict is not None else self.config.return_dict

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

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

				if attention_mask is None:
					_UpperCAmelCase      :							Optional[int]       =   jnp.ones_like(_UpperCamelCase       )

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

				return self.module.apply(
				    {"""params""": params or self.params} ,						jnp.array(_UpperCamelCase ,						dtype="""i4"""       ) ,						jnp.array(_UpperCamelCase ,						dtype=jnp.floataa       ) ,						jnp.array(_UpperCamelCase ,						dtype="""i4"""       ) ,						jnp.array(_UpperCamelCase ,						dtype="""i4"""       ) ,						jnp.array(_UpperCamelCase ,						dtype="""i4"""       ) ,						not train ,						_UpperCamelCase ,						_UpperCamelCase ,						_UpperCamelCase ,						rngs=_UpperCamelCase ,						)
			def 							__snake_case	(						self ,						_A ,						_A=None ,						_A=None ,						_A=None ,						_A = None ,						_A = None ,						_A=False ,						)      ->      Union[str, Any]:






				'''simple docstring'''







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

				if token_type_ids is None:
					_UpperCAmelCase      :							Optional[int]       =   jnp.zeros_like(_UpperCamelCase       )

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

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

				def _get_features(_A ,						_A ,						_A ,						_A ,						_A ,						_A       ):
					_UpperCAmelCase      :							Any       =   module.text_model(
					    input_ids=_UpperCamelCase ,						attention_mask=_UpperCamelCase ,						position_ids=_UpperCamelCase ,						token_type_ids=_UpperCamelCase ,						deterministic=_UpperCamelCase ,						)
					_UpperCAmelCase      :							Any       =   text_outputs[1]
					_UpperCAmelCase      :							Optional[Any]       =   module.text_projection(_UpperCamelCase       )
					return text_features

				return self.module.apply(
				    {"""params""": params or self.params} ,						jnp.array(_UpperCamelCase ,						dtype="""i4"""       ) ,						jnp.array(_UpperCamelCase ,						dtype="""i4"""       ) ,						jnp.array(_UpperCamelCase ,						dtype="""i4"""       ) ,						jnp.array(_UpperCamelCase ,						dtype="""i4"""       ) ,						not train ,						method=_get_features ,						rngs=_UpperCamelCase ,						)
			def 							__snake_case	(						self ,						_A ,						_A = None ,						_A = None ,						_A=False       )      ->      List[str]:






				'''simple docstring'''







				_UpperCAmelCase      :							List[str]       =   {}
				if dropout_rng is not None:
					_UpperCAmelCase      :							str       =   dropout_rng

				def _get_features(_A ,						_A ,						_A       ):
					_UpperCAmelCase      :							List[str]       =   module.vision_model(pixel_values=_UpperCamelCase ,						deterministic=_UpperCamelCase       )
					_UpperCAmelCase      :							int       =   vision_outputs[1]  # pooled_output
					_UpperCAmelCase      :							int       =   module.visual_projection(_UpperCamelCase       )
					return image_features

				return self.module.apply(
				    {"""params""": params or self.params} ,						jnp.array(_UpperCamelCase ,						dtype=jnp.floataa       ) ,						not train ,						method=_get_features ,						rngs=_UpperCamelCase ,						)
			@classmethod
			def 							__snake_case	(						cls ,						_A = None ,						_A = None ,						*_A ,						**_A ,						)      ->      FlaxPreTrainedModel:






				'''simple docstring'''







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

				_UpperCAmelCase      :							Dict       =   {
				    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      :							Union[str, Any]       =   kwargs_text.pop("""model""" ,						_UpperCamelCase       )
				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      :							List[str]       =   AutoConfig.from_pretrained(_UpperCamelCase       )
						_UpperCAmelCase      :							List[str]       =   text_config

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

				_UpperCAmelCase      :							Union[str, Any]       =   kwargs_vision.pop("""model""" ,						_UpperCamelCase       )
				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      :							Any       =   AutoConfig.from_pretrained(_UpperCamelCase       )
						_UpperCAmelCase      :							Union[str, Any]       =   vision_config

					_UpperCAmelCase      :							int       =   FlaxAutoModel.from_pretrained(_UpperCamelCase ,						*_UpperCamelCase ,						**_UpperCamelCase       )

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

				# init model
				_UpperCAmelCase      :							Dict       =   cls(_UpperCamelCase ,						*_UpperCamelCase ,						dtype=_UpperCamelCase ,						**_UpperCamelCase       )

				if vision_config.model_type == "clip":
					_UpperCAmelCase      :							Dict       =   vision_model.params["""vision_model"""]
					_UpperCAmelCase      :							Union[str, Any]       =   vision_model.params["""visual_projection"""]["""kernel"""]
				else:
					_UpperCAmelCase      :							Dict       =   vision_model.params

				_UpperCAmelCase      :							List[Any]       =   text_model.params

				return model








"""simple docstring"""






import unittest

import numpy as np
import torch
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()






class       _UpperCAmelCase     (						__a      ,       unittest.TestCase):
			__a	:    List[Any]			   =							KandinskyPriorPipeline
			__a	:    List[Any]			   =							["""prompt"""]
			__a	:    List[str]			   =							["""prompt""", """negative_prompt"""]
			__a	:    Tuple			   =							[
			    """num_images_per_prompt""",
			    """generator""",
			    """num_inference_steps""",
			    """latents""",
			    """negative_prompt""",
			    """guidance_scale""",
			    """output_type""",
			    """return_dict""",
			]
			__a	:    str			   =							False
			@property
			def 							__snake_case	(						self       )      ->      Any:






				'''simple docstring'''







				return 32
			@property
			def 							__snake_case	(						self       )      ->      Union[str, Any]:






				'''simple docstring'''







				return 32
			@property
			def 							__snake_case	(						self       )      ->      Dict:






				'''simple docstring'''







				return self.time_input_dim
			@property
			def 							__snake_case	(						self       )      ->      Optional[Any]:






				'''simple docstring'''







				return self.time_input_dim * 4
			@property
			def 							__snake_case	(						self       )      ->      List[Any]:






				'''simple docstring'''







				return 1_00
			@property
			def 							__snake_case	(						self       )      ->      Optional[int]:






				'''simple docstring'''







				_UpperCAmelCase      :							List[str]       =   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""       )
				return tokenizer
			@property
			def 							__snake_case	(						self       )      ->      int:






				'''simple docstring'''







				torch.manual_seed(0       )
				_UpperCAmelCase      :							List[Any]       =   CLIPTextConfig(
				    bos_token_id=0 ,						eos_token_id=2 ,						hidden_size=self.text_embedder_hidden_size ,						projection_dim=self.text_embedder_hidden_size ,						intermediate_size=37 ,						layer_norm_eps=1e-05 ,						num_attention_heads=4 ,						num_hidden_layers=5 ,						pad_token_id=1 ,						vocab_size=10_00 ,						)
				return CLIPTextModelWithProjection(_A       )
			@property
			def 							__snake_case	(						self       )      ->      List[Any]:






				'''simple docstring'''







				torch.manual_seed(0       )

				_UpperCAmelCase      :							Optional[int]       =   {
				    """num_attention_heads""": 2,
				    """attention_head_dim""": 12,
				    """embedding_dim""": self.text_embedder_hidden_size,
				    """num_layers""": 1,
				}

				_UpperCAmelCase      :							List[Any]       =   PriorTransformer(**_A       )
				# clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
				_UpperCAmelCase      :							List[Any]       =   nn.Parameter(torch.ones(model.clip_std.shape       )       )
				return model
			@property
			def 							__snake_case	(						self       )      ->      Dict:






				'''simple docstring'''







				torch.manual_seed(0       )
				_UpperCAmelCase      :							Optional[Any]       =   CLIPVisionConfig(
				    hidden_size=self.text_embedder_hidden_size ,						image_size=2_24 ,						projection_dim=self.text_embedder_hidden_size ,						intermediate_size=37 ,						num_attention_heads=4 ,						num_channels=3 ,						num_hidden_layers=5 ,						patch_size=14 ,						)

				_UpperCAmelCase      :							Tuple       =   CLIPVisionModelWithProjection(_A       )
				return model
			@property
			def 							__snake_case	(						self       )      ->      List[str]:






				'''simple docstring'''







				_UpperCAmelCase      :							str       =   CLIPImageProcessor(
				    crop_size=2_24 ,						do_center_crop=_A ,						do_normalize=_A ,						do_resize=_A ,						image_mean=[0.48145466, 0.4578275, 0.40821073] ,						image_std=[0.26862954, 0.26130258, 0.27577711] ,						resample=3 ,						size=2_24 ,						)

				return image_processor
			def 							__snake_case	(						self       )      ->      Dict:






				'''simple docstring'''







				_UpperCAmelCase      :							Optional[int]       =   self.dummy_prior
				_UpperCAmelCase      :							int       =   self.dummy_image_encoder
				_UpperCAmelCase      :							int       =   self.dummy_text_encoder
				_UpperCAmelCase      :							str       =   self.dummy_tokenizer
				_UpperCAmelCase      :							Any       =   self.dummy_image_processor

				_UpperCAmelCase      :							Any       =   UnCLIPScheduler(
				    variance_type="""fixed_small_log""" ,						prediction_type="""sample""" ,						num_train_timesteps=10_00 ,						clip_sample=_A ,						clip_sample_range=10.0 ,						)

				_UpperCAmelCase      :							Any       =   {
				    """prior""": prior,
				    """image_encoder""": image_encoder,
				    """text_encoder""": text_encoder,
				    """tokenizer""": tokenizer,
				    """scheduler""": scheduler,
				    """image_processor""": image_processor,
				}

				return components
			def 							__snake_case	(						self ,						_A ,						_A=0       )      ->      Optional[Any]:






				'''simple docstring'''







				if str(_A       ).startswith("""mps"""       ):
					_UpperCAmelCase      :							Any       =   torch.manual_seed(_A       )
				else:
					_UpperCAmelCase      :							Optional[int]       =   torch.Generator(device=_A       ).manual_seed(_A       )
				_UpperCAmelCase      :							Union[str, Any]       =   {
				    """prompt""": """horse""",
				    """generator""": generator,
				    """guidance_scale""": 4.0,
				    """num_inference_steps""": 2,
				    """output_type""": """np""",
				}
				return inputs
			def 							__snake_case	(						self       )      ->      Union[str, Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							List[str]       =   """cpu"""

				_UpperCAmelCase      :							int       =   self.get_dummy_components()

				_UpperCAmelCase      :							Optional[Any]       =   self.pipeline_class(**_A       )
				_UpperCAmelCase      :							List[Any]       =   pipe.to(_A       )

				pipe.set_progress_bar_config(disable=_A       )

				_UpperCAmelCase      :							str       =   pipe(**self.get_dummy_inputs(_A       )       )
				_UpperCAmelCase      :							Any       =   output.image_embeds

				_UpperCAmelCase      :							Optional[int]       =   pipe(
				    **self.get_dummy_inputs(_A       ) ,						return_dict=_A ,						)[0]

				_UpperCAmelCase      :							int       =   image[0, -10:]
				_UpperCAmelCase      :							Optional[Any]       =   image_from_tuple[0, -10:]

				assert image.shape == (1, 32)

				_UpperCAmelCase      :							Optional[Any]       =   np.array(
				    [-0.0532, 1.7120, 0.3656, -1.0852, -0.8946, -1.1756, 0.4348, 0.2482, 0.5146, -0.1156]       )

				assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2
				assert np.abs(image_from_tuple_slice.flatten() - expected_slice       ).max() < 1e-2
			@skip_mps
			def 							__snake_case	(						self       )      ->      Tuple:






				'''simple docstring'''







				_UpperCAmelCase      :							str       =   torch_device == """cpu"""
				_UpperCAmelCase      :							Optional[Any]       =   True
				_UpperCAmelCase      :							List[Any]       =   False

				self._test_inference_batch_single_identical(
				    test_max_difference=_A ,						relax_max_difference=_A ,						test_mean_pixel_difference=_A ,						)
			@skip_mps
			def 							__snake_case	(						self       )      ->      Any:






				'''simple docstring'''







				_UpperCAmelCase      :							Any       =   torch_device == """cpu"""
				_UpperCAmelCase      :							List[str]       =   False

				self._test_attention_slicing_forward_pass(
				    test_max_difference=_A ,						test_mean_pixel_difference=_A ,						)





