Duplicate from feng2022/Time-TravelRephotography

Co-authored-by: Time-travelRephotography <feng2022@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,27 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,133 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+wandb/
+*.lmdb/
+*.pkl

.gitmodules

@@ -0,0 +1,9 @@
+[submodule "third_party/face_parsing"]
+	path = third_party/face_parsing
+	url = https://github.com/Time-Travel-Rephotography/face-parsing.PyTorch.git
+[submodule "models/encoder4editing"]
+	path = models/encoder4editing
+	url = https://github.com/Time-Travel-Rephotography/encoder4editing.git
+[submodule "losses/contextual_loss"]
+	path = losses/contextual_loss
+	url = https://github.com/Time-Travel-Rephotography/contextual_loss_pytorch.git

README.md

@@ -0,0 +1,14 @@
+---
+title: Time TravelRephotography
+emoji: 🦀
+colorFrom: yellow
+colorTo: red
+sdk: gradio
+sdk_version: 2.9.4
+app_file: app.py
+pinned: false
+license: mit
+duplicated_from: feng2022/Time-TravelRephotography
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference

Time_TravelRephotography/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2020 Time-Travel-Rephotography
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

Time_TravelRephotography/LICENSE-NVIDIA

@@ -0,0 +1,101 @@
+Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+
+
+Nvidia Source Code License-NC
+
+=======================================================================
+
+1. Definitions
+
+"Licensor" means any person or entity that distributes its Work.
+
+"Software" means the original work of authorship made available under
+this License.
+
+"Work" means the Software and any additions to or derivative works of
+the Software that are made available under this License.
+
+"Nvidia Processors" means any central processing unit (CPU), graphics
+processing unit (GPU), field-programmable gate array (FPGA),
+application-specific integrated circuit (ASIC) or any combination
+thereof designed, made, sold, or provided by Nvidia or its affiliates.
+
+The terms "reproduce," "reproduction," "derivative works," and
+"distribution" have the meaning as provided under U.S. copyright law;
+provided, however, that for the purposes of this License, derivative
+works shall not include works that remain separable from, or merely
+link (or bind by name) to the interfaces of, the Work.
+
+Works, including the Software, are "made available" under this License
+by including in or with the Work either (a) a copyright notice
+referencing the applicability of this License to the Work, or (b) a
+copy of this License.
+
+2. License Grants
+
+    2.1 Copyright Grant. Subject to the terms and conditions of this
+    License, each Licensor grants to you a perpetual, worldwide,
+    non-exclusive, royalty-free, copyright license to reproduce,
+    prepare derivative works of, publicly display, publicly perform,
+    sublicense and distribute its Work and any resulting derivative
+    works in any form.
+
+3. Limitations
+
+    3.1 Redistribution. You may reproduce or distribute the Work only
+    if (a) you do so under this License, (b) you include a complete
+    copy of this License with your distribution, and (c) you retain
+    without modification any copyright, patent, trademark, or
+    attribution notices that are present in the Work.
+
+    3.2 Derivative Works. You may specify that additional or different
+    terms apply to the use, reproduction, and distribution of your
+    derivative works of the Work ("Your Terms") only if (a) Your Terms
+    provide that the use limitation in Section 3.3 applies to your
+    derivative works, and (b) you identify the specific derivative
+    works that are subject to Your Terms. Notwithstanding Your Terms,
+    this License (including the redistribution requirements in Section
+    3.1) will continue to apply to the Work itself.
+
+    3.3 Use Limitation. The Work and any derivative works thereof only
+    may be used or intended for use non-commercially. The Work or
+    derivative works thereof may be used or intended for use by Nvidia
+    or its affiliates commercially or non-commercially. As used herein,
+    "non-commercially" means for research or evaluation purposes only.
+
+    3.4 Patent Claims. If you bring or threaten to bring a patent claim
+    against any Licensor (including any claim, cross-claim or
+    counterclaim in a lawsuit) to enforce any patents that you allege
+    are infringed by any Work, then your rights under this License from
+    such Licensor (including the grants in Sections 2.1 and 2.2) will
+    terminate immediately.
+
+    3.5 Trademarks. This License does not grant any rights to use any
+    Licensor's or its affiliates' names, logos, or trademarks, except
+    as necessary to reproduce the notices described in this License.
+
+    3.6 Termination. If you violate any term of this License, then your
+    rights under this License (including the grants in Sections 2.1 and
+    2.2) will terminate immediately.
+
+4. Disclaimer of Warranty.
+
+THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR
+NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER
+THIS LICENSE. 
+
+5. Limitation of Liability.
+
+EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL
+THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE
+SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT,
+INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF
+OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK
+(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION,
+LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER
+COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF
+THE POSSIBILITY OF SUCH DAMAGES.
+
+=======================================================================

Time_TravelRephotography/LICENSE-STYLEGAN2

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 Kim Seonghyeon
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

Time_TravelRephotography/losses/color_transfer_loss.py

@@ -0,0 +1,60 @@
+from typing import List, Optional
+
+import torch
+from torch import nn
+from torch.nn.functional import (
+    smooth_l1_loss,
+)
+
+
+def flatten_CHW(im: torch.Tensor) -> torch.Tensor:
+    """
+    (B, C, H, W) -> (B, -1)
+    """
+    B = im.shape[0]
+    return im.reshape(B, -1)
+
+
+def stddev(x: torch.Tensor) -> torch.Tensor:
+    """
+    x: (B, -1), assume with mean normalized
+    Retuens:
+        stddev: (B)
+    """
+    return torch.sqrt(torch.mean(x * x, dim=-1))
+
+
+def gram_matrix(input_):
+    B, C = input_.shape[:2]
+    features = input_.view(B, C, -1)
+    N = features.shape[-1]
+    G = torch.bmm(features, features.transpose(1, 2))  # C x C
+    return G.div(C * N)
+
+
+class ColorTransferLoss(nn.Module):
+    """Penalize the gram matrix difference between StyleGAN2's ToRGB outputs"""
+    def __init__(
+        self,
+        init_rgbs,
+        scale_rgb: bool = False
+    ):
+        super().__init__()
+
+        with torch.no_grad():
+            init_feats = [x.detach() for x in init_rgbs]
+            self.stds = [stddev(flatten_CHW(rgb)) if scale_rgb else 1 for rgb in init_feats]  # (B, 1, 1, 1) or scalar
+            self.grams = [gram_matrix(rgb / std) for rgb, std in zip(init_feats, self.stds)]
+
+    def forward(self, rgbs: List[torch.Tensor], level: int = None):
+        if level is None:
+            level = len(self.grams)
+
+        feats = rgbs
+        loss = 0
+        for i, (rgb, std) in enumerate(zip(feats[:level], self.stds[:level])):
+            G = gram_matrix(rgb / std)
+            loss = loss + smooth_l1_loss(G, self.grams[i])
+
+        return loss
+

Time_TravelRephotography/losses/contextual_loss/.gitignore

@@ -0,0 +1,104 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/

Time_TravelRephotography/losses/contextual_loss/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 Sou Uchida
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

Time_TravelRephotography/losses/contextual_loss/__init__.py

@@ -0,0 +1 @@
+from .modules import *

Time_TravelRephotography/losses/contextual_loss/config.py

@@ -0,0 +1,2 @@
+# TODO: add supports for L1, L2 etc.
+LOSS_TYPES = ['cosine', 'l1', 'l2']

Time_TravelRephotography/losses/contextual_loss/functional.py

@@ -0,0 +1,198 @@
+import torch
+import torch.nn.functional as F
+
+from .config import LOSS_TYPES
+
+__all__ = ['contextual_loss', 'contextual_bilateral_loss']
+
+
+def contextual_loss(x: torch.Tensor,
+                    y: torch.Tensor,
+                    band_width: float = 0.5,
+                    loss_type: str = 'cosine',
+                    all_dist: bool = False):
+    """
+    Computes contextual loss between x and y.
+    The most of this code is copied from
+        https://gist.github.com/yunjey/3105146c736f9c1055463c33b4c989da.
+
+    Parameters
+    ---
+    x : torch.Tensor
+        features of shape (N, C, H, W).
+    y : torch.Tensor
+        features of shape (N, C, H, W).
+    band_width : float, optional
+        a band-width parameter used to convert distance to similarity.
+        in the paper, this is described as :math:`h`.
+    loss_type : str, optional
+        a loss type to measure the distance between features.
+        Note: `l1` and `l2` frequently raises OOM.
+
+    Returns
+    ---
+    cx_loss : torch.Tensor
+        contextual loss between x and y (Eq (1) in the paper)
+    """
+
+    assert x.size() == y.size(), 'input tensor must have the same size.'
+    assert loss_type in LOSS_TYPES, f'select a loss type from {LOSS_TYPES}.'
+
+    N, C, H, W = x.size()
+
+    if loss_type == 'cosine':
+        dist_raw = compute_cosine_distance(x, y)
+    elif loss_type == 'l1':
+        dist_raw = compute_l1_distance(x, y)
+    elif loss_type == 'l2':
+        dist_raw = compute_l2_distance(x, y)
+
+    dist_tilde = compute_relative_distance(dist_raw)
+    cx = compute_cx(dist_tilde, band_width)
+    if all_dist:
+        return cx
+
+    cx = torch.mean(torch.max(cx, dim=1)[0], dim=1)  # Eq(1)
+    cx_loss = torch.mean(-torch.log(cx + 1e-5))  # Eq(5)
+
+    return cx_loss
+
+
+# TODO: Operation check
+def contextual_bilateral_loss(x: torch.Tensor,
+                              y: torch.Tensor,
+                              weight_sp: float = 0.1,
+                              band_width: float = 1.,
+                              loss_type: str = 'cosine'):
+    """
+    Computes Contextual Bilateral (CoBi) Loss between x and y,
+        proposed in https://arxiv.org/pdf/1905.05169.pdf.
+
+    Parameters
+    ---
+    x : torch.Tensor
+        features of shape (N, C, H, W).
+    y : torch.Tensor
+        features of shape (N, C, H, W).
+    band_width : float, optional
+        a band-width parameter used to convert distance to similarity.
+        in the paper, this is described as :math:`h`.
+    loss_type : str, optional
+        a loss type to measure the distance between features.
+        Note: `l1` and `l2` frequently raises OOM.
+
+    Returns
+    ---
+    cx_loss : torch.Tensor
+        contextual loss between x and y (Eq (1) in the paper).
+    k_arg_max_NC : torch.Tensor
+        indices to maximize similarity over channels.
+    """
+
+    assert x.size() == y.size(), 'input tensor must have the same size.'
+    assert loss_type in LOSS_TYPES, f'select a loss type from {LOSS_TYPES}.'
+
+    # spatial loss
+    grid = compute_meshgrid(x.shape).to(x.device)
+    dist_raw = compute_l2_distance(grid, grid)
+    dist_tilde = compute_relative_distance(dist_raw)
+    cx_sp = compute_cx(dist_tilde, band_width)
+
+    # feature loss
+    if loss_type == 'cosine':
+        dist_raw = compute_cosine_distance(x, y)
+    elif loss_type == 'l1':
+        dist_raw = compute_l1_distance(x, y)
+    elif loss_type == 'l2':
+        dist_raw = compute_l2_distance(x, y)
+    dist_tilde = compute_relative_distance(dist_raw)
+    cx_feat = compute_cx(dist_tilde, band_width)
+
+    # combined loss
+    cx_combine = (1. - weight_sp) * cx_feat + weight_sp * cx_sp
+
+    k_max_NC, _ = torch.max(cx_combine, dim=2, keepdim=True)
+
+    cx = k_max_NC.mean(dim=1)
+    cx_loss = torch.mean(-torch.log(cx + 1e-5))
+
+    return cx_loss
+
+
+def compute_cx(dist_tilde, band_width):
+    w = torch.exp((1 - dist_tilde) / band_width)  # Eq(3)
+    cx = w / torch.sum(w, dim=2, keepdim=True)  # Eq(4)
+    return cx
+
+
+def compute_relative_distance(dist_raw):
+    dist_min, _ = torch.min(dist_raw, dim=2, keepdim=True)
+    dist_tilde = dist_raw / (dist_min + 1e-5)
+    return dist_tilde
+
+
+def compute_cosine_distance(x, y):
+    # mean shifting by channel-wise mean of `y`.
+    y_mu = y.mean(dim=(0, 2, 3), keepdim=True)
+    x_centered = x - y_mu
+    y_centered = y - y_mu
+
+    # L2 normalization
+    x_normalized = F.normalize(x_centered, p=2, dim=1)
+    y_normalized = F.normalize(y_centered, p=2, dim=1)
+
+    # channel-wise vectorization
+    N, C, *_ = x.size()
+    x_normalized = x_normalized.reshape(N, C, -1)  # (N, C, H*W)
+    y_normalized = y_normalized.reshape(N, C, -1)  # (N, C, H*W)
+
+    # consine similarity
+    cosine_sim = torch.bmm(x_normalized.transpose(1, 2),
+                           y_normalized)  # (N, H*W, H*W)
+
+    # convert to distance
+    dist = 1 - cosine_sim
+
+    return dist
+
+
+# TODO: Considering avoiding OOM.
+def compute_l1_distance(x: torch.Tensor, y: torch.Tensor):
+    N, C, H, W = x.size()
+    x_vec = x.view(N, C, -1)
+    y_vec = y.view(N, C, -1)
+
+    dist = x_vec.unsqueeze(2) - y_vec.unsqueeze(3)
+    dist = dist.abs().sum(dim=1)
+    dist = dist.transpose(1, 2).reshape(N, H*W, H*W)
+    dist = dist.clamp(min=0.)
+
+    return dist
+
+
+# TODO: Considering avoiding OOM.
+def compute_l2_distance(x, y):
+    N, C, H, W = x.size()
+    x_vec = x.view(N, C, -1)
+    y_vec = y.view(N, C, -1)
+    x_s = torch.sum(x_vec ** 2, dim=1)
+    y_s = torch.sum(y_vec ** 2, dim=1)
+
+    A = y_vec.transpose(1, 2) @ x_vec
+    dist = y_s - 2 * A + x_s.transpose(0, 1)
+    dist = dist.transpose(1, 2).reshape(N, H*W, H*W)
+    dist = dist.clamp(min=0.)
+
+    return dist
+
+
+def compute_meshgrid(shape):
+    N, C, H, W = shape
+    rows = torch.arange(0, H, dtype=torch.float32) / (H + 1)
+    cols = torch.arange(0, W, dtype=torch.float32) / (W + 1)
+
+    feature_grid = torch.meshgrid(rows, cols)
+    feature_grid = torch.stack(feature_grid).unsqueeze(0)
+    feature_grid = torch.cat([feature_grid for _ in range(N)], dim=0)
+
+    return feature_grid

Time_TravelRephotography/losses/contextual_loss/modules/__init__.py

@@ -0,0 +1,4 @@
+from .contextual import ContextualLoss
+from .contextual_bilateral import ContextualBilateralLoss
+
+__all__ = ['ContextualLoss', 'ContextualBilateralLoss']

Time_TravelRephotography/losses/contextual_loss/modules/contextual.py

@@ -0,0 +1,121 @@
+import random
+from typing import (
+    Iterable,
+    List,
+    Optional,
+)
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from .vgg import VGG19
+from .. import functional as F
+from ..config import LOSS_TYPES
+
+
+class ContextualLoss(nn.Module):
+    """
+    Creates a criterion that measures the contextual loss.
+
+    Parameters
+    ---
+    band_width : int, optional
+        a band_width parameter described as :math:`h` in the paper.
+    use_vgg : bool, optional
+        if you want to use VGG feature, set this `True`.
+    vgg_layer : str, optional
+        intermidiate layer name for VGG feature.
+        Now we support layer names:
+            `['relu1_2', 'relu2_2', 'relu3_4', 'relu4_4', 'relu5_4']`
+    """
+
+    def __init__(
+            self,
+            band_width: float = 0.5,
+            loss_type: str = 'cosine',
+            use_vgg: bool = False,
+            vgg_layers: List[str] = ['relu3_4'],
+            feature_1d_size: int = 64,
+    ):
+
+        super().__init__()
+
+        assert band_width > 0, 'band_width parameter must be positive.'
+        assert loss_type in LOSS_TYPES,\
+            f'select a loss type from {LOSS_TYPES}.'
+
+        self.loss_type = loss_type
+        self.band_width = band_width
+        self.feature_1d_size = feature_1d_size
+
+        if use_vgg:
+            self.vgg_model = VGG19()
+            self.vgg_layers = vgg_layers
+            self.register_buffer(
+                name='vgg_mean',
+                tensor=torch.tensor(
+                    [[[0.485]], [[0.456]], [[0.406]]], requires_grad=False)
+            )
+            self.register_buffer(
+                name='vgg_std',
+                tensor=torch.tensor(
+                    [[[0.229]], [[0.224]], [[0.225]]], requires_grad=False)
+            )
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor, all_dist: bool = False):
+        if not hasattr(self, 'vgg_model'):
+            return self.contextual_loss(x, y, self.feature_1d_size, self.band_width, all_dist=all_dist)
+
+
+        x = self.forward_vgg(x)
+        y = self.forward_vgg(y)
+
+        loss = 0
+        for layer in self.vgg_layers:
+            # picking up vgg feature maps
+            fx = getattr(x, layer)
+            fy = getattr(y, layer)
+            loss = loss + self.contextual_loss(
+                fx, fy, self.feature_1d_size, self.band_width, all_dist=all_dist, loss_type=self.loss_type
+            )
+        return loss
+
+    def forward_vgg(self, x: torch.Tensor):
+        assert x.shape[1] == 3, 'VGG model takes 3 chennel images.'
+        # [-1, 1] -> [0, 1]
+        x = (x + 1) * 0.5
+
+        # normalization
+        x = x.sub(self.vgg_mean.detach()).div(self.vgg_std)
+        return self.vgg_model(x)
+
+    @classmethod
+    def contextual_loss(
+            cls,
+            x: torch.Tensor, y: torch.Tensor,
+            feature_1d_size: int,
+            band_width: int,
+            all_dist: bool = False,
+            loss_type: str = 'cosine',
+    ) -> torch.Tensor:
+        feature_size = feature_1d_size ** 2
+        if np.prod(x.shape[2:]) > feature_size or np.prod(y.shape[2:]) > feature_size:
+            x, indices = cls.random_sampling(x, feature_1d_size=feature_1d_size)
+            y, _ = cls.random_sampling(y, feature_1d_size=feature_1d_size, indices=indices)
+
+        return F.contextual_loss(x, y, band_width, all_dist=all_dist, loss_type=loss_type)
+
+    @staticmethod
+    def random_sampling(
+            tensor_NCHW: torch.Tensor, feature_1d_size: int, indices: Optional[List] = None
+    ):
+        N, C, H, W = tensor_NCHW.shape
+        S = H * W
+        tensor_NCS = tensor_NCHW.reshape([N, C, S])
+        if indices is None:
+            all_indices = list(range(S))
+            random.shuffle(all_indices)
+            indices = all_indices[:feature_1d_size**2]
+        res = tensor_NCS[:, :, indices].reshape(N, -1, feature_1d_size, feature_1d_size)
+        return res, indices

Time_TravelRephotography/losses/contextual_loss/modules/contextual_bilateral.py

@@ -0,0 +1,69 @@
+import torch
+import torch.nn as nn
+
+from .vgg import VGG19
+from .. import functional as F
+from ..config import LOSS_TYPES
+
+
+class ContextualBilateralLoss(nn.Module):
+    """
+    Creates a criterion that measures the contextual bilateral loss.
+
+    Parameters
+    ---
+    weight_sp : float, optional
+        a balancing weight between spatial and feature loss.
+    band_width : int, optional
+        a band_width parameter described as :math:`h` in the paper.
+    use_vgg : bool, optional
+        if you want to use VGG feature, set this `True`.
+    vgg_layer : str, optional
+        intermidiate layer name for VGG feature.
+        Now we support layer names:
+            `['relu1_2', 'relu2_2', 'relu3_4', 'relu4_4', 'relu5_4']`
+    """
+
+    def __init__(self,
+                 weight_sp: float = 0.1,
+                 band_width: float = 0.5,
+                 loss_type: str = 'cosine',
+                 use_vgg: bool = False,
+                 vgg_layer: str = 'relu3_4'):
+
+        super(ContextualBilateralLoss, self).__init__()
+
+        assert band_width > 0, 'band_width parameter must be positive.'
+        assert loss_type in LOSS_TYPES,\
+            f'select a loss type from {LOSS_TYPES}.'
+
+        self.band_width = band_width
+
+        if use_vgg:
+            self.vgg_model = VGG19()
+            self.vgg_layer = vgg_layer
+            self.register_buffer(
+                name='vgg_mean',
+                tensor=torch.tensor(
+                    [[[0.485]], [[0.456]], [[0.406]]], requires_grad=False)
+            )
+            self.register_buffer(
+                name='vgg_std',
+                tensor=torch.tensor(
+                    [[[0.229]], [[0.224]], [[0.225]]], requires_grad=False)
+            )
+
+    def forward(self, x, y):
+        if hasattr(self, 'vgg_model'):
+            assert x.shape[1] == 3 and y.shape[1] == 3,\
+                'VGG model takes 3 chennel images.'
+
+            # normalization
+            x = x.sub(self.vgg_mean.detach()).div(self.vgg_std.detach())
+            y = y.sub(self.vgg_mean.detach()).div(self.vgg_std.detach())
+
+            # picking up vgg feature maps
+            x = getattr(self.vgg_model(x), self.vgg_layer)
+            y = getattr(self.vgg_model(y), self.vgg_layer)
+
+        return F.contextual_bilateral_loss(x, y, self.band_width)

Time_TravelRephotography/losses/contextual_loss/modules/vgg.py

@@ -0,0 +1,48 @@
+from collections import namedtuple
+
+import torch.nn as nn
+import torchvision.models.vgg as vgg
+
+
+class VGG19(nn.Module):
+    def __init__(self, requires_grad=False):
+        super(VGG19, self).__init__()
+        vgg_pretrained_features = vgg.vgg19(pretrained=True).features
+        self.slice1 = nn.Sequential()
+        self.slice2 = nn.Sequential()
+        self.slice3 = nn.Sequential()
+        self.slice4 = nn.Sequential()
+        self.slice5 = nn.Sequential()
+        for x in range(4):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(4, 9):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(9, 18):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(18, 27):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(27, 36):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1_2 = h
+        h = self.slice2(h)
+        h_relu2_2 = h
+        h = self.slice3(h)
+        h_relu3_4 = h
+        h = self.slice4(h)
+        h_relu4_4 = h
+        h = self.slice5(h)
+        h_relu5_4 = h
+
+        vgg_outputs = namedtuple(
+            "VggOutputs", ['relu1_2', 'relu2_2',
+                           'relu3_4', 'relu4_4', 'relu5_4'])
+        out = vgg_outputs(h_relu1_2, h_relu2_2,
+                          h_relu3_4, h_relu4_4, h_relu5_4)
+
+        return out

Time_TravelRephotography/losses/joint_loss.py

@@ -0,0 +1,167 @@
+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:
+    @staticmethod
+    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'])
+
+    @staticmethod
+    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()

Time_TravelRephotography/losses/perceptual_loss.py

@@ -0,0 +1,111 @@
+"""
+Code borrowed from https://gist.github.com/alper111/8233cdb0414b4cb5853f2f730ab95a49#file-vgg_perceptual_loss-py-L5
+"""
+import torch
+import torchvision
+from models.vggface import VGGFaceFeats
+
+
+def cos_loss(fi, ft):
+    return 1 - torch.nn.functional.cosine_similarity(fi, ft).mean()
+
+
+class VGGPerceptualLoss(torch.nn.Module):
+    def __init__(self, resize=False):
+        super(VGGPerceptualLoss, self).__init__()
+        blocks = []
+        blocks.append(torchvision.models.vgg16(pretrained=True).features[:4].eval())
+        blocks.append(torchvision.models.vgg16(pretrained=True).features[4:9].eval())
+        blocks.append(torchvision.models.vgg16(pretrained=True).features[9:16].eval())
+        blocks.append(torchvision.models.vgg16(pretrained=True).features[16:23].eval())
+        for bl in blocks:
+            for p in bl:
+                p.requires_grad = False
+        self.blocks = torch.nn.ModuleList(blocks)
+        self.transform = torch.nn.functional.interpolate
+        self.mean = torch.nn.Parameter(torch.tensor([0.485, 0.456, 0.406]).view(1,3,1,1))
+        self.std = torch.nn.Parameter(torch.tensor([0.229, 0.224, 0.225]).view(1,3,1,1))
+        self.resize = resize
+
+    def forward(self, input, target, max_layer=4, cos_dist: bool = False):
+        target = (target + 1) * 0.5
+        input = (input + 1) * 0.5
+
+        if input.shape[1] != 3:
+            input = input.repeat(1, 3, 1, 1)
+            target = target.repeat(1, 3, 1, 1)
+        input = (input-self.mean) / self.std
+        target = (target-self.mean) / self.std
+        if self.resize:
+            input = self.transform(input, mode='bilinear', size=(224, 224), align_corners=False)
+            target = self.transform(target, mode='bilinear', size=(224, 224), align_corners=False)
+        x = input
+        y = target
+        loss = 0.0
+        loss_func = cos_loss if cos_dist else torch.nn.functional.l1_loss
+        for bi, block in enumerate(self.blocks[:max_layer]):
+            x = block(x)
+            y = block(y)
+            loss += loss_func(x, y.detach())
+        return loss
+
+
+class VGGFacePerceptualLoss(torch.nn.Module):
+    def __init__(self, weight_path: str = "checkpoint/vgg_face_dag.pt", resize: bool = False):
+        super().__init__()
+        self.vgg = VGGFaceFeats()
+        self.vgg.load_state_dict(torch.load(weight_path))
+
+        mean = torch.tensor(self.vgg.meta["mean"]).view(1, 3, 1, 1) / 255.0
+        self.register_buffer("mean", mean)
+
+        self.transform = torch.nn.functional.interpolate
+        self.resize = resize
+
+    def forward(self, input, target, max_layer: int = 4, cos_dist: bool = False):
+        target = (target + 1) * 0.5
+        input = (input + 1) * 0.5
+
+        # preprocessing
+        if input.shape[1] != 3:
+            input = input.repeat(1, 3, 1, 1)
+            target = target.repeat(1, 3, 1, 1)
+        input = input - self.mean
+        target = target - self.mean
+        if self.resize:
+            input = self.transform(input, mode='bilinear', size=(224, 224), align_corners=False)
+            target = self.transform(target, mode='bilinear', size=(224, 224), align_corners=False)
+
+        input_feats = self.vgg(input)
+        target_feats = self.vgg(target)
+
+        loss_func = cos_loss if cos_dist else torch.nn.functional.l1_loss
+        # calc perceptual loss
+        loss = 0.0
+        for fi, ft in zip(input_feats[:max_layer], target_feats[:max_layer]):
+            loss = loss + loss_func(fi, ft.detach())
+        return loss
+
+
+class PerceptualLoss(torch.nn.Module):
+    def __init__(
+            self, lambda_vggface: float = 0.025 / 0.15, lambda_vgg: float = 1,  eps: float = 1e-8, cos_dist: bool = False
+    ):
+        super().__init__()
+        self.register_buffer("lambda_vggface", torch.tensor(lambda_vggface))
+        self.register_buffer("lambda_vgg", torch.tensor(lambda_vgg))
+        self.cos_dist = cos_dist
+
+        if lambda_vgg > eps:
+            self.vgg = VGGPerceptualLoss()
+        if lambda_vggface > eps:
+            self.vggface = VGGFacePerceptualLoss()
+
+    def forward(self, input, target, eps=1e-8, use_vggface: bool = True, use_vgg=True, max_vgg_layer=4):
+        loss = 0.0
+        if self.lambda_vgg > eps and use_vgg:
+            loss = loss + self.lambda_vgg * self.vgg(input, target, max_layer=max_vgg_layer)
+        if self.lambda_vggface > eps and use_vggface:
+            loss = loss + self.lambda_vggface * self.vggface(input, target, cos_dist=self.cos_dist)
+        return loss
+

Time_TravelRephotography/losses/reconstruction.py

@@ -0,0 +1,119 @@
+from argparse import (
+    ArgumentParser,
+    Namespace,
+)
+from typing import Optional
+
+import numpy as np
+import torch
+from torch import nn
+
+from losses.perceptual_loss import PerceptualLoss
+from models.degrade import Downsample
+from utils.misc import optional_string
+
+
+class ReconstructionArguments:
+    @staticmethod
+    def add_arguments(parser: ArgumentParser):
+        parser.add_argument("--vggface", type=float, default=0.3, help="vggface")
+        parser.add_argument("--vgg", type=float, default=1, help="vgg")
+        parser.add_argument('--recon_size', type=int, default=256, help="size for face reconstruction loss")
+
+    @staticmethod
+    def to_string(args: Namespace) -> str:
+        return (
+            f"s{args.recon_size}"
+            + optional_string(args.vgg > 0, f"-vgg{args.vgg}")
+            + optional_string(args.vggface > 0, f"-vggface{args.vggface}")
+        )
+
+
+def create_perceptual_loss(args: Namespace):
+    return PerceptualLoss(lambda_vgg=args.vgg, lambda_vggface=args.vggface, cos_dist=False)
+
+
+class EyeLoss(nn.Module):
+    def __init__(
+            self,
+            target: torch.Tensor,
+            input_size: int = 1024,
+            input_channels: int = 3,
+            percept: Optional[nn.Module] = None,
+            args: Optional[Namespace] = None
+    ):
+        """
+        target: target image
+        """
+        assert not (percept is None and args is None)
+
+        super().__init__()
+
+        self.target = target
+
+        target_size = target.shape[-1]
+        self.downsample = Downsample(input_size, target_size, input_channels) \
+                if target_size != input_size else (lambda x: x)
+
+        self.percept = percept if percept is not None else create_perceptual_loss(args)
+
+        eye_size = np.array((224, 224))
+        btlrs = []
+        for sgn in [1, -1]:
+            center = np.array((480, 384 * sgn))   # (y, x)
+            b, t = center[0] - eye_size[0] // 2, center[0] + eye_size[0] // 2
+            l, r = center[1] - eye_size[1] // 2, center[1] + eye_size[1] // 2
+            btlrs.append((np.array((b, t, l, r)) / 1024 * target_size).astype(int))
+        self.btlrs = np.stack(btlrs, axis=0)
+
+    def forward(self, img: torch.Tensor, degrade: nn.Module = None):
+        """
+        img: it should be the degraded version of the generated image
+        """
+        if degrade is not None:
+            img = degrade(img, downsample=self.downsample)
+
+        loss = 0
+        for (b, t, l, r) in self.btlrs:
+            loss = loss + self.percept(
+                img[:, :, b:t, l:r], self.target[:, :, b:t, l:r],
+                use_vggface=False, max_vgg_layer=4,
+                # use_vgg=False,
+            )
+        return loss
+
+
+class FaceLoss(nn.Module):
+    def __init__(
+            self,
+            target: torch.Tensor,
+            input_size: int = 1024,
+            input_channels: int = 3,
+            size: int = 256,
+            percept: Optional[nn.Module] = None,
+            args: Optional[Namespace] = None
+    ):
+        """
+        target: target image
+        """
+        assert not (percept is None and args is None)
+
+        super().__init__()
+
+        target_size = target.shape[-1]
+        self.target = target if target_size == size \
+                else Downsample(target_size, size, target.shape[1]).to(target.device)(target)
+
+        self.downsample = Downsample(input_size, size, input_channels) \
+                if size != input_size else (lambda x: x)
+
+        self.percept = percept if percept is not None else create_perceptual_loss(args)
+
+    def forward(self, img: torch.Tensor, degrade: nn.Module = None):
+        """
+        img: it should be the degraded version of the generated image
+        """
+        if degrade is not None:
+            img = degrade(img, downsample=self.downsample)
+        loss = self.percept(img, self.target)
+        return loss

Time_TravelRephotography/losses/regularize_noise.py

@@ -0,0 +1,37 @@
+from typing import Iterable
+
+import torch
+from torch import nn
+
+
+class NoiseRegularizer(nn.Module):
+    def forward(self, noises: Iterable[torch.Tensor]):
+        loss = 0
+
+        for noise in noises:
+            size = noise.shape[2]
+
+            while True:
+                loss = (
+                    loss
+                    + (noise * torch.roll(noise, shifts=1, dims=3)).mean().pow(2)
+                    + (noise * torch.roll(noise, shifts=1, dims=2)).mean().pow(2)
+                )
+
+                if size <= 8:
+                    break
+
+                noise = noise.reshape([1, 1, size // 2, 2, size // 2, 2])
+                noise = noise.mean([3, 5])
+                size //= 2
+
+        return loss
+
+    @staticmethod
+    def normalize(noises: Iterable[torch.Tensor]):
+        for noise in noises:
+            mean = noise.mean()
+            std = noise.std()
+
+            noise.data.add_(-mean).div_(std)
+

Time_TravelRephotography/model.py

@@ -0,0 +1,697 @@
+import math
+import random
+import functools
+import operator
+import numpy as np
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.autograd import Function
+
+from op import FusedLeakyReLU, fused_leaky_relu, upfirdn2d
+
+
+class PixelNorm(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, input):
+        return input * torch.rsqrt(torch.mean(input ** 2, dim=1, keepdim=True) + 1e-8)
+
+
+def make_kernel(k):
+    k = torch.tensor(k, dtype=torch.float32)
+
+    if k.ndim == 1:
+        k = k[None, :] * k[:, None]
+
+    k /= k.sum()
+
+    return k
+
+
+class Upsample(nn.Module):
+    def __init__(self, kernel, factor=2):
+        super().__init__()
+
+        self.factor = factor
+        kernel = make_kernel(kernel) * (factor ** 2)
+        self.register_buffer('kernel', kernel)
+
+        p = kernel.shape[0] - factor
+
+        pad0 = (p + 1) // 2 + factor - 1
+        pad1 = p // 2
+
+        self.pad = (pad0, pad1)
+
+    def forward(self, input):
+        out = upfirdn2d(input, self.kernel, up=self.factor, down=1, pad=self.pad)
+
+        return out
+
+
+class Downsample(nn.Module):
+    def __init__(self, kernel, factor=2):
+        super().__init__()
+
+        self.factor = factor
+        kernel = make_kernel(kernel)
+        self.register_buffer('kernel', kernel)
+
+        p = kernel.shape[0] - factor
+
+        pad0 = (p + 1) // 2
+        pad1 = p // 2
+
+        self.pad = (pad0, pad1)
+
+    def forward(self, input):
+        out = upfirdn2d(input, self.kernel, up=1, down=self.factor, pad=self.pad)
+
+        return out
+
+
+class Blur(nn.Module):
+    def __init__(self, kernel, pad, upsample_factor=1):
+        super().__init__()
+
+        kernel = make_kernel(kernel)
+
+        if upsample_factor > 1:
+            kernel = kernel * (upsample_factor ** 2)
+
+        self.register_buffer('kernel', kernel)
+
+        self.pad = pad
+
+    def forward(self, input):
+        out = upfirdn2d(input, self.kernel, pad=self.pad)
+
+        return out
+
+
+class EqualConv2d(nn.Module):
+    def __init__(
+        self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True
+    ):
+        super().__init__()
+
+        self.weight = nn.Parameter(
+            torch.randn(out_channel, in_channel, kernel_size, kernel_size)
+        )
+        self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2)
+
+        self.stride = stride
+        self.padding = padding
+
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(out_channel))
+
+        else:
+            self.bias = None
+
+    def forward(self, input):
+        out = F.conv2d(
+            input,
+            self.weight * self.scale,
+            bias=self.bias,
+            stride=self.stride,
+            padding=self.padding,
+        )
+
+        return out
+
+    def __repr__(self):
+        return (
+            f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},'
+            f' {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})'
+        )
+
+
+class EqualLinear(nn.Module):
+    def __init__(
+        self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None
+    ):
+        super().__init__()
+
+        self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul))
+
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init))
+
+        else:
+            self.bias = None
+
+        self.activation = activation
+
+        self.scale = (1 / math.sqrt(in_dim)) * lr_mul
+        self.lr_mul = lr_mul
+
+    def forward(self, input):
+        if self.activation:
+            out = F.linear(input, self.weight * self.scale)
+            out = fused_leaky_relu(out, self.bias * self.lr_mul)
+
+        else:
+            out = F.linear(
+                input, self.weight * self.scale, bias=self.bias * self.lr_mul
+            )
+
+        return out
+
+    def __repr__(self):
+        return (
+            f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})'
+        )
+
+
+class ScaledLeakyReLU(nn.Module):
+    def __init__(self, negative_slope=0.2):
+        super().__init__()
+
+        self.negative_slope = negative_slope
+
+    def forward(self, input):
+        out = F.leaky_relu(input, negative_slope=self.negative_slope)
+
+        return out * math.sqrt(2)
+
+
+class ModulatedConv2d(nn.Module):
+    def __init__(
+        self,
+        in_channel,
+        out_channel,
+        kernel_size,
+        style_dim,
+        demodulate=True,
+        upsample=False,
+        downsample=False,
+        blur_kernel=[1, 3, 3, 1],
+    ):
+        super().__init__()
+
+        self.eps = 1e-8
+        self.kernel_size = kernel_size
+        self.in_channel = in_channel
+        self.out_channel = out_channel
+        self.upsample = upsample
+        self.downsample = downsample
+
+        if upsample:
+            factor = 2
+            p = (len(blur_kernel) - factor) - (kernel_size - 1)
+            pad0 = (p + 1) // 2 + factor - 1
+            pad1 = p // 2 + 1
+
+            self.blur = Blur(blur_kernel, pad=(pad0, pad1), upsample_factor=factor)
+
+        if downsample:
+            factor = 2
+            p = (len(blur_kernel) - factor) + (kernel_size - 1)
+            pad0 = (p + 1) // 2
+            pad1 = p // 2
+
+            self.blur = Blur(blur_kernel, pad=(pad0, pad1))
+
+        fan_in = in_channel * kernel_size ** 2
+        self.scale = 1 / math.sqrt(fan_in)
+        self.padding = kernel_size // 2
+
+        self.weight = nn.Parameter(
+            torch.randn(1, out_channel, in_channel, kernel_size, kernel_size)
+        )
+
+        self.modulation = EqualLinear(style_dim, in_channel, bias_init=1)
+
+        self.demodulate = demodulate
+
+    def __repr__(self):
+        return (
+            f'{self.__class__.__name__}({self.in_channel}, {self.out_channel}, {self.kernel_size}, '
+            f'upsample={self.upsample}, downsample={self.downsample})'
+        )
+
+    def forward(self, input, style):
+        batch, in_channel, height, width = input.shape
+
+        style = self.modulation(style).view(batch, 1, in_channel, 1, 1)
+        weight = self.scale * self.weight * style
+
+        if self.demodulate:
+            demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + 1e-8)
+            weight = weight * demod.view(batch, self.out_channel, 1, 1, 1)
+
+        weight = weight.view(
+            batch * self.out_channel, in_channel, self.kernel_size, self.kernel_size
+        )
+
+        if self.upsample:
+            input = input.view(1, batch * in_channel, height, width)
+            weight = weight.view(
+                batch, self.out_channel, in_channel, self.kernel_size, self.kernel_size
+            )
+            weight = weight.transpose(1, 2).reshape(
+                batch * in_channel, self.out_channel, self.kernel_size, self.kernel_size
+            )
+            out = F.conv_transpose2d(input, weight, padding=0, stride=2, groups=batch)
+            _, _, height, width = out.shape
+            out = out.view(batch, self.out_channel, height, width)
+            out = self.blur(out)
+
+        elif self.downsample:
+            input = self.blur(input)
+            _, _, height, width = input.shape
+            input = input.view(1, batch * in_channel, height, width)
+            out = F.conv2d(input, weight, padding=0, stride=2, groups=batch)
+            _, _, height, width = out.shape
+            out = out.view(batch, self.out_channel, height, width)
+
+        else:
+            input = input.view(1, batch * in_channel, height, width)
+            out = F.conv2d(input, weight, padding=self.padding, groups=batch)
+            _, _, height, width = out.shape
+            out = out.view(batch, self.out_channel, height, width)
+
+        return out
+
+
+class NoiseInjection(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.weight = nn.Parameter(torch.zeros(1))
+
+    def forward(self, image, noise=None):
+        if noise is None:
+            batch, _, height, width = image.shape
+            noise = image.new_empty(batch, 1, height, width).normal_()
+
+        return image + self.weight * noise
+
+
+class ConstantInput(nn.Module):
+    def __init__(self, channel, size=4):
+        super().__init__()
+
+        self.input = nn.Parameter(torch.randn(1, channel, size, size))
+
+    def forward(self, input):
+        batch = input.shape[0]
+        out = self.input.repeat(batch, 1, 1, 1)
+
+        return out
+
+
+class StyledConv(nn.Module):
+    def __init__(
+        self,
+        in_channel,
+        out_channel,
+        kernel_size,
+        style_dim,
+        upsample=False,
+        blur_kernel=[1, 3, 3, 1],
+        demodulate=True,
+    ):
+        super().__init__()
+
+        self.conv = ModulatedConv2d(
+            in_channel,
+            out_channel,
+            kernel_size,
+            style_dim,
+            upsample=upsample,
+            blur_kernel=blur_kernel,
+            demodulate=demodulate,
+        )
+
+        self.noise = NoiseInjection()
+        # self.bias = nn.Parameter(torch.zeros(1, out_channel, 1, 1))
+        # self.activate = ScaledLeakyReLU(0.2)
+        self.activate = FusedLeakyReLU(out_channel)
+
+    def forward(self, input, style, noise=None):
+        out = self.conv(input, style)
+        out = self.noise(out, noise=noise)
+        # out = out + self.bias
+        out = self.activate(out)
+
+        return out
+
+
+class ToRGB(nn.Module):
+    def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[1, 3, 3, 1]):
+        super().__init__()
+
+        if upsample:
+            self.upsample = Upsample(blur_kernel)
+
+        self.conv = ModulatedConv2d(in_channel, 3, 1, style_dim, demodulate=False)
+        self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1))
+
+    def forward(self, input, style, skip=None):
+        out = self.conv(input, style)
+        style_modulated = out
+        out = out + self.bias
+
+        if skip is not None:
+            skip = self.upsample(skip)
+
+            out = out + skip
+
+        return out, style_modulated
+
+
+class Generator(nn.Module):
+    def __init__(
+        self,
+        size,
+        style_dim,
+        n_mlp,
+        channel_multiplier=2,
+        blur_kernel=[1, 3, 3, 1],
+        lr_mlp=0.01,
+    ):
+        super().__init__()
+
+        self.size = size
+
+        self.style_dim = style_dim
+
+        layers = [PixelNorm()]
+
+        for i in range(n_mlp):
+            layers.append(
+                EqualLinear(
+                    style_dim, style_dim, lr_mul=lr_mlp, activation='fused_lrelu'
+                )
+            )
+
+        self.style = nn.Sequential(*layers)
+
+        self.channels = {
+            4: 512,
+            8: 512,
+            16: 512,
+            32: 512,
+            64: 256 * channel_multiplier,
+            128: 128 * channel_multiplier,
+            256: 64 * channel_multiplier,
+            512: 32 * channel_multiplier,
+            1024: 16 * channel_multiplier,
+        }
+
+        self.input = ConstantInput(self.channels[4])
+        self.conv1 = StyledConv(
+            self.channels[4], self.channels[4], 3, style_dim, blur_kernel=blur_kernel
+        )
+        self.to_rgb1 = ToRGB(self.channels[4], style_dim, upsample=False)
+
+        self.log_size = int(math.log(size, 2))
+        self.num_layers = (self.log_size - 2) * 2 + 1
+
+        self.convs = nn.ModuleList()
+        self.upsamples = nn.ModuleList()
+        self.to_rgbs = nn.ModuleList()
+        self.noises = nn.Module()
+
+        in_channel = self.channels[4]
+
+        for layer_idx in range(self.num_layers):
+            res = (layer_idx + 5) // 2
+            shape = [1, 1, 2 ** res, 2 ** res]
+            self.noises.register_buffer(f'noise_{layer_idx}', torch.randn(*shape))
+
+        for i in range(3, self.log_size + 1):
+            out_channel = self.channels[2 ** i]
+
+            self.convs.append(
+                StyledConv(
+                    in_channel,
+                    out_channel,
+                    3,
+                    style_dim,
+                    upsample=True,
+                    blur_kernel=blur_kernel,
+                )
+            )
+
+            self.convs.append(
+                StyledConv(
+                    out_channel, out_channel, 3, style_dim, blur_kernel=blur_kernel
+                )
+            )
+
+            self.to_rgbs.append(ToRGB(out_channel, style_dim))
+
+            in_channel = out_channel
+
+        self.n_latent = self.log_size * 2 - 2
+
+    @property
+    def device(self):
+        # TODO if multi-gpu is expected, could use the following more expensive version
+        #device, = list(set(p.device for p in self.parameters()))
+        return next(self.parameters()).device
+
+    @staticmethod
+    def get_latent_size(size):
+        log_size = int(math.log(size, 2))
+        return log_size * 2 - 2
+
+    @staticmethod
+    def make_noise_by_size(size: int, device: torch.device):
+        log_size = int(math.log(size, 2))
+        noises = [torch.randn(1, 1, 2 ** 2, 2 ** 2, device=device)]
+
+        for i in range(3, log_size + 1):
+            for _ in range(2):
+                noises.append(torch.randn(1, 1, 2 ** i, 2 ** i, device=device))
+
+        return noises
+
+
+    def make_noise(self):
+        return self.make_noise_by_size(self.size, self.input.input.device)
+
+    def mean_latent(self, n_latent):
+        latent_in = torch.randn(
+            n_latent, self.style_dim, device=self.input.input.device
+        )
+        latent = self.style(latent_in).mean(0, keepdim=True)
+
+        return latent
+
+    def get_latent(self, input):
+        return self.style(input)
+
+    def forward(
+        self,
+        styles,
+        return_latents=False,
+        inject_index=None,
+        truncation=1,
+        truncation_latent=None,
+        input_is_latent=False,
+        noise=None,
+        randomize_noise=True,
+    ):
+        if not input_is_latent:
+            styles = [self.style(s) for s in styles]
+
+        if noise is None:
+            if randomize_noise:
+                noise = [None] * self.num_layers
+            else:
+                noise = [
+                    getattr(self.noises, f'noise_{i}') for i in range(self.num_layers)
+                ]
+
+        if truncation < 1:
+            style_t = []
+
+            for style in styles:
+                style_t.append(
+                    truncation_latent + truncation * (style - truncation_latent)
+                )
+
+            styles = style_t
+
+        if len(styles) < 2:
+            inject_index = self.n_latent
+
+            if styles[0].ndim < 3:
+                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
+
+            else:
+                latent = styles[0]
+
+        else:
+            if inject_index is None:
+                inject_index = random.randint(1, self.n_latent - 1)
+
+            latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
+            latent2 = styles[1].unsqueeze(1).repeat(1, self.n_latent - inject_index, 1)
+
+            latent = torch.cat([latent, latent2], 1)
+
+        out = self.input(latent)
+        out = self.conv1(out, latent[:, 0], noise=noise[0])
+
+        skip, rgb_mod = self.to_rgb1(out, latent[:, 1])
+
+
+        rgbs = [rgb_mod] # all but the last skip
+        i = 1
+        for conv1, conv2, noise1, noise2, to_rgb in zip(
+            self.convs[::2], self.convs[1::2], noise[1::2], noise[2::2], self.to_rgbs
+        ):
+            out = conv1(out, latent[:, i], noise=noise1)
+            out = conv2(out, latent[:, i + 1], noise=noise2)
+            skip, rgb_mod = to_rgb(out, latent[:, i + 2], skip)
+            rgbs.append(rgb_mod)
+
+            i += 2
+
+        image = skip
+
+        if return_latents:
+            return image, latent, rgbs
+
+        else:
+            return image, None, rgbs
+
+
+class ConvLayer(nn.Sequential):
+    def __init__(
+        self,
+        in_channel,
+        out_channel,
+        kernel_size,
+        downsample=False,
+        blur_kernel=[1, 3, 3, 1],
+        bias=True,
+        activate=True,
+    ):
+        layers = []
+
+        if downsample:
+            factor = 2
+            p = (len(blur_kernel) - factor) + (kernel_size - 1)
+            pad0 = (p + 1) // 2
+            pad1 = p // 2
+
+            layers.append(Blur(blur_kernel, pad=(pad0, pad1)))
+
+            stride = 2
+            self.padding = 0
+
+        else:
+            stride = 1
+            self.padding = kernel_size // 2
+
+        layers.append(
+            EqualConv2d(
+                in_channel,
+                out_channel,
+                kernel_size,
+                padding=self.padding,
+                stride=stride,
+                bias=bias and not activate,
+            )
+        )
+
+        if activate:
+            if bias:
+                layers.append(FusedLeakyReLU(out_channel))
+
+            else:
+                layers.append(ScaledLeakyReLU(0.2))
+
+        super().__init__(*layers)
+
+
+class ResBlock(nn.Module):
+    def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]):
+        super().__init__()
+
+        self.conv1 = ConvLayer(in_channel, in_channel, 3)
+        self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True)
+
+        self.skip = ConvLayer(
+            in_channel, out_channel, 1, downsample=True, activate=False, bias=False
+        )
+
+    def forward(self, input):
+        out = self.conv1(input)
+        out = self.conv2(out)
+
+        skip = self.skip(input)
+        out = (out + skip) / math.sqrt(2)
+
+        return out
+
+
+class Discriminator(nn.Module):
+    def __init__(self, size, channel_multiplier=2, blur_kernel=[1, 3, 3, 1]):
+        super().__init__()
+
+        channels = {
+            4: 512,
+            8: 512,
+            16: 512,
+            32: 512,
+            64: 256 * channel_multiplier,
+            128: 128 * channel_multiplier,
+            256: 64 * channel_multiplier,
+            512: 32 * channel_multiplier,
+            1024: 16 * channel_multiplier,
+        }
+
+        convs = [ConvLayer(3, channels[size], 1)]
+
+        log_size = int(math.log(size, 2))
+
+        in_channel = channels[size]
+
+        for i in range(log_size, 2, -1):
+            out_channel = channels[2 ** (i - 1)]
+
+            convs.append(ResBlock(in_channel, out_channel, blur_kernel))
+
+            in_channel = out_channel
+
+        self.convs = nn.Sequential(*convs)
+
+        self.stddev_group = 4
+        self.stddev_feat = 1
+
+        self.final_conv = ConvLayer(in_channel + 1, channels[4], 3)
+        self.final_linear = nn.Sequential(
+            EqualLinear(channels[4] * 4 * 4, channels[4], activation='fused_lrelu'),
+            EqualLinear(channels[4], 1),
+        )
+
+    def forward(self, input):
+        out = self.convs(input)
+
+        batch, channel, height, width = out.shape
+        group = min(batch, self.stddev_group)
+        stddev = out.view(
+            group, -1, self.stddev_feat, channel // self.stddev_feat, height, width
+        )
+        stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-8)
+        stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
+        stddev = stddev.repeat(group, 1, height, width)
+        out = torch.cat([out, stddev], 1)
+
+        out = self.final_conv(out)
+
+        out = out.view(batch, -1)
+        out = self.final_linear(out)
+
+        return out
+

Time_TravelRephotography/models/degrade.py

@@ -0,0 +1,122 @@
+from argparse import (
+    ArgumentParser,
+    Namespace,
+)
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from utils.misc import optional_string
+
+from .gaussian_smoothing import GaussianSmoothing
+
+
+class DegradeArguments:
+    @staticmethod
+    def add_arguments(parser: ArgumentParser):
+        parser.add_argument('--spectral_sensitivity', choices=["g", "b", "gb"], default="g",
+            help="Type of spectral sensitivity. g: grayscale (panchromatic), b: blue-sensitive, gb: green+blue (orthochromatic)")
+        parser.add_argument('--gaussian', type=float, default=0,
+            help="estimated blur radius in pixels of the input photo if it is scaled to 1024x1024")
+
+    @staticmethod
+    def to_string(args: Namespace) -> str:
+        return (
+            f"{args.spectral_sensitivity}"
+            + optional_string(args.gaussian > 0, f"-G{args.gaussian}")
+        )
+
+
+class CameraResponse(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.register_parameter("gamma", nn.Parameter(torch.ones(1)))
+        self.register_parameter("offset", nn.Parameter(torch.zeros(1)))
+        self.register_parameter("gain", nn.Parameter(torch.ones(1)))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = torch.clamp(x, max=1, min=-1+1e-2)
+        x = (1 + x) * 0.5
+        x = self.offset + self.gain * torch.pow(x, self.gamma)
+        x = (x - 0.5) * 2
+        # b = torch.clamp(b, max=1, min=-1)
+        return x
+
+
+class SpectralResponse(nn.Module):
+    # TODO: use enum instead for color mode
+    def __init__(self, spectral_sensitivity: str = 'b'):
+        assert spectral_sensitivity in ("g", "b", "gb"), f"spectral_sensitivity {spectral_sensitivity} is not implemented."
+
+        super().__init__()
+
+        self.spectral_sensitivity = spectral_sensitivity
+
+        if self.spectral_sensitivity == "g":
+            self.register_buffer("to_gray", torch.tensor([0.299, 0.587, 0.114]).reshape(1, -1, 1, 1))
+
+    def forward(self, rgb: torch.Tensor) -> torch.Tensor:
+        if self.spectral_sensitivity == "b":
+            x = rgb[:, -1:]
+        elif self.spectral_sensitivity == "gb":
+            x = (rgb[:, 1:2] + rgb[:, -1:]) * 0.5
+        else:
+            assert self.spectral_sensitivity == "g"
+            x = (rgb * self.to_gray).sum(dim=1, keepdim=True)
+        return x
+
+
+class Downsample(nn.Module):
+    """Antialiasing downsampling"""
+    def __init__(self, input_size: int, output_size: int, channels: int):
+        super().__init__()
+        if input_size % output_size == 0:
+            self.stride = input_size // output_size
+            self.grid = None
+        else:
+            self.stride = 1
+            step = input_size / output_size
+            x = torch.arange(output_size) * step
+            Y, X = torch.meshgrid(x, x)
+            grid = torch.stack((X, Y), dim=-1)
+            grid /= torch.Tensor((input_size - 1, input_size - 1)).view(1, 1, -1)
+            grid = grid * 2 - 1
+            self.register_buffer("grid", grid)
+        sigma = 0.5 * input_size / output_size
+        #print(f"{input_size} -> {output_size}: sigma={sigma}")
+        self.blur = GaussianSmoothing(channels, int(2 * (sigma * 2) + 1 + 0.5), sigma)
+
+    def forward(self, im: torch.Tensor):
+        out = self.blur(im, stride=self.stride)
+        if self.grid is not None:
+            out = F.grid_sample(out, self.grid[None].expand(im.shape[0], -1, -1, -1))
+        return out
+
+
+
+class Degrade(nn.Module):
+    """
+    Simulate the degradation of antique film
+    """
+    def __init__(self, args:Namespace):
+        super().__init__()
+        self.srf = SpectralResponse(args.spectral_sensitivity)
+        self.crf = CameraResponse()
+        self.gaussian = None
+        if args.gaussian is not None and args.gaussian > 0:
+            self.gaussian = GaussianSmoothing(3, 2 * int(args.gaussian * 2 + 0.5) + 1, args.gaussian)
+
+    def forward(self, img: torch.Tensor, downsample: nn.Module = None):
+        if self.gaussian is not None:
+            img = self.gaussian(img)
+        if downsample is not None:
+            img = downsample(img)
+        img = self.srf(img)
+        img = self.crf(img)
+        # Note that I changed it back to 3 channels
+        return img.repeat((1, 3, 1, 1)) if img.shape[1] == 1 else img
+
+
+

Time_TravelRephotography/models/encoder.py

@@ -0,0 +1,66 @@
+from argparse import Namespace, ArgumentParser
+from functools import partial
+
+from torch import nn
+
+from .resnet import ResNetBasicBlock, activation_func, norm_module, Conv2dAuto
+
+
+def add_arguments(parser: ArgumentParser) -> ArgumentParser:
+    parser.add_argument("--latent_size", type=int, default=512, help="latent size")
+    return parser
+
+
+def create_model(args) -> nn.Module:
+    in_channels = 3 if "rgb" in args and args.rgb else 1
+    return Encoder(in_channels, args.encoder_size, latent_size=args.latent_size)
+
+
+class Flatten(nn.Module):
+    def forward(self, input_):
+        return input_.view(input_.size(0), -1)
+
+
+class Encoder(nn.Module):
+    def __init__(
+            self, in_channels: int, size: int, latent_size: int = 512,
+            activation: str = 'leaky_relu', norm: str = "instance"
+    ):
+        super().__init__()
+
+        out_channels0 = 64
+        norm_m = norm_module(norm)
+        self.conv0 = nn.Sequential(
+            Conv2dAuto(in_channels, out_channels0, kernel_size=5),
+            norm_m(out_channels0),
+            activation_func(activation),
+        )
+
+        pool_kernel = 2
+        self.pool = nn.AvgPool2d(pool_kernel)
+
+        num_channels = [128, 256, 512, 512]
+        # FIXME: this is a hack
+        if size >= 256:
+            num_channels.append(512)
+
+        residual = partial(ResNetBasicBlock, activation=activation, norm=norm, bias=True)
+        residual_blocks = nn.ModuleList()
+        for in_channel, out_channel in zip([out_channels0] + num_channels[:-1], num_channels):
+            residual_blocks.append(residual(in_channel, out_channel))
+            residual_blocks.append(nn.AvgPool2d(pool_kernel))
+        self.residual_blocks = nn.Sequential(*residual_blocks)
+
+        self.last = nn.Sequential(
+            nn.ReLU(),
+            nn.AvgPool2d(4),    # TODO: not sure whehter this would cause problem
+            Flatten(),
+            nn.Linear(num_channels[-1], latent_size, bias=True)
+        )
+
+    def forward(self, input_):
+        out = self.conv0(input_)
+        out = self.pool(out)
+        out = self.residual_blocks(out)
+        out = self.last(out)
+        return out

Time_TravelRephotography/models/encoder4editing/.gitignore

@@ -0,0 +1,133 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# Custom dataset
+pretrained_models
+results_test

Time_TravelRephotography/models/encoder4editing/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 omertov
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

Time_TravelRephotography/models/encoder4editing/README.md

@@ -0,0 +1,143 @@
+# Designing an Encoder for StyleGAN Image Manipulation (SIGGRAPH 2021)
+  <a href="https://arxiv.org/abs/2102.02766"><img src="https://img.shields.io/badge/arXiv-2008.00951-b31b1b.svg"></a>
+  <a href="https://opensource.org/licenses/MIT"><img src="https://img.shields.io/badge/License-MIT-yellow.svg"></a>
+  [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](http://colab.research.google.com/github/omertov/encoder4editing/blob/main/notebooks/inference_playground.ipynb)
+
+> Recently, there has been a surge of diverse methods for performing image editing by employing pre-trained unconditional generators. Applying these methods on real images, however, remains a challenge, as it necessarily requires the inversion of the images into their latent space. To successfully invert a real image, one needs to find a latent code that reconstructs the input image accurately, and more importantly, allows for its meaningful manipulation. In this paper, we carefully study the latent space of StyleGAN, the state-of-the-art unconditional generator. We identify and analyze the existence of a distortion-editability tradeoff and a distortion-perception tradeoff within the StyleGAN latent space. We then suggest two principles for designing encoders in a manner that allows one to control the proximity of the inversions to regions that StyleGAN was originally trained on. We present an encoder based on our two principles that is specifically designed for facilitating editing on real images by balancing these tradeoffs. By evaluating its performance qualitatively and quantitatively on numerous challenging domains, including cars and horses, we show that our inversion method, followed by common editing techniques, achieves superior real-image editing quality, with only a small reconstruction accuracy drop.
+
+<p align="center">
+<img src="docs/teaser.jpg" width="800px"/>
+</p>
+
+## Description   
+Official Implementation of "<a href="https://arxiv.org/abs/2102.02766">Designing an Encoder for StyleGAN Image Manipulation</a>" paper for both training and evaluation. 
+The e4e encoder is specifically designed to complement existing image manipulation techniques performed over StyleGAN's latent space.
+
+## Recent Updates
+`2021.08.17`: Add single style code encoder (use `--encoder_type SingleStyleCodeEncoder`). <br />
+`2021.03.25`: Add pose editing direction.
+
+## Getting Started
+### Prerequisites
+- Linux or macOS
+- NVIDIA GPU + CUDA CuDNN (CPU may be possible with some modifications, but is not inherently supported)
+- Python 3
+
+### Installation
+- Clone the repository:
+``` 
+git clone https://github.com/omertov/encoder4editing.git
+cd encoder4editing
+```
+- Dependencies:  
+We recommend running this repository using [Anaconda](https://docs.anaconda.com/anaconda/install/). 
+All dependencies for defining the environment are provided in `environment/e4e_env.yaml`.
+
+### Inference Notebook
+We provide a Jupyter notebook found in `notebooks/inference_playground.ipynb` that allows one to encode and perform several editings on real images using StyleGAN.   
+
+### Pretrained Models
+Please download the pre-trained models from the following links. Each e4e model contains the entire pSp framework architecture, including the encoder and decoder weights.
+| Path | Description
+| :--- | :----------
+|[FFHQ Inversion](https://drive.google.com/file/d/1cUv_reLE6k3604or78EranS7XzuVMWeO/view?usp=sharing)  | FFHQ e4e encoder.
+|[Cars Inversion](https://drive.google.com/file/d/17faPqBce2m1AQeLCLHUVXaDfxMRU2QcV/view?usp=sharing)  | Cars e4e encoder.
+|[Horse Inversion](https://drive.google.com/file/d/1TkLLnuX86B_BMo2ocYD0kX9kWh53rUVX/view?usp=sharing)  | Horse e4e encoder.
+|[Church Inversion](https://drive.google.com/file/d/1-L0ZdnQLwtdy6-A_Ccgq5uNJGTqE7qBa/view?usp=sharing) | Church e4e encoder.
+
+If you wish to use one of the pretrained models for training or inference, you may do so using the flag `--checkpoint_path`.
+
+In addition, we provide various auxiliary models needed for training your own e4e model from scratch.
+| Path | Description
+| :--- | :----------
+|[FFHQ StyleGAN](https://drive.google.com/file/d/1EM87UquaoQmk17Q8d5kYIAHqu0dkYqdT/view?usp=sharing) | StyleGAN model pretrained on FFHQ taken from [rosinality](https://github.com/rosinality/stylegan2-pytorch) with 1024x1024 output resolution.
+|[IR-SE50 Model](https://drive.google.com/file/d/1KW7bjndL3QG3sxBbZxreGHigcCCpsDgn/view?usp=sharing) | Pretrained IR-SE50 model taken from [TreB1eN](https://github.com/TreB1eN/InsightFace_Pytorch) for use in our ID loss during training.
+|[MOCOv2 Model](https://drive.google.com/file/d/18rLcNGdteX5LwT7sv_F7HWr12HpVEzVe/view?usp=sharing) | Pretrained ResNet-50 model trained using MOCOv2 for use in our simmilarity loss for domains other then human faces during training.
+
+By default, we assume that all auxiliary models are downloaded and saved to the directory `pretrained_models`. However, you may use your own paths by changing the necessary values in `configs/path_configs.py`. 
+
+## Training
+To train the e4e encoder, make sure the paths to the required models, as well as training and testing data is configured in `configs/path_configs.py` and `configs/data_configs.py`.
+#### **Training the e4e Encoder**
+```
+python scripts/train.py \
+--dataset_type cars_encode \
+--exp_dir new/experiment/directory \
+--start_from_latent_avg \
+--use_w_pool \
+--w_discriminator_lambda 0.1 \
+--progressive_start 20000 \
+--id_lambda 0.5 \
+--val_interval 10000 \
+--max_steps 200000 \
+--stylegan_size 512 \
+--stylegan_weights path/to/pretrained/stylegan.pt \
+--workers 8 \
+--batch_size 8 \
+--test_batch_size 4 \
+--test_workers 4 
+```
+
+#### Training on your own dataset
+In order to train the e4e encoder on a custom dataset, perform the following adjustments:
+1. Insert the paths to your train and test data into the `dataset_paths` variable defined in `configs/paths_config.py`:
+```
+dataset_paths = {
+    'my_train_data': '/path/to/train/images/directory',
+    'my_test_data': '/path/to/test/images/directory'
+}
+```
+2. Configure a new dataset under the DATASETS variable defined in `configs/data_configs.py`:
+```
+DATASETS = {
+   'my_data_encode': {
+        'transforms': transforms_config.EncodeTransforms,
+        'train_source_root': dataset_paths['my_train_data'],
+        'train_target_root': dataset_paths['my_train_data'],
+        'test_source_root': dataset_paths['my_test_data'],
+        'test_target_root': dataset_paths['my_test_data']
+    }
+}
+```
+Refer to `configs/transforms_config.py` for the transformations applied to the train and test images during training. 
+
+3. Finally, run a training session with `--dataset_type my_data_encode`.
+
+## Inference
+Having trained your model, you can use `scripts/inference.py` to apply the model on a set of images.   
+For example, 
+```
+python scripts/inference.py \
+--images_dir=/path/to/images/directory \
+--save_dir=/path/to/saving/directory \
+path/to/checkpoint.pt 
+```
+
+## Latent Editing Consistency (LEC)
+As described in the paper, we suggest a new metric, Latent Editing Consistency (LEC), for evaluating the encoder's 
+performance.
+We provide an example for calculating the metric over the FFHQ StyleGAN using the aging editing direction in 
+`metrics/LEC.py`.     
+
+To run the example: 
+```
+cd metrics
+python LEC.py \
+--images_dir=/path/to/images/directory \
+path/to/checkpoint.pt 
+```
+
+## Acknowledgments
+This code borrows heavily from [pixel2style2pixel](https://github.com/eladrich/pixel2style2pixel)
+
+## Citation
+If you use this code for your research, please cite our paper <a href="https://arxiv.org/abs/2102.02766">Designing an Encoder for StyleGAN Image Manipulation</a>:
+
+```
+@article{tov2021designing,
+  title={Designing an Encoder for StyleGAN Image Manipulation},
+  author={Tov, Omer and Alaluf, Yuval and Nitzan, Yotam and Patashnik, Or and Cohen-Or, Daniel},
+  journal={arXiv preprint arXiv:2102.02766},
+  year={2021}
+}
+```

Time_TravelRephotography/models/encoder4editing/__init__.py

@@ -0,0 +1,15 @@
+from .utils.model_utils import setup_model
+
+
+def get_latents(net, x, is_cars=False):
+    codes = net.encoder(x)
+    if net.opts.start_from_latent_avg:
+        if codes.ndim == 2:
+            codes = codes + net.latent_avg.repeat(codes.shape[0], 1, 1)[:, 0, :]
+        else:
+            codes = codes + net.latent_avg.repeat(codes.shape[0], 1, 1)
+    if codes.shape[1] == 18 and is_cars:
+        codes = codes[:, :16, :]
+    return codes
+
+

Time_TravelRephotography/models/encoder4editing/bash_scripts/inference.sh

@@ -0,0 +1,15 @@
+set -exo
+
+list="$1"
+ckpt="${2:-pretrained_models/e4e_ffhq_encode.pt}"
+
+base_dir="$REPHOTO/dataset/historically_interesting/aligned/manual_celebrity_in_19th_century/tier1/${list}/"
+save_dir="results_test/${list}/"
+
+
+TORCH_EXTENSIONS_DIR=/tmp/torch_extensions
+PYTHONPATH="" \
+python scripts/inference.py \
+    --images_dir="${base_dir}" \
+    --save_dir="${save_dir}" \
+    "${ckpt}"

Time_TravelRephotography/models/encoder4editing/configs/data_configs.py

@@ -0,0 +1,41 @@
+from configs import transforms_config
+from configs.paths_config import dataset_paths
+
+
+DATASETS = {
+	'ffhq_encode': {
+		'transforms': transforms_config.EncodeTransforms,
+		'train_source_root': dataset_paths['ffhq'],
+		'train_target_root': dataset_paths['ffhq'],
+		'test_source_root': dataset_paths['celeba_test'],
+		'test_target_root': dataset_paths['celeba_test'],
+	},
+	'cars_encode': {
+		'transforms': transforms_config.CarsEncodeTransforms,
+		'train_source_root': dataset_paths['cars_train'],
+		'train_target_root': dataset_paths['cars_train'],
+		'test_source_root': dataset_paths['cars_test'],
+		'test_target_root': dataset_paths['cars_test'],
+	},
+	'horse_encode': {
+		'transforms': transforms_config.EncodeTransforms,
+		'train_source_root': dataset_paths['horse_train'],
+		'train_target_root': dataset_paths['horse_train'],
+		'test_source_root': dataset_paths['horse_test'],
+		'test_target_root': dataset_paths['horse_test'],
+	},
+	'church_encode': {
+		'transforms': transforms_config.EncodeTransforms,
+		'train_source_root': dataset_paths['church_train'],
+		'train_target_root': dataset_paths['church_train'],
+		'test_source_root': dataset_paths['church_test'],
+		'test_target_root': dataset_paths['church_test'],
+	},
+	'cats_encode': {
+		'transforms': transforms_config.EncodeTransforms,
+		'train_source_root': dataset_paths['cats_train'],
+		'train_target_root': dataset_paths['cats_train'],
+		'test_source_root': dataset_paths['cats_test'],
+		'test_target_root': dataset_paths['cats_test'],
+	}
+}

Time_TravelRephotography/models/encoder4editing/configs/paths_config.py

@@ -0,0 +1,28 @@
+dataset_paths = {
+	#  Face Datasets (In the paper: FFHQ - train, CelebAHQ - test)
+	'ffhq': '',
+	'celeba_test': '',
+
+	#  Cars Dataset (In the paper: Stanford cars)
+	'cars_train': '',
+	'cars_test': '',
+
+	#  Horse Dataset (In the paper: LSUN Horse)
+	'horse_train': '',
+	'horse_test': '',
+
+	#  Church Dataset (In the paper: LSUN Church)
+	'church_train': '',
+	'church_test': '',
+
+	#  Cats Dataset (In the paper: LSUN Cat)
+	'cats_train': '',
+	'cats_test': ''
+}
+
+model_paths = {
+	'stylegan_ffhq': 'pretrained_models/stylegan2-ffhq-config-f.pt',
+	'ir_se50': 'pretrained_models/model_ir_se50.pth',
+	'shape_predictor': 'pretrained_models/shape_predictor_68_face_landmarks.dat',
+	'moco': 'pretrained_models/moco_v2_800ep_pretrain.pth'
+}

Time_TravelRephotography/models/encoder4editing/configs/transforms_config.py

@@ -0,0 +1,62 @@
+from abc import abstractmethod
+import torchvision.transforms as transforms
+
+
+class TransformsConfig(object):
+
+	def __init__(self, opts):
+		self.opts = opts
+
+	@abstractmethod
+	def get_transforms(self):
+		pass
+
+
+class EncodeTransforms(TransformsConfig):
+
+	def __init__(self, opts):
+		super(EncodeTransforms, self).__init__(opts)
+
+	def get_transforms(self):
+		transforms_dict = {
+			'transform_gt_train': transforms.Compose([
+				transforms.Resize((256, 256)),
+				transforms.RandomHorizontalFlip(0.5),
+				transforms.ToTensor(),
+				transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+			'transform_source': None,
+			'transform_test': transforms.Compose([
+				transforms.Resize((256, 256)),
+				transforms.ToTensor(),
+				transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+			'transform_inference': transforms.Compose([
+				transforms.Resize((256, 256)),
+				transforms.ToTensor(),
+				transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+		}
+		return transforms_dict
+
+
+class CarsEncodeTransforms(TransformsConfig):
+
+	def __init__(self, opts):
+		super(CarsEncodeTransforms, self).__init__(opts)
+
+	def get_transforms(self):
+		transforms_dict = {
+			'transform_gt_train': transforms.Compose([
+				transforms.Resize((192, 256)),
+				transforms.RandomHorizontalFlip(0.5),
+				transforms.ToTensor(),
+				transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+			'transform_source': None,
+			'transform_test': transforms.Compose([
+				transforms.Resize((192, 256)),
+				transforms.ToTensor(),
+				transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+			'transform_inference': transforms.Compose([
+				transforms.Resize((192, 256)),
+				transforms.ToTensor(),
+				transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+		}
+		return transforms_dict

Time_TravelRephotography/models/encoder4editing/criteria/id_loss.py

@@ -0,0 +1,47 @@
+import torch
+from torch import nn
+from configs.paths_config import model_paths
+from models.encoders.model_irse import Backbone
+
+
+class IDLoss(nn.Module):
+    def __init__(self):
+        super(IDLoss, self).__init__()
+        print('Loading ResNet ArcFace')
+        self.facenet = Backbone(input_size=112, num_layers=50, drop_ratio=0.6, mode='ir_se')
+        self.facenet.load_state_dict(torch.load(model_paths['ir_se50']))
+        self.face_pool = torch.nn.AdaptiveAvgPool2d((112, 112))
+        self.facenet.eval()
+        for module in [self.facenet, self.face_pool]:
+            for param in module.parameters():
+                param.requires_grad = False
+
+    def extract_feats(self, x):
+        x = x[:, :, 35:223, 32:220]  # Crop interesting region
+        x = self.face_pool(x)
+        x_feats = self.facenet(x)
+        return x_feats
+
+    def forward(self, y_hat, y, x):
+        n_samples = x.shape[0]
+        x_feats = self.extract_feats(x)
+        y_feats = self.extract_feats(y)  # Otherwise use the feature from there
+        y_hat_feats = self.extract_feats(y_hat)
+        y_feats = y_feats.detach()
+        loss = 0
+        sim_improvement = 0
+        id_logs = []
+        count = 0
+        for i in range(n_samples):
+            diff_target = y_hat_feats[i].dot(y_feats[i])
+            diff_input = y_hat_feats[i].dot(x_feats[i])
+            diff_views = y_feats[i].dot(x_feats[i])
+            id_logs.append({'diff_target': float(diff_target),
+                            'diff_input': float(diff_input),
+                            'diff_views': float(diff_views)})
+            loss += 1 - diff_target
+            id_diff = float(diff_target) - float(diff_views)
+            sim_improvement += id_diff
+            count += 1
+
+        return loss / count, sim_improvement / count, id_logs

Time_TravelRephotography/models/encoder4editing/criteria/lpips/lpips.py

@@ -0,0 +1,35 @@
+import torch
+import torch.nn as nn
+
+from criteria.lpips.networks import get_network, LinLayers
+from criteria.lpips.utils import get_state_dict
+
+
+class LPIPS(nn.Module):
+    r"""Creates a criterion that measures
+    Learned Perceptual Image Patch Similarity (LPIPS).
+    Arguments:
+        net_type (str): the network type to compare the features:
+                        'alex' | 'squeeze' | 'vgg'. Default: 'alex'.
+        version (str): the version of LPIPS. Default: 0.1.
+    """
+    def __init__(self, net_type: str = 'alex', version: str = '0.1'):
+
+        assert version in ['0.1'], 'v0.1 is only supported now'
+
+        super(LPIPS, self).__init__()
+
+        # pretrained network
+        self.net = get_network(net_type).to("cuda")
+
+        # linear layers
+        self.lin = LinLayers(self.net.n_channels_list).to("cuda")
+        self.lin.load_state_dict(get_state_dict(net_type, version))
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        feat_x, feat_y = self.net(x), self.net(y)
+
+        diff = [(fx - fy) ** 2 for fx, fy in zip(feat_x, feat_y)]
+        res = [l(d).mean((2, 3), True) for d, l in zip(diff, self.lin)]
+
+        return torch.sum(torch.cat(res, 0)) / x.shape[0]

Time_TravelRephotography/models/encoder4editing/criteria/lpips/networks.py

@@ -0,0 +1,96 @@
+from typing import Sequence
+
+from itertools import chain
+
+import torch
+import torch.nn as nn
+from torchvision import models
+
+from criteria.lpips.utils import normalize_activation
+
+
+def get_network(net_type: str):
+    if net_type == 'alex':
+        return AlexNet()
+    elif net_type == 'squeeze':
+        return SqueezeNet()
+    elif net_type == 'vgg':
+        return VGG16()
+    else:
+        raise NotImplementedError('choose net_type from [alex, squeeze, vgg].')
+
+
+class LinLayers(nn.ModuleList):
+    def __init__(self, n_channels_list: Sequence[int]):
+        super(LinLayers, self).__init__([
+            nn.Sequential(
+                nn.Identity(),
+                nn.Conv2d(nc, 1, 1, 1, 0, bias=False)
+            ) for nc in n_channels_list
+        ])
+
+        for param in self.parameters():
+            param.requires_grad = False
+
+
+class BaseNet(nn.Module):
+    def __init__(self):
+        super(BaseNet, self).__init__()
+
+        # register buffer
+        self.register_buffer(
+            'mean', torch.Tensor([-.030, -.088, -.188])[None, :, None, None])
+        self.register_buffer(
+            'std', torch.Tensor([.458, .448, .450])[None, :, None, None])
+
+    def set_requires_grad(self, state: bool):
+        for param in chain(self.parameters(), self.buffers()):
+            param.requires_grad = state
+
+    def z_score(self, x: torch.Tensor):
+        return (x - self.mean) / self.std
+
+    def forward(self, x: torch.Tensor):
+        x = self.z_score(x)
+
+        output = []
+        for i, (_, layer) in enumerate(self.layers._modules.items(), 1):
+            x = layer(x)
+            if i in self.target_layers:
+                output.append(normalize_activation(x))
+            if len(output) == len(self.target_layers):
+                break
+        return output
+
+
+class SqueezeNet(BaseNet):
+    def __init__(self):
+        super(SqueezeNet, self).__init__()
+
+        self.layers = models.squeezenet1_1(True).features
+        self.target_layers = [2, 5, 8, 10, 11, 12, 13]
+        self.n_channels_list = [64, 128, 256, 384, 384, 512, 512]
+
+        self.set_requires_grad(False)
+
+
+class AlexNet(BaseNet):
+    def __init__(self):
+        super(AlexNet, self).__init__()
+
+        self.layers = models.alexnet(True).features
+        self.target_layers = [2, 5, 8, 10, 12]
+        self.n_channels_list = [64, 192, 384, 256, 256]
+
+        self.set_requires_grad(False)
+
+
+class VGG16(BaseNet):
+    def __init__(self):
+        super(VGG16, self).__init__()
+
+        self.layers = models.vgg16(True).features
+        self.target_layers = [4, 9, 16, 23, 30]
+        self.n_channels_list = [64, 128, 256, 512, 512]
+
+        self.set_requires_grad(False)

Time_TravelRephotography/models/encoder4editing/criteria/lpips/utils.py

@@ -0,0 +1,30 @@
+from collections import OrderedDict
+
+import torch
+
+
+def normalize_activation(x, eps=1e-10):
+    norm_factor = torch.sqrt(torch.sum(x ** 2, dim=1, keepdim=True))
+    return x / (norm_factor + eps)
+
+
+def get_state_dict(net_type: str = 'alex', version: str = '0.1'):
+    # build url
+    url = 'https://raw.githubusercontent.com/richzhang/PerceptualSimilarity/' \
+        + f'master/lpips/weights/v{version}/{net_type}.pth'
+
+    # download
+    old_state_dict = torch.hub.load_state_dict_from_url(
+        url, progress=True,
+        map_location=None if torch.cuda.is_available() else torch.device('cpu')
+    )
+
+    # rename keys
+    new_state_dict = OrderedDict()
+    for key, val in old_state_dict.items():
+        new_key = key
+        new_key = new_key.replace('lin', '')
+        new_key = new_key.replace('model.', '')
+        new_state_dict[new_key] = val
+
+    return new_state_dict

Time_TravelRephotography/models/encoder4editing/criteria/moco_loss.py

@@ -0,0 +1,71 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from configs.paths_config import model_paths
+
+
+class MocoLoss(nn.Module):
+
+    def __init__(self, opts):
+        super(MocoLoss, self).__init__()
+        print("Loading MOCO model from path: {}".format(model_paths["moco"]))
+        self.model = self.__load_model()
+        self.model.eval()
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+    @staticmethod
+    def __load_model():
+        import torchvision.models as models
+        model = models.__dict__["resnet50"]()
+        # freeze all layers but the last fc
+        for name, param in model.named_parameters():
+            if name not in ['fc.weight', 'fc.bias']:
+                param.requires_grad = False
+        checkpoint = torch.load(model_paths['moco'], map_location="cpu")
+        state_dict = checkpoint['state_dict']
+        # rename moco pre-trained keys
+        for k in list(state_dict.keys()):
+            # retain only encoder_q up to before the embedding layer
+            if k.startswith('module.encoder_q') and not k.startswith('module.encoder_q.fc'):
+                # remove prefix
+                state_dict[k[len("module.encoder_q."):]] = state_dict[k]
+            # delete renamed or unused k
+            del state_dict[k]
+        msg = model.load_state_dict(state_dict, strict=False)
+        assert set(msg.missing_keys) == {"fc.weight", "fc.bias"}
+        # remove output layer
+        model = nn.Sequential(*list(model.children())[:-1]).cuda()
+        return model
+
+    def extract_feats(self, x):
+        x = F.interpolate(x, size=224)
+        x_feats = self.model(x)
+        x_feats = nn.functional.normalize(x_feats, dim=1)
+        x_feats = x_feats.squeeze()
+        return x_feats
+
+    def forward(self, y_hat, y, x):
+        n_samples = x.shape[0]
+        x_feats = self.extract_feats(x)
+        y_feats = self.extract_feats(y)
+        y_hat_feats = self.extract_feats(y_hat)
+        y_feats = y_feats.detach()
+        loss = 0
+        sim_improvement = 0
+        sim_logs = []
+        count = 0
+        for i in range(n_samples):
+            diff_target = y_hat_feats[i].dot(y_feats[i])
+            diff_input = y_hat_feats[i].dot(x_feats[i])
+            diff_views = y_feats[i].dot(x_feats[i])
+            sim_logs.append({'diff_target': float(diff_target),
+                             'diff_input': float(diff_input),
+                             'diff_views': float(diff_views)})
+            loss += 1 - diff_target
+            sim_diff = float(diff_target) - float(diff_views)
+            sim_improvement += sim_diff
+            count += 1
+
+        return loss / count, sim_improvement / count, sim_logs

Time_TravelRephotography/models/encoder4editing/criteria/w_norm.py

@@ -0,0 +1,14 @@
+import torch
+from torch import nn
+
+
+class WNormLoss(nn.Module):
+
+	def __init__(self, start_from_latent_avg=True):
+		super(WNormLoss, self).__init__()
+		self.start_from_latent_avg = start_from_latent_avg
+
+	def forward(self, latent, latent_avg=None):
+		if self.start_from_latent_avg:
+			latent = latent - latent_avg
+		return torch.sum(latent.norm(2, dim=(1, 2))) / latent.shape[0]

Time_TravelRephotography/models/encoder4editing/datasets/gt_res_dataset.py

@@ -0,0 +1,32 @@
+#!/usr/bin/python
+# encoding: utf-8
+import os
+from torch.utils.data import Dataset
+from PIL import Image
+import torch
+
+class GTResDataset(Dataset):
+
+	def __init__(self, root_path, gt_dir=None, transform=None, transform_train=None):
+		self.pairs = []
+		for f in os.listdir(root_path):
+			image_path = os.path.join(root_path, f)
+			gt_path = os.path.join(gt_dir, f)
+			if f.endswith(".jpg") or f.endswith(".png"):
+				self.pairs.append([image_path, gt_path.replace('.png', '.jpg'), None])
+		self.transform = transform
+		self.transform_train = transform_train
+
+	def __len__(self):
+		return len(self.pairs)
+
+	def __getitem__(self, index):
+		from_path, to_path, _ = self.pairs[index]
+		from_im = Image.open(from_path).convert('RGB')
+		to_im = Image.open(to_path).convert('RGB')
+
+		if self.transform:
+			to_im = self.transform(to_im)
+			from_im = self.transform(from_im)
+
+		return from_im, to_im

Time_TravelRephotography/models/encoder4editing/datasets/images_dataset.py

@@ -0,0 +1,33 @@
+from torch.utils.data import Dataset
+from PIL import Image
+from utils import data_utils
+
+
+class ImagesDataset(Dataset):
+
+	def __init__(self, source_root, target_root, opts, target_transform=None, source_transform=None):
+		self.source_paths = sorted(data_utils.make_dataset(source_root))
+		self.target_paths = sorted(data_utils.make_dataset(target_root))
+		self.source_transform = source_transform
+		self.target_transform = target_transform
+		self.opts = opts
+
+	def __len__(self):
+		return len(self.source_paths)
+
+	def __getitem__(self, index):
+		from_path = self.source_paths[index]
+		from_im = Image.open(from_path)
+		from_im = from_im.convert('RGB')
+
+		to_path = self.target_paths[index]
+		to_im = Image.open(to_path).convert('RGB')
+		if self.target_transform:
+			to_im = self.target_transform(to_im)
+
+		if self.source_transform:
+			from_im = self.source_transform(from_im)
+		else:
+			from_im = to_im
+
+		return from_im, to_im

Time_TravelRephotography/models/encoder4editing/datasets/inference_dataset.py

@@ -0,0 +1,25 @@
+from torch.utils.data import Dataset
+from PIL import Image
+from utils import data_utils
+
+
+class InferenceDataset(Dataset):
+
+	def __init__(self, root, opts, transform=None, preprocess=None):
+		self.paths = sorted(data_utils.make_dataset(root))
+		self.transform = transform
+		self.preprocess = preprocess
+		self.opts = opts
+
+	def __len__(self):
+		return len(self.paths)
+
+	def __getitem__(self, index):
+		from_path = self.paths[index]
+		if self.preprocess is not None:
+			from_im = self.preprocess(from_path)
+		else:
+			from_im = Image.open(from_path).convert('RGB')
+		if self.transform:
+			from_im = self.transform(from_im)
+		return from_im

Time_TravelRephotography/models/encoder4editing/editings/ganspace.py

@@ -0,0 +1,22 @@
+import torch
+
+
+def edit(latents, pca, edit_directions):
+    edit_latents = []
+    for latent in latents:
+        for pca_idx, start, end, strength in edit_directions:
+            delta = get_delta(pca, latent, pca_idx, strength)
+            delta_padded = torch.zeros(latent.shape).to('cuda')
+            delta_padded[start:end] += delta.repeat(end - start, 1)
+            edit_latents.append(latent + delta_padded)
+    return torch.stack(edit_latents)
+
+
+def get_delta(pca, latent, idx, strength):
+    # pca: ganspace checkpoint. latent: (16, 512) w+
+    w_centered = latent - pca['mean'].to('cuda')
+    lat_comp = pca['comp'].to('cuda')
+    lat_std = pca['std'].to('cuda')
+    w_coord = torch.sum(w_centered[0].reshape(-1)*lat_comp[idx].reshape(-1)) / lat_std[idx]
+    delta = (strength - w_coord)*lat_comp[idx]*lat_std[idx]
+    return delta

Time_TravelRephotography/models/encoder4editing/editings/ganspace_pca/cars_pca.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5c3bae61ecd85de077fbbf103f5f30cf4b7676fe23a8508166eaf2ce73c8392
+size 167562

Time_TravelRephotography/models/encoder4editing/editings/ganspace_pca/ffhq_pca.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d7f9df1c96180d9026b9cb8d04753579fbf385f321a9d0e263641601c5e5d36
+size 167562

Time_TravelRephotography/models/encoder4editing/editings/interfacegan_directions/age.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:50074516b1629707d89b5e43d6b8abd1792212fa3b961a87a11323d6a5222ae0
+size 2808