[LoRA] feat: support loading regular Flux LoRAs into Flux Control, and Fill

[sayakpaul]

What does this PR do?

Fixes #10180, #10227, #10184

In short, this PR enables few-steps inference for Flux Control, Fill, Redux, etc.

Fill + Turbo LoRA

from diffusers import FluxFillPipeline
from diffusers.utils import load_image
import torch

pipe = FluxFillPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16
).to("cuda")

adapter_id = "alimama-creative/FLUX.1-Turbo-Alpha"
pipe.load_lora_weights(adapter_id)

image = load_image("https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/cup.png")
mask = load_image("https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/cup_mask.png")

image = pipe(
    prompt="a white paper cup",
    image=image,
    mask_image=mask,
    height=1632,
    width=1232,
    guidance_scale=30,
    num_inference_steps=8,
    max_sequence_length=512,
    generator=torch.Generator("cpu").manual_seed(0)
).images[0]
image.save("flux-fill-dev.png")

Flux Control LoRA + Turbo LoRA (different from the previous one)

from diffusers import FluxControlPipeline
from image_gen_aux import DepthPreprocessor
from diffusers.utils import load_image
from huggingface_hub import hf_hub_download
import torch

control_pipe = FluxControlPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
control_pipe.load_lora_weights("black-forest-labs/FLUX.1-Depth-dev-lora", adapter_name="depth")
control_pipe.load_lora_weights(
    hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"), adapter_name="hyper-sd"
)
control_pipe.set_adapters(["depth", "hyper-sd"], adapter_weights=[0.85, 0.125])
control_pipe.enable_model_cpu_offload()

prompt = "A robot made of exotic candies and chocolates of different kinds. The background is filled with confetti and celebratory gifts."
control_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/robot.png")

processor = DepthPreprocessor.from_pretrained("LiheYoung/depth-anything-large-hf")
control_image = processor(control_image)[0].convert("RGB")

image = control_pipe(
    prompt=prompt,
    control_image=control_image,
    height=1024,
    width=1024,
    num_inference_steps=8,
    guidance_scale=10.0,
    generator=torch.Generator().manual_seed(42),
).images[0]
image.save("output.png")

Todods

• Integration tests
• Docs

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[sayakpaul]

Requesting for a review from @BenjaminBossan for initial stuff. Then will request reviews from others.

[BenjaminBossan]

Thanks for extending the functionality of loading LoRA adapters when shapes need to be expanded. The PR LGTM, I only have a nit.

One question that came up (maybe it was already discussed and I just missed it or forgot): Right now, this type of expansion is permanent, right? I.e. even after unloading the LoRA that made the expansion necessary in the first place, the expansion is not undone. Probably that would be quite hard to add and not worth the effort, I'm just curious.

[sayakpaul]

@BenjaminBossan

Right now, this type of expansion is permanent, right? I.e. even after unloading the LoRA that made the expansion necessary in the first place, the expansion is not undone. Probably that would be quite hard to add and not worth the effort, I'm just curious.

The LoRA state dict expansion is permanent. But model-level state dict expansion can be undone is being added in #10206

[sayakpaul]

Removing this part of the test because in case LoRA input feature dimensions are less than the original, we expand it.

This is tested below with test_lora_expanding_shape_with_normal_lora() and test_load_regular_lora().

[BenjaminBossan]

The LoRA state dict expansion is permanent. But model-level state dict expansion can be undone is being added in #10206

Yes, something similar for LoRA would be nice, but it's not as important, as the overhead for LoRA should be relatively small.