[LoRA] feat: support loading loras into 4bit quantized Flux models. by sayakpaul · Pull Request #10578 · huggingface/diffusers (original) (raw)

What does this PR do?

We broke Flux LoRA (not Control LoRA) loading for 4bit BnB Flux in 0.32.0, when supporting Flux Control LoRAs (yeah only applies to Flux).

To reproduce:

Code

import torch from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig, FluxTransformer2DModel, FluxPipeline from huggingface_hub import hf_hub_download

transformer_4bit = FluxTransformer2DModel.from_pretrained( "black-forest-labs/FLUX.1-dev", subfolder="transformer", quantization_config=DiffusersBitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16 ), torch_dtype=torch.bfloat16, ) pipe = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", transformer=transformer_4bit, torch_dtype=torch.bfloat16, ).to("cuda")

pipe.load_lora_weights( hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"), adapter_name="hyper-sd" ) pipe.set_adapters("hyper-sd", adapter_weights=0.125)

prompt = "A robot made of exotic candies and chocolates of different kinds. The background is filled with confetti and celebratory gifts."

image = pipe( prompt=prompt, height=1024, width=1024, max_sequence_length=512, num_inference_steps=8, guidance_scale=50, generator=torch.Generator().manual_seed(42), ).images[0] image.save("out.jpg")

This above code will run with v0.31.0-release branch but will fail with v0.32.0-release along with main. This went uncaught because we don't test for it.

This PR attempts to partially fix the problem so that we can at least resort to a behavior similar to what was happening in the v0.31.0-release branch. Want to use this PR to refine how we're doing that. I want to ship this PR first and tackle the TODOs in a follow-up. Once this PR is done, we might have to do a patch release.

Related issue: #10550