problem in denoising and generated image

[Ariz-Mohammadi]

Hello, explainable AI,
First of all, thank you so much for the videos and codes you shared. I have learned lots of things from you.
I tried to generate new images (unconditional)using the Mammography images dataset(3100 images). I trained the VQVAE for 100 epochs, and the results were good, as you can see below.

Then, I trained the diffusion part. I started with 100 epochs, and when I ran the sample_ddpm_vqvae.py the output image was like this

I thought maybe the number of epochs in train_ddpm_vqvae.py was not enough, and it was still learning. So I changed to 500, but I got the same result. I thought there might be another problem. Below is the generated (denoised) image for 500 epochs.

As you can see, nothing has changed.
"I would like to share my configuration here and would greatly appreciate it if you could take a look and provide your comments on the issue I'm facing. Do you think increasing the learning rate might help?"

dataset_config:
im_path: '/cta/users/undergrad2/StableDiffusion-PyTorch/data/Breast/Breast-img'
im_channels : 3
im_size : 256
name: 'Breast'

diffusion_params:
num_timesteps : 1000
beta_start: 0.00085
beta_end: 0.012

ldm_params:
down_channels: [ 256, 384, 512, 768 ]
mid_channels: [ 768, 512 ]
down_sample: [ True, True, True ]
attn_down : [True, True, True]
time_emb_dim: 512
norm_channels: 32
num_heads: 16
conv_out_channels : 128
num_down_layers : 2
num_mid_layers : 2
num_up_layers : 2
condition_config:
condition_types: None #[ 'classes' ]
class_condition_config:
cond_drop_prob : 0.1
class_condition_l : 11
text_condition_config:
text_embed_model: 'clip'
train_text_embed_model: False
text_embed_dim: 512
cond_drop_prob: 0.1
image_condition_config:
image_condition_input_channels: 18
image_condition_output_channels: 3
image_condition_h : 512
image_condition_w : 512
cond_drop_prob: 0.1

autoencoder_params:
z_channels: 3
codebook_size : 8192
down_channels : [64, 128, 256, 256]
mid_channels : [256, 256]
down_sample : [True, True, True]
attn_down : [False, False, False]
norm_channels: 32
num_heads: 4
num_down_layers : 2
num_mid_layers : 2
num_up_layers : 2
num_headblocks : 1

train_params:
seed : 1111
task_name: '/cta/users/undergrad2/StableDiffusion-PyTorch/Breast_training'
ldm_batch_size: 16
autoencoder_batch_size: 8
disc_start: 10
disc_weight: 0.5
codebook_weight: 1
commitment_beta: 0.2
perceptual_weight: 1
kl_weight: 0.000005
ldm_epochs: 500 # 100 it was 100 at first but didnt work
autoencoder_epochs: 100
num_samples: 1
num_grid_rows: 1
ldm_lr: 0.000005
autoencoder_lr: 0.00001
autoencoder_acc_steps: 4
autoencoder_img_save_steps: 256
use_latents: True
save_latents : False #in vqvae if u set it true, in vqvae_latents it will save latents (which are for batches) also u can generate using code
cf_guidance_scale : 1.0
load_ckpt: True

vqvae_latent_dir_name: '/cta/users/undergrad2/StableDiffusion-PyTorch/Breast_training/vqvae_latents'
vae_latent_dir_name: '/cta/users/undergrad2/StableDiffusion-PyTorch/Breast_training/vae_latents'
vqvae_autoencoder_ckpt_name: '/cta/users/undergrad2/StableDiffusion-PyTorch/Breast_training/vqvae_autoencoder_ckpt.pth'
vae_autoencoder_ckpt_name: '/cta/users/undergrad2/StableDiffusion-PyTorch/Breast_training/vae_autoencoder_ckpt.pth'
vqvae_discriminator_ckpt_name: '/cta/users/undergrad2/StableDiffusion-PyTorch/Breast_training/vqvae_discriminator_ckpt.pth'
vae_discriminator_ckpt_name: '/cta/users/undergrad2/StableDiffusion-PyTorch/Breast_training/vae_discriminator_ckpt.pth'
ldm_ckpt_name: '/cta/users/undergrad2/StableDiffusion-PyTorch/Breast_training/ddpm_ckpt_epoch_495.pth'

Thanks a million!

[Ariz-Mohammadi]

I fixed the problem, thank you!

[MahanVeisi8]

I fixed the problem, thank you!

Hi dear Ariz Mohammadi,

May I kindly ask how you managed to solve the issue you mentioned? I am currently facing a somewhat similar problem and would deeply appreciate your insights.

I’m training a model on the CIFAR-10 dataset. The VQVAE part works well, and I’ve checked the restoration—it looks fine. However, when it comes to training the UNet LDM, I initially struggled to find a learning rate that suited my dataset. After some trial and error, I used the Celeb configuration as a reference, adjusted it for my case, and managed to achieve a loss of around 0.08.
Actaully it may even be better to show you the config clas:
`
class Config:
def init(self):
self.config = {
"dataset_params": {
"im_path": "/kaggle/input/cifar10-64x64-resized-via-cai-super-resolution/cifar10-64/train",
"im_channels": 3,
"im_size": 64,
"name": "CIFAR10"
},
"diffusion_params": {
"num_timesteps": 1000,
"beta_start": 0.00085,
"beta_end": 0.012
},
"ldm_params": {
"down_channels": [256, 384, 512, 768],
"mid_channels": [768, 512],
"down_sample": [True, False, False],
"attn_down": [True, True, True],
"time_emb_dim": 512,
"norm_channels": 32,
"num_heads": 16,
"conv_out_channels": 128,
"num_down_layers": 3,
"num_mid_layers": 3,
"num_up_layers": 3,
"condition_config": {
"condition_types": ["class"],
"class_condition_config": {
"num_classes": 10,
"cond_drop_prob": 0.1
}
}
},
"autoencoder_params": {
"z_channels": 128,
"codebook_size": 100,
"down_channels": [32, 64, 128],
"mid_channels": [128, 128],
"down_sample": [True, True],
"attn_down": [False, False],
"norm_channels": 32,
"num_heads": 16,
"num_down_layers": 1,
"num_mid_layers": 1,
"num_up_layers": 1
},
"train_params": {
"seed": 1111,
"task_name": "/kaggle/working/VQVAE/models",
"ldm_batch_size": 64,
"autoencoder_batch_size": 64,
"disc_start": 1000,
"disc_weight": 0.5,
"codebook_weight": 1,
"commitment_beta": 0.25,
"perceptual_weight": 1,
"kl_weight": 0.000005,
"ldm_epochs": 100,
"autoencoder_epochs": 20,
"num_samples": 36,
"num_grid_rows": 6,
"ldm_lr": 0.00005,
"autoencoder_lr": 0.0001,
"autoencoder_acc_steps": 1,
"autoencoder_img_save_steps": 8,
"save_latents": True,
"cf_guidance_scale": 1.0,
"vae_latent_dir_name": "/vae_latents",
"vqvae_latent_dir_name": "/vqvae_latents",
"ldm_ckpt_name": "cifar10_ddpm_ckpt_class_cond.pth",
"vqvae_autoencoder_ckpt_name": "/kaggle/input/vqvae/pytorch/default/1/cifar10_vqvae_autoencoder_ckpt.pth",
"vqvae_discriminator_ckpt_name": "/kaggle/input/vqvae/pytorch/default/1/cifar10_vqvae_discriminator_ckpt.pth",
"lpips_weights_path": "/kaggle/input/kaggleinputdata/pytorch/default/1/vgg.pth",
}
}

`

Despite this, the denoising and sampling results are far from satisfactory. It has been quite discouraging, especially after dedicating about three weeks to this work.

As you can see, i think my results of X0 are somehow alike yourself's even if the datasets are different.
Here is the sampling code:
`py
import torch
import torchvision
import argparse
import yaml
import os
from torchvision.utils import make_grid
from PIL import Image
from tqdm import tqdm

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def sample(model, scheduler, train_config, diffusion_model_config,
autoencoder_model_config, diffusion_config, dataset_config, vae):
r"""
Sample stepwise by going backward one timestep at a time.
We save the x0 predictions
"""
im_size = dataset_config['im_size'] // 2 ** sum(autoencoder_model_config['down_sample'])

########### Sample random noise latent ##########
xt = torch.randn((train_config['num_samples'],
                  autoencoder_model_config['z_channels'],
                  im_size,
                  im_size)).to(device)
###############################################

############# Validate the config #################
condition_config = get_config_value(diffusion_model_config, key='condition_config', default_value=None)
assert condition_config is not None, ("This sampling script is for class conditional "
                                      "but no conditioning config found")
condition_types = get_config_value(condition_config, 'condition_types', [])
assert 'class' in condition_types, ("This sampling script is for class conditional "
                                      "but no class condition found in config")
validate_class_config(condition_config)
###############################################

############ Create Conditional input ###############
num_classes = condition_config['class_condition_config']['num_classes']
sample_classes = torch.randint(0, num_classes, (train_config['num_samples'], ))
print('Generating images for {}'.format(list(sample_classes.numpy())))
cond_input = {
    'class': torch.nn.functional.one_hot(sample_classes, num_classes).to(device)
}
# Unconditional input for classifier free guidance
uncond_input = {
    'class': cond_input['class'] * 0
}
###############################################

# By default classifier free guidance is disabled
# Change value in config or change default value here to enable it
cf_guidance_scale = get_config_value(train_config, 'cf_guidance_scale', 1.0)

################# Sampling Loop ########################
for i in tqdm(reversed(range(diffusion_config['num_timesteps']))):
    # Get prediction of noise
    t = (torch.ones((xt.shape[0],))*i).long().to(device)
    noise_pred_cond = model(xt, t, cond_input)
    
    if cf_guidance_scale > 1:
        noise_pred_uncond = model(xt, t, uncond_input)
        noise_pred = noise_pred_uncond + cf_guidance_scale*(noise_pred_cond - noise_pred_uncond)
    else:
        noise_pred = noise_pred_cond
    
    # Use scheduler to get x0 and xt-1
    xt, x0_pred = scheduler.sample_prev_timestep(xt, noise_pred, torch.as_tensor(i).to(device))
    
    # if i == 0:
    #     # Decode ONLY the final image to save time
    #     ims = vae.decode(xt)
    # else:
    #     ims = x0_pred
    
    ims = vae.decode(xt)
    ims = torch.clamp(ims, -1., 1.).detach().cpu()
    ims = (ims + 1) / 2
    grid = make_grid(ims, nrow=1)
    img = torchvision.transforms.ToPILImage()(grid)
    
    if not os.path.exists(os.path.join(train_config['task_name'], 'cond_class_samples')):
        os.mkdir(os.path.join(train_config['task_name'], 'cond_class_samples'))
    img.save(os.path.join(train_config['task_name'], 'cond_class_samples', 'x0_{}.png'.format(i)))
    img.close()
##############################################################

def infer():
# Read the config file #
config_instance = Config()
config = config_instance.config

diffusion_config = config['diffusion_params']
dataset_config = config['dataset_params']
diffusion_model_config = config['ldm_params']
autoencoder_model_config = config['autoencoder_params']
train_config = config['train_params']
print(config)

########## Create the noise scheduler #############
scheduler = LinearNoiseScheduler(num_timesteps=diffusion_config['num_timesteps'],
                                 beta_start=diffusion_config['beta_start'],
                                 beta_end=diffusion_config['beta_end'])
###############################################

########## Load Unet #############
model = Unet(im_channels=autoencoder_model_config['z_channels'],
             model_config=diffusion_model_config).to(device)
model.eval()
if os.path.exists(os.path.join(train_config['task_name'],
                               train_config['ldm_ckpt_name'])):
    print('Loaded unet checkpoint')
    model.load_state_dict(torch.load(os.path.join(train_config['task_name'],
                                                  train_config['ldm_ckpt_name']),
                                     map_location=device))
else:
    raise Exception('Model checkpoint {} not found'.format(os.path.join(train_config['task_name'],
                                                                        train_config['ldm_ckpt_name'])))
#####################################

# Create output directories
if not os.path.exists(train_config['task_name']):
    os.mkdir(train_config['task_name'])

########## Load VQVAE #############
vae = VQVAE(im_channels=dataset_config['im_channels'],
            model_config=autoencoder_model_config).to(device)
vae.eval()

# Load vae if found
if os.path.exists(os.path.join(train_config['vqvae_autoencoder_ckpt_name'])):
    print('Loaded vae checkpoint')
    # Load the checkpoint first
    checkpoint = torch.load(train_config['vqvae_autoencoder_ckpt_name'], map_location=device)
    # Now load the state_dict into the model
    vae.load_state_dict(checkpoint)
else:
    raise Exception('VAE checkpoint {} not found'.format(os.path.join(train_config['vqvae_autoencoder_ckpt_name'])))
#####################################

with torch.no_grad():
    sample(model, scheduler, train_config, diffusion_model_config,
           autoencoder_model_config, diffusion_config, dataset_config, vae)

infer()

`

I’ve attached an example image generated after training using my current setup. It doesn’t look good at all. Any guidance you can offer—whether related to your loss values or how you resolved your own issue—would mean a lot to me.

Thank you so much in advance for your time and help!

[MahanVeisi8]

@Ariz-Mohammadi