Torch.compile & multiple LoRAs with flux

[jonluca]

I am trying to speed up inference on an image generation pipeline that will swap in many differenet loras.

Is it possible to compile the base Flux model and then load a lora into it, create an image, and unload the lora without needing to recompile the model every time?

from diffusers import FluxPipeline


my_loras = [None, "lora1.safetensors", "lora2.safetensors"]
model = FluxPipeline.from_pretrained(base_model_id, torch_dtype=torch.bfloat16).to('cuda')
model.transformer.to(memory_format=torch.channels_last)
model.transformer = torch.compile(model.transformer)
# initial compilation
model(
            "A photo of a cat",
            num_inference_steps=8,
            guidance_scale=3.5,
            height=1024,
            width=1024,
            max_sequence_length=512,
        ).images[0]
for lora in my_loras:
    model.load_lora_weights(lora_model)
    model(
            "A photo of a cat",
            num_inference_steps=8,
            guidance_scale=3.5,
            height=1024,
            width=1024,
            max_sequence_length=512,
        ).images[0]

In this example, it seems like it has to recompile every time I add in a Lora. Each lora will be used exactly once, so I'd like to take advantage of the speed improvements to the base model.

Each lora will be discarded and theres no guarantee that the sizes of the loras will be the same.

Does fusing help here? or is there a way to tell pytorch to reuse the compiled information? I didnt think that a LoRA would change the compilation learnings.

[yiyixuxu]

we have a WIP here #9453

[rodneyviana]

Not for Flux. The discussion doesn't mention Flux