train.py

import argparse
import os
import math
import random
from pathlib import Path
import h5py
from PIL import Image

import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

from transformers import CLIPTokenizer, CLIPTextModel
from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler
from diffusers import logging as diffusers_logging
from accelerate import Accelerator

diffusers_logging.set_verbosity_error()


def modify_vae_to_rgba(vae):
    enc = vae.encoder
    dec = vae.decoder

    old_conv_in = enc.conv_in
    new_conv_in = torch.nn.Conv2d(
        in_channels=4,
        out_channels=old_conv_in.out_channels,
        kernel_size=old_conv_in.kernel_size,
        stride=old_conv_in.stride,
        padding=old_conv_in.padding,
        bias=(old_conv_in.bias is not None),
    )

    with torch.no_grad():
        new_conv_in.weight[:, :3, :, :] = old_conv_in.weight
        new_conv_in.weight[:, 3:4, :, :] = old_conv_in.weight.mean(dim=1, keepdim=True)
        if old_conv_in.bias is not None:
            new_conv_in.bias.copy_(old_conv_in.bias)

    enc.conv_in = new_conv_in

    old_conv_out = dec.conv_out
    new_conv_out = torch.nn.Conv2d(
        in_channels=old_conv_out.in_channels,
        out_channels=4,
        kernel_size=old_conv_out.kernel_size,
        stride=old_conv_out.stride,
        padding=old_conv_out.padding,
        bias=(old_conv_out.bias is not None),
    )

    with torch.no_grad():
        new_conv_out.weight[:3, :, :, :] = old_conv_out.weight
        new_conv_out.weight[3:4, :, :, :] = old_conv_out.weight.mean(
            dim=0, keepdim=True
        )
        if old_conv_out.bias is not None:
            new_conv_out.bias[:3].copy_(old_conv_out.bias)
            new_conv_out.bias[3] = old_conv_out.bias.mean()

    dec.conv_out = new_conv_out

    return vae


class H5RGBAKeyDataset(Dataset):
    def __init__(self, h5path, resolution=64):
        self.h5 = h5py.File(h5path, "r")
        self.images = self.h5["images"]
        self.names = self.h5["name"]
        self.descs = self.h5["description"]
        self.resolution = resolution
        self.length = len(self.images)

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        img = self.images[idx]
        name = self.names[idx]
        desc = self.descs[idx]

        if isinstance(name, bytes):
            name = name.decode("utf-8", errors="ignore")
        if isinstance(desc, bytes):
            desc = desc.decode("utf-8", errors="ignore")

        prompt = f"{name}, {desc}"

        rgba = torch.from_numpy(img).permute(2, 0, 1).float() / 127.5 - 1.0
        rgba = F.interpolate(
            rgba.unsqueeze(0), size=(self.resolution, self.resolution), mode="nearest"
        ).squeeze(0)

        return rgba, prompt


def collate_fn(batch):
    imgs = torch.stack([b[0] for b in batch])
    prompts = [b[1] for b in batch]
    return imgs, prompts


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--h5", required=True)
    parser.add_argument("--pretrained_model", default="runwayml/stable-diffusion-v1-5")
    parser.add_argument("--output_dir", default="./sd_rgba_out")
    parser.add_argument("--resolution", type=int, default=64)
    parser.add_argument("--batch_size", type=int, default=64)
    parser.add_argument("--lr", type=float, default=1e-5)
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--device", default=None)
    parser.add_argument("--save_every_steps", type=int, default=5000)
    args = parser.parse_args()

    accelerator = Accelerator(mixed_precision="fp16")
    device = accelerator.device

    tokenizer = CLIPTokenizer.from_pretrained(
        args.pretrained_model, subfolder="tokenizer", use_fast=True
    )
    text_encoder = CLIPTextModel.from_pretrained(
        args.pretrained_model, subfolder="text_encoder"
    )
    text_encoder.eval()
    for p in text_encoder.parameters():
        p.requires_grad = False

    vae = AutoencoderKL.from_pretrained(args.pretrained_model, subfolder="vae")
    vae = modify_vae_to_rgba(vae)
    vae.requires_grad_(False)
    vae.to(device)

    unet = UNet2DConditionModel.from_pretrained(args.pretrained_model, subfolder="unet")
    unet.to(device)

    noise_scheduler = DDPMScheduler.from_pretrained(
        args.pretrained_model, subfolder="scheduler"
    )

    dataset = H5RGBAKeyDataset(args.h5, resolution=args.resolution)
    dataloader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=True,
        collate_fn=collate_fn,
        drop_last=True,
        num_workers=4,
    )

    optimizer = torch.optim.AdamW(unet.parameters(), lr=args.lr, weight_decay=0.01)

    unet, text_encoder, optimizer, dataloader = accelerator.prepare(
        unet, text_encoder, optimizer, dataloader
    )

    global_step = 0
    total_steps = math.ceil(len(dataloader) * args.epochs)
    for epoch in range(args.epochs):
        for batch in dataloader:
            rgba_imgs, prompts = batch
            rgba_imgs = rgba_imgs.to(device, dtype=torch.float32)
            with torch.no_grad():
                vae.eval()
                latents = vae.encode(rgba_imgs.float()).latent_dist.sample()
                latents = latents * vae.config.scaling_factor

            batch_size = latents.shape[0]
            timesteps = torch.randint(
                0,
                noise_scheduler.config.num_train_timesteps,
                (batch_size,),
                device=device,
                dtype=torch.long,
            )

            noise = torch.randn_like(latents)
            noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)

            text_inputs = tokenizer(
                prompts,
                padding="max_length",
                truncation=True,
                max_length=tokenizer.model_max_length,
                return_tensors="pt",
            )
            input_ids = text_inputs.input_ids.to(device)
            attention_mask = text_inputs.attention_mask.to(device)
            with torch.no_grad():
                text_encoder.eval()
                encoder_hidden_states = text_encoder(
                    input_ids=input_ids, attention_mask=attention_mask
                )[0]

            model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
            loss = F.mse_loss(model_pred, noise)

            accelerator.backward(loss)
            optimizer.step()
            optimizer.zero_grad()
            global_step += 1

            if global_step % 100 == 0:
                print(f"step {global_step}/{total_steps} loss {loss.item():.6f}")

            if global_step % args.save_every_steps == 0:
                accelerator.wait_for_everyone()
                if accelerator.is_main_process:
                    outdir = Path(args.output_dir) / f"checkpoint-{global_step}"
                    outdir.mkdir(parents=True, exist_ok=True)
                    unet_to_save = accelerator.unwrap_model(unet)
                    unet_to_save.save_pretrained(outdir)
                    vae.save_pretrained(outdir / "vae")
                    text_encoder.save_pretrained(outdir / "text_encoder")
                    tokenizer.save_pretrained(outdir / "tokenizer")
                    print("Saved checkpoint to", outdir)

    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
        final_dir = Path(args.output_dir) / "final_pipeline"
        final_dir.mkdir(parents=True, exist_ok=True)
        unet_to_save = accelerator.unwrap_model(unet)
        unet_to_save.save_pretrained(final_dir / "unet")
        vae.save_pretrained(final_dir / "vae")
        text_encoder.save_pretrained(final_dir / "text_encoder")
        tokenizer.save_pretrained(final_dir / "tokenizer")
        print("Saved final pipeline to", final_dir)


if __name__ == "__main__":
    main()