Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

✨ [Core] Add FreeU mechanism #5164

Merged
merged 63 commits into from
Oct 5, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
63 commits
Select commit Hold shift + click to select a range
eec915a
✨ Added Fourier filter function to upsample blocks
kadirnar Sep 24, 2023
8f2ee6b
🔧 Update Fourier_filter for float16 support
kadirnar Sep 24, 2023
ea33fc5
✨ Added UNetFreeUConfig to UNet model for FreeU adaptation 🛠️
kadirnar Sep 24, 2023
6d77fff
move unet to its original form and add fourier_filter to torch_utils.
sayakpaul Sep 27, 2023
7c4676b
implement freeU enable mechanism
sayakpaul Sep 27, 2023
a152d91
implement disable mechanism
sayakpaul Sep 27, 2023
c0de18d
resolution index.
sayakpaul Sep 27, 2023
a0ced0e
correct resolution idx condition.
sayakpaul Sep 27, 2023
19d5ab0
fix copies.
sayakpaul Sep 27, 2023
f3f9441
no need to use resolution_idx in vae.
sayakpaul Sep 27, 2023
266987c
spell out the kwargs
sayakpaul Sep 27, 2023
83846f5
proper config property
sayakpaul Sep 27, 2023
cb6ebe8
fix attribution setting
sayakpaul Sep 27, 2023
9c8eacd
place unet hasattr properly.
sayakpaul Sep 27, 2023
988ef76
fix: attribute access.
sayakpaul Sep 27, 2023
755a5c8
proper disable
sayakpaul Sep 27, 2023
37b091e
remove validation method.
sayakpaul Sep 27, 2023
4fd4adf
debug
sayakpaul Sep 27, 2023
40291ba
debug
sayakpaul Sep 27, 2023
2e15e94
debug
sayakpaul Sep 27, 2023
1e9c79a
debug
sayakpaul Sep 27, 2023
883fc9b
debug
sayakpaul Sep 27, 2023
ba08f30
debug
sayakpaul Sep 27, 2023
540974a
potential fix.
sayakpaul Sep 27, 2023
d38c251
add: doc.
sayakpaul Sep 27, 2023
785c0a0
fix copies
sayakpaul Sep 27, 2023
466e054
add: tests.
sayakpaul Sep 27, 2023
80e560e
add: support freeU in SDXL.
sayakpaul Sep 28, 2023
670b34b
set default value of resolution idx.
sayakpaul Sep 28, 2023
c5fc938
set default values for resolution_idx.
sayakpaul Sep 28, 2023
05ea56e
fix copies
sayakpaul Sep 28, 2023
93ee867
fix rest.
sayakpaul Sep 28, 2023
535eb59
fix copies
sayakpaul Sep 28, 2023
d277e64
address PR comments.
sayakpaul Sep 28, 2023
bb2d368
run fix-copies
sayakpaul Sep 28, 2023
e95b186
move apply_free_u to utils and other minors.
sayakpaul Sep 30, 2023
902cf7d
introduce support for video (unet3D)
sayakpaul Sep 30, 2023
15b1052
minor ups
sayakpaul Oct 2, 2023
7dcc939
consistent fix-copies.
sayakpaul Oct 2, 2023
3bf28bb
consistent stuff
sayakpaul Oct 2, 2023
64ade67
Merge branch 'main' into add-freeU
sayakpaul Oct 2, 2023
08d61fa
Merge branch 'main' into add-freeU
sayakpaul Oct 2, 2023
d68663c
fix-copies
sayakpaul Oct 2, 2023
0e0af08
add: rest
sayakpaul Oct 2, 2023
8d0a204
add: docs.
sayakpaul Oct 2, 2023
55ad535
fix: tests
sayakpaul Oct 2, 2023
797b4b9
fix: doc path
sayakpaul Oct 2, 2023
86419b9
Merge branch 'main' into add-freeU
sayakpaul Oct 2, 2023
5e27ff7
Merge branch 'main' into add-freeU
patrickvonplaten Oct 4, 2023
2ad4953
Merge branch 'main' into add-freeU
sayakpaul Oct 5, 2023
1a8e5d1
Apply suggestions from code review
sayakpaul Oct 5, 2023
adc9d5c
style up
sayakpaul Oct 5, 2023
c4f99d4
move to techniques.
sayakpaul Oct 5, 2023
518e4b1
add: slow test for sd freeu.
sayakpaul Oct 5, 2023
dee3781
add: slow test for sd freeu.
sayakpaul Oct 5, 2023
fc39d22
add: slow test for sd freeu.
sayakpaul Oct 5, 2023
053f3ed
add: slow test for sd freeu.
sayakpaul Oct 5, 2023
d8ef3a1
add: slow test for sd freeu.
sayakpaul Oct 5, 2023
3da96e2
add: slow test for sd freeu.
sayakpaul Oct 5, 2023
d8c8771
add: slow test for video with freeu
sayakpaul Oct 5, 2023
0d34cf0
add: slow test for video with freeu
sayakpaul Oct 5, 2023
8e72f85
add: slow test for video with freeu
sayakpaul Oct 5, 2023
aa1a061
style
sayakpaul Oct 5, 2023
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions docs/source/en/_toctree.yml
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,8 @@
title: Control image brightness
- local: using-diffusers/weighted_prompts
title: Prompt weighting
- local: using-diffusers/freeu
title: Improve generation quality with FreeU
title: Techniques
- sections:
- local: using-diffusers/pipeline_overview
Expand Down
123 changes: 123 additions & 0 deletions docs/source/en/using-diffusers/freeu.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,123 @@
# Improve generation quality with FreeU

[[open-in-colab]]

The UNet is responsible for denoising during the reverse diffusion process, and there are two distinct features in its architecture:

1. Backbone features primarily contribute to the denoising process
2. Skip features mainly introduce high-frequency features into the decoder module and can make the network overlook the semantics in the backbone features

However, the skip connection can sometimes introduce unnatural image details. [FreeU](https://hf.co/papers/2309.11497) is a technique for improving image quality by rebalancing the contributions from the UNet’s skip connections and backbone feature maps.

FreeU is applied during inference and it does not require any additional training. The technique works for different tasks such as text-to-image, image-to-image, and text-to-video.

In this guide, you will apply FreeU to the [`StableDiffusionPipeline`], [`StableDiffusionXLPipeline`], and [`TextToVideoSDPipeline`].

## StableDiffusionPipeline

Load the pipeline:

```py
from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, safety_checker=None
).to("cuda")
```

Then enable the FreeU mechanism with the FreeU-specific hyperparameters. These values are scaling factors for the backbone and skip features.

```py
pipeline.enable_freeu(s1=0.9, s2=0.2, b1=1.2, b2=1.4)
```

The values above are from the official FreeU [code repository](https://github.com/ChenyangSi/FreeU) where you can also find [reference hyperparameters](https://github.com/ChenyangSi/FreeU#range-for-more-parameters) for different models.

<Tip>

Disable the FreeU mechanism by calling `disable_freeu()` on a pipeline.

</Tip>

And then run inference:

```py
prompt = "A squirrel eating a burger"
seed = 2023
image = pipeline(prompt, generator=torch.manual_seed(seed)).images[0]
```

The figure below compares non-FreeU and FreeU results respectively for the same hyperparameters used above (`prompt` and `seed`):

![](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/freeu/sdv1_5_freeu.jpg)


Let's see how Stable Diffusion 2 results are impacted:

```py
from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16, safety_checker=None
).to("cuda")

prompt = "A squirrel eating a burger"
seed = 2023

pipeline.enable_freeu(s1=0.9, s2=0.2, b1=1.1, b2=1.2)
image = pipeline(prompt, generator=torch.manual_seed(seed)).images[0]
```


![](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/freeu/sdv2_1_freeu.jpg)

## Stable Diffusion XL

Finally, let's take a look at how FreeU affects Stable Diffusion XL results:

```py
from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16,
).to("cuda")

prompt = "A squirrel eating a burger"
seed = 2023

# Comes from
# https://wandb.ai/nasirk24/UNET-FreeU-SDXL/reports/FreeU-SDXL-Optimal-Parameters--Vmlldzo1NDg4NTUw
pipeline.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
image = pipeline(prompt, generator=torch.manual_seed(seed)).images[0]
```


![](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/freeu/sdxl_freeu.jpg)

## Text-to-video generation

FreeU can also be used to improve video quality:

```python
from diffusers import DiffusionPipeline
from diffusers.utils import export_to_video
import torch

model_id = "cerspense/zeroscope_v2_576w"
pipe = DiffusionPipeline.from_pretrained("cerspense/zeroscope_v2_576w", torch_dtype=torch.float16).to("cuda")
pipe = pipe.to("cuda")

prompt = "an astronaut riding a horse on mars"
seed = 2023

# The values come from
# https://github.com/lyn-rgb/FreeU_Diffusers#video-pipelines
pipe.enable_freeu(b1=1.2, b2=1.4, s1=0.9, s2=0.2)
video_frames = pipe(prompt, height=320, width=576, num_frames=30, generator=torch.manual_seed(seed)).frames
export_to_video(video_frames, "astronaut_rides_horse.mp4")
```

Thanks to [kadirnar](https://github.com/kadirnar/) for helping to integrate the feature, and to [justindujardin](https://github.com/justindujardin) for the helpful discussions.
Loading
Loading