|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from models import iresnet
|
|
from lpips.lpips import LPIPS
|
|
from pytorch_msssim import SSIM
|
|
|
|
|
|
def adopt_weight(weight, global_step, threshold=0, value=0.):
|
|
if global_step < threshold:
|
|
weight = value
|
|
return weight
|
|
|
|
|
|
def hinge_d_loss(logits_real, logits_fake):
|
|
loss_real = torch.mean(F.relu(1. - logits_real))
|
|
loss_fake = torch.mean(F.relu(1. + logits_fake))
|
|
d_loss = 0.5 * (loss_real + loss_fake)
|
|
return d_loss
|
|
|
|
|
|
def mse_d_loss(logits_real, logits_fake):
|
|
loss_real = torch.mean((logits_real - 1.) ** 2)
|
|
loss_fake = torch.mean(logits_fake ** 2)
|
|
d_loss = 0.5 * (loss_real + loss_fake)
|
|
return d_loss
|
|
|
|
|
|
def vanilla_d_loss(logits_real, logits_fake):
|
|
d_loss = 0.5 * (
|
|
torch.mean(torch.nn.functional.softplus(-logits_real)) +
|
|
torch.mean(torch.nn.functional.softplus(logits_fake)))
|
|
return d_loss
|
|
|
|
|
|
def create_fr_model(model_path, depth="100"):
|
|
model = iresnet(depth)
|
|
model.load_state_dict(torch.load(model_path))
|
|
|
|
return model
|
|
|
|
|
|
def downscale(img: torch.tensor):
|
|
half_size = img.shape[-1] // 8
|
|
img = F.interpolate(img, size=(half_size, half_size), mode='bicubic', align_corners=False)
|
|
return img
|
|
|
|
|
|
class VQLPIPSWithDiscriminator(nn.Module):
|
|
def __init__(self, disc_start=1000, disc_factor=1.0, disc_weight=1.0,
|
|
disc_conditional=False, disc_loss="mse", id_loss="mse",
|
|
fr_model="./models/arcface-r100-glint360k.pth"):
|
|
super().__init__()
|
|
assert disc_loss in ["hinge", "vanilla", "mse", "smooth"]
|
|
self.loss_name = disc_loss
|
|
self.perceptual_loss = LPIPS().eval()
|
|
self.discriminator_iter_start = disc_start
|
|
if disc_loss == "hinge":
|
|
self.disc_loss = hinge_d_loss
|
|
elif disc_loss == "vanilla":
|
|
self.disc_loss = vanilla_d_loss
|
|
elif disc_loss == "mse":
|
|
self.disc_loss = mse_d_loss
|
|
else:
|
|
raise ValueError(f"Unknown GAN loss '{disc_loss}'.")
|
|
print(f"VQLPIPSWithDiscriminator running with {disc_loss} loss.")
|
|
self.fr_model = create_fr_model(fr_model).eval()
|
|
if id_loss == "mse":
|
|
self.feature_loss = nn.MSELoss()
|
|
elif id_loss == "cosine":
|
|
self.feature_loss = nn.CosineSimilarity()
|
|
self.disc_factor = disc_factor
|
|
self.discriminator_weight = disc_weight
|
|
self.disc_conditional = disc_conditional
|
|
self.ssim_loss = SSIM(data_range=1, size_average=True, channel=3)
|
|
|
|
def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer=None):
|
|
if last_layer is not None:
|
|
nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0]
|
|
g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
|
|
else:
|
|
nll_grads = torch.autograd.grad(nll_loss, self.last_layer[0], retain_graph=True)[0]
|
|
g_grads = torch.autograd.grad(g_loss, self.last_layer[0], retain_graph=True)[0]
|
|
|
|
d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4)
|
|
d_weight = torch.clamp(d_weight, 0.0, 1e4).detach()
|
|
d_weight = d_weight * self.discriminator_weight
|
|
return d_weight
|
|
|
|
def forward(self, im_features, gt_indices, logits, gt_img, image, discriminator, emb_loss,
|
|
epoch, last_layer=None, cond=None, mask=None):
|
|
rec_loss = (image - gt_img) ** 2
|
|
|
|
if epoch >= 0:
|
|
gen_feature = self.fr_model(image)
|
|
feature_loss = torch.mean(1 - torch.cosine_similarity(im_features, gen_feature))
|
|
else:
|
|
feature_loss = 0
|
|
|
|
p_loss = self.perceptual_loss(image, gt_img) * 2
|
|
|
|
with torch.cuda.amp.autocast(enabled=False):
|
|
ssim_loss = 1 - self.ssim_loss((image.float() + 1) / 2, (gt_img + 1) / 2)
|
|
logits_fake = discriminator(image)
|
|
logits_real_d = discriminator(gt_img.detach())
|
|
logits_fake_d = discriminator(image.detach())
|
|
|
|
if mask is None:
|
|
token_loss = (logits[:, 1:, :] - gt_indices[:, 1:, :])
|
|
token_loss = torch.mean(token_loss ** 2)
|
|
else:
|
|
token_loss = torch.abs((logits[:, 1:, :] - gt_indices[:, 1:, :])) * mask[:, 1:, None]
|
|
token_loss = token_loss.sum() / mask[:, 1:].sum()
|
|
|
|
nll_loss = torch.mean(rec_loss + p_loss) + \
|
|
ssim_loss + \
|
|
token_loss + feature_loss + emb_loss
|
|
|
|
g_loss = -torch.mean(logits_fake)
|
|
|
|
d_weight = self.calculate_adaptive_weight(nll_loss, g_loss, last_layer=last_layer)
|
|
disc_factor = adopt_weight(self.disc_factor, epoch, threshold=self.discriminator_iter_start)
|
|
ae_loss = nll_loss + d_weight * disc_factor * g_loss
|
|
|
|
|
|
disc_factor = adopt_weight(self.disc_factor, epoch, threshold=self.discriminator_iter_start)
|
|
d_loss = disc_factor * self.disc_loss(logits_real_d, logits_fake_d)
|
|
return ae_loss, d_loss, token_loss, rec_loss, ssim_loss, p_loss, feature_loss
|
|
|
