


'''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 							__magic_name__(						lowerCamelCase,				lowerCamelCase,				lowerCamelCase=1E-12):
						__lowerCAmelCase							       =				jnp.divide(emb_a.T,				jnp.clip(jnp.linalg.norm(lowerCamelCase,				axis=1),				a_min=lowerCamelCase)).T
						__lowerCAmelCase							       =				jnp.divide(emb_a.T,				jnp.clip(jnp.linalg.norm(lowerCamelCase,				axis=1),				a_min=lowerCamelCase)).T
						return jnp.matmul(lowerCamelCase,				norm_emb_a.T)


class       a__				( nn.Module      ):







				"""simple docstring"""



				__UpperCamelCase    :		CLIPConfig
				__UpperCamelCase    :		jnp.dtype	=	jnp.floataa




				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				FlaxCLIPVisionModule(self.config.vision_config		)
										__lowerCAmelCase							       =				nn.Dense(self.config.projection_dim       ,			use_bias=__lowercase       ,			dtype=self.dtype		)

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

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




				def __call__(self       ,			__lowercase		):
										__lowerCAmelCase							       =				self.vision_model(__lowercase		)[1]
										__lowerCAmelCase							       =				self.visual_projection(__lowercase		)

										__lowerCAmelCase							       =				jax_cosine_distance(__lowercase       ,			self.special_care_embeds		)
										__lowerCAmelCase							       =				jax_cosine_distance(__lowercase       ,			self.concept_embeds		)

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

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

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

										return has_nsfw_concepts


class       a__				( __A      ):







				"""simple docstring"""



				__UpperCamelCase    :		List[str]	=	CLIPConfig
				__UpperCamelCase    :		Optional[Any]	=	'clip_input'
				__UpperCamelCase    :		Any	=	FlaxStableDiffusionSafetyCheckerModule




				def __init__(self       ,			__lowercase       ,			__lowercase = None       ,			__lowercase = 0       ,			__lowercase = jnp.floataa       ,			__lowercase = True       ,			**__lowercase       ,			):
										if input_shape is None:
																__lowerCAmelCase							       =				(1, 2_24, 2_24, 3)
										__lowerCAmelCase							       =				self.module_class(config=__lowercase       ,			dtype=__lowercase       ,			**__lowercase		)
										super().__init__(__lowercase       ,			__lowercase       ,			input_shape=__lowercase       ,			seed=__lowercase       ,			dtype=__lowercase       ,			_do_init=_do_init		)




				def 					_snake_case  (self       ,			__lowercase       ,			__lowercase       ,			__lowercase = None		):
										# init input tensor
										__lowerCAmelCase							       =				jax.random.normal(__lowercase       ,			__lowercase		)

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

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

										return random_params




				def __call__(self       ,			__lowercase       ,			__lowercase = None       ,			):
										__lowerCAmelCase							       =				jnp.transpose(__lowercase       ,			(0, 2, 3, 1)		)

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





'''simple docstring'''

from PIL import Image

def 							__magic_name__(						lowerCamelCase,				lowerCamelCase):

						def brightness(lowerCamelCase) -> float:
												return 1_2_8 + level + (c - 1_2_8)

						if not -2_55.0 <= level <= 2_55.0:
												raise ValueError('''level must be between -255.0 (black) and 255.0 (white)''')
						return img.point(lowerCamelCase)


if __name__ == "__main__":
	# Load image
	with Image.open("""image_data/lena.jpg""") as img:
		# Change brightness to 100
		_UpperCAmelCase				: str							    =				change_brightness(img, 1_0_0)
		brigt_img.save("""image_data/lena_brightness.png""", format="""png""")

