latest release breaks Flax pipes for Stable Diffusion when using bfloat16 and/or non-default scheduler

[krahnikblis]

Describe the bug

using Colab TPU environment, and bfloat16, blows up pipes if using anything but a packaged-built pipeline and its built-in default scheduler, even then it can unpredictably produce NaNs instead of images (can't reliably replicate the NaN thing, but it happens more often than not - see logs/output attached). the errors are often about data type mismatch, but in different places within the pipe or schedulers. one of the schedulers is missing its "key" argument in the pipe implementation, so that errors for that reason.

prior to recent release, i had a functional workbook which would build components separately, basically to make it easier to implement custom VAE and apply my LoRA-combinator model to the Unet and Text Encoder params. this now breaks 100% of the time (the usage of component-built pipe). however, a separate issue i logged re: img2img whitewashed outputs demonstrated that an all-in-one pipeline functions fine except for that issue, so i tried that as a fix in this case... in trying that way, outputs can be ok, but not reliably (NaNs issue) - in attached outputs, i was lucky enough to get image outputs in 3 out of 10 tests, but cannot reliably use the combo of scheduler + dtype that happened to work (because more NaNs)

tests: since the errors are consistently in the scheduler itself, or in the pipe's application/usage of scheduler, and since they regard the data type, the test reproduction script below is a loop through all Flax schedulers, and through either default or bfloat16 type when building them. the first time i tried this, i was using from_config for the schedulers (because there's no weights file, thus schedulers aren't "pretrained"); that was a whole mess of other problems, which were reduced when i switched to from_pretrained (again, kinda nonsense to use that method in this way). i used the JAX config to debug NaNs, and had to run the pipe with jit=False; it makes the run very slow, but that isolated the problem to this line in the pipe: latents, scheduler_state = self.scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple() which again is to do with the scheduler...

i also found that some of the combos with data type errors can be resolved by using custom latents as input to pipe, but in those cases, the outputs were all NaNs (swap one problem for another), so i didn't include that test method here. i'm also not sure what's the utility of using custom latents at all, and imo it's not something a user should need to [even be expected to know how to] do to defeat data type errors...

Reproduction

cell 1:

#@title bfloat + scheduler issue reproduction p1 (installs + maybe colab runtime restart)
!pip install -qq -U jax jaxlib
!pip install -qq flax optax transformers ftfy diffusers
### sometimes requires restart at this point, because google.

cell 2:

#@title bfloat + scheduler issue reproduction p2 (imports + setup)
# setup TPU
import jax.tools.colab_tpu
jax.tools.colab_tpu.setup_tpu()
# jax.tools.colab_tpu.setup_tpu(tpu_driver_version='tpu_driver_nightly')
import os
num_devices = jax.device_count()
device_type = jax.devices()[0].device_kind
print(f"Found {num_devices} JAX devices of type {device_type}.")
# imports
import PIL, jax, requests
import jax.numpy as jnp
import numpy as np
from diffusers import FlaxStableDiffusionPipeline, FlaxLMSDiscreteScheduler, FlaxDDIMScheduler, FlaxDPMSolverMultistepScheduler, FlaxPNDMScheduler, FlaxDDPMScheduler
from flax import jax_utils
# make a handy display fn:
def display_jnp_image(image):
    image = jnp.clip(np.array(image * 255),0,255).astype(np.uint8)
    if image.shape[0] == 3:
        image = np.transpose(image, axes=(1, 2, 0))
    image = PIL.Image.fromarray(np.array(image),mode="RGB")
    display(image)
# get pretrained
pretrained_model_name_or_path = "stable-diffusion-v1-5/"
if not os.path.exists(pretrained_model_name_or_path):
    !git lfs install 
    !git clone -b bf16 https://huggingface.co/runwayml/stable-diffusion-v1-5
# variables and pipe
dtype = jnp.bfloat16 
pipe, params = FlaxStableDiffusionPipeline.from_pretrained(pretrained_model_name_or_path, dtype=dtype)
orig_scheduler = pipe.scheduler._class_name
rng = jax.random.PRNGKey(42)
rng_repl = jax.random.split(rng,jax.device_count())
guidance_scale = jax_utils.replicate(jnp.array([7.5]))
prompt = "a pretty goose for you, and a pretty racing car"
prompt_ids = jax_utils.replicate(pipe.tokenizer(prompt,padding="max_length",max_length=pipe.tokenizer.model_max_length,truncation=True,return_tensors="np").input_ids)

cell 3:

#@title bfloat + scheduler issue reproduction p3 (run + scheduler tests)
# loop scheduler tests
for dtype_test in [0,1]:
    print("Testing scheduler data type" + (" default" if dtype_test == 0 else " bfloat16"))
    settings = {} if dtype_test == 0 else {"dtype": jnp.bfloat16} # edit: oops, had reversed.
    for lib_scheduler in [FlaxPNDMScheduler, FlaxLMSDiscreteScheduler, FlaxDDIMScheduler, FlaxDPMSolverMultistepScheduler, FlaxDDPMScheduler]:
        # build and implant
        pipe.scheduler, params["scheduler"] = lib_scheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler", **settings) # nonsense to need to use from_pretrained instead of from_config here... 
        print("Scheduler = " + lib_scheduler.__name__ + (": DEFAULT" if orig_scheduler == lib_scheduler.__name__ else ""))
        try:
            # run
            images_jnp = pipe(
                        prompt_ids=prompt_ids,
                        params=jax_utils.replicate(params), 
                        prng_seed=rng_repl,
                        num_inference_steps=50,
                        guidance_scale=guidance_scale,
                        jit=True
                        ).images
            # show me what you're made of
            if jnp.any(jnp.isnan(images_jnp)):
                raise ValueError("output image is NaNs")
            display_jnp_image(images_jnp[0].squeeze())
        except Exception as e:
            print("ERROR: " + str(e))

attaching outputs in place of logs...

Logs

don't know how to get logs, but here's output:

Testing scheduler data type bfloat16
Scheduler = FlaxPNDMScheduler: DEFAULT
ERROR: lax.select requires arguments to have the same dtypes, got bfloat16, float32. (Tip: jnp.where is a similar function that does automatic type promotion on inputs).
Scheduler = FlaxLMSDiscreteScheduler
ERROR: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[])>with<DynamicJaxprTrace(level=1/1)>
The problem arose with the `bool` function. 
The error occurred while tracing the function scanned_fun at /usr/local/lib/python3.8/dist-packages/jax/_src/lax/control_flow/loops.py:1619 for scan. This concrete value was not available in Python because it depends on the values of the argument 'loop_carry'.

See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError
Scheduler = FlaxDDIMScheduler
ERROR: output image is NaNs
Scheduler = FlaxDPMSolverMultistepScheduler
ERROR: scan carry output and input must have identical types, got
('ShapedArray(int32[], weak_type=True)', ('ShapedArray(float32[1,4,64,64])', DPMSolverMultistepSchedulerState(common=CommonSchedulerState(alphas='ShapedArray(bfloat16[1000])', betas='ShapedArray(bfloat16[1000])', alphas_cumprod='ShapedArray(bfloat16[1000])'), alpha_t='ShapedArray(bfloat16[1000])', sigma_t='ShapedArray(bfloat16[1000])', lambda_t='ShapedArray(bfloat16[1000])', init_noise_sigma='ShapedArray(bfloat16[])', timesteps='ShapedArray(int32[50])', num_inference_steps='ShapedArray(int32[], weak_type=True)', model_outputs='ShapedArray(bfloat16[2,1,4,64,64])', lower_order_nums='ShapedArray(int32[])', prev_timestep='ShapedArray(int32[])', cur_sample='DIFFERENT ShapedArray(float32[1,4,64,64]) vs. ShapedArray(bfloat16[1,4,64,64])'))).
Scheduler = FlaxDDPMScheduler
ERROR: step() missing 1 required positional argument: 'key'

Testing scheduler data type default
Scheduler = FlaxPNDMScheduler: DEFAULT

IMAGE WAS OK

Scheduler = FlaxLMSDiscreteScheduler
ERROR: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[])>with<DynamicJaxprTrace(level=1/1)>
The problem arose with the `bool` function. 
The error occurred while tracing the function scanned_fun at /usr/local/lib/python3.8/dist-packages/jax/_src/lax/control_flow/loops.py:1619 for scan. This concrete value was not available in Python because it depends on the values of the argument 'loop_carry'.

See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError
Scheduler = FlaxDDIMScheduler

IMAGE WAS OK

Scheduler = FlaxDPMSolverMultistepScheduler

IMAGE WAS OK

Scheduler = FlaxDDPMScheduler
ERROR: step() missing 1 required positional argument: 'key'

System Info

Google Colab Pro, TPU environment with High RAM, default setup except what is included in reproduction script

[patrickvonplaten]

Thanks a lot for the ping @krahnikblis :-)

@pcuenca do you have some time today/tomorrow to check this?

[pcuenca]

Absolutely! Thanks @krahnikblis for the great report :)

[pcuenca]

Hi @krahnikblis! I've been looking into this and there are several issues here.

• Regarding the use of from_pretrained to load schedulers, this was done for symmetry with the models and the PyTorch implementation. Some users may not be aware of the differences between a scheduler and a model, and it's useful to provide a consistent API to load any pipeline component using the same mechanism. I do agree that from_config should work too, maybe we can tackle that in a new PR.

• The main issue was that inference does not work in bfloat16 because scheduler weights need to remain in float32. You can check that out if you load a pipeline in bfloat16 and then take a look at the scheduler params. So, for your code to work, you need to omit the dtype when loading the scheduler with from_pretrained.

• Having said that, your testing surfaced two problems with the following schedulers: FlaxDDPMScheduler, FlaxLMSDiscreteScheduler:

  • The problem with FlaxDDPMScheduler was introduced by myself in 6043838. I just opened [Flax DDPM] Make key optional so default pipelines don't fail #2176 to fix it. However, generation with that scheduler does not work after fixing the bug - all generated images are black. I verified this was also the case before 6043838, so we need to look into this further.
  • FlaxLMSDiscreteScheduler is indeed going through a non-compilable path. This scheduler was made stateless in [Flax] Stateless schedulers, fixes and refactors #1661. Perhaps @skirsten would like to take a look? Please, let us know if you can't right now :)

• Our documentation could do with some indications about how to change the default scheduler in Flax. I'll open a PR to add an example to https://huggingface.co/docs/diffusers/using-diffusers/schedulers. For reference, this is how I usually perform inference in Flax with a custom scheduler:

import jax
import numpy as np
from flax.jax_utils import replicate
from flax.training.common_utils import shard

from diffusers import FlaxStableDiffusionPipeline, FlaxDPMSolverMultistepScheduler

model_id = "runwayml/stable-diffusion-v1-5"
scheduler, scheduler_state = FlaxDDPMScheduler.from_pretrained(
    model_id,
    subfolder="scheduler"
)
pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(
    model_id,
    scheduler=scheduler,
    revision="bf16",
    dtype=jax.numpy.bfloat16,
)
params["scheduler"] = scheduler_state

prompt = "a photo of an astronaut riding a horse on mars"
num_samples = jax.device_count()
prompt_ids = pipeline.prepare_inputs([prompt] * num_samples)

prng_seed = jax.random.PRNGKey(0)
num_inference_steps = 25

# shard inputs and rng
params = replicate(params)
prng_seed = jax.random.split(prng_seed, jax.device_count())
prompt_ids = shard(prompt_ids)

images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images
images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:])))

[pcuenca]

Thanks again for your detailed report, @krahnikblis!

[skirsten]

Hi, thanks for pinging me. I just had a quick look:

Regarding ERROR: lax.select requires arguments to have the same dtypes, got bfloat16, float32. (Tip: jnp.where is a similar function that does automatic type promotion on inputs). and ERROR: scan carry output and input must have identical types, ... DIFFERENT ShapedArray(float32[1,4,64,64]) vs. ShapedArray(bfloat16[1,4,64,64]):

Like @pcuenca said, by default the scheduler is loaded in float32 but what you tried should actually work.
I was able to test this some time ago using scheduler and pipeline in float16 and pipeline in float16 and scheduler in float32.

Regardless of the scheduler dtype, the outputs should always have the same dtype as the inputs.
So there is something going on regarding type promotion etc.
I will have a look at the code.

---

ERROR: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[])>with<DynamicJaxprTrace(level=1/1)> in FlaxLMSDiscreteScheduler:
This is expected. This scheduler uses scipy.integrate so its not jitable and thus does not work with the default pipeline. I refactored it a bit in #1661 but did not touch the actual logic. So I honestly have no idea if it would work in a non-jitted pipeline 😅 .

[pcuenca]

Thanks a lot, @skirsten! Ah, so it sounds like the FlaxLMSDiscreteScheduler has never been jittable then. I'll open a separate issue for that :)

[krahnikblis]

updating my position from before - i did get the DDIMS to work consistently, but i think i had to create it separately from the bf16 pipe, with the default data type, and assign it to the pipe? i mean, i know i did that, but i think i had to, because the bf16 pipe with scheduler built in gave me some grief or the other...

regarding the data types requiring to be fp32 for scheduler, is that "it has to be that way because that's how we copied the code from the authors" (i see several hard-coded fp32 defs in the source), or is that "it has to be that way because with lower precision the model produces crap outputs"?

regarding the LMSDS, i love it when someone has already solved a problem when i run into one... looks like this could be helpful in rewriting/approximating that quad function. i still don't grok the JVP/VJP stuff, but if that doesn't work, wondering if the integral can be estimated using a scan function over the sigmas/timesteps? hey i'm a marketing data scientist; no formal training here, just try connecting seemingly related dots till something works!!

[skirsten]

Okay, so I ran some tests for the different dtypes:

compute_dtype	weights_dtype	scheduler	scheduler_dtype	works?
bfloat16	bfloat16	DDIM	bfloat16	❌
bfloat16	bfloat16	DDIM	float16	✅
bfloat16	bfloat16	DDIM	float32	✅
bfloat16	bfloat16	DPMSolverMultistep	bfloat16	❌
bfloat16	bfloat16	DPMSolverMultistep	float16	❌
bfloat16	bfloat16	DPMSolverMultistep	float32	✅
bfloat16	bfloat16	PNDM	bfloat16	❌
bfloat16	bfloat16	PNDM	float16	❌
bfloat16	bfloat16	PNDM	float32	✅
bfloat16	float16	DDIM	bfloat16	❌
bfloat16	float16	DDIM	float16	✅
bfloat16	float16	DDIM	float32	✅
bfloat16	float16	DPMSolverMultistep	bfloat16	❌
bfloat16	float16	DPMSolverMultistep	float16	❌
bfloat16	float16	DPMSolverMultistep	float32	✅
bfloat16	float16	PNDM	bfloat16	❌
bfloat16	float16	PNDM	float16	❌
bfloat16	float16	PNDM	float32	✅
bfloat16	float32	DDIM	bfloat16	❌
bfloat16	float32	DDIM	float16	✅
bfloat16	float32	DDIM	float32	✅
bfloat16	float32	DPMSolverMultistep	bfloat16	❌
bfloat16	float32	DPMSolverMultistep	float16	❌
bfloat16	float32	DPMSolverMultistep	float32	✅
bfloat16	float32	PNDM	bfloat16	❌
bfloat16	float32	PNDM	float16	❌
bfloat16	float32	PNDM	float32	✅
float16	bfloat16	DDIM	bfloat16	❌
float16	bfloat16	DDIM	float16	✅
float16	bfloat16	DDIM	float32	✅
float16	bfloat16	DPMSolverMultistep	bfloat16	❌
float16	bfloat16	DPMSolverMultistep	float16	❌
float16	bfloat16	DPMSolverMultistep	float32	✅
float16	bfloat16	PNDM	bfloat16	❌
float16	bfloat16	PNDM	float16	❌
float16	bfloat16	PNDM	float32	✅
float16	float16	DDIM	bfloat16	❌
float16	float16	DDIM	float16	✅
float16	float16	DDIM	float32	✅
float16	float16	DPMSolverMultistep	bfloat16	❌
float16	float16	DPMSolverMultistep	float16	❌
float16	float16	DPMSolverMultistep	float32	✅
float16	float16	PNDM	bfloat16	❌
float16	float16	PNDM	float16	❌
float16	float16	PNDM	float32	✅
float16	float32	DDIM	bfloat16	❌
float16	float32	DDIM	float16	✅
float16	float32	DDIM	float32	✅
float16	float32	DPMSolverMultistep	bfloat16	❌
float16	float32	DPMSolverMultistep	float16	❌
float16	float32	DPMSolverMultistep	float32	✅
float16	float32	PNDM	bfloat16	❌
float16	float32	PNDM	float16	❌
float16	float32	PNDM	float32	✅
float32	bfloat16	DDIM	bfloat16	❌
float32	bfloat16	DDIM	float16	✅
float32	bfloat16	DDIM	float32	✅
float32	bfloat16	DPMSolverMultistep	bfloat16	❌
float32	bfloat16	DPMSolverMultistep	float16	❌
float32	bfloat16	DPMSolverMultistep	float32	✅
float32	bfloat16	PNDM	bfloat16	❌
float32	bfloat16	PNDM	float16	❌
float32	bfloat16	PNDM	float32	✅
float32	float16	DDIM	bfloat16	❌
float32	float16	DDIM	float16	✅
float32	float16	DDIM	float32	✅
float32	float16	DPMSolverMultistep	bfloat16	❌
float32	float16	DPMSolverMultistep	float16	❌
float32	float16	DPMSolverMultistep	float32	✅
float32	float16	PNDM	bfloat16	❌
float32	float16	PNDM	float16	❌
float32	float16	PNDM	float32	✅
float32	float32	DDIM	bfloat16	❌
float32	float32	DDIM	float16	✅
float32	float32	DDIM	float32	✅
float32	float32	DPMSolverMultistep	bfloat16	❌
float32	float32	DPMSolverMultistep	float16	❌
float32	float32	DPMSolverMultistep	float32	✅
float32	float32	PNDM	bfloat16	❌
float32	float32	PNDM	float16	❌
float32	float32	PNDM	float32	✅

Code

import jax
import numpy as np
import jax.numpy as jnp
from flax.jax_utils import replicate, unreplicate
from flax.training.common_utils import shard
from diffusers import (
    FlaxStableDiffusionPipeline,
    FlaxDPMSolverMultistepScheduler,
    FlaxDDIMScheduler,
    FlaxPNDMScheduler,
)

for compute_dtype in [jnp.float32, jnp.float16, jnp.bfloat16]:
    pipe, params = FlaxStableDiffusionPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5",
        revision="flax",  # weights have fp32 dtype
        dtype=compute_dtype,
    )
    params = replicate(params)

    for weights_dtype in [jnp.float32, jnp.float16, jnp.bfloat16]:

        # At this point the weights have weights_dtype dtype and the scheduler has fp32 dtype

        casted_params = {
            # Use to_bf16, to_fp16 and to_fp32 instead of _cast_floating_to
            "unet": pipe.unet._cast_floating_to(params["unet"], weights_dtype, None),
            "vae": pipe.vae._cast_floating_to(params["vae"], weights_dtype, None),
            "text_encoder": pipe.text_encoder._cast_floating_to(
                params["text_encoder"], weights_dtype, None
            ),
            "safety_checker": pipe.safety_checker._cast_floating_to(
                params["safety_checker"], weights_dtype, None
            ),
        }

        for scheduler in [
            FlaxDPMSolverMultistepScheduler,
            FlaxDDIMScheduler,
            FlaxPNDMScheduler,
        ]:
            for scheduler_dtype in [jnp.float32, jnp.float16, jnp.bfloat16]:
                pipe.scheduler = scheduler.from_config(
                    pipe.scheduler.config, dtype=scheduler_dtype
                )
                casted_params["scheduler"] = replicate(pipe.scheduler.create_state())

                try:
                    prompt = "a photo of an astronaut riding a horse on mars"
                    num_samples = jax.device_count()
                    prompt_ids = pipe.prepare_inputs([prompt] * num_samples)

                    prng_seed = jax.random.PRNGKey(0)
                    num_inference_steps = 25

                    # shard inputs and rng
                    prng_seed = jax.random.split(prng_seed, jax.device_count())
                    prompt_ids = shard(prompt_ids)

                    images = pipe(
                        prompt_ids,
                        casted_params,
                        prng_seed,
                        num_inference_steps,
                        jit=True,
                    ).images
                    images = pipe.numpy_to_pil(
                        np.asarray(images.reshape((num_samples,) + images.shape[-3:]))
                    )

                    print("works: ", end="")
                except Exception as e:
                    print(e)
                    print("failed: ", end="")

                print(
                    f"compute_dtype: {compute_dtype.__name__}, weights_dtype: {weights_dtype.__name__}, scheduler: {scheduler.__name__}, scheduler_dtype: {scheduler_dtype.__name__}"
                )

        del casted_params

    del params

Not great, not terrible, certainly could be improved, but at least as long as the scheduler_dtype is float32 it should work ;)

Also here are some examples of loading the scheduler in different ways:

pipe, params = FlaxStableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    revision="bf16", 
    dtype=jnp.bfloat16,
)

# Test 1:
# pipe.scheduler, params["scheduler"] = FlaxDPMSolverMultistepScheduler.from_pretrained(
#     "runwayml/stable-diffusion-v1-5", subfolder="scheduler"
# )

# Test 2:
# pipe.scheduler = FlaxDPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
# params["scheduler"] = pipe.scheduler.create_state()

# Test 3:
# pipe.scheduler = FlaxDPMSolverMultistepScheduler.from_config(
#     {
#         "beta_end": 0.012,
#         "beta_schedule": "scaled_linear",
#         "beta_start": 0.00085,
#         "num_train_timesteps": 1000,
#         "set_alpha_to_one": False,
#         "skip_prk_steps": True,
#         "steps_offset": 1,
#         "trained_betas": None,
#     }
# )
# params["scheduler"] = pipe.scheduler.create_state()

# Test 4:
# pipe.scheduler = FlaxDPMSolverMultistepScheduler(
#     beta_end=0.012,
#     beta_schedule="scaled_linear",
#     beta_start=0.00085,
#     num_train_timesteps=1000,
#     # set_alpha_to_one=False,
#     # skip_prk_steps=True,
#     # steps_offset=1,
#     trained_betas=None,
# )
# params["scheduler"] = pipe.scheduler.create_state()

d_params = replicate(params)
# ...

Note that I never set the scheduler dtype (which means it defaults to float32).
Hope this helps.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.