Quantization is slow with FLUX.1-dev AND no effective LoRA support

[cryptexis]

Hi everyone,

I have the following use-case. I have one base model e.g. flux-dev and I would like to dynamically load/unload different LoRAs.
I have tried different quantization methods - and to be honest it was a pain use them.
Quanto was generating pure noise half of the time and not reliable at all. Also when I finally made it work, LoRAs won't load because of layers or keys were missing (depending on the LoRA).
Torchao was very promising, but on the LoRA loading step was reporting that it supports only int8.

While one wants to replicate a simple behaviour of comfyUI and load fp8 checkpoint and have a similar speed and memory usage, my journey literally turned into hell.

Now the only quantization method described in the documentation was BitsAndBytes method. Let me post the code first:

import torch
from diffusers import FluxTransformer2DModel, FluxPipeline

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import T5EncoderModel,

from rich import print
from time import time

torch_dtype = torch.float16

quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
transformer_8bit = FluxTransformer2DModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    torch_dtype=torch_dtype,
    cache_dir="./flux"
)

quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
text_encoder_8bit = T5EncoderModel.from_pretrained(
    model_id,
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    torch_dtype=torch_dtype,
    cache_dir="./flux",

)

pipe = FluxPipeline.from_pretrained(
    model_id, 
    transformer=transformer_8bit, 
    text_encoder_2=text_encoder_8bit,
    torch_dtype=torch_dtype,
    device_map="balanced",
    cache_dir="./flux",
)

with torch.no_grad():
    prompt = "A cat holding a sign that says hello world"
    image = pipe(
        prompt,
        output_type="pil",
        num_inference_steps=20,
        guidance_scale=3.5,
        height=1280,
        width=720
    ).images[0]

pipe.load_lora_weights(
    "./",  
    weight_name="Hyper-FLUX.1-dev-8steps-lora.safetensors",
    adapter_name="speed"
)
pipe.set_adapters(["speed"], adapter_weights=[0.125])

with torch.no_grad():
    prompt = "A cat holding a sign that says hello world"
    image = pipe(
        prompt,
        output_type="pil",
        num_inference_steps=8,
        guidance_scale=3.5,
        height=1280,
        width=720,
        generator=torch.Generator("cpu").manual_seed(0)
    ).images[0]

When I create pipeline:

pipe = FluxPipeline.from_pretrained(
    model_id, 
    torch_dtype=torch_dtype,
    device_map="balanced",
    cache_dir="./flux",
)

the vanilla version with memory usage between 30-40GB usage.

Without LoRA - I have 2.12it/sec
With LoRA - it drops to 1.87it/sec (12% decrease in speed) - well understandable (more parameters -> slower the model)

Now if I run it with quantization above, memory usage is between 22-26GB, however

Without LoRA - I have 1.56it/sec (26% decrease from vanilla)
With LoRA - it drops to 1.05it/sec (another 32% decrease in speed) - well understandable (more parameters -> slower the model)

Usage with LoRA also sums up to be 43% slower than with LoRA in vanilla mode. I am using L40s GPU

torch==2.6.0
diffussers==0.32.2
accelerate==1.6.0
transformers==4.50.3
peft==0.15.1
protobuf==6.30.2
sentencepiece=0.2.0

Is there anything could be done in general for dynamically loading LoRAs to quantized models?
And also for not losing so much performance?

[bghira]

#10512 for part 1

fuse LoRA to base model for part 2 question (performance).

unfortunately you'll have the base model change permanently when you fuse and unfuse weights due to precision issues, especially when quantising models.

there's no free lunch, ComfyUI stores a whole second copy of the model weights so that it can reset them after fuse.

[DN6]

Hi @cryptexis if you're trying reproduce loading the model in FP8 you might not need to use quantization. Could you please provide a link to the ComfyUI checkpoint you'd like to use?

[Eamymao]

I would like to recommend this speedup framework for you, which support both fp8 and lora while speeding up!
As someone who runs SD and FLUX models daily, this useful framework to speed up my image generation is so impressive!

As described, it only cost half the time to generate a 1024*1024 image than the original diffusers. At the same time, the image quality is not lost due to acceleration. Moreover, the code is very similar to diffusers and is easy to use.

Here are the github link: https://github.com/TMElyralab/lyraDiff
I really regretted not knowing this tool earlier. I would like to recommend it to all my friends in the community if it is useful.

[cryptexis]

hi @DN6
this is the one I am using in comfyUI:
FP8

@Eamymao it looks promising, I will check it out for sure. How does it work with LoRAs that reduce the number of steps e.g. https://huggingface.co/ByteDance/Hyper-SD/blob/main/Hyper-FLUX.1-dev-8steps-lora.safetensors ? Did you have chance to check ?

@bghira so the pipeline you propose is:

1. Load Lora
2. Fuse with base
3. Quantize
4. Run
5. when there is a need of new LoRA repeat the process ?

[DN6]

@cryptexis Could you try this snippet

import torch
from diffusers import FluxTransformer2DModel, FluxPipeline

ckpt_path= "https://huggingface.co/lllyasviel/flux1_dev/blob/main/flux1-dev-fp8.safetensors"
torch_dtype = torch.bfloat16

pipe = FluxPipeline.from_single_file(ckpt_path, torch_dtype=torch_dtype)
pipe.load_lora_weights("ByteDance/Hyper-SD", weight_name="Hyper-FLUX.1-dev-8steps-lora.safetensors", adapter_name="speed")
pipe.set_adapters(["speed"], adapter_weights=[0.125])
pipe.transformer.enable_layerwise_casting(storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16)

pipe.to("cuda")

prompt = "A cute cat hold a sign that says hello, surreal style"
image = pipe(prompt, num_inference_steps=8).images[0]
image.save("flux-layerwise.png")

[cryptexis]

@DN6 worked! thanks a lot!

[liorRabkin]

Hi, I have the same problem as described at the beginning of this issue. I tried @DN6 suggestion but got the error - 'FluxTransformer2DModel' object has no attribute 'enable_layerwise_casting'.
Any idea what the problem might be?

[bghira]

you'll need a newer Diffusers. install from git main :)

[Amitg1]

Hi, we tested the suggested solution.
Several questions:

1. @DN6 solution does not speed up the inference; it only makes it worse (A100) , if I understand correctly, it only helps with reading the FP8 file type. What do we gain from loading FP8?
2. The problem still stands, How can we use FP8 quantization in memory to speed inference + fast LoRa exchange?
3. Why enable_layerwise_casting was not released on 33.0?

Thank you!

[bghira]

a100 doesnt have fp8

[asomoza]

also layer wise casting was released on 0.33.0

[DN6]

@Amitg1 Layerwise casting is intended to save memory without having to use quantization. Regarding performance, please refer to @bghira's earlier comment here for more context about the slower speed.
#11215 (comment)

One thing you could try if you're going down the quantization route is to use a torch.compile friendly quant framework (TorchAO, Quanto, GGUF) to quantize the transformer and text encoders, compile them, and then use pipe.load_lora_weights. After a warm up run you should see a speed boost during inference.