Repository navigation
latest release breaks Flax pipes for Stable Diffusion when using bfloat16 and/or non-default scheduler #2155
Description
Activity
Thanks a lot for the ping @krahnikblis :-)
@pcuenca do you have some time today/tomorrow to check this?
Absolutely! Thanks @krahnikblis for the great report :)
Reacted by Patrick von PlatenHi @krahnikblis! I've been looking into this and there are several issues here.
-
Regarding the use of
from_pretrainedto 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 thatfrom_configshould work too, maybe we can tackle that in a new PR. -
The main issue was that inference does not work in
bfloat16because scheduler weights need to remain infloat32. You can check that out if you load a pipeline inbfloat16and then take a look at the scheduler params. So, for your code to work, you need to omit thedtypewhen loading the scheduler withfrom_pretrained. -
Having said that, your testing surfaced two problems with the following schedulers:
FlaxDDPMScheduler,FlaxLMSDiscreteScheduler:- The problem with
FlaxDDPMSchedulerwas introduced by myself in 6043838. I just opened [Flax DDPM] Makekeyoptional 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. FlaxLMSDiscreteScheduleris 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 :)
- The problem with
-
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:])))
Reacted by Xin Kong-
Thanks again for your detailed report, @krahnikblis!
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).andERROR: 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 usesscipy.integrateso 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 😅 .Thanks a lot, @skirsten! Ah, so it sounds like the
FlaxLMSDiscreteSchedulerhas never been jittable then. I'll open a separate issue for that :)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!!
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 paramsNot great, not terrible, certainly could be improved, but at least as long as the
scheduler_dtypeisfloat32it 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.Reacted by Álvaro TorralbaThis 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.
- addedstaleIssues that haven't received updatesIssues that haven't received updates
on Mar 1, 2023 so... closed due to stale means... not important enough to ever fix?
Hi @krahnikblis, this issue originated 3 separate PRs and a very thorough investigation by @skirsten showing that inference works when the scheduler (any scheduler) is loaded in
float32. Two of the three PRs have been merged, and the remaining one affects one scheduler and is being tracked separately in that PR. So I strongly disagree with this being "not important enough".I consider this to having been adequately addressed and limitations explained. If you think there are particular use cases that should be considered, please feel free to open a new issue explaining them so we can discuss there.
Reacted by Patrick von Platen
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_configfor 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 tofrom_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 withjit=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:
cell 2:
cell 3:
attaching outputs in place of logs...
Logs
System Info
Google Colab Pro, TPU environment with High RAM, default setup except what is included in reproduction script