Files for the app

.gitattributes

@@ -32,3 +32,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+deps/windows/glumpy-1.1.0-cp37-cp37m-win_amd64.whl filter=lfs diff=lfs merge=lfs -text
+deps/windows/PyOpenGL-3.1.4-cp37-cp37m-win_amd64.whl filter=lfs diff=lfs merge=lfs -text
+deps/windows/triangle-20190115.3-cp37-cp37m-win_amd64.whl filter=lfs diff=lfs merge=lfs -text
+models/stylegan/stylegan_tf/stylegan-teaser.png filter=lfs diff=lfs merge=lfs -text
+models/stylegan2/stylegan2-pytorch/doc/sample.png filter=lfs diff=lfs merge=lfs -text
+models/stylegan2/stylegan2-pytorch/doc/stylegan2-church-config-f.png filter=lfs diff=lfs merge=lfs -text

FashionGen.py

@@ -0,0 +1,265 @@
+import random
+import streamlit as st
+import torch
+import PIL
+import numpy as np
+from PIL import Image
+import imageio
+from models import get_instrumented_model
+from decomposition import get_or_compute
+from config import Config
+from skimage import img_as_ubyte
+import clip
+from torchvision.transforms import Resize, Normalize, Compose, CenterCrop
+from torch.optim import Adam
+from stqdm import stqdm
+
+torch.set_num_threads(8)
+
+# Speed up computation
+torch.autograd.set_grad_enabled(True)
+#torch.backends.cudnn.benchmark = True
+
+# Specify model to use
+config = Config(
+  model='StyleGAN2',
+  layer='style',
+  output_class= 'lookbook',
+  components=80,
+  use_w=True,
+  batch_size=5_000, # style layer quite small
+)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+preprocess = Compose([
+    Resize(224),
+    CenterCrop(224),
+    Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711)),
+])
+
+@st.cache_data
+def clip_optimized_latent(text, seed, iterations=25, lr=1e-2):
+    seed = int(seed)
+    text_input = clip.tokenize([text]).to(device)
+
+    # Initialize a random latent vector
+    latent_vector = model.sample_latent(1,seed=seed).detach()
+    latent_vector.requires_grad = True
+    latent_vector = [latent_vector]*model.get_max_latents()
+    params = [torch.nn.Parameter(latent_vector[i], requires_grad=True) for i in range(len(latent_vector))]
+    optimizer = Adam(params, lr=lr)
+
+    with torch.no_grad():
+        text_features = clip_model.encode_text(text_input)
+        
+    #pbar = tqdm(range(iterations), dynamic_ncols=True)
+    
+    for iteration in stqdm(range(iterations)):
+        optimizer.zero_grad()
+
+        # Generate an image from the latent vector
+        image = model.sample(params)
+        image = image.to(device)
+        
+        # Preprocess the image for the CLIP model
+        image = preprocess(image)
+        #image = clip_preprocess(Image.fromarray((image_np * 255).astype(np.uint8))).unsqueeze(0).to(device)
+        
+        # Extract features from the image
+        image_features = clip_model.encode_image(image)
+
+        # Calculate the loss and backpropagate
+        loss = -torch.cosine_similarity(text_features, image_features).mean()
+        loss.backward()
+        optimizer.step()
+        
+        #pbar.set_description(f"Loss: {loss.item()}")  # Update the progress bar to show the current loss
+        w = [param.detach().cpu().numpy() for param in params]
+    
+    return w
+          
+def mix_w(w1, w2, content, style):
+    for i in range(0,5):
+        w2[i] = w1[i] * (1 - content) + w2[i] * content
+
+    for i in range(5, 16):
+        w2[i] = w1[i] * (1 - style) + w2[i] * style
+    
+    return w2
+    
+def display_sample_pytorch(seed, truncation, directions, distances, scale, start, end, w=None, disp=True, save=None, noise_spec=None):
+    # blockPrint()
+    model.truncation = truncation
+    if w is None:
+        w = model.sample_latent(1, seed=seed).detach().cpu().numpy()
+        w = [w]*model.get_max_latents() # one per layer
+    else:
+        w_numpy = [x.detach().numpy() for x in w]
+        w = [np.expand_dims(x, 0) for x in w_numpy]
+        #w = [x.unsqueeze(0) for x in w]
+
+    
+    for l in range(start, end):
+      for i in range(len(directions)):
+        w[l] = w[l] + directions[i] * distances[i] * scale
+    
+    torch.cuda.empty_cache()
+    #save image and display
+    out = model.sample(w)
+    out = out.permute(0, 2, 3, 1).cpu().detach().numpy()
+    out = np.clip(out, 0.0, 1.0).squeeze()
+    
+    final_im = Image.fromarray((out * 255).astype(np.uint8)).resize((500,500),Image.LANCZOS)
+    
+    
+    if save is not None:
+      if disp == False:
+        print(save)
+      final_im.save(f'out/{seed}_{save:05}.png')
+    if disp:
+      display(final_im)
+    
+    return final_im
+
+## Generate image for app
+def generate_image(content, style, truncation, c0, c1, c2, c3, c4, c5, c6, start_layer, end_layer,w1,w2):
+
+    scale = 1
+    params = {'c0': c0,
+          'c1': c1,
+          'c2': c2,
+          'c3': c3,
+          'c4': c4,
+          'c5': c5,
+          'c6': c6}
+
+    param_indexes = {'c0': 0,
+              'c1': 1,
+              'c2': 2,
+              'c3': 3,
+              'c4': 4,
+              'c5': 5,
+              'c6': 6}
+
+    directions = []
+    distances = []
+    for k, v in params.items():
+        directions.append(latent_dirs[param_indexes[k]])
+        distances.append(v)
+            
+    if w1 is not None and w2 is not None:
+        w1 = [torch.from_numpy(x).to(device) for x in w1]
+        w2 = [torch.from_numpy(x).to(device) for x in w2]
+
+
+    #w1 = clip_optimized_latent(text1, seed1, iters)
+    im1 = model.sample(w1)
+    im1_np = im1.permute(0, 2, 3, 1).cpu().detach().numpy()
+    im1_np = np.clip(im1_np, 0.0, 1.0).squeeze()
+    
+    #w2 = clip_optimized_latent(text2, seed2, iters)
+    im2 = model.sample(w2)
+    im2_np = im2.permute(0, 2, 3, 1).cpu().detach().numpy()
+    im2_np = np.clip(im2_np, 0.0, 1.0).squeeze()
+
+    combined_im = np.concatenate([im1_np, im2_np], axis=1)
+    input_im = Image.fromarray((combined_im * 255).astype(np.uint8))
+    
+
+    mixed_w = mix_w(w1, w2, content, style)
+    return input_im, display_sample_pytorch(seed1, truncation, directions, distances, scale, int(start_layer), int(end_layer), w=mixed_w, disp=False)
+    
+    
+# Streamlit app title
+st.title("FashionGen Demo - AI assisted fashion design")
+"""This application employs the StyleGAN framework, CLIP and GANSpace exploration techniques to synthesize images of garments from textual inputs. With training based on the comprehensive LookBook dataset, it supports an efficient fashion design process by transforming text into visual concepts, showcasing the practical application of Generative Adversarial Networks (GANs) in the realm of creative design."""
+
+@st.cache_resource
+def load_model():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # Load the pre-trained CLIP model
+    clip_model, clip_preprocess = clip.load("ViT-B/32", device=device)
+    inst = get_instrumented_model(config.model, config.output_class,
+                              config.layer, torch.device('cpu'), use_w=config.use_w)
+    return clip_model, inst
+
+# Then, to load your models, call this function:
+clip_model, inst = load_model()
+model = inst.model
+
+
+path_to_components = get_or_compute(config, inst)
+comps = np.load(path_to_components)
+lst = comps.files
+latent_dirs = []
+latent_stdevs = []
+
+load_activations = False
+
+for item in lst:
+    if load_activations:
+      if item == 'act_comp':
+        for i in range(comps[item].shape[0]):
+          latent_dirs.append(comps[item][i])
+      if item == 'act_stdev':
+        for i in range(comps[item].shape[0]):
+          latent_stdevs.append(comps[item][i])
+    else:
+      if item == 'lat_comp':
+        for i in range(comps[item].shape[0]):
+          latent_dirs.append(comps[item][i])
+      if item == 'lat_stdev':
+        for i in range(comps[item].shape[0]):
+          latent_stdevs.append(comps[item][i])
+
+## Side bar texts
+st.sidebar.title('Tuning Parameters')
+st.sidebar.subheader('(CLIP + GANSpace)')
+
+
+# Create UI widgets
+
+if 'seed1' not in st.session_state and 'seed2' not in st.session_state:
+    st.session_state['seed1'] = random.randint(1, 1000)
+    st.session_state['seed2'] = random.randint(1, 1000)
+seed1 = st.sidebar.number_input("Seed 1", value= st.session_state['seed1'])
+seed2 = st.sidebar.number_input("Seed 2", value= st.session_state['seed2'])
+text1 = st.sidebar.text_input("Text Description 1")
+text2 = st.sidebar.text_input("Text Description 2")
+iters = st.sidebar.number_input("Iterations for CLIP Optimization", value = 25)
+submit_button = st.sidebar.button("Submit")
+content = st.sidebar.slider("Structural Composition", min_value=0.0, max_value=1.0, value=0.5)
+style = st.sidebar.slider("Style", min_value=0.0, max_value=1.0, value=0.5)
+truncation = st.sidebar.slider("Dimensional Scaling", min_value=0.0, max_value=1.0, value=0.5)
+
+slider_min_val = -20
+slider_max_val = 20
+slider_step = 1
+
+c0 = st.sidebar.slider("Sleeve Size Scaling", min_value=slider_min_val, max_value=slider_max_val, value=0)
+c1 = st.sidebar.slider("Jacket Features", min_value=slider_min_val, max_value=slider_max_val, value=0)
+c2 = st.sidebar.slider("Women's Overcoat", min_value=slider_min_val, max_value=slider_max_val, value=0)
+c3 = st.sidebar.slider("Coat", min_value=slider_min_val, max_value=slider_max_val, value=0)
+c4 = st.sidebar.slider("Graphic Elements", min_value=slider_min_val, max_value=slider_max_val, value=0)
+c5 = st.sidebar.slider("Darker Color", min_value=slider_min_val, max_value=slider_max_val, value=0)
+c6 = st.sidebar.slider("Modest Neckline", min_value=slider_min_val, max_value=slider_max_val, value=0)
+start_layer = st.sidebar.number_input("Start Layer", value=0)
+end_layer = st.sidebar.number_input("End Layer", value=14)
+
+
+
+if submit_button:  # Execute when the submit button is pressed
+    w1 = clip_optimized_latent(text1, seed1, iters)
+    st.session_state['w1-np'] = w1
+    w2 = clip_optimized_latent(text2, seed2, iters)
+    st.session_state['w2-np'] = w2
+
+try:
+    input_im, output_im = generate_image(content, style, truncation, c0, c1, c2, c3, c4, c5, c6, start_layer, end_layer,st.session_state['w1-np'],st.session_state['w2-np'])
+    st.image(input_im, caption="Input Image")
+    st.image(output_im, caption="Output Image")
+except:
+    pass
+
+        

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 prathmeshdahikar
+
+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.

README.md

@@ -1,13 +1,2 @@
----
-title: FashionGen
-emoji: 🦀
-colorFrom: indigo
-colorTo: indigo
-sdk: streamlit
-sdk_version: 1.19.0
-app_file: app.py
-pinned: false
-license: afl-3.0
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# FashionGen
+AI assisted fashion design

TkTorchWindow.py

@@ -0,0 +1,208 @@
+# Copyright 2020 Erik Härkönen. All rights reserved.
+# This file is licensed to you under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License. You may obtain a copy
+# of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software distributed under
+# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
+# OF ANY KIND, either express or implied. See the License for the specific language
+# governing permissions and limitations under the License.
+
+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.height = 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)
+
+    # Called by pyopengltk.BaseOpenGLFrame
+    # when the frame goes onto the screen
+    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()

bpe_simple_vocab_16e6.txt.gz

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

cache/components/stylegan2-lookbook_style_ipca_c80_n300000_w.npz

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

clip.py

@@ -0,0 +1,237 @@
+import hashlib
+import os
+import urllib
+import warnings
+from typing import Any, Union, List
+from pkg_resources import packaging
+
+import torch
+from PIL import Image
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
+from tqdm import tqdm
+
+from model_clip import build_model
+from simple_tokenizer import SimpleTokenizer as _Tokenizer
+
+try:
+    from torchvision.transforms import InterpolationMode
+    BICUBIC = InterpolationMode.BICUBIC
+except ImportError:
+    BICUBIC = Image.BICUBIC
+
+
+if packaging.version.parse(torch.__version__) < packaging.version.parse("1.7.1"):
+    warnings.warn("PyTorch version 1.7.1 or higher is recommended")
+
+
+__all__ = ["available_models", "load", "tokenize"]
+_tokenizer = _Tokenizer()
+
+_MODELS = {
+    "RN50": "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt",
+    "RN101": "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt",
+    "RN50x4": "https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt",
+    "RN50x16": "https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt",
+    "RN50x64": "https://openaipublic.azureedge.net/clip/models/be1cfb55d75a9666199fb2206c106743da0f6468c9d327f3e0d0a543a9919d9c/RN50x64.pt",
+    "ViT-B/32": "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt",
+    "ViT-B/16": "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt",
+    "ViT-L/14": "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt",
+    "ViT-L/14@336px": "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt",
+}
+
+
+def _download(url: str, root: str):
+    os.makedirs(root, exist_ok=True)
+    filename = os.path.basename(url)
+
+    expected_sha256 = url.split("/")[-2]
+    download_target = os.path.join(root, filename)
+
+    if os.path.exists(download_target) and not os.path.isfile(download_target):
+        raise RuntimeError(f"{download_target} exists and is not a regular file")
+
+    if os.path.isfile(download_target):
+        if hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256:
+            return download_target
+        else:
+            warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
+
+    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
+        with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True, unit_divisor=1024) as loop:
+            while True:
+                buffer = source.read(8192)
+                if not buffer:
+                    break
+
+                output.write(buffer)
+                loop.update(len(buffer))
+
+    if hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256:
+        raise RuntimeError("Model has been downloaded but the SHA256 checksum does not not match")
+
+    return download_target
+
+
+def _convert_image_to_rgb(image):
+    return image.convert("RGB")
+
+
+def _transform(n_px):
+    return Compose([
+        Resize(n_px, interpolation=BICUBIC),
+        CenterCrop(n_px),
+        _convert_image_to_rgb,
+        ToTensor(),
+        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+    ])
+
+
+def available_models() -> List[str]:
+    """Returns the names of available CLIP models"""
+    return list(_MODELS.keys())
+
+
+def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit: bool = False, download_root: str = None):
+    """Load a CLIP model
+
+    Parameters
+    ----------
+    name : str
+        A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
+
+    device : Union[str, torch.device]
+        The device to put the loaded model
+
+    jit : bool
+        Whether to load the optimized JIT model or more hackable non-JIT model (default).
+
+    download_root: str
+        path to download the model files; by default, it uses "~/.cache/clip"
+
+    Returns
+    -------
+    model : torch.nn.Module
+        The CLIP model
+
+    preprocess : Callable[[PIL.Image], torch.Tensor]
+        A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
+    """
+    if name in _MODELS:
+        model_path = _download(_MODELS[name], download_root or os.path.expanduser("~/.cache/clip"))
+    elif os.path.isfile(name):
+        model_path = name
+    else:
+        raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
+
+    with open(model_path, 'rb') as opened_file:
+        try:
+            # loading JIT archive
+            model = torch.jit.load(opened_file, map_location=device if jit else "cpu").eval()
+            state_dict = None
+        except RuntimeError:
+            # loading saved state dict
+            if jit:
+                warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
+                jit = False
+            state_dict = torch.load(opened_file, map_location="cpu")
+
+    if not jit:
+        model = build_model(state_dict or model.state_dict()).to(device)
+        if str(device) == "cpu":
+            model.float()
+        return model, _transform(model.visual.input_resolution)
+
+    # patch the device names
+    device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])
+    device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1]
+
+    def patch_device(module):
+        try:
+            graphs = [module.graph] if hasattr(module, "graph") else []
+        except RuntimeError:
+            graphs = []
+
+        if hasattr(module, "forward1"):
+            graphs.append(module.forward1.graph)
+
+        for graph in graphs:
+            for node in graph.findAllNodes("prim::Constant"):
+                if "value" in node.attributeNames() and str(node["value"]).startswith("cuda"):
+                    node.copyAttributes(device_node)
+
+    model.apply(patch_device)
+    patch_device(model.encode_image)
+    patch_device(model.encode_text)
+
+    # patch dtype to float32 on CPU
+    if str(device) == "cpu":
+        float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])
+        float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]
+        float_node = float_input.node()
+
+        def patch_float(module):
+            try:
+                graphs = [module.graph] if hasattr(module, "graph") else []
+            except RuntimeError:
+                graphs = []
+
+            if hasattr(module, "forward1"):
+                graphs.append(module.forward1.graph)
+
+            for graph in graphs:
+                for node in graph.findAllNodes("aten::to"):
+                    inputs = list(node.inputs())
+                    for i in [1, 2]:  # dtype can be the second or third argument to aten::to()
+                        if inputs[i].node()["value"] == 5:
+                            inputs[i].node().copyAttributes(float_node)
+
+        model.apply(patch_float)
+        patch_float(model.encode_image)
+        patch_float(model.encode_text)
+
+        model.float()
+
+    return model, _transform(model.input_resolution.item())
+
+
+def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> Union[torch.IntTensor, torch.LongTensor]:
+    """
+    Returns the tokenized representation of given input string(s)
+
+    Parameters
+    ----------
+    texts : Union[str, List[str]]
+        An input string or a list of input strings to tokenize
+
+    context_length : int
+        The context length to use; all CLIP models use 77 as the context length
+
+    truncate: bool
+        Whether to truncate the text in case its encoding is longer than the context length
+
+    Returns
+    -------
+    A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length].
+    We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long.
+    """
+    if isinstance(texts, str):
+        texts = [texts]
+
+    sot_token = _tokenizer.encoder["<|startoftext|>"]
+    eot_token = _tokenizer.encoder["<|endoftext|>"]
+    all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
+    if packaging.version.parse(torch.__version__) < packaging.version.parse("1.8.0"):
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+    else:
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)
+
+    for i, tokens in enumerate(all_tokens):
+        if len(tokens) > context_length:
+            if truncate:
+                tokens = tokens[:context_length]
+                tokens[-1] = eot_token
+            else:
+                raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}")
+        result[i, :len(tokens)] = torch.tensor(tokens)
+
+    return result

config.py

@@ -0,0 +1,72 @@
+# Copyright 2020 Erik Härkönen. All rights reserved.
+# This file is licensed to you under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License. You may obtain a copy
+# of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software distributed under
+# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
+# OF ANY KIND, either express or implied. See the License for the specific language
+# governing permissions and limitations under the License.
+
+import sys
+import argparse
+import json
+from copy import deepcopy
+
+class Config:
+    def __init__(self, **kwargs):
+        self.from_args([]) # set all defaults
+        self.default_args = deepcopy(self.__dict__)
+        self.from_dict(kwargs) # override
+
+    def __str__(self):
+        custom = {}
+        default = {}
+
+        # Find non-default arguments
+        for k, v in self.__dict__.items():
+            if k == 'default_args':
+                continue
+            
+            in_default = k in self.default_args
+            same_value = self.default_args.get(k) == v
+            
+            if in_default and same_value:
+                default[k] = v
+            else:
+                custom[k] = v
+
+        config = {
+            'custom': custom,
+            'default': default
+        }
+
+        return json.dumps(config, indent=4)
+    
+    def __repr__(self):
+        return self.__str__()
+    
+    def from_dict(self, dictionary):
+        for k, v in dictionary.items():
+            setattr(self, k, v)
+        return self
+    
+    def from_args(self, args=sys.argv[1:]):
+        parser = argparse.ArgumentParser(description='GAN component analysis config')
+        parser.add_argument('--model', dest='model', type=str, default='StyleGAN', help='The network to analyze') # StyleGAN, DCGAN, ProGAN, BigGAN-XYZ
+        parser.add_argument('--layer', dest='layer', type=str, default='g_mapping', help='The layer to analyze')
+        parser.add_argument('--class', dest='output_class', type=str, default=None, help='Output class to generate (BigGAN: Imagenet, ProGAN: LSUN)')
+        parser.add_argument('--est', dest='estimator', type=str, default='ipca', help='The algorithm to use [pca, fbpca, cupca, spca, ica]')
+        parser.add_argument('--sparsity', type=float, default=1.0, help='Sparsity parameter of SPCA')
+        parser.add_argument('--video', dest='make_video', action='store_true', help='Generate output videos (MP4s)')
+        parser.add_argument('--batch', dest='batch_mode', action='store_true', help="Don't open windows, instead save results to file")
+        parser.add_argument('-b', dest='batch_size', type=int, default=None, help='Minibatch size, leave empty for automatic detection')
+        parser.add_argument('-c', dest='components', type=int, default=80, help='Number of components to keep')
+        parser.add_argument('-n', type=int, default=300_000, help='Number of examples to use in decomposition')
+        parser.add_argument('--use_w', action='store_true', help='Use W latent space (StyleGAN(2))')
+        parser.add_argument('--sigma', type=float, default=2.0, help='Number of stdevs to walk in visualize.py')
+        parser.add_argument('--inputs', type=str, default=None, help='Path to directory with named components')
+        parser.add_argument('--seed', type=int, default=None, help='Seed used in decomposition')
+        args = parser.parse_args(args)
+
+        return self.from_dict(args.__dict__)

decomposition.py

@@ -0,0 +1,402 @@
+# Copyright 2020 Erik Härkönen. All rights reserved.
+# This file is licensed to you under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License. You may obtain a copy
+# of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software distributed under
+# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
+# OF ANY KIND, either express or implied. See the License for the specific language
+# governing permissions and limitations under the License.
+
+# Patch for broken CTRL+C handler
+# https://github.com/ContinuumIO/anaconda-issues/issues/905
+import os
+os.environ['FOR_DISABLE_CONSOLE_CTRL_HANDLER'] = '1'
+
+import numpy as np
+import os
+from pathlib import Path
+import re
+import sys
+import datetime
+import argparse
+import torch
+import json
+from types import SimpleNamespace
+import scipy
+from scipy.cluster.vq import kmeans
+from tqdm import trange
+from netdissect.nethook import InstrumentedModel
+from config import Config
+from estimators import get_estimator
+from models import get_instrumented_model
+
+SEED_SAMPLING = 1
+SEED_RANDOM_DIRS = 2
+SEED_LINREG = 3
+SEED_VISUALIZATION = 5
+
+B = 20
+n_clusters = 500
+
+def get_random_dirs(components, dimensions):
+    gen = np.random.RandomState(seed=SEED_RANDOM_DIRS)
+    dirs = gen.normal(size=(components, dimensions))
+    dirs /= np.sqrt(np.sum(dirs**2, axis=1, keepdims=True))
+    return dirs.astype(np.float32)
+
+# Compute maximum batch size for given VRAM and network
+def get_max_batch_size(inst, device, layer_name=None):
+    inst.remove_edits()
+
+    # Reset statistics
+    torch.cuda.reset_max_memory_cached(device)
+    torch.cuda.reset_max_memory_allocated(device)
+    total_mem = torch.cuda.get_device_properties(device).total_memory
+
+    B_max = 20
+
+    # Measure actual usage
+    for i in range(2, B_max, 2):
+        z = inst.model.sample_latent(n_samples=i)
+        if layer_name:
+            inst.model.partial_forward(z, layer_name)
+        else:
+            inst.model.forward(z)
+
+        maxmem = torch.cuda.max_memory_allocated(device)
+        del z
+
+        if maxmem > 0.5*total_mem:
+            print('Batch size {:d}: memory usage {:.0f}MB'.format(i, maxmem / 1e6))
+            return i
+
+    return B_max
+
+# Solve for directions in latent space that match PCs in activaiton space
+def linreg_lstsq(comp_np, mean_np, stdev_np, inst, config):
+    print('Performing least squares regression', flush=True)
+
+    torch.manual_seed(SEED_LINREG)
+    np.random.seed(SEED_LINREG)
+
+    comp = torch.from_numpy(comp_np).float().to(inst.model.device)
+    mean = torch.from_numpy(mean_np).float().to(inst.model.device)
+    stdev = torch.from_numpy(stdev_np).float().to(inst.model.device)
+
+    n_samp = max(10_000, config.n) // B * B # make divisible
+    n_comp = comp.shape[0]
+    latent_dims = inst.model.get_latent_dims()
+    
+    # We're looking for M s.t. M*P*G'(Z) = Z => M*A = Z
+    #   Z = batch of latent vectors (n_samples x latent_dims)
+    #   G'(Z) = batch of activations at intermediate layer
+    #   A = P*G'(Z) = projected activations (n_samples x pca_coords)
+    #   M = linear mapping (pca_coords x latent_dims)
+
+    # Minimization min_M ||MA - Z||_l2 rewritten as min_M.T ||A.T*M.T - Z.T||_l2
+    # to match format expected by pytorch.lstsq
+
+    # TODO: regression on pixel-space outputs? (using nonlinear optimizer)
+    # min_M lpips(G_full(MA), G_full(Z))
+
+    # Tensors to fill with data
+    # Dimensions other way around, so these are actually the transposes
+    A = np.zeros((n_samp, n_comp), dtype=np.float32)
+    Z = np.zeros((n_samp, latent_dims), dtype=np.float32)
+    
+    # Project tensor X onto PCs, return coordinates
+    def project(X, comp):
+        N = X.shape[0]
+        K = comp.shape[0]
+        coords = torch.bmm(comp.expand([N]+[-1]*comp.ndim), X.view(N, -1, 1))
+        return coords.reshape(N, K)
+
+    for i in trange(n_samp // B, desc='Collecting samples', ascii=True):
+        z = inst.model.sample_latent(B)
+        inst.model.partial_forward(z, config.layer)
+        act = inst.retained_features()[config.layer].reshape(B, -1)
+
+        # Project onto basis
+        act = act - mean
+        coords = project(act, comp)
+        coords_scaled = coords / stdev
+
+        A[i*B:(i+1)*B] = coords_scaled.detach().cpu().numpy()
+        Z[i*B:(i+1)*B] = z.detach().cpu().numpy().reshape(B, -1)
+
+    # Solve least squares fit
+
+    # gelsd = divide-and-conquer SVD; good default
+    # gelsy = complete orthogonal factorization; sometimes faster
+    # gelss = SVD; slow but less memory hungry
+    M_t = scipy.linalg.lstsq(A, Z, lapack_driver='gelsd')[0] # torch.lstsq(Z, A)[0][:n_comp, :]
+    
+    # Solution given by rows of M_t
+    Z_comp = M_t[:n_comp, :]
+    Z_mean = np.mean(Z, axis=0, keepdims=True)
+
+    return Z_comp, Z_mean
+
+def regression(comp, mean, stdev, inst, config):
+    # Sanity check: verify orthonormality
+    M = np.dot(comp, comp.T)
+    if not np.allclose(M, np.identity(M.shape[0])):
+        det = np.linalg.det(M)
+        print(f'WARNING: Computed basis is not orthonormal (determinant={det})')
+
+    return linreg_lstsq(comp, mean, stdev, inst, config)
+
+def compute(config, dump_name, instrumented_model):
+    global B
+
+    timestamp = lambda : datetime.datetime.now().strftime("%d.%m %H:%M")
+    print(f'[{timestamp()}] Computing', dump_name.name)
+
+    # Ensure reproducibility
+    torch.manual_seed(0) # also sets cuda seeds
+    np.random.seed(0)
+
+    # Speed up backend
+    torch.backends.cudnn.benchmark = True
+
+    has_gpu = torch.cuda.is_available()
+    device = torch.device('cuda' if has_gpu else 'cpu')
+    layer_key = config.layer
+
+    if instrumented_model is None:
+        inst = get_instrumented_model(config.model, config.output_class, layer_key, device)
+        model = inst.model
+    else:
+        print('Reusing InstrumentedModel instance')
+        inst = instrumented_model
+        model = inst.model
+        inst.remove_edits()
+        model.set_output_class(config.output_class)
+
+    # Regress back to w space
+    if config.use_w:
+        print('Using W latent space')
+        model.use_w()
+
+    inst.retain_layer(layer_key)
+    model.partial_forward(model.sample_latent(1), layer_key)
+    sample_shape = inst.retained_features()[layer_key].shape
+    sample_dims = np.prod(sample_shape)
+    print('Feature shape:', sample_shape)
+
+    input_shape = inst.model.get_latent_shape()
+    input_dims = inst.model.get_latent_dims()
+
+    config.components = min(config.components, sample_dims)
+    transformer = get_estimator(config.estimator, config.components, config.sparsity)
+
+    X = None
+    X_global_mean = None
+
+    # Figure out batch size if not provided
+    B = config.batch_size or get_max_batch_size(inst, device, layer_key)
+
+    # Divisible by B (ignored in output name)
+    N = config.n // B * B
+
+    # Compute maximum batch size based on RAM + pagefile budget
+    target_bytes = 20 * 1_000_000_000 # GB
+    feat_size_bytes = sample_dims * np.dtype('float64').itemsize
+    N_limit_RAM = np.floor_divide(target_bytes, feat_size_bytes)
+    if not transformer.batch_support and N > N_limit_RAM:
+        print('WARNING: estimator does not support batching, ' \
+            'given config will use {:.1f} GB memory.'.format(feat_size_bytes / 1_000_000_000 * N))
+
+    # 32-bit LAPACK gets very unhappy about huge matrices (in linalg.svd)
+    if config.estimator == 'ica':
+        lapack_max_N = np.floor_divide(np.iinfo(np.int32).max // 4, sample_dims) # 4x extra buffer
+        if N > lapack_max_N:
+            raise RuntimeError(f'Matrices too large for ICA, please use N <= {lapack_max_N}')
+
+    print('B={}, N={}, dims={}, N/dims={:.1f}'.format(B, N, sample_dims, N/sample_dims), flush=True)
+
+    # Must not depend on chosen batch size (reproducibility)
+    NB = max(B, max(2_000, 3*config.components)) # ipca: as large as possible!
+    
+    samples = None
+    if not transformer.batch_support:
+        samples = np.zeros((N + NB, sample_dims), dtype=np.float32)
+
+    torch.manual_seed(config.seed or SEED_SAMPLING)
+    np.random.seed(config.seed or SEED_SAMPLING)
+
+    # Use exactly the same latents regardless of batch size
+    # Store in main memory, since N might be huge (1M+)
+    # Run in batches, since sample_latent() might perform Z -> W mapping
+    n_lat = ((N + NB - 1) // B + 1) * B
+    latents = np.zeros((n_lat, *input_shape[1:]), dtype=np.float32)
+    with torch.no_grad():
+        for i in trange(n_lat // B, desc='Sampling latents'):
+            latents[i*B:(i+1)*B] = model.sample_latent(n_samples=B).cpu().numpy()
+
+    # Decomposition on non-Gaussian latent space
+    samples_are_latents = layer_key in ['g_mapping', 'style'] and inst.model.latent_space_name() == 'W'
+
+    canceled = False
+    try:
+        X = np.ones((NB, sample_dims), dtype=np.float32)
+        action = 'Fitting' if transformer.batch_support else 'Collecting'
+        for gi in trange(0, N, NB, desc=f'{action} batches (NB={NB})', ascii=True):
+            for mb in range(0, NB, B):
+                z = torch.from_numpy(latents[gi+mb:gi+mb+B]).to(device)
+                
+                if samples_are_latents:
+                    # Decomposition on latents directly (e.g. StyleGAN W)
+                    batch = z.reshape((B, -1))
+                else:
+                    # Decomposition on intermediate layer
+                    with torch.no_grad():
+                        model.partial_forward(z, layer_key)
+                    
+                    # Permuted to place PCA dimensions last
+                    batch = inst.retained_features()[layer_key].reshape((B, -1))
+
+                space_left = min(B, NB - mb)
+                X[mb:mb+space_left] = batch.cpu().numpy()[:space_left]
+
+            if transformer.batch_support:
+                if not transformer.fit_partial(X.reshape(-1, sample_dims)):
+                    break
+            else:
+                samples[gi:gi+NB, :] = X.copy()
+    except KeyboardInterrupt:
+        if not transformer.batch_support:
+            sys.exit(1) # no progress yet
+        
+        dump_name = dump_name.parent / dump_name.name.replace(f'n{N}', f'n{gi}')
+        print(f'Saving current state to "{dump_name.name}" before exiting')
+        canceled = True
+        
+    if not transformer.batch_support:
+        X = samples # Use all samples
+        X_global_mean = X.mean(axis=0, keepdims=True, dtype=np.float32) # TODO: activations surely multi-modal...!
+        X -= X_global_mean
+        
+        print(f'[{timestamp()}] Fitting whole batch')
+        t_start_fit = datetime.datetime.now()
+
+        transformer.fit(X)
+        
+        print(f'[{timestamp()}] Done in {datetime.datetime.now() - t_start_fit}')
+        assert np.all(transformer.transformer.mean_ < 1e-3), 'Mean of normalized data should be zero'
+    else:
+        X_global_mean = transformer.transformer.mean_.reshape((1, sample_dims))
+        X = X.reshape(-1, sample_dims)
+        X -= X_global_mean
+
+    X_comp, X_stdev, X_var_ratio = transformer.get_components()
+    
+    assert X_comp.shape[1] == sample_dims \
+        and X_comp.shape[0] == config.components \
+        and X_global_mean.shape[1] == sample_dims \
+        and X_stdev.shape[0] == config.components, 'Invalid shape'
+
+    # 'Activations' are really latents in a secondary latent space
+    if samples_are_latents:
+        Z_comp = X_comp
+        Z_global_mean = X_global_mean
+    else:
+        Z_comp, Z_global_mean = regression(X_comp, X_global_mean, X_stdev, inst, config)
+
+    # Normalize
+    Z_comp /= np.linalg.norm(Z_comp, axis=-1, keepdims=True)
+
+    # Random projections
+    # We expect these to explain much less of the variance
+    random_dirs = get_random_dirs(config.components, np.prod(sample_shape))
+    n_rand_samples = min(5000, X.shape[0])
+    X_view = X[:n_rand_samples, :].T
+    assert np.shares_memory(X_view, X), "Error: slice produced copy"
+    X_stdev_random = np.dot(random_dirs, X_view).std(axis=1)
+
+    # Inflate back to proper shapes (for easier broadcasting)
+    X_comp = X_comp.reshape(-1, *sample_shape)
+    X_global_mean = X_global_mean.reshape(sample_shape)
+    Z_comp = Z_comp.reshape(-1, *input_shape)
+    Z_global_mean = Z_global_mean.reshape(input_shape)
+
+    # Compute stdev in latent space if non-Gaussian
+    lat_stdev = np.ones_like(X_stdev)
+    if config.use_w:
+        samples = model.sample_latent(5000).reshape(5000, input_dims).detach().cpu().numpy()
+        coords = np.dot(Z_comp.reshape(-1, input_dims), samples.T)
+        lat_stdev = coords.std(axis=1)
+
+    os.makedirs(dump_name.parent, exist_ok=True)
+    np.savez_compressed(dump_name, **{
+        'act_comp': X_comp.astype(np.float32),
+        'act_mean': X_global_mean.astype(np.float32),
+        'act_stdev': X_stdev.astype(np.float32),
+        'lat_comp': Z_comp.astype(np.float32),
+        'lat_mean': Z_global_mean.astype(np.float32),
+        'lat_stdev': lat_stdev.astype(np.float32),
+        'var_ratio': X_var_ratio.astype(np.float32),
+        'random_stdevs': X_stdev_random.astype(np.float32),
+    })
+
+    if canceled:
+        sys.exit(1)
+
+    # Don't shutdown if passed as param
+    if instrumented_model is None:
+        inst.close()
+        del inst
+        del model
+
+    del X
+    del X_comp
+    del random_dirs
+    del batch
+    del samples
+    del latents
+    torch.cuda.empty_cache()
+
+# Return cached results or commpute if needed
+# Pass existing InstrumentedModel instance to reuse it
+def get_or_compute(config, model=None, submit_config=None, force_recompute=False):
+    if submit_config is None:
+        wrkdir = str(Path(__file__).parent.resolve())
+        submit_config = SimpleNamespace(run_dir_root = wrkdir, run_dir = wrkdir)
+    
+    # Called directly by run.py
+    return _compute(submit_config, config, model, force_recompute)
+
+def _compute(submit_config, config, model=None, force_recompute=False):
+    basedir = Path(submit_config.run_dir)
+    outdir = basedir / 'out'
+    
+    if config.n is None:
+        raise RuntimeError('Must specify number of samples with -n=XXX')
+
+    if model and not isinstance(model, InstrumentedModel):
+        raise RuntimeError('Passed model has to be wrapped in "InstrumentedModel"')
+    
+    if config.use_w and not 'StyleGAN' in config.model:
+        raise RuntimeError(f'Cannot change latent space of non-StyleGAN model {config.model}')
+
+    transformer = get_estimator(config.estimator, config.components, config.sparsity)
+    dump_name = "{}-{}_{}_{}_n{}{}{}.npz".format(
+        config.model.lower(),
+        config.output_class.replace(' ', '_'),
+        config.layer.lower(),
+        transformer.get_param_str(),
+        config.n,
+        '_w' if config.use_w else '',
+        f'_seed{config.seed}' if config.seed else ''
+    )
+
+    dump_path = basedir / 'cache' / 'components' / dump_name
+
+    if not dump_path.is_file() or force_recompute:
+        print('Not cached')
+        t_start = datetime.datetime.now()
+        compute(config, dump_path, model)
+        print('Total time:', datetime.datetime.now() - t_start)
+    
+    return dump_path

deps/windows/PyOpenGL-3.1.4-cp37-cp37m-win_amd64.whl

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

deps/windows/glumpy-1.1.0-cp37-cp37m-win_amd64.whl

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

deps/windows/triangle-20190115.3-cp37-cp37m-win_amd64.whl

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

environment.yml

@@ -0,0 +1,25 @@
+name: ganspace
+channels:
+  - defaults
+  - conda-forge
+  - pytorch
+dependencies:
+  - python=3.7
+  - pytorch::pytorch=1.3
+  - pytorch::torchvision
+  - cudatoolkit=10.1
+  - pillow=6.2
+  - ffmpeg
+  - tqdm
+  - scipy
+  - scikit-learn
+  - scikit-image
+  - boto3
+  - requests
+  - nltk
+  - pip
+  - pip:
+    - fbpca
+    - pyopengltk
+
+# conda env update -f environment.yml --prune

estimators.py

@@ -0,0 +1,218 @@
+# Copyright 2020 Erik Härkönen. All rights reserved.
+# This file is licensed to you under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License. You may obtain a copy
+# of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software distributed under
+# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
+# OF ANY KIND, either express or implied. See the License for the specific language
+# governing permissions and limitations under the License.
+
+from sklearn.decomposition import FastICA, PCA, IncrementalPCA, MiniBatchSparsePCA, SparsePCA, KernelPCA
+import fbpca
+import numpy as np
+import itertools
+from types import SimpleNamespace
+
+# ICA
+class ICAEstimator():
+    def __init__(self, n_components):
+        self.n_components = n_components
+        self.maxiter = 10000
+        self.whiten = True # ICA: whitening is essential, should not be skipped
+        self.transformer = FastICA(n_components, random_state=0, whiten=self.whiten, max_iter=self.maxiter)
+        self.batch_support = False
+        self.stdev = np.zeros((n_components,))
+        self.total_var = 0.0
+
+    def get_param_str(self):
+        return "ica_c{}{}".format(self.n_components, '_w' if self.whiten else '')
+    
+    def fit(self, X):
+        self.transformer.fit(X)
+        if self.transformer.n_iter_ >= self.maxiter:
+            raise RuntimeError(f'FastICA did not converge (N={X.shape[0]}, it={self.maxiter})')
+
+        # Normalize components
+        self.transformer.components_ /= np.sqrt(np.sum(self.transformer.components_**2, axis=-1, keepdims=True))
+
+        # Save variance for later
+        self.total_var = X.var(axis=0).sum()
+
+        # Compute projected standard deviations
+        self.stdev = np.dot(self.transformer.components_, X.T).std(axis=1)
+
+        # Sort components based on explained variance
+        idx = np.argsort(self.stdev)[::-1]
+        self.stdev = self.stdev[idx]
+        self.transformer.components_[:] = self.transformer.components_[idx]
+
+    def get_components(self):
+        var_ratio = self.stdev**2 / self.total_var
+        return self.transformer.components_, self.stdev, var_ratio # ICA outputs are not normalized
+
+# Incremental PCA
+class IPCAEstimator():
+    def __init__(self, n_components):
+        self.n_components = n_components
+        self.whiten = False
+        self.transformer = IncrementalPCA(n_components, whiten=self.whiten, batch_size=max(100, 2*n_components))
+        self.batch_support = True
+
+    def get_param_str(self):
+        return "ipca_c{}{}".format(self.n_components, '_w' if self.whiten else '')
+
+    def fit(self, X):
+        self.transformer.fit(X)
+
+    def fit_partial(self, X):
+        try:
+            self.transformer.partial_fit(X)
+            self.transformer.n_samples_seen_ = \
+                self.transformer.n_samples_seen_.astype(np.int64) # avoid overflow
+            return True
+        except ValueError as e:
+            print(f'\nIPCA error:', e)
+            return False
+
+    def get_components(self):
+        stdev = np.sqrt(self.transformer.explained_variance_) # already sorted
+        var_ratio = self.transformer.explained_variance_ratio_
+        return self.transformer.components_, stdev, var_ratio # PCA outputs are normalized
+
+# Standard PCA
+class PCAEstimator():
+    def __init__(self, n_components):
+        self.n_components = n_components
+        self.solver = 'full'
+        self.transformer = PCA(n_components, svd_solver=self.solver)
+        self.batch_support = False
+
+    def get_param_str(self):
+        return f"pca-{self.solver}_c{self.n_components}"
+
+    def fit(self, X):
+        self.transformer.fit(X)
+
+        # Save variance for later
+        self.total_var = X.var(axis=0).sum()
+
+        # Compute projected standard deviations
+        self.stdev = np.dot(self.transformer.components_, X.T).std(axis=1)
+
+        # Sort components based on explained variance
+        idx = np.argsort(self.stdev)[::-1]
+        self.stdev = self.stdev[idx]
+        self.transformer.components_[:] = self.transformer.components_[idx]
+
+        # Check orthogonality
+        dotps = [np.dot(*self.transformer.components_[[i, j]])
+            for (i, j) in itertools.combinations(range(self.n_components), 2)]
+        if not np.allclose(dotps, 0, atol=1e-4):
+            print('IPCA components not orghogonal, max dot', np.abs(dotps).max())
+
+        self.transformer.mean_ = X.mean(axis=0, keepdims=True)
+
+    def get_components(self):
+        var_ratio = self.stdev**2 / self.total_var
+        return self.transformer.components_, self.stdev, var_ratio
+
+# Facebook's PCA
+# Good default choice: very fast and accurate.
+# Very high sample counts won't fit into RAM,
+# in which case IncrementalPCA must be used.
+class FacebookPCAEstimator():
+    def __init__(self, n_components):
+        self.n_components = n_components
+        self.transformer = SimpleNamespace()
+        self.batch_support = False
+        self.n_iter = 2
+        self.l = 2*self.n_components
+
+    def get_param_str(self):
+        return "fbpca_c{}_it{}_l{}".format(self.n_components, self.n_iter, self.l)
+
+    def fit(self, X):
+        U, s, Va = fbpca.pca(X, k=self.n_components, n_iter=self.n_iter, raw=True, l=self.l)
+        self.transformer.components_ = Va
+        
+        # Save variance for later
+        self.total_var = X.var(axis=0).sum()
+
+        # Compute projected standard deviations
+        self.stdev = np.dot(self.transformer.components_, X.T).std(axis=1)
+
+        # Sort components based on explained variance
+        idx = np.argsort(self.stdev)[::-1]
+        self.stdev = self.stdev[idx]
+        self.transformer.components_[:] = self.transformer.components_[idx]
+
+        # Check orthogonality
+        dotps = [np.dot(*self.transformer.components_[[i, j]])
+            for (i, j) in itertools.combinations(range(self.n_components), 2)]
+        if not np.allclose(dotps, 0, atol=1e-4):
+            print('FBPCA components not orghogonal, max dot', np.abs(dotps).max())
+
+        self.transformer.mean_ = X.mean(axis=0, keepdims=True)
+        
+    def get_components(self):
+        var_ratio = self.stdev**2 / self.total_var
+        return self.transformer.components_, self.stdev, var_ratio
+
+# Sparse PCA
+# The algorithm is online along the features direction, not the samples direction
+#   => no partial_fit
+class SPCAEstimator():
+    def __init__(self, n_components, alpha=10.0):
+        self.n_components = n_components
+        self.whiten = False
+        self.alpha = alpha  # higher alpha => sparser components
+        #self.transformer = MiniBatchSparsePCA(n_components, alpha=alpha, n_iter=100,
+        #    batch_size=max(20, n_components//5), random_state=0, normalize_components=True)
+        self.transformer = SparsePCA(n_components, alpha=alpha, ridge_alpha=0.01,
+            max_iter=100, random_state=0, n_jobs=-1, normalize_components=True) # TODO: warm start using PCA result?
+        self.batch_support = False # maybe through memmap and HDD-stored tensor
+        self.stdev = np.zeros((n_components,))
+        self.total_var = 0.0
+
+    def get_param_str(self):
+        return "spca_c{}_a{}{}".format(self.n_components, self.alpha, '_w' if self.whiten else '')
+        
+    def fit(self, X):
+        self.transformer.fit(X)
+
+        # Save variance for later
+        self.total_var = X.var(axis=0).sum()
+
+        # Compute projected standard deviations
+        # NB: cannot simply project with dot product!
+        self.stdev = self.transformer.transform(X).std(axis=0) # X = (n_samples, n_features)
+
+        # Sort components based on explained variance
+        idx = np.argsort(self.stdev)[::-1]
+        self.stdev = self.stdev[idx]
+        self.transformer.components_[:] = self.transformer.components_[idx]
+
+        # Check orthogonality
+        dotps = [np.dot(*self.transformer.components_[[i, j]])
+            for (i, j) in itertools.combinations(range(self.n_components), 2)]
+        if not np.allclose(dotps, 0, atol=1e-4):
+            print('SPCA components not orghogonal, max dot', np.abs(dotps).max())
+
+    def get_components(self):
+        var_ratio = self.stdev**2 / self.total_var
+        return self.transformer.components_, self.stdev, var_ratio # SPCA outputs are normalized
+
+def get_estimator(name, n_components, alpha):
+    if name == 'pca':
+        return PCAEstimator(n_components)
+    if name == 'ipca':
+        return IPCAEstimator(n_components)
+    elif name == 'fbpca':
+        return FacebookPCAEstimator(n_components)
+    elif name == 'ica':
+        return ICAEstimator(n_components)
+    elif name == 'spca':
+        return SPCAEstimator(n_components, alpha)
+    else:
+        raise RuntimeError('Unknown estimator')

interactive.py

@@ -0,0 +1,655 @@
+# Copyright 2020 Erik Härkönen. All rights reserved.
+# This file is licensed to you under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License. You may obtain a copy
+# of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software distributed under
+# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
+# OF ANY KIND, either express or implied. See the License for the specific language
+# governing permissions and limitations under the License.
+
+# An interactive glumpy (OpenGL) + tkinter viewer for interacting with principal components.
+# Requires OpenGL and CUDA support for rendering.
+
+import torch
+import numpy as np
+import tkinter as tk
+from tkinter import ttk
+from types import SimpleNamespace
+import matplotlib.pyplot as plt
+from pathlib import Path
+from os import makedirs
+from models import get_instrumented_model
+from config import Config
+from decomposition import get_or_compute
+from torch.nn.functional import interpolate
+from TkTorchWindow import TorchImageView
+from functools import partial
+from platform import system
+from PIL import Image
+from utils import pad_frames, prettify_name
+import pickle
+
+# For platform specific UI tweaks
+is_windows = 'Windows' in system()
+is_linux = 'Linux' in system()
+is_mac = 'Darwin' in system()
+
+# Read input parameters
+args = Config().from_args()
+
+# Don't bother without GPU
+assert torch.cuda.is_available(), 'Interactive mode requires CUDA'
+
+# Use syntax from paper
+def get_edit_name(idx, s, e, name=None):
+    return 'E({comp}, {edit_range}){edit_name}'.format(
+        comp = idx,
+        edit_range = f'{s}-{e}' if e > s else s,
+        edit_name = f': {name}' if name else ''
+    )
+
+# Load or compute PCA basis vectors
+def load_components(class_name, inst):
+    global components, state, use_named_latents
+
+    config = args.from_dict({ 'output_class': class_name })
+    dump_name = get_or_compute(config, inst)
+    data = np.load(dump_name, allow_pickle=False)
+    X_comp = data['act_comp']
+    X_mean = data['act_mean']
+    X_stdev = data['act_stdev']
+    Z_comp = data['lat_comp']
+    Z_mean = data['lat_mean']
+    Z_stdev = data['lat_stdev']
+    random_stdev_act = np.mean(data['random_stdevs'])
+    n_comp = X_comp.shape[0]
+    data.close()
+
+    # Transfer to GPU
+    components = SimpleNamespace(
+        X_comp = torch.from_numpy(X_comp).cuda().float(),
+        X_mean = torch.from_numpy(X_mean).cuda().float(),
+        X_stdev = torch.from_numpy(X_stdev).cuda().float(),
+        Z_comp = torch.from_numpy(Z_comp).cuda().float(),
+        Z_stdev = torch.from_numpy(Z_stdev).cuda().float(),
+        Z_mean = torch.from_numpy(Z_mean).cuda().float(),
+        names = [f'Component {i}' for i in range(n_comp)],
+        latent_types = [model.latent_space_name()]*n_comp,
+        ranges = [(0, model.get_max_latents())]*n_comp,
+    )
+    
+    state.component_class = class_name # invalidates cache
+    use_named_latents = False
+    print('Loaded components for', class_name, 'from', dump_name)
+
+# Load previously exported named components from
+# directory specified with '--inputs=path/to/comp'
+def load_named_components(path, class_name):
+    global components, state, use_named_latents
+
+    import glob
+    matches = glob.glob(f'{path}/*.pkl')
+
+    selected = []
+    for dump_path in matches:
+        with open(dump_path, 'rb') as f:
+            data = pickle.load(f)
+            if data['model_name'] != model_name or data['output_class'] != class_name:
+                continue
+
+            if data['latent_space'] != model.latent_space_name():
+                print('Skipping', dump_path, '(wrong latent space)')
+                continue
+            
+            selected.append(data)
+            print('Using', dump_path)
+
+    if len(selected) == 0:
+        raise RuntimeError('No valid components in given path.')
+
+    comp_dict = { k : [] for k in ['X_comp', 'Z_comp', 'X_stdev', 'Z_stdev', 'names', 'types', 'layer_names', 'ranges', 'latent_types'] }
+    components = SimpleNamespace(**comp_dict)
+
+    for d in selected:
+        s = d['edit_start']
+        e = d['edit_end']
+        title = get_edit_name(d['component_index'], s, e - 1, d['name']) # show inclusive
+        components.X_comp.append(torch.from_numpy(d['act_comp']).cuda())
+        components.Z_comp.append(torch.from_numpy(d['lat_comp']).cuda())
+        components.X_stdev.append(d['act_stdev'])
+        components.Z_stdev.append(d['lat_stdev'])
+        components.names.append(title)
+        components.types.append(d['edit_type'])
+        components.layer_names.append(d['decomposition']['layer']) # only for act
+        components.ranges.append((s, e))
+        components.latent_types.append(d['latent_space']) # W or Z
+    
+    use_named_latents = True
+    print('Loaded named components')
+
+def setup_model():
+    global model, inst, layer_name, model_name, feat_shape, args, class_name
+
+    model_name = args.model
+    layer_name = args.layer
+    class_name = args.output_class
+
+    # Speed up pytorch
+    torch.autograd.set_grad_enabled(False)
+    torch.backends.cudnn.benchmark = True
+
+    # Load model
+    inst = get_instrumented_model(model_name, class_name, layer_name, torch.device('cuda'), use_w=args.use_w)
+    model = inst.model
+
+    feat_shape = inst.feature_shape[layer_name]
+    sample_dims = np.prod(feat_shape)
+
+    # Initialize 
+    if args.inputs:
+        load_named_components(args.inputs, class_name)
+    else:
+        load_components(class_name, inst)
+
+# Project tensor 'X' onto orthonormal basis 'comp', return coordinates
+def project_ortho(X, comp):
+    N = comp.shape[0]
+    coords = (comp.reshape(N, -1) * X.reshape(-1)).sum(dim=1)
+    return coords.reshape([N]+[1]*X.ndim)
+
+def zero_sliders():
+    for v in ui_state.sliders:
+        v.set(0.0)
+
+def reset_sliders(zero_on_failure=True):
+    global ui_state
+
+    mode = ui_state.mode.get()
+
+    # Not orthogonal: need to solve least-norm problem
+    # Not batch size 1: one set of sliders not enough
+    # Not principal components: unsupported format
+    is_ortho = not (mode == 'latent' and model.latent_space_name() == 'Z')
+    is_single = state.z.shape[0] == 1
+    is_pcs = not use_named_latents
+
+    state.lat_slider_offset = 0
+    state.act_slider_offset = 0
+
+    enabled = False
+    if not (enabled and is_ortho and is_single and is_pcs):
+        if zero_on_failure:
+            zero_sliders()
+        return
+
+    if  mode == 'activation':
+        val = state.base_act
+        mean = components.X_mean
+        comp = components.X_comp
+        stdev = components.X_stdev
+    else:
+        val = state.z
+        mean = components.Z_mean
+        comp = components.Z_comp
+        stdev = components.Z_stdev
+
+    n_sliders = len(ui_state.sliders)
+    coords = project_ortho(val - mean, comp)
+    offset = torch.sum(coords[:n_sliders] * comp[:n_sliders], dim=0)
+    scaled_coords = (coords.view(-1) / stdev).detach().cpu().numpy()
+    
+    # Part representable by sliders
+    if mode == 'activation':
+        state.act_slider_offset = offset
+    else:
+        state.lat_slider_offset = offset
+
+    for i in range(n_sliders):
+        ui_state.sliders[i].set(round(scaled_coords[i], ndigits=1))
+
+def setup_ui():
+    global root, toolbar, ui_state, app, canvas
+
+    root = tk.Tk()
+    scale = 1.0
+    app = TorchImageView(root, width=int(scale*1024), height=int(scale*1024), show_fps=False)
+    app.pack(fill=tk.BOTH, expand=tk.YES)
+    root.protocol("WM_DELETE_WINDOW", shutdown)
+    root.title('GANspace')
+
+    toolbar = tk.Toplevel(root)
+    toolbar.protocol("WM_DELETE_WINDOW", shutdown)
+    toolbar.geometry("215x800+0+0")
+    toolbar.title('')
+
+    N_COMPONENTS = min(70, len(components.names))
+    ui_state = SimpleNamespace(
+        sliders = [tk.DoubleVar(value=0.0) for _ in range(N_COMPONENTS)],
+        scales = [],
+        truncation = tk.DoubleVar(value=0.9),
+        outclass = tk.StringVar(value=class_name),
+        random_seed = tk.StringVar(value='0'),
+        mode = tk.StringVar(value='latent'),
+        batch_size = tk.IntVar(value=1), # how many images to show in window
+        edit_layer_start = tk.IntVar(value=0),
+        edit_layer_end = tk.IntVar(value=model.get_max_latents() - 1),
+        slider_max_val = 10.0
+    )
+
+    # Z vs activation mode button
+    #tk.Radiobutton(toolbar, text=f"Latent ({model.latent_space_name()})", variable=ui_state.mode, command=reset_sliders, value='latent').pack(fill="x")
+    #tk.Radiobutton(toolbar, text="Activation", variable=ui_state.mode, command=reset_sliders, value='activation').pack(fill="x")
+
+    # Choose range where latents are modified
+    def set_min(val):
+        ui_state.edit_layer_start.set(min(int(val), ui_state.edit_layer_end.get()))
+    def set_max(val):
+        ui_state.edit_layer_end.set(max(int(val), ui_state.edit_layer_start.get()))
+    max_latent_idx = model.get_max_latents() - 1
+    
+    if not use_named_latents:
+        slider_min = tk.Scale(toolbar, command=set_min, variable=ui_state.edit_layer_start,
+            label='Layer start', from_=0, to=max_latent_idx, orient=tk.HORIZONTAL).pack(fill="x")
+        slider_max = tk.Scale(toolbar, command=set_max, variable=ui_state.edit_layer_end,
+            label='Layer end', from_=0, to=max_latent_idx, orient=tk.HORIZONTAL).pack(fill="x")
+
+    # Scrollable list of components
+    outer_frame = tk.Frame(toolbar, borderwidth=2, relief=tk.SUNKEN)
+    canvas = tk.Canvas(outer_frame, highlightthickness=0, borderwidth=0)
+    frame = tk.Frame(canvas)
+    vsb = tk.Scrollbar(outer_frame, orient="vertical", command=canvas.yview)
+    canvas.configure(yscrollcommand=vsb.set)
+
+    vsb.pack(side="right", fill="y")
+    canvas.pack(side="left", fill="both", expand=True)
+    canvas.create_window((4,4), window=frame, anchor="nw")
+
+    def onCanvasConfigure(event):
+        canvas.itemconfigure("all", width=event.width)
+        canvas.configure(scrollregion=canvas.bbox("all"))
+    canvas.bind("<Configure>", onCanvasConfigure)
+
+    def on_scroll(event):
+        delta = 1 if (event.num == 5 or event.delta < 0) else -1
+        canvas.yview_scroll(delta, "units")
+
+    canvas.bind_all("<Button-4>", on_scroll)
+    canvas.bind_all("<Button-5>", on_scroll)
+    canvas.bind_all("<MouseWheel>", on_scroll)
+    canvas.bind_all("<Key>", lambda event : handle_keypress(event.keysym_num))
+
+    # Sliders and buttons
+    for i in range(N_COMPONENTS):
+        inner = tk.Frame(frame, borderwidth=1, background="#aaaaaa")
+        scale = tk.Scale(inner, variable=ui_state.sliders[i], from_=-ui_state.slider_max_val,
+            to=ui_state.slider_max_val, resolution=0.1, orient=tk.HORIZONTAL, label=components.names[i])
+        scale.pack(fill=tk.X, side=tk.LEFT, expand=True)
+        ui_state.scales.append(scale) # for changing label later
+        if not use_named_latents:
+            tk.Button(inner, text=f"Save", command=partial(export_direction, i, inner)).pack(fill=tk.Y, side=tk.RIGHT)
+        inner.pack(fill=tk.X)
+
+    outer_frame.pack(fill="both", expand=True, pady=0)
+
+    tk.Button(toolbar, text="Reset", command=reset_sliders).pack(anchor=tk.CENTER, fill=tk.X, padx=4, pady=4)
+
+    tk.Scale(toolbar, variable=ui_state.truncation, from_=0.01, to=1.0,
+        resolution=0.01, orient=tk.HORIZONTAL, label='Truncation').pack(fill="x")
+
+    tk.Scale(toolbar, variable=ui_state.batch_size, from_=1, to=9,
+        resolution=1, orient=tk.HORIZONTAL, label='Batch size').pack(fill="x")
+    
+    # Output class
+    frame = tk.Frame(toolbar)
+    tk.Label(frame, text="Class name").pack(fill="x", side="left")
+    tk.Entry(frame, textvariable=ui_state.outclass).pack(fill="x", side="right", expand=True, padx=5)
+    frame.pack(fill=tk.X, pady=3)
+
+    # Random seed
+    def update_seed():
+        seed_str = ui_state.random_seed.get()
+        if seed_str.isdigit():
+            resample_latent(int(seed_str))
+    frame = tk.Frame(toolbar)
+    tk.Label(frame, text="Seed").pack(fill="x", side="left")
+    tk.Entry(frame, textvariable=ui_state.random_seed, width=12).pack(fill="x", side="left", expand=True, padx=2)
+    tk.Button(frame, text="Update", command=update_seed).pack(fill="y", side="right", padx=3)
+    frame.pack(fill=tk.X, pady=3)
+    
+    # Get new latent or new components
+    tk.Button(toolbar, text="Resample latent", command=partial(resample_latent, None, False)).pack(anchor=tk.CENTER, fill=tk.X, padx=4, pady=4)
+    #tk.Button(toolbar, text="Recompute", command=recompute_components).pack(anchor=tk.CENTER, fill=tk.X)
+
+# App state
+state = SimpleNamespace(
+    z=None,                  # current latent(s)
+    lat_slider_offset = 0,   # part of lat that is explained by sliders
+    act_slider_offset = 0,   # part of act that is explained by sliders
+    component_class=None,    # name of current PCs' image class
+    seed=0,                  # Latent z_i generated by seed+i
+    base_act = None,         # activation of considered layer given z
+)
+
+def resample_latent(seed=None, only_style=False):
+    class_name = ui_state.outclass.get()
+    if class_name.isnumeric():
+        class_name = int(class_name)
+    
+    if hasattr(model, 'is_valid_class'):
+        if not model.is_valid_class(class_name):
+            return
+
+    model.set_output_class(class_name)
+    
+    B = ui_state.batch_size.get()
+    state.seed = np.random.randint(np.iinfo(np.int32).max - B) if seed is None else seed
+    ui_state.random_seed.set(str(state.seed))
+    
+    # Use consecutive seeds along batch dimension (for easier reproducibility)
+    trunc = ui_state.truncation.get()
+    latents = [model.sample_latent(1, seed=state.seed + i, truncation=trunc) for i in range(B)]
+
+    state.z = torch.cat(latents).clone().detach() # make leaf node
+    assert state.z.is_leaf, 'Latent is not leaf node!'
+    
+    if hasattr(model, 'truncation'):
+        model.truncation = ui_state.truncation.get()
+    print(f'Seeds: {state.seed} -> {state.seed + B - 1}' if B > 1 else f'Seed: {state.seed}')
+
+    torch.manual_seed(state.seed)
+    model.partial_forward(state.z, layer_name)
+    state.base_act = inst.retained_features()[layer_name]
+    
+    reset_sliders(zero_on_failure=False)
+
+    # Remove focus from text entry
+    canvas.focus_set()
+
+# Used to recompute after changing class of conditional model
+def recompute_components():
+    class_name = ui_state.outclass.get()
+    if class_name.isnumeric():
+        class_name = int(class_name)
+    
+    if hasattr(model, 'is_valid_class'):
+        if not model.is_valid_class(class_name):
+            return
+
+    if hasattr(model, 'set_output_class'):
+        model.set_output_class(class_name)
+    
+    load_components(class_name, inst)
+
+# Used to detect parameter changes for lazy recomputation
+class ParamCache():
+    def update(self, **kwargs):
+        dirty = False
+        for argname, val in kwargs.items():
+            # Check pointer, then value
+            current = getattr(self, argname, 0)
+            if current is not val and pickle.dumps(current) != pickle.dumps(val):
+                setattr(self, argname, val)
+                dirty = True
+        return dirty
+
+cache = ParamCache()
+
+def l2norm(t):
+    return torch.norm(t.view(t.shape[0], -1), p=2, dim=1, keepdim=True)
+
+def apply_edit(z0, delta):
+    return z0 + delta
+
+def reposition_toolbar():
+    size, X, Y = root.winfo_geometry().split('+')
+    W, H = size.split('x')
+    toolbar_W = toolbar.winfo_geometry().split('x')[0]
+    offset_y = -30 if is_linux else 0 # window title bar
+    toolbar.geometry(f'{toolbar_W}x{H}+{int(X)-int(toolbar_W)}+{int(Y)+offset_y}')
+    toolbar.update()
+
+def on_draw():
+    global img
+
+    n_comp = len(ui_state.sliders)
+    slider_vals = np.array([s.get() for s in ui_state.sliders], dtype=np.float32)
+    
+    # Run model sparingly
+    mode = ui_state.mode.get()
+    latent_start = ui_state.edit_layer_start.get()
+    latent_end = ui_state.edit_layer_end.get() + 1 # save as exclusive, show as inclusive
+
+    if cache.update(coords=slider_vals, comp=state.component_class, mode=mode, z=state.z, s=latent_start, e=latent_end):
+        with torch.no_grad():
+            z_base = state.z - state.lat_slider_offset
+            z_deltas = [0.0]*model.get_max_latents()
+            z_delta_global = 0.0
+            
+            n_comp = slider_vals.size
+            act_deltas = {}
+            
+            if torch.is_tensor(state.act_slider_offset):
+                act_deltas[layer_name] = -state.act_slider_offset
+
+            for space in components.latent_types:
+                assert space == model.latent_space_name(), \
+                    'Cannot mix latent spaces (for now)'
+
+            for c in range(n_comp):
+                coord = slider_vals[c]
+                if coord == 0:
+                    continue
+
+                edit_mode = components.types[c] if use_named_latents else mode
+                
+                # Activation offset
+                if edit_mode in ['activation', 'both']:
+                    delta = components.X_comp[c] * components.X_stdev[c] * coord
+                    name = components.layer_names[c] if use_named_latents else layer_name
+                    act_deltas[name] = act_deltas.get(name, 0.0) + delta
+
+                # Latent offset
+                if edit_mode in ['latent', 'both']:
+                    delta = components.Z_comp[c] * components.Z_stdev[c] * coord
+                    edit_range = components.ranges[c] if use_named_latents else (latent_start, latent_end)
+                    full_range = (edit_range == (0, model.get_max_latents()))
+                    
+                    # Single or multiple offsets?
+                    if full_range:
+                        z_delta_global = z_delta_global + delta
+                    else:
+                        for l in range(*edit_range):
+                            z_deltas[l] = z_deltas[l] + delta
+
+            # Apply activation deltas
+            inst.remove_edits()
+            for layer, delta in act_deltas.items():
+                inst.edit_layer(layer, offset=delta)
+            
+            # Evaluate
+            has_offsets = any(torch.is_tensor(t) for t in z_deltas)
+            z_final = apply_edit(z_base, z_delta_global)
+            if has_offsets:
+                z_final = [apply_edit(z_final, d) for d in z_deltas]
+            img = model.forward(z_final).clamp(0.0, 1.0)
+
+    app.draw(img)
+
+# Save necessary data to disk for later loading
+def export_direction(idx, button_frame):
+    name = tk.StringVar(value='')
+    num_strips = tk.IntVar(value=0)
+    strip_width = tk.IntVar(value=5)
+
+    slider_values = np.array([s.get() for s in ui_state.sliders])
+    slider_value = slider_values[idx]
+    if (slider_values != 0).sum() > 1:
+        print('Please modify only one slider')
+        return
+    elif slider_value == 0:
+        print('Modify selected slider to set usable range (currently 0)')
+        return
+    
+    popup = tk.Toplevel(root)
+    popup.geometry("200x200+0+0")
+    tk.Label(popup, text="Edit name").pack()
+    tk.Entry(popup, textvariable=name).pack(pady=5)
+    # tk.Scale(popup, from_=0, to=30, variable=num_strips,
+    #    resolution=1, orient=tk.HORIZONTAL, length=200, label='Image strips to export').pack()
+    # tk.Scale(popup, from_=3, to=15, variable=strip_width,
+    #    resolution=1, orient=tk.HORIZONTAL, length=200, label='Image strip width').pack()
+    tk.Button(popup, text='OK', command=popup.quit).pack()
+    
+    canceled = False
+    def on_close():
+        nonlocal canceled
+        canceled = True
+        popup.quit()
+
+    popup.protocol("WM_DELETE_WINDOW", on_close)
+    x = button_frame.winfo_rootx()
+    y = button_frame.winfo_rooty()
+    w = int(button_frame.winfo_geometry().split('x')[0])
+    popup.geometry('%dx%d+%d+%d' % (180, 90, x + w, y))
+    popup.mainloop()
+    popup.destroy()
+
+    # Update slider name
+    label = get_edit_name(idx, ui_state.edit_layer_start.get(),
+        ui_state.edit_layer_end.get(), name.get())
+    ui_state.scales[idx].config(label=label)
+
+    if canceled:
+        return
+
+    params = {
+        'name': name.get(),
+        'sigma_range': slider_value,
+        'component_index': idx,
+        'act_comp': components.X_comp[idx].detach().cpu().numpy(),
+        'lat_comp': components.Z_comp[idx].detach().cpu().numpy(), # either Z or W
+        'latent_space': model.latent_space_name(),
+        'act_stdev': components.X_stdev[idx].item(),
+        'lat_stdev': components.Z_stdev[idx].item(),
+        'model_name': model_name,
+        'output_class': ui_state.outclass.get(), # applied onto
+        'decomposition': {
+            'name': args.estimator,
+            'components': args.components,
+            'samples': args.n,
+            'layer': args.layer,
+            'class_name': state.component_class # computed from
+        },
+        'edit_type': ui_state.mode.get(),
+        'truncation': ui_state.truncation.get(),
+        'edit_start': ui_state.edit_layer_start.get(),
+        'edit_end': ui_state.edit_layer_end.get() + 1, # show as inclusive, save as exclusive
+        'example_seed': state.seed,
+    }
+
+    edit_mode_str = params['edit_type']
+    if edit_mode_str == 'latent':
+        edit_mode_str = model.latent_space_name().lower()
+
+    comp_class = state.component_class
+    appl_class = params['output_class']
+    if comp_class != appl_class:
+        comp_class = f'{comp_class}_onto_{appl_class}'
+
+    file_ident = "{model}-{name}-{cls}-{est}-{mode}-{layer}-comp{idx}-range{start}-{end}".format(
+        model=model_name,
+        name=prettify_name(params['name']),
+        cls=comp_class,
+        est=args.estimator,
+        mode=edit_mode_str,
+        layer=args.layer,
+        idx=idx,
+        start=params['edit_start'],
+        end=params['edit_end'],
+    )
+
+    out_dir = Path(__file__).parent / 'out' / 'directions'
+    makedirs(out_dir / file_ident, exist_ok=True)
+
+    with open(out_dir / f"{file_ident}.pkl", 'wb') as outfile:
+        pickle.dump(params, outfile)
+
+    print(f'Direction "{name.get()}" saved as "{file_ident}.pkl"')
+
+    batch_size = ui_state.batch_size.get()
+    len_padded = ((num_strips.get() - 1) // batch_size + 1) * batch_size
+    orig_seed = state.seed
+
+    reset_sliders()
+
+    # Limit max resolution
+    max_H = 512
+    ratio = min(1.0, max_H / inst.output_shape[2])
+
+    strips = [[] for _ in range(len_padded)]
+    for b in range(0, len_padded, batch_size):
+        # Resample
+        resample_latent((orig_seed + b) % np.iinfo(np.int32).max)
+
+        sigmas = np.linspace(slider_value, -slider_value, strip_width.get(), dtype=np.float32)
+        for sid, sigma in enumerate(sigmas):
+            ui_state.sliders[idx].set(sigma)
+            
+            # Advance and show results on screen
+            on_draw()
+            root.update()
+            app.update()
+            
+            batch_res = (255*img).byte().permute(0, 2, 3, 1).detach().cpu().numpy()
+
+            for i, data in enumerate(batch_res):
+                # Save individual
+                name_nodots = file_ident.replace('.', '_')
+                outname = out_dir / file_ident / f"{name_nodots}_ex{b+i}_{sid}.png"
+                im = Image.fromarray(data)
+                im = im.resize((int(ratio*im.size[0]), int(ratio*im.size[1])), Image.ANTIALIAS)
+                im.save(outname)
+                strips[b+i].append(data)
+
+    for i, strip in enumerate(strips[:num_strips.get()]):
+        print(f'Saving strip {i + 1}/{num_strips.get()}', end='\r', flush=True)
+        data = np.hstack(pad_frames(strip))
+        im = Image.fromarray(data)
+        im = im.resize((int(ratio*im.size[0]), int(ratio*im.size[1])), Image.ANTIALIAS)
+        im.save(out_dir / file_ident / f"{file_ident}_ex{i}.png")
+
+    # Reset to original state
+    resample_latent(orig_seed)
+    ui_state.sliders[idx].set(slider_value)
+
+
+# Shared by glumpy and tkinter
+def handle_keypress(code):
+    if code == 65307: # ESC
+        shutdown()
+    elif code == 65360: # HOME
+        reset_sliders()
+    elif code == 114: # R
+        pass #reset_sliders()
+    
+def shutdown():
+    global pending_close
+    pending_close = True
+
+def on_key_release(symbol, modifiers):
+    handle_keypress(symbol)
+
+if __name__=='__main__':
+    setup_model()
+    setup_ui()
+    resample_latent()
+
+    pending_close = False
+    while not pending_close:
+        root.update()
+        app.update()
+        on_draw()
+        reposition_toolbar()
+
+    root.destroy()

model_clip.py

@@ -0,0 +1,436 @@
+from collections import OrderedDict
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu2 = nn.ReLU(inplace=True)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu3 = nn.ReLU(inplace=True)
+
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(OrderedDict([
+                ("-1", nn.AvgPool2d(stride)),
+                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
+                ("1", nn.BatchNorm2d(planes * self.expansion))
+            ]))
+
+    def forward(self, x: torch.Tensor):
+        identity = x
+
+        out = self.relu1(self.bn1(self.conv1(x)))
+        out = self.relu2(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu3(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x):
+        x = x.flatten(start_dim=2).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
+        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x[:1], key=x, value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False
+        )
+        return x.squeeze(0)
+
+
+class ModifiedResNet(nn.Module):
+    """
+    A ResNet class that is similar to torchvision's but contains the following changes:
+    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+    - The final pooling layer is a QKV attention instead of an average pool
+    """
+
+    def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
+        super().__init__()
+        self.output_dim = output_dim
+        self.input_resolution = input_resolution
+
+        # the 3-layer stem
+        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(width // 2)
+        self.relu1 = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.relu2 = nn.ReLU(inplace=True)
+        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.relu3 = nn.ReLU(inplace=True)
+        self.avgpool = nn.AvgPool2d(2)
+
+        # residual layers
+        self._inplanes = width  # this is a *mutable* variable used during construction
+        self.layer1 = self._make_layer(width, layers[0])
+        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
+
+        embed_dim = width * 32  # the ResNet feature dimension
+        self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim)
+
+    def _make_layer(self, planes, blocks, stride=1):
+        layers = [Bottleneck(self._inplanes, planes, stride)]
+
+        self._inplanes = planes * Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(Bottleneck(self._inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        def stem(x):
+            x = self.relu1(self.bn1(self.conv1(x)))
+            x = self.relu2(self.bn2(self.conv2(x)))
+            x = self.relu3(self.bn3(self.conv3(x)))
+            x = self.avgpool(x)
+            return x
+
+        x = x.type(self.conv1.weight.dtype)
+        x = stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.attnpool(x)
+
+        return x
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("c_proj", nn.Linear(d_model * 4, d_model))
+        ]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+
+    def attention(self, x: torch.Tensor):
+        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
+        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.attention(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)])
+
+    def forward(self, x: torch.Tensor):
+        return self.resblocks(x)
+
+
+class VisionTransformer(nn.Module):
+    def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.output_dim = output_dim
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
+
+        scale = width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
+        self.ln_pre = LayerNorm(width)
+
+        self.transformer = Transformer(width, layers, heads)
+
+        self.ln_post = LayerNorm(width)
+        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+
+    def forward(self, x: torch.Tensor):
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        x = self.ln_post(x[:, 0, :])
+
+        if self.proj is not None:
+            x = x @ self.proj
+
+        return x
+
+
+class CLIP(nn.Module):
+    def __init__(self,
+                 embed_dim: int,
+                 # vision
+                 image_resolution: int,
+                 vision_layers: Union[Tuple[int, int, int, int], int],
+                 vision_width: int,
+                 vision_patch_size: int,
+                 # text
+                 context_length: int,
+                 vocab_size: int,
+                 transformer_width: int,
+                 transformer_heads: int,
+                 transformer_layers: int
+                 ):
+        super().__init__()
+
+        self.context_length = context_length
+
+        if isinstance(vision_layers, (tuple, list)):
+            vision_heads = vision_width * 32 // 64
+            self.visual = ModifiedResNet(
+                layers=vision_layers,
+                output_dim=embed_dim,
+                heads=vision_heads,
+                input_resolution=image_resolution,
+                width=vision_width
+            )
+        else:
+            vision_heads = vision_width // 64
+            self.visual = VisionTransformer(
+                input_resolution=image_resolution,
+                patch_size=vision_patch_size,
+                width=vision_width,
+                layers=vision_layers,
+                heads=vision_heads,
+                output_dim=embed_dim
+            )
+
+        self.transformer = Transformer(
+            width=transformer_width,
+            layers=transformer_layers,
+            heads=transformer_heads,
+            attn_mask=self.build_attention_mask()
+        )
+
+        self.vocab_size = vocab_size
+        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+        self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width))
+        self.ln_final = LayerNorm(transformer_width)
+
+        self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.initialize_parameters()
+
+    def initialize_parameters(self):
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.positional_embedding, std=0.01)
+
+        if isinstance(self.visual, ModifiedResNet):
+            if self.visual.attnpool is not None:
+                std = self.visual.attnpool.c_proj.in_features ** -0.5
+                nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
+
+            for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]:
+                for name, param in resnet_block.named_parameters():
+                    if name.endswith("bn3.weight"):
+                        nn.init.zeros_(param)
+
+        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
+        attn_std = self.transformer.width ** -0.5
+        fc_std = (2 * self.transformer.width) ** -0.5
+        for block in self.transformer.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the vision tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.context_length, self.context_length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    @property
+    def dtype(self):
+        return self.visual.conv1.weight.dtype
+
+    def encode_image(self, image):
+        return self.visual(image.type(self.dtype))
+
+    def encode_text(self, text):
+        x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.type(self.dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)
+
+        # x.shape = [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+
+        return x
+
+    def forward(self, image, text):
+        image_features = self.encode_image(image)
+        text_features = self.encode_text(text)
+
+        # normalized features
+        image_features = image_features / image_features.norm(dim=1, keepdim=True)
+        text_features = text_features / text_features.norm(dim=1, keepdim=True)
+
+        # cosine similarity as logits
+        logit_scale = self.logit_scale.exp()
+        logits_per_image = logit_scale * image_features @ text_features.t()
+        logits_per_text = logits_per_image.t()
+
+        # shape = [global_batch_size, global_batch_size]
+        return logits_per_image, logits_per_text
+
+
+def convert_weights(model: nn.Module):
+    """Convert applicable model parameters to fp16"""
+
+    def _convert_weights_to_fp16(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.half()
+            if l.bias is not None:
+                l.bias.data = l.bias.data.half()
+
+        if isinstance(l, nn.MultiheadAttention):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.half()
+
+        for name in ["text_projection", "proj"]:
+            if hasattr(l, name):
+                attr = getattr(l, name)
+                if attr is not None:
+                    attr.data = attr.data.half()
+
+    model.apply(_convert_weights_to_fp16)
+
+
+def build_model(state_dict: dict):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_resolution = vision_patch_size * grid_size
+    else:
+        counts: list = [len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_resolution = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith("transformer.resblocks")))
+
+    model = CLIP(
+        embed_dim,
+        image_resolution, vision_layers, vision_width, vision_patch_size,
+        context_length, vocab_size, transformer_width, transformer_heads, transformer_layers
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        if key in state_dict:
+            del state_dict[key]
+
+    convert_weights(model)
+    model.load_state_dict(state_dict)
+    return model.eval()

models/__init__.py

@@ -0,0 +1,11 @@
+# Copyright 2020 Erik Härkönen. All rights reserved.
+# This file is licensed to you under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License. You may obtain a copy
+# of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software distributed under
+# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
+# OF ANY KIND, either express or implied. See the License for the specific language
+# governing permissions and limitations under the License.
+
+from .wrappers import *

models/biggan/__init__.py

@@ -0,0 +1,8 @@
+from pathlib import Path
+import sys
+
+module_path = Path(__file__).parent / 'pytorch_biggan'
+sys.path.append(str(module_path.resolve()))
+from pytorch_pretrained_biggan import *
+from pytorch_pretrained_biggan.model import GenBlock
+from pytorch_pretrained_biggan.file_utils import http_get, s3_get

models/biggan/pytorch_biggan/.gitignore

@@ -0,0 +1,110 @@
+# 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/
+
+# vscode
+.vscode/
+
+# models
+models/

models/biggan/pytorch_biggan/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2020 Erik Härkönen
+
+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.

models/biggan/pytorch_biggan/MANIFEST.in

@@ -0,0 +1 @@
+include LICENSE

models/biggan/pytorch_biggan/README.md

@@ -0,0 +1,227 @@
+# BigStyleGAN
+This is a copy of HuggingFace's BigGAN implementation, with the addition of layerwise latent inputs.
+
+# PyTorch pretrained BigGAN
+An op-for-op PyTorch reimplementation of DeepMind's BigGAN model with the pre-trained weights from DeepMind.
+
+## Introduction
+
+This repository contains an op-for-op PyTorch reimplementation of DeepMind's BigGAN that was released with the paper [Large Scale GAN Training for High Fidelity Natural Image Synthesis](https://openreview.net/forum?id=B1xsqj09Fm) by Andrew Brock, Jeff Donahue and Karen Simonyan.
+
+This PyTorch implementation of BigGAN is provided with the [pretrained 128x128, 256x256 and 512x512 models by DeepMind](https://tfhub.dev/deepmind/biggan-deep-128/1). We also provide the scripts used to download and convert these models from the TensorFlow Hub models.
+
+This reimplementation was done from the raw computation graph of the Tensorflow version and behave similarly to the TensorFlow version (variance of the output difference of the order of 1e-5).
+
+This implementation currently only contains the generator as the weights of the discriminator were not released (although the structure of the discriminator is very similar to the generator so it could be added pretty easily. Tell me if you want to do a PR on that, I would be happy to help.)
+
+## Installation
+
+This repo was tested on Python 3.6 and PyTorch 1.0.1
+
+PyTorch pretrained BigGAN can be installed from pip as follows:
+```bash
+pip install pytorch-pretrained-biggan
+```
+
+If you simply want to play with the GAN this should be enough.
+
+If you want to use the conversion scripts and the imagenet utilities, additional requirements are needed, in particular TensorFlow and NLTK. To install all the requirements please use the `full_requirements.txt` file:
+```bash
+git clone https://github.com/huggingface/pytorch-pretrained-BigGAN.git
+cd pytorch-pretrained-BigGAN
+pip install -r full_requirements.txt
+```
+
+## Models
+
+This repository provide direct and simple access to the pretrained "deep" versions of BigGAN for 128, 256 and 512 pixels resolutions as described in the [associated publication](https://openreview.net/forum?id=B1xsqj09Fm).
+Here are some details on the models:
+
+- `BigGAN-deep-128`: a 50.4M parameters model generating 128x128 pixels images, the model dump weights 201 MB,
+- `BigGAN-deep-256`: a 55.9M parameters model generating 256x256 pixels images, the model dump weights 224 MB,
+- `BigGAN-deep-512`: a 56.2M parameters model generating 512x512 pixels images, the model dump weights 225 MB.
+
+Please refer to Appendix B of the paper for details on the architectures.
+
+All models comprise pre-computed batch norm statistics for 51 truncation values between 0 and 1 (see Appendix C.1 in the paper for details).
+
+## Usage
+
+Here is a quick-start example using `BigGAN` with a pre-trained model.
+
+See the [doc section](#doc) below for details on these classes and methods.
+
+```python
+import torch
+from pytorch_pretrained_biggan import (BigGAN, one_hot_from_names, truncated_noise_sample,
+                                       save_as_images, display_in_terminal)
+
+# OPTIONAL: if you want to have more information on what's happening, activate the logger as follows
+import logging
+logging.basicConfig(level=logging.INFO)
+
+# Load pre-trained model tokenizer (vocabulary)
+model = BigGAN.from_pretrained('biggan-deep-256')
+
+# Prepare a input
+truncation = 0.4
+class_vector = one_hot_from_names(['soap bubble', 'coffee', 'mushroom'], batch_size=3)
+noise_vector = truncated_noise_sample(truncation=truncation, batch_size=3)
+
+# All in tensors
+noise_vector = torch.from_numpy(noise_vector)
+class_vector = torch.from_numpy(class_vector)
+
+# If you have a GPU, put everything on cuda
+noise_vector = noise_vector.to('cuda')
+class_vector = class_vector.to('cuda')
+model.to('cuda')
+
+# Generate an image
+with torch.no_grad():
+    output = model(noise_vector, class_vector, truncation)
+
+# If you have a GPU put back on CPU
+output = output.to('cpu')
+
+# If you have a sixtel compatible terminal you can display the images in the terminal
+# (see https://github.com/saitoha/libsixel for details)
+display_in_terminal(output)
+
+# Save results as png images
+save_as_images(output)
+```
+
+![output_0](assets/output_0.png)
+![output_1](assets/output_1.png)
+![output_2](assets/output_2.png)
+
+## Doc
+
+### Loading DeepMind's pre-trained weights
+
+To load one of DeepMind's pre-trained models, instantiate a `BigGAN` model with `from_pretrained()` as:
+
+```python
+model = BigGAN.from_pretrained(PRE_TRAINED_MODEL_NAME_OR_PATH, cache_dir=None)
+```
+
+where
+
+- `PRE_TRAINED_MODEL_NAME_OR_PATH` is either:
+
+  - the shortcut name of a Google AI's or OpenAI's pre-trained model selected in the list:
+
+    - `biggan-deep-128`: 12-layer, 768-hidden, 12-heads, 110M parameters
+    - `biggan-deep-256`: 24-layer, 1024-hidden, 16-heads, 340M parameters
+    - `biggan-deep-512`: 12-layer, 768-hidden, 12-heads , 110M parameters
+
+  - a path or url to a pretrained model archive containing:
+
+    - `config.json`: a configuration file for the model, and
+    - `pytorch_model.bin` a PyTorch dump of a pre-trained instance of `BigGAN` (saved with the usual `torch.save()`).
+
+  If `PRE_TRAINED_MODEL_NAME_OR_PATH` is a shortcut name, the pre-trained weights will be downloaded from AWS S3 (see the links [here](pytorch_pretrained_biggan/model.py)) and stored in a cache folder to avoid future download (the cache folder can be found at `~/.pytorch_pretrained_biggan/`).
+- `cache_dir` can be an optional path to a specific directory to download and cache the pre-trained model weights.
+
+### Configuration
+
+`BigGANConfig` is a class to store and load BigGAN configurations. It's defined in [`config.py`](./pytorch_pretrained_biggan/config.py).
+
+Here are some details on the attributes:
+
+- `output_dim`: output resolution of the GAN (128, 256 or 512) for the pre-trained models,
+- `z_dim`: size of the noise vector (128 for the pre-trained models).
+- `class_embed_dim`: size of the class embedding vectors (128 for the pre-trained models).
+- `channel_width`: size of each channel (128 for the pre-trained models).
+- `num_classes`: number of classes in the training dataset, like imagenet (1000 for the pre-trained models).
+- `layers`: A list of layers definition. Each definition for a layer is a triple of [up-sample in the layer ? (bool), number of input channels (int), number of output channels (int)]
+- `attention_layer_position`: Position of the self-attention layer in the layer hierarchy (8 for the pre-trained models).
+- `eps`: epsilon value to use for spectral and batch normalization layers (1e-4 for the pre-trained models).
+- `n_stats`: number of pre-computed statistics for the batch normalization layers associated to various truncation values between 0 and 1 (51 for the pre-trained models).
+
+### Model
+
+`BigGAN` is a PyTorch model (`torch.nn.Module`) of BigGAN defined in [`model.py`](./pytorch_pretrained_biggan/model.py). This model comprises the class embeddings (a linear layer) and the generator with a series of convolutions and conditional batch norms. The discriminator is currently not implemented since pre-trained weights have not been released for it.
+
+The inputs and output are **identical to the TensorFlow model inputs and outputs**.
+
+We detail them here.
+
+`BigGAN` takes as *inputs*:
+
+- `z`: a torch.FloatTensor of shape [batch_size, config.z_dim] with noise sampled from a truncated normal distribution, and
+- `class_label`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to a `sentence B` token (see BERT paper for more details).
+- `truncation`: a float between 0 (not comprised) and 1. The truncation of the truncated normal used for creating the noise vector. This truncation value is used to selecte between a set of pre-computed statistics (means and variances) for the batch norm layers.
+
+`BigGAN` *outputs* an array of shape [batch_size, 3, resolution, resolution] where resolution is 128, 256 or 512 depending of the model:
+
+### Utilities: Images, Noise, Imagenet classes
+
+We provide a few utility method to use the model. They are defined in [`utils.py`](./pytorch_pretrained_biggan/utils.py).
+
+Here are some details on these methods:
+
+- `truncated_noise_sample(batch_size=1, dim_z=128, truncation=1., seed=None)`:
+
+    Create a truncated noise vector.
+    - Params:
+        - batch_size: batch size.
+        - dim_z: dimension of z
+        - truncation: truncation value to use
+        - seed: seed for the random generator
+    - Output:
+        array of shape (batch_size, dim_z)
+
+- `convert_to_images(obj)`:
+
+    Convert an output tensor from BigGAN in a list of images.
+    - Params:
+        - obj: tensor or numpy array of shape (batch_size, channels, height, width)
+    - Output:
+        - list of Pillow Images of size (height, width)
+
+- `save_as_images(obj, file_name='output')`:
+
+    Convert and save an output tensor from BigGAN in a list of saved images.
+    - Params:
+        - obj: tensor or numpy array of shape (batch_size, channels, height, width)
+        - file_name: path and beggingin of filename to save.
+            Images will be saved as `file_name_{image_number}.png`
+
+- `display_in_terminal(obj)`:
+
+    Convert and display an output tensor from BigGAN in the terminal. This function use `libsixel` and will only work in a libsixel-compatible terminal. Please refer to https://github.com/saitoha/libsixel for more details.
+    - Params:
+        - obj: tensor or numpy array of shape (batch_size, channels, height, width)
+        - file_name: path and beggingin of filename to save.
+            Images will be saved as `file_name_{image_number}.png`
+
+- `one_hot_from_int(int_or_list, batch_size=1)`:
+
+    Create a one-hot vector from a class index or a list of class indices.
+    - Params:
+        - int_or_list: int, or list of int, of the imagenet classes (between 0 and 999)
+        - batch_size: batch size.
+            - If int_or_list is an int create a batch of identical classes.
+            - If int_or_list is a list, we should have `len(int_or_list) == batch_size`
+    - Output:
+        - array of shape (batch_size, 1000)
+
+- `one_hot_from_names(class_name, batch_size=1)`:
+
+    Create a one-hot vector from the name of an imagenet class ('tennis ball', 'daisy', ...). We use NLTK's wordnet search to try to find the relevant synset of ImageNet and take the first one. If we can't find it direcly, we look at the hyponyms and hypernyms of the class name.
+    - Params:
+        - class_name: string containing the name of an imagenet object.
+    - Output:
+        - array of shape (batch_size, 1000)
+
+## Download and conversion scripts
+
+Scripts to download and convert the TensorFlow models from TensorFlow Hub are provided in [./scripts](./scripts/).
+
+The scripts can be used directly as:
+```bash
+./scripts/download_tf_hub_models.sh
+./scripts/convert_tf_hub_models.sh
+```

models/biggan/pytorch_biggan/full_requirements.txt

@@ -0,0 +1,5 @@
+tensorflow
+tensorflow-hub
+Pillow
+nltk
+libsixel-python

models/biggan/pytorch_biggan/pytorch_pretrained_biggan/__init__.py

@@ -0,0 +1,6 @@
+from .config import BigGANConfig
+from .model import BigGAN
+from .file_utils import PYTORCH_PRETRAINED_BIGGAN_CACHE, cached_path
+from .utils import (truncated_noise_sample, save_as_images,
+                    convert_to_images, display_in_terminal,
+                    one_hot_from_int, one_hot_from_names)

models/biggan/pytorch_biggan/pytorch_pretrained_biggan/config.py

@@ -0,0 +1,70 @@
+# coding: utf-8
+"""
+BigGAN config.
+"""
+from __future__ import (absolute_import, division, print_function, unicode_literals)
+
+import copy
+import json
+
+class BigGANConfig(object):
+    """ Configuration class to store the configuration of a `BigGAN`. 
+        Defaults are for the 128x128 model.
+        layers tuple are (up-sample in the layer ?, input channels, output channels)
+    """
+    def __init__(self,
+                 output_dim=128,
+                 z_dim=128,
+                 class_embed_dim=128,
+                 channel_width=128,
+                 num_classes=1000,
+                 layers=[(False, 16, 16),
+                         (True, 16, 16),
+                         (False, 16, 16),
+                         (True, 16, 8),
+                         (False, 8, 8),
+                         (True, 8, 4),
+                         (False, 4, 4),
+                         (True, 4, 2),
+                         (False, 2, 2),
+                         (True, 2, 1)],
+                 attention_layer_position=8,
+                 eps=1e-4,
+                 n_stats=51):
+        """Constructs BigGANConfig. """
+        self.output_dim = output_dim
+        self.z_dim = z_dim
+        self.class_embed_dim = class_embed_dim
+        self.channel_width = channel_width
+        self.num_classes = num_classes
+        self.layers = layers
+        self.attention_layer_position = attention_layer_position
+        self.eps = eps
+        self.n_stats = n_stats
+
+    @classmethod
+    def from_dict(cls, json_object):
+        """Constructs a `BigGANConfig` from a Python dictionary of parameters."""
+        config = BigGANConfig()
+        for key, value in json_object.items():
+            config.__dict__[key] = value
+        return config
+
+    @classmethod
+    def from_json_file(cls, json_file):
+        """Constructs a `BigGANConfig` from a json file of parameters."""
+        with open(json_file, "r", encoding='utf-8') as reader:
+            text = reader.read()
+        return cls.from_dict(json.loads(text))
+
+    def __repr__(self):
+        return str(self.to_json_string())
+
+    def to_dict(self):
+        """Serializes this instance to a Python dictionary."""
+        output = copy.deepcopy(self.__dict__)
+        return output
+
+    def to_json_string(self):
+        """Serializes this instance to a JSON string."""
+        return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"

models/biggan/pytorch_biggan/pytorch_pretrained_biggan/convert_tf_to_pytorch.py

@@ -0,0 +1,312 @@
+# coding: utf-8
+"""
+Convert a TF Hub model for BigGAN in a PT one.
+"""
+from __future__ import (absolute_import, division, print_function, unicode_literals)
+
+from itertools import chain
+
+import os
+import argparse
+import logging
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.functional import normalize
+
+from .model import BigGAN, WEIGHTS_NAME, CONFIG_NAME
+from .config import BigGANConfig
+
+logger = logging.getLogger(__name__)
+
+
+def extract_batch_norm_stats(tf_model_path, batch_norm_stats_path=None):
+    try:
+        import numpy as np
+        import tensorflow as tf
+        import tensorflow_hub as hub
+    except ImportError:
+        raise ImportError("Loading a TensorFlow models in PyTorch, requires TensorFlow and TF Hub to be installed. "
+                          "Please see https://www.tensorflow.org/install/ for installation instructions for TensorFlow. "
+                          "And see https://github.com/tensorflow/hub for installing Hub. "
+                          "Probably pip install tensorflow tensorflow-hub")
+    tf.reset_default_graph()
+    logger.info('Loading BigGAN module from: {}'.format(tf_model_path))
+    module = hub.Module(tf_model_path)
+    inputs = {k: tf.placeholder(v.dtype, v.get_shape().as_list(), k)
+              for k, v in module.get_input_info_dict().items()}
+    output = module(inputs)
+
+    initializer = tf.global_variables_initializer()
+    sess = tf.Session()
+    stacks = sum(((i*10 + 1, i*10 + 3, i*10 + 6, i*10 + 8) for i in range(50)), ())
+    numpy_stacks = []
+    for i in stacks:
+        logger.info("Retrieving module_apply_default/stack_{}".format(i))
+        try:
+            stack_var = tf.get_default_graph().get_tensor_by_name("module_apply_default/stack_%d:0" % i)
+        except KeyError:
+            break  # We have all the stats
+        numpy_stacks.append(sess.run(stack_var))
+
+    if batch_norm_stats_path is not None:
+        torch.save(numpy_stacks, batch_norm_stats_path)
+    else:
+        return numpy_stacks
+
+
+def build_tf_to_pytorch_map(model, config):
+    """ Build a map from TF variables to PyTorch modules. """
+    tf_to_pt_map = {}
+
+    # Embeddings and GenZ
+    tf_to_pt_map.update({'linear/w/ema_0.9999': model.embeddings.weight,
+                         'Generator/GenZ/G_linear/b/ema_0.9999': model.generator.gen_z.bias,
+                         'Generator/GenZ/G_linear/w/ema_0.9999': model.generator.gen_z.weight_orig,
+                         'Generator/GenZ/G_linear/u0': model.generator.gen_z.weight_u})
+
+    # GBlock blocks
+    model_layer_idx = 0
+    for i, (up, in_channels, out_channels) in enumerate(config.layers):
+        if i == config.attention_layer_position:
+            model_layer_idx += 1
+        layer_str = "Generator/GBlock_%d/" % i if i > 0 else "Generator/GBlock/"
+        layer_pnt = model.generator.layers[model_layer_idx]
+        for i in range(4):  #  Batchnorms
+            batch_str = layer_str + ("BatchNorm_%d/" % i if i > 0 else "BatchNorm/")
+            batch_pnt = getattr(layer_pnt, 'bn_%d' % i)
+            for name in ('offset', 'scale'):
+                sub_module_str = batch_str + name + "/"
+                sub_module_pnt = getattr(batch_pnt, name)
+                tf_to_pt_map.update({sub_module_str + "w/ema_0.9999": sub_module_pnt.weight_orig,
+                                     sub_module_str + "u0": sub_module_pnt.weight_u})
+        for i in range(4):  # Convolutions
+            conv_str = layer_str + "conv%d/" % i
+            conv_pnt = getattr(layer_pnt, 'conv_%d' % i)
+            tf_to_pt_map.update({conv_str + "b/ema_0.9999": conv_pnt.bias,
+                                 conv_str + "w/ema_0.9999": conv_pnt.weight_orig,
+                                 conv_str + "u0": conv_pnt.weight_u})
+        model_layer_idx += 1
+
+    # Attention block
+    layer_str = "Generator/attention/"
+    layer_pnt = model.generator.layers[config.attention_layer_position]
+    tf_to_pt_map.update({layer_str + "gamma/ema_0.9999": layer_pnt.gamma})
+    for pt_name, tf_name in zip(['snconv1x1_g', 'snconv1x1_o_conv', 'snconv1x1_phi', 'snconv1x1_theta'],
+                                ['g/', 'o_conv/', 'phi/', 'theta/']):
+        sub_module_str = layer_str + tf_name
+        sub_module_pnt = getattr(layer_pnt, pt_name)
+        tf_to_pt_map.update({sub_module_str + "w/ema_0.9999": sub_module_pnt.weight_orig,
+                             sub_module_str + "u0": sub_module_pnt.weight_u})
+
+    # final batch norm and conv to rgb
+    layer_str = "Generator/BatchNorm/"
+    layer_pnt = model.generator.bn
+    tf_to_pt_map.update({layer_str + "offset/ema_0.9999": layer_pnt.bias,
+                         layer_str + "scale/ema_0.9999": layer_pnt.weight})
+    layer_str = "Generator/conv_to_rgb/"
+    layer_pnt = model.generator.conv_to_rgb
+    tf_to_pt_map.update({layer_str + "b/ema_0.9999": layer_pnt.bias,
+                         layer_str + "w/ema_0.9999": layer_pnt.weight_orig,
+                         layer_str + "u0": layer_pnt.weight_u})
+    return tf_to_pt_map
+
+
+def load_tf_weights_in_biggan(model, config, tf_model_path, batch_norm_stats_path=None):
+    """ Load tf checkpoints and standing statistics in a pytorch model
+    """
+    try:
+        import numpy as np
+        import tensorflow as tf
+    except ImportError:
+        raise ImportError("Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see "
+            "https://www.tensorflow.org/install/ for installation instructions.")
+    # Load weights from TF model
+    checkpoint_path = tf_model_path + "/variables/variables"
+    init_vars = tf.train.list_variables(checkpoint_path)
+    from pprint import pprint
+    pprint(init_vars)
+
+    # Extract batch norm statistics from model if needed
+    if batch_norm_stats_path:
+        stats = torch.load(batch_norm_stats_path)
+    else:
+        logger.info("Extracting batch norm stats")
+        stats = extract_batch_norm_stats(tf_model_path)
+
+    # Build TF to PyTorch weights loading map
+    tf_to_pt_map = build_tf_to_pytorch_map(model, config)
+
+    tf_weights = {}
+    for name in tf_to_pt_map.keys():
+        array = tf.train.load_variable(checkpoint_path, name)
+        tf_weights[name] = array
+        # logger.info("Loading TF weight {} with shape {}".format(name, array.shape))
+
+    # Load parameters
+    with torch.no_grad():
+        pt_params_pnt = set()
+        for name, pointer in tf_to_pt_map.items():
+            array = tf_weights[name]
+            if pointer.dim() == 1:
+                if pointer.dim() < array.ndim:
+                    array = np.squeeze(array)
+            elif pointer.dim() == 2:  # Weights
+                array = np.transpose(array)
+            elif pointer.dim() == 4:  # Convolutions
+                array = np.transpose(array, (3, 2, 0, 1))
+            else:
+                raise "Wrong dimensions to adjust: " + str((pointer.shape, array.shape))
+            if pointer.shape != array.shape:
+                raise ValueError("Wrong dimensions: " + str((pointer.shape, array.shape)))
+            logger.info("Initialize PyTorch weight {} with shape {}".format(name, pointer.shape))
+            pointer.data = torch.from_numpy(array) if isinstance(array, np.ndarray) else torch.tensor(array)
+            tf_weights.pop(name, None)
+            pt_params_pnt.add(pointer.data_ptr())
+
+        # Prepare SpectralNorm buffers by running one step of Spectral Norm (no need to train the model):
+        for module in model.modules():
+            for n, buffer in module.named_buffers():
+                if n == 'weight_v':
+                    weight_mat = module.weight_orig
+                    weight_mat = weight_mat.reshape(weight_mat.size(0), -1)
+                    u = module.weight_u
+
+                    v = normalize(torch.mv(weight_mat.t(), u), dim=0, eps=config.eps)
+                    buffer.data = v
+                    pt_params_pnt.add(buffer.data_ptr())
+
+                    u = normalize(torch.mv(weight_mat, v), dim=0, eps=config.eps)
+                    module.weight_u.data = u
+                    pt_params_pnt.add(module.weight_u.data_ptr())
+
+        # Load batch norm statistics
+        index = 0
+        for layer in model.generator.layers:
+            if not hasattr(layer, 'bn_0'):
+                continue
+            for i in range(4):  #  Batchnorms
+                bn_pointer = getattr(layer, 'bn_%d' % i)
+                pointer = bn_pointer.running_means
+                if pointer.shape != stats[index].shape:
+                    raise "Wrong dimensions: " + str((pointer.shape, stats[index].shape))
+                pointer.data = torch.from_numpy(stats[index])
+                pt_params_pnt.add(pointer.data_ptr())
+
+                pointer = bn_pointer.running_vars
+                if pointer.shape != stats[index+1].shape:
+                    raise "Wrong dimensions: " + str((pointer.shape, stats[index].shape))
+                pointer.data = torch.from_numpy(stats[index+1])
+                pt_params_pnt.add(pointer.data_ptr())
+
+                index += 2
+
+        bn_pointer = model.generator.bn
+        pointer = bn_pointer.running_means
+        if pointer.shape != stats[index].shape:
+            raise "Wrong dimensions: " + str((pointer.shape, stats[index].shape))
+        pointer.data = torch.from_numpy(stats[index])
+        pt_params_pnt.add(pointer.data_ptr())
+
+        pointer = bn_pointer.running_vars
+        if pointer.shape != stats[index+1].shape:
+            raise "Wrong dimensions: " + str((pointer.shape, stats[index].shape))
+        pointer.data = torch.from_numpy(stats[index+1])
+        pt_params_pnt.add(pointer.data_ptr())
+
+    remaining_params = list(n for n, t in chain(model.named_parameters(), model.named_buffers()) \
+                            if t.data_ptr() not in pt_params_pnt)
+
+    logger.info("TF Weights not copied to PyTorch model: {} -".format(', '.join(tf_weights.keys())))
+    logger.info("Remanining parameters/buffers from PyTorch model: {} -".format(', '.join(remaining_params)))
+
+    return model
+
+
+BigGAN128 = BigGANConfig(output_dim=128, z_dim=128, class_embed_dim=128, channel_width=128, num_classes=1000,
+                         layers=[(False, 16, 16),
+                                 (True, 16, 16),
+                                 (False, 16, 16),
+                                 (True, 16, 8),
+                                 (False, 8, 8),
+                                 (True, 8, 4),
+                                 (False, 4, 4),
+                                 (True, 4, 2),
+                                 (False, 2, 2),
+                                 (True, 2, 1)],
+                         attention_layer_position=8, eps=1e-4, n_stats=51)
+
+BigGAN256 = BigGANConfig(output_dim=256, z_dim=128, class_embed_dim=128, channel_width=128, num_classes=1000,
+                         layers=[(False, 16, 16),
+                                 (True, 16, 16),
+                                 (False, 16, 16),
+                                 (True, 16, 8),
+                                 (False, 8, 8),
+                                 (True, 8, 8),
+                                 (False, 8, 8),
+                                 (True, 8, 4),
+                                 (False, 4, 4),
+                                 (True, 4, 2),
+                                 (False, 2, 2),
+                                 (True, 2, 1)],
+                         attention_layer_position=8, eps=1e-4, n_stats=51)
+
+BigGAN512 = BigGANConfig(output_dim=512, z_dim=128, class_embed_dim=128, channel_width=128, num_classes=1000,
+                         layers=[(False, 16, 16),
+                                 (True, 16, 16),
+                                 (False, 16, 16),
+                                 (True, 16, 8),
+                                 (False, 8, 8),
+                                 (True, 8, 8),
+                                 (False, 8, 8),
+                                 (True, 8, 4),
+                                 (False, 4, 4),
+                                 (True, 4, 2),
+                                 (False, 2, 2),
+                                 (True, 2, 1),
+                                 (False, 1, 1),
+                                 (True, 1, 1)],
+                         attention_layer_position=8, eps=1e-4, n_stats=51)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Convert a BigGAN TF Hub model in a PyTorch model")
+    parser.add_argument("--model_type", type=str, default="", required=True,
+                        help="BigGAN model type (128, 256, 512)")
+    parser.add_argument("--tf_model_path", type=str, default="", required=True,
+                        help="Path of the downloaded TF Hub model")
+    parser.add_argument("--pt_save_path", type=str, default="",
+                        help="Folder to save the PyTorch model (default: Folder of the TF Hub model)")
+    parser.add_argument("--batch_norm_stats_path", type=str, default="",
+                        help="Path of previously extracted batch norm statistics")
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.INFO)
+
+    if not args.pt_save_path:
+        args.pt_save_path = args.tf_model_path
+
+    if args.model_type == "128":
+        config = BigGAN128
+    elif args.model_type == "256":
+        config = BigGAN256
+    elif args.model_type == "512":
+        config = BigGAN512
+    else:
+        raise ValueError("model_type should be one of 128, 256 or 512")
+
+    model = BigGAN(config)
+    model = load_tf_weights_in_biggan(model, config, args.tf_model_path, args.batch_norm_stats_path)
+
+    model_save_path = os.path.join(args.pt_save_path, WEIGHTS_NAME)
+    config_save_path = os.path.join(args.pt_save_path, CONFIG_NAME)
+
+    logger.info("Save model dump to {}".format(model_save_path))
+    torch.save(model.state_dict(), model_save_path)
+    logger.info("Save configuration file to {}".format(config_save_path))
+    with open(config_save_path, "w", encoding="utf-8") as f:
+        f.write(config.to_json_string())
+
+if __name__ == "__main__":
+    main()

models/biggan/pytorch_biggan/pytorch_pretrained_biggan/file_utils.py

@@ -0,0 +1,249 @@
+"""
+Utilities for working with the local dataset cache.
+This file is adapted from the AllenNLP library at https://github.com/allenai/allennlp
+Copyright by the AllenNLP authors.
+"""
+from __future__ import (absolute_import, division, print_function, unicode_literals)
+
+import json
+import logging
+import os
+import shutil
+import tempfile
+from functools import wraps
+from hashlib import sha256
+import sys
+from io import open
+
+import boto3
+import requests
+from botocore.exceptions import ClientError
+from tqdm import tqdm
+
+try:
+    from urllib.parse import urlparse
+except ImportError:
+    from urlparse import urlparse
+
+try:
+    from pathlib import Path
+    PYTORCH_PRETRAINED_BIGGAN_CACHE = Path(os.getenv('PYTORCH_PRETRAINED_BIGGAN_CACHE',
+                                                   Path.home() / '.pytorch_pretrained_biggan'))
+except (AttributeError, ImportError):
+    PYTORCH_PRETRAINED_BIGGAN_CACHE = os.getenv('PYTORCH_PRETRAINED_BIGGAN_CACHE',
+                                              os.path.join(os.path.expanduser("~"), '.pytorch_pretrained_biggan'))
+
+logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
+
+
+def url_to_filename(url, etag=None):
+    """
+    Convert `url` into a hashed filename in a repeatable way.
+    If `etag` is specified, append its hash to the url's, delimited
+    by a period.
+    """
+    url_bytes = url.encode('utf-8')
+    url_hash = sha256(url_bytes)
+    filename = url_hash.hexdigest()
+
+    if etag:
+        etag_bytes = etag.encode('utf-8')
+        etag_hash = sha256(etag_bytes)
+        filename += '.' + etag_hash.hexdigest()
+
+    return filename
+
+
+def filename_to_url(filename, cache_dir=None):
+    """
+    Return the url and etag (which may be ``None``) stored for `filename`.
+    Raise ``EnvironmentError`` if `filename` or its stored metadata do not exist.
+    """
+    if cache_dir is None:
+        cache_dir = PYTORCH_PRETRAINED_BIGGAN_CACHE
+    if sys.version_info[0] == 3 and isinstance(cache_dir, Path):
+        cache_dir = str(cache_dir)
+
+    cache_path = os.path.join(cache_dir, filename)
+    if not os.path.exists(cache_path):
+        raise EnvironmentError("file {} not found".format(cache_path))
+
+    meta_path = cache_path + '.json'
+    if not os.path.exists(meta_path):
+        raise EnvironmentError("file {} not found".format(meta_path))
+
+    with open(meta_path, encoding="utf-8") as meta_file:
+        metadata = json.load(meta_file)
+    url = metadata['url']
+    etag = metadata['etag']
+
+    return url, etag
+
+
+def cached_path(url_or_filename, cache_dir=None):
+    """
+    Given something that might be a URL (or might be a local path),
+    determine which. If it's a URL, download the file and cache it, and
+    return the path to the cached file. If it's already a local path,
+    make sure the file exists and then return the path.
+    """
+    if cache_dir is None:
+        cache_dir = PYTORCH_PRETRAINED_BIGGAN_CACHE
+    if sys.version_info[0] == 3 and isinstance(url_or_filename, Path):
+        url_or_filename = str(url_or_filename)
+    if sys.version_info[0] == 3 and isinstance(cache_dir, Path):
+        cache_dir = str(cache_dir)
+
+    parsed = urlparse(url_or_filename)
+
+    if parsed.scheme in ('http', 'https', 's3'):
+        # URL, so get it from the cache (downloading if necessary)
+        return get_from_cache(url_or_filename, cache_dir)
+    elif os.path.exists(url_or_filename):
+        # File, and it exists.
+        return url_or_filename
+    elif parsed.scheme == '':
+        # File, but it doesn't exist.
+        raise EnvironmentError("file {} not found".format(url_or_filename))
+    else:
+        # Something unknown
+        raise ValueError("unable to parse {} as a URL or as a local path".format(url_or_filename))
+
+
+def split_s3_path(url):
+    """Split a full s3 path into the bucket name and path."""
+    parsed = urlparse(url)
+    if not parsed.netloc or not parsed.path:
+        raise ValueError("bad s3 path {}".format(url))
+    bucket_name = parsed.netloc
+    s3_path = parsed.path
+    # Remove '/' at beginning of path.
+    if s3_path.startswith("/"):
+        s3_path = s3_path[1:]
+    return bucket_name, s3_path
+
+
+def s3_request(func):
+    """
+    Wrapper function for s3 requests in order to create more helpful error
+    messages.
+    """
+
+    @wraps(func)
+    def wrapper(url, *args, **kwargs):
+        try:
+            return func(url, *args, **kwargs)
+        except ClientError as exc:
+            if int(exc.response["Error"]["Code"]) == 404:
+                raise EnvironmentError("file {} not found".format(url))
+            else:
+                raise
+
+    return wrapper
+
+
+@s3_request
+def s3_etag(url):
+    """Check ETag on S3 object."""
+    s3_resource = boto3.resource("s3")
+    bucket_name, s3_path = split_s3_path(url)
+    s3_object = s3_resource.Object(bucket_name, s3_path)
+    return s3_object.e_tag
+
+
+@s3_request
+def s3_get(url, temp_file):
+    """Pull a file directly from S3."""
+    s3_resource = boto3.resource("s3")
+    bucket_name, s3_path = split_s3_path(url)
+    s3_resource.Bucket(bucket_name).download_fileobj(s3_path, temp_file)
+
+
+def http_get(url, temp_file):
+    req = requests.get(url, stream=True)
+    content_length = req.headers.get('Content-Length')
+    total = int(content_length) if content_length is not None else None
+    progress = tqdm(unit="B", total=total)
+    for chunk in req.iter_content(chunk_size=1024):
+        if chunk: # filter out keep-alive new chunks
+            progress.update(len(chunk))
+            temp_file.write(chunk)
+    progress.close()
+
+
+def get_from_cache(url, cache_dir=None):
+    """
+    Given a URL, look for the corresponding dataset in the local cache.
+    If it's not there, download it. Then return the path to the cached file.
+    """
+    if cache_dir is None:
+        cache_dir = PYTORCH_PRETRAINED_BIGGAN_CACHE
+    if sys.version_info[0] == 3 and isinstance(cache_dir, Path):
+        cache_dir = str(cache_dir)
+
+    if not os.path.exists(cache_dir):
+        os.makedirs(cache_dir)
+
+    # Get eTag to add to filename, if it exists.
+    if url.startswith("s3://"):
+        etag = s3_etag(url)
+    else:
+        response = requests.head(url, allow_redirects=True)
+        if response.status_code != 200:
+            raise IOError("HEAD request failed for url {} with status code {}"
+                          .format(url, response.status_code))
+        etag = response.headers.get("ETag")
+
+    filename = url_to_filename(url, etag)
+
+    # get cache path to put the file
+    cache_path = os.path.join(cache_dir, filename)
+
+    if not os.path.exists(cache_path):
+        # Download to temporary file, then copy to cache dir once finished.
+        # Otherwise you get corrupt cache entries if the download gets interrupted.
+        with tempfile.NamedTemporaryFile() as temp_file:
+            logger.info("%s not found in cache, downloading to %s", url, temp_file.name)
+
+            # GET file object
+            if url.startswith("s3://"):
+                s3_get(url, temp_file)
+            else:
+                http_get(url, temp_file)
+
+            # we are copying the file before closing it, so flush to avoid truncation
+            temp_file.flush()
+            # shutil.copyfileobj() starts at the current position, so go to the start
+            temp_file.seek(0)
+
+            logger.info("copying %s to cache at %s", temp_file.name, cache_path)
+            with open(cache_path, 'wb') as cache_file:
+                shutil.copyfileobj(temp_file, cache_file)
+
+            logger.info("creating metadata file for %s", cache_path)
+            meta = {'url': url, 'etag': etag}
+            meta_path = cache_path + '.json'
+            with open(meta_path, 'w', encoding="utf-8") as meta_file:
+                json.dump(meta, meta_file)
+
+            logger.info("removing temp file %s", temp_file.name)
+
+    return cache_path
+
+
+def read_set_from_file(filename):
+    '''
+    Extract a de-duped collection (set) of text from a file.
+    Expected file format is one item per line.
+    '''
+    collection = set()
+    with open(filename, 'r', encoding='utf-8') as file_:
+        for line in file_:
+            collection.add(line.rstrip())
+    return collection
+
+
+def get_file_extension(path, dot=True, lower=True):
+    ext = os.path.splitext(path)[1]
+    ext = ext if dot else ext[1:]
+    return ext.lower() if lower else ext

models/biggan/pytorch_biggan/pytorch_pretrained_biggan/model.py

@@ -0,0 +1,345 @@
+# coding: utf-8
+""" BigGAN PyTorch model.
+    From "Large Scale GAN Training for High Fidelity Natural Image Synthesis"
+    By Andrew Brocky, Jeff Donahuey and Karen Simonyan.
+    https://openreview.net/forum?id=B1xsqj09Fm
+
+    PyTorch version implemented from the computational graph of the TF Hub module for BigGAN.
+    Some part of the code are adapted from https://github.com/brain-research/self-attention-gan
+
+    This version only comprises the generator (since the discriminator's weights are not released).
+    This version only comprises the "deep" version of BigGAN (see publication).
+
+    Modified by Erik Härkönen:
+    * Added support for per-layer latent vectors
+"""
+from __future__ import (absolute_import, division, print_function, unicode_literals)
+
+import os
+import logging
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .config import BigGANConfig
+from .file_utils import cached_path
+
+logger = logging.getLogger(__name__)
+
+PRETRAINED_MODEL_ARCHIVE_MAP = {
+    'biggan-deep-128': "https://s3.amazonaws.com/models.huggingface.co/biggan/biggan-deep-128-pytorch_model.bin",
+    'biggan-deep-256': "https://s3.amazonaws.com/models.huggingface.co/biggan/biggan-deep-256-pytorch_model.bin",
+    'biggan-deep-512': "https://s3.amazonaws.com/models.huggingface.co/biggan/biggan-deep-512-pytorch_model.bin",
+}
+
+PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    'biggan-deep-128': "https://s3.amazonaws.com/models.huggingface.co/biggan/biggan-deep-128-config.json",
+    'biggan-deep-256': "https://s3.amazonaws.com/models.huggingface.co/biggan/biggan-deep-256-config.json",
+    'biggan-deep-512': "https://s3.amazonaws.com/models.huggingface.co/biggan/biggan-deep-512-config.json",
+}
+
+WEIGHTS_NAME = 'pytorch_model.bin'
+CONFIG_NAME = 'config.json'
+
+
+def snconv2d(eps=1e-12, **kwargs):
+    return nn.utils.spectral_norm(nn.Conv2d(**kwargs), eps=eps)
+
+def snlinear(eps=1e-12, **kwargs):
+    return nn.utils.spectral_norm(nn.Linear(**kwargs), eps=eps)
+
+def sn_embedding(eps=1e-12, **kwargs):
+    return nn.utils.spectral_norm(nn.Embedding(**kwargs), eps=eps)
+
+class SelfAttn(nn.Module):
+    """ Self attention Layer"""
+    def __init__(self, in_channels, eps=1e-12):
+        super(SelfAttn, self).__init__()
+        self.in_channels = in_channels
+        self.snconv1x1_theta = snconv2d(in_channels=in_channels, out_channels=in_channels//8,
+                                        kernel_size=1, bias=False, eps=eps)
+        self.snconv1x1_phi = snconv2d(in_channels=in_channels, out_channels=in_channels//8,
+                                      kernel_size=1, bias=False, eps=eps)
+        self.snconv1x1_g = snconv2d(in_channels=in_channels, out_channels=in_channels//2,
+                                    kernel_size=1, bias=False, eps=eps)
+        self.snconv1x1_o_conv = snconv2d(in_channels=in_channels//2, out_channels=in_channels,
+                                         kernel_size=1, bias=False, eps=eps)
+        self.maxpool = nn.MaxPool2d(2, stride=2, padding=0)
+        self.softmax  = nn.Softmax(dim=-1)
+        self.gamma = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x):
+        _, ch, h, w = x.size()
+        # Theta path
+        theta = self.snconv1x1_theta(x)
+        theta = theta.view(-1, ch//8, h*w)
+        # Phi path
+        phi = self.snconv1x1_phi(x)
+        phi = self.maxpool(phi)
+        phi = phi.view(-1, ch//8, h*w//4)
+        # Attn map
+        attn = torch.bmm(theta.permute(0, 2, 1), phi)
+        attn = self.softmax(attn)
+        # g path
+        g = self.snconv1x1_g(x)
+        g = self.maxpool(g)
+        g = g.view(-1, ch//2, h*w//4)
+        # Attn_g - o_conv
+        attn_g = torch.bmm(g, attn.permute(0, 2, 1))
+        attn_g = attn_g.view(-1, ch//2, h, w)
+        attn_g = self.snconv1x1_o_conv(attn_g)
+        # Out
+        out = x + self.gamma*attn_g
+        return out
+
+
+class BigGANBatchNorm(nn.Module):
+    """ This is a batch norm module that can handle conditional input and can be provided with pre-computed
+        activation means and variances for various truncation parameters.
+
+        We cannot just rely on torch.batch_norm since it cannot handle
+        batched weights (pytorch 1.0.1). We computate batch_norm our-self without updating running means and variances.
+        If you want to train this model you should add running means and variance computation logic.
+    """
+    def __init__(self, num_features, condition_vector_dim=None, n_stats=51, eps=1e-4, conditional=True):
+        super(BigGANBatchNorm, self).__init__()
+        self.num_features = num_features
+        self.eps = eps
+        self.conditional = conditional
+
+        # We use pre-computed statistics for n_stats values of truncation between 0 and 1
+        self.register_buffer('running_means', torch.zeros(n_stats, num_features))
+        self.register_buffer('running_vars', torch.ones(n_stats, num_features))
+        self.step_size = 1.0 / (n_stats - 1)
+
+        if conditional:
+            assert condition_vector_dim is not None
+            self.scale = snlinear(in_features=condition_vector_dim, out_features=num_features, bias=False, eps=eps)
+            self.offset = snlinear(in_features=condition_vector_dim, out_features=num_features, bias=False, eps=eps)
+        else:
+            self.weight = torch.nn.Parameter(torch.Tensor(num_features))
+            self.bias = torch.nn.Parameter(torch.Tensor(num_features))
+
+    def forward(self, x, truncation, condition_vector=None):
+        # Retreive pre-computed statistics associated to this truncation
+        coef, start_idx = math.modf(truncation / self.step_size)
+        start_idx = int(start_idx)
+        if coef != 0.0:  # Interpolate
+            running_mean = self.running_means[start_idx] * coef + self.running_means[start_idx + 1] * (1 - coef)
+            running_var = self.running_vars[start_idx] * coef + self.running_vars[start_idx + 1] * (1 - coef)
+        else:
+            running_mean = self.running_means[start_idx]
+            running_var = self.running_vars[start_idx]
+
+        if self.conditional:
+            running_mean = running_mean.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+            running_var = running_var.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+
+            weight = 1 + self.scale(condition_vector).unsqueeze(-1).unsqueeze(-1)
+            bias = self.offset(condition_vector).unsqueeze(-1).unsqueeze(-1)
+
+            out = (x - running_mean) / torch.sqrt(running_var + self.eps) * weight + bias
+        else:
+            out = F.batch_norm(x, running_mean, running_var, self.weight, self.bias,
+                               training=False, momentum=0.0, eps=self.eps)
+
+        return out
+
+
+class GenBlock(nn.Module):
+    def __init__(self, in_size, out_size, condition_vector_dim, reduction_factor=4, up_sample=False,
+                 n_stats=51, eps=1e-12):
+        super(GenBlock, self).__init__()
+        self.up_sample = up_sample
+        self.drop_channels = (in_size != out_size)
+        middle_size = in_size // reduction_factor
+
+        self.bn_0 = BigGANBatchNorm(in_size, condition_vector_dim, n_stats=n_stats, eps=eps, conditional=True)
+        self.conv_0 = snconv2d(in_channels=in_size, out_channels=middle_size, kernel_size=1, eps=eps)
+
+        self.bn_1 = BigGANBatchNorm(middle_size, condition_vector_dim, n_stats=n_stats, eps=eps, conditional=True)
+        self.conv_1 = snconv2d(in_channels=middle_size, out_channels=middle_size, kernel_size=3, padding=1, eps=eps)
+
+        self.bn_2 = BigGANBatchNorm(middle_size, condition_vector_dim, n_stats=n_stats, eps=eps, conditional=True)
+        self.conv_2 = snconv2d(in_channels=middle_size, out_channels=middle_size, kernel_size=3, padding=1, eps=eps)
+
+        self.bn_3 = BigGANBatchNorm(middle_size, condition_vector_dim, n_stats=n_stats, eps=eps, conditional=True)
+        self.conv_3 = snconv2d(in_channels=middle_size, out_channels=out_size, kernel_size=1, eps=eps)
+
+        self.relu = nn.ReLU()
+
+    def forward(self, x, cond_vector, truncation):
+        x0 = x
+
+        x = self.bn_0(x, truncation, cond_vector)
+        x = self.relu(x)
+        x = self.conv_0(x)
+
+        x = self.bn_1(x, truncation, cond_vector)
+        x = self.relu(x)
+        if self.up_sample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv_1(x)
+
+        x = self.bn_2(x, truncation, cond_vector)
+        x = self.relu(x)
+        x = self.conv_2(x)
+
+        x = self.bn_3(x, truncation, cond_vector)
+        x = self.relu(x)
+        x = self.conv_3(x)
+
+        if self.drop_channels:
+            new_channels = x0.shape[1] // 2
+            x0 = x0[:, :new_channels, ...]
+        if self.up_sample:
+            x0 = F.interpolate(x0, scale_factor=2, mode='nearest')
+
+        out = x + x0
+        return out
+
+class Generator(nn.Module):
+    def __init__(self, config):
+        super(Generator, self).__init__()
+        self.config = config
+        ch = config.channel_width
+        condition_vector_dim = config.z_dim * 2
+
+        self.gen_z = snlinear(in_features=condition_vector_dim,
+                              out_features=4 * 4 * 16 * ch, eps=config.eps)
+
+        layers = []
+        for i, layer in enumerate(config.layers):
+            if i == config.attention_layer_position:
+                layers.append(SelfAttn(ch*layer[1], eps=config.eps))
+            layers.append(GenBlock(ch*layer[1],
+                                   ch*layer[2],
+                                   condition_vector_dim,
+                                   up_sample=layer[0],
+                                   n_stats=config.n_stats,
+                                   eps=config.eps))
+        self.layers = nn.ModuleList(layers)
+
+        self.bn = BigGANBatchNorm(ch, n_stats=config.n_stats, eps=config.eps, conditional=False)
+        self.relu = nn.ReLU()
+        self.conv_to_rgb = snconv2d(in_channels=ch, out_channels=ch, kernel_size=3, padding=1, eps=config.eps)
+        self.tanh = nn.Tanh()
+
+    def forward(self, cond_vector, truncation):
+        z = self.gen_z(cond_vector[0])
+
+        # We use this conversion step to be able to use TF weights:
+        # TF convention on shape is [batch, height, width, channels]
+        # PT convention on shape is [batch, channels, height, width]
+        z = z.view(-1, 4, 4, 16 * self.config.channel_width)
+        z = z.permute(0, 3, 1, 2).contiguous()
+
+        cond_idx = 1
+        for i, layer in enumerate(self.layers):
+            if isinstance(layer, GenBlock):
+                z = layer(z, cond_vector[cond_idx], truncation)
+                cond_idx += 1
+            else:
+                z = layer(z)
+
+        z = self.bn(z, truncation)
+        z = self.relu(z)
+        z = self.conv_to_rgb(z)
+        z = z[:, :3, ...]
+        z = self.tanh(z)
+        return z
+
+class BigGAN(nn.Module):
+    """BigGAN Generator."""
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, cache_dir=None, *inputs, **kwargs):
+        if pretrained_model_name_or_path in PRETRAINED_MODEL_ARCHIVE_MAP:
+            model_file = PRETRAINED_MODEL_ARCHIVE_MAP[pretrained_model_name_or_path]
+            config_file = PRETRAINED_CONFIG_ARCHIVE_MAP[pretrained_model_name_or_path]
+        else:
+            model_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
+            config_file = os.path.join(pretrained_model_name_or_path, CONFIG_NAME)
+
+        try:
+            resolved_model_file = cached_path(model_file, cache_dir=cache_dir)
+            resolved_config_file = cached_path(config_file, cache_dir=cache_dir)
+        except EnvironmentError:
+            logger.error("Wrong model name, should be a valid path to a folder containing "
+                         "a {} file and a {} file or a model name in {}".format(
+                         WEIGHTS_NAME, CONFIG_NAME, PRETRAINED_MODEL_ARCHIVE_MAP.keys()))
+            raise
+
+        logger.info("loading model {} from cache at {}".format(pretrained_model_name_or_path, resolved_model_file))
+
+        # Load config
+        config = BigGANConfig.from_json_file(resolved_config_file)
+        logger.info("Model config {}".format(config))
+
+        # Instantiate model.
+        model = cls(config, *inputs, **kwargs)
+        state_dict = torch.load(resolved_model_file, map_location='cpu' if not torch.cuda.is_available() else None)
+        model.load_state_dict(state_dict, strict=False)
+        return model
+
+    def __init__(self, config):
+        super(BigGAN, self).__init__()
+        self.config = config
+        self.embeddings = nn.Linear(config.num_classes, config.z_dim, bias=False)
+        self.generator = Generator(config)
+        self.n_latents = len(config.layers) + 1 # one for gen_z + one per layer
+
+    def forward(self, z, class_label, truncation):
+        assert 0 < truncation <= 1
+
+        if not isinstance(z, list):
+            z = self.n_latents*[z]
+        
+        if isinstance(class_label, list):
+            embed = [self.embeddings(l) for l in class_label]
+        else:
+            embed = self.n_latents*[self.embeddings(class_label)]
+        
+        assert len(z) == self.n_latents, f'Expected {self.n_latents} latents, got {len(z)}'
+        assert len(embed) == self.n_latents, f'Expected {self.n_latents} class vectors, got {len(class_label)}'
+
+        cond_vectors = [torch.cat((z, e), dim=1) for (z, e) in zip(z, embed)]
+        z = self.generator(cond_vectors, truncation)
+        return z
+
+
+if __name__ == "__main__":
+    import PIL
+    from .utils import truncated_noise_sample, save_as_images, one_hot_from_names
+    from .convert_tf_to_pytorch import load_tf_weights_in_biggan
+
+    load_cache = False
+    cache_path = './saved_model.pt'
+    config = BigGANConfig()
+    model = BigGAN(config)
+    if not load_cache:
+        model = load_tf_weights_in_biggan(model, config, './models/model_128/', './models/model_128/batchnorms_stats.bin')
+        torch.save(model.state_dict(), cache_path)
+    else:
+        model.load_state_dict(torch.load(cache_path))
+
+    model.eval()
+
+    truncation = 0.4
+    noise = truncated_noise_sample(batch_size=2, truncation=truncation)
+    label = one_hot_from_names('diver', batch_size=2)
+
+    # Tests
+    # noise = np.zeros((1, 128))
+    # label = [983]
+
+    noise = torch.tensor(noise, dtype=torch.float)
+    label = torch.tensor(label, dtype=torch.float)
+    with torch.no_grad():
+        outputs = model(noise, label, truncation)
+    print(outputs.shape)
+
+    save_as_images(outputs)

models/biggan/pytorch_biggan/pytorch_pretrained_biggan/utils.py

@@ -0,0 +1,216 @@
+# coding: utf-8
+""" BigGAN utilities to prepare truncated noise samples and convert/save/display output images.
+    Also comprise ImageNet utilities to prepare one hot input vectors for ImageNet classes.
+    We use Wordnet so you can just input a name in a string and automatically get a corresponding
+    imagenet class if it exists (or a hypo/hypernym exists in imagenet).
+"""
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import json
+import logging
+from io import BytesIO
+
+import numpy as np
+from scipy.stats import truncnorm
+
+logger = logging.getLogger(__name__)
+
+NUM_CLASSES = 1000
+
+
+def truncated_noise_sample(batch_size=1, dim_z=128, truncation=1., seed=None):
+    """ Create a truncated noise vector.
+        Params:
+            batch_size: batch size.
+            dim_z: dimension of z
+            truncation: truncation value to use
+            seed: seed for the random generator
+        Output:
+            array of shape (batch_size, dim_z)
+    """
+    state = None if seed is None else np.random.RandomState(seed)
+    values = truncnorm.rvs(-2, 2, size=(batch_size, dim_z), random_state=state).astype(np.float32)
+    return truncation * values
+
+
+def convert_to_images(obj):
+    """ Convert an output tensor from BigGAN in a list of images.
+        Params:
+            obj: tensor or numpy array of shape (batch_size, channels, height, width)
+        Output:
+            list of Pillow Images of size (height, width)
+    """
+    try:
+        import PIL
+    except ImportError:
+        raise ImportError("Please install Pillow to use images: pip install Pillow")
+
+    if not isinstance(obj, np.ndarray):
+        obj = obj.detach().numpy()
+
+    obj = obj.transpose((0, 2, 3, 1))
+    obj = np.clip(((obj + 1) / 2.0) * 256, 0, 255)
+
+    img = []
+    for i, out in enumerate(obj):
+        out_array = np.asarray(np.uint8(out), dtype=np.uint8)
+        img.append(PIL.Image.fromarray(out_array))
+    return img
+
+
+def save_as_images(obj, file_name='output'):
+    """ Convert and save an output tensor from BigGAN in a list of saved images.
+        Params:
+            obj: tensor or numpy array of shape (batch_size, channels, height, width)
+            file_name: path and beggingin of filename to save.
+                Images will be saved as `file_name_{image_number}.png`
+    """
+    img = convert_to_images(obj)
+
+    for i, out in enumerate(img):
+        current_file_name = file_name + '_%d.png' % i
+        logger.info("Saving image to {}".format(current_file_name))
+        out.save(current_file_name, 'png')
+
+
+def display_in_terminal(obj):
+    """ Convert and display an output tensor from BigGAN in the terminal.
+        This function use `libsixel` and will only work in a libsixel-compatible terminal.
+        Please refer to https://github.com/saitoha/libsixel for more details.
+
+        Params:
+            obj: tensor or numpy array of shape (batch_size, channels, height, width)
+            file_name: path and beggingin of filename to save.
+                Images will be saved as `file_name_{image_number}.png`
+    """
+    try:
+        import PIL
+        from libsixel import (sixel_output_new, sixel_dither_new, sixel_dither_initialize,
+                              sixel_dither_set_palette, sixel_dither_set_pixelformat,
+                              sixel_dither_get, sixel_encode, sixel_dither_unref,
+                              sixel_output_unref, SIXEL_PIXELFORMAT_RGBA8888,
+                              SIXEL_PIXELFORMAT_RGB888, SIXEL_PIXELFORMAT_PAL8,
+                              SIXEL_PIXELFORMAT_G8, SIXEL_PIXELFORMAT_G1)
+    except ImportError:
+        raise ImportError("Display in Terminal requires Pillow, libsixel "
+                          "and a libsixel compatible terminal. "
+                          "Please read info at https://github.com/saitoha/libsixel "
+                          "and install with pip install Pillow libsixel-python")
+
+    s = BytesIO()
+
+    images = convert_to_images(obj)
+    widths, heights = zip(*(i.size for i in images))
+
+    output_width = sum(widths)
+    output_height = max(heights)
+
+    output_image = PIL.Image.new('RGB', (output_width, output_height))
+
+    x_offset = 0
+    for im in images:
+        output_image.paste(im, (x_offset,0))
+        x_offset += im.size[0]
+
+    try:
+        data = output_image.tobytes()
+    except NotImplementedError:
+        data = output_image.tostring()
+    output = sixel_output_new(lambda data, s: s.write(data), s)
+
+    try:
+        if output_image.mode == 'RGBA':
+            dither = sixel_dither_new(256)
+            sixel_dither_initialize(dither, data, output_width, output_height, SIXEL_PIXELFORMAT_RGBA8888)
+        elif output_image.mode == 'RGB':
+            dither = sixel_dither_new(256)
+            sixel_dither_initialize(dither, data, output_width, output_height, SIXEL_PIXELFORMAT_RGB888)
+        elif output_image.mode == 'P':
+            palette = output_image.getpalette()
+            dither = sixel_dither_new(256)
+            sixel_dither_set_palette(dither, palette)
+            sixel_dither_set_pixelformat(dither, SIXEL_PIXELFORMAT_PAL8)
+        elif output_image.mode == 'L':
+            dither = sixel_dither_get(SIXEL_BUILTIN_G8)
+            sixel_dither_set_pixelformat(dither, SIXEL_PIXELFORMAT_G8)
+        elif output_image.mode == '1':
+            dither = sixel_dither_get(SIXEL_BUILTIN_G1)
+            sixel_dither_set_pixelformat(dither, SIXEL_PIXELFORMAT_G1)
+        else:
+            raise RuntimeError('unexpected output_image mode')
+        try:
+            sixel_encode(data, output_width, output_height, 1, dither, output)
+            print(s.getvalue().decode('ascii'))
+        finally:
+            sixel_dither_unref(dither)
+    finally:
+        sixel_output_unref(output)
+
+
+def one_hot_from_int(int_or_list, batch_size=1):
+    """ Create a one-hot vector from a class index or a list of class indices.
+        Params:
+            int_or_list: int, or list of int, of the imagenet classes (between 0 and 999)
+            batch_size: batch size.
+                If int_or_list is an int create a batch of identical classes.
+                If int_or_list is a list, we should have `len(int_or_list) == batch_size`
+        Output:
+            array of shape (batch_size, 1000)
+    """
+    if isinstance(int_or_list, int):
+        int_or_list = [int_or_list]
+
+    if len(int_or_list) == 1 and batch_size > 1:
+        int_or_list = [int_or_list[0]] * batch_size
+
+    assert batch_size == len(int_or_list)
+
+    array = np.zeros((batch_size, NUM_CLASSES), dtype=np.float32)
+    for i, j in enumerate(int_or_list):
+        array[i, j] = 1.0
+    return array
+
+
+def one_hot_from_names(class_name_or_list, batch_size=1):
+    """ Create a one-hot vector from the name of an imagenet class ('tennis ball', 'daisy', ...).
+        We use NLTK's wordnet search to try to find the relevant synset of ImageNet and take the first one.
+        If we can't find it direcly, we look at the hyponyms and hypernyms of the class name.
+
+        Params:
+            class_name_or_list: string containing the name of an imagenet object or a list of such strings (for a batch).
+        Output:
+            array of shape (batch_size, 1000)
+    """
+    try:
+        from nltk.corpus import wordnet as wn
+    except ImportError:
+        raise ImportError("You need to install nltk to use this function")
+
+    if not isinstance(class_name_or_list, (list, tuple)):
+        class_name_or_list = [class_name_or_list]
+    else:
+        batch_size = max(batch_size, len(class_name_or_list))
+
+    classes = []
+    for class_name in class_name_or_list:
+        class_name = class_name.replace(" ", "_")
+
+        original_synsets = wn.synsets(class_name)
+        original_synsets = list(filter(lambda s: s.pos() == 'n', original_synsets))  # keep only names
+        if not original_synsets:
+            return None
+
+        possible_synsets = list(filter(lambda s: s.offset() in IMAGENET, original_synsets))
+        if possible_synsets:
+            classes.append(IMAGENET[possible_synsets[0].offset()])
+        else:
+            # try hypernyms and hyponyms
+            possible_synsets = sum([s.hypernyms() + s.hyponyms() for s in original_synsets], [])
+            possible_synsets = list(filter(lambda s: s.offset() in IMAGENET, possible_synsets))
+            if possible_synsets:
+                classes.append(IMAGENET[possible_synsets[0].offset()])
+
+    return one_hot_from_int(classes, batch_size=batch_size)
+
+
+IMAGENET = {1440764: 0, 1443537: 1, 1484850: 2, 1491361: 3, 1494475: 4, 1496331: 5, 1498041: 6, 1514668: 7, 1514859: 8, 1518878: 9, 1530575: 10, 1531178: 11, 1532829: 12, 1534433: 13, 1537544: 14, 1558993: 15, 1560419: 16, 1580077: 17, 1582220: 18, 1592084: 19, 1601694: 20, 1608432: 21, 1614925: 22, 1616318: 23, 1622779: 24, 1629819: 25, 1630670: 26, 1631663: 27, 1632458: 28, 1632777: 29, 1641577: 30, 1644373: 31, 1644900: 32, 1664065: 33, 1665541: 34, 1667114: 35, 1667778: 36, 1669191: 37, 1675722: 38, 1677366: 39, 1682714: 40, 1685808: 41, 1687978: 42, 1688243: 43, 1689811: 44, 1692333: 45, 1693334: 46, 1694178: 47, 1695060: 48, 1697457: 49, 1698640: 50, 1704323: 51, 1728572: 52, 1728920: 53, 1729322: 54, 1729977: 55, 1734418: 56, 1735189: 57, 1737021: 58, 1739381: 59, 1740131: 60, 1742172: 61, 1744401: 62, 1748264: 63, 1749939: 64, 1751748: 65, 1753488: 66, 1755581: 67, 1756291: 68, 1768244: 69, 1770081: 70, 1770393: 71, 1773157: 72, 1773549: 73, 1773797: 74, 1774384: 75, 1774750: 76, 1775062: 77, 1776313: 78, 1784675: 79, 1795545: 80, 1796340: 81, 1797886: 82, 1798484: 83, 1806143: 84, 1806567: 85, 1807496: 86, 1817953: 87, 1818515: 88, 1819313: 89, 1820546: 90, 1824575: 91, 1828970: 92, 1829413: 93, 1833805: 94, 1843065: 95, 1843383: 96, 1847000: 97, 1855032: 98, 1855672: 99, 1860187: 100, 1871265: 101, 1872401: 102, 1873310: 103, 1877812: 104, 1882714: 105, 1883070: 106, 1910747: 107, 1914609: 108, 1917289: 109, 1924916: 110, 1930112: 111, 1943899: 112, 1944390: 113, 1945685: 114, 1950731: 115, 1955084: 116, 1968897: 117, 1978287: 118, 1978455: 119, 1980166: 120, 1981276: 121, 1983481: 122, 1984695: 123, 1985128: 124, 1986214: 125, 1990800: 126, 2002556: 127, 2002724: 128, 2006656: 129, 2007558: 130, 2009229: 131, 2009912: 132, 2011460: 133, 2012849: 134, 2013706: 135, 2017213: 136, 2018207: 137, 2018795: 138, 2025239: 139, 2027492: 140, 2028035: 141, 2033041: 142, 2037110: 143, 2051845: 144, 2056570: 145, 2058221: 146, 2066245: 147, 2071294: 148, 2074367: 149, 2077923: 150, 2085620: 151, 2085782: 152, 2085936: 153, 2086079: 154, 2086240: 155, 2086646: 156, 2086910: 157, 2087046: 158, 2087394: 159, 2088094: 160, 2088238: 161, 2088364: 162, 2088466: 163, 2088632: 164, 2089078: 165, 2089867: 166, 2089973: 167, 2090379: 168, 2090622: 169, 2090721: 170, 2091032: 171, 2091134: 172, 2091244: 173, 2091467: 174, 2091635: 175, 2091831: 176, 2092002: 177, 2092339: 178, 2093256: 179, 2093428: 180, 2093647: 181, 2093754: 182, 2093859: 183, 2093991: 184, 2094114: 185, 2094258: 186, 2094433: 187, 2095314: 188, 2095570: 189, 2095889: 190, 2096051: 191, 2096177: 192, 2096294: 193, 2096437: 194, 2096585: 195, 2097047: 196, 2097130: 197, 2097209: 198, 2097298: 199, 2097474: 200, 2097658: 201, 2098105: 202, 2098286: 203, 2098413: 204, 2099267: 205, 2099429: 206, 2099601: 207, 2099712: 208, 2099849: 209, 2100236: 210, 2100583: 211, 2100735: 212, 2100877: 213, 2101006: 214, 2101388: 215, 2101556: 216, 2102040: 217, 2102177: 218, 2102318: 219, 2102480: 220, 2102973: 221, 2104029: 222, 2104365: 223, 2105056: 224, 2105162: 225, 2105251: 226, 2105412: 227, 2105505: 228, 2105641: 229, 2105855: 230, 2106030: 231, 2106166: 232, 2106382: 233, 2106550: 234, 2106662: 235, 2107142: 236, 2107312: 237, 2107574: 238, 2107683: 239, 2107908: 240, 2108000: 241, 2108089: 242, 2108422: 243, 2108551: 244, 2108915: 245, 2109047: 246, 2109525: 247, 2109961: 248, 2110063: 249, 2110185: 250, 2110341: 251, 2110627: 252, 2110806: 253, 2110958: 254, 2111129: 255, 2111277: 256, 2111500: 257, 2111889: 258, 2112018: 259, 2112137: 260, 2112350: 261, 2112706: 262, 2113023: 263, 2113186: 264, 2113624: 265, 2113712: 266, 2113799: 267, 2113978: 268, 2114367: 269, 2114548: 270, 2114712: 271, 2114855: 272, 2115641: 273, 2115913: 274, 2116738: 275, 2117135: 276, 2119022: 277, 2119789: 278, 2120079: 279, 2120505: 280, 2123045: 281, 2123159: 282, 2123394: 283, 2123597: 284, 2124075: 285, 2125311: 286, 2127052: 287, 2128385: 288, 2128757: 289, 2128925: 290, 2129165: 291, 2129604: 292, 2130308: 293, 2132136: 294, 2133161: 295, 2134084: 296, 2134418: 297, 2137549: 298, 2138441: 299, 2165105: 300, 2165456: 301, 2167151: 302, 2168699: 303, 2169497: 304, 2172182: 305, 2174001: 306, 2177972: 307, 2190166: 308, 2206856: 309, 2219486: 310, 2226429: 311, 2229544: 312, 2231487: 313, 2233338: 314, 2236044: 315, 2256656: 316, 2259212: 317, 2264363: 318, 2268443: 319, 2268853: 320, 2276258: 321, 2277742: 322, 2279972: 323, 2280649: 324, 2281406: 325, 2281787: 326, 2317335: 327, 2319095: 328, 2321529: 329, 2325366: 330, 2326432: 331, 2328150: 332, 2342885: 333, 2346627: 334, 2356798: 335, 2361337: 336, 2363005: 337, 2364673: 338, 2389026: 339, 2391049: 340, 2395406: 341, 2396427: 342, 2397096: 343, 2398521: 344, 2403003: 345, 2408429: 346, 2410509: 347, 2412080: 348, 2415577: 349, 2417914: 350, 2422106: 351, 2422699: 352, 2423022: 353, 2437312: 354, 2437616: 355, 2441942: 356, 2442845: 357, 2443114: 358, 2443484: 359, 2444819: 360, 2445715: 361, 2447366: 362, 2454379: 363, 2457408: 364, 2480495: 365, 2480855: 366, 2481823: 367, 2483362: 368, 2483708: 369, 2484975: 370, 2486261: 371, 2486410: 372, 2487347: 373, 2488291: 374, 2488702: 375, 2489166: 376, 2490219: 377, 2492035: 378, 2492660: 379, 2493509: 380, 2493793: 381, 2494079: 382, 2497673: 383, 2500267: 384, 2504013: 385, 2504458: 386, 2509815: 387, 2510455: 388, 2514041: 389, 2526121: 390, 2536864: 391, 2606052: 392, 2607072: 393, 2640242: 394, 2641379: 395, 2643566: 396, 2655020: 397, 2666196: 398, 2667093: 399, 2669723: 400, 2672831: 401, 2676566: 402, 2687172: 403, 2690373: 404, 2692877: 405, 2699494: 406, 2701002: 407, 2704792: 408, 2708093: 409, 2727426: 410, 2730930: 411, 2747177: 412, 2749479: 413, 2769748: 414, 2776631: 415, 2777292: 416, 2782093: 417, 2783161: 418, 2786058: 419, 2787622: 420, 2788148: 421, 2790996: 422, 2791124: 423, 2791270: 424, 2793495: 425, 2794156: 426, 2795169: 427, 2797295: 428, 2799071: 429, 2802426: 430, 2804414: 431, 2804610: 432, 2807133: 433, 2808304: 434, 2808440: 435, 2814533: 436, 2814860: 437, 2815834: 438, 2817516: 439, 2823428: 440, 2823750: 441, 2825657: 442, 2834397: 443, 2835271: 444, 2837789: 445, 2840245: 446, 2841315: 447, 2843684: 448, 2859443: 449, 2860847: 450, 2865351: 451, 2869837: 452, 2870880: 453, 2871525: 454, 2877765: 455, 2879718: 456, 2883205: 457, 2892201: 458, 2892767: 459, 2894605: 460, 2895154: 461, 2906734: 462, 2909870: 463, 2910353: 464, 2916936: 465, 2917067: 466, 2927161: 467, 2930766: 468, 2939185: 469, 2948072: 470, 2950826: 471, 2951358: 472, 2951585: 473, 2963159: 474, 2965783: 475, 2966193: 476, 2966687: 477, 2971356: 478, 2974003: 479, 2977058: 480, 2978881: 481, 2979186: 482, 2980441: 483, 2981792: 484, 2988304: 485, 2992211: 486, 2992529: 487, 2999410: 488, 3000134: 489, 3000247: 490, 3000684: 491, 3014705: 492, 3016953: 493, 3017168: 494, 3018349: 495, 3026506: 496, 3028079: 497, 3032252: 498, 3041632: 499, 3042490: 500, 3045698: 501, 3047690: 502, 3062245: 503, 3063599: 504, 3063689: 505, 3065424: 506, 3075370: 507, 3085013: 508, 3089624: 509, 3095699: 510, 3100240: 511, 3109150: 512, 3110669: 513, 3124043: 514, 3124170: 515, 3125729: 516, 3126707: 517, 3127747: 518, 3127925: 519, 3131574: 520, 3133878: 521, 3134739: 522, 3141823: 523, 3146219: 524, 3160309: 525, 3179701: 526, 3180011: 527, 3187595: 528, 3188531: 529, 3196217: 530, 3197337: 531, 3201208: 532, 3207743: 533, 3207941: 534, 3208938: 535, 3216828: 536, 3218198: 537, 3220513: 538, 3223299: 539, 3240683: 540, 3249569: 541, 3250847: 542, 3255030: 543, 3259280: 544, 3271574: 545, 3272010: 546, 3272562: 547, 3290653: 548, 3291819: 549, 3297495: 550, 3314780: 551, 3325584: 552, 3337140: 553, 3344393: 554, 3345487: 555, 3347037: 556, 3355925: 557, 3372029: 558, 3376595: 559, 3379051: 560, 3384352: 561, 3388043: 562, 3388183: 563, 3388549: 564, 3393912: 565, 3394916: 566, 3400231: 567, 3404251: 568, 3417042: 569, 3424325: 570, 3425413: 571, 3443371: 572, 3444034: 573, 3445777: 574, 3445924: 575, 3447447: 576, 3447721: 577, 3450230: 578, 3452741: 579, 3457902: 580, 3459775: 581, 3461385: 582, 3467068: 583, 3476684: 584, 3476991: 585, 3478589: 586, 3481172: 587, 3482405: 588, 3483316: 589, 3485407: 590, 3485794: 591, 3492542: 592, 3494278: 593, 3495258: 594, 3496892: 595, 3498962: 596, 3527444: 597, 3529860: 598, 3530642: 599, 3532672: 600, 3534580: 601, 3535780: 602, 3538406: 603, 3544143: 604, 3584254: 605, 3584829: 606, 3590841: 607, 3594734: 608, 3594945: 609, 3595614: 610, 3598930: 611, 3599486: 612, 3602883: 613, 3617480: 614, 3623198: 615, 3627232: 616, 3630383: 617, 3633091: 618, 3637318: 619, 3642806: 620, 3649909: 621, 3657121: 622, 3658185: 623, 3661043: 624, 3662601: 625, 3666591: 626, 3670208: 627, 3673027: 628, 3676483: 629, 3680355: 630, 3690938: 631, 3691459: 632, 3692522: 633, 3697007: 634, 3706229: 635, 3709823: 636, 3710193: 637, 3710637: 638, 3710721: 639, 3717622: 640, 3720891: 641, 3721384: 642, 3724870: 643, 3729826: 644, 3733131: 645, 3733281: 646, 3733805: 647, 3742115: 648, 3743016: 649, 3759954: 650, 3761084: 651, 3763968: 652, 3764736: 653, 3769881: 654, 3770439: 655, 3770679: 656, 3773504: 657, 3775071: 658, 3775546: 659, 3776460: 660, 3777568: 661, 3777754: 662, 3781244: 663, 3782006: 664, 3785016: 665, 3786901: 666, 3787032: 667, 3788195: 668, 3788365: 669, 3791053: 670, 3792782: 671, 3792972: 672, 3793489: 673, 3794056: 674, 3796401: 675, 3803284: 676, 3804744: 677, 3814639: 678, 3814906: 679, 3825788: 680, 3832673: 681, 3837869: 682, 3838899: 683, 3840681: 684, 3841143: 685, 3843555: 686, 3854065: 687, 3857828: 688, 3866082: 689, 3868242: 690, 3868863: 691, 3871628: 692, 3873416: 693, 3874293: 694, 3874599: 695, 3876231: 696, 3877472: 697, 3877845: 698, 3884397: 699, 3887697: 700, 3888257: 701, 3888605: 702, 3891251: 703, 3891332: 704, 3895866: 705, 3899768: 706, 3902125: 707, 3903868: 708, 3908618: 709, 3908714: 710, 3916031: 711, 3920288: 712, 3924679: 713, 3929660: 714, 3929855: 715, 3930313: 716, 3930630: 717, 3933933: 718, 3935335: 719, 3937543: 720, 3938244: 721, 3942813: 722, 3944341: 723, 3947888: 724, 3950228: 725, 3954731: 726, 3956157: 727, 3958227: 728, 3961711: 729, 3967562: 730, 3970156: 731, 3976467: 732, 3976657: 733, 3977966: 734, 3980874: 735, 3982430: 736, 3983396: 737, 3991062: 738, 3992509: 739, 3995372: 740, 3998194: 741, 4004767: 742, 4005630: 743, 4008634: 744, 4009552: 745, 4019541: 746, 4023962: 747, 4026417: 748, 4033901: 749, 4033995: 750, 4037443: 751, 4039381: 752, 4040759: 753, 4041544: 754, 4044716: 755, 4049303: 756, 4065272: 757, 4067472: 758, 4069434: 759, 4070727: 760, 4074963: 761, 4081281: 762, 4086273: 763, 4090263: 764, 4099969: 765, 4111531: 766, 4116512: 767, 4118538: 768, 4118776: 769, 4120489: 770, 4125021: 771, 4127249: 772, 4131690: 773, 4133789: 774, 4136333: 775, 4141076: 776, 4141327: 777, 4141975: 778, 4146614: 779, 4147183: 780, 4149813: 781, 4152593: 782, 4153751: 783, 4154565: 784, 4162706: 785, 4179913: 786, 4192698: 787, 4200800: 788, 4201297: 789, 4204238: 790, 4204347: 791, 4208210: 792, 4209133: 793, 4209239: 794, 4228054: 795, 4229816: 796, 4235860: 797, 4238763: 798, 4239074: 799, 4243546: 800, 4251144: 801, 4252077: 802, 4252225: 803, 4254120: 804, 4254680: 805, 4254777: 806, 4258138: 807, 4259630: 808, 4263257: 809, 4264628: 810, 4265275: 811, 4266014: 812, 4270147: 813, 4273569: 814, 4275548: 815, 4277352: 816, 4285008: 817, 4286575: 818, 4296562: 819, 4310018: 820, 4311004: 821, 4311174: 822, 4317175: 823, 4325704: 824, 4326547: 825, 4328186: 826, 4330267: 827, 4332243: 828, 4335435: 829, 4336792: 830, 4344873: 831, 4346328: 832, 4347754: 833, 4350905: 834, 4355338: 835, 4355933: 836, 4356056: 837, 4357314: 838, 4366367: 839, 4367480: 840, 4370456: 841, 4371430: 842, 4371774: 843, 4372370: 844, 4376876: 845, 4380533: 846, 4389033: 847, 4392985: 848, 4398044: 849, 4399382: 850, 4404412: 851, 4409515: 852, 4417672: 853, 4418357: 854, 4423845: 855, 4428191: 856, 4429376: 857, 4435653: 858, 4442312: 859, 4443257: 860, 4447861: 861, 4456115: 862, 4458633: 863, 4461696: 864, 4462240: 865, 4465501: 866, 4467665: 867, 4476259: 868, 4479046: 869, 4482393: 870, 4483307: 871, 4485082: 872, 4486054: 873, 4487081: 874, 4487394: 875, 4493381: 876, 4501370: 877, 4505470: 878, 4507155: 879, 4509417: 880, 4515003: 881, 4517823: 882, 4522168: 883, 4523525: 884, 4525038: 885, 4525305: 886, 4532106: 887, 4532670: 888, 4536866: 889, 4540053: 890, 4542943: 891, 4548280: 892, 4548362: 893, 4550184: 894, 4552348: 895, 4553703: 896, 4554684: 897, 4557648: 898, 4560804: 899, 4562935: 900, 4579145: 901, 4579432: 902, 4584207: 903, 4589890: 904, 4590129: 905, 4591157: 906, 4591713: 907, 4592741: 908, 4596742: 909, 4597913: 910, 4599235: 911, 4604644: 912, 4606251: 913, 4612504: 914, 4613696: 915, 6359193: 916, 6596364: 917, 6785654: 918, 6794110: 919, 6874185: 920, 7248320: 921, 7565083: 922, 7579787: 923, 7583066: 924, 7584110: 925, 7590611: 926, 7613480: 927, 7614500: 928, 7615774: 929, 7684084: 930, 7693725: 931, 7695742: 932, 7697313: 933, 7697537: 934, 7711569: 935, 7714571: 936, 7714990: 937, 7715103: 938, 7716358: 939, 7716906: 940, 7717410: 941, 7717556: 942, 7718472: 943, 7718747: 944, 7720875: 945, 7730033: 946, 7734744: 947, 7742313: 948, 7745940: 949, 7747607: 950, 7749582: 951, 7753113: 952, 7753275: 953, 7753592: 954, 7754684: 955, 7760859: 956, 7768694: 957, 7802026: 958, 7831146: 959, 7836838: 960, 7860988: 961, 7871810: 962, 7873807: 963, 7875152: 964, 7880968: 965, 7892512: 966, 7920052: 967, 7930864: 968, 7932039: 969, 9193705: 970, 9229709: 971, 9246464: 972, 9256479: 973, 9288635: 974, 9332890: 975, 9399592: 976, 9421951: 977, 9428293: 978, 9468604: 979, 9472597: 980, 9835506: 981, 10148035: 982, 10565667: 983, 11879895: 984, 11939491: 985, 12057211: 986, 12144580: 987, 12267677: 988, 12620546: 989, 12768682: 990, 12985857: 991, 12998815: 992, 13037406: 993, 13040303: 994, 13044778: 995, 13052670: 996, 13054560: 997, 13133613: 998, 15075141: 999}

models/biggan/pytorch_biggan/requirements.txt

@@ -0,0 +1,8 @@
+# PyTorch
+torch>=0.4.1
+# progress bars in model download and training scripts
+tqdm
+# Accessing files from S3 directly.
+boto3
+# Used for downloading models over HTTP
+requests

models/biggan/pytorch_biggan/scripts/convert_tf_hub_models.sh

@@ -0,0 +1,21 @@
+# Copyright (c) 2019-present, Thomas Wolf, Huggingface Inc.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+
+set -e
+set -x
+
+models="128 256 512"
+
+mkdir -p models/model_128
+mkdir -p models/model_256
+mkdir -p models/model_512
+
+# Convert TF Hub models.
+for model in $models
+do
+    pytorch_pretrained_biggan --model_type $model --tf_model_path models/model_$model --pt_save_path models/model_$model
+done

models/biggan/pytorch_biggan/scripts/download_tf_hub_models.sh

@@ -0,0 +1,21 @@
+# Copyright (c) 2019-present, Thomas Wolf, Huggingface Inc.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+
+set -e
+set -x
+
+models="128 256 512"
+
+mkdir -p models/model_128
+mkdir -p models/model_256
+mkdir -p models/model_512
+
+# Download TF Hub models.
+for model in $models
+do
+    curl -L "https://tfhub.dev/deepmind/biggan-deep-$model/1?tf-hub-format=compressed" | tar -zxvC models/model_$model
+done

models/biggan/pytorch_biggan/setup.py

@@ -0,0 +1,69 @@
+"""
+Simple check list from AllenNLP repo: https://github.com/allenai/allennlp/blob/master/setup.py
+
+To create the package for pypi.
+
+1. Change the version in __init__.py and setup.py.
+
+2. Commit these changes with the message: "Release: VERSION"
+
+3. Add a tag in git to mark the release: "git tag VERSION -m'Adds tag VERSION for pypi' "
+   Push the tag to git: git push --tags origin master
+
+4. Build both the sources and the wheel. Do not change anything in setup.py between
+   creating the wheel and the source distribution (obviously).
+
+   For the wheel, run: "python setup.py bdist_wheel" in the top level allennlp directory.
+   (this will build a wheel for the python version you use to build it - make sure you use python 3.x).
+
+   For the sources, run: "python setup.py sdist"
+   You should now have a /dist directory with both .whl and .tar.gz source versions of allennlp.
+
+5. Check that everything looks correct by uploading the package to the pypi test server:
+
+   twine upload dist/* -r pypitest
+   (pypi suggest using twine as other methods upload files via plaintext.)
+
+   Check that you can install it in a virtualenv by running:
+   pip install -i https://testpypi.python.org/pypi allennlp
+
+6. Upload the final version to actual pypi:
+   twine upload dist/* -r pypi
+
+7. Copy the release notes from RELEASE.md to the tag in github once everything is looking hunky-dory.
+
+"""
+from io import open
+from setuptools import find_packages, setup
+
+setup(
+    name="pytorch_pretrained_biggan",
+    version="0.1.0",
+    author="Thomas Wolf",
+    author_email="thomas@huggingface.co",
+    description="PyTorch version of DeepMind's BigGAN model with pre-trained models",
+    long_description=open("README.md", "r", encoding='utf-8').read(),
+    long_description_content_type="text/markdown",
+    keywords='BIGGAN GAN deep learning google deepmind',
+    license='Apache',
+    url="https://github.com/huggingface/pytorch-pretrained-BigGAN",
+    packages=find_packages(exclude=["*.tests", "*.tests.*",
+                                    "tests.*", "tests"]),
+    install_requires=['torch>=0.4.1',
+                      'numpy',
+                      'boto3',
+                      'requests',
+                      'tqdm'],
+    tests_require=['pytest'],
+    entry_points={
+      'console_scripts': [
+        "pytorch_pretrained_biggan=pytorch_pretrained_biggan.convert_tf_to_pytorch:main",
+      ]
+    },
+    classifiers=[
+          'Intended Audience :: Science/Research',
+          'License :: OSI Approved :: Apache Software License',
+          'Programming Language :: Python :: 3',
+          'Topic :: Scientific/Engineering :: Artificial Intelligence',
+    ],
+)

models/stylegan/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 2020 Erik Härkönen. All rights reserved.
+# This file is licensed to you under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License. You may obtain a copy
+# of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software distributed under
+# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
+# OF ANY KIND, either express or implied. See the License for the specific language
+# governing permissions and limitations under the License.
+
+from pathlib import Path
+import sys
+
+#module_path = Path(__file__).parent / 'pytorch_biggan'
+#sys.path.append(str(module_path.resolve()))
+
+from .model import StyleGAN_G, NoiseLayer

models/stylegan/model.py

@@ -0,0 +1,456 @@
+# Copyright 2020 Erik Härkönen. All rights reserved.
+# This file is licensed to you under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License. You may obtain a copy
+# of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software distributed under
+# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
+# OF ANY KIND, either express or implied. See the License for the specific language
+# governing permissions and limitations under the License.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from collections import OrderedDict
+from pathlib import Path
+import requests
+import pickle
+import sys
+
+import numpy as np
+
+# Reimplementation of StyleGAN in PyTorch
+# Source: https://github.com/lernapparat/lernapparat/blob/master/style_gan/pytorch_style_gan.ipynb
+
+class MyLinear(nn.Module):
+    """Linear layer with equalized learning rate and custom learning rate multiplier."""
+    def __init__(self, input_size, output_size, gain=2**(0.5), use_wscale=False, lrmul=1, bias=True):
+        super().__init__()
+        he_std = gain * input_size**(-0.5) # He init
+        # Equalized learning rate and custom learning rate multiplier.
+        if use_wscale:
+            init_std = 1.0 / lrmul
+            self.w_mul = he_std * lrmul
+        else:
+            init_std = he_std / lrmul
+            self.w_mul = lrmul
+        self.weight = torch.nn.Parameter(torch.randn(output_size, input_size) * init_std)
+        if bias:
+            self.bias = torch.nn.Parameter(torch.zeros(output_size))
+            self.b_mul = lrmul
+        else:
+            self.bias = None
+
+    def forward(self, x):
+        bias = self.bias
+        if bias is not None:
+            bias = bias * self.b_mul
+        return F.linear(x, self.weight * self.w_mul, bias)
+
+class MyConv2d(nn.Module):
+    """Conv layer with equalized learning rate and custom learning rate multiplier."""
+    def __init__(self, input_channels, output_channels, kernel_size, gain=2**(0.5), use_wscale=False, lrmul=1, bias=True,
+                intermediate=None, upscale=False):
+        super().__init__()
+        if upscale:
+            self.upscale = Upscale2d()
+        else:
+            self.upscale = None
+        he_std = gain * (input_channels * kernel_size ** 2) ** (-0.5) # He init
+        self.kernel_size = kernel_size
+        if use_wscale:
+            init_std = 1.0 / lrmul
+            self.w_mul = he_std * lrmul
+        else:
+            init_std = he_std / lrmul
+            self.w_mul = lrmul
+        self.weight = torch.nn.Parameter(torch.randn(output_channels, input_channels, kernel_size, kernel_size) * init_std)
+        if bias:
+            self.bias = torch.nn.Parameter(torch.zeros(output_channels))
+            self.b_mul = lrmul
+        else:
+            self.bias = None
+        self.intermediate = intermediate
+
+    def forward(self, x):
+        bias = self.bias
+        if bias is not None:
+            bias = bias * self.b_mul
+        
+        have_convolution = False
+        if self.upscale is not None and min(x.shape[2:]) * 2 >= 128:
+            # this is the fused upscale + conv from StyleGAN, sadly this seems incompatible with the non-fused way
+            # this really needs to be cleaned up and go into the conv...
+            w = self.weight * self.w_mul
+            w = w.permute(1, 0, 2, 3)
+            # probably applying a conv on w would be more efficient. also this quadruples the weight (average)?!
+            w = F.pad(w, (1,1,1,1))
+            w = w[:, :, 1:, 1:]+ w[:, :, :-1, 1:] + w[:, :, 1:, :-1] + w[:, :, :-1, :-1]
+            x = F.conv_transpose2d(x, w, stride=2, padding=(w.size(-1)-1)//2)
+            have_convolution = True
+        elif self.upscale is not None:
+            x = self.upscale(x)
+    
+        if not have_convolution and self.intermediate is None:
+            return F.conv2d(x, self.weight * self.w_mul, bias, padding=self.kernel_size//2)
+        elif not have_convolution:
+            x = F.conv2d(x, self.weight * self.w_mul, None, padding=self.kernel_size//2)
+        
+        if self.intermediate is not None:
+            x = self.intermediate(x)
+        if bias is not None:
+            x = x + bias.view(1, -1, 1, 1)
+        return x
+
+class NoiseLayer(nn.Module):
+    """adds noise. noise is per pixel (constant over channels) with per-channel weight"""
+    def __init__(self, channels):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(channels))
+        self.noise = None
+    
+    def forward(self, x, noise=None):
+        if noise is None and self.noise is None:
+            noise = torch.randn(x.size(0), 1, x.size(2), x.size(3), device=x.device, dtype=x.dtype)
+        elif noise is None:
+            # here is a little trick: if you get all the noiselayers and set each
+            # modules .noise attribute, you can have pre-defined noise.
+            # Very useful for analysis
+            noise = self.noise
+        x = x + self.weight.view(1, -1, 1, 1) * noise
+        return x
+
+class StyleMod(nn.Module):
+    def __init__(self, latent_size, channels, use_wscale):
+        super(StyleMod, self).__init__()
+        self.lin = MyLinear(latent_size,
+                            channels * 2,
+                            gain=1.0, use_wscale=use_wscale)
+        
+    def forward(self, x, latent):
+        style = self.lin(latent) # style => [batch_size, n_channels*2]
+        shape = [-1, 2, x.size(1)] + (x.dim() - 2) * [1]
+        style = style.view(shape)  # [batch_size, 2, n_channels, ...]
+        x = x * (style[:, 0] + 1.) + style[:, 1]
+        return x
+
+class PixelNormLayer(nn.Module):
+    def __init__(self, epsilon=1e-8):
+        super().__init__()
+        self.epsilon = epsilon
+    def forward(self, x):
+        return x * torch.rsqrt(torch.mean(x**2, dim=1, keepdim=True) + self.epsilon)
+
+class BlurLayer(nn.Module):
+    def __init__(self, kernel=[1, 2, 1], normalize=True, flip=False, stride=1):
+        super(BlurLayer, self).__init__()
+        kernel=[1, 2, 1]
+        kernel = torch.tensor(kernel, dtype=torch.float32)
+        kernel = kernel[:, None] * kernel[None, :]
+        kernel = kernel[None, None]
+        if normalize:
+            kernel = kernel / kernel.sum()
+        if flip:
+            kernel = kernel[:, :, ::-1, ::-1]
+        self.register_buffer('kernel', kernel)
+        self.stride = stride
+    
+    def forward(self, x):
+        # expand kernel channels
+        kernel = self.kernel.expand(x.size(1), -1, -1, -1)
+        x = F.conv2d(
+            x,
+            kernel,
+            stride=self.stride,
+            padding=int((self.kernel.size(2)-1)/2),
+            groups=x.size(1)
+        )
+        return x
+
+def upscale2d(x, factor=2, gain=1):
+    assert x.dim() == 4
+    if gain != 1:
+        x = x * gain
+    if factor != 1:
+        shape = x.shape
+        x = x.view(shape[0], shape[1], shape[2], 1, shape[3], 1).expand(-1, -1, -1, factor, -1, factor)
+        x = x.contiguous().view(shape[0], shape[1], factor * shape[2], factor * shape[3])
+    return x
+
+class Upscale2d(nn.Module):
+    def __init__(self, factor=2, gain=1):
+        super().__init__()
+        assert isinstance(factor, int) and factor >= 1
+        self.gain = gain
+        self.factor = factor
+    def forward(self, x):
+        return upscale2d(x, factor=self.factor, gain=self.gain)
+
+class G_mapping(nn.Sequential):
+    def __init__(self, nonlinearity='lrelu', use_wscale=True):
+        act, gain = {'relu': (torch.relu, np.sqrt(2)),
+                     'lrelu': (nn.LeakyReLU(negative_slope=0.2), np.sqrt(2))}[nonlinearity]
+        layers = [
+            ('pixel_norm', PixelNormLayer()),
+            ('dense0', MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale)),
+            ('dense0_act', act),
+            ('dense1', MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale)),
+            ('dense1_act', act),
+            ('dense2', MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale)),
+            ('dense2_act', act),
+            ('dense3', MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale)),
+            ('dense3_act', act),
+            ('dense4', MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale)),
+            ('dense4_act', act),
+            ('dense5', MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale)),
+            ('dense5_act', act),
+            ('dense6', MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale)),
+            ('dense6_act', act),
+            ('dense7', MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale)),
+            ('dense7_act', act)
+        ]
+        super().__init__(OrderedDict(layers))
+        
+    def forward(self, x):
+        return super().forward(x)
+
+class Truncation(nn.Module):
+    def __init__(self, avg_latent, max_layer=8, threshold=0.7):
+        super().__init__()
+        self.max_layer = max_layer
+        self.threshold = threshold
+        self.register_buffer('avg_latent', avg_latent)
+    def forward(self, x):
+        assert x.dim() == 3
+        interp = torch.lerp(self.avg_latent, x, self.threshold)
+        do_trunc = (torch.arange(x.size(1)) < self.max_layer).view(1, -1, 1)
+        return torch.where(do_trunc, interp, x)
+
+class LayerEpilogue(nn.Module):
+    """Things to do at the end of each layer."""
+    def __init__(self, channels, dlatent_size, use_wscale, use_noise, use_pixel_norm, use_instance_norm, use_styles, activation_layer):
+        super().__init__()
+        layers = []
+        if use_noise:
+            layers.append(('noise', NoiseLayer(channels)))
+        layers.append(('activation', activation_layer))
+        if use_pixel_norm:
+            layers.append(('pixel_norm', PixelNorm()))
+        if use_instance_norm:
+            layers.append(('instance_norm', nn.InstanceNorm2d(channels)))
+        self.top_epi = nn.Sequential(OrderedDict(layers))
+        if use_styles:
+            self.style_mod = StyleMod(dlatent_size, channels, use_wscale=use_wscale)
+        else:
+            self.style_mod = None
+    def forward(self, x, dlatents_in_slice=None):
+        x = self.top_epi(x)
+        if self.style_mod is not None:
+            x = self.style_mod(x, dlatents_in_slice)
+        else:
+            assert dlatents_in_slice is None
+        return x
+
+
+class InputBlock(nn.Module):
+    def __init__(self, nf, dlatent_size, const_input_layer, gain, use_wscale, use_noise, use_pixel_norm, use_instance_norm, use_styles, activation_layer):
+        super().__init__()
+        self.const_input_layer = const_input_layer
+        self.nf = nf
+        if self.const_input_layer:
+            # called 'const' in tf
+            self.const = nn.Parameter(torch.ones(1, nf, 4, 4))
+            self.bias = nn.Parameter(torch.ones(nf))
+        else:
+            self.dense = MyLinear(dlatent_size, nf*16, gain=gain/4, use_wscale=use_wscale) # tweak gain to match the official implementation of Progressing GAN
+        self.epi1 = LayerEpilogue(nf, dlatent_size, use_wscale, use_noise, use_pixel_norm, use_instance_norm, use_styles, activation_layer)
+        self.conv = MyConv2d(nf, nf, 3, gain=gain, use_wscale=use_wscale)
+        self.epi2 = LayerEpilogue(nf, dlatent_size, use_wscale, use_noise, use_pixel_norm, use_instance_norm, use_styles, activation_layer)
+        
+    def forward(self, dlatents_in_range):
+        batch_size = dlatents_in_range.size(0)
+        if self.const_input_layer:
+            x = self.const.expand(batch_size, -1, -1, -1)
+            x = x + self.bias.view(1, -1, 1, 1)
+        else:
+            x = self.dense(dlatents_in_range[:, 0]).view(batch_size, self.nf, 4, 4)
+        x = self.epi1(x, dlatents_in_range[:, 0])
+        x = self.conv(x)
+        x = self.epi2(x, dlatents_in_range[:, 1])
+        return x
+
+
+class GSynthesisBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, blur_filter, dlatent_size, gain, use_wscale, use_noise, use_pixel_norm, use_instance_norm, use_styles, activation_layer):
+        # 2**res x 2**res # res = 3..resolution_log2
+        super().__init__()
+        if blur_filter:
+            blur = BlurLayer(blur_filter)
+        else:
+            blur = None
+        self.conv0_up = MyConv2d(in_channels, out_channels, kernel_size=3, gain=gain, use_wscale=use_wscale,
+                                 intermediate=blur, upscale=True)
+        self.epi1 = LayerEpilogue(out_channels, dlatent_size, use_wscale, use_noise, use_pixel_norm, use_instance_norm, use_styles, activation_layer)
+        self.conv1 = MyConv2d(out_channels, out_channels, kernel_size=3, gain=gain, use_wscale=use_wscale)
+        self.epi2 = LayerEpilogue(out_channels, dlatent_size, use_wscale, use_noise, use_pixel_norm, use_instance_norm, use_styles, activation_layer)
+            
+    def forward(self, x, dlatents_in_range):
+        x = self.conv0_up(x)
+        x = self.epi1(x, dlatents_in_range[:, 0])
+        x = self.conv1(x)
+        x = self.epi2(x, dlatents_in_range[:, 1])
+        return x
+
+class G_synthesis(nn.Module):
+    def __init__(self,
+            dlatent_size        = 512,          # Disentangled latent (W) dimensionality.
+            num_channels        = 3,            # Number of output color channels.
+            resolution          = 1024,         # Output resolution.
+            fmap_base           = 8192,         # Overall multiplier for the number of feature maps.
+            fmap_decay          = 1.0,          # log2 feature map reduction when doubling the resolution.
+            fmap_max            = 512,          # Maximum number of feature maps in any layer.
+            use_styles          = True,         # Enable style inputs?
+            const_input_layer   = True,         # First layer is a learned constant?
+            use_noise           = True,         # Enable noise inputs?
+            randomize_noise     = True,         # True = randomize noise inputs every time (non-deterministic), False = read noise inputs from variables.
+            nonlinearity        = 'lrelu',      # Activation function: 'relu', 'lrelu'
+            use_wscale          = True,         # Enable equalized learning rate?
+            use_pixel_norm      = False,        # Enable pixelwise feature vector normalization?
+            use_instance_norm   = True,         # Enable instance normalization?
+            dtype               = torch.float32,  # Data type to use for activations and outputs.
+            blur_filter         = [1,2,1],      # Low-pass filter to apply when resampling activations. None = no filtering.
+        ):
+        
+        super().__init__()
+        def nf(stage):
+            return min(int(fmap_base / (2.0 ** (stage * fmap_decay))), fmap_max)
+        self.dlatent_size = dlatent_size
+        resolution_log2 = int(np.log2(resolution))
+        assert resolution == 2**resolution_log2 and resolution >= 4
+
+        act, gain = {'relu': (torch.relu, np.sqrt(2)),
+                     'lrelu': (nn.LeakyReLU(negative_slope=0.2), np.sqrt(2))}[nonlinearity]
+        num_layers = resolution_log2 * 2 - 2
+        num_styles = num_layers if use_styles else 1
+        torgbs = []
+        blocks = []
+        for res in range(2, resolution_log2 + 1):
+            channels = nf(res-1)
+            name = '{s}x{s}'.format(s=2**res)
+            if res == 2:
+                blocks.append((name,
+                               InputBlock(channels, dlatent_size, const_input_layer, gain, use_wscale,
+                                      use_noise, use_pixel_norm, use_instance_norm, use_styles, act)))
+                
+            else:
+                blocks.append((name,
+                               GSynthesisBlock(last_channels, channels, blur_filter, dlatent_size, gain, use_wscale, use_noise, use_pixel_norm, use_instance_norm, use_styles, act)))
+            last_channels = channels
+        self.torgb = MyConv2d(channels, num_channels, 1, gain=1, use_wscale=use_wscale)
+        self.blocks = nn.ModuleDict(OrderedDict(blocks))
+        
+    def forward(self, dlatents_in):
+        # Input: Disentangled latents (W) [minibatch, num_layers, dlatent_size].
+        # lod_in = tf.cast(tf.get_variable('lod', initializer=np.float32(0), trainable=False), dtype)
+        batch_size = dlatents_in.size(0)       
+        for i, m in enumerate(self.blocks.values()):
+            if i == 0:
+                x = m(dlatents_in[:, 2*i:2*i+2])
+            else:
+                x = m(x, dlatents_in[:, 2*i:2*i+2])
+        rgb = self.torgb(x)
+        return rgb
+
+
+class StyleGAN_G(nn.Sequential):
+    def __init__(self, resolution, truncation=1.0):
+        self.resolution = resolution
+        self.layers = OrderedDict([
+            ('g_mapping', G_mapping()),
+            #('truncation', Truncation(avg_latent)),
+            ('g_synthesis', G_synthesis(resolution=resolution)),
+        ])
+        super().__init__(self.layers)
+
+    def forward(self, x, latent_is_w=False):
+        if isinstance(x, list):
+            assert len(x) == 18, 'Must provide 1 or 18 latents'
+            if not latent_is_w:
+                x = [self.layers['g_mapping'].forward(l) for l in x]
+            x = torch.stack(x, dim=1)
+        else:
+            if not latent_is_w:
+                x = self.layers['g_mapping'].forward(x)
+            x = x.unsqueeze(1).expand(-1, 18, -1)
+
+        x = self.layers['g_synthesis'].forward(x)
+
+        return x
+
+    # From: https://github.com/lernapparat/lernapparat/releases/download/v2019-02-01/
+    def load_weights(self, checkpoint):    
+        self.load_state_dict(torch.load(checkpoint))
+
+    def export_from_tf(self, pickle_path):
+        module_path = Path(__file__).parent / 'stylegan_tf'
+        sys.path.append(str(module_path.resolve()))
+
+        import dnnlib, dnnlib.tflib, pickle, torch, collections
+        dnnlib.tflib.init_tf()
+
+        weights = pickle.load(open(pickle_path,'rb'))
+        weights_pt = [collections.OrderedDict([(k, torch.from_numpy(v.value().eval())) for k,v in w.trainables.items()]) for w in weights]
+        #torch.save(weights_pt, pytorch_name)
+
+        # then on the PyTorch side run
+        state_G, state_D, state_Gs = weights_pt #torch.load('./karras2019stylegan-ffhq-1024x1024.pt')
+        def key_translate(k):
+            k = k.lower().split('/')
+            if k[0] == 'g_synthesis':
+                if not k[1].startswith('torgb'):
+                    k.insert(1, 'blocks')
+                k = '.'.join(k)
+                k = (k.replace('const.const','const').replace('const.bias','bias').replace('const.stylemod','epi1.style_mod.lin')
+                    .replace('const.noise.weight','epi1.top_epi.noise.weight')
+                    .replace('conv.noise.weight','epi2.top_epi.noise.weight')
+                    .replace('conv.stylemod','epi2.style_mod.lin')
+                    .replace('conv0_up.noise.weight', 'epi1.top_epi.noise.weight')
+                    .replace('conv0_up.stylemod','epi1.style_mod.lin')
+                    .replace('conv1.noise.weight', 'epi2.top_epi.noise.weight')
+                    .replace('conv1.stylemod','epi2.style_mod.lin')
+                    .replace('torgb_lod0','torgb'))
+            else:
+                k = '.'.join(k)
+            return k
+
+        def weight_translate(k, w):
+            k = key_translate(k)
+            if k.endswith('.weight'):
+                if w.dim() == 2:
+                    w = w.t()
+                elif w.dim() == 1:
+                    pass
+                else:
+                    assert w.dim() == 4
+                    w = w.permute(3, 2, 0, 1)
+            return w
+
+        # we delete the useless torgb filters
+        param_dict = {key_translate(k) : weight_translate(k, v) for k,v in state_Gs.items() if 'torgb_lod' not in key_translate(k)}
+        if 1:
+            sd_shapes = {k : v.shape for k,v in self.state_dict().items()}
+            param_shapes = {k : v.shape for k,v in param_dict.items() }
+
+            for k in list(sd_shapes)+list(param_shapes):
+                pds = param_shapes.get(k)
+                sds = sd_shapes.get(k)
+                if pds is None:
+                    print ("sd only", k, sds)
+                elif sds is None:
+                    print ("pd only", k, pds)
+                elif sds != pds:
+                    print ("mismatch!", k, pds, sds)
+
+        self.load_state_dict(param_dict, strict=False) # needed for the blur kernels
+        torch.save(self.state_dict(), Path(pickle_path).with_suffix('.pt'))

models/stylegan/stylegan_tf/LICENSE.txt

@@ -0,0 +1,410 @@
+Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+
+
+Attribution-NonCommercial 4.0 International
+
+=======================================================================
+
+Creative Commons Corporation ("Creative Commons") is not a law firm and
+does not provide legal services or legal advice. Distribution of
+Creative Commons public licenses does not create a lawyer-client or
+other relationship. Creative Commons makes its licenses and related
+information available on an "as-is" basis. Creative Commons gives no
+warranties regarding its licenses, any material licensed under their
+terms and conditions, or any related information. Creative Commons
+disclaims all liability for damages resulting from their use to the
+fullest extent possible.
+
+Using Creative Commons Public Licenses
+
+Creative Commons public licenses provide a standard set of terms and
+conditions that creators and other rights holders may use to share
+original works of authorship and other material subject to copyright
+and certain other rights specified in the public license below. The
+following considerations are for informational purposes only, are not
+exhaustive, and do not form part of our licenses.
+
+     Considerations for licensors: Our public licenses are
+     intended for use by those authorized to give the public
+     permission to use material in ways otherwise restricted by
+     copyright and certain other rights. Our licenses are
+     irrevocable. Licensors should read and understand the terms
+     and conditions of the license they choose before applying it.
+     Licensors should also secure all rights necessary before
+     applying our licenses so that the public can reuse the
+     material as expected. Licensors should clearly mark any
+     material not subject to the license. This includes other CC-
+     licensed material, or material used under an exception or
+     limitation to copyright. More considerations for licensors:
+    wiki.creativecommons.org/Considerations_for_licensors
+
+     Considerations for the public: By using one of our public
+     licenses, a licensor grants the public permission to use the
+     licensed material under specified terms and conditions. If
+     the licensor's permission is not necessary for any reason--for
+     example, because of any applicable exception or limitation to
+     copyright--then that use is not regulated by the license. Our
+     licenses grant only permissions under copyright and certain
+     other rights that a licensor has authority to grant. Use of
+     the licensed material may still be restricted for other
+     reasons, including because others have copyright or other
+     rights in the material. A licensor may make special requests,
+     such as asking that all changes be marked or described.
+     Although not required by our licenses, you are encouraged to
+     respect those requests where reasonable. More_considerations
+     for the public: 
+    wiki.creativecommons.org/Considerations_for_licensees
+
+=======================================================================
+
+Creative Commons Attribution-NonCommercial 4.0 International Public
+License
+
+By exercising the Licensed Rights (defined below), You accept and agree
+to be bound by the terms and conditions of this Creative Commons
+Attribution-NonCommercial 4.0 International Public License ("Public
+License"). To the extent this Public License may be interpreted as a
+contract, You are granted the Licensed Rights in consideration of Your
+acceptance of these terms and conditions, and the Licensor grants You
+such rights in consideration of benefits the Licensor receives from
+making the Licensed Material available under these terms and
+conditions.
+
+
+Section 1 -- Definitions.
+
+  a. Adapted Material means material subject to Copyright and Similar
+     Rights that is derived from or based upon the Licensed Material
+     and in which the Licensed Material is translated, altered,
+     arranged, transformed, or otherwise modified in a manner requiring
+     permission under the Copyright and Similar Rights held by the
+     Licensor. For purposes of this Public License, where the Licensed
+     Material is a musical work, performance, or sound recording,
+     Adapted Material is always produced where the Licensed Material is
+     synched in timed relation with a moving image.
+
+  b. Adapter's License means the license You apply to Your Copyright
+     and Similar Rights in Your contributions to Adapted Material in
+     accordance with the terms and conditions of this Public License.
+
+  c. Copyright and Similar Rights means copyright and/or similar rights
+     closely related to copyright including, without limitation,
+     performance, broadcast, sound recording, and Sui Generis Database
+     Rights, without regard to how the rights are labeled or
+     categorized. For purposes of this Public License, the rights
+     specified in Section 2(b)(1)-(2) are not Copyright and Similar
+     Rights.
+  d. Effective Technological Measures means those measures that, in the
+     absence of proper authority, may not be circumvented under laws
+     fulfilling obligations under Article 11 of the WIPO Copyright
+     Treaty adopted on December 20, 1996, and/or similar international
+     agreements.
+
+  e. Exceptions and Limitations means fair use, fair dealing, and/or
+     any other exception or limitation to Copyright and Similar Rights
+     that applies to Your use of the Licensed Material.
+
+  f. Licensed Material means the artistic or literary work, database,
+     or other material to which the Licensor applied this Public
+     License.
+
+  g. Licensed Rights means the rights granted to You subject to the
+     terms and conditions of this Public License, which are limited to
+     all Copyright and Similar Rights that apply to Your use of the
+     Licensed Material and that the Licensor has authority to license.
+
+  h. Licensor means the individual(s) or entity(ies) granting rights
+     under this Public License.
+
+  i. NonCommercial means not primarily intended for or directed towards
+     commercial advantage or monetary compensation. For purposes of
+     this Public License, the exchange of the Licensed Material for
+     other material subject to Copyright and Similar Rights by digital
+     file-sharing or similar means is NonCommercial provided there is
+     no payment of monetary compensation in connection with the
+     exchange.
+
+  j. Share means to provide material to the public by any means or
+     process that requires permission under the Licensed Rights, such
+     as reproduction, public display, public performance, distribution,
+     dissemination, communication, or importation, and to make material
+     available to the public including in ways that members of the
+     public may access the material from a place and at a time
+     individually chosen by them.
+
+  k. Sui Generis Database Rights means rights other than copyright
+     resulting from Directive 96/9/EC of the European Parliament and of
+     the Council of 11 March 1996 on the legal protection of databases,
+     as amended and/or succeeded, as well as other essentially
+     equivalent rights anywhere in the world.
+
+  l. You means the individual or entity exercising the Licensed Rights
+     under this Public License. Your has a corresponding meaning.
+
+
+Section 2 -- Scope.
+
+  a. License grant.
+
+       1. Subject to the terms and conditions of this Public License,
+          the Licensor hereby grants You a worldwide, royalty-free,
+          non-sublicensable, non-exclusive, irrevocable license to
+          exercise the Licensed Rights in the Licensed Material to:
+
+            a. reproduce and Share the Licensed Material, in whole or
+               in part, for NonCommercial purposes only; and
+
+            b. produce, reproduce, and Share Adapted Material for
+               NonCommercial purposes only.
+
+       2. Exceptions and Limitations. For the avoidance of doubt, where
+          Exceptions and Limitations apply to Your use, this Public
+          License does not apply, and You do not need to comply with
+          its terms and conditions.
+
+       3. Term. The term of this Public License is specified in Section
+          6(a).
+
+       4. Media and formats; technical modifications allowed. The
+          Licensor authorizes You to exercise the Licensed Rights in
+          all media and formats whether now known or hereafter created,
+          and to make technical modifications necessary to do so. The
+          Licensor waives and/or agrees not to assert any right or
+          authority to forbid You from making technical modifications
+          necessary to exercise the Licensed Rights, including
+          technical modifications necessary to circumvent Effective
+          Technological Measures. For purposes of this Public License,
+          simply making modifications authorized by this Section 2(a)
+          (4) never produces Adapted Material.
+
+       5. Downstream recipients.
+
+            a. Offer from the Licensor -- Licensed Material. Every
+               recipient of the Licensed Material automatically
+               receives an offer from the Licensor to exercise the
+               Licensed Rights under the terms and conditions of this
+               Public License.
+
+            b. No downstream restrictions. You may not offer or impose
+               any additional or different terms or conditions on, or
+               apply any Effective Technological Measures to, the
+               Licensed Material if doing so restricts exercise of the
+               Licensed Rights by any recipient of the Licensed
+               Material.
+
+       6. No endorsement. Nothing in this Public License constitutes or
+          may be construed as permission to assert or imply that You
+          are, or that Your use of the Licensed Material is, connected
+          with, or sponsored, endorsed, or granted official status by,
+          the Licensor or others designated to receive attribution as
+          provided in Section 3(a)(1)(A)(i).
+
+  b. Other rights.
+
+       1. Moral rights, such as the right of integrity, are not
+          licensed under this Public License, nor are publicity,
+          privacy, and/or other similar personality rights; however, to
+          the extent possible, the Licensor waives and/or agrees not to
+          assert any such rights held by the Licensor to the limited
+          extent necessary to allow You to exercise the Licensed
+          Rights, but not otherwise.
+
+       2. Patent and trademark rights are not licensed under this
+          Public License.
+
+       3. To the extent possible, the Licensor waives any right to
+          collect royalties from You for the exercise of the Licensed
+          Rights, whether directly or through a collecting society
+          under any voluntary or waivable statutory or compulsory
+          licensing scheme. In all other cases the Licensor expressly
+          reserves any right to collect such royalties, including when
+          the Licensed Material is used other than for NonCommercial
+          purposes.
+
+
+Section 3 -- License Conditions.
+
+Your exercise of the Licensed Rights is expressly made subject to the
+following conditions.
+
+  a. Attribution.
+
+       1. If You Share the Licensed Material (including in modified
+          form), You must:
+
+            a. retain the following if it is supplied by the Licensor
+               with the Licensed Material:
+
+                 i. identification of the creator(s) of the Licensed
+                    Material and any others designated to receive
+                    attribution, in any reasonable manner requested by
+                    the Licensor (including by pseudonym if
+                    designated);
+
+                ii. a copyright notice;
+
+               iii. a notice that refers to this Public License;
+
+                iv. a notice that refers to the disclaimer of
+                    warranties;
+
+                 v. a URI or hyperlink to the Licensed Material to the
+                    extent reasonably practicable;
+
+            b. indicate if You modified the Licensed Material and
+               retain an indication of any previous modifications; and
+
+            c. indicate the Licensed Material is licensed under this
+               Public License, and include the text of, or the URI or
+               hyperlink to, this Public License.
+
+       2. You may satisfy the conditions in Section 3(a)(1) in any
+          reasonable manner based on the medium, means, and context in
+          which You Share the Licensed Material. For example, it may be
+          reasonable to satisfy the conditions by providing a URI or
+          hyperlink to a resource that includes the required
+          information.
+
+       3. If requested by the Licensor, You must remove any of the
+          information required by Section 3(a)(1)(A) to the extent
+          reasonably practicable.
+
+       4. If You Share Adapted Material You produce, the Adapter's
+          License You apply must not prevent recipients of the Adapted
+          Material from complying with this Public License.
+
+
+Section 4 -- Sui Generis Database Rights.
+
+Where the Licensed Rights include Sui Generis Database Rights that
+apply to Your use of the Licensed Material:
+
+  a. for the avoidance of doubt, Section 2(a)(1) grants You the right
+     to extract, reuse, reproduce, and Share all or a substantial
+     portion of the contents of the database for NonCommercial purposes
+     only;
+
+  b. if You include all or a substantial portion of the database
+     contents in a database in which You have Sui Generis Database
+     Rights, then the database in which You have Sui Generis Database
+     Rights (but not its individual contents) is Adapted Material; and
+
+  c. You must comply with the conditions in Section 3(a) if You Share
+     all or a substantial portion of the contents of the database.
+
+For the avoidance of doubt, this Section 4 supplements and does not
+replace Your obligations under this Public License where the Licensed
+Rights include other Copyright and Similar Rights.
+
+
+Section 5 -- Disclaimer of Warranties and Limitation of Liability.
+
+  a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
+     EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
+     AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
+     ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
+     IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
+     WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
+     PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
+     ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
+     KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
+     ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
+
+  b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
+     TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
+     NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
+     INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
+     COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
+     USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
+     ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
+     DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
+     IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
+
+  c. The disclaimer of warranties and limitation of liability provided
+     above shall be interpreted in a manner that, to the extent
+     possible, most closely approximates an absolute disclaimer and
+     waiver of all liability.
+
+
+Section 6 -- Term and Termination.
+
+  a. This Public License applies for the term of the Copyright and
+     Similar Rights licensed here. However, if You fail to comply with
+     this Public License, then Your rights under this Public License
+     terminate automatically.
+
+  b. Where Your right to use the Licensed Material has terminated under
+     Section 6(a), it reinstates:
+
+       1. automatically as of the date the violation is cured, provided
+          it is cured within 30 days of Your discovery of the
+          violation; or
+
+       2. upon express reinstatement by the Licensor.
+
+     For the avoidance of doubt, this Section 6(b) does not affect any
+     right the Licensor may have to seek remedies for Your violations
+     of this Public License.
+
+  c. For the avoidance of doubt, the Licensor may also offer the
+     Licensed Material under separate terms or conditions or stop
+     distributing the Licensed Material at any time; however, doing so
+     will not terminate this Public License.
+
+  d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
+     License.
+
+
+Section 7 -- Other Terms and Conditions.
+
+  a. The Licensor shall not be bound by any additional or different
+     terms or conditions communicated by You unless expressly agreed.
+
+  b. Any arrangements, understandings, or agreements regarding the
+     Licensed Material not stated herein are separate from and
+     independent of the terms and conditions of this Public License.
+
+
+Section 8 -- Interpretation.
+
+  a. For the avoidance of doubt, this Public License does not, and
+     shall not be interpreted to, reduce, limit, restrict, or impose
+     conditions on any use of the Licensed Material that could lawfully
+     be made without permission under this Public License.
+
+  b. To the extent possible, if any provision of this Public License is
+     deemed unenforceable, it shall be automatically reformed to the
+     minimum extent necessary to make it enforceable. If the provision
+     cannot be reformed, it shall be severed from this Public License
+     without affecting the enforceability of the remaining terms and
+     conditions.
+
+  c. No term or condition of this Public License will be waived and no
+     failure to comply consented to unless expressly agreed to by the
+     Licensor.
+
+  d. Nothing in this Public License constitutes or may be interpreted
+     as a limitation upon, or waiver of, any privileges and immunities
+     that apply to the Licensor or You, including from the legal
+     processes of any jurisdiction or authority.
+
+=======================================================================
+
+Creative Commons is not a party to its public
+licenses. Notwithstanding, Creative Commons may elect to apply one of
+its public licenses to material it publishes and in those instances
+will be considered the "Licensor." The text of the Creative Commons
+public licenses is dedicated to the public domain under the CC0 Public
+Domain Dedication. Except for the limited purpose of indicating that
+material is shared under a Creative Commons public license or as
+otherwise permitted by the Creative Commons policies published at
+creativecommons.org/policies, Creative Commons does not authorize the
+use of the trademark "Creative Commons" or any other trademark or logo
+of Creative Commons without its prior written consent including,
+without limitation, in connection with any unauthorized modifications
+to any of its public licenses or any other arrangements,
+understandings, or agreements concerning use of licensed material. For
+the avoidance of doubt, this paragraph does not form part of the
+public licenses.
+
+Creative Commons may be contacted at creativecommons.org.

models/stylegan/stylegan_tf/README.md

@@ -0,0 +1,232 @@
+## StyleGAN — Official TensorFlow Implementation
+![Python 3.6](https://img.shields.io/badge/python-3.6-green.svg?style=plastic)
+![TensorFlow 1.10](https://img.shields.io/badge/tensorflow-1.10-green.svg?style=plastic)
+![cuDNN 7.3.1](https://img.shields.io/badge/cudnn-7.3.1-green.svg?style=plastic)
+![License CC BY-NC](https://img.shields.io/badge/license-CC_BY--NC-green.svg?style=plastic)
+
+![Teaser image](./stylegan-teaser.png)
+**Picture:** *These people are not real – they were produced by our generator that allows control over different aspects of the image.*
+
+This repository contains the official TensorFlow implementation of the following paper:
+
+> **A Style-Based Generator Architecture for Generative Adversarial Networks**<br>
+> Tero Karras (NVIDIA), Samuli Laine (NVIDIA), Timo Aila (NVIDIA)<br>
+> https://arxiv.org/abs/1812.04948
+>
+> **Abstract:** *We propose an alternative generator architecture for generative adversarial networks, borrowing from style transfer literature. The new architecture leads to an automatically learned, unsupervised separation of high-level attributes (e.g., pose and identity when trained on human faces) and stochastic variation in the generated images (e.g., freckles, hair), and it enables intuitive, scale-specific control of the synthesis. The new generator improves the state-of-the-art in terms of traditional distribution quality metrics, leads to demonstrably better interpolation properties, and also better disentangles the latent factors of variation. To quantify interpolation quality and disentanglement, we propose two new, automated methods that are applicable to any generator architecture. Finally, we introduce a new, highly varied and high-quality dataset of human faces.*
+
+For business inquiries, please contact [researchinquiries@nvidia.com](mailto:researchinquiries@nvidia.com)<br>
+For press and other inquiries, please contact Hector Marinez at [hmarinez@nvidia.com](mailto:hmarinez@nvidia.com)<br>
+
+**&#9733;&#9733;&#9733; NEW: StyleGAN2 is available at [https://github.com/NVlabs/stylegan2](https://github.com/NVlabs/stylegan2) &#9733;&#9733;&#9733;**
+
+## Resources
+
+Material related to our paper is available via the following links:
+
+- Paper: https://arxiv.org/abs/1812.04948
+- Video: https://youtu.be/kSLJriaOumA
+- Code: https://github.com/NVlabs/stylegan
+- FFHQ: https://github.com/NVlabs/ffhq-dataset
+
+Additional material can be found on Google Drive:
+
+| Path | Description
+| :--- | :----------
+| [StyleGAN](https://drive.google.com/open?id=1uka3a1noXHAydRPRbknqwKVGODvnmUBX) | Main folder.
+| &boxvr;&nbsp; [stylegan-paper.pdf](https://drive.google.com/open?id=1v-HkF3Ehrpon7wVIx4r5DLcko_U_V6Lt) | High-quality version of the paper PDF.
+| &boxvr;&nbsp; [stylegan-video.mp4](https://drive.google.com/open?id=1uzwkZHQX_9pYg1i0d1Nbe3D9xPO8-qBf) | High-quality version of the result video.
+| &boxvr;&nbsp; [images](https://drive.google.com/open?id=1-l46akONUWF6LCpDoeq63H53rD7MeiTd) | Example images produced using our generator.
+| &boxv;&nbsp; &boxvr;&nbsp; [representative-images](https://drive.google.com/open?id=1ToY5P4Vvf5_c3TyUizQ8fckFFoFtBvD8) | High-quality images to be used in articles, blog posts, etc.
+| &boxv;&nbsp; &boxur;&nbsp; [100k-generated-images](https://drive.google.com/open?id=100DJ0QXyG89HZzB4w2Cbyf4xjNK54cQ1) | 100,000 generated images for different amounts of truncation.
+| &boxv;&nbsp; &ensp;&ensp; &boxvr;&nbsp; [ffhq-1024x1024](https://drive.google.com/open?id=14lm8VRN1pr4g_KVe6_LvyDX1PObst6d4) | Generated using Flickr-Faces-HQ dataset at 1024&times;1024.
+| &boxv;&nbsp; &ensp;&ensp; &boxvr;&nbsp; [bedrooms-256x256](https://drive.google.com/open?id=1Vxz9fksw4kgjiHrvHkX4Hze4dyThFW6t) | Generated using LSUN Bedroom dataset at 256&times;256.
+| &boxv;&nbsp; &ensp;&ensp; &boxvr;&nbsp; [cars-512x384](https://drive.google.com/open?id=1MFCvOMdLE2_mpeLPTiDw5dxc2CRuKkzS) | Generated using LSUN Car dataset at 512&times;384.
+| &boxv;&nbsp; &ensp;&ensp; &boxur;&nbsp; [cats-256x256](https://drive.google.com/open?id=1gq-Gj3GRFiyghTPKhp8uDMA9HV_0ZFWQ) | Generated using LSUN Cat dataset at 256&times;256.
+| &boxvr;&nbsp; [videos](https://drive.google.com/open?id=1N8pOd_Bf8v89NGUaROdbD8-ayLPgyRRo) | Example videos produced using our generator.
+| &boxv;&nbsp; &boxur;&nbsp; [high-quality-video-clips](https://drive.google.com/open?id=1NFO7_vH0t98J13ckJYFd7kuaTkyeRJ86) | Individual segments of the result video as high-quality MP4.
+| &boxvr;&nbsp; [ffhq-dataset](https://drive.google.com/open?id=1u2xu7bSrWxrbUxk-dT-UvEJq8IjdmNTP) | Raw data for the [Flickr-Faces-HQ dataset](https://github.com/NVlabs/ffhq-dataset).
+| &boxur;&nbsp; [networks](https://drive.google.com/open?id=1MASQyN5m0voPcx7-9K0r5gObhvvPups7) | Pre-trained networks as pickled instances of [dnnlib.tflib.Network](./dnnlib/tflib/network.py).
+| &ensp;&ensp; &boxvr;&nbsp; [stylegan-ffhq-1024x1024.pkl](https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ) | StyleGAN trained with Flickr-Faces-HQ dataset at 1024&times;1024.
+| &ensp;&ensp; &boxvr;&nbsp; [stylegan-celebahq-1024x1024.pkl](https://drive.google.com/uc?id=1MGqJl28pN4t7SAtSrPdSRJSQJqahkzUf) | StyleGAN trained with CelebA-HQ dataset at 1024&times;1024.
+| &ensp;&ensp; &boxvr;&nbsp; [stylegan-bedrooms-256x256.pkl](https://drive.google.com/uc?id=1MOSKeGF0FJcivpBI7s63V9YHloUTORiF) | StyleGAN trained with LSUN Bedroom dataset at 256&times;256.
+| &ensp;&ensp; &boxvr;&nbsp; [stylegan-cars-512x384.pkl](https://drive.google.com/uc?id=1MJ6iCfNtMIRicihwRorsM3b7mmtmK9c3) | StyleGAN trained with LSUN Car dataset at 512&times;384.
+| &ensp;&ensp; &boxvr;&nbsp; [stylegan-cats-256x256.pkl](https://drive.google.com/uc?id=1MQywl0FNt6lHu8E_EUqnRbviagS7fbiJ) | StyleGAN trained with LSUN Cat dataset at 256&times;256.
+| &ensp;&ensp; &boxur;&nbsp; [metrics](https://drive.google.com/open?id=1MvYdWCBuMfnoYGptRH-AgKLbPTsIQLhl) | Auxiliary networks for the quality and disentanglement metrics.
+| &ensp;&ensp; &ensp;&ensp; &boxvr;&nbsp; [inception_v3_features.pkl](https://drive.google.com/uc?id=1MzTY44rLToO5APn8TZmfR7_ENSe5aZUn) | Standard [Inception-v3](https://arxiv.org/abs/1512.00567) classifier that outputs a raw feature vector.
+| &ensp;&ensp; &ensp;&ensp; &boxvr;&nbsp; [vgg16_zhang_perceptual.pkl](https://drive.google.com/uc?id=1N2-m9qszOeVC9Tq77WxsLnuWwOedQiD2) | Standard [LPIPS](https://arxiv.org/abs/1801.03924) metric to estimate perceptual similarity.
+| &ensp;&ensp; &ensp;&ensp; &boxvr;&nbsp; [celebahq-classifier-00-male.pkl](https://drive.google.com/uc?id=1Q5-AI6TwWhCVM7Muu4tBM7rp5nG_gmCX) | Binary classifier trained to detect a single attribute of CelebA-HQ.
+| &ensp;&ensp; &ensp;&ensp; &boxur;&nbsp;&#x22ef; | Please see the file listing for remaining networks.
+
+## Licenses
+
+All material, excluding the Flickr-Faces-HQ dataset, is made available under [Creative Commons BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/) license by NVIDIA Corporation. You can **use, redistribute, and adapt** the material for **non-commercial purposes**, as long as you give appropriate credit by **citing our paper** and **indicating any changes** that you've made.
+
+For license information regarding the FFHQ dataset, please refer to the [Flickr-Faces-HQ repository](https://github.com/NVlabs/ffhq-dataset).
+
+`inception_v3_features.pkl` and `inception_v3_softmax.pkl` are derived from the pre-trained [Inception-v3](https://arxiv.org/abs/1512.00567) network by Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jonathon Shlens, and Zbigniew Wojna. The network was originally shared under [Apache 2.0](https://github.com/tensorflow/models/blob/master/LICENSE) license on the [TensorFlow Models](https://github.com/tensorflow/models) repository.
+
+`vgg16.pkl` and `vgg16_zhang_perceptual.pkl` are derived from the pre-trained [VGG-16](https://arxiv.org/abs/1409.1556) network by Karen Simonyan and Andrew Zisserman. The network was originally shared under [Creative Commons BY 4.0](https://creativecommons.org/licenses/by/4.0/) license on the [Very Deep Convolutional Networks for Large-Scale Visual Recognition](http://www.robots.ox.ac.uk/~vgg/research/very_deep/) project page.
+
+`vgg16_zhang_perceptual.pkl` is further derived from the pre-trained [LPIPS](https://arxiv.org/abs/1801.03924) weights by Richard Zhang, Phillip Isola, Alexei A. Efros, Eli Shechtman, and Oliver Wang. The weights were originally shared under [BSD 2-Clause "Simplified" License](https://github.com/richzhang/PerceptualSimilarity/blob/master/LICENSE) on the [PerceptualSimilarity](https://github.com/richzhang/PerceptualSimilarity) repository.
+
+## System requirements
+
+* Both Linux and Windows are supported, but we strongly recommend Linux for performance and compatibility reasons.
+* 64-bit Python 3.6 installation. We recommend Anaconda3 with numpy 1.14.3 or newer.
+* TensorFlow 1.10.0 or newer with GPU support.
+* One or more high-end NVIDIA GPUs with at least 11GB of DRAM. We recommend NVIDIA DGX-1 with 8 Tesla V100 GPUs.
+* NVIDIA driver 391.35 or newer, CUDA toolkit 9.0 or newer, cuDNN 7.3.1 or newer.
+
+## Using pre-trained networks
+
+A minimal example of using a pre-trained StyleGAN generator is given in [pretrained_example.py](./pretrained_example.py). When executed, the script downloads a pre-trained StyleGAN generator from Google Drive and uses it to generate an image:
+
+```
+> python pretrained_example.py
+Downloading https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ .... done
+
+Gs                              Params    OutputShape          WeightShape
+---                             ---       ---                  ---
+latents_in                      -         (?, 512)             -
+...
+images_out                      -         (?, 3, 1024, 1024)   -
+---                             ---       ---                  ---
+Total                           26219627
+
+> ls results
+example.png # https://drive.google.com/uc?id=1UDLT_zb-rof9kKH0GwiJW_bS9MoZi8oP
+```
+
+A more advanced example is given in [generate_figures.py](./generate_figures.py). The script reproduces the figures from our paper in order to illustrate style mixing, noise inputs, and truncation:
+```
+> python generate_figures.py
+results/figure02-uncurated-ffhq.png     # https://drive.google.com/uc?id=1U3r1xgcD7o-Fd0SBRpq8PXYajm7_30cu
+results/figure03-style-mixing.png       # https://drive.google.com/uc?id=1U-nlMDtpnf1RcYkaFQtbh5oxnhA97hy6
+results/figure04-noise-detail.png       # https://drive.google.com/uc?id=1UX3m39u_DTU6eLnEW6MqGzbwPFt2R9cG
+results/figure05-noise-components.png   # https://drive.google.com/uc?id=1UQKPcvYVeWMRccGMbs2pPD9PVv1QDyp_
+results/figure08-truncation-trick.png   # https://drive.google.com/uc?id=1ULea0C12zGlxdDQFNLXOWZCHi3QNfk_v
+results/figure10-uncurated-bedrooms.png # https://drive.google.com/uc?id=1UEBnms1XMfj78OHj3_cx80mUf_m9DUJr
+results/figure11-uncurated-cars.png     # https://drive.google.com/uc?id=1UO-4JtAs64Kun5vIj10UXqAJ1d5Ir1Ke
+results/figure12-uncurated-cats.png     # https://drive.google.com/uc?id=1USnJc14prlu3QAYxstrtlfXC9sDWPA-W
+```
+
+The pre-trained networks are stored as standard pickle files on Google Drive:
+
+```
+# Load pre-trained network.
+url = 'https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ' # karras2019stylegan-ffhq-1024x1024.pkl
+with dnnlib.util.open_url(url, cache_dir=config.cache_dir) as f:
+    _G, _D, Gs = pickle.load(f)
+    # _G = Instantaneous snapshot of the generator. Mainly useful for resuming a previous training run.
+    # _D = Instantaneous snapshot of the discriminator. Mainly useful for resuming a previous training run.
+    # Gs = Long-term average of the generator. Yields higher-quality results than the instantaneous snapshot.
+```
+
+The above code downloads the file and unpickles it to yield 3 instances of [dnnlib.tflib.Network](./dnnlib/tflib/network.py). To generate images, you will typically want to use `Gs` &ndash; the other two networks are provided for completeness. In order for `pickle.load()` to work, you will need to have the `dnnlib` source directory in your PYTHONPATH and a `tf.Session` set as default. The session can initialized by calling `dnnlib.tflib.init_tf()`.
+
+There are three ways to use the pre-trained generator:
+
+1. Use `Gs.run()` for immediate-mode operation where the inputs and outputs are numpy arrays:
+   ```
+   # Pick latent vector.
+   rnd = np.random.RandomState(5)
+   latents = rnd.randn(1, Gs.input_shape[1])
+
+   # Generate image.
+   fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
+   images = Gs.run(latents, None, truncation_psi=0.7, randomize_noise=True, output_transform=fmt)
+   ```
+   The first argument is a batch of latent vectors of shape `[num, 512]`. The second argument is reserved for class labels (not used by StyleGAN). The remaining keyword arguments are optional and can be used to further modify the operation (see below). The output is a batch of images, whose format is dictated by the `output_transform` argument.
+
+2. Use `Gs.get_output_for()` to incorporate the generator as a part of a larger TensorFlow expression:
+   ```
+   latents = tf.random_normal([self.minibatch_per_gpu] + Gs_clone.input_shape[1:])
+   images = Gs_clone.get_output_for(latents, None, is_validation=True, randomize_noise=True)
+   images = tflib.convert_images_to_uint8(images)
+   result_expr.append(inception_clone.get_output_for(images))
+   ```
+   The above code is from [metrics/frechet_inception_distance.py](./metrics/frechet_inception_distance.py). It generates a batch of random images and feeds them directly to the [Inception-v3](https://arxiv.org/abs/1512.00567) network without having to convert the data to numpy arrays in between.
+
+3. Look up `Gs.components.mapping` and `Gs.components.synthesis` to access individual sub-networks of the generator. Similar to `Gs`, the sub-networks are represented as independent instances of [dnnlib.tflib.Network](./dnnlib/tflib/network.py):
+   ```
+   src_latents = np.stack(np.random.RandomState(seed).randn(Gs.input_shape[1]) for seed in src_seeds)
+   src_dlatents = Gs.components.mapping.run(src_latents, None) # [seed, layer, component]
+   src_images = Gs.components.synthesis.run(src_dlatents, randomize_noise=False, **synthesis_kwargs)
+   ```
+   The above code is from [generate_figures.py](./generate_figures.py). It first transforms a batch of latent vectors into the intermediate *W* space using the mapping network and then turns these vectors into a batch of images using the synthesis network. The `dlatents` array stores a separate copy of the same *w* vector for each layer of the synthesis network to facilitate style mixing.
+
+The exact details of the generator are defined in [training/networks_stylegan.py](./training/networks_stylegan.py) (see `G_style`, `G_mapping`, and `G_synthesis`). The following keyword arguments can be specified to modify the behavior when calling `run()` and `get_output_for()`:
+
+* `truncation_psi` and `truncation_cutoff` control the truncation trick that that is performed by default when using `Gs` (&psi;=0.7, cutoff=8). It can be disabled by setting `truncation_psi=1` or `is_validation=True`, and the image quality can be further improved at the cost of variation by setting e.g. `truncation_psi=0.5`. Note that truncation is always disabled when using the sub-networks directly. The average *w* needed to manually perform the truncation trick can be looked up using `Gs.get_var('dlatent_avg')`.
+
+* `randomize_noise` determines whether to use re-randomize the noise inputs for each generated image (`True`, default) or whether to use specific noise values for the entire minibatch (`False`). The specific values can be accessed via the `tf.Variable` instances that are found using `[var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]`.
+
+* When using the mapping network directly, you can specify `dlatent_broadcast=None` to disable the automatic duplication of `dlatents` over the layers of the synthesis network.
+
+* Runtime performance can be fine-tuned via `structure='fixed'` and `dtype='float16'`. The former disables support for progressive growing, which is not needed for a fully-trained generator, and the latter performs all computation using half-precision floating point arithmetic.
+
+## Preparing datasets for training
+
+The training and evaluation scripts operate on datasets stored as multi-resolution TFRecords. Each dataset is represented by a directory containing the same image data in several resolutions to enable efficient streaming. There is a separate *.tfrecords file for each resolution, and if the dataset contains labels, they are stored in a separate file as well. By default, the scripts expect to find the datasets at `datasets/<NAME>/<NAME>-<RESOLUTION>.tfrecords`. The directory can be changed by editing [config.py](./config.py):
+
+```
+result_dir = 'results'
+data_dir = 'datasets'
+cache_dir = 'cache'
+```
+
+To obtain the FFHQ dataset (`datasets/ffhq`), please refer to the [Flickr-Faces-HQ repository](https://github.com/NVlabs/ffhq-dataset).
+
+To obtain the CelebA-HQ dataset (`datasets/celebahq`), please refer to the [Progressive GAN repository](https://github.com/tkarras/progressive_growing_of_gans).
+
+To obtain other datasets, including LSUN, please consult their corresponding project pages. The datasets can be converted to multi-resolution TFRecords using the provided [dataset_tool.py](./dataset_tool.py):
+
+```
+> python dataset_tool.py create_lsun datasets/lsun-bedroom-full ~/lsun/bedroom_lmdb --resolution 256
+> python dataset_tool.py create_lsun_wide datasets/lsun-car-512x384 ~/lsun/car_lmdb --width 512 --height 384
+> python dataset_tool.py create_lsun datasets/lsun-cat-full ~/lsun/cat_lmdb --resolution 256
+> python dataset_tool.py create_cifar10 datasets/cifar10 ~/cifar10
+> python dataset_tool.py create_from_images datasets/custom-dataset ~/custom-images
+```
+
+## Training networks
+
+Once the datasets are set up, you can train your own StyleGAN networks as follows:
+
+1. Edit [train.py](./train.py) to specify the dataset and training configuration by uncommenting or editing specific lines.
+2. Run the training script with `python train.py`.
+3. The results are written to a newly created directory `results/<ID>-<DESCRIPTION>`.
+4. The training may take several days (or weeks) to complete, depending on the configuration.
+
+By default, `train.py` is configured to train the highest-quality StyleGAN (configuration F in Table 1) for the FFHQ dataset at 1024&times;1024 resolution using 8 GPUs. Please note that we have used 8 GPUs in all of our experiments. Training with fewer GPUs may not produce identical results &ndash; if you wish to compare against our technique, we strongly recommend using the same number of GPUs.
+
+Expected training times for the default configuration using Tesla V100 GPUs:
+
+| GPUs | 1024&times;1024  | 512&times;512    | 256&times;256    |
+| :--- | :--------------  | :------------    | :------------    |
+| 1    | 41 days 4 hours  | 24 days 21 hours | 14 days 22 hours |
+| 2    | 21 days 22 hours | 13 days 7 hours  | 9 days 5 hours   |
+| 4    | 11 days 8 hours  | 7 days 0 hours   | 4 days 21 hours  |
+| 8    | 6 days 14 hours  | 4 days 10 hours  | 3 days 8 hours   |
+
+## Evaluating quality and disentanglement
+
+The quality and disentanglement metrics used in our paper can be evaluated using [run_metrics.py](./run_metrics.py). By default, the script will evaluate the Fr&eacute;chet Inception Distance (`fid50k`) for the pre-trained FFHQ generator and write the results into a newly created directory under `results`. The exact behavior can be changed by uncommenting or editing specific lines in [run_metrics.py](./run_metrics.py).
+
+Expected evaluation time and results for the pre-trained FFHQ generator using one Tesla V100 GPU:
+
+| Metric    | Time      | Result   | Description
+| :-----    | :---      | :-----   | :----------
+| fid50k    | 16 min    | 4.4159   | Fr&eacute;chet Inception Distance using 50,000 images.
+| ppl_zfull | 55 min    | 664.8854 | Perceptual Path Length for full paths in *Z*.
+| ppl_wfull | 55 min    | 233.3059 | Perceptual Path Length for full paths in *W*.
+| ppl_zend  | 55 min    | 666.1057 | Perceptual Path Length for path endpoints in *Z*.
+| ppl_wend  | 55 min    | 197.2266 | Perceptual Path Length for path endpoints in *W*.
+| ls        | 10 hours  | z: 165.0106<br>w: 3.7447 | Linear Separability in *Z* and *W*.
+
+Please note that the exact results may vary from run to run due to the non-deterministic nature of TensorFlow.
+
+## Acknowledgements
+
+We thank Jaakko Lehtinen, David Luebke, and Tuomas Kynk&auml;&auml;nniemi for in-depth discussions and helpful comments; Janne Hellsten, Tero Kuosmanen, and Pekka J&auml;nis for compute infrastructure and help with the code release.

models/stylegan/stylegan_tf/config.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial
+# 4.0 International License. To view a copy of this license, visit
+# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
+# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
+
+"""Global configuration."""
+
+#----------------------------------------------------------------------------
+# Paths.
+
+result_dir = 'results'
+data_dir = 'datasets'
+cache_dir = 'cache'
+run_dir_ignore = ['results', 'datasets', 'cache']
+
+#----------------------------------------------------------------------------

models/stylegan/stylegan_tf/dataset_tool.py

@@ -0,0 +1,645 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial
+# 4.0 International License. To view a copy of this license, visit
+# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
+# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
+
+"""Tool for creating multi-resolution TFRecords datasets for StyleGAN and ProGAN."""
+
+# pylint: disable=too-many-lines
+import os
+import sys
+import glob
+import argparse
+import threading
+import six.moves.queue as Queue # pylint: disable=import-error
+import traceback
+import numpy as np
+import tensorflow as tf
+import PIL.Image
+import dnnlib.tflib as tflib
+
+from training import dataset
+
+#----------------------------------------------------------------------------
+
+def error(msg):
+    print('Error: ' + msg)
+    exit(1)
+
+#----------------------------------------------------------------------------
+
+class TFRecordExporter:
+    def __init__(self, tfrecord_dir, expected_images, print_progress=True, progress_interval=10):
+        self.tfrecord_dir       = tfrecord_dir
+        self.tfr_prefix         = os.path.join(self.tfrecord_dir, os.path.basename(self.tfrecord_dir))
+        self.expected_images    = expected_images
+        self.cur_images         = 0
+        self.shape              = None
+        self.resolution_log2    = None
+        self.tfr_writers        = []
+        self.print_progress     = print_progress
+        self.progress_interval  = progress_interval
+
+        if self.print_progress:
+            print('Creating dataset "%s"' % tfrecord_dir)
+        if not os.path.isdir(self.tfrecord_dir):
+            os.makedirs(self.tfrecord_dir)
+        assert os.path.isdir(self.tfrecord_dir)
+
+    def close(self):
+        if self.print_progress:
+            print('%-40s\r' % 'Flushing data...', end='', flush=True)
+        for tfr_writer in self.tfr_writers:
+            tfr_writer.close()
+        self.tfr_writers = []
+        if self.print_progress:
+            print('%-40s\r' % '', end='', flush=True)
+            print('Added %d images.' % self.cur_images)
+
+    def choose_shuffled_order(self): # Note: Images and labels must be added in shuffled order.
+        order = np.arange(self.expected_images)
+        np.random.RandomState(123).shuffle(order)
+        return order
+
+    def add_image(self, img):
+        if self.print_progress and self.cur_images % self.progress_interval == 0:
+            print('%d / %d\r' % (self.cur_images, self.expected_images), end='', flush=True)
+        if self.shape is None:
+            self.shape = img.shape
+            self.resolution_log2 = int(np.log2(self.shape[1]))
+            assert self.shape[0] in [1, 3]
+            assert self.shape[1] == self.shape[2]
+            assert self.shape[1] == 2**self.resolution_log2
+            tfr_opt = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.NONE)
+            for lod in range(self.resolution_log2 - 1):
+                tfr_file = self.tfr_prefix + '-r%02d.tfrecords' % (self.resolution_log2 - lod)
+                self.tfr_writers.append(tf.python_io.TFRecordWriter(tfr_file, tfr_opt))
+        assert img.shape == self.shape
+        for lod, tfr_writer in enumerate(self.tfr_writers):
+            if lod:
+                img = img.astype(np.float32)
+                img = (img[:, 0::2, 0::2] + img[:, 0::2, 1::2] + img[:, 1::2, 0::2] + img[:, 1::2, 1::2]) * 0.25
+            quant = np.rint(img).clip(0, 255).astype(np.uint8)
+            ex = tf.train.Example(features=tf.train.Features(feature={
+                'shape': tf.train.Feature(int64_list=tf.train.Int64List(value=quant.shape)),
+                'data': tf.train.Feature(bytes_list=tf.train.BytesList(value=[quant.tostring()]))}))
+            tfr_writer.write(ex.SerializeToString())
+        self.cur_images += 1
+
+    def add_labels(self, labels):
+        if self.print_progress:
+            print('%-40s\r' % 'Saving labels...', end='', flush=True)
+        assert labels.shape[0] == self.cur_images
+        with open(self.tfr_prefix + '-rxx.labels', 'wb') as f:
+            np.save(f, labels.astype(np.float32))
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
+#----------------------------------------------------------------------------
+
+class ExceptionInfo(object):
+    def __init__(self):
+        self.value = sys.exc_info()[1]
+        self.traceback = traceback.format_exc()
+
+#----------------------------------------------------------------------------
+
+class WorkerThread(threading.Thread):
+    def __init__(self, task_queue):
+        threading.Thread.__init__(self)
+        self.task_queue = task_queue
+
+    def run(self):
+        while True:
+            func, args, result_queue = self.task_queue.get()
+            if func is None:
+                break
+            try:
+                result = func(*args)
+            except:
+                result = ExceptionInfo()
+            result_queue.put((result, args))
+
+#----------------------------------------------------------------------------
+
+class ThreadPool(object):
+    def __init__(self, num_threads):
+        assert num_threads >= 1
+        self.task_queue = Queue.Queue()
+        self.result_queues = dict()
+        self.num_threads = num_threads
+        for _idx in range(self.num_threads):
+            thread = WorkerThread(self.task_queue)
+            thread.daemon = True
+            thread.start()
+
+    def add_task(self, func, args=()):
+        assert hasattr(func, '__call__') # must be a function
+        if func not in self.result_queues:
+            self.result_queues[func] = Queue.Queue()
+        self.task_queue.put((func, args, self.result_queues[func]))
+
+    def get_result(self, func): # returns (result, args)
+        result, args = self.result_queues[func].get()
+        if isinstance(result, ExceptionInfo):
+            print('\n\nWorker thread caught an exception:\n' + result.traceback)
+            raise result.value
+        return result, args
+
+    def finish(self):
+        for _idx in range(self.num_threads):
+            self.task_queue.put((None, (), None))
+
+    def __enter__(self): # for 'with' statement
+        return self
+
+    def __exit__(self, *excinfo):
+        self.finish()
+
+    def process_items_concurrently(self, item_iterator, process_func=lambda x: x, pre_func=lambda x: x, post_func=lambda x: x, max_items_in_flight=None):
+        if max_items_in_flight is None: max_items_in_flight = self.num_threads * 4
+        assert max_items_in_flight >= 1
+        results = []
+        retire_idx = [0]
+
+        def task_func(prepared, _idx):
+            return process_func(prepared)
+
+        def retire_result():
+            processed, (_prepared, idx) = self.get_result(task_func)
+            results[idx] = processed
+            while retire_idx[0] < len(results) and results[retire_idx[0]] is not None:
+                yield post_func(results[retire_idx[0]])
+                results[retire_idx[0]] = None
+                retire_idx[0] += 1
+
+        for idx, item in enumerate(item_iterator):
+            prepared = pre_func(item)
+            results.append(None)
+            self.add_task(func=task_func, args=(prepared, idx))
+            while retire_idx[0] < idx - max_items_in_flight + 2:
+                for res in retire_result(): yield res
+        while retire_idx[0] < len(results):
+            for res in retire_result(): yield res
+
+#----------------------------------------------------------------------------
+
+def display(tfrecord_dir):
+    print('Loading dataset "%s"' % tfrecord_dir)
+    tflib.init_tf({'gpu_options.allow_growth': True})
+    dset = dataset.TFRecordDataset(tfrecord_dir, max_label_size='full', repeat=False, shuffle_mb=0)
+    tflib.init_uninitialized_vars()
+    import cv2  # pip install opencv-python
+
+    idx = 0
+    while True:
+        try:
+            images, labels = dset.get_minibatch_np(1)
+        except tf.errors.OutOfRangeError:
+            break
+        if idx == 0:
+            print('Displaying images')
+            cv2.namedWindow('dataset_tool')
+            print('Press SPACE or ENTER to advance, ESC to exit')
+        print('\nidx = %-8d\nlabel = %s' % (idx, labels[0].tolist()))
+        cv2.imshow('dataset_tool', images[0].transpose(1, 2, 0)[:, :, ::-1]) # CHW => HWC, RGB => BGR
+        idx += 1
+        if cv2.waitKey() == 27:
+            break
+    print('\nDisplayed %d images.' % idx)
+
+#----------------------------------------------------------------------------
+
+def extract(tfrecord_dir, output_dir):
+    print('Loading dataset "%s"' % tfrecord_dir)
+    tflib.init_tf({'gpu_options.allow_growth': True})
+    dset = dataset.TFRecordDataset(tfrecord_dir, max_label_size=0, repeat=False, shuffle_mb=0)
+    tflib.init_uninitialized_vars()
+
+    print('Extracting images to "%s"' % output_dir)
+    if not os.path.isdir(output_dir):
+        os.makedirs(output_dir)
+    idx = 0
+    while True:
+        if idx % 10 == 0:
+            print('%d\r' % idx, end='', flush=True)
+        try:
+            images, _labels = dset.get_minibatch_np(1)
+        except tf.errors.OutOfRangeError:
+            break
+        if images.shape[1] == 1:
+            img = PIL.Image.fromarray(images[0][0], 'L')
+        else:
+            img = PIL.Image.fromarray(images[0].transpose(1, 2, 0), 'RGB')
+        img.save(os.path.join(output_dir, 'img%08d.png' % idx))
+        idx += 1
+    print('Extracted %d images.' % idx)
+
+#----------------------------------------------------------------------------
+
+def compare(tfrecord_dir_a, tfrecord_dir_b, ignore_labels):
+    max_label_size = 0 if ignore_labels else 'full'
+    print('Loading dataset "%s"' % tfrecord_dir_a)
+    tflib.init_tf({'gpu_options.allow_growth': True})
+    dset_a = dataset.TFRecordDataset(tfrecord_dir_a, max_label_size=max_label_size, repeat=False, shuffle_mb=0)
+    print('Loading dataset "%s"' % tfrecord_dir_b)
+    dset_b = dataset.TFRecordDataset(tfrecord_dir_b, max_label_size=max_label_size, repeat=False, shuffle_mb=0)
+    tflib.init_uninitialized_vars()
+
+    print('Comparing datasets')
+    idx = 0
+    identical_images = 0
+    identical_labels = 0
+    while True:
+        if idx % 100 == 0:
+            print('%d\r' % idx, end='', flush=True)
+        try:
+            images_a, labels_a = dset_a.get_minibatch_np(1)
+        except tf.errors.OutOfRangeError:
+            images_a, labels_a = None, None
+        try:
+            images_b, labels_b = dset_b.get_minibatch_np(1)
+        except tf.errors.OutOfRangeError:
+            images_b, labels_b = None, None
+        if images_a is None or images_b is None:
+            if images_a is not None or images_b is not None:
+                print('Datasets contain different number of images')
+            break
+        if images_a.shape == images_b.shape and np.all(images_a == images_b):
+            identical_images += 1
+        else:
+            print('Image %d is different' % idx)
+        if labels_a.shape == labels_b.shape and np.all(labels_a == labels_b):
+            identical_labels += 1
+        else:
+            print('Label %d is different' % idx)
+        idx += 1
+    print('Identical images: %d / %d' % (identical_images, idx))
+    if not ignore_labels:
+        print('Identical labels: %d / %d' % (identical_labels, idx))
+
+#----------------------------------------------------------------------------
+
+def create_mnist(tfrecord_dir, mnist_dir):
+    print('Loading MNIST from "%s"' % mnist_dir)
+    import gzip
+    with gzip.open(os.path.join(mnist_dir, 'train-images-idx3-ubyte.gz'), 'rb') as file:
+        images = np.frombuffer(file.read(), np.uint8, offset=16)
+    with gzip.open(os.path.join(mnist_dir, 'train-labels-idx1-ubyte.gz'), 'rb') as file:
+        labels = np.frombuffer(file.read(), np.uint8, offset=8)
+    images = images.reshape(-1, 1, 28, 28)
+    images = np.pad(images, [(0,0), (0,0), (2,2), (2,2)], 'constant', constant_values=0)
+    assert images.shape == (60000, 1, 32, 32) and images.dtype == np.uint8
+    assert labels.shape == (60000,) and labels.dtype == np.uint8
+    assert np.min(images) == 0 and np.max(images) == 255
+    assert np.min(labels) == 0 and np.max(labels) == 9
+    onehot = np.zeros((labels.size, np.max(labels) + 1), dtype=np.float32)
+    onehot[np.arange(labels.size), labels] = 1.0
+
+    with TFRecordExporter(tfrecord_dir, images.shape[0]) as tfr:
+        order = tfr.choose_shuffled_order()
+        for idx in range(order.size):
+            tfr.add_image(images[order[idx]])
+        tfr.add_labels(onehot[order])
+
+#----------------------------------------------------------------------------
+
+def create_mnistrgb(tfrecord_dir, mnist_dir, num_images=1000000, random_seed=123):
+    print('Loading MNIST from "%s"' % mnist_dir)
+    import gzip
+    with gzip.open(os.path.join(mnist_dir, 'train-images-idx3-ubyte.gz'), 'rb') as file:
+        images = np.frombuffer(file.read(), np.uint8, offset=16)
+    images = images.reshape(-1, 28, 28)
+    images = np.pad(images, [(0,0), (2,2), (2,2)], 'constant', constant_values=0)
+    assert images.shape == (60000, 32, 32) and images.dtype == np.uint8
+    assert np.min(images) == 0 and np.max(images) == 255
+
+    with TFRecordExporter(tfrecord_dir, num_images) as tfr:
+        rnd = np.random.RandomState(random_seed)
+        for _idx in range(num_images):
+            tfr.add_image(images[rnd.randint(images.shape[0], size=3)])
+
+#----------------------------------------------------------------------------
+
+def create_cifar10(tfrecord_dir, cifar10_dir):
+    print('Loading CIFAR-10 from "%s"' % cifar10_dir)
+    import pickle
+    images = []
+    labels = []
+    for batch in range(1, 6):
+        with open(os.path.join(cifar10_dir, 'data_batch_%d' % batch), 'rb') as file:
+            data = pickle.load(file, encoding='latin1')
+        images.append(data['data'].reshape(-1, 3, 32, 32))
+        labels.append(data['labels'])
+    images = np.concatenate(images)
+    labels = np.concatenate(labels)
+    assert images.shape == (50000, 3, 32, 32) and images.dtype == np.uint8
+    assert labels.shape == (50000,) and labels.dtype == np.int32
+    assert np.min(images) == 0 and np.max(images) == 255
+    assert np.min(labels) == 0 and np.max(labels) == 9
+    onehot = np.zeros((labels.size, np.max(labels) + 1), dtype=np.float32)
+    onehot[np.arange(labels.size), labels] = 1.0
+
+    with TFRecordExporter(tfrecord_dir, images.shape[0]) as tfr:
+        order = tfr.choose_shuffled_order()
+        for idx in range(order.size):
+            tfr.add_image(images[order[idx]])
+        tfr.add_labels(onehot[order])
+
+#----------------------------------------------------------------------------
+
+def create_cifar100(tfrecord_dir, cifar100_dir):
+    print('Loading CIFAR-100 from "%s"' % cifar100_dir)
+    import pickle
+    with open(os.path.join(cifar100_dir, 'train'), 'rb') as file:
+        data = pickle.load(file, encoding='latin1')
+    images = data['data'].reshape(-1, 3, 32, 32)
+    labels = np.array(data['fine_labels'])
+    assert images.shape == (50000, 3, 32, 32) and images.dtype == np.uint8
+    assert labels.shape == (50000,) and labels.dtype == np.int32
+    assert np.min(images) == 0 and np.max(images) == 255
+    assert np.min(labels) == 0 and np.max(labels) == 99
+    onehot = np.zeros((labels.size, np.max(labels) + 1), dtype=np.float32)
+    onehot[np.arange(labels.size), labels] = 1.0
+
+    with TFRecordExporter(tfrecord_dir, images.shape[0]) as tfr:
+        order = tfr.choose_shuffled_order()
+        for idx in range(order.size):
+            tfr.add_image(images[order[idx]])
+        tfr.add_labels(onehot[order])
+
+#----------------------------------------------------------------------------
+
+def create_svhn(tfrecord_dir, svhn_dir):
+    print('Loading SVHN from "%s"' % svhn_dir)
+    import pickle
+    images = []
+    labels = []
+    for batch in range(1, 4):
+        with open(os.path.join(svhn_dir, 'train_%d.pkl' % batch), 'rb') as file:
+            data = pickle.load(file, encoding='latin1')
+        images.append(data[0])
+        labels.append(data[1])
+    images = np.concatenate(images)
+    labels = np.concatenate(labels)
+    assert images.shape == (73257, 3, 32, 32) and images.dtype == np.uint8
+    assert labels.shape == (73257,) and labels.dtype == np.uint8
+    assert np.min(images) == 0 and np.max(images) == 255
+    assert np.min(labels) == 0 and np.max(labels) == 9
+    onehot = np.zeros((labels.size, np.max(labels) + 1), dtype=np.float32)
+    onehot[np.arange(labels.size), labels] = 1.0
+
+    with TFRecordExporter(tfrecord_dir, images.shape[0]) as tfr:
+        order = tfr.choose_shuffled_order()
+        for idx in range(order.size):
+            tfr.add_image(images[order[idx]])
+        tfr.add_labels(onehot[order])
+
+#----------------------------------------------------------------------------
+
+def create_lsun(tfrecord_dir, lmdb_dir, resolution=256, max_images=None):
+    print('Loading LSUN dataset from "%s"' % lmdb_dir)
+    import lmdb # pip install lmdb # pylint: disable=import-error
+    import cv2 # pip install opencv-python
+    import io
+    with lmdb.open(lmdb_dir, readonly=True).begin(write=False) as txn:
+        total_images = txn.stat()['entries'] # pylint: disable=no-value-for-parameter
+        if max_images is None:
+            max_images = total_images
+        with TFRecordExporter(tfrecord_dir, max_images) as tfr:
+            for _idx, (_key, value) in enumerate(txn.cursor()):
+                try:
+                    try:
+                        img = cv2.imdecode(np.fromstring(value, dtype=np.uint8), 1)
+                        if img is None:
+                            raise IOError('cv2.imdecode failed')
+                        img = img[:, :, ::-1] # BGR => RGB
+                    except IOError:
+                        img = np.asarray(PIL.Image.open(io.BytesIO(value)))
+                    crop = np.min(img.shape[:2])
+                    img = img[(img.shape[0] - crop) // 2 : (img.shape[0] + crop) // 2, (img.shape[1] - crop) // 2 : (img.shape[1] + crop) // 2]
+                    img = PIL.Image.fromarray(img, 'RGB')
+                    img = img.resize((resolution, resolution), PIL.Image.ANTIALIAS)
+                    img = np.asarray(img)
+                    img = img.transpose([2, 0, 1]) # HWC => CHW
+                    tfr.add_image(img)
+                except:
+                    print(sys.exc_info()[1])
+                if tfr.cur_images == max_images:
+                    break
+
+#----------------------------------------------------------------------------
+
+def create_lsun_wide(tfrecord_dir, lmdb_dir, width=512, height=384, max_images=None):
+    assert width == 2 ** int(np.round(np.log2(width)))
+    assert height <= width
+    print('Loading LSUN dataset from "%s"' % lmdb_dir)
+    import lmdb # pip install lmdb # pylint: disable=import-error
+    import cv2 # pip install opencv-python
+    import io
+    with lmdb.open(lmdb_dir, readonly=True).begin(write=False) as txn:
+        total_images = txn.stat()['entries'] # pylint: disable=no-value-for-parameter
+        if max_images is None:
+            max_images = total_images
+        with TFRecordExporter(tfrecord_dir, max_images, print_progress=False) as tfr:
+            for idx, (_key, value) in enumerate(txn.cursor()):
+                try:
+                    try:
+                        img = cv2.imdecode(np.fromstring(value, dtype=np.uint8), 1)
+                        if img is None:
+                            raise IOError('cv2.imdecode failed')
+                        img = img[:, :, ::-1] # BGR => RGB
+                    except IOError:
+                        img = np.asarray(PIL.Image.open(io.BytesIO(value)))
+
+                    ch = int(np.round(width * img.shape[0] / img.shape[1]))
+                    if img.shape[1] < width or ch < height:
+                        continue
+
+                    img = img[(img.shape[0] - ch) // 2 : (img.shape[0] + ch) // 2]
+                    img = PIL.Image.fromarray(img, 'RGB')
+                    img = img.resize((width, height), PIL.Image.ANTIALIAS)
+                    img = np.asarray(img)
+                    img = img.transpose([2, 0, 1]) # HWC => CHW
+
+                    canvas = np.zeros([3, width, width], dtype=np.uint8)
+                    canvas[:, (width - height) // 2 : (width + height) // 2] = img
+                    tfr.add_image(canvas)
+                    print('\r%d / %d => %d ' % (idx + 1, total_images, tfr.cur_images), end='')
+
+                except:
+                    print(sys.exc_info()[1])
+                if tfr.cur_images == max_images:
+                    break
+    print()
+
+#----------------------------------------------------------------------------
+
+def create_celeba(tfrecord_dir, celeba_dir, cx=89, cy=121):
+    print('Loading CelebA from "%s"' % celeba_dir)
+    glob_pattern = os.path.join(celeba_dir, 'img_align_celeba_png', '*.png')
+    image_filenames = sorted(glob.glob(glob_pattern))
+    expected_images = 202599
+    if len(image_filenames) != expected_images:
+        error('Expected to find %d images' % expected_images)
+
+    with TFRecordExporter(tfrecord_dir, len(image_filenames)) as tfr:
+        order = tfr.choose_shuffled_order()
+        for idx in range(order.size):
+            img = np.asarray(PIL.Image.open(image_filenames[order[idx]]))
+            assert img.shape == (218, 178, 3)
+            img = img[cy - 64 : cy + 64, cx - 64 : cx + 64]
+            img = img.transpose(2, 0, 1) # HWC => CHW
+            tfr.add_image(img)
+
+#----------------------------------------------------------------------------
+
+def create_from_images(tfrecord_dir, image_dir, shuffle):
+    print('Loading images from "%s"' % image_dir)
+    image_filenames = sorted(glob.glob(os.path.join(image_dir, '*')))
+    if len(image_filenames) == 0:
+        error('No input images found')
+
+    img = np.asarray(PIL.Image.open(image_filenames[0]))
+    resolution = img.shape[0]
+    channels = img.shape[2] if img.ndim == 3 else 1
+    if img.shape[1] != resolution:
+        error('Input images must have the same width and height')
+    if resolution != 2 ** int(np.floor(np.log2(resolution))):
+        error('Input image resolution must be a power-of-two')
+    if channels not in [1, 3]:
+        error('Input images must be stored as RGB or grayscale')
+
+    with TFRecordExporter(tfrecord_dir, len(image_filenames)) as tfr:
+        order = tfr.choose_shuffled_order() if shuffle else np.arange(len(image_filenames))
+        for idx in range(order.size):
+            img = np.asarray(PIL.Image.open(image_filenames[order[idx]]))
+            if channels == 1:
+                img = img[np.newaxis, :, :] # HW => CHW
+            else:
+                img = img.transpose([2, 0, 1]) # HWC => CHW
+            tfr.add_image(img)
+
+#----------------------------------------------------------------------------
+
+def create_from_hdf5(tfrecord_dir, hdf5_filename, shuffle):
+    print('Loading HDF5 archive from "%s"' % hdf5_filename)
+    import h5py # conda install h5py
+    with h5py.File(hdf5_filename, 'r') as hdf5_file:
+        hdf5_data = max([value for key, value in hdf5_file.items() if key.startswith('data')], key=lambda lod: lod.shape[3])
+        with TFRecordExporter(tfrecord_dir, hdf5_data.shape[0]) as tfr:
+            order = tfr.choose_shuffled_order() if shuffle else np.arange(hdf5_data.shape[0])
+            for idx in range(order.size):
+                tfr.add_image(hdf5_data[order[idx]])
+            npy_filename = os.path.splitext(hdf5_filename)[0] + '-labels.npy'
+            if os.path.isfile(npy_filename):
+                tfr.add_labels(np.load(npy_filename)[order])
+
+#----------------------------------------------------------------------------
+
+def execute_cmdline(argv):
+    prog = argv[0]
+    parser = argparse.ArgumentParser(
+        prog        = prog,
+        description = 'Tool for creating multi-resolution TFRecords datasets for StyleGAN and ProGAN.',
+        epilog      = 'Type "%s <command> -h" for more information.' % prog)
+
+    subparsers = parser.add_subparsers(dest='command')
+    subparsers.required = True
+    def add_command(cmd, desc, example=None):
+        epilog = 'Example: %s %s' % (prog, example) if example is not None else None
+        return subparsers.add_parser(cmd, description=desc, help=desc, epilog=epilog)
+
+    p = add_command(    'display',          'Display images in dataset.',
+                                            'display datasets/mnist')
+    p.add_argument(     'tfrecord_dir',     help='Directory containing dataset')
+
+    p = add_command(    'extract',          'Extract images from dataset.',
+                                            'extract datasets/mnist mnist-images')
+    p.add_argument(     'tfrecord_dir',     help='Directory containing dataset')
+    p.add_argument(     'output_dir',       help='Directory to extract the images into')
+
+    p = add_command(    'compare',          'Compare two datasets.',
+                                            'compare datasets/mydataset datasets/mnist')
+    p.add_argument(     'tfrecord_dir_a',   help='Directory containing first dataset')
+    p.add_argument(     'tfrecord_dir_b',   help='Directory containing second dataset')
+    p.add_argument(     '--ignore_labels',  help='Ignore labels (default: 0)', type=int, default=0)
+
+    p = add_command(    'create_mnist',     'Create dataset for MNIST.',
+                                            'create_mnist datasets/mnist ~/downloads/mnist')
+    p.add_argument(     'tfrecord_dir',     help='New dataset directory to be created')
+    p.add_argument(     'mnist_dir',        help='Directory containing MNIST')
+
+    p = add_command(    'create_mnistrgb',  'Create dataset for MNIST-RGB.',
+                                            'create_mnistrgb datasets/mnistrgb ~/downloads/mnist')
+    p.add_argument(     'tfrecord_dir',     help='New dataset directory to be created')
+    p.add_argument(     'mnist_dir',        help='Directory containing MNIST')
+    p.add_argument(     '--num_images',     help='Number of composite images to create (default: 1000000)', type=int, default=1000000)
+    p.add_argument(     '--random_seed',    help='Random seed (default: 123)', type=int, default=123)
+
+    p = add_command(    'create_cifar10',   'Create dataset for CIFAR-10.',
+                                            'create_cifar10 datasets/cifar10 ~/downloads/cifar10')
+    p.add_argument(     'tfrecord_dir',     help='New dataset directory to be created')
+    p.add_argument(     'cifar10_dir',      help='Directory containing CIFAR-10')
+
+    p = add_command(    'create_cifar100',  'Create dataset for CIFAR-100.',
+                                            'create_cifar100 datasets/cifar100 ~/downloads/cifar100')
+    p.add_argument(     'tfrecord_dir',     help='New dataset directory to be created')
+    p.add_argument(     'cifar100_dir',     help='Directory containing CIFAR-100')
+
+    p = add_command(    'create_svhn',      'Create dataset for SVHN.',
+                                            'create_svhn datasets/svhn ~/downloads/svhn')
+    p.add_argument(     'tfrecord_dir',     help='New dataset directory to be created')
+    p.add_argument(     'svhn_dir',         help='Directory containing SVHN')
+
+    p = add_command(    'create_lsun',      'Create dataset for single LSUN category.',
+                                            'create_lsun datasets/lsun-car-100k ~/downloads/lsun/car_lmdb --resolution 256 --max_images 100000')
+    p.add_argument(     'tfrecord_dir',     help='New dataset directory to be created')
+    p.add_argument(     'lmdb_dir',         help='Directory containing LMDB database')
+    p.add_argument(     '--resolution',     help='Output resolution (default: 256)', type=int, default=256)
+    p.add_argument(     '--max_images',     help='Maximum number of images (default: none)', type=int, default=None)
+
+    p = add_command(    'create_lsun_wide', 'Create LSUN dataset with non-square aspect ratio.',
+                                            'create_lsun_wide datasets/lsun-car-512x384 ~/downloads/lsun/car_lmdb --width 512 --height 384')
+    p.add_argument(     'tfrecord_dir',     help='New dataset directory to be created')
+    p.add_argument(     'lmdb_dir',         help='Directory containing LMDB database')
+    p.add_argument(     '--width',          help='Output width (default: 512)', type=int, default=512)
+    p.add_argument(     '--height',         help='Output height (default: 384)', type=int, default=384)
+    p.add_argument(     '--max_images',     help='Maximum number of images (default: none)', type=int, default=None)
+
+    p = add_command(    'create_celeba',    'Create dataset for CelebA.',
+                                            'create_celeba datasets/celeba ~/downloads/celeba')
+    p.add_argument(     'tfrecord_dir',     help='New dataset directory to be created')
+    p.add_argument(     'celeba_dir',       help='Directory containing CelebA')
+    p.add_argument(     '--cx',             help='Center X coordinate (default: 89)', type=int, default=89)
+    p.add_argument(     '--cy',             help='Center Y coordinate (default: 121)', type=int, default=121)
+
+    p = add_command(    'create_from_images', 'Create dataset from a directory full of images.',
+                                            'create_from_images datasets/mydataset myimagedir')
+    p.add_argument(     'tfrecord_dir',     help='New dataset directory to be created')
+    p.add_argument(     'image_dir',        help='Directory containing the images')
+    p.add_argument(     '--shuffle',        help='Randomize image order (default: 1)', type=int, default=1)
+
+    p = add_command(    'create_from_hdf5', 'Create dataset from legacy HDF5 archive.',
+                                            'create_from_hdf5 datasets/celebahq ~/downloads/celeba-hq-1024x1024.h5')
+    p.add_argument(     'tfrecord_dir',     help='New dataset directory to be created')
+    p.add_argument(     'hdf5_filename',    help='HDF5 archive containing the images')
+    p.add_argument(     '--shuffle',        help='Randomize image order (default: 1)', type=int, default=1)
+
+    args = parser.parse_args(argv[1:] if len(argv) > 1 else ['-h'])
+    func = globals()[args.command]
+    del args.command
+    func(**vars(args))
+
+#----------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    execute_cmdline(sys.argv)
+
+#----------------------------------------------------------------------------

models/stylegan/stylegan_tf/dnnlib/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial
+# 4.0 International License. To view a copy of this license, visit
+# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
+# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
+
+from . import submission
+
+from .submission.run_context import RunContext
+
+from .submission.submit import SubmitTarget
+from .submission.submit import PathType
+from .submission.submit import SubmitConfig
+from .submission.submit import get_path_from_template
+from .submission.submit import submit_run
+
+from .util import EasyDict
+
+submit_config: SubmitConfig = None # Package level variable for SubmitConfig which is only valid when inside the run function.