Initial refactor

Author: alex-damian

.gitignore

@@ -0,0 +1,5 @@
+.DS_Store
+.idea/*
+runs/*
+input/*
+cache/*

PULSE.py

@@ -0,0 +1,176 @@
+from stylegan import G_synthesis,G_mapping
+from dataclasses import dataclass
+from SphericalOptimizer import SphericalOptimizer
+from pathlib import Path
+import numpy as np
+import time
+import torch
+from loss import LossBuilder
+from functools import partial
+from drive import open_url
+
+
+class PULSE(torch.nn.Module):
+    def __init__(self, cache_dir):
+        super(PULSE, self).__init__()
+        self.synthesis = G_synthesis().cuda()
+
+        cache_dir = Path(cache_dir)
+        cache_dir.mkdir(parents=True, exist_ok = True)
+
+        print("Loading Synthesis Network")
+        with open_url("https://drive.google.com/uc?id=1TCViX1YpQyRsklTVYEJwdbmK91vklCo8", cache_dir=cache_dir) as f:
+            self.synthesis.load_state_dict(torch.load(f))
+
+        for param in self.synthesis.parameters():
+            param.requires_grad = False
+
+        self.lrelu = torch.nn.LeakyReLU(negative_slope=0.2)
+
+        if Path("gaussian_fit.pt").exists():
+            self.gaussian_fit = torch.load("gaussian_fit.pt")
+        else:
+            print("Fitting Linear Layer to Mapping Network")
+            print("\tLoading Mapping Network")
+            mapping = G_mapping().cuda()
+
+            with open_url("https://drive.google.com/uc?id=14R6iHGf5iuVx3DMNsACAl7eBr7Vdpd0k", cache_dir=cache_dir) as f:
+                    mapping.load_state_dict(torch.load(f))
+
+            print("\tRunning Mapping Network")
+            with torch.no_grad():
+                torch.manual_seed(0)
+                latent = torch.randn((1000000,512),dtype=torch.float32, device="cuda")
+                latent_out = torch.nn.LeakyReLU(5)(mapping(latent))
+                self.gaussian_fit = {"mean": latent_out.mean(0), "std": latent_out.std(0)}
+                torch.save(self.gaussian_fit,"gaussian_fit.pt")
+                print("\tSaved \"gaussian_fit.pt\"")
+
+    def forward(self, ref_im,
+                seed,
+                loss_str,
+                eps,
+                noise_type,
+                num_trainable_noise_layers,
+                tile_latent,
+                bad_noise_layers,
+                opt_name,
+                learning_rate,
+                steps,
+                lr_schedule,
+                save_intermediate,
+                **kwargs):
+
+        if seed:
+            torch.manual_seed(seed)
+            torch.cuda.manual_seed(seed)
+            torch.backends.cudnn.deterministic = True
+
+        batch_size = ref_im.shape[0]
+
+        # Generate latent tensor
+        if(tile_latent):
+            latent = torch.randn(
+                (batch_size, 1, 512), dtype=torch.float, requires_grad=True, device='cuda')
+        else:
+            latent = torch.randn(
+                (batch_size, 18, 512), dtype=torch.float, requires_grad=True, device='cuda')
+
+        # Generate list of noise tensors
+        noise = [] # stores all of the noise tensors
+        noise_vars = []  # stores the noise tensors that we want to optimize on
+
+        for i in range(18):
+            # dimension of the ith noise tensor
+            res = (batch_size, 1, 2**(i//2+2), 2**(i//2+2))
+
+            if(noise_type == 'zero' or i in [int(layer) for layer in bad_noise_layers.split('.')]):
+                new_noise = torch.zeros(res, dtype=torch.float, device='cuda')
+                new_noise.requires_grad = False
+            elif(noise_type == 'fixed'):
+                new_noise = torch.randn(res, dtype=torch.float, device='cuda')
+                new_noise.requires_grad = False
+            elif (noise_type == 'trainable'):
+                new_noise = torch.randn(res, dtype=torch.float, device='cuda')
+                if (i < num_trainable_noise_layers):
+                    new_noise.requires_grad = True
+                    noise_vars.append(new_noise)
+                else:
+                    new_noise.requires_grad = False
+            else:
+                raise Exception("unknown noise type")
+
+            noise.append(new_noise)
+
+        var_list = [latent]+noise_vars
+
+        opt_dict = {
+            'sgd': torch.optim.SGD,
+            'adam': torch.optim.Adam,
+            'sgdm': partial(torch.optim.SGD, momentum=0.9),
+            'adamax': torch.optim.Adamax
+        }
+        opt_func = opt_dict[opt_name]
+        opt = SphericalOptimizer(opt_func, var_list, lr=learning_rate)
+
+        schedule_dict = {
+            'fixed': lambda x: 1,
+            'linear1cycle': lambda x: (9*(1-np.abs(x/steps-1/2)*2)+1)/10,
+            'linear1cycledrop': lambda x: (9*(1-np.abs(x/(0.9*steps)-1/2)*2)+1)/10 if x < 0.9*steps else 1/10 + (x-0.9*steps)/(0.1*steps)*(1/1000-1/10),
+        }
+        schedule_func = schedule_dict[lr_schedule]
+        scheduler = torch.optim.lr_scheduler.LambdaLR(opt.opt, schedule_func)
+        
+        loss_builder = LossBuilder(ref_im, loss_str, eps).cuda()
+
+        min_loss = np.inf
+        best_summary = ""
+        start_t = time.time()
+        if(save_intermediate):
+            int_HR = []
+            int_LR = []
+
+        print("Optimizing")
+        for j in range(steps):
+            opt.opt.zero_grad()
+
+            # Duplicate latent in case tile_latent = True
+            if (tile_latent):
+                latent_in = latent.expand(-1, 18, -1)
+            else:
+                latent_in = latent
+
+            # Apply learned linear mapping to match latent distribution to that of the mapping network
+            latent_in = self.lrelu(latent_in*self.gaussian_fit["std"] + self.gaussian_fit["mean"])
+
+            # Normalize image to [0,1] instead of [-1,1]
+            gen_im = (self.synthesis(latent_in, noise)+1)/2
+
+            # Calculate Losses
+            loss, loss_dict = loss_builder(latent_in, gen_im)
+            loss_dict['TOTAL'] = loss
+
+            # Save intermediate HR and LR images
+            if(save_intermediate):
+                int_HR.append(gen_im.cpu().detach().clamp(0, 1))
+                int_LR.append(loss_builder.D(gen_im).cpu().detach().clamp(0, 1))
+
+            # Save best summary for log
+            if(loss < min_loss):
+                min_loss = loss
+                best_summary = f'BEST ({j+1}) | '+' | '.join(
+                [f'{x}: {y:.4f}' for x, y in loss_dict.items()])
+                best_im = gen_im.clone()
+
+            loss.backward()
+            opt.step()
+            scheduler.step()
+
+        total_t = time.time()-start_t
+        current_info = f' | time: {total_t:.1f} | it/s: {(j+1)/total_t:.2f} | batchsize: {batch_size}'
+        print(best_summary+current_info)
+
+        if(save_intermediate):
+            return best_im.cpu().detach().clamp(0,1), int_HR, int_LR
+        else:
+            return best_im.cpu().detach().clamp(0,1)

SphericalOptimizer.py

@@ -0,0 +1,26 @@
+import math
+import torch
+from torch.optim import Optimizer
+
+# Spherical Optimizer Class
+# Uses the first two dimensions as batch information
+# Optimizes over the surface of a sphere using the initial radius throughout
+#
+# Example Usage:
+# opt = SphericalOptimizer(torch.optim.SGD, [x], lr=0.01)
+
+class SphericalOptimizer(Optimizer):
+    def __init__(self, optimizer, params, **kwargs):
+        self.opt = optimizer(params, **kwargs)
+        self.params = params
+        with torch.no_grad():
+            self.radii = {param: (param.pow(2).sum(tuple(range(2,param.ndim)),keepdim=True)+1e-9).sqrt() for param in params}
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = self.opt.step(closure)
+        for param in self.params:
+            param.data.div_((param.pow(2).sum(tuple(range(2,param.ndim)),keepdim=True)+1e-9).sqrt())
+            param.mul_(self.radii[param])
+
+        return loss

bicubic.py

@@ -0,0 +1,75 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class BicubicDownSample(nn.Module):
+    def bicubic_kernel(self, x, a=-0.50):
+        """
+        This equation is exactly copied from the website below:
+        https://clouard.users.greyc.fr/Pantheon/experiments/rescaling/index-en.html#bicubic
+        """
+        abs_x = torch.abs(x)
+        if abs_x <= 1.:
+            return (a + 2.) * torch.pow(abs_x, 3.) - (a + 3.) * torch.pow(abs_x, 2.) + 1
+        elif 1. < abs_x < 2.:
+            return a * torch.pow(abs_x, 3) - 5. * a * torch.pow(abs_x, 2.) + 8. * a * abs_x - 4. * a
+        else:
+            return 0.0
+
+    def __init__(self, factor=4, cuda=True, padding='reflect'):
+        super().__init__()
+        self.factor = factor
+        size = factor * 4
+        k = torch.tensor([self.bicubic_kernel((i - torch.floor(torch.tensor(size / 2)) + 0.5) / factor)
+                          for i in range(size)], dtype=torch.float32)
+        k = k / torch.sum(k)
+        # k = torch.einsum('i,j->ij', (k, k))
+        k1 = torch.reshape(k, shape=(1, 1, size, 1))
+        self.k1 = torch.cat([k1, k1, k1], dim=0)
+        k2 = torch.reshape(k, shape=(1, 1, 1, size))
+        self.k2 = torch.cat([k2, k2, k2], dim=0)
+        self.cuda = '.cuda' if cuda else ''
+        self.padding = padding
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, x, nhwc=False, clip_round=False, byte_output=False):
+        # x = torch.from_numpy(x).type('torch.FloatTensor')
+        filter_height = self.factor * 4
+        filter_width = self.factor * 4
+        stride = self.factor
+
+        pad_along_height = max(filter_height - stride, 0)
+        pad_along_width = max(filter_width - stride, 0)
+        filters1 = self.k1.type('torch{}.FloatTensor'.format(self.cuda))
+        filters2 = self.k2.type('torch{}.FloatTensor'.format(self.cuda))
+
+        # compute actual padding values for each side
+        pad_top = pad_along_height // 2
+        pad_bottom = pad_along_height - pad_top
+        pad_left = pad_along_width // 2
+        pad_right = pad_along_width - pad_left
+
+        # apply mirror padding
+        if nhwc:
+            x = torch.transpose(torch.transpose(
+                x, 2, 3), 1, 2)   # NHWC to NCHW
+
+        # downscaling performed by 1-d convolution
+        x = F.pad(x, (0, 0, pad_top, pad_bottom), self.padding)
+        x = F.conv2d(input=x, weight=filters1, stride=(stride, 1), groups=3)
+        if clip_round:
+            x = torch.clamp(torch.round(x), 0.0, 255.)
+
+        x = F.pad(x, (pad_left, pad_right, 0, 0), self.padding)
+        x = F.conv2d(input=x, weight=filters2, stride=(1, stride), groups=3)
+        if clip_round:
+            x = torch.clamp(torch.round(x), 0.0, 255.)
+
+        if nhwc:
+            x = torch.transpose(torch.transpose(x, 1, 3), 1, 2)
+        if byte_output:
+            return x.type('torch.ByteTensor'.format(self.cuda))
+        else:
+            return x

drive.py

@@ -0,0 +1,90 @@
+# URL helpers, see https://github.com/NVlabs/stylegan
+# ------------------------------------------------------------------------------------------
+
+import requests
+import html
+import hashlib
+import glob
+import os
+import io
+from typing import Any
+import re
+import uuid
+
+def is_url(obj: Any) -> bool:
+    """Determine whether the given object is a valid URL string."""
+    if not isinstance(obj, str) or not "://" in obj:
+        return False
+    try:
+        res = requests.compat.urlparse(obj)
+        if not res.scheme or not res.netloc or not "." in res.netloc:
+            return False
+        res = requests.compat.urlparse(requests.compat.urljoin(obj, "/"))
+        if not res.scheme or not res.netloc or not "." in res.netloc:
+            return False
+    except:
+        return False
+    return True
+
+
+def open_url(url: str, cache_dir: str = None, num_attempts: int = 10, verbose: bool = True) -> Any:
+    """Download the given URL and return a binary-mode file object to access the data."""
+    assert is_url(url)
+    assert num_attempts >= 1
+
+    # Lookup from cache.
+    url_md5 = hashlib.md5(url.encode("utf-8")).hexdigest()
+    if cache_dir is not None:
+        cache_files = glob.glob(os.path.join(cache_dir, url_md5 + "_*"))
+        if len(cache_files) == 1:
+            return open(cache_files[0], "rb")
+
+    # Download.
+    url_name = None
+    url_data = None
+    with requests.Session() as session:
+        if verbose:
+            print("Downloading %s ..." % url, end="", flush=True)
+        for attempts_left in reversed(range(num_attempts)):
+            try:
+                with session.get(url) as res:
+                    res.raise_for_status()
+                    if len(res.content) == 0:
+                        raise IOError("No data received")
+
+                    if len(res.content) < 8192:
+                        content_str = res.content.decode("utf-8")
+                        if "download_warning" in res.headers.get("Set-Cookie", ""):
+                            links = [html.unescape(link) for link in content_str.split('"') if "export=download" in link]
+                            if len(links) == 1:
+                                url = requests.compat.urljoin(url, links[0])
+                                raise IOError("Google Drive virus checker nag")
+                        if "Google Drive - Quota exceeded" in content_str:
+                            raise IOError("Google Drive quota exceeded")
+
+                    match = re.search(r'filename="([^"]*)"', res.headers.get("Content-Disposition", ""))
+                    url_name = match[1] if match else url
+                    url_data = res.content
+                    if verbose:
+                        print(" done")
+                    break
+            except:
+                if not attempts_left:
+                    if verbose:
+                        print(" failed")
+                    raise
+                if verbose:
+                    print(".", end="", flush=True)
+
+    # Save to cache.
+    if cache_dir is not None:
+        safe_name = re.sub(r"[^0-9a-zA-Z-._]", "_", url_name)
+        cache_file = os.path.join(cache_dir, url_md5 + "_" + safe_name)
+        temp_file = os.path.join(cache_dir, "tmp_" + uuid.uuid4().hex + "_" + url_md5 + "_" + safe_name)
+        os.makedirs(cache_dir, exist_ok=True)
+        with open(temp_file, "wb") as f:
+            f.write(url_data)
+        os.replace(temp_file, cache_file) # atomic
+
+    # Return data as file object.
+    return io.BytesIO(url_data)