[Cross-Post] MAISI ControlNet Semantic Mask Conditioning Training Issues

[mazniashvili]

Hello all,
I will cross-post this from the main repository ( #8451 )
I am currently working on my master's project and it is a bit urgent:

firstly I would like to thank the MAISI team for their amazing work.
I have been playing around with it and I am currently trying to reproduce the mask-conditioned synthesis on a coronary CT dataset. I have trained the VAE and the MAISI diffusion model from scratch according to the tutorials, with great success.

Unfortunately, I cannot get the ControlNet to converge, I have tried different learning rates (1e-5 and 1e-6), let it train
for more than 20000 steps, however, as the images below show the conditioning is not applied well enough.
I am also comparing the conditionally generated images at every epoch, but they yield similar results.

Does someone have some pointers on how to properly train the ControlNet with semantic masks and how to achieve comparable results?

[Can-Zhao]

Did you see it get better after training it longer?

[mazniashvili]

Hello, thank you for reaching out!
Yes, I have gone up to 750,000 steps, but unfortunately it produced similar loss patterns and inference with the semantic mask was similar to the one above.

[Can-Zhao]

May I ask if you have binarized the mask as the controlnet training script shows? Thank you!

[mazniashvili]

Hi! Yes, we used the infer_controlnet.py script. Our current hypothesis is that the dataset is too small at
about 750 images for training, which might not enable stronger alignment between the unconditional ldm and the controlnet.
But the synthesized images are very good for the most part, here are some more samples:

[sara-create]

Hello @Can-Zhao ! I have the same problem , my loss doesn't decrease even after a long time and I don't understand why . I binarize the masks and follow the script ( maisi_train_controlnet_tutorial.ipynb ) but it doesn't converge after a lot of trials. For the context , I have 299 CT-scans with their masks.