Spaces:
Runtime error
Runtime error
| from argparse import ( | |
| ArgumentParser, | |
| Namespace, | |
| ) | |
| from typing import ( | |
| Dict, | |
| Iterable, | |
| Optional, | |
| Tuple, | |
| ) | |
| import numpy as np | |
| import torch | |
| from torch import nn | |
| from utils.misc import ( | |
| optional_string, | |
| iterable_to_str, | |
| ) | |
| from .contextual_loss import ContextualLoss | |
| from .color_transfer_loss import ColorTransferLoss | |
| from .regularize_noise import NoiseRegularizer | |
| from .reconstruction import ( | |
| EyeLoss, | |
| FaceLoss, | |
| create_perceptual_loss, | |
| ReconstructionArguments, | |
| ) | |
| class LossArguments: | |
| def add_arguments(parser: ArgumentParser): | |
| ReconstructionArguments.add_arguments(parser) | |
| parser.add_argument("--color_transfer", type=float, default=1e10, help="color transfer loss weight") | |
| parser.add_argument("--eye", type=float, default=0.1, help="eye loss weight") | |
| parser.add_argument('--noise_regularize', type=float, default=5e4) | |
| # contextual loss | |
| parser.add_argument("--contextual", type=float, default=0.1, help="contextual loss weight") | |
| parser.add_argument("--cx_layers", nargs='*', help="contextual loss layers", | |
| choices=['relu1_2', 'relu2_2', 'relu3_4', 'relu4_4', 'relu5_4'], | |
| default=['relu3_4', 'relu2_2', 'relu1_2']) | |
| def to_string(args: Namespace) -> str: | |
| return ( | |
| ReconstructionArguments.to_string(args) | |
| + optional_string(args.eye > 0, f"-eye{args.eye}") | |
| + optional_string(args.color_transfer, f"-color{args.color_transfer:.1e}") | |
| + optional_string( | |
| args.contextual, | |
| f"-cx{args.contextual}({iterable_to_str(args.cx_layers)})" | |
| ) | |
| #+ optional_string(args.mse, f"-mse{args.mse}") | |
| + optional_string(args.noise_regularize, f"-NR{args.noise_regularize:.1e}") | |
| ) | |
| class BakedMultiContextualLoss(nn.Module): | |
| """Random sample different image patches for different vgg layers.""" | |
| def __init__(self, sibling: torch.Tensor, args: Namespace, size: int = 256): | |
| super().__init__() | |
| self.cxs = nn.ModuleList([ContextualLoss(use_vgg=True, vgg_layers=[layer]) | |
| for layer in args.cx_layers]) | |
| self.size = size | |
| self.sibling = sibling.detach() | |
| def forward(self, img: torch.Tensor): | |
| cx_loss = 0 | |
| for cx in self.cxs: | |
| h, w = np.random.randint(0, high=img.shape[-1] - self.size, size=2) | |
| cx_loss = cx(self.sibling[..., h:h+self.size, w:w+self.size], img[..., h:h+self.size, w:w+self.size]) + cx_loss | |
| return cx_loss | |
| class BakedContextualLoss(ContextualLoss): | |
| def __init__(self, sibling: torch.Tensor, args: Namespace, size: int = 256): | |
| super().__init__(use_vgg=True, vgg_layers=args.cx_layers) | |
| self.size = size | |
| self.sibling = sibling.detach() | |
| def forward(self, img: torch.Tensor): | |
| h, w = np.random.randint(0, high=img.shape[-1] - self.size, size=2) | |
| return super().forward(self.sibling[..., h:h+self.size, w:w+self.size], img[..., h:h+self.size, w:w+self.size]) | |
| class JointLoss(nn.Module): | |
| def __init__( | |
| self, | |
| args: Namespace, | |
| target: torch.Tensor, | |
| sibling: Optional[torch.Tensor], | |
| sibling_rgbs: Optional[Iterable[torch.Tensor]] = None, | |
| ): | |
| super().__init__() | |
| self.weights = { | |
| "face": 1., "eye": args.eye, | |
| "contextual": args.contextual, "color_transfer": args.color_transfer, | |
| "noise": args.noise_regularize, | |
| } | |
| reconstruction = {} | |
| if args.vgg > 0 or args.vggface > 0: | |
| percept = create_perceptual_loss(args) | |
| reconstruction.update( | |
| {"face": FaceLoss(target, input_size=args.generator_size, size=args.recon_size, percept=percept)} | |
| ) | |
| if args.eye > 0: | |
| reconstruction.update( | |
| {"eye": EyeLoss(target, input_size=args.generator_size, percept=percept)} | |
| ) | |
| self.reconstruction = nn.ModuleDict(reconstruction) | |
| exemplar = {} | |
| if args.contextual > 0 and len(args.cx_layers) > 0: | |
| assert sibling is not None | |
| exemplar.update( | |
| {"contextual": BakedContextualLoss(sibling, args)} | |
| ) | |
| if args.color_transfer > 0: | |
| assert sibling_rgbs is not None | |
| self.sibling_rgbs = sibling_rgbs | |
| exemplar.update( | |
| {"color_transfer": ColorTransferLoss(init_rgbs=sibling_rgbs)} | |
| ) | |
| self.exemplar = nn.ModuleDict(exemplar) | |
| if args.noise_regularize > 0: | |
| self.noise_criterion = NoiseRegularizer() | |
| def forward( | |
| self, img, degrade=None, noises=None, rgbs=None, rgb_level: Optional[int] = None | |
| ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: | |
| """ | |
| Args: | |
| rgbs: results from the ToRGB layers | |
| """ | |
| # TODO: add current optimization resolution for noises | |
| losses = {} | |
| # reconstruction losses | |
| for name, criterion in self.reconstruction.items(): | |
| losses[name] = criterion(img, degrade=degrade) | |
| # exemplar losses | |
| if 'contextual' in self.exemplar: | |
| losses["contextual"] = self.exemplar["contextual"](img) | |
| if "color_transfer" in self.exemplar: | |
| assert rgbs is not None | |
| losses["color_transfer"] = self.exemplar["color_transfer"](rgbs, level=rgb_level) | |
| # noise regularizer | |
| if self.weights["noise"] > 0: | |
| losses["noise"] = self.noise_criterion(noises) | |
| total_loss = 0 | |
| for name, loss in losses.items(): | |
| total_loss = total_loss + self.weights[name] * loss | |
| return total_loss, losses | |
| def update_sibling(self, sibling: torch.Tensor): | |
| assert "contextual" in self.exemplar | |
| self.exemplar["contextual"].sibling = sibling.detach() | |