| import torch, os, glob, random, copy
|
| import torch.nn.functional as F
|
| from torch.utils.data import DataLoader
|
| import torch.distributed as dist
|
| from torch.nn.parallel import DistributedDataParallel as DDP
|
| import numpy as np
|
| from argparse import ArgumentParser
|
| from time import time
|
| from tqdm import tqdm
|
| from omegaconf import OmegaConf
|
| from dataset import RealESRGANDataset, RealESRGANDegrader
|
| from model import Net
|
| from ram.models.ram_lora import ram
|
| from torchvision import transforms
|
| from utils import add_lora_to_unet
|
|
|
| dist.init_process_group(backend="nccl", init_method="env://")
|
| rank = dist.get_rank()
|
| world_size = dist.get_world_size()
|
|
|
| parser = ArgumentParser()
|
| parser.add_argument("--epoch", type=int, default=200)
|
| parser.add_argument("--batch_size", type=int, default=12)
|
| parser.add_argument("--learning_rate", type=float, default=1e-4)
|
| parser.add_argument("--model_dir", type=str, default="weight")
|
| parser.add_argument("--log_dir", type=str, default="log")
|
| parser.add_argument("--save_interval", type=int, default=10)
|
|
|
| args = parser.parse_args()
|
|
|
|
|
| seed = rank
|
| random.seed(seed)
|
| np.random.seed(seed)
|
| torch.manual_seed(seed)
|
| torch.cuda.manual_seed_all(seed)
|
|
|
| config = OmegaConf.load("config.yml")
|
|
|
| epoch = args.epoch
|
| learning_rate = args.learning_rate
|
| bsz = args.batch_size
|
|
|
| device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu")
|
| torch.backends.cudnn.allow_tf32 = True
|
| torch.backends.cuda.matmul.allow_tf32 = True
|
|
|
| if rank == 0:
|
| print("batch size per gpu =", bsz)
|
|
|
| from diffusers import StableDiffusionPipeline
|
| model_id = "stabilityai/stable-diffusion-2-1-base"
|
| pipe = StableDiffusionPipeline.from_pretrained(model_id).to(device)
|
|
|
| vae = pipe.vae
|
| tokenizer = pipe.tokenizer
|
| unet = pipe.unet
|
| text_encoder = pipe.text_encoder
|
|
|
| unet_D = copy.deepcopy(unet)
|
| new_conv_in = torch.nn.Conv2d(256, 320, 3, padding=1).to(device)
|
| new_conv_in.weight.data = unet_D.conv_in.weight.data.repeat(1, 64, 1, 1) / 64
|
| new_conv_in.bias.data = unet_D.conv_in.bias.data
|
| unet_D.conv_in = new_conv_in
|
| unet_D = add_lora_to_unet(unet_D)
|
| unet_D.set_adapters(["default_encoder", "default_decoder", "default_others"])
|
|
|
| vae_teacher = copy.deepcopy(vae)
|
| unet_teacher = copy.deepcopy(unet)
|
|
|
| osediff = torch.load("./weight/pretrained/osediff.pkl", weights_only=False)
|
| vae_teacher.load_state_dict(osediff["vae"])
|
| unet_teacher.load_state_dict(osediff["unet"])
|
|
|
| from diffusers.models.autoencoders.vae import Decoder
|
| ckpt_halfdecoder = torch.load("./weight/pretrained/halfDecoder.ckpt", weights_only=False)
|
| decoder = Decoder(in_channels=4,
|
| out_channels=3,
|
| up_block_types=["UpDecoderBlock2D" for _ in range(4)],
|
| block_out_channels=[64, 128, 256, 256],
|
| layers_per_block=2,
|
| norm_num_groups=32,
|
| act_fn="silu",
|
| norm_type="group",
|
| mid_block_add_attention=True).to(device)
|
| decoder_ckpt = {}
|
| for k, v in ckpt_halfdecoder["state_dict"].items():
|
| if "decoder" in k:
|
| new_k = k.replace("decoder.", "")
|
| decoder_ckpt[new_k] = v
|
| decoder.load_state_dict(decoder_ckpt, strict=True)
|
|
|
| ram_transforms = transforms.Compose([
|
| transforms.Resize((384, 384)),
|
| transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
| ])
|
|
|
| DAPE = ram(pretrained="./weight/pretrained/ram_swin_large_14m.pth",
|
| pretrained_condition="./weight/pretrained/DAPE.pth",
|
| image_size=384,
|
| vit="swin_l").eval().to(device)
|
|
|
| vae.requires_grad_(False)
|
| unet.requires_grad_(False)
|
| text_encoder.requires_grad_(False)
|
| vae_teacher.requires_grad_(False)
|
| unet_teacher.requires_grad_(False)
|
| decoder.requires_grad_(False)
|
| DAPE.requires_grad_(False)
|
|
|
| model = DDP(Net(unet, copy.deepcopy(decoder)).to(device), device_ids=[rank])
|
| model_D = DDP(unet_D.to(device), device_ids=[rank])
|
| model.requires_grad_(True)
|
| model_D.requires_grad_(False)
|
| params_to_opt = []
|
| for n, p in model_D.named_parameters():
|
| if "lora" in n or "conv_in" in n:
|
| p.requires_grad = True
|
| params_to_opt.append(p)
|
|
|
| if rank == 0:
|
| param_cnt = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
| print("#Param.", param_cnt/1e6, "M")
|
|
|
| dataset = RealESRGANDataset(config, bsz)
|
| degrader = RealESRGANDegrader(config, device)
|
| dataloader = DataLoader(dataset, batch_size=bsz, num_workers=8)
|
| optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
|
| optimizer_D = torch.optim.Adam(params_to_opt, lr=1e-6)
|
| scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[100,], gamma=0.5)
|
| scaler = torch.cuda.amp.GradScaler()
|
|
|
| model_dir = "./%s" % (args.model_dir,)
|
| log_path = "./%s/log.txt" % (args.log_dir,)
|
| os.makedirs(model_dir, exist_ok=True)
|
| os.makedirs(args.log_dir, exist_ok=True)
|
|
|
| print("start training...")
|
| timesteps = torch.tensor([999], device=device).long().expand(bsz,)
|
| alpha = pipe.scheduler.alphas_cumprod[999]
|
| for epoch_i in range(1, epoch + 1):
|
| start_time = time()
|
| loss_avg = 0.0
|
| loss_distil_avg = 0.0
|
| loss_adv_avg = 0.0
|
| loss_D_avg = 0.0
|
| iter_num = 0
|
| dist.barrier()
|
| for batch in tqdm(dataloader):
|
| with torch.cuda.amp.autocast(enabled=True):
|
| with torch.no_grad():
|
| LR, HR = degrader.degrade(batch)
|
| text_input = tokenizer(DAPE.generate_tag(ram_transforms(LR))[0],
|
| max_length=tokenizer.model_max_length,
|
| padding="max_length", return_tensors="pt").to(device)
|
| encoder_hidden_states = text_encoder(text_input.input_ids, return_dict=False)[0]
|
| LR, HR = LR * 2 - 1, HR * 2 - 1
|
| LR_ = F.interpolate(LR, scale_factor=4, mode="bicubic")
|
| LR_latents = vae_teacher.encode(LR_).latent_dist.mean * vae_teacher.config.scaling_factor
|
| HR_latents = vae.encode(HR).latent_dist.mean
|
| pred_teacher = unet_teacher(
|
| LR_latents,
|
| timesteps,
|
| encoder_hidden_states=encoder_hidden_states,
|
| return_dict=False,
|
| )[0]
|
| z0_teacher = (LR_latents-((1-alpha)**0.5)*pred_teacher)/(alpha**0.5)
|
| z0_teacher = vae_teacher.post_quant_conv(z0_teacher / vae_teacher.config.scaling_factor)
|
| z0_teacher = decoder.conv_in(z0_teacher)
|
| z0_teacher = decoder.mid_block(z0_teacher)
|
| z0_gt = vae.post_quant_conv(HR_latents)
|
| z0_gt = decoder.conv_in(z0_gt)
|
| z0_gt = decoder.mid_block(z0_gt)
|
| z0_student = model(LR)
|
| loss_distil = (z0_student - z0_teacher).abs().mean()
|
| loss_adv = F.softplus(-model_D(
|
| z0_student,
|
| timesteps,
|
| encoder_hidden_states=encoder_hidden_states,
|
| return_dict=False,
|
| )[0]).mean()
|
| loss = loss_distil + loss_adv
|
| optimizer.zero_grad(set_to_none=True)
|
| scaler.scale(loss).backward()
|
| scaler.step(optimizer)
|
| scaler.update()
|
| with torch.cuda.amp.autocast(enabled=True):
|
| pred_real = model_D(
|
| z0_gt.detach(),
|
| timesteps,
|
| encoder_hidden_states=encoder_hidden_states,
|
| return_dict=False,
|
| )[0]
|
| pred_fake = model_D(
|
| z0_student.detach(),
|
| timesteps,
|
| encoder_hidden_states=encoder_hidden_states,
|
| return_dict=False,
|
| )[0]
|
| loss_D = F.softplus(pred_fake).mean() + F.softplus(-pred_real).mean()
|
| optimizer_D.zero_grad(set_to_none=True)
|
| scaler.scale(loss_D).backward()
|
| scaler.step(optimizer_D)
|
| scaler.update()
|
| loss_avg += loss.item()
|
| loss_distil_avg += loss_distil.item()
|
| loss_adv_avg += loss_adv.item()
|
| loss_D_avg += loss_D.item()
|
| iter_num += 1
|
|
|
|
|
|
|
|
|
| scheduler.step()
|
| loss_avg /= iter_num
|
| loss_distil_avg /= iter_num
|
| loss_adv_avg /= iter_num
|
| loss_D_avg /= iter_num
|
| log_data = "[%d/%d] Average loss: %f, distil loss: %f, adv loss: %f, D loss: %f, time cost: %.2fs, cur lr is %f." % (epoch_i, epoch, loss_avg, loss_distil_avg, loss_adv_avg, loss_D_avg, time() - start_time, scheduler.get_last_lr()[0])
|
| if rank == 0:
|
| print(log_data)
|
| with open(log_path, "a") as log_file:
|
| log_file.write(log_data + "\n")
|
| if epoch_i % args.save_interval == 0:
|
| torch.save(model.state_dict(), "./%s/net_params_%d.pkl" % (model_dir, epoch_i))
|
|
|
