class FlaxStableDiffusionSafetyCheckerModule(nn.Module):
    config: CLIPConfig
    dtype: jnp.dtype = jnp.float32

    def setup(self):
        self.vision_model = FlaxCLIPVisionModule(self.config.vision_config)
        self.visual_projection = nn.Dense(self.config.projection_dim, use_bias=False, dtype=self.dtype)

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

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

    def __call__(self, clip_input):
        pooled_output = self.vision_model(clip_input)[1]
        image_embeds = self.visual_projection(pooled_output)

        special_cos_dist = jax_cosine_distance(image_embeds, self.special_care_embeds)
        cos_dist = jax_cosine_distance(image_embeds, self.concept_embeds)

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

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

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

        return has_nsfw_concepts
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/pipelines/stable_diffusion/safety_checker_flax.py