Initial commit

.gitignore

@@ -0,0 +1,7 @@
+**/__pycache__
+**/.vscode
+**/.DS_Store
+cache/
+out/
+checkpoints/
+**/.ipynb_checkpoints/

.gitmodules

@@ -0,0 +1,8 @@
+[submodule "stylegan/stylegan_tf"]
+	path = models/stylegan/stylegan_tf
+	url = https://github.com/NVlabs/stylegan.git
+	ignore = untracked
+[submodule "stylegan2/stylegan2-pytorch"]
+	path = models/stylegan2/stylegan2-pytorch
+	url = https://github.com/harskish/stylegan2-pytorch.git
+	ignore = untracked

LICENSE.txt

@@ -0,0 +1,6 @@
+Exclusive Copyright 2020, Erik Härkönen
+
+This code is released for the purpose of academic reproducibility.
+No license is granted for derivative works or other uses, besides non-commercial experimentation.
+
+We are working on adding a real non-commercial license.

README.md

@@ -0,0 +1,115 @@
+# GANSpace: Discovering Interpretable GAN Controls
+![Python 3.7](https://img.shields.io/badge/python-3.7-green.svg)
+![PyTorch 1.3](https://img.shields.io/badge/pytorch-1.3-green.svg)
+![teaser](teaser.jpg)
+<p align="justify"><b>Figure 1:</b> Sequences of image edits performed using control discovered with our method, applied to three different GANs. The white insets specify the particular edits using notation explained in Section 3.4 ('Layer-wise Edits').</p>
+
+
+> **GANSpace: Discovering Interpretable GAN Controls**<br>
+> Erik Härkönen<sup>1,2</sup>, Aaron Hertzmann<sup>2</sup>, Jaakko Lehtinen<sup>1,3</sup>, Sylvain Paris<sup>2</sup><br>
+> <sup>1</sup>Aalto University, <sup>2</sup>Adobe Research, <sup>3</sup>NVIDIA<br>
+> https://arxiv.org/abs/TODO
+>
+> <p align="justify"><b>Abstract:</b> <i>This paper describes a simple technique to analyze Generative Adversarial Networks (GANs) and create interpretable controls for image synthesis, such as change of viewpoint, aging, lighting, and time of day. We identify important latent directions based on Principal Components Analysis (PCA) applied in activation space. Then, we show that interpretable edits can be defined based on layer-wise application of these edit directions. Moreover, we show that BigGAN can be controlled with layer-wise inputs in a StyleGAN-like manner. A user may identify a large number of interpretable controls with these mechanisms. We demonstrate results on GANs from various datasets.</i></p>
+> <p align="justify"><b>Video:</b> 
+> https://youtu.be/jdTICDa_eAI
+
+## Setup
+See the [setup instructions](SETUP.md).
+
+## Usage
+This repository includes versions of BigGAN, StyleGAN, and StyleGAN2 modified to support per-layer latent vectors.
+
+**Interactive model exploration**
+```
+# Explore BigGAN-deep husky
+python interactive.py --model=BigGAN-512 --class=husky --layer=generator.gen_z -n=1000000
+
+# Explore StyleGAN2 ffhq in W space
+python interactive.py --model=StyleGAN2 --class=ffhq --layer=style --use_w -n=1000000 -b=10000
+
+# Explore StyleGAN2 cars in Z space
+python interactive.py --model=StyleGAN2 --class=car --layer=style -n=1000000 -b=10000
+```
+```
+# Apply previously saved edits interactively
+python interactive.py --model=StyleGAN2 --class=ffhq --layer=style --use_w --inputs=out/directions
+```
+
+**Visualize principal components**
+```
+# Visualize StyleGAN2 ffhq W principal components
+python visualize.py --model=StyleGAN2 --class=ffhq --use_w --layer=style -b=10000
+
+# Create videos of StyleGAN wikiart components (saved to ./out)
+python visualize.py --model=StyleGAN --class=wikiart --use_w --layer=g_mapping -b=10000 --batch --video
+```
+
+**Options**
+```
+Command line paramaters:
+  --model      one of [ProGAN, BigGAN-512, BigGAN-256, BigGAN-128, StyleGAN, StyleGAN2]
+  --class      class name; leave empty to list options
+  --layer      layer at which to perform PCA; leave empty to list options
+  --use_w      treat W as the main latent space (StyleGAN / StyleGAN2)
+  --inputs     load previously exported edits from directory
+  --sigma      number of stdevs to use in visualize.py
+  -n           number of PCA samples
+  -b           override automatic minibatch size detection
+  -c           number of components to keep
+```
+
+## Reproducibility
+All the figures presented in the paper can be recreated using the included Jupyter notebooks:
+* Figure 1: `figure_teaser.ipynb`
+* Figure 2: `figure_pca_illustration.ipynb`
+* Figure 3: `figure_pca_cleanup.ipynb`
+* Figure 4: `figure_steerability_comp.ipynb`
+* Figure 5: `figure_biggan_edit_transferability.ipynb`
+* Figure 6: `figure_biggan_style_mixing.ipybb`
+* Figure 7: `figure_biggan_style_resampling.ipynb`
+* Figure 8: `figure_edit_zoo.ipynb`
+
+## Known issues
+* The interactive viewer sometimes freezes on startup on Ubuntu 18.04. The freeze is resolved by clicking on the terminal window and pressing the control key. Any insight into the issue would be greatly appreciated!
+
+## Integrating a new model
+1. Create a wrapper for the model in `models/wrappers.py` using the `BaseModel` interface.
+2. Add the model to `get_model()` in `models/wrappers.py`.
+
+## Importing StyleGAN checkpoints from TensorFlow
+It is possible to import trained StyleGAN and StyleGAN2 weights from TensorFlow into GANSpace.
+
+### StyleGAN
+1. Install TensorFlow: `conda install tensorflow-gpu=1.*`.
+2. Modify methods `__init__()`, `load_model()` in `models/wrappers.py` under class StyleGAN.
+
+### StyleGAN2
+1. Follow instructions in stylegan2-pytorch's [instructions][stylegan2_pytorch].
+2. Save the converted checkpoint as `checkpoints/stylegan2/<dataset>_<resolution>.pt`.
+3. Modify methods `__init__()`, `download_checkpoint()` in `models/wrappers.py` under class StyleGAN2.
+
+## Acknowledgements
+We would like to thank:
+
+* The authors of the PyTorch implementations of [BigGAN][biggan_pytorch], [StyleGAN][stylegan_pytorch], and [StyleGAN2][stylegan2_pytorch]:<br>Thomas Wolf, Piotr Bialecki, Thomas Viehmann, and Kim Seonghyeon.
+* Joel Simon from ArtBreeder for providing us with the landscape model for StyleGAN.
+* David Bau and colleagues for the excellent [GAN Dissection][gandissect] project.
+* Justin Pinkney for the [Awesome Pretrained StyleGAN][pretrained_stylegan] collection.
+* Toumas Kynkäänniemi for giving us a helping hand with the experiments.
+* The Aalto Science-IT project for providing computational resources for this project.
+
+## License
+
+This code is released for the purpose of academic reproducibility.
+No license is granted for derivative works or other uses, besides non-commercial experimentation. <b>We are working on adding a real non-commercial license.</b>
+
+The directory `netdissect` is a derivative of the [GAN Dissection][gandissect] project, and is provided under the MIT license.<br>
+The directory `models/biggan` is provided under the MIT license.
+
+
+[biggan_pytorch]: https://github.com/huggingface/pytorch-pretrained-BigGAN
+[stylegan_pytorch]: https://github.com/lernapparat/lernapparat/blob/master/style_gan/pytorch_style_gan.ipynb
+[stylegan2_pytorch]: https://github.com/rosinality/stylegan2-pytorch
+[gandissect]: https://github.com/CSAILVision/GANDissect
+[pretrained_stylegan]: https://github.com/justinpinkney/awesome-pretrained-stylegan

SETUP.md

@@ -0,0 +1,33 @@
+## Setup
+1. Install anaconda or miniconda
+2. Create environment: `conda create -n ganspace python=3.7`
+3. Activate environment: `conda activate ganspace`
+4. Install dependencies: `conda env update -f environment.yml --prune`
+5. Setup submodules: `git submodule update --init --recursive`
+6. Run command `python -c "import nltk; nltk.download('wordnet')"`
+
+### Interactive viewer
+The interactive viewer (<i>interactive.py</i>) has the following dependencies:
+- Glumpy
+- PyCUDA with OpenGL support
+
+#### Windows
+Install included dependencies (downloaded from https://www.lfd.uci.edu/~gohlke/pythonlibs/):<br/> 
+`pip install deps/windows/*`
+
+#### Linux
+1. Install CUDA toolkit (match the version in environment.yml)
+2. Download pycuda sources from: https://pypi.org/project/pycuda/#files
+3. Extract files: `tar -xzf pycuda-VERSION.tar.gz`
+4. Configure: `python configure.py --cuda-enable-gl --cuda-root=/path/to/cuda`
+5. Compile and install: `make install`
+6. Install Glumpy: `pip install setuptools cython glumpy`
+
+### StyleGAN2
+The bundled StyleGAN2 model requires additional setup steps due to the custom CUDA kernels involved.
+1. Install CUDA toolkit (match the version in environment.yml)
+2. On Windows: install and open 'x64 Native Tools Command Prompt for VS 2017'
+3. `conda activate ganspace`
+4. `cd models/stylegan2/stylegan2-pytorch/op`
+5. `python setup.py install`
+6. Test: `python -c "import torch; import upfirdn2d_op; import fused; print('OK')"`

TkTorchWindow.py

@@ -0,0 +1,196 @@
+import tkinter as tk
+import numpy as np
+import time
+from contextlib import contextmanager
+import pycuda.driver
+from pycuda.gl import graphics_map_flags
+from glumpy import gloo, gl
+from pyopengltk import OpenGLFrame
+import torch
+from torch.autograd import Variable
+
+# TkInter widget that can draw torch tensors directly from GPU memory
+
+@contextmanager
+def cuda_activate(img):
+    """Context manager simplifying use of pycuda.gl.RegisteredImage"""
+    mapping = img.map()
+    yield mapping.array(0,0)
+    mapping.unmap()
+
+def create_shared_texture(w, h, c=4,
+        map_flags=graphics_map_flags.WRITE_DISCARD,
+        dtype=np.uint8):
+    """Create and return a Texture2D with gloo and pycuda views."""
+    tex = np.zeros((h,w,c), dtype).view(gloo.Texture2D)
+    tex.activate() # force gloo to create on GPU
+    tex.deactivate()
+    cuda_buffer = pycuda.gl.RegisteredImage(
+        int(tex.handle), tex.target, map_flags)
+    return tex, cuda_buffer
+
+# Shape batch as square if possible
+def get_grid_dims(B):
+    S = int(B**0.5 + 0.5)
+    while B % S != 0:
+        S -= 1
+    return (B // S, S)
+
+def create_gl_texture(tensor_shape):
+    if len(tensor_shape) != 4:
+        raise RuntimeError('Please provide a tensor of shape NCHW')
+
+    N, C, H, W = tensor_shape
+
+    cols, rows = get_grid_dims(N)
+    tex, cuda_buffer = create_shared_texture(W*cols, H*rows, 4)
+
+    return tex, cuda_buffer
+
+# Create window with OpenGL context
+class TorchImageView(OpenGLFrame):
+    def __init__(self, root = None, show_fps=True, **kwargs):
+        self.root = root or tk.Tk()
+        self.width = kwargs.get('width', 512)
+        self.width = kwargs.get('height', 512)
+        self.show_fps = show_fps
+        self.pycuda_initialized = False
+        self.animate = 0 # disable internal main loop
+        OpenGLFrame.__init__(self, root, **kwargs)
+
+    def initgl(self):
+        if not self.pycuda_initialized:
+            self.setup_gl(self.width, self.height)
+            self.pycuda_initialized = True
+
+        """Initalize gl states when the frame is created"""
+        gl.glViewport(0, 0, self.width, self.height)
+        gl.glClearColor(0.0, 0.0, 0.0, 0.0)
+        self.dt_history = [1000/60]
+        self.t0 = time.time()
+        self.t_last = self.t0
+        self.nframes = 0
+
+    def setup_gl(self, width, height):
+        # setup pycuda and torch
+        import pycuda.gl.autoinit
+        import pycuda.gl
+
+        assert torch.cuda.is_available(), "PyTorch: CUDA is not available"
+        print('Using GPU {}'.format(torch.cuda.current_device()))
+
+        # Create tensor to be shared between GL and CUDA
+        # Always overwritten so no sharing is necessary
+        dummy = torch.cuda.FloatTensor((1))
+        dummy.uniform_()
+        dummy = Variable(dummy)
+
+        # Create a buffer with pycuda and gloo views, using tensor created above
+        self.tex, self.cuda_buffer = create_gl_texture((1, 3, width, height))
+
+        # create a shader to program to draw to the screen
+        vertex = """
+        uniform float scale;
+        attribute vec2 position;
+        attribute vec2 texcoord;
+        varying vec2 v_texcoord;
+        void main()
+        {
+            v_texcoord = texcoord;
+            gl_Position = vec4(scale*position, 0.0, 1.0);
+        } """
+        fragment = """
+        uniform sampler2D tex;
+        varying vec2 v_texcoord;
+        void main()
+        {
+            gl_FragColor = texture2D(tex, v_texcoord);
+        } """
+        # Build the program and corresponding buffers (with 4 vertices)
+        self.screen = gloo.Program(vertex, fragment, count=4)
+
+        # NDC coordinates:         Texcoords:          Vertex order,
+        # (-1, +1)       (+1, +1)   (0,0)     (1,0)    triangle strip:
+        #        +-------+               +----+          1----3
+        #        |  NDC  |               |    |          |  / | 
+        #        | SPACE |               |    |          | /  |
+        #        +-------+               +----+          2----4
+        # (-1, -1)       (+1, -1)   (0,1)     (1,1)
+
+        # Upload data to GPU
+        self.screen['position'] = [(-1,+1), (-1,-1), (+1,+1), (+1,-1)]
+        self.screen['texcoord'] = [(0,0), (0,1), (1,0), (1,1)]
+        self.screen['scale'] = 1.0
+        self.screen['tex'] = self.tex
+
+    # Don't call directly, use update() instead
+    def redraw(self):
+        t_now = time.time()
+        dt = t_now - self.t_last
+        self.t_last = t_now
+
+        self.dt_history = ([dt] + self.dt_history)[:50]
+        dt_mean = sum(self.dt_history) / len(self.dt_history)
+
+        if self.show_fps and self.nframes % 60 == 0:
+            self.master.title('FPS: {:.0f}'.format(1 / dt_mean))
+
+    def draw(self, img):
+        assert len(img.shape) == 4, "Please provide an NCHW image tensor"
+        assert img.device.type == "cuda", "Please provide a CUDA tensor"
+
+        if img.dtype.is_floating_point:
+            img = (255*img).byte()
+
+        # Tile images
+        N, C, H, W = img.shape
+
+        if N > 1:
+            cols, rows = get_grid_dims(N)
+            img = img.reshape(cols, rows, C, H, W)
+            img = img.permute(2, 1, 3, 0, 4) # [C, rows, H, cols, W]
+            img = img.reshape(1, C, rows*H, cols*W)
+
+        tensor = img.squeeze().permute(1, 2, 0).data # CHW => HWC
+        if C == 3:
+            tensor = torch.cat((tensor, tensor[:,:,0:1]),2) # add the alpha channel
+            tensor[:,:,3] = 1 # set alpha
+
+        tensor = tensor.contiguous()
+
+        tex_h, tex_w, _ = self.tex.shape
+        tensor_h, tensor_w, _ = tensor.shape
+
+        if (tex_h, tex_w) != (tensor_h, tensor_w):
+            print(f'Resizing texture to {tensor_w}*{tensor_h}')
+            self.tex, self.cuda_buffer = create_gl_texture((N, C, H, W)) # original shape
+            self.screen['tex'] = self.tex
+
+        # copy from torch into buffer
+        assert self.tex.nbytes == tensor.numel()*tensor.element_size(), "Tensor and texture shape mismatch!"
+        with cuda_activate(self.cuda_buffer) as ary:
+            cpy = pycuda.driver.Memcpy2D()
+            cpy.set_src_device(tensor.data_ptr())
+            cpy.set_dst_array(ary)
+            cpy.width_in_bytes = cpy.src_pitch = cpy.dst_pitch = self.tex.nbytes//tensor_h
+            cpy.height = tensor_h
+            cpy(aligned=False)
+            torch.cuda.synchronize()
+
+        # draw to screen
+        self.screen.draw(gl.GL_TRIANGLE_STRIP)
+
+    def update(self):
+        self.update_idletasks()
+        self.tkMakeCurrent()
+        self.redraw()
+        self.tkSwapBuffers()
+
+# USAGE:
+# root = tk.Tk()
+# iv = TorchImageView(root, width=512, height=512)
+# iv.pack(fill='both', expand=True)
+# while True:
+#     iv.draw(nchw_tensor)
+#     root.update()
+#     iv.update()