How to load lora weight in distribution applications?

[Johnson-yue]

If I want to use xDiT with 2 GPU inference FluxControlPipeline, how should I do

I write a xFuserFluxControlPipeline class, but it can not load lora weight with right way
xFuserFluxTransformer in 1GPU have some parameters and another GPU have others.
How should I do ??

[asomoza]

this is an interesting question but sadly I don't have a clear answer since I haven't experimenting with what you're doing.

The first solution that comes to mind is to fuse the lora with the model before splitting the model but that's not really a solution if you plan to load and unload multiple loras.

Maybe @a-r-r-o-w has a more solid answer about this.

[a-r-r-o-w]

I think this question would be better asked to the maintainers of the xDiT repository (I see a duplicate here: xdit-project/xDiT#520).

From the looks of it, it seems like you have a custom pipeline implementation based on xDiT workflow for parallelization (I assume so because I couldn't find relevant implementation in their repo). Since it is custom code, it will be hard for us to debug/point out anything useful.

What should work is the same suggestion @asomoza. Simply adding the functionality of load_lora_weights on the pipeline, by making your class derive the DiffusionPipeline class, should work. You just need to make sure that both GPUs load the same weights and lora individually and correctly, and then you can apply any parallelism as you like.