
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__			(  SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_=1e-12     )				->		Union[str, Any]:
			lowerCAmelCase__      :      Optional[int]         =		jnp.divide(emb_a.T       ,							jnp.clip(jnp.linalg.norm(SCREAMING_SNAKE_CASE_       ,							axis=1     )       ,							a_min=SCREAMING_SNAKE_CASE_     )     ).T
			lowerCAmelCase__      :      Tuple         =		jnp.divide(emb_a.T       ,							jnp.clip(jnp.linalg.norm(SCREAMING_SNAKE_CASE_       ,							axis=1     )       ,							a_min=SCREAMING_SNAKE_CASE_     )     ).T
			return jnp.matmul(SCREAMING_SNAKE_CASE_       ,							norm_emb_a.T     )







class      A__							(     nn.Module	):
				lowercase			    =  42
				lowercase			    =  jnp.floataa






				def 							_lowerCamelCase		(						self					:      Any   ):






							'''simple docstring'''







							lowerCAmelCase__      :      Optional[int]         =		FlaxCLIPVisionModule(self.config.vision_config   )
							lowerCAmelCase__      :      List[Any]         =		nn.Dense(self.config.projection_dim    ,     use_bias=a    ,     dtype=self.dtype   )

							lowerCAmelCase__      :      Union[str, Any]         =		self.param('concept_embeds'    ,     jax.nn.initializers.ones    ,     (17, self.config.projection_dim)   )
							lowerCAmelCase__      :      Optional[Any]         =		self.param(
							    'special_care_embeds'    ,     jax.nn.initializers.ones    ,     (3, self.config.projection_dim)   )

							lowerCAmelCase__      :      List[Any]         =		self.param('concept_embeds_weights'    ,     jax.nn.initializers.ones    ,     (17,)   )
							lowerCAmelCase__      :      str         =		self.param('special_care_embeds_weights'    ,     jax.nn.initializers.ones    ,     (3,)   )






				def __call__(						self					:      Dict    ,     a					:      Optional[Any]   ):






							'''simple docstring'''







							lowerCAmelCase__      :      Tuple         =		self.vision_model(a   )[1]
							lowerCAmelCase__      :      Dict         =		self.visual_projection(a   )

							lowerCAmelCase__      :      List[str]         =		jax_cosine_distance(a    ,     self.special_care_embeds   )
							lowerCAmelCase__      :      str         =		jax_cosine_distance(a    ,     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__      :      Tuple         =		0.0

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

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

							return has_nsfw_concepts







class      A__							(     __magic_name__	):
				lowercase			    =  CLIPConfig
				lowercase			    =  'clip_input'
				lowercase			    =  FlaxStableDiffusionSafetyCheckerModule






				def __init__(						self					:      Optional[Any]    ,     a					:      CLIPConfig    ,     a					:      Optional[Tuple] = None    ,     a					:      int = 0    ,     a					:      jnp.dtype = jnp.floataa    ,     a					:      bool = True    ,     **a					:      Any    ,     ):






							'''simple docstring'''







							if input_shape is None:
										lowerCAmelCase__      :      str         =		(1, 224, 224, 3)
							lowerCAmelCase__      :      Any         =		self.module_class(config=a    ,     dtype=a    ,     **a   )
							super().__init__(a    ,     a    ,     input_shape=a    ,     seed=a    ,     dtype=a    ,     _do_init=_do_init   )






				def 							_lowerCamelCase		(						self					:      Dict    ,     a					:      jax.random.KeyArray    ,     a					:      Tuple    ,     a					:      FrozenDict = None   ):






							'''simple docstring'''







							lowerCAmelCase__      :      str         =		jax.random.normal(a    ,     a   )

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

							lowerCAmelCase__      :      Any         =		self.module.init(a    ,     a   )['params']

							return random_params






				def __call__(						self					:      List[str]    ,     a					:      Dict    ,     a					:      dict = None    ,     ):






							'''simple docstring'''







							lowerCAmelCase__      :      Tuple         =		jnp.transpose(a    ,     (0, 2, 3, 1)   )

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

import datetime
import platform
import subprocess
from typing import Optional, Tuple, Union

import numpy as np
def  lowerCAmelCase__			(  SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_     )				->		np.array:
			lowerCAmelCase__      :      Dict         =		F'''{sampling_rate}'''
			lowerCAmelCase__      :      Any         =		'1'
			lowerCAmelCase__      :      Optional[Any]         =		'f32le'
			lowerCAmelCase__      :      Any         =		[
			    'ffmpeg',
			    '-i',
			    'pipe:0',
			    '-ac',
			    ac,
			    '-ar',
			    ar,
			    '-f',
			    format_for_conversion,
			    '-hide_banner',
			    '-loglevel',
			    'quiet',
			    'pipe:1',
			]

			try:
						with subprocess.Popen(SCREAMING_SNAKE_CASE_       ,							stdin=subprocess.PIPE       ,							stdout=subprocess.PIPE     ) as ffmpeg_process:
									lowerCAmelCase__      :      List[Any]         =		ffmpeg_process.communicate(SCREAMING_SNAKE_CASE_     )
			except FileNotFoundError as error:
						raise ValueError('ffmpeg was not found but is required to load audio files from filename'     ) from error
			lowerCAmelCase__      :      List[str]         =		output_stream[0]
			lowerCAmelCase__      :      str         =		np.frombuffer(SCREAMING_SNAKE_CASE_       ,							np.floataa     )
			if audio.shape[0] == 0:
						raise ValueError('Malformed soundfile'     )
			return audio
def  lowerCAmelCase__			(  SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_ = "f32le"       ,							)				->		Dict:
			lowerCAmelCase__      :      Optional[Any]         =		F'''{sampling_rate}'''
			lowerCAmelCase__      :      Any         =		'1'
			if format_for_conversion == "s16le":
						lowerCAmelCase__      :      Dict         =		2
			elif format_for_conversion == "f32le":
						lowerCAmelCase__      :      List[str]         =		4
			else:
						raise ValueError(F'''Unhandled format `{format_for_conversion}`. Please use `s16le` or `f32le`'''     )

			lowerCAmelCase__      :      Tuple         =		platform.system()
			if system == "Linux":
						lowerCAmelCase__      :      str         =		'alsa'
						lowerCAmelCase__      :      str         =		'default'
			elif system == "Darwin":
						lowerCAmelCase__      :      Any         =		'avfoundation'
						lowerCAmelCase__      :      Tuple         =		':0'
			elif system == "Windows":
						lowerCAmelCase__      :      Any         =		'dshow'
						lowerCAmelCase__      :      int         =		'default'

			lowerCAmelCase__      :      Any         =		[
			    'ffmpeg',
			    '-f',
			    format_,
			    '-i',
			    input_,
			    '-ac',
			    ac,
			    '-ar',
			    ar,
			    '-f',
			    format_for_conversion,
			    '-fflags',
			    'nobuffer',
			    '-hide_banner',
			    '-loglevel',
			    'quiet',
			    'pipe:1',
			]
			lowerCAmelCase__      :      List[str]         =		int(round(sampling_rate * chunk_length_s     )     ) * size_of_sample
			lowerCAmelCase__      :      str         =		_ffmpeg_stream(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_     )
			for item in iterator:
						yield item
def  lowerCAmelCase__			(  SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_ = None       ,							SCREAMING_SNAKE_CASE_ = None       ,							SCREAMING_SNAKE_CASE_ = "f32le"       ,							)				->		str:
			if stream_chunk_s is not None:
						lowerCAmelCase__      :      Union[str, Any]         =		stream_chunk_s
			else:
						lowerCAmelCase__      :      Tuple         =		chunk_length_s

			lowerCAmelCase__      :      Any         =		ffmpeg_microphone(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							format_for_conversion=SCREAMING_SNAKE_CASE_     )
			if format_for_conversion == "s16le":
						lowerCAmelCase__      :      Optional[Any]         =		np.intaa
						lowerCAmelCase__      :      Optional[Any]         =		2
			elif format_for_conversion == "f32le":
						lowerCAmelCase__      :      Optional[Any]         =		np.floataa
						lowerCAmelCase__      :      Optional[Any]         =		4
			else:
						raise ValueError(F'''Unhandled format `{format_for_conversion}`. Please use `s16le` or `f32le`'''     )

			if stride_length_s is None:
						lowerCAmelCase__      :      Dict         =		chunk_length_s / 6
			lowerCAmelCase__      :      int         =		int(round(sampling_rate * chunk_length_s     )     ) * size_of_sample
			if isinstance(SCREAMING_SNAKE_CASE_       ,							(int, float)     ):
						lowerCAmelCase__      :      Dict         =		[stride_length_s, stride_length_s]

			lowerCAmelCase__      :      Union[str, Any]         =		int(round(sampling_rate * stride_length_s[0]     )     ) * size_of_sample
			lowerCAmelCase__      :      List[Any]         =		int(round(sampling_rate * stride_length_s[1]     )     ) * size_of_sample
			lowerCAmelCase__      :      Any         =		datetime.datetime.now()
			lowerCAmelCase__      :      Any         =		datetime.timedelta(seconds=SCREAMING_SNAKE_CASE_     )
			for item in chunk_bytes_iter(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							stride=(stride_left, stride_right)       ,							stream=SCREAMING_SNAKE_CASE_     ):
						# Put everything back in numpy scale
						lowerCAmelCase__      :      Any         =		np.frombuffer(item['raw']       ,							dtype=SCREAMING_SNAKE_CASE_     )
						lowerCAmelCase__      :      Optional[Any]         =		(
						    item['stride'][0] // size_of_sample,
						    item['stride'][1] // size_of_sample,
						)
						lowerCAmelCase__      :      Optional[int]         =		sampling_rate
						audio_time += delta
						if datetime.datetime.now() > audio_time + 10 * delta:
									# We're late !! SKIP
									continue
						yield item
def  lowerCAmelCase__			(  SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_ = False     )				->		Optional[int]:
			lowerCAmelCase__      :      Union[str, Any]         =		b''
			lowerCAmelCase__      ,				lowerCAmelCase__      :      Union[str, Any]         =		stride
			if stride_left + stride_right >= chunk_len:
						raise ValueError(
						    F'''Stride needs to be strictly smaller than chunk_len: ({stride_left}, {stride_right}) vs {chunk_len}'''     )
			lowerCAmelCase__      :      List[str]         =		0
			for raw in iterator:
						acc += raw
						if stream and len(SCREAMING_SNAKE_CASE_     ) < chunk_len:
									lowerCAmelCase__      :      Tuple         =		(_stride_left, 0)
									yield {"raw": acc[:chunk_len], "stride": stride, "partial": True}
						else:
									while len(SCREAMING_SNAKE_CASE_     ) >= chunk_len:
												# We are flushing the accumulator
												lowerCAmelCase__      :      Dict         =		(_stride_left, stride_right)
												lowerCAmelCase__      :      Any         =		{'raw': acc[:chunk_len], 'stride': stride}
												if stream:
															lowerCAmelCase__      :      Optional[int]         =		False
												yield item
												lowerCAmelCase__      :      Optional[int]         =		stride_left
												lowerCAmelCase__      :      Optional[int]         =		acc[chunk_len - stride_left - stride_right :]
    # Last chunk
			if len(SCREAMING_SNAKE_CASE_     ) > stride_left:
						lowerCAmelCase__      :      Tuple         =		{'raw': acc, 'stride': (_stride_left, 0)}
						if stream:
									lowerCAmelCase__      :      Any         =		False
						yield item







def  lowerCAmelCase__			(  SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_     )				->		str:
			lowerCAmelCase__      :      int         =		2**24  # 16Mo
			try:
						with subprocess.Popen(SCREAMING_SNAKE_CASE_       ,							stdout=subprocess.PIPE       ,							bufsize=SCREAMING_SNAKE_CASE_     ) as ffmpeg_process:
									while True:
												lowerCAmelCase__      :      List[str]         =		ffmpeg_process.stdout.read(SCREAMING_SNAKE_CASE_     )
												if raw == b"":
															break
												yield raw
			except FileNotFoundError as error:
						raise ValueError('ffmpeg was not found but is required to stream audio files from filename'     ) from error