Safety Checker jitting

[pcuenca]

I did some changes to make it easier to jit the safety checker:

• Do not transpose (assumes a certain number of dimensions). The array should be prepared externally.
• Create a method (which is not jitted) to get the final flags from the distances.

Example use:

feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-large-patch14")
safety_checker, params = StableDiffusionSafetyCheckerModel.from_pretrained(
    "/sddata/sd-v1-4-flax/safety_checker", _do_init=False, dtype=dtype,
)

params = replicate(params)
unshard = lambda x: einops.rearrange(x, 'd b ... -> (d b) ...')

@jax.pmap
def get_safety_scores(features, params):
    special_cos_dist, cos_dist = safety_checker(features, params)
    return (special_cos_dist, cos_dist)

def run_safety_checker(images):
    pil_images = [Image.fromarray(image) for image in images]
    jnp_images = jnp.array(images)
    jnp_images = shard(jnp_images)
    features = feature_extractor(pil_images, return_tensors="np").pixel_values
    features = jnp.transpose(features, (0, 2, 3, 1))
    features = shard(features)
    special_cos_dist, cos_dist = get_safety_scores(features, params)
    images, has_nsfw = safety_checker.filtered_with_scores(
        unshard(special_cos_dist),
        unshard(cos_dist),
        images,
        params=unreplicate(params),
    )
    return images, has_nsfw

shape = (8, 512, 512, 3)
images = np.random.rand(*shape)
images = (images * 255).round().astype("uint8")
_, has_nsfw = run_safety_checker(images)