FlaxDiffusionPipeline & FlaxStableDiffusionPipeline #559

mishig25 · 2022-09-19T09:32:36Z

Implement FlaxDiffusionPipeline & FlaxStableDiffusionPipeline

Based on https://github.com/patil-suraj/stable-diffusion-jax/blob/stateless-scheduler/stable_diffusion_jax/pipeline_stable_diffusion.py

From this comment, we have decided to create FlaxDiffusionPipeline, rather than try to reuse DiffusionPipeline (which currently handles pytorch & onnx).

Design decisions & questions

Just like DiffusionPipeline does not inherit from torch.nn.Module, FlaxDiffusionPipeline should not inherit from flax.linen.Module either. Wdyt?
Every pipeline (for example, FlaxStableDiffusionPipeline ) needs to implement InferenceState(flax.struct.dataclass) so that pmap can consume the pipeline. Wdyt?
If the first two points above holds, then the implementation of DiffusionPipeline & FlaxDiffusionPipeline are quite similar except with one major difference. Since flax pretrained models are initialized as model, params = xyz.from_pretrained(), FlaxDiffusionPipeline from_pretrained, save_pretrained methods needs to handle inference_state. See example here

TODOS:

handle all the TODO comments I left in the implementation
test the entire pipeline

HuggingFaceDocBuilderDev · 2022-09-19T09:35:51Z

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

kashif · 2022-09-19T09:55:17Z

@mishig25 I believe we will need to use the FlaxDDIMScheduler, ...

pcuenca · 2022-09-19T10:17:23Z

Assumptions 1 and 2 sound reasonable to me.

Regarding 3, if you want to override, say, the scheduler, then you'd need to do something like this:

scheduler, scheduler_state = SomeFlaxScheduler.from_pretrained(...)
inference_state = InferenceState(scheduler_state=scheduler_state)
pipe = FlaxStableDiffusionPipeline.from_pretrained(
    model_path,
    scheduler=scheduler,
    inference_state=inference_state,
)

Is that correct?

If so, my first instinct would be to return the final InferenceState too. However I haven't played with the code yet or made myself familiar with it.

patrickvonplaten · 2022-09-19T15:13:57Z

src/diffusers/models/__init__.py

@@ -14,4 +14,5 @@

 from .unet_2d import UNet2DModel
 from .unet_2d_condition import UNet2DConditionModel
+from .unet_2d_condition_flax import FlaxUNet2DConditionModel


We need to wrap this into a if flax_available_... statement I think

patrickvonplaten · 2022-09-19T15:14:57Z

src/diffusers/pipeline_flax_utils.py

+from .utils import CONFIG_NAME, DIFFUSERS_CACHE, BaseOutput, logging
+
+
+INDEX_FILE = "diffusion_flax_model.bin"


Actually this is never used (just like in PyTorch) we can remove it I think

into flax_pipeline

patrickvonplaten · 2022-09-19T20:47:48Z

src/diffusers/models/__init__.py

@@ -14,4 +14,6 @@

 from .unet_2d import UNet2DModel


these should be wrapped into is_available(...)

patrickvonplaten · 2022-09-19T20:48:20Z

src/diffusers/pipelines/stable_diffusion/__init__.py

+
+
+@flax.struct.dataclass
+class InferenceState:


this should be removed - let's just make it an inference state

This could potentially be helpful to override pipeline modules, as in my code snippet above #559 (comment).

We can do the same with a dictionary, but it's uglier in my opinion. Or with a helper function that returns a dict.

For now I think it can just be a dict no? dicts are more universal and it means that not every pipeline has to have a data class state

Removing for now -> let's maybe add later again if necessary

patrickvonplaten · 2022-09-19T20:49:10Z

src/diffusers/pipeline_flax_utils.py

+                    params[name] = loaded_params
+                elif issubclass(class_obj, SchedulerMixin):
+                    loaded_sub_model = load_method(loadable_folder, **loading_kwargs)
+                    params[name] = loaded_sub_model.create_state()


@pcuenca @kashif @patil-suraj this means that every flax scheduler needs a create_state() function, but I think the design is ok/makes sense

i think it does, I don't see a problem with this approach.

patrickvonplaten · 2022-09-19T20:50:19Z

src/diffusers/schedulers/scheduling_ddim_flax.py

-        self.alphas_cumprod = jnp.cumprod(self.alphas, axis=0)
+
+        # HACK for now - clean up later (PVP)
+        self._alphas_cumprod = jnp.cumprod(self.alphas, axis=0)


that's a hack for now - in the future IMO we should move all this logic to create_state so that the scheduler is fully stateless

patrickvonplaten · 2022-09-19T20:50:36Z

src/diffusers/schedulers/scheduling_ddim_flax.py

        beta_prod_t = 1 - alpha_prod_t

        # 3. compute predicted original sample from predicted noise also called
        # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
        pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)

-        # 4. Clip "predicted x_0"


Let's for now remove all the unimportant things

patrickvonplaten · 2022-09-19T20:50:51Z

src/diffusers/schedulers/scheduling_pndm_flax.py

@@ -148,7 +148,8 @@ def __init__(
        # mainly at formula (9), (12), (13) and the Algorithm 2.
        self.pndm_order = 4

-        self.state = PNDMSchedulerState.create(num_train_timesteps=num_train_timesteps)
+    def create_state(self):


cc @pcuenca @kashif what do you think about the design?

having a look thanks! I am also adding the scheduler tests too so this will be helpful I believe

patrickvonplaten · 2022-09-19T20:51:07Z

src/diffusers/schedulers/scheduling_ddim_flax.py

-        alpha_prod_t = self.alphas_cumprod[timestep]
-        alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
+        alpha_prod_t = alphas_cumprod[timestep]
+        alpha_prod_t_prev = jnp.where(prev_timestep >= 0, alphas_cumprod[prev_timestep], self.final_alpha_cumprod)


As @pcuenca mentioned we need jnp.where methods everywhere

yup! onto it

patrickvonplaten · 2022-09-19T20:53:32Z

For now the API looks as follows (on a TPUv3-8)

#!/usr/bin/env python3
from diffusers import FlaxStableDiffusionPipeline
from jax import pmap
import numpy as np
import jax
from flax.jax_utils import replicate
from flax.training.common_utils import shard

pipeline, params = FlaxStableDiffusionPipeline.from_pretrained("fusing/sd-v1-4-flax", use_auth_token=True)

prompt = "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of field, close up, split lighting, cinematic"
prng_seed = jax.random.PRNGKey(0)

p_sample = pmap(pipeline.__call__, static_broadcasted_argnums=(3,))

# shard inputs and rng
params = replicate(params)
prng_seed = jax.random.split(prng_seed, 8)
num_samples = jax.device_count()
prompt = num_samples * [prompt]
prompt_ids = pipeline.prepare_prompts(prompt)
prompt_ids = shard(prompt_ids)

# set inference steps
num_inference_steps = 50

images = p_sample(prompt_ids, params, prng_seed, num_inference_steps).images

images_pil = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:])))
# Problem: resulting images don't look good

@patil-suraj , @pcuenca , @kashif , @mishig25 very keen to get your feedback on the API.

Note: We need to pass tensors into the forward call which is why I've added a prepare_prompts function.
Note: Right now the pipeline generates incorrect images and needs debugging.

pcuenca · 2022-09-20T06:01:39Z

The pipeline runs out of memory in v2-8. Using dtype=jnp.bfloat16 still loads everything in float32. I think this is in part because of #565: the dtype is passed as part of the kwargs but then gets ignored. I only wanted to exclude it when saving the configuration, as it couldn't be serialized, but it's always ignored on load even when we want to override it.

What would be the best way to deal with this?

(This comment applies to the model instance, params are also loaded in float32 and have to be converted if necessary)

pcuenca · 2022-09-20T08:36:37Z

Questions about schedulers (and overridden pipeline modules):

If the user provides their own scheduler to pipeline from_pretrained, how is the state going to be handled and added to the params dict?
a. Invoke scheduler.create_state() inside from_pretrained anyway.
b. Have the user pass it to the pipeline using a new kw arg called scheduler_params.
c. Let them provide a dictionary with params for all overridden modules.
d. Go back to using the InferenceState so we can pass it instead of a dictionary.
PNDM requires the latents shape in set_timesteps. This is because it reserves space in the state for the 4 samples that are used in the step computations. Perhaps there’s a better way to resolve it without this information, but assuming there isn’t, what should we do?
a. Always send the shape to all schedulers when invoking set_timesteps() and let them ignore it if they don’t have a use for it. This is a departure from the PyTorch version, but we already had to add the state argument anyway.
b. Only send it in specific cases, checking types or signatures.

For now I’m going with 1.c and 2.a, would love to hear other opinions.

src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py

patil-suraj · 2022-09-20T09:39:07Z

The pipeline runs out of memory in v2-8. Using dtype=jnp.bfloat16 still loads everything in float32. I think this is in part because of #565: the dtype is passed as part of the kwargs but then gets ignored. I only wanted to exclude it when saving the configuration, as it couldn't be serialized, but it's always ignored on load even when we want to override it.

What would be the best way to deal with this?

(This comment applies to the model instance, params are also loaded in float32 and have to be converted if necessary)

@pcuenca the dtype only specifies the dtype of computation and not of params , so all params will be loaded in fp32 by default. It's analogues to with autocast("cuda") in PT.

We can create fp16/bf16 branch for flax weights, the same way we do in pt , so params get loaded in that dtype.

pcuenca · 2022-09-20T10:23:48Z

The pipeline runs out of memory in v2-8. Using dtype=jnp.bfloat16 still loads everything in float32. I think this is in part because of #565: the dtype is passed as part of the kwargs but then gets ignored. I only wanted to exclude it when saving the configuration, as it couldn't be serialized, but it's always ignored on load even when we want to override it.
What would be the best way to deal with this?
(This comment applies to the model instance, params are also loaded in float32 and have to be converted if necessary)

@pcuenca the dtype only specifies the dtype of computation and not of params , so all params will be loaded in fp32 by default. It's analogues to with autocast("cuda") in PT.

We can create fp16/bf16 branch for flax weights, the same way we do in pt , so params get loaded in that dtype.

You are right. In that case we could convert the weights to bfloat16 externally to save memory; for example in the notebook demo. Or save a specific branch as you propose.

However the pipeline in your repo @patil-suraj runs fine in v2-8 TPUs (I've been running an inference backend since last Friday, including the safety checker). Any idea why this one doesn't fit?

mishig25 · 2022-09-20T15:15:42Z

src/diffusers/pipeline_flax_utils.py

+        init_kwargs = {}
+
+        # inference_params
+        params = {}


does it not have to be special flax data structure to avoid memory fragmentation and other issues with pmap?
or it being regular python dict is okay?
cc: @patil-suraj @patrickvonplaten @pcuenca

I don't know tbh, very interesting if that's the case! Do you have a reference?

I was in an assumption that everytime you shard a data, it needs to be @flux.struct.xyz?
going over the docstring her

def shard(xs): """Helper for pmap to shard a pytree of arrays by local_device_count.

I guess params is a valid pytree since it is just a dict that contains valid pytree nodes. So it should be fine?

Yes, a dict is a valid pytree that can be sharded :)

Dict works - we could make it a frozen dict to be super certain! Will update to a frozen dict so that the pipeline is not allowed to change it internally -> then we're fully in Jaxistan

sounds perfect!

Will do this in a future PR

…flax_pipeline

into flax_pipeline

patrickvonplaten · 2022-09-20T19:25:09Z

src/diffusers/modeling_flax_utils.py

                if not os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME):
                    raise EnvironmentError(
                        f"Error no file named {WEIGHTS_NAME} found in directory {pretrained_model_name_or_path} "
                    )
                model_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
+            elif os.path.isfile(os.path.join(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME)):


@patil-suraj @younesbelkada we need to first check from_pt here in case there are both Flax and PT files, otherwise it breaks

patrickvonplaten · 2022-09-20T19:28:02Z

Pipeline can easily be checked with: #592

skirsten · 2022-10-24T20:49:42Z

Off-topic

@patil-suraj

We can create fp16/bf16 branch for flax weights, the same way we do in pt , so params get loaded in that dtype.

Considering there is only a bf16 and fp32 (flax) branch, how can I convert the model to fp16 for flax?
Probably using FlaxModelMixin.to_bf16 but I did not figure out how to apply it. Maybe somebody could upload the scripts used for the conversion to the script folder?

patil-suraj · 2022-10-25T10:59:41Z

There is a FlaxModelMixin.to_fp16 also, and we need to convert each individial model(unet, text_encoder etc) to fp16 and then save and load the pipeline.

For example:

pipe, params = FlaxStableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", revision="flax",
)
params["unet"] = pipe.unet.to_fp16(params["unet"])
params["vae"] = pipe.vae.to_fp16(params["vae"])
params["text_encoder"] = pipe.text_encoder.to_fp16(params["text_encoder"])
params["safety_checker"] = pipe.safety_checker.to_fp16(params["safety_checker"])

pipe.save_pretrained(pipe, params=params)

* WIP: flax FlaxDiffusionPipeline & FlaxStableDiffusionPipeline * todo comment * Fix imports * Fix imports * add dummies * Fix empty init * make pipeline work * up * Use Flax schedulers (typing, docstring) * Wrap model imports inside availability checks. * more updates * make sure flax is not broken * make style * more fixes * up Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com> Co-authored-by: Pedro Cuenca <pedro@latenitesoft.com>

mishig25 added 2 commits September 19, 2022 09:28

WIP: flax FlaxDiffusionPipeline & FlaxStableDiffusionPipeline

b9ca406

todo comment

30abc63

mishig25 changed the title ~~Flax_pipeline~~ WIP: flax FlaxDiffusionPipeline & FlaxStableDiffusionPipeline Sep 19, 2022

mishig25 requested review from kashif, patil-suraj, patrickvonplaten and pcuenca September 19, 2022 09:53

Merge branch 'main' into flax_pipeline

9b54559

patrickvonplaten reviewed Sep 19, 2022

View reviewed changes

mishig25 and others added 7 commits September 19, 2022 15:38

Fix imports

4b2becb

Fix imports

7f0e429

add dummies

d9e2ae1

Fix empty init

d51e881

Merge branch 'flax_pipeline' of https://github.com/huggingface/diffusers

741046d

into flax_pipeline

make pipeline work

7aab68d

merge conflict

7d3fff6

patrickvonplaten marked this pull request as ready for review September 19, 2022 20:45

up

47d7739

patrickvonplaten reviewed Sep 19, 2022

View reviewed changes

pcuenca mentioned this pull request Sep 20, 2022

Allow pipeline to run in bfloat16 #581

Closed

Use Flax schedulers (typing, docstring)

0c2a868

Wrap model imports inside availability checks.

69b1d7a

patil-suraj reviewed Sep 20, 2022

View reviewed changes

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

mishig25 commented Sep 20, 2022

View reviewed changes

patrickvonplaten added 6 commits September 20, 2022 16:37

Merge branch 'main' of https://github.com/huggingface/diffusers into …

0da4b8d

…flax_pipeline

more updates

82a5cf3

Merge branch 'flax_pipeline' of https://github.com/huggingface/diffusers

c00d98f

into flax_pipeline

make sure flax is not broken

2e9e523

make style

abb2250

more fixes

61342a2

patrickvonplaten changed the title ~~WIP: flax FlaxDiffusionPipeline & FlaxStableDiffusionPipeline~~ FlaxDiffusionPipeline & FlaxStableDiffusionPipeline Sep 20, 2022

up

182e485

patrickvonplaten reviewed Sep 20, 2022

View reviewed changes

patrickvonplaten merged commit d934d3d into main Sep 20, 2022

patrickvonplaten deleted the flax_pipeline branch September 20, 2022 19:30

PhaneeshB pushed a commit to nod-ai/diffusers that referenced this pull request Mar 1, 2023

remove gsutil_flags and fix download (huggingface#559)

6be5926

		from .utils import CONFIG_NAME, DIFFUSERS_CACHE, BaseOutput, logging


		INDEX_FILE = "diffusion_flax_model.bin"

FlaxDiffusionPipeline & FlaxStableDiffusionPipeline #559

FlaxDiffusionPipeline & FlaxStableDiffusionPipeline #559

Uh oh!

Conversation

mishig25 commented Sep 19, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Implement FlaxDiffusionPipeline & FlaxStableDiffusionPipeline

Design decisions & questions

Uh oh!

HuggingFaceDocBuilderDev commented Sep 19, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kashif commented Sep 19, 2022

Uh oh!

pcuenca commented Sep 19, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

patrickvonplaten commented Sep 19, 2022

Uh oh!

pcuenca commented Sep 20, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pcuenca commented Sep 20, 2022

Uh oh!

Uh oh!

patil-suraj commented Sep 20, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pcuenca commented Sep 20, 2022

Uh oh!

mishig25 Sep 20, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

mishig25 commented Sep 19, 2022 •

edited

Loading

HuggingFaceDocBuilderDev commented Sep 19, 2022 •

edited

Loading

pcuenca commented Sep 20, 2022 •

edited

Loading

patil-suraj commented Sep 20, 2022 •

edited

Loading

mishig25 Sep 20, 2022 •

edited

Loading