|
--- |
|
language: |
|
- en |
|
library_name: diffusers |
|
license: other |
|
license_name: flux-1-dev-non-commercial-license |
|
license_link: LICENSE.md |
|
--- |
|
|
|
LoRA is the de-facto technique for quickly adapting a pre-trained large model on custom use cases. Typically, LoRA matrices are low-rank in nature. Now, the word “low” can vary depending on the context, but usually, for a large diffusion model like [Flux](https://huggingface.co/black-forest-labs/FLUX.1-dev), a rank of 128 can be considered high. This is because users may often need to keep multiple LoRAs unfused in memory to be able to quickly switch between them. So, the higher the rank, the higher the memory on top of the volume of the base model. |
|
|
|
So, what if we could take an existing LoRA checkpoint with a high rank and reduce its rank even further to: |
|
|
|
- Reduce the memory requirements |
|
- Enable use cases like `torch.compile()` (which require all the LoRAs to be of the same rank to avoid re-compilation) |
|
|
|
This project explores two options to reduce the original LoRA checkpoint into an even smaller one: |
|
|
|
* Random projections |
|
* SVD |
|
|
|
## Random projections |
|
|
|
Basic idea: |
|
|
|
1. Generate a random projection matrix: `R = torch.randn(new_rank, original_rank, dtype=torch.float32) / torch.sqrt(torch.tensor(new_rank, dtype=torch.float32))`. |
|
2. Then compute the new LoRA up and down matrices: |
|
|
|
```python |
|
# We keep R in torch.float32 for numerical stability. |
|
lora_A_new = (R @ lora_A.to(R.dtype)).to(lora_A.dtype) |
|
lora_B_new = (lora_B.to(R.dtype) @ R.T).to(lora_B.dtype) |
|
``` |
|
|
|
If `lora_A` and `lora_B` had shapes of (42, 3072) and (3072, 42) respectively, `lora_A_new` and `lora_B_new` will have (4, 3072) and (3072, 4), respectively. |
|
|
|
|
|
### Results |
|
|
|
Tried on this LoRA: [https://huggingface.co/glif/how2draw](https://huggingface.co/glif/how2draw). Unless explicitly specified, a rank of 4 was used for all experiments. Here’s a side-by-side comparison of the original and the reduced LoRAs (on the same seed). |
|
|
|
<details> |
|
<summary>Inference code</summary> |
|
|
|
```python |
|
from diffusers import DiffusionPipeline |
|
import torch |
|
|
|
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to("cuda") |
|
# Change accordingly. |
|
lora_id = "How2Draw-V2_000002800_svd.safetensors" |
|
pipe.load_lora_weights(lora_id) |
|
|
|
prompts = [ |
|
"Yorkshire Terrier with smile, How2Draw", |
|
"a dolphin, How2Draw", |
|
"an owl, How3Draw", |
|
"A silhouette of a girl performing a ballet pose, with elegant lines to suggest grace and movement. The background can include simple outlines of ballet shoes and a music note. The image should convey elegance and poise in a minimalistic style, How2Draw" |
|
] |
|
images = pipe( |
|
prompts, num_inference_steps=50, max_sequence_length=512, guidance_scale=3.5, generator=torch.manual_seed(0) |
|
).images |
|
``` |
|
|
|
</details> |
|
|
|
<table style="border-collapse: collapse;"> |
|
<tbody> |
|
<tr> |
|
<td align="center"><img src="https://huggingface.co/sayakpaul/lower-rank-flux-lora/resolve/main/images/collage_0.png" alt="Image 1"></td> |
|
<td align="center">Yorkshire Terrier with smile, How2Draw</td> |
|
</tr> |
|
<tr> |
|
<td align="center"><img src="https://huggingface.co/sayakpaul/lower-rank-flux-lora/resolve/main/images/collage_1.png" alt="Image 2"></td> |
|
<td align="center">a dolphin, How2Draw</td> |
|
</tr> |
|
<tr> |
|
<td align="center"><img src="https://huggingface.co/sayakpaul/lower-rank-flux-lora/resolve/main/images/collage_2.png" alt="Image 3"></td> |
|
<td align="center">an owl, How3Draw</td> |
|
</tr> |
|
<tr> |
|
<td align="center"><img src="https://huggingface.co/sayakpaul/lower-rank-flux-lora/resolve/main/images/collage_3.png" alt="Image 4"></td> |
|
<td align="center" style="padding: 0; margin: 0;"> |
|
A silhouette of a girl performing a ballet pose, with elegant lines to suggest grace and movement. |
|
The background can include simple outlines of ballet shoes and a music note. |
|
The image should convey elegance and poise in a minimalistic style, How2Draw |
|
</td> |
|
</tr> |
|
</tbody> |
|
</table> |
|
|
|
|
|
Code: [`low_rank_lora.py`](https://huggingface.co/sayakpaul/lower-rank-flux-lora/blob/main/low_rank_lora.py) |
|
|
|
### Notes |
|
|
|
* One should experiment with the `new_rank` parameter to obtain the desired trade-off between performance and memory. With a `new_rank` of 4, we reduce the size of the LoRA from 451MB to 42MB. |
|
* There is a `use_sparse` option in the script above for using sparse random projection matrices. |
|
|
|
## SVD |
|
|
|
<details> |
|
<summary>Results</summary> |
|
|
|
 |
|
|
|
 |
|
|
|
 |
|
|
|
 |
|
|
|
</details> |
|
|
|
### Randomized SVD |
|
|
|
Full SVD can be time-consuming. Truncated SVD is useful very large sparse matrices. We can use randomized SVD for none-to-negligible loss in quality but significantly faster speed. |
|
|
|
<details> |
|
<summary>Results</summary> |
|
|
|
 |
|
|
|
 |
|
|
|
 |
|
|
|
 |
|
|
|
</details> |
|
|
|
Code: [`svd_low_rank_lora.py`](https://huggingface.co/sayakpaul/lower-rank-flux-lora/blob/main/svd_low_rank_lora.py) |
|
|
|
### Tune the knobs in SVD |
|
|
|
- `new_rank` as always |
|
- `niter` when using randomized SVD |
|
|
|
## Reduced checkpoints |
|
|
|
* Randomized SVD: [How2Draw-V2_000002800_rand_svd.safetensors](./How2Draw-V2_000002800_rand_svd.safetensors) |
|
* Full SVD: [How2Draw-V2_000002800_svd.safetensors](./How2Draw-V2_000002800_svd.safetensors) |
|
* Random projections: [How2Draw-V2_000002800_reduced.safetensors](./How2Draw-V2_000002800_reduced.safetensors) |
