Flax safety checker #825

pcuenca · 2022-10-13T12:22:57Z

I managed to change the pipeline to run the slow portion of the safety checker in pmap mode. The way it works is that __call__ now invokes a generate function first, and then computes the safety scores. Both generate and get_safety_scores are explicitly pmapped inside __call__, and therefore arguments to __call__ are sharded too.

This is how it works from a user's point of view:

from diffusers import FlaxStableDiffusionPipeline

pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(
    "/home/pedro/code/diffusers/sd-v1-4-flax",
    dtype=jnp.bfloat16
)

prompt = "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of field, close up, split lighting, cinematic"
prompt = [prompt] * jax.device_count()
prompt_ids = pipeline.prepare_inputs(prompt)

prng_seed = jax.random.PRNGKey(0)

# Replication done by the pipeline
output = pipeline(prompt_ids, params, prng_seed)

This breaks the test recently added by @patrickvonplaten.

Please, let me know if this is acceptable and I'll fix the test and finalize a couple of TODOs.

src/diffusers/pipeline_flax_utils.py

src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py

HuggingFaceDocBuilderDev · 2022-10-13T12:57:32Z

The documentation is not available anymore as the PR was closed or merged.

We could have decorated `generate` with `pmap`, but I wanted to keep it in case someone wants to invoke it in non-parallel mode.

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

…s into flax-safety-checker

src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py

patrickvonplaten

Super cool that you got it to work @pcuenca !

In my opinion in the long run we could / should strive for a solution that would allow us to wrap everything into a single pmap(...) call, e.g. not having to call pmap(...) from inside.

For today's release I think it's totally fine, however! Just two things I'd advocate for to change:

Let's move all sharding, replicate and random.split outside of generate => the user should have control over these (it's also safer in terms of backwards comp). Only internal shard/replicate should be called with if-statements if something like if jit and jax.device_count > 1 => then unshard/shard in def run_safety_checker`
Also I'd advocate to not run pmap by default but only if a flag called jit=False is passed via jit=True by the user because it fits better with JAX API (e.g. things are not jitted by default). In the long run we will then remove this jit=True flag and change the internals so that one can pmap(...) the whole function end-to-end.
Let's make generate and run_safety_checker private methods

pcuenca · 2022-10-13T13:11:02Z

Super cool that you got it to work @pcuenca !

In my opinion in the long run we could / should strive for a solution that would allow us to wrap everything into a single pmap(...) call, e.g. not having to call pmap(...) from inside.

For today's release I think it's totally fine, however! Just two things I'd advocate for to change:

Let's move all sharding, replicate and random.split outside of generate => the user should have control over these (it's also safer in terms of backwards comp)

Also I'd advocate to not run pmap by default but only if a flag called jit=False is passed via jit=True by the user because it fits better with JAX API (e.g. things are not jitted by default). In the long run we will then remove this jit=True flag and change the internals so that one can pmap(...) the whole function end-to-end.

Sounds great, thanks a lot for the fast review!

Making it work in a single function was much harder, so I opted for this intermediate solution. This way we don't require users to run the two steps themselves (generation and safety checker). Totally agree that we should try to wrap everything inside a single pmap call.

I also agree with the other comments, it didn't feel right to take decision to use pmap on our own. Thanks!

src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py

@kashif

As suggested by @kashif.

patil-suraj

Looks good, very cool !

patil-suraj · 2022-10-13T14:51:28Z

src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py

+            special_cos_dist, cos_dist = _p_get_safety_scores(self, features, safety_model_params)
+            special_cos_dist = unshard(special_cos_dist)
+            cos_dist = unshard(cos_dist)
+            safety_model_params = unreplicate(safety_model_params)


why do we need to do unreplicate here ?

Because if we are using jit, safety_model_params is extracted from the params dict which is already replicated. We use the replicated version in _p_get_safety_scores a couple of lines above, but then we need the unreplicated one to compute the scores in self.safety_checker.filtered_with_scores

patil-suraj · 2022-10-13T14:56:29Z

src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py

+# TODO: maybe use a config dict instead of so many static argnums
+@partial(jax.pmap, static_broadcasted_argnums=(0, 4, 5, 6, 7, 9))
+def _p_generate(
+    pipe, prompt_ids, params, prng_seed, num_inference_steps, height, width, guidance_scale, latents, debug
+):
+    return pipe._generate(
+        prompt_ids, params, prng_seed, num_inference_steps, height, width, guidance_scale, latents, debug
+    )
+
+
+@partial(jax.pmap, static_broadcasted_argnums=(0,))
+def _p_get_safety_scores(pipe, features, params):
+    return pipe._get_safety_scores(features, params)


(nit)

maybe have this as pipeline methods.

patil-suraj · 2022-10-13T14:56:55Z

src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py


-        return FlaxStableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
+def unshard(x: jnp.ndarray):


Let's maybe also make it private

@kashif

* Remove set_format in Flax pipeline. * Remove DummyChecker. * Run safety_checker in pipeline. * Don't pmap on every call. We could have decorated `generate` with `pmap`, but I wanted to keep it in case someone wants to invoke it in non-parallel mode. * Remove commented line Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com> * Replicate outside __call__, prepare for optional jitting. * Remove unnecessary clipping. As suggested by @kashif. * Do not jit unless requested. * Send all args to generate. * make style * Remove unused imports. * Fix docstring. Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

@kashif

* Remove set_format in Flax pipeline. * Remove DummyChecker. * Run safety_checker in pipeline. * Don't pmap on every call. We could have decorated `generate` with `pmap`, but I wanted to keep it in case someone wants to invoke it in non-parallel mode. * Remove commented line Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com> * Replicate outside __call__, prepare for optional jitting. * Remove unnecessary clipping. As suggested by @kashif. * Do not jit unless requested. * Send all args to generate. * make style * Remove unused imports. * Fix docstring. Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

pcuenca and others added 3 commits October 12, 2022 13:23

Remove set_format in Flax pipeline.

3c838a2

Remove DummyChecker.

a444010

Run safety_checker in pipeline.

3ca68c4

pcuenca requested review from patrickvonplaten and patil-suraj October 13, 2022 12:23

Merge branch 'main' into flax-safety-checker

9d84107

patrickvonplaten reviewed Oct 13, 2022

View reviewed changes

src/diffusers/pipeline_flax_utils.py Show resolved Hide resolved

patrickvonplaten reviewed Oct 13, 2022

View reviewed changes

src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py Outdated Show resolved Hide resolved

kashif reviewed Oct 13, 2022

View reviewed changes

src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py Outdated Show resolved Hide resolved

pcuenca and others added 3 commits October 13, 2022 12:57

Don't pmap on every call.

750e20f

We could have decorated `generate` with `pmap`, but I wanted to keep it in case someone wants to invoke it in non-parallel mode.

Remove commented line

dcd27fd

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

Merge branch 'flax-safety-checker' of github.com:huggingface/diffuser…

a0680ed

…s into flax-safety-checker

patrickvonplaten reviewed Oct 13, 2022

View reviewed changes

src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py Outdated Show resolved Hide resolved

patrickvonplaten reviewed Oct 13, 2022

View reviewed changes

patil-suraj reviewed Oct 13, 2022

View reviewed changes

pcuenca added 6 commits October 13, 2022 13:50

Replicate outside __call__, prepare for optional jitting.

d65d1a2

Remove unnecessary clipping.

4239bff

As suggested by @kashif.

Do not jit unless requested.

866600b

Send all args to generate.

86cb5a1

Merge remote-tracking branch 'origin/main' into flax-safety-checker

1cd8bb5

make style

fe2817b

pcuenca requested review from kashif, patrickvonplaten and patil-suraj October 13, 2022 14:25

pcuenca added 2 commits October 13, 2022 14:27

Remove unused imports.

b255e9a

Fix docstring.

2533c50

patil-suraj approved these changes Oct 13, 2022

View reviewed changes

patrickvonplaten approved these changes Oct 13, 2022

View reviewed changes

patrickvonplaten merged commit 78db11d into main Oct 13, 2022

patil-suraj deleted the flax-safety-checker branch October 13, 2022 15:05

patrickvonplaten mentioned this pull request Mar 13, 2023

[Community] Make safety model end-to-end compileable - Inference time of JAX / Flax pipeline #927

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Flax safety checker #825

Flax safety checker #825

pcuenca commented Oct 13, 2022

HuggingFaceDocBuilderDev commented Oct 13, 2022 •

edited

Loading

patrickvonplaten left a comment •

edited

Loading

pcuenca commented Oct 13, 2022

patil-suraj left a comment

patil-suraj Oct 13, 2022

pcuenca Oct 13, 2022

patil-suraj Oct 13, 2022

patil-suraj Oct 13, 2022


		return FlaxStableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
		def unshard(x: jnp.ndarray):

Flax safety checker #825

Flax safety checker #825

Conversation

pcuenca commented Oct 13, 2022

HuggingFaceDocBuilderDev commented Oct 13, 2022 • edited Loading

patrickvonplaten left a comment • edited Loading

Choose a reason for hiding this comment

pcuenca commented Oct 13, 2022

patil-suraj left a comment

Choose a reason for hiding this comment

patil-suraj Oct 13, 2022

Choose a reason for hiding this comment

pcuenca Oct 13, 2022

Choose a reason for hiding this comment

patil-suraj Oct 13, 2022

Choose a reason for hiding this comment

patil-suraj Oct 13, 2022

Choose a reason for hiding this comment

HuggingFaceDocBuilderDev commented Oct 13, 2022 •

edited

Loading

patrickvonplaten left a comment •

edited

Loading