MAISI ControlNet Semantic Mask Conditioning Training Issues

[mazniashvili]

Hello all,
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?