





'''simple docstring'''







import tempfile
import unittest

import torch

from diffusers import UNetaDConditionModel
from diffusers.training_utils import EMAModel
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps, torch_device


enable_full_determinism()


class    __a (unittest.TestCase       ):
							__a   :				Tuple       =    "hf-internal-testing/tiny-stable-diffusion-pipe"
							__a   :				List[Any]       =    1
							__a   :				Any       =    77
							__a   :				List[str]       =    32
							__a   :				Optional[int]       =    4
							__a   :				Optional[int]       =    64
							__a   :				Optional[Any]       =    torch.manual_seed(0       )

							def       UpperCAmelCase__		(							self							:       int    ,       __magic_name__							:       Any=0.9_9_9_9				)   ->    Union[str, Any]:







													"""simple docstring"""
													UpperCAmelCase_     : Any         =					UNetaDConditionModel.from_pretrained(self.model_id    ,       subfolder='''unet'''				)
													UpperCAmelCase_     : Optional[Any]         =					unet.to(__magic_name__				)
													UpperCAmelCase_     : Optional[Any]         =					EMAModel(unet.parameters()    ,       decay=__magic_name__    ,       model_cls=__magic_name__    ,       model_config=unet.config				)
													return unet, ema_unet

							def       UpperCAmelCase__		(							self							:       Optional[int]				)   ->    int:







													"""simple docstring"""
													UpperCAmelCase_     : str         =					torch.randn(
													    self.batch_size    ,       self.num_in_channels    ,       self.latent_height    ,       self.latent_width    ,       generator=self.generator				).to(__magic_name__				)
													UpperCAmelCase_     : List[str]         =					torch.randint(0    ,       10_00    ,       size=(self.batch_size,)    ,       generator=self.generator				).to(__magic_name__				)
													UpperCAmelCase_     : Optional[int]         =					torch.randn(
													    self.batch_size    ,       self.prompt_length    ,       self.text_encoder_hidden_dim    ,       generator=self.generator				).to(__magic_name__				)
													return noisy_latents, timesteps, encoder_hidden_states

							def       UpperCAmelCase__		(							self							:       List[Any]    ,       __magic_name__							:       str				)   ->    Tuple:







													"""simple docstring"""
													UpperCAmelCase_     : int         =					{}
													for k, param in unet.state_dict().items():
																			UpperCAmelCase_     : str         =					torch.randn_like(__magic_name__				) + (param * torch.randn_like(__magic_name__				))
																			updated_state_dict.update({k: updated_param}				)
													unet.load_state_dict(__magic_name__				)
													return unet

							def       UpperCAmelCase__		(							self							:       Any				)   ->    str:







													"""simple docstring"""
													UpperCAmelCase_     : List[str]         =					self.get_models()
													# Take the first (hypothetical) EMA step.
													ema_unet.step(unet.parameters()				)
													assert ema_unet.optimization_step == 1

													# Take two more.
													for _ in range(2				):
																			ema_unet.step(unet.parameters()				)
													assert ema_unet.optimization_step == 3

							def       UpperCAmelCase__		(							self							:       Tuple				)   ->    str:







													"""simple docstring"""
													UpperCAmelCase_     : List[str]         =					self.get_models()
													# Since the `unet` is not being updated (i.e., backprop'd)
													# there won't be any difference between the `params` of `unet`
													# and `ema_unet` even if we call `ema_unet.step(unet.parameters())`.
													ema_unet.step(unet.parameters()				)
													UpperCAmelCase_     : Optional[Any]         =					list(unet.parameters()				)
													for s_param, param in zip(ema_unet.shadow_params    ,       __magic_name__				):
																			assert torch.allclose(__magic_name__    ,       __magic_name__				)

													# The above holds true even if we call `ema.step()` multiple times since
													# `unet` params are still not being updated.
													for _ in range(4				):
																			ema_unet.step(unet.parameters()				)
													for s_param, param in zip(ema_unet.shadow_params    ,       __magic_name__				):
																			assert torch.allclose(__magic_name__    ,       __magic_name__				)

							def       UpperCAmelCase__		(							self							:       str				)   ->    Tuple:







													"""simple docstring"""
													UpperCAmelCase_     : Dict         =					self.get_models()
													# Here we simulate the parameter updates for `unet`. Since there might
													# be some parameters which are initialized to zero we take extra care to
													# initialize their values to something non-zero before the multiplication.
													UpperCAmelCase_     : int         =					self.simulate_backprop(__magic_name__				)

													# Take the EMA step.
													ema_unet.step(unet_pseudo_updated_step_one.parameters()				)

													# Now the EMA'd parameters won't be equal to the original model parameters.
													UpperCAmelCase_     : List[str]         =					list(unet_pseudo_updated_step_one.parameters()				)
													for s_param, param in zip(ema_unet.shadow_params    ,       __magic_name__				):
																			assert ~torch.allclose(__magic_name__    ,       __magic_name__				)

													# Ensure this is the case when we take multiple EMA steps.
													for _ in range(4				):
																			ema_unet.step(unet.parameters()				)
													for s_param, param in zip(ema_unet.shadow_params    ,       __magic_name__				):
																			assert ~torch.allclose(__magic_name__    ,       __magic_name__				)

							def       UpperCAmelCase__		(							self							:       Dict				)   ->    Tuple:







													"""simple docstring"""
													UpperCAmelCase_     : str         =					self.get_models()

													# First backprop + EMA
													UpperCAmelCase_     : List[Any]         =					self.simulate_backprop(__magic_name__				)
													ema_unet.step(unet_step_one.parameters()				)
													UpperCAmelCase_     : Optional[int]         =					ema_unet.shadow_params

													# Second backprop + EMA
													UpperCAmelCase_     : Optional[int]         =					self.simulate_backprop(__magic_name__				)
													ema_unet.step(unet_step_two.parameters()				)
													UpperCAmelCase_     : Any         =					ema_unet.shadow_params

													for step_one, step_two in zip(__magic_name__    ,       __magic_name__				):
																			assert ~torch.allclose(__magic_name__    ,       __magic_name__				)

							def       UpperCAmelCase__		(							self							:       Tuple				)   ->    List[str]:







													"""simple docstring"""
													UpperCAmelCase_     : List[Any]         =					self.get_models(decay=0.0				)
													UpperCAmelCase_     : List[Any]         =					self.simulate_backprop(__magic_name__				)
													ema_unet.step(unet_step_one.parameters()				)
													UpperCAmelCase_     : str         =					ema_unet.shadow_params

													UpperCAmelCase_     : List[Any]         =					self.simulate_backprop(__magic_name__				)
													ema_unet.step(unet_step_two.parameters()				)
													UpperCAmelCase_     : Optional[Any]         =					ema_unet.shadow_params

													for step_one, step_two in zip(__magic_name__    ,       __magic_name__				):
																			assert torch.allclose(__magic_name__    ,       __magic_name__				)



							@skip_mps
							def       UpperCAmelCase__		(							self							:       Optional[int]				)   ->    List[str]:







													"""simple docstring"""
													UpperCAmelCase_     : Optional[int]         =					self.get_models()
													UpperCAmelCase_     : Union[str, Any]         =					self.get_dummy_inputs()

													with tempfile.TemporaryDirectory() as tmpdir:
																			ema_unet.save_pretrained(__magic_name__				)
																			UpperCAmelCase_     : Any         =					UNetaDConditionModel.from_pretrained(__magic_name__    ,       model_cls=__magic_name__				)
																			UpperCAmelCase_     : Any         =					loaded_unet.to(unet.device				)

													# Since no EMA step has been performed the outputs should match.
													UpperCAmelCase_     : List[str]         =					unet(__magic_name__    ,       __magic_name__    ,       __magic_name__				).sample
													UpperCAmelCase_     : Optional[int]         =					loaded_unet(__magic_name__    ,       __magic_name__    ,       __magic_name__				).sample

													assert torch.allclose(__magic_name__    ,       __magic_name__    ,       atol=1E-4				)












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


snake_case_       :    int          =					logging.get_logger(__name__)


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

							def       UpperCAmelCase__		(							self							:       Optional[Any]				)   ->    Tuple:







													"""simple docstring"""
													UpperCAmelCase_     : Dict         =					self.config.text_config
													UpperCAmelCase_     : Dict         =					self.config.vision_config

													UpperCAmelCase_     : Tuple         =					self.config.projection_dim
													UpperCAmelCase_     : str         =					text_config.hidden_size
													UpperCAmelCase_     : List[Any]         =					vision_config.hidden_size

													UpperCAmelCase_     : int         =					FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
													UpperCAmelCase_     : int         =					FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__    ,       __magic_name__				).module_class

													UpperCAmelCase_     : str         =					text_module(__magic_name__    ,       dtype=self.dtype				)
													UpperCAmelCase_     : List[Any]         =					vision_module(__magic_name__    ,       dtype=self.dtype				)

													UpperCAmelCase_     : int         =					nn.Dense(
													    self.projection_dim    ,       dtype=self.dtype    ,       kernel_init=jax.nn.initializers.normal(0.0_2				)    ,       use_bias=__magic_name__    ,       )
													UpperCAmelCase_     : Union[str, Any]         =					nn.Dense(
													    self.projection_dim    ,       dtype=self.dtype    ,       kernel_init=jax.nn.initializers.normal(0.0_2				)    ,       use_bias=__magic_name__    ,       )
													UpperCAmelCase_     : Tuple         =					self.param('''logit_scale'''    ,       jax.nn.initializers.ones    ,       []				)



							def __call__(							self							:       Tuple    ,       __magic_name__							:       Optional[Any]=None    ,       __magic_name__							:       List[Any]=None    ,       __magic_name__							:       Union[str, Any]=None    ,       __magic_name__							:       str=None    ,       __magic_name__							:       List[Any]=None    ,       __magic_name__							:       bool = True    ,       __magic_name__							:       List[str]=None    ,       __magic_name__							:       Optional[int]=None    ,       __magic_name__							:       Any=None    ,       )   ->    Optional[Any]:







													"""simple docstring"""
													UpperCAmelCase_     : int         =					return_dict if return_dict is not None else self.config.return_dict

													UpperCAmelCase_     : Union[str, Any]         =					self.vision_model(
													    pixel_values=__magic_name__    ,       deterministic=__magic_name__    ,       output_attentions=__magic_name__    ,       output_hidden_states=__magic_name__    ,       return_dict=__magic_name__    ,       )

													UpperCAmelCase_     : str         =					self.text_model(
													    input_ids=__magic_name__    ,       attention_mask=__magic_name__    ,       token_type_ids=__magic_name__    ,       position_ids=__magic_name__    ,       deterministic=__magic_name__    ,       output_attentions=__magic_name__    ,       output_hidden_states=__magic_name__    ,       return_dict=__magic_name__    ,       )

													UpperCAmelCase_     : List[str]         =					vision_outputs[1]
													UpperCAmelCase_     : Union[str, Any]         =					self.visual_projection(__magic_name__				)

													UpperCAmelCase_     : List[Any]         =					text_outputs[1]
													UpperCAmelCase_     : Optional[Any]         =					self.text_projection(__magic_name__				)

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

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





class    __a (lowerCamelCase       ):
							__a   :				int       =    HybridCLIPConfig
							__a   :				Optional[Any]       =    FlaxHybridCLIPModule

							def __init__(							self							:       Dict    ,       __magic_name__							:       HybridCLIPConfig    ,       __magic_name__							:       Optional[Tuple] = None    ,       __magic_name__							:       int = 0    ,       __magic_name__							:       jnp.dtype = jnp.floataa    ,       **__magic_name__							:       Optional[Any]    ,       )   ->    str:







													"""simple docstring"""
													if input_shape is None:
																			UpperCAmelCase_     : Optional[Any]         =					((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))

													UpperCAmelCase_     : Any         =					self.module_class(config=__magic_name__    ,       dtype=__magic_name__    ,       **__magic_name__				)
													super().__init__(__magic_name__    ,       __magic_name__    ,       input_shape=__magic_name__    ,       seed=__magic_name__    ,       dtype=__magic_name__				)

							def       UpperCAmelCase__		(							self							:       List[Any]    ,       __magic_name__							:       jax.random.PRNGKey    ,       __magic_name__							:       Tuple    ,       __magic_name__							:       FrozenDict = None				)   ->    FrozenDict:







													"""simple docstring"""
													# init input tensor
													UpperCAmelCase_     : List[Any]         =					jnp.zeros(input_shape[0]    ,       dtype='''i4'''				)
													UpperCAmelCase_     : List[Any]         =					jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__magic_name__				).shape[-1]				)    ,       input_shape[0]				)
													UpperCAmelCase_     : Optional[int]         =					jnp.ones_like(__magic_name__				)
													UpperCAmelCase_     : str         =					jnp.ones_like(__magic_name__				)

													UpperCAmelCase_     : Optional[int]         =					jax.random.normal(__magic_name__    ,       input_shape[1]				)

													UpperCAmelCase_  ,    UpperCAmelCase_     : Optional[int]         =					jax.random.split(__magic_name__				)
													UpperCAmelCase_     : List[Any]         =					{'''params''': params_rng, '''dropout''': dropout_rng}

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

							def __call__(							self							:       Optional[Any]    ,       __magic_name__							:       List[str]    ,       __magic_name__							:       int    ,       __magic_name__							:       List[str]=None    ,       __magic_name__							:       Union[str, Any]=None    ,       __magic_name__							:       List[str]=None    ,       __magic_name__							:       dict = None    ,       __magic_name__							:       jax.random.PRNGKey = None    ,       __magic_name__							:       bool = False    ,       __magic_name__							:       Optional[bool] = None    ,       __magic_name__							:       Optional[bool] = None    ,       __magic_name__							:       Optional[bool] = None    ,       )   ->    List[Any]:







													"""simple docstring"""
													UpperCAmelCase_     : List[Any]         =					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_     : List[Any]         =					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(__magic_name__				).shape[-1]				)    ,       input_ids.shape				)

													if token_type_ids is None:
																			UpperCAmelCase_     : Optional[Any]         =					jnp.zeros_like(__magic_name__				)

													if attention_mask is None:
																			UpperCAmelCase_     : Tuple         =					jnp.ones_like(__magic_name__				)

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

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

							def       UpperCAmelCase__		(							self							:       str    ,       __magic_name__							:       Union[str, Any]    ,       __magic_name__							:       int=None    ,       __magic_name__							:       List[str]=None    ,       __magic_name__							:       int=None    ,       __magic_name__							:       dict = None    ,       __magic_name__							:       jax.random.PRNGKey = None    ,       __magic_name__							:       Tuple=False    ,       )   ->    Tuple:







													"""simple docstring"""
													if position_ids is None:
																			UpperCAmelCase_     : Optional[int]         =					jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__magic_name__				).shape[-1]				)    ,       input_ids.shape				)

													if token_type_ids is None:
																			UpperCAmelCase_     : Tuple         =					jnp.zeros_like(__magic_name__				)

													if attention_mask is None:
																			UpperCAmelCase_     : Tuple         =					jnp.ones_like(__magic_name__				)

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

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

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

							def       UpperCAmelCase__		(							self							:       List[Any]    ,       __magic_name__							:       str    ,       __magic_name__							:       dict = None    ,       __magic_name__							:       jax.random.PRNGKey = None    ,       __magic_name__							:       Dict=False				)   ->    Optional[Any]:







													"""simple docstring"""
													UpperCAmelCase_     : List[str]         =					{}
													if dropout_rng is not None:
																			UpperCAmelCase_     : Union[str, Any]         =					dropout_rng

													def _get_features(__magic_name__							:       List[str]    ,       __magic_name__							:       List[Any]    ,       __magic_name__							:       List[Any]				):
																			UpperCAmelCase_     : Optional[int]         =					module.vision_model(pixel_values=__magic_name__    ,       deterministic=__magic_name__				)
																			UpperCAmelCase_     : Tuple         =					vision_outputs[1]  # pooled_output
																			UpperCAmelCase_     : List[Any]         =					module.visual_projection(__magic_name__				)
																			return image_features

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



							@classmethod
							def       UpperCAmelCase__		(							cls							:       Tuple    ,       __magic_name__							:       str = None    ,       __magic_name__							:       str = None    ,       *__magic_name__							:       Union[str, Any]    ,       **__magic_name__							:       Union[str, Any]    ,       )   ->    FlaxPreTrainedModel:







													"""simple docstring"""
													UpperCAmelCase_     : Union[str, Any]         =					{
													    argument[len('''text_'''				) :]: value for argument, value in kwargs.items() if argument.startswith('''text_'''				)
													}

													UpperCAmelCase_     : List[Any]         =					{
													    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[str]         =					kwargs_text.pop('''model'''    ,       __magic_name__				)
													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_     : Optional[int]         =					AutoConfig.from_pretrained(__magic_name__				)
																									UpperCAmelCase_     : Optional[int]         =					text_config

																			UpperCAmelCase_     : Dict         =					FlaxAutoModel.from_pretrained(__magic_name__    ,       *__magic_name__    ,       **__magic_name__				)

													UpperCAmelCase_     : Union[str, Any]         =					kwargs_vision.pop('''model'''    ,       __magic_name__				)
													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_     : Tuple         =					AutoConfig.from_pretrained(__magic_name__				)
																									UpperCAmelCase_     : Union[str, Any]         =					vision_config

																			UpperCAmelCase_     : Union[str, Any]         =					FlaxAutoModel.from_pretrained(__magic_name__    ,       *__magic_name__    ,       **__magic_name__				)

													# instantiate config with corresponding kwargs
													UpperCAmelCase_     : Union[str, Any]         =					kwargs.pop('''dtype'''    ,       jnp.floataa				)
													UpperCAmelCase_     : Tuple         =					HybridCLIPConfig.from_text_vision_configs(text_model.config    ,       vision_model.config    ,       **__magic_name__				)

													# init model
													UpperCAmelCase_     : Union[str, Any]         =					cls(__magic_name__    ,       *__magic_name__    ,       dtype=__magic_name__    ,       **__magic_name__				)

													if vision_config.model_type == "clip":
																			UpperCAmelCase_     : Union[str, Any]         =					vision_model.params['''vision_model''']
																			UpperCAmelCase_     : Tuple         =					vision_model.params['''visual_projection''']['''kernel''']
													else:
																			UpperCAmelCase_     : Union[str, Any]         =					vision_model.params

													UpperCAmelCase_     : Optional[Any]         =					text_model.params

													return model





