
"""simple docstring"""




from typing import Optional, Tuple

import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.core.frozen_dict import FrozenDict
from transformers import CLIPConfig, FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPVisionModule



def        _lowerCamelCase ( __a,							__a,							__a=1E-12    ):
		SCREAMING_SNAKE_CASE_						 = jnp.divide(emb_a.T,							jnp.clip(jnp.linalg.norm(__a,							axis=1    ),							a_min=__a    )    ).T
		SCREAMING_SNAKE_CASE_						 = jnp.divide(emb_a.T,							jnp.clip(jnp.linalg.norm(__a,							axis=1    ),							a_min=__a    )    ).T
		return jnp.matmul(__a,							norm_emb_a.T    )





class  snake_case     (							nn.Module   ):
							UpperCAmelCase__						     =      42
							UpperCAmelCase__						     =      jnp.floataa

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = FlaxCLIPVisionModule(self.config.vision_config    )
									SCREAMING_SNAKE_CASE_						 = nn.Dense(self.config.projection_dim       ,       use_bias=SCREAMING_SNAKE_CASE_       ,       dtype=self.dtype    )

									SCREAMING_SNAKE_CASE_						 = self.param('''concept_embeds'''       ,       jax.nn.initializers.ones       ,       (17, self.config.projection_dim)    )
									SCREAMING_SNAKE_CASE_						 = self.param(
									    '''special_care_embeds'''       ,       jax.nn.initializers.ones       ,       (3, self.config.projection_dim)    )

									SCREAMING_SNAKE_CASE_						 = self.param('''concept_embeds_weights'''       ,       jax.nn.initializers.ones       ,       (17,)    )
									SCREAMING_SNAKE_CASE_						 = self.param('''special_care_embeds_weights'''       ,       jax.nn.initializers.ones       ,       (3,)    )







							def __call__(self       ,       SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.vision_model(SCREAMING_SNAKE_CASE_    )[1]
									SCREAMING_SNAKE_CASE_						 = self.visual_projection(SCREAMING_SNAKE_CASE_    )

									SCREAMING_SNAKE_CASE_						 = jax_cosine_distance(SCREAMING_SNAKE_CASE_       ,       self.special_care_embeds    )
									SCREAMING_SNAKE_CASE_						 = jax_cosine_distance(SCREAMING_SNAKE_CASE_       ,       self.concept_embeds    )

									# increase this value to create a stronger `nfsw` filter
									# at the cost of increasing the possibility of filtering benign image inputs
									SCREAMING_SNAKE_CASE_						 = 0.0

									SCREAMING_SNAKE_CASE_						 = special_cos_dist - self.special_care_embeds_weights[None, :] + adjustment
									SCREAMING_SNAKE_CASE_						 = jnp.round(SCREAMING_SNAKE_CASE_       ,       3    )
									SCREAMING_SNAKE_CASE_						 = jnp.any(special_scores > 0       ,       axis=1       ,       keepdims=SCREAMING_SNAKE_CASE_    )
									# Use a lower threshold if an image has any special care concept
									SCREAMING_SNAKE_CASE_						 = is_special_care * 0.01

									SCREAMING_SNAKE_CASE_						 = cos_dist - self.concept_embeds_weights[None, :] + special_adjustment
									SCREAMING_SNAKE_CASE_						 = jnp.round(SCREAMING_SNAKE_CASE_       ,       3    )
									SCREAMING_SNAKE_CASE_						 = jnp.any(concept_scores > 0       ,       axis=1    )

									return has_nsfw_concepts







class  snake_case     (							__lowercase   ):
							UpperCAmelCase__						     =      CLIPConfig
							UpperCAmelCase__						     =      '''clip_input'''
							UpperCAmelCase__						     =      FlaxStableDiffusionSafetyCheckerModule

							def __init__(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_ = None       ,       SCREAMING_SNAKE_CASE_ = 0       ,       SCREAMING_SNAKE_CASE_ = jnp.floataa       ,       SCREAMING_SNAKE_CASE_ = True       ,       **SCREAMING_SNAKE_CASE_       ,       ):

									"""simple docstring"""







									if input_shape is None:
											SCREAMING_SNAKE_CASE_						 = (1, 2_24, 2_24, 3)
									SCREAMING_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_       ,       _do_init=_do_init    )

							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_ = None    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = jax.random.normal(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )

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

									SCREAMING_SNAKE_CASE_						 = self.module.init(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )['''params''']

									return random_params







							def __call__(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_ = None       ,       ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = jnp.transpose(SCREAMING_SNAKE_CASE_       ,       (0, 2, 3, 1)    )

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

"""simple docstring"""




import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionPipeline
from diffusers.utils.testing_utils import load_image, nightly, require_torch_gpu, torch_device


lowerCAmelCase__							=				False





class  snake_case     (							unittest.TestCase   ):
							pass







@nightly
@require_torch_gpu
class  snake_case     (							unittest.TestCase   ):

							def _lowercase					(self    ):

									"""simple docstring"""







									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = VersatileDiffusionPipeline.from_pretrained('''shi-labs/versatile-diffusion'''       ,       torch_dtype=torch.floataa    )
									pipe.to(SCREAMING_SNAKE_CASE_    )
									pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_    )

									SCREAMING_SNAKE_CASE_						 = load_image(
									    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'''    )

									SCREAMING_SNAKE_CASE_						 = torch.manual_seed(0    )
									SCREAMING_SNAKE_CASE_						 = pipe.dual_guided(
									    prompt='''first prompt'''       ,       image=SCREAMING_SNAKE_CASE_       ,       text_to_image_strength=0.75       ,       generator=SCREAMING_SNAKE_CASE_       ,       guidance_scale=7.5       ,       num_inference_steps=2       ,       output_type='''numpy'''       ,       ).images

									with tempfile.TemporaryDirectory() as tmpdirname:
											pipe.save_pretrained(SCREAMING_SNAKE_CASE_    )
											SCREAMING_SNAKE_CASE_						 = VersatileDiffusionPipeline.from_pretrained(SCREAMING_SNAKE_CASE_       ,       torch_dtype=torch.floataa    )
									pipe.to(SCREAMING_SNAKE_CASE_    )
									pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_    )

									SCREAMING_SNAKE_CASE_						 = generator.manual_seed(0    )
									SCREAMING_SNAKE_CASE_						 = pipe.dual_guided(
									    prompt='''first prompt'''       ,       image=SCREAMING_SNAKE_CASE_       ,       text_to_image_strength=0.75       ,       generator=SCREAMING_SNAKE_CASE_       ,       guidance_scale=7.5       ,       num_inference_steps=2       ,       output_type='''numpy'''       ,       ).images

									assert np.abs(image - new_image    ).sum() < 1e-5, "Models don't have the same forward pass"







							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = VersatileDiffusionPipeline.from_pretrained('''shi-labs/versatile-diffusion'''       ,       torch_dtype=torch.floataa    )
									pipe.to(SCREAMING_SNAKE_CASE_    )
									pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_    )

									SCREAMING_SNAKE_CASE_						 = '''cyberpunk 2077'''
									SCREAMING_SNAKE_CASE_						 = load_image(
									    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'''    )
									SCREAMING_SNAKE_CASE_						 = torch.manual_seed(0    )
									SCREAMING_SNAKE_CASE_						 = pipe.dual_guided(
									    prompt=SCREAMING_SNAKE_CASE_       ,       image=SCREAMING_SNAKE_CASE_       ,       text_to_image_strength=0.75       ,       generator=SCREAMING_SNAKE_CASE_       ,       guidance_scale=7.5       ,       num_inference_steps=50       ,       output_type='''numpy'''       ,       ).images

									SCREAMING_SNAKE_CASE_						 = image[0, 2_53:2_56, 2_53:2_56, -1]

									assert image.shape == (1, 5_12, 5_12, 3)
									SCREAMING_SNAKE_CASE_						 = np.array([0.14_48, 0.16_19, 0.17_41, 0.10_86, 0.11_47, 0.11_28, 0.11_99, 0.11_65, 0.10_01]    )

									assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1e-1

									SCREAMING_SNAKE_CASE_						 = '''A painting of a squirrel eating a burger '''
									SCREAMING_SNAKE_CASE_						 = torch.manual_seed(0    )
									SCREAMING_SNAKE_CASE_						 = pipe.text_to_image(
									    prompt=SCREAMING_SNAKE_CASE_       ,       generator=SCREAMING_SNAKE_CASE_       ,       guidance_scale=7.5       ,       num_inference_steps=50       ,       output_type='''numpy'''    ).images

									SCREAMING_SNAKE_CASE_						 = image[0, 2_53:2_56, 2_53:2_56, -1]

									assert image.shape == (1, 5_12, 5_12, 3)
									SCREAMING_SNAKE_CASE_						 = np.array([0.33_67, 0.31_69, 0.26_56, 0.38_70, 0.47_90, 0.37_96, 0.40_09, 0.48_78, 0.47_78]    )

									assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1e-1

									SCREAMING_SNAKE_CASE_						 = pipe.image_variation(SCREAMING_SNAKE_CASE_       ,       generator=SCREAMING_SNAKE_CASE_       ,       output_type='''numpy'''    ).images

									SCREAMING_SNAKE_CASE_						 = image[0, 2_53:2_56, 2_53:2_56, -1]

									assert image.shape == (1, 5_12, 5_12, 3)
									SCREAMING_SNAKE_CASE_						 = np.array([0.30_76, 0.31_23, 0.32_84, 0.37_82, 0.37_70, 0.38_94, 0.42_97, 0.43_31, 0.44_56]    )

									assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1e-1