train_model.py

from datasets import load_dataset
from PIL.Image import Image
import PIL
from PIL.Image import Resampling
import numpy as np
from rct_diffusion_pipeline import RCTDiffusionPipeline
import torch
import torchvision.transforms as T
import torch.nn.functional as F
from diffusers.optimization import get_cosine_schedule_with_warmup
from tqdm.auto import tqdm
from accelerate import Accelerator
from diffusers import DDPMScheduler, UNet2DConditionModel, AutoencoderKL

SAMPLE_SIZE = 256
LATENT_SIZE = 32
SAMPLE_NUM_CHANNELS = 3
LATENT_NUM_CHANNELS = 4

def save_and_test(pipeline, epoch):
    outputs = pipeline([[('aleppo pine tree', 1.0)]], [[('dark green', 1.0)]])
    for image_index in range(len(outputs)):
        file_name = f'out{image_index}_{epoch}.png'
        outputs[image_index].save(file_name)
    
    model_file = f'rct_foliage_{epoch}.pth'
    pipeline.save_pretrained(model_file)

def convert_images(dataset):
        # let's get all the entries for the 4 views split in four lists
    views = []
    num_images = int(dataset.num_rows / 4)

    for view_index in range(4):
        entries = [entry for entry in dataset if entry['view'] == view_index]
        views.append(entries)
    
    # convert those images to 256x256 by cropping and scaling up the image
    image_views = []
    for view_index in range(4):
        images = []
        for entry in views[view_index]:
            image = entry['image']
            
            scale_factor = int(np.minimum(SAMPLE_SIZE / image.width, SAMPLE_SIZE / image.height))
            image = Image.resize(image, size=(scale_factor * image.width, scale_factor * image.height), resample=Resampling.NEAREST)
            
            new_image = PIL.Image.new('RGB', (SAMPLE_SIZE, SAMPLE_SIZE))
            new_image.paste(image, box=(int((SAMPLE_SIZE - image.width)/2), int((SAMPLE_SIZE - image.height)/2)))
            images.append(new_image)
        image_views.append(images)
    
    del views
    
    # convert those views in tensors
    targets = torch.Tensor(size=(num_images, 4, SAMPLE_NUM_CHANNELS, SAMPLE_SIZE, SAMPLE_SIZE)).to(dtype=torch.float16)
    pillow_to_tensor = T.ToTensor()

    for image_index in range(num_images):
        for view_index in range(4):
            targets[image_index, view_index] = pillow_to_tensor(image_views[view_index][image_index]).to(dtype=torch.float16)
    del image_views
    del entries

    return torch.reshape(targets, (num_images, 4 * SAMPLE_NUM_CHANNELS, SAMPLE_SIZE, SAMPLE_SIZE))

def convert_labels(dataset, model, num_images):
    # get the labels
    view0_entries = [row for row in dataset if row['view'] == 0]
    obj_descriptions = [row['object_description'] for row in view0_entries]
    colors1 = [row['color1'] for row in view0_entries]
    colors2 = [row['color2'] for row in view0_entries]
    colors3 = [row['color3'] for row in view0_entries]

    del view0_entries

    # convert those descriptions, color1, color2 and color3 to a list of tuple with label and weight=1.0
    obj_descriptions = [[(obj_desc, 1.0)] for obj_desc in obj_descriptions]
    colors1 = [[(color1, 1.0)] for color1 in colors1]
    colors2 = [[(color2, 1.0)] for color2 in colors2]
    colors3 = [[(color3, 1.0)] for color3 in colors3]

    # convert those tuples in numpy arrays using the helper function of the model
    
    obj_descriptions = [model.get_object_description_weights(obj_desc) for obj_desc in obj_descriptions]
    colors1 = [model.get_color1_weights(color1) for color1 in colors1]
    colors2 = [model.get_color2_weights(color2) for color2 in colors2]
    colors3 = [model.get_color3_weights(color3) for color3 in colors3]

    # finally, convert those numpy arrays to a tensor
    class_labels = model.pack_labels_to_tensor(num_images, obj_descriptions, colors1, colors2, colors3)
    del obj_descriptions
    del colors1
    del colors2
    del colors3
    del dataset
    return class_labels.to(dtype=torch.float16, device='cuda')

def train_model(batch_size=4, epochs=100, scheduler_num_timesteps=20, save_model_interval=10, start_learning_rate=1e-3, lr_warmup_steps=500):
    dataset = load_dataset('frutiemax/rct_dataset')
    dataset = dataset['train']

    targets = convert_images(dataset)
    num_images = int(dataset.num_rows / 4)

    unet = UNet2DConditionModel(sample_size=LATENT_SIZE, in_channels=LATENT_NUM_CHANNELS * 4, out_channels=LATENT_NUM_CHANNELS * 4, \
                        down_block_types=('CrossAttnDownBlock2D', 'CrossAttnDownBlock2D', 'DownBlock2D'),\
                              up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D', 'CrossAttnUpBlock2D'), cross_attention_dim=160,
                            block_out_channels=(64, 128, 256), norm_num_groups=32)
    unet = unet.to(dtype=torch.float16)
    scheduler = DDPMScheduler(num_train_timesteps=20)
    vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", use_safetensors=True)
    vae = vae.to(dtype=torch.float16)

    optimizer = torch.optim.Adam(unet.parameters(), lr=start_learning_rate)
    lr_scheduler = get_cosine_schedule_with_warmup(
        optimizer=optimizer,
        num_warmup_steps=lr_warmup_steps,
        num_training_steps=num_images * epochs
    )
    model = RCTDiffusionPipeline(unet, scheduler, vae)
    model.load_dictionaries_from_dataset()
    labels = convert_labels(dataset, model, num_images)
    del model

    # lets train for 100 epoch for each sprite in the dataset with a random noise level
    progress_bar = tqdm(total=epochs)
    accelerator = Accelerator(mixed_precision='fp16')
    unet, scheduler, lr_scheduler, vae = accelerator.prepare(unet, scheduler, lr_scheduler, vae)

    for epoch in range(epochs):
        # create a noisy version of each sprite
        for batch_index in range(0, num_images, batch_size):
            progress_bar.set_description(f'epoch={epoch}, batch_index={batch_index}')
            batch_end = np.minimum(num_images, batch_index + batch_size)
            clean_images = targets[batch_index:batch_end]
            clean_images = torch.reshape(clean_images, ((batch_end - batch_index), SAMPLE_NUM_CHANNELS * 4, SAMPLE_SIZE, SAMPLE_SIZE)).to(device='cuda', dtype=torch.float16)

            noise = torch.randn(clean_images.shape, dtype=torch.float16, device='cuda')
            timesteps = torch.randint(0, scheduler.config.num_train_timesteps, (batch_end - batch_index, )).to(device='cuda')
            #timesteps = timesteps.to(dtype=torch.int, device='cuda')
            noisy_images = scheduler.add_noise(clean_images, noise, timesteps)
            del clean_images

            # encode through the vae
            with accelerator.accumulate(unet):
                latent_images = torch.Tensor(size=(batch_end - batch_index, LATENT_NUM_CHANNELS * 4, LATENT_SIZE, LATENT_SIZE)).to(device='cuda', dtype=torch.float16)
                latent_noises = torch.Tensor(size=(batch_end - batch_index, LATENT_NUM_CHANNELS * 4, LATENT_SIZE, LATENT_SIZE)).to(device='cuda', dtype=torch.float16)
                for view_index in range(4):
                    images = noisy_images[:, view_index*SAMPLE_NUM_CHANNELS:(view_index+1)*SAMPLE_NUM_CHANNELS]
                    result = vae.encode(images).latent_dist.sample()
                    latent_images[:, view_index*LATENT_NUM_CHANNELS:(view_index+1)*LATENT_NUM_CHANNELS] = result

                    images = noise[:, view_index*SAMPLE_NUM_CHANNELS:(view_index+1)*SAMPLE_NUM_CHANNELS]
                    result = vae.encode(images).latent_dist.sample()
                    latent_noises[:, view_index*LATENT_NUM_CHANNELS:(view_index+1)*LATENT_NUM_CHANNELS] = result

                del noise
                del noisy_images
                unet_results = unet(latent_images, timesteps, labels[batch_index:batch_end])[0]
                unet_results = unet_results.to(dtype=torch.float16)

                loss = F.mse_loss(unet_results, latent_noises)
                accelerator.backward(loss)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()
        
        if (epoch + 1) % save_model_interval == 0:
            model = RCTDiffusionPipeline(accelerator.unwrap_model(unet), scheduler, vae)
            save_and_test(model, epoch)
        progress_bar.update(1)


if __name__ == '__main__':
    train_model(1, save_model_interval=1)