train.py

import os
import torch
import random
import torch.nn.functional as F

from torch.optim import AdamW
from accelerate import Accelerator
from diffusers import AutoencoderKL, DDIMScheduler, UNet2DConditionModel
from tqdm import tqdm
from typing import Optional
from torch.utils.data import DataLoader
from torchvision.utils import save_image

from model.attn_processor import SkipAttnProcessor
from model.utils import init_adapter
from args import get_args
from data.vitonhd import VITHONHD
from utils import prepare_image, prepare_mask_image, compute_vae_encodings, compute_dream_and_update_latents_for_inpaint
from model.pipeline import CatVTONPipeline


def init_models(model_root: str,
                weight_dtype: str = "no",
                vae_subfolder: str = "vae",
                device = "cpu"):
    if weight_dtype == "no":
        weight_dtype = torch.float32
    elif weight_dtype == "fp16":
        weight_dtype = torch.float16
    elif weight_dtype == "bf16":
        weight_dtype = torch.bfloat16
    else:
        raise NotImplemented

    print(f"load vae from {vae_subfolder}")
    vae = AutoencoderKL.from_pretrained(model_root, subfolder=vae_subfolder)
    unet = UNet2DConditionModel.from_pretrained(model_root, subfolder="unet")
    try:
        noise_scheduler = DDIMScheduler.from_pretrained(model_root, subfolder="scheduler")
    except Exception as e:
        noise_scheduler = DDIMScheduler.from_pretrained(model_root, subfolder="noise_scheduler")

    init_adapter(unet, cross_attn_cls=SkipAttnProcessor)
    
    vae.to(device)
    unet.to(device)

    vae.requires_grad_(False)
    unet.requires_grad_(False)

    for name, param in unet.named_modules():
        if "attn1" in name:
            param.requires_grad_(True)

    unet.train()
    # unet.enable_gradient_checkpointing()

    optimizer_path = os.path.join(model_root, "optim.pth")
    if os.path.exists(optimizer_path):
        optimizer_state_dict = torch.load(optimizer_path)
    else:
        optimizer_state_dict = None

    return noise_scheduler, vae, unet, optimizer_state_dict


def train_one_step(batch,
                   noise_scheduler,
                   vae,
                   unet,
                   device,
                   extra_condition_key):
    person = prepare_image(batch['person'])
    cloth = prepare_image(batch['cloth'])
    mask = prepare_mask_image(batch['mask'])

    masked_person = person * (mask < 0.5)

    person_latent = compute_vae_encodings(person, vae)  # 加噪
    masked_person_latent = compute_vae_encodings(masked_person, vae)

    if random.random() < 0.15:
        # for cfg
        cloth_latent = torch.zeros_like(masked_person_latent).to(device).to(masked_person_latent.dtype)
    else:
        cloth_latent = compute_vae_encodings(cloth, vae)

    mask_latent = F.interpolate(mask, size=masked_person_latent.shape[-2:], mode="nearest")

    bsz = masked_person_latent.shape[0]

    timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz, )).to(device).long()
    
    first_input_latent = torch.concat([person_latent, cloth_latent], dim=-2)

    noise = torch.randn_like(first_input_latent)

    noisy_first_latent = noise_scheduler.add_noise(first_input_latent, noise, timesteps)

    masked_latent_concat = torch.cat([masked_person_latent, cloth_latent], dim=-2)

    extra_condition = batch.get(extra_condition_key, None)
    extra_condition = F.interpolate(extra_condition, size=mask_latent.shape[-2:], mode="nearest")

    mask_latent_concat = torch.cat([mask_latent, extra_condition], dim=-2)


    inpainting_latent_model_input = torch.cat([noisy_first_latent, mask_latent_concat, masked_latent_concat], dim=1)

    noise_pred = unet(
        inpainting_latent_model_input,
        timesteps,
        encoder_hidden_states=None,
        return_dict=False
    )[0]

    loss = F.mse_loss(noise_pred.float(), noise.float(), reduction="mean")

    return loss


def main():
    args = get_args()

    accelerator = Accelerator(
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        mixed_precision=args.weight_dtype
    )

    device = accelerator.device

    train_dataset = VITHONHD(args.train_data_record_path, args.height, args.width, extra_condition_key=args.extra_condition_key)
    train_dataloader = DataLoader(dataset=train_dataset, batch_size=args.batch_size, num_workers=args.num_workers, shuffle=True)

    noise_scheduler, vae, unet, optimizer_state_dict = init_models(args.model_root, device=device, vae_subfolder=args.vae_subfolder)
    optimizer = AdamW(unet.parameters(), lr=args.lr, betas=(args.beta1, args.beta2), eps=args.eps, weight_decay=args.weight_decay)
    if optimizer_state_dict:
        print("加载优化器状态")
        optimizer.load_state_dict(optimizer_state_dict)

    if accelerator.is_main_process:
        eval_dataset = VITHONHD(args.eval_data_record_path, args.height, args.width, is_train=False, extra_condition_key=args.extra_condition_key)
        eval_dataloader = DataLoader(dataset=eval_dataset, batch_size=1, num_workers=args.num_workers)
    else:
        eval_dataloader = None

    (
        unet,
        optimizer,
        train_dataloader
    ) = accelerator.prepare(
        unet,
        optimizer,
        train_dataloader
    )

    global_step = args.global_steps
    reach_max_steps = False

    progress_bar = tqdm(initial=global_step, total=args.max_steps, disable=not accelerator.is_main_process)
    progress_bar.set_description("train catvton")

    while True:
        if reach_max_steps:
            print("到达最大训练步数，停止训练")
            break
        avg_loss = 0.
        for batch in train_dataloader:
            with accelerator.accumulate(unet):
                with accelerator.autocast():
                    loss = train_one_step(
                        batch,
                        noise_scheduler,
                        vae,
                        unet,
                        device,
                        args.extra_condition_key
                    )

                avg_loss += loss.item()

                accelerator.backward(loss)

                # TODO: 需要关注
                if accelerator.sync_gradients:
                    accelerator.clip_grad_norm_(unet.parameters(), args.max_grad_norm)

                optimizer.step()
                optimizer.zero_grad()
                
            if accelerator.sync_gradients:
                avg_loss = torch.tensor(avg_loss).to(device)
                avg_loss = accelerator.gather(avg_loss).mean().item() / accelerator.gradient_accumulation_steps
                progress_bar.update(1)
                logs = {"step_loss": avg_loss, "global_steps": global_step}
                progress_bar.set_postfix(**logs)

                global_step += 1
                avg_loss = 0.

                # 验证并保存模型
                if global_step % args.logging_steps == 0 or global_step >= args.max_steps:
                    if accelerator.is_main_process:
                        unwrap_unet = accelerator.unwrap_model(unet)
                        unwrap_unet.eval()
                        pipeline = CatVTONPipeline(noise_scheduler, vae, unwrap_unet)
                        
                        os.makedirs(f"../eval_outputs/{args.eval_output_dir}/{global_step}", exist_ok=True)
                        
                        with torch.no_grad():
                            for idx, batch in enumerate(eval_dataloader):
                                if args.extra_condition_key:
                                    sample = pipeline(
                                        image=batch['person'],
                                        condition_image=batch['cloth'],
                                        mask=batch['mask'],
                                        extra_condition=batch[args.extra_condition_key]
                                    )[0]
                                else:
                                    sample = pipeline(
                                        image=batch['person'],
                                        condition_image=batch['cloth'],
                                        mask=batch['mask']
                                    )[0]

                                sample.save(f"../eval_outputs/{args.eval_output_dir}/{global_step}/{idx}.png")

                        save_path = os.path.join(args.output_dir, args.checkpoint_dir)
                        pipeline.save_pretrained(save_path)
                        torch.save(optimizer.state_dict(), f"{save_path}/optim.pth")

                        del pipeline
                        del unwrap_unet
                        torch.cuda.empty_cache()
                        
                    if global_step >= args.max_steps:
                        reach_max_steps = True
                        break


if __name__ == "__main__":
    main()