






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    UpperCamelCase						(	_A,   _A,   _A=1e-12			):






						"""simple docstring"""







						__magic_name__      :       Optional[int]       						= jnp.divide(emb_a.T,   jnp.clip(jnp.linalg.norm(_A,   axis=1			),   a_min=_A			)			).T
						__magic_name__      :       List[Any]       						= 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				):
	lowercase__						:					CLIPConfig
	lowercase__						:					jnp.dtype											=   jnp.floataa
	def       __magic_name__    (    self			)							->				Tuple:
							__magic_name__      :       Optional[Any]       						= FlaxCLIPVisionModule(self.config.vision_config			)
							__magic_name__      :       List[Any]       						= nn.Dense(self.config.projection_dim      ,       use_bias=lowerCAmelCase__      ,       dtype=self.dtype			)

							__magic_name__      :       Optional[Any]       						= self.param("""concept_embeds"""      ,       jax.nn.initializers.ones      ,       (17, self.config.projection_dim)			)
							__magic_name__      :       Optional[int]       						= self.param(
							    """special_care_embeds"""      ,       jax.nn.initializers.ones      ,       (3, self.config.projection_dim)			)

							__magic_name__      :       Optional[int]       						= self.param("""concept_embeds_weights"""      ,       jax.nn.initializers.ones      ,       (17,)			)
							__magic_name__      :       List[Any]       						= self.param("""special_care_embeds_weights"""      ,       jax.nn.initializers.ones      ,       (3,)			)







	def __call__(    self      ,       lowerCAmelCase__			)							->				Optional[Any]:
							__magic_name__      :       Union[str, Any]       						= self.vision_model(lowerCAmelCase__			)[1]
							__magic_name__      :       List[str]       						= self.visual_projection(lowerCAmelCase__			)

							__magic_name__      :       int       						= jax_cosine_distance(lowerCAmelCase__      ,       self.special_care_embeds			)
							__magic_name__      :       str       						= jax_cosine_distance(lowerCAmelCase__      ,       self.concept_embeds			)

							# increase this value to create a stronger `nfsw` filter
							# at the cost of increasing the possibility of filtering benign image inputs
							__magic_name__      :       Optional[Any]       						= 0.0

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

							__magic_name__      :       Optional[int]       						= cos_dist - self.concept_embeds_weights[None, :] + special_adjustment
							__magic_name__      :       Optional[Any]       						= jnp.round(lowerCAmelCase__      ,       3			)
							__magic_name__      :       Union[str, Any]       						= jnp.any(concept_scores > 0      ,       axis=1			)

							return has_nsfw_concepts





class 	snake_case__ (							_lowerCAmelCase				):
	lowercase__						:					Any											=   CLIPConfig
	lowercase__						:					List[str]											=   '''clip_input'''
	lowercase__						:					Tuple											=   FlaxStableDiffusionSafetyCheckerModule
	def __init__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__ = None      ,       lowerCAmelCase__ = 0      ,       lowerCAmelCase__ = jnp.floataa      ,       lowerCAmelCase__ = True      ,       **lowerCAmelCase__      ,       )							->				List[str]:
							if input_shape is None:
													__magic_name__      :       int       						= (1, 2_24, 2_24, 3)
							__magic_name__      :       List[str]       						= self.module_class(config=lowerCAmelCase__      ,       dtype=lowerCAmelCase__      ,       **lowerCAmelCase__			)
							super().__init__(lowerCAmelCase__      ,       lowerCAmelCase__      ,       input_shape=lowerCAmelCase__      ,       seed=lowerCAmelCase__      ,       dtype=lowerCAmelCase__      ,       _do_init=_do_init			)
	def       __magic_name__    (    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__      ,       lowerCAmelCase__ = None			)							->				FrozenDict:
							# init input tensor
							__magic_name__      :       Union[str, Any]       						= jax.random.normal(lowerCAmelCase__      ,       lowerCAmelCase__			)

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

							__magic_name__      :       List[str]       						= self.module.init(lowerCAmelCase__      ,       lowerCAmelCase__			)["""params"""]

							return random_params







	def __call__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__ = None      ,       )							->				Any:
							__magic_name__      :       Dict       						= jnp.transpose(lowerCAmelCase__      ,       (0, 2, 3, 1)			)

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








from manim import *





class 	snake_case__ (							_lowerCAmelCase				):
	def       __magic_name__    (    self			)							->				Dict:
							__magic_name__      :       int       						= Rectangle(height=0.5      ,       width=0.5			)
							__magic_name__      :       Optional[int]       						= Rectangle(height=0.2_5      ,       width=0.2_5			)
							__magic_name__      :       str       						= Rectangle(height=0.4_6      ,       width=0.4_6			).set_stroke(width=0			)

							__magic_name__      :       List[Any]       						= [mem.copy() for i in range(6			)]
							__magic_name__      :       int       						= [mem.copy() for i in range(6			)]
							__magic_name__      :       Tuple       						= VGroup(*lowerCAmelCase__			).arrange(lowerCAmelCase__      ,       buff=0			)
							__magic_name__      :       List[str]       						= VGroup(*lowerCAmelCase__			).arrange(lowerCAmelCase__      ,       buff=0			)
							__magic_name__      :       str       						= VGroup(lowerCAmelCase__      ,       lowerCAmelCase__			).arrange(lowerCAmelCase__      ,       buff=0			)
							__magic_name__      :       Union[str, Any]       						= Text("""CPU"""      ,       font_size=24			)
							__magic_name__      :       Tuple       						= Group(lowerCAmelCase__      ,       lowerCAmelCase__			).arrange(lowerCAmelCase__      ,       buff=0.5      ,       aligned_edge=lowerCAmelCase__			)
							cpu.move_to([-2.5, -0.5, 0]			)
							self.add(lowerCAmelCase__			)

							__magic_name__      :       Any       						= [mem.copy() for i in range(4			)]
							__magic_name__      :       List[Any]       						= VGroup(*lowerCAmelCase__			).arrange(lowerCAmelCase__      ,       buff=0			)
							__magic_name__      :       Tuple       						= Text("""GPU"""      ,       font_size=24			)
							__magic_name__      :       Tuple       						= Group(lowerCAmelCase__      ,       lowerCAmelCase__			).arrange(lowerCAmelCase__      ,       buff=0.5      ,       aligned_edge=lowerCAmelCase__			)
							gpu.move_to([-1, -1, 0]			)
							self.add(lowerCAmelCase__			)

							__magic_name__      :       Union[str, Any]       						= [mem.copy() for i in range(6			)]
							__magic_name__      :       Union[str, Any]       						= VGroup(*lowerCAmelCase__			).arrange(lowerCAmelCase__      ,       buff=0			)

							__magic_name__      :       str       						= Text("""Model"""      ,       font_size=24			)
							__magic_name__      :       Optional[int]       						= Group(lowerCAmelCase__      ,       lowerCAmelCase__			).arrange(lowerCAmelCase__      ,       buff=0.5      ,       aligned_edge=lowerCAmelCase__			)
							model.move_to([3, -1.0, 0]			)
							self.add(lowerCAmelCase__			)

							__magic_name__      :       str       						= []
							__magic_name__      :       Tuple       						= []
							__magic_name__      :       Union[str, Any]       						= []

							for i, rect in enumerate(lowerCAmelCase__			):
													rect.set_stroke(lowerCAmelCase__			)

													__magic_name__      :       Optional[Any]       						= Rectangle(height=0.4_6 / 4      ,       width=0.4_6 / 3			).set_stroke(width=0.0			).set_fill(lowerCAmelCase__      ,       opacity=0.7			)

													if i == 0:
																			cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT			)      ,       buff=0.0_2      ,       direction=lowerCAmelCase__			)
																			cpu_target.set_x(cpu_target.get_x() + 0.1			)
													elif i == 3:
																			cpu_target.next_to(model_cpu_arr[0]      ,       direction=lowerCAmelCase__      ,       buff=0.0			)
													else:
																			cpu_target.next_to(model_cpu_arr[i - 1]      ,       direction=lowerCAmelCase__      ,       buff=0.0			)
													self.add(lowerCAmelCase__			)
													model_cpu_arr.append(lowerCAmelCase__			)

							self.add(*lowerCAmelCase__      ,       *lowerCAmelCase__      ,       *lowerCAmelCase__			)

							__magic_name__      :       Optional[Any]       						= [mem.copy() for i in range(6			)]
							__magic_name__      :       Optional[Any]       						= VGroup(*lowerCAmelCase__			).arrange(lowerCAmelCase__      ,       buff=0			)

							__magic_name__      :       Any       						= Text("""Loaded Checkpoint"""      ,       font_size=24			)
							__magic_name__      :       Optional[int]       						= Group(lowerCAmelCase__      ,       lowerCAmelCase__			).arrange(lowerCAmelCase__      ,       buff=0.5      ,       aligned_edge=lowerCAmelCase__			)
							checkpoint.move_to([3, 0.5, 0]			)

							self.add(lowerCAmelCase__			)

							__magic_name__      :       Optional[int]       						= []
							__magic_name__      :       Tuple       						= []

							for i, rect in enumerate(lowerCAmelCase__			):
													__magic_name__      :       Dict       						= fill.copy().set_fill(lowerCAmelCase__      ,       opacity=0.7			)
													target.move_to(lowerCAmelCase__			)
													ckpt_arr.append(lowerCAmelCase__			)

													__magic_name__      :       int       						= target.copy()
													if i < 5:
																			cpu_target.move_to(cpu_left_col_base[i + 1]			)
													else:
																			cpu_target.move_to(cpu_right_col_base[i - 5]			)
													ckpt_cpu_arr.append(lowerCAmelCase__			)
							self.add(*lowerCAmelCase__      ,       *lowerCAmelCase__			)

							__magic_name__      :       Tuple       						= Square(side_length=2.2			)
							key.move_to([-5, 2, 0]			)

							__magic_name__      :       str       						= MarkupText(
							    F'<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'      ,       font_size=18      ,       )

							key_text.move_to([-5, 2.4, 0]			)

							self.add(lowerCAmelCase__      ,       lowerCAmelCase__			)

							__magic_name__      :       Any       						= MarkupText(
							    F'<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'      ,       font_size=18      ,       )

							blue_text.next_to(lowerCAmelCase__      ,       DOWN * 2.4      ,       aligned_edge=key_text.get_left()			)
							self.add(lowerCAmelCase__			)

							__magic_name__      :       Optional[Any]       						= MarkupText(
							    F'Based on the passed in configuration, weights are stored in\na variety of np.memmaps on disk or to a particular device.'      ,       font_size=24      ,       )
							step_a.move_to([2, 2, 0]			)

							__magic_name__      :       int       						= [meta_mem.copy() for i in range(6			)]
							__magic_name__      :       Union[str, Any]       						= [meta_mem.copy() for i in range(6			)]
							__magic_name__      :       Any       						= VGroup(*lowerCAmelCase__			).arrange(lowerCAmelCase__      ,       buff=0			)
							__magic_name__      :       str       						= VGroup(*lowerCAmelCase__			).arrange(lowerCAmelCase__      ,       buff=0			)
							__magic_name__      :       Tuple       						= VGroup(lowerCAmelCase__      ,       lowerCAmelCase__			).arrange(lowerCAmelCase__      ,       buff=0			)
							__magic_name__      :       int       						= Text("""Disk"""      ,       font_size=24			)
							__magic_name__      :       Union[str, Any]       						= Group(lowerCAmelCase__      ,       lowerCAmelCase__			).arrange(lowerCAmelCase__      ,       buff=0.5      ,       aligned_edge=lowerCAmelCase__			)
							disk.move_to([-4.0, -1.2_5, 0]			)
							self.play(Write(lowerCAmelCase__      ,       run_time=3			)      ,       Write(lowerCAmelCase__      ,       run_time=1			)      ,       Create(lowerCAmelCase__      ,       run_time=1			)			)

							__magic_name__      :       List[Any]       						= []
							for i, rect in enumerate(lowerCAmelCase__			):
													__magic_name__      :       Dict       						= rect.copy()
													target.generate_target()
													target.target.move_to(disk_left_col_base[i]			).scale(0.5			)
													animations.append(MoveToTarget(lowerCAmelCase__      ,       run_time=1.5			)			)
							self.play(*lowerCAmelCase__			)

							self.play(FadeOut(lowerCAmelCase__			)			)

							__magic_name__      :       str       						= MarkupText(F'Then, the checkpoint is removed from memory\nthrough garbage collection.'      ,       font_size=24			)
							step_a.move_to([2, 2, 0]			)

							self.play(Write(lowerCAmelCase__      ,       run_time=3			)			)

							self.play(
							    FadeOut(lowerCAmelCase__      ,       lowerCAmelCase__      ,       *lowerCAmelCase__      ,       *lowerCAmelCase__			)      ,       )

							self.wait()
