Unverified Commit 9b5180cb authored by latentCall145's avatar latentCall145 Committed by GitHub
Browse files

Flux fp16 inference fix (#9097)



* clipping for fp16

* fix typo

* added fp16 inference to docs

* fix docs typo

* include link for fp16 investigation

---------
Co-authored-by: default avatarSayak Paul <spsayakpaul@gmail.com>
parent 16a93f1a
...@@ -37,7 +37,7 @@ Both checkpoints have slightly difference usage which we detail below. ...@@ -37,7 +37,7 @@ Both checkpoints have slightly difference usage which we detail below.
```python ```python
import torch import torch
from diffusers import FluxPipeline from diffusers import FluxPipeline
pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16) pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload() pipe.enable_model_cpu_offload()
...@@ -61,7 +61,7 @@ out.save("image.png") ...@@ -61,7 +61,7 @@ out.save("image.png")
```python ```python
import torch import torch
from diffusers import FluxPipeline from diffusers import FluxPipeline
pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16) pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload() pipe.enable_model_cpu_offload()
...@@ -77,6 +77,34 @@ out = pipe( ...@@ -77,6 +77,34 @@ out = pipe(
out.save("image.png") out.save("image.png")
``` ```
## Running FP16 inference
Flux can generate high-quality images with FP16 (i.e. to accelerate inference on Turing/Volta GPUs) but produces different outputs compared to FP32/BF16. The issue is that some activations in the text encoders have to be clipped when running in FP16, which affects the overall image. Forcing text encoders to run with FP32 inference thus removes this output difference. See [here](https://github.com/huggingface/diffusers/pull/9097#issuecomment-2272292516) for details.
FP16 inference code:
```python
import torch
from diffusers import FluxPipeline
pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16) # can replace schnell with dev
# to run on low vram GPUs (i.e. between 4 and 32 GB VRAM)
pipe.enable_sequential_cpu_offload()
pipe.vae.enable_slicing()
pipe.vae.enable_tiling()
pipe.to(torch.float16) # casting here instead of in the pipeline constructor because doing so in the constructor loads all models into CPU memory at once
prompt = "A cat holding a sign that says hello world"
out = pipe(
prompt=prompt,
guidance_scale=0.,
height=768,
width=1360,
num_inference_steps=4,
max_sequence_length=256,
).images[0]
out.save("image.png")
```
## Single File Loading for the `FluxTransformer2DModel` ## Single File Loading for the `FluxTransformer2DModel`
The `FluxTransformer2DModel` supports loading checkpoints in the original format shipped by Black Forest Labs. This is also useful when trying to load finetunes or quantized versions of the models that have been published by the community. The `FluxTransformer2DModel` supports loading checkpoints in the original format shipped by Black Forest Labs. This is also useful when trying to load finetunes or quantized versions of the models that have been published by the community.
...@@ -134,4 +162,4 @@ image.save("flux-fp8-dev.png") ...@@ -134,4 +162,4 @@ image.save("flux-fp8-dev.png")
[[autodoc]] FluxPipeline [[autodoc]] FluxPipeline
- all - all
- __call__ - __call__
\ No newline at end of file
...@@ -125,6 +125,8 @@ class FluxSingleTransformerBlock(nn.Module): ...@@ -125,6 +125,8 @@ class FluxSingleTransformerBlock(nn.Module):
gate = gate.unsqueeze(1) gate = gate.unsqueeze(1)
hidden_states = gate * self.proj_out(hidden_states) hidden_states = gate * self.proj_out(hidden_states)
hidden_states = residual + hidden_states hidden_states = residual + hidden_states
if hidden_states.dtype == torch.float16:
hidden_states = hidden_states.clip(-65504, 65504)
return hidden_states return hidden_states
...@@ -223,6 +225,8 @@ class FluxTransformerBlock(nn.Module): ...@@ -223,6 +225,8 @@ class FluxTransformerBlock(nn.Module):
context_ff_output = self.ff_context(norm_encoder_hidden_states) context_ff_output = self.ff_context(norm_encoder_hidden_states)
encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
if encoder_hidden_states.dtype == torch.float16:
encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
return encoder_hidden_states, hidden_states return encoder_hidden_states, hidden_states
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment