Diffusion with spade for image to image translation

[OdedRotem314]

Hi all

I want to train a diffusion model (preferably not a ldm) between two image modalities (imgs1+noise ==> imgs2).
I want to use SPADE so that imgs1 are also injected to the middle layers of the network.
Images are volumetric by the way.
I trained successfully using DiffusionModelUNet but I want to see if SPADEDiffusionModelUNet can improve details.

Can anyone write down the main code:

1. SPADEDiffusionModelUNet hyperparameters
2. Training script - is partial needed here? is context needed.
   The below code did not seem to train anything.

noise = torch.randn_like(images2).to(device)
timesteps = torch.randint(0, 1000, (len(images1),)).to(device)
noisy_seg = scheduler.add_noise(original_samples=images2, noise=noise, timesteps=timesteps)
combined = torch.cat((noisy_seg, images1), dim=1).to(device)
model = partial(model, seg=images1)
noise_prediction = model(x=combined, timesteps=timesteps)

Thanks
Oded

[KumoLiu]

Hi @OdedRotem314, looks like you are interested in this tutorial, https://github.com/Project-MONAI/tutorials/blob/main/generation/spade_ldm/spade_ldm_brats.ipynb.
For more tutorials, you can refer to this folder: https://github.com/Project-MONAI/tutorials/tree/main/generation

Hope it helps, thanks.

[OdedRotem314]

Hi KumoLiu
Thank you for your response!
I used the spade ldm tutorial as a starting point but ran into some difficulties as mentioned above.
I am not training an autoencoder and doing 3D image to image translation so I'm not sure how to adapt the inferer.

This is what I do but the inference predictions at x0 turn up as pure noise.

model = SPADEDiffusionModelUNet(spatial_dims=3, in_channels=2, out_channels=1, num_channels=(64, 128, 256), attention_levels=(False, False, True) , num_res_blocks=2, num_head_channels=64, norm_num_groups = 8, with_conditioning = False, label_nc = 1, spade_intermediate_channels=64) # resblock_updown = True)

train step

noise = torch.randn_like(imgs2).to(device)
timesteps = torch.randint(0, 1000, (len(imgs1),)).to(device)
noisy_imgs2 = scheduler.add_noise(original_samples=imgs2, noise=noise, timesteps=timesteps)
combined = torch.cat((noisy_imgs2 , imgs1), dim=1).to(device)
noise_prediction = model(x=combined, timesteps=timesteps, seg=imgs1)

What do you think is wrong here?