Training ControlNet-Xs

[nguyendung1612]

Thank you @UmerHA for contributing the ControlNet demo. I tried to create a training script so I could train the model on my data but had some problems using the UNetControlNetXSModel.
Can you provide training or give me some suggestions?

[UmerHA]

What were the problems you ran into?

[nguyendung1612]

What were the problems you ran into?

I tried 2 ways:

1. I trained the UNetControlNetXSModel and saved the pretrained model. But when I checked your Pipeline and yours pretrain, I saw that you only saved the ControlNetXSAdapter class. How can you save ControlNetXSAdapter while training UNetControlNetXSModel?

2. I trained ControlNetXSAdapter class (optimizer and scheduler for this), but my loss did not change during training. I think it's because I only used UNetControlNetXSModel for forward without optimizing it.

[UmerHA]

Your first approach is correct. You can train and save a UNetControlNetXSModel, let's call it unet_cnxs. Then use it in a pipeline with StableDiffusionControlNetXSPipeline(unet=unet_cnxs, controlnet=None).

Let me know if that works!

[rardz]

But saving a UNetControlNetXSModel after training results in storing an additional set of UNet weights, which will be extremely wasteful. In fact, the ControlNetXSAdapter class already contains all the parameters of controlnet-xs part. Why not put all the controlling processing into ControlNetXSAdapter's forward as the controlnet module do.

[UmerHA]

2 reasons:

1. the user might decide to unfreeze the unet also, and train the entire model. so it needs to be possible to save the entire model (unet+adapter).

2. extracting the adapter from the UNetControlNetXSModel is currently just not implemented. but it's 100% doable. feel free to do it! happy to guide you!

[rardz]

but:

1. UNet' parameters can be independently unfreezed for training in a non-rewriting layer approach, which might even be more convenient, also as the way controlnet do.
2. Without the new layer structure, there will be no need to implement the subsequent adapter weights extracting part.

[UmerHA]

It's not that easy, as controlnet and controlnet-xs differs by quite a lot. A controlnet addon can be independently run from the unet, and its output is added to the unet at some point. But a controlnet-xs can not be run independently. It needs to be fused with the unet.

Here's the relevant discussion from the PR: #6772 (comment). If you're questions can't be answered, I'd suggest you can ping Yiyi.

Honestly, though, I think for the use case in this discussion, the easiest way is to extract the cnxs adapter from the UNetControlNetXSModel.

[sayakpaul]

It should be turned into a discussion.