flux fill cannot use lora(flux turbo lora)

[Suprhimp]

Describe the bug

I want to use flux fill pipeline with turbo lora, but when I load pipeline and load lora model, than gives error

Reproduction

from diffusers import FluxFillPipeline

def model_fn(model_dir: str) -> FluxFillPipeline:

    pipe = FluxFillPipeline.from_pretrained(
       "black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16
    ).to("cuda")
    
    pipe.load_lora_weights(f"alimama-creative/FLUX.1-Turbo-Alpha")
    pipe.fuse_lora()
    
    
    return pipe

Logs

NotImplementedError: Only LoRAs with input/output features higher than the current module's input/output features are currently supported. The provided LoRA contains in_features=64 and out_features=3072, which are lower than module_in_features=384 and module_out_features=3072. If you require support for this please open an issue at https://github.com/huggingface/diffusers/issues.

System Info

latest(github version diffusers), python3.10, ubuntu with nvidia gpu

Who can help?

@sayakpaul

[sayakpaul]

I am not sure you can. You're trying to load a LoRA that was obtained on Flux.1-Dev into something different.

[Suprhimp]

I thought that I can use it becuase I tried with comfyui first. it worked well with comfyui

[sayakpaul]

Cc: @a-r-r-o-w could be nice to try the LoRA expansion idea here as well.

[a-r-r-o-w]

Will have to take a look at what Comfy is doing to be sure they're not simply dropping the input lora layer (which requires expansion/shrinking) and using all the other layers - in this case it will work easily. Or perhaps they are not passing the channelwise-concatenated latents through the lora layer and just passing the denoising latents. I currently don't have the bandwidth to try this out, but will be sure to look into it soon!

[sayakpaul]

Or perhaps they are not passing the channelwise-concatenated latents through the lora layer and just passing the denoising latents.

Valid.

[Suprhimp]

I hope I can use turbo for flux fill soon :)

[kadirnar]

@sayakpaul
This code works. Is it useful for you?

https://github.com/nftblackmagic/catvton-flux/blob/main/tryon_inference_lora.py#L26

Code:

        transformer = FluxTransformer2DModel.from_pretrained(
            "xiaozaa/flux1-fill-dev-diffusers",   ## The official Flux-Fill weights
            torch_dtype=torch.bfloat16
        )
        print("Start loading LoRA weights")
        state_dict, network_alphas = FluxFillPipeline.lora_state_dict(
            pretrained_model_name_or_path_or_dict="xiaozaa/catvton-flux-lora-alpha",     ## The tryon Lora weights
            weight_name="pytorch_lora_weights.safetensors",
            return_alphas=True
        )
        is_correct_format = all("lora" in key or "dora_scale" in key for key in state_dict.keys())
        if not is_correct_format:
            raise ValueError("Invalid LoRA checkpoint.")

        FluxFillPipeline.load_lora_into_transformer(
            state_dict=state_dict,
            network_alphas=network_alphas,
            transformer=transformer,
        )
        

[Suprhimp]

got error ;)

def model_fn(model_dir: str) -> FluxFillPipeline:
    
    transformer = FluxTransformer2DModel.from_pretrained(
         "xiaozaa/flux1-fill-dev-diffusers",   ## The official Flux-Fill weights
        torch_dtype=torch.bfloat16
    )
    print("Start loading LoRA weights")
    state_dict, network_alphas = FluxFillPipeline.lora_state_dict(
        pretrained_model_name_or_path_or_dict=f"alimama-creative/FLUX.1-Turbo-Alpha",     ## The tryon Lora weights
        return_alphas=True
    )
    is_correct_format = all("lora" in key or "dora_scale" in key for key in state_dict.keys())
    if not is_correct_format:
        raise ValueError("Invalid LoRA checkpoint.")

    FluxFillPipeline.load_lora_into_transformer(
        state_dict=state_dict,
        network_alphas=network_alphas,
        transformer=transformer,
    )

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

        
    # pipe.load_lora_weights(f"alimama-creative/FLUX.1-Turbo-Alpha")
    # pipe.fuse_lora()
    return pipe

error message:

RuntimeError: Error(s) in loading state_dict for FluxTransformer2DModel:
        size mismatch for x_embedder.lora_A.default_0.weight: copying a param with shape torch.Size([64, 64]) from checkpoint, the shape in current model is torch.Size([64, 384]).

But loading xiaozaa/catvton-flux-lora-alpha works.

Also I tried with this way but it is not work too.

def model_fn(model_dir: str) -> FluxFillPipeline:
    
    transformer = FluxTransformer2DModel.from_pretrained(
         "xiaozaa/flux1-fill-dev-diffusers",   ## The official Flux-Fill weights
        torch_dtype=torch.bfloat16
    )

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

        
    pipe.load_lora_weights(f"alimama-creative/FLUX.1-Turbo-Alpha")
    pipe.fuse_lora()
    return pipe

[sayakpaul]

@kadirnar it won't work as the LoRA you showed in the example was obtained on the Flux fill checkpoint itself.

[sayakpaul]

I tried this huggingface:0d96a89...huggingface:ed91c533f but the outputs are pure garbage. Wondering what the strategy here is.

[sayakpaul]

Oh we had to set the scales accordingly:

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")

^ should work:

Make sure to install from expand-flux-lora branch.