Upload 18 files

annotated_directions.py

@@ -0,0 +1,171 @@
+annotated_directions = {
+    'stylegan2_ffhq1024': {
+        # Directions used in paper with a single decomposition:
+        'big_eyes': {
+            'parameters': [7, 6, 30],  # used in main paper
+            'layer': 5,
+            'ranks': [512, 8],
+            'checkpoints_path': [
+                './checkpoints/Us-name_stylegan2_ffhq1024-layer_5-rank_8.npy',
+                './checkpoints/Uc-name_stylegan2_ffhq1024-layer_5-rank_512.npy',
+            ],
+        },
+        'long_nose': {
+            'parameters': [5, 82, 30],  # used in main paper
+            'layer': 5,
+            'ranks': [512, 8],
+            'checkpoints_path': [
+                './checkpoints/Us-name_stylegan2_ffhq1024-layer_5-rank_8.npy',
+                './checkpoints/Uc-name_stylegan2_ffhq1024-layer_5-rank_512.npy',
+            ],
+        },
+        'smile': {
+            'parameters': [4, 46, -30],  # used in sup. material
+            'layer': 5,
+            'ranks': [512, 8],
+            'checkpoints_path': [
+                './checkpoints/Us-name_stylegan2_ffhq1024-layer_5-rank_8.npy',
+                './checkpoints/Uc-name_stylegan2_ffhq1024-layer_5-rank_512.npy',
+            ],
+        },
+        'open_mouth': {
+            'parameters': [4, 39, 30],  # used in sup. material
+            'layer': 5,
+            'ranks': [512, 8],
+            'checkpoints_path': [
+                './checkpoints/Us-name_stylegan2_ffhq1024-layer_5-rank_8.npy',
+                './checkpoints/Uc-name_stylegan2_ffhq1024-layer_5-rank_512.npy',
+            ],
+        },
+
+        # Additional directions
+        'big_eyeballs': {
+            'parameters': [8, 27, 100],
+            'layer': 6,
+            'ranks': [512, 16],
+            'checkpoints_path': [
+                './checkpoints/Us-name_stylegan2_ffhq1024-layer_6-rank_16.npy',
+                './checkpoints/Uc-name_stylegan2_ffhq1024-layer_6-rank_512.npy',
+            ],
+        },
+        'wide_nose': {
+            'parameters': [15, 13, 100],
+            'layer': 6,
+            'ranks': [512, 16],
+            'checkpoints_path': [
+                './checkpoints/Us-name_stylegan2_ffhq1024-layer_6-rank_16.npy',
+                './checkpoints/Uc-name_stylegan2_ffhq1024-layer_6-rank_512.npy',
+            ],
+        },
+        'glance_left': {
+            'parameters': [8, 281, 50],
+            'layer': 6,
+            'ranks': [512, 16],
+            'checkpoints_path': [
+                './checkpoints/Us-name_stylegan2_ffhq1024-layer_6-rank_16.npy',
+                './checkpoints/Uc-name_stylegan2_ffhq1024-layer_6-rank_512.npy',
+            ],
+        },
+        'glance_right': {
+            'parameters': [8, 281, -70],
+            'layer': 6,
+            'ranks': [512, 16],
+            'checkpoints_path': [
+                './checkpoints/Us-name_stylegan2_ffhq1024-layer_6-rank_16.npy',
+                './checkpoints/Uc-name_stylegan2_ffhq1024-layer_6-rank_512.npy',
+            ],
+        },
+        'bald_forehead': {
+            'parameters': [3, 25, 100],
+            'layer': 6,
+            'ranks': [512, 16],
+            'checkpoints_path': [
+                './checkpoints/Us-name_stylegan2_ffhq1024-layer_6-rank_16.npy',
+                './checkpoints/Uc-name_stylegan2_ffhq1024-layer_6-rank_512.npy',
+            ],
+        },
+        'light_eyebrows': {
+            'parameters': [8, 4, 30],
+            'layer': 7,
+            'ranks': [512, 16],
+            'checkpoints_path': [
+                './checkpoints/Us-name_stylegan2_ffhq1024-layer_7-rank_16.npy',
+                './checkpoints/Uc-name_stylegan2_ffhq1024-layer_7-rank_512.npy',
+            ],
+        },
+        'dark_eyebrows': {
+            'parameters': [8, 9, 30],
+            'layer': 7,
+            'ranks': [512, 16],
+            'checkpoints_path': [
+                './checkpoints/Us-name_stylegan2_ffhq1024-layer_7-rank_16.npy',
+                './checkpoints/Uc-name_stylegan2_ffhq1024-layer_7-rank_512.npy',
+            ],
+        },
+        'no_eyebrows': {
+            'parameters': [8, 4, 50],
+            'layer': 7,
+            'ranks': [512, 16],
+            'checkpoints_path': [
+                './checkpoints/Us-name_stylegan2_ffhq1024-layer_7-rank_16.npy',
+                './checkpoints/Uc-name_stylegan2_ffhq1024-layer_7-rank_512.npy',
+            ],
+        },
+        'dark_eyes': {
+            'parameters': [11, 176, 50],
+            'layer': 7,
+            'ranks': [512, 16],
+            'checkpoints_path': [
+                './checkpoints/Us-name_stylegan2_ffhq1024-layer_7-rank_16.npy',
+                './checkpoints/Uc-name_stylegan2_ffhq1024-layer_7-rank_512.npy',
+            ],
+        },
+        'red_eyes': {
+            'parameters': [11, 109, 60],
+            'layer': 7,
+            'ranks': [512, 16],
+            'checkpoints_path': [
+                './checkpoints/Us-name_stylegan2_ffhq1024-layer_7-rank_16.npy',
+                './checkpoints/Uc-name_stylegan2_ffhq1024-layer_7-rank_512.npy',
+            ],
+        },
+        'eyes_short': {
+            'parameters': [11, 262, 70],
+            'layer': 7,
+            'ranks': [512, 16],
+            'checkpoints_path': [
+                './checkpoints/Us-name_stylegan2_ffhq1024-layer_7-rank_16.npy',
+                './checkpoints/Uc-name_stylegan2_ffhq1024-layer_7-rank_512.npy',
+            ],
+        },
+        'eyes_open': {
+            'parameters': [11, 28, 50],
+            'layer': 7,
+            'ranks': [512, 16],
+            'checkpoints_path': [
+                './checkpoints/Us-name_stylegan2_ffhq1024-layer_7-rank_16.npy',
+                './checkpoints/Uc-name_stylegan2_ffhq1024-layer_7-rank_512.npy',
+            ],
+        },
+        'eyes_close': {
+            'parameters': [11, 398, 80],
+            'layer': 7,
+            'ranks': [512, 16],
+            'checkpoints_path': [
+                './checkpoints/Us-name_stylegan2_ffhq1024-layer_7-rank_16.npy',
+                './checkpoints/Uc-name_stylegan2_ffhq1024-layer_7-rank_512.npy',
+            ],
+        },
+        'no_eyes': {
+            'parameters': [11, 0, -200],
+            'layer': 7,
+            'ranks': [512, 16],
+            'checkpoints_path': [
+                './checkpoints/Us-name_stylegan2_ffhq1024-layer_7-rank_16.npy',
+                './checkpoints/Uc-name_stylegan2_ffhq1024-layer_7-rank_512.npy',
+            ],
+        },
+
+    },
+
+}

app.py

@@ -0,0 +1,44 @@
+import numpy as np
+import gradio as gr
+import torch
+from PIL import Image
+from model import Model as Model
+from annotated_directions import annotated_directions
+device = torch.device('cpu')
+
+torch.set_grad_enabled(False)
+model_name = "stylegan2_ffhq1024"
+
+directions = list(annotated_directions[model_name].keys())
+
+
+def inference(seed, direction):
+    layer = annotated_directions[model_name][direction]['layer']
+    M = Model(model_name, trunc_psi=1.0, device=device, layer=layer)
+    M.ranks = annotated_directions[model_name][direction]['ranks']
+
+    # load the checkpoint
+    try:
+        M.Us = torch.Tensor(np.load(annotated_directions[model_name][direction]['checkpoints_path'][0])).to(device)
+        M.Uc = torch.Tensor(np.load(annotated_directions[model_name][direction]['checkpoints_path'][1])).to(device)
+    except KeyError:
+        raise KeyError('ERROR: No directions specified in ./annotated_directions.py for this model')
+
+    part, appearance, lam = annotated_directions[model_name][direction]['parameters']
+
+    Z, image, image2, part_img = M.edit_at_layer([[part]], [appearance], [lam], t=seed, Uc=M.Uc, Us=M.Us, noise=None)
+
+    dif = np.tile(((np.mean((image - image2)**2, -1)))[:,:,None], [1,1,3]).astype(np.uint8)
+
+    return Image.fromarray(np.concatenate([image, image2, dif], 1))
+
+
+demo = gr.Interface(
+    fn=inference,
+    inputs=[gr.Slider(0, 1000, value=64), gr.Dropdown(directions, value='no_eyebrows')],
+    outputs=[gr.Image(type="pil", label="original | edited | mean-squared difference")],
+    title="PandA (ICLR'23) - FFHQ edit zoo",
+    description="Provides a quick interface to manipulate pre-annotated directions with pre-trained global parts and appearances factors. Note that we use the free CPU tier, so synthesis takes about 10 seconds.",
+    article="Check out the full demo and paper at: <a href='https://github.com/james-oldfield/PandA'>https://github.com/james-oldfield/PandA</a>"
+)
+demo.launch()

checkpoints/Uc-name_stylegan2_afhqdog512-layer_8-rank_512.npy

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

checkpoints/Uc-name_stylegan2_ffhq1024-layer_5-rank_512.npy

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

checkpoints/Uc-name_stylegan2_ffhq1024-layer_6-rank_512.npy

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

checkpoints/Uc-name_stylegan2_ffhq1024-layer_7-rank_512.npy

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

checkpoints/Us-name_stylegan2_afhqdog512-layer_8-rank_8.npy

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

checkpoints/Us-name_stylegan2_ffhq1024-layer_5-rank_8.npy

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

checkpoints/Us-name_stylegan2_ffhq1024-layer_6-rank_16.npy

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

checkpoints/Us-name_stylegan2_ffhq1024-layer_7-rank_16.npy

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

environment.yml

@@ -0,0 +1,25 @@
+name: PandA
+channels:
+  - pytorch
+  - defaults
+dependencies:
+  - python=3.8.5
+  - pip=20.3
+  - cudatoolkit=11.3
+  - pytorch=1.11.0
+  - torchvision=0.12.0
+  - numpy=1.19.2
+  - pip:
+    - tensorly==0.7.0
+    - opencv-python==4.1.2.30
+    - boto3==1.21.8
+    - botocore==1.24.8
+    - imageio==2.14.1
+    - matplotlib==3.5.1
+    - nltk==3.7
+    - numpy
+    - Pillow==9.0.1
+    - requests==2.27.0
+    - bs4
+    - ipykernel
+    - jupyterlab

model.py

@@ -0,0 +1,510 @@
+from networks.load_generator import load_generator
+from networks.genforce.utils.visualizer import postprocess_image as postprocess
+from networks.biggan import one_hot_from_names, truncated_noise_sample
+from networks.stylegan3.load_stylegan3 import make_transform
+
+from matplotlib import pyplot as plt
+
+from utils import plot_masks, plot_colours, mapRange
+
+import torch
+import numpy as np
+from PIL import Image
+import tensorly as tl
+tl.set_backend('pytorch')
+
+
+class Model():
+    def __init__(self, model_name, t=0, layer=5, trunc_psi=1.0, trunc_layers=18, device='cuda', biggan_classes=['fox']):
+        """
+        Instantiate the model for decomposition and/or local image editing.
+
+        Parameters
+        ----------
+        model_name : string
+            Name of architecture and dataset--one of the items in ./networks/genforce/models/model_zoo.py.
+        t : int
+            Random seed for the generator (to generate a sample image).
+        layer : int
+            Intermediate layer at which to perform the decomposition.
+        trunc_psi : float
+            Truncation value in [0, 1].
+        trunc_layers : int
+            Number of layers at which to apply truncation.
+        device : string
+            Device to store the tensors on.
+        biggan_classes : list
+            List of strings specifying imagenet classes of interest (e.g. ['alp', 'breakwater']).
+        """
+        self.gan_type = model_name.split('_')[0]
+        self.model_name = model_name
+        self.randomize_noise = False
+        self.device = device
+        self.biggan_classes = biggan_classes
+        self.layer = layer  # layer to decompose
+        self.start = 0 if 'stylegan2' in self.gan_type else 2
+        self.trunc_psi = trunc_psi
+        self.trunc_layers = trunc_layers
+
+        self.generator = load_generator(model_name, device)
+        noise = torch.Tensor(np.random.randn(1, self.generator.z_space_dim)).to(self.device)
+        z, image = self.sample(noise, layer=layer, trunc_psi=trunc_psi, trunc_layers=trunc_layers, verbose=True)
+
+        self.c = z.shape[1]
+        self.s = z.shape[2]
+        self.image = image
+
+    def HOSVD(self, batch_size=10, n_iters=100):
+        """
+        Initialises the appearance basis A. In particular, computes the left-singular vectors of the channel mode's scatter matrix.
+
+        Note: total samples used is batch_size * n_iters
+
+        Parameters
+        ----------
+        batch_size : int
+            Number of activations to sample in a single go.
+        n_iters : int
+            Number of times to sample `batch_size`-many activations.
+        """
+        np.random.seed(0)
+        torch.manual_seed(0)
+
+        with torch.no_grad():
+            Z = torch.zeros((batch_size * n_iters, self.c, self.s, self.s), device=self.device)
+
+            # note: perform in loops to have a larger effective batch size
+            print('Starting loops...')
+            for i in range(n_iters):
+                np.random.seed(i)
+                torch.manual_seed(i)
+                noise = torch.Tensor(np.random.randn(batch_size, self.generator.z_space_dim)).to(self.device)
+                z, _ = self.sample(noise, layer=self.layer, partial=True)
+
+                Z[(batch_size * i):(batch_size * (i + 1))] = z
+
+            Z = Z.view([-1, self.c, self.s**2])
+            print(f'Generated {batch_size * n_iters} gan samples...')
+
+            scat = 0
+            for _, x in enumerate(Z):
+                # mode-3 unfolding in the paper, but in PyTorch channel mode is first.
+                m_unfold = tl.unfold(x, 0)
+                scat += m_unfold @ m_unfold.T
+
+            self.Uc_init, _, _ = np.linalg.svd((scat / len(Z)).cpu().numpy())
+            self.Uc_init = torch.Tensor(self.Uc_init).to(self.device)
+
+            print('... HOSVD done')
+
+    def decompose(self, ranks=[512, 8], lr=1e-8, batch_size=1, its=10000, log_modulo=1000, hosvd_init=True, stochastic=True, n_iters=1, verbose=True):
+        """
+        Performs the decomposition in the paper. In particular, Algorithm 1.,
+        either with a non-fixed batch of samples (stochastic=True), or descends the full gradients.
+
+        Parameters
+        ----------
+        ranks : list
+            List of integers specifying the R_C and R_S, the ranks--i.e. number of parts and appearances respectively.
+        lr : float
+            Learning rate the projected gradient descent.
+        batch_size : int
+            Number of samples in each batch.
+        its : int
+            Total number of iterations.
+        log_modulo : int
+            Parameter used to control how often "training" information is displayed.
+        hosvd_init : bool
+            Initialise appearance factors from HOSVD? (else from random normal).
+        stochastic : bool
+            Sample the batch again each iteration? Else descent full gradients
+        n_iters : int
+            Number of `batch_size`-many samples to take (for full gradient).
+            The total activations are sampled in batches in a loop to enable it to fit in memory.
+        verbose : bool
+            Prints extra information.
+        """
+        self.ranks = ranks
+        np.random.seed(0)
+        torch.manual_seed(0)
+
+        #######################
+        # init from HOSVD, else random normal
+        Uc = self.Uc_init[:, :ranks[0]].detach().clone().to(self.device) if hosvd_init else torch.randn(self.Uc_init.shape[0], ranks[0]).detach().clone().to(self.device) * 0.01
+        Us = torch.Tensor(np.random.uniform(0, 0.01, size=[self.s**2, ranks[1]])).to(self.device)
+        #######################
+
+        print(f'Uc shape: {Uc.shape}, Us shape: {Us.shape}')
+
+        with torch.no_grad():
+            zeros = torch.zeros_like(Us, device=self.device)
+            Us = torch.maximum(Us, zeros)
+
+            # use a fixed batch (i.e. descend the full gradient)
+            if not stochastic:
+                Z = torch.zeros((batch_size * n_iters, self.c, self.s, self.s), device=self.device)
+
+                # note: perform in loops to have a larger effective batch size
+                print(f'Starting loops, total Z shape: {Z.shape}...')
+                for i in range(n_iters):
+                    np.random.seed(i)
+                    torch.manual_seed(i)
+                    noise = torch.Tensor(np.random.randn(batch_size, self.generator.z_space_dim)).to(self.device)
+                    z, _ = self.sample(noise, layer=self.layer, partial=True)
+
+                    Z[(batch_size * i):(batch_size * (i + 1))] = z
+
+            for t in range(its):
+                np.random.seed(t)
+                torch.manual_seed(t)
+
+                # resample the batch, if stochastic
+                if stochastic:
+                    noise = torch.Tensor(np.random.randn(batch_size, self.generator.z_space_dim)).to(self.device)
+                    Z, _ = self.sample(noise, layer=self.layer, partial=True)
+
+                if verbose:
+                    # reconstruct (for visualisation)
+                    coords = tl.tenalg.multi_mode_dot(Z.view(-1, self.c, self.s**2).float(), [Uc.T, Us.T], transpose=False, modes=[1, 2])
+                    Z_rec = tl.tenalg.multi_mode_dot(coords, [Uc, Us], transpose=False, modes=[1, 2])
+
+                    self.rec_loss = torch.mean(torch.norm(Z.view(-1, self.c, self.s**2).float() - Z_rec, p='fro', dim=[1, 2]) ** 2)
+
+                # Update S
+                z = Z.view(-1, self.c, self.s**2).float()
+                Us_g = -4 * (torch.transpose(z,1,2)@Uc@Uc.T@z@Us) + \
+                    2 * (Us@Us.T@torch.transpose(z,1,2)@Uc@Uc.T@Uc@Uc.T@z@Us + torch.transpose(z,1,2)@Uc@Uc.T@Uc@Uc.T@z@Us@Us.T@Us)
+                Us_g = torch.sum(Us_g, 0)
+
+                Us = Us - lr * Us_g
+                # --- projection step ---a
+                Us = torch.maximum(Us, zeros)
+
+                # Update C
+                Uc_g = -4 * (z@Us@Us.T@torch.transpose(z,1,2)@Uc) + \
+                    2 * (Uc@Uc.T@z@Us@Us.T@Us@Us.T@torch.transpose(z,1,2)@Uc + z@Us@Us.T@Us@Us.T@torch.transpose(z,1,2)@Uc@Uc.T@Uc)
+                Uc_g = torch.sum(Uc_g, 0)
+                Uc = Uc - lr * Uc_g
+
+                self.Us = Us
+                self.Uc = Uc
+
+                if t % log_modulo == 0 and verbose:
+                    print(f'ITERATION: {t}')
+                    z, x = self.sample(noise, layer=self.layer, partial=False)
+
+                    # here we display the learnt parts factors and also overlay them over the images to visualise.
+                    plot_masks(Us.T, r=min(ranks[-1], 32), s=self.s)
+                    plt.show()
+                    plot_colours(x, Us.T, r=ranks[-1], s=self.s, seed=-1)
+                    plt.show()
+
+    def decompose_autograd(self, ranks=[512, 8], lr=1e-8, batch_size=1, its=10000, log_modulo=1000, verbose=True, hosvd_init=True):
+        """
+        Performs the same decomposition in the paper, only uses autograd with Adam optimizer (and projected gradient descent).
+
+        Parameters
+        ----------
+        ranks : list
+            List of integers specifying the R_C and R_S, the ranks--i.e. number of parts and appearances respectively.
+        lr : float
+            Learning rate the projected gradient descent.
+        batch_size : int
+            Number of samples in each batch.
+        its : int
+            Total number of iterations.
+        log_modulo : int
+            Parameter used to control how often "training" information is displayed.
+        hosvd_init : bool
+            Initialise appearance factors from HOSVD? (else from random normal).
+        verbose : bool
+            Prints extra information.
+        """
+        self.ranks = ranks
+        np.random.seed(0)
+        torch.manual_seed(0)
+
+        #######################
+        # init from HOSVD, else random normal
+        Uc = torch.nn.Parameter(self.Uc_init[:, :ranks[0]].detach().clone().to(self.device), requires_grad=True) \
+            if hosvd_init else torch.nn.Parameter(torch.randn(self.Uc_init.shape[0], ranks[0]).detach().clone().to(self.device) * 0.01)
+        Us = torch.nn.Parameter(torch.Tensor(np.random.uniform(0, 0.01, size=[self.s**2, ranks[1]])).to(self.device), requires_grad=True)
+        #######################
+        optimizerS = torch.optim.Adam([Us], lr=lr)
+        optimizerC = torch.optim.Adam([Uc], lr=lr)
+
+        print(f'Uc shape: {Uc.shape}, Us shape: {Us.shape}')
+
+        zeros = torch.zeros_like(Us, device=self.device)
+        for t in range(its):
+            np.random.seed(t)
+            torch.manual_seed(t)
+
+            noise = torch.Tensor(np.random.randn(batch_size, self.generator.z_space_dim)).to(self.device)
+            Z, _ = self.sample(noise, layer=self.layer, partial=True)
+
+            # Update S
+            # reconstruct
+            coords = tl.tenalg.multi_mode_dot(Z.view(-1, self.c, self.s**2).float(), [Uc.T, Us.T], transpose=False, modes=[1, 2])
+            Z_rec = tl.tenalg.multi_mode_dot(coords, [Uc, Us], transpose=False, modes=[1, 2])
+
+            rec_loss = torch.mean(torch.norm(Z.view(-1, self.c, self.s**2).float() - Z_rec, p='fro', dim=[1, 2]) ** 2)
+            rec_loss.backward(retain_graph=True)
+
+            optimizerS.step()
+            # --- projection step ---
+            Us.data = torch.maximum(Us.data, zeros)
+            optimizerS.zero_grad()
+            optimizerC.zero_grad()
+
+            # Update C
+            # reconstruct with updated Us
+            coords = tl.tenalg.multi_mode_dot(Z.view(-1, self.c, self.s**2).float(), [Uc.T, Us.T], transpose=False, modes=[1, 2])
+            Z_rec = tl.tenalg.multi_mode_dot(coords, [Uc, Us], transpose=False, modes=[1, 2])
+
+            rec_loss = torch.mean(torch.norm(Z.view(-1, self.c, self.s**2).float() - Z_rec, p='fro', dim=[1, 2]) ** 2)
+            rec_loss.backward()
+            optimizerC.step()
+            optimizerS.zero_grad()
+            optimizerC.zero_grad()
+
+            self.Us = Us
+            self.Uc = Uc
+
+            with torch.no_grad():
+                if t % log_modulo == 0 and verbose:
+                    print(f'Iteration {t} -- rec {rec_loss}')
+
+                    noise = torch.Tensor(np.random.randn(batch_size, self.generator.z_space_dim)).to(self.device)
+                    Z, x = self.sample(noise, layer=self.layer, partial=False)
+
+                    plot_masks(Us.T, r=min(ranks[-1], 32), s=self.s)
+                    plt.show()
+                    plot_colours(x, Us.T, r=ranks[-1], s=self.s, seed=-1)
+                    plt.show()
+
+    def refine(self, Z, image, lr=1e-8, its=1000, log_modulo=250, verbose=True):
+        """
+        Performs the "refinement" step described in the paper, for a given sample Z.
+
+        Parameters
+        ----------
+        Z : torch.Tensor
+            Intermediate activations for target refinement.
+        image : np.array
+            Corresponding image for Z (purely for visualisation purposes).
+        lr : float
+            Learning rate the projected gradient descent.
+        its : int
+            Total number of iterations.
+        log_modulo : int
+            Parameter used to control how often "training" information is displayed.
+        verbose : bool
+            Prints extra information.
+
+        Returns
+        -------
+        UsR : torch.Tensor
+            The refined factors \tilde{P}_i.
+        """
+        np.random.seed(0)
+        torch.manual_seed(0)
+
+        #######################
+        # init from global spatial factors
+        UsR = self.Us.clone()
+        Uc = self.Uc
+        #######################
+
+        zeros = torch.zeros_like(self.Us, device=self.device)
+        for t in range(its):
+            with torch.no_grad():
+                z = Z.view(-1, self.c, self.s**2).float()
+                # descend refinement term's gradient
+                UsR_g = -4 * (torch.transpose(z,1,2)@Uc@Uc.T@z@UsR) + \
+                    2 * (UsR@UsR.T@torch.transpose(z,1,2)@Uc@Uc.T@Uc@Uc.T@z@UsR + torch.transpose(z,1,2)@Uc@Uc.T@Uc@Uc.T@z@UsR@UsR.T@UsR)
+                UsR_g = torch.sum(UsR_g, 0)
+
+                # Update S
+                UsR = UsR - lr * UsR_g
+                # PGD step
+                UsR = torch.maximum(UsR, zeros)
+
+                if ((t + 1) % log_modulo == 0 and verbose):
+                    print(f'iteration {t}')
+
+                    plot_masks(UsR.T, s=self.s, r=min(self.ranks[-1], 16))
+                    plt.show()
+                    plot_colours(image, UsR.T, s=self.s, r=self.ranks[-1], seed=-1, alpha=0.9)
+                    plt.show()
+
+        return UsR
+
+    def edit_at_layer(self, part, appearance, lam, t, Uc, Us, noise=None, b_idx=0):
+        """
+        Performs the "refinement" step described in the paper, for a given sample Z.
+
+        Parameters
+        ----------
+        part : list
+            List of ints containing the part(s) (column of Us) at which to edit.
+        appearance : list
+            List of ints containing the appearance (column of Uc) to apply at the corresponding part(s).
+        lam : list
+            List of ints containing the magnitude for each edit.
+        t : int
+            Random seed to edit
+        Uc : np.array
+            Learnt appearance factors
+        Us : np.array
+            Learnt parts factors
+        noise : np.array
+            If specified, the target latent code itself to edit (i.e. instead of providing than a random seed number).
+        b_idx : int
+            Index of biggan categories to use.
+
+        Returns
+        -------
+        Z : torch.Tensor
+            The intermediate activation at layer self.L
+        image : np.array
+            The original image for sample `t` or from latent code `noise`.
+        image2 : np.array
+            The edited image.
+        part : np.array
+            The part used to edit.
+        """
+        with torch.no_grad():
+            if noise is None:
+                np.random.seed(t)
+                torch.manual_seed(t)
+                noise = torch.Tensor(np.random.randn(1, self.generator.z_space_dim)).to(self.device)
+            else:
+                np.random.seed(0)
+                torch.manual_seed(0)
+
+            direc = 0
+            for i in range(len(appearance)):
+                a = Uc[:, appearance[i]]
+                p = torch.sum(Us[:, part[i]], dim=-1).reshape([self.s, self.s])
+                p = mapRange(p, torch.min(p), torch.max(p), 0.0, 1.0)
+
+                # here, we basically form a rank-1 "tensor", to add to the target sample's activations.
+                # intuitively, the non-zero spatial positions of the part are filled with the appearance vector.
+                direc += lam[i] * tl.tenalg.outer([a, p])
+
+            if self.gan_type in ['stylegan', 'stylegan2']:
+                noise = self.generator.mapping(noise)['w']
+                noise_trunc = self.generator.truncation(noise, trunc_psi=self.trunc_psi, trunc_layers=self.trunc_layers)
+
+                Z = self.generator.synthesis(noise_trunc, start=self.start, stop=self.layer)['x']
+
+                x = self.generator.synthesis(noise_trunc, x=Z, start=self.layer)['image']
+                x_prime = self.generator.synthesis(noise_trunc, x=Z + direc, start=self.layer)['image']
+            elif 'pggan' in self.gan_type:
+                Z = self.generator(noise, start=self.start, stop=self.layer)['x']
+
+                x = self.generator(Z, start=self.layer)['image']
+                x_prime = self.generator(Z + direc, start=self.layer)['image']
+            elif 'biggan' in self.gan_type:
+                print(f'Choosing a {self.biggan_classes[b_idx]}')
+                class_vector = torch.tensor(one_hot_from_names([self.biggan_classes[b_idx]]), device=self.device)
+                noise_vector = torch.tensor(truncated_noise_sample(truncation=self.trunc_psi, batch_size=1, seed=t), device=self.device)
+
+                result = self.generator(noise_vector, class_vector, self.trunc_psi, stop=self.layer)
+                Z, cond_vector = result['z'], result['cond_vector']
+                x = self.generator(Z, class_vector, self.trunc_psi, cond_vector=cond_vector, start=self.layer)['z']
+                x_prime = self.generator(Z + direc, class_vector, self.trunc_psi, cond_vector=cond_vector, start=self.layer)['z']
+            elif 'stylegan3' in self.gan_type:
+                label = torch.zeros([1, 0], device=self.device)
+                Z = self.generator(noise, label, stop=self.layer, truncation_psi=self.trunc_psi, noise_mode='const')
+
+                x = self.generator(noise, label, x=Z, start=self.layer, stop=None, truncation_psi=self.trunc_psi, noise_mode='const')
+                x_prime = self.generator(noise, label, x=Z + direc, start=self.layer, stop=None, truncation_psi=self.trunc_psi, noise_mode='const')
+
+            image = np.array(Image.fromarray(postprocess(x.cpu().numpy())[0]).resize((256, 256)))
+            image2 = np.array(Image.fromarray(postprocess(x_prime.cpu().numpy())[0]).resize((256, 256)))
+
+            part = np.array(Image.fromarray(p.detach().cpu().numpy() * 255).convert('RGB').resize((256, 256), Image.NEAREST))
+            return Z, image, image2, part
+
+    def sample(self, noise, layer=5, partial=False, trunc_psi=1.0, trunc_layers=18, verbose=False):
+        """
+        Samples intermediate feature maps and resulting image the desired generator.
+
+        Parameters
+        ----------
+        noise : np.array
+            (batch_size, z_dim)-dim random standard gaussian noise.
+        layer : int
+            Intermediate layer at which to return intermediate features.
+        partial : bool
+            Perform full forward pass, and return image too? or just intermediate activations at layer number `layer`?
+        trunc_psi : float
+            Truncation value in [0, 1].
+        trunc_layers : int
+            Number of layers at which to apply truncation.
+        biggan_classes : list
+            List of strings specifying imagenet classes of interest (e.g. ['alp', 'breakwater']).
+        verbose : bool
+            Print out additional information?
+
+        Returns
+        -------
+        Z : torch.Tensor
+            The intermediate activations of shape [C, H, W].
+        image : np.array
+            Output RGB image.
+        """
+        with torch.no_grad():
+            if self.gan_type in ['stylegan', 'stylegan2']:
+                noise = self.generator.mapping(noise)['w']
+                noise_trunc = self.generator.truncation(noise, trunc_psi=trunc_psi, trunc_layers=trunc_layers)
+                Z = self.generator.synthesis(noise_trunc, start=self.start, stop=layer)['x']
+                if not partial:
+                    x = self.generator.synthesis(noise_trunc, x=Z, start=layer)['image']
+            elif 'pggan' in self.gan_type:
+                Z = self.generator(noise, start=self.start, stop=layer)['x']
+                if not partial:
+                    x = self.generator(Z, start=layer)['image']
+            elif 'biggan' in self.gan_type:
+                if verbose:
+                    print(f'Using BigGAN class names: {", ".join(self.biggan_classes)}')
+
+                class_vector = torch.tensor(one_hot_from_names(list(np.random.choice(self.biggan_classes, noise.shape[0])), batch_size=noise.shape[0]), device=self.device)
+                noise_vector = torch.tensor(truncated_noise_sample(truncation=self.trunc_psi, batch_size=noise.shape[0]), device=self.device)
+
+                result = self.generator(noise_vector, class_vector, self.trunc_psi, stop=layer)
+                Z = result['z']
+                cond_vector = result['cond_vector']
+
+                if not partial:
+                    x = self.generator(Z, class_vector, self.trunc_psi, cond_vector=cond_vector, start=layer)['z']
+            elif 'stylegan3' in self.gan_type:
+                label = torch.zeros([noise.shape[0], 0], device=self.device)
+                if hasattr(self.generator.synthesis, 'input'):
+                    m = np.linalg.inv(make_transform((0,0), 0))
+                    self.generator.synthesis.input.transform.copy_(torch.from_numpy(m))
+
+                Z = self.generator(noise, label, x=None, start=0, stop=layer, truncation_psi=trunc_psi, noise_mode='const')
+                if not partial:
+                    x = self.generator(noise, label, x=Z, start=layer, stop=None, truncation_psi=trunc_psi, noise_mode='const')
+
+            if verbose:
+                print(f'-- Partial Z shape at layer {layer}: {Z.shape}')
+
+            if partial:
+                return Z, None
+            else:
+                image = postprocess(x.detach().cpu().numpy())
+                image = np.array(Image.fromarray(image[0]).resize((256, 256)))
+                return Z, image
+
+    def save(self):
+        Uc_path = f'./checkpoints/Uc-name_{self.model_name}-layer_{self.layer}-rank_{self.ranks[0]}.npy'
+        Us_path = f'./checkpoints/Us-name_{self.model_name}-layer_{self.layer}-rank_{self.ranks[1]}.npy'
+
+        np.save(Us_path, self.Us.detach().cpu().numpy())
+        np.save(Uc_path, self.Uc.detach().cpu().numpy())
+
+        print(f'Saved factors to {Us_path}, {Uc_path}')

readme.md

@@ -0,0 +1,82 @@
+# PandA: Unsupervised Learning of Parts and Appearances in the Feature Maps of GANs
+
+## [ [paper](https://openreview.net/pdf?id=iUdSB2kK9GY) | [project page](http://eecs.qmul.ac.uk/~jo001/PandA/) | [video](https://www.youtube.com/watch?v=1KY055goKP0) | [edit zoo](https://colab.research.google.com/github/james-oldfield/PandA/blob/main/ffhq-edit-zoo.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/james-oldfield/PandA/blob/main/demo.ipynb) ]
+
+![main.jpg](./images/main.jpg)
+
+> **PandA: Unsupervised Learning of Parts and Appearances in the Feature Maps of GANs**<br>
+> James Oldfield, Christos Tzelepis, Yannis Panagakis, Mihalis A. Nicolaou, and Ioannis Patras<br>
+> *International Conference on Learning Representations (ICLR)*, 2023 <br>
+> https://arxiv.org/abs/2206.00048 <br>
+>
+> **Abstract**: Recent advances in the understanding of Generative Adversarial Networks (GANs) have led to remarkable progress in visual editing and synthesis tasks, capitalizing on the rich semantics that are embedded in the latent spaces of pre-trained GANs. However, existing methods are often tailored to specific GAN architectures and are limited to either discovering global semantic directions that do not facilitate localized control, or require some form of supervision through manually provided regions or segmentation masks. In this light, we present an architecture-agnostic approach that jointly discovers factors representing spatial parts and their appearances in an entirely unsupervised fashion. These factors are obtained by applying a semi-nonnegative tensor factorization on the feature maps, which in turn enables context-aware local image editing with pixel-level control. In addition, we show that the discovered appearance factors correspond to saliency maps that localize concepts of interest, without using any labels. Experiments on a wide range of GAN architectures and datasets show that, in comparison to the state of the art, our method is far more efficient in terms of training time and, most importantly, provides much more accurate localized control.
+
+![cat-gif](./images/cat-eye-control.gif)
+> An example of using our learnt appearances and semantic parts for local image editing.
+
+## Experiments
+
+We provide a number of notebooks to reproduce the experiments in the paper and to explore the model. Please see the following notebooks:
+
+# [`./demo.ipynb`](./demo.ipynb)
+
+This notebook contains the code to learn the parts and appearance factors at a target layer in a target GAN. Contains code for local image editing using the learnt parts, and provides code for refining the parts factors.
+
+| Local image editing (at the learnt semantic parts) | |
+| :-- | :-- |
+| ![image](./images/l8-t645-Rs16-Rc512-rTrue-lam[100]-p[6]-start-end.gif) | ![image](./images/l8-t16-Rs8-Rc512-rTrue-lam-150-p2-start-end.gif)
+
+# [`./localize-concepts.ipynb`](./localize-concepts.ipynb)
+
+Provides code to localize/visualize concepts of interest for a model/dataset of interest (setup for the "background" concept in `stylegan2_afhqdog512` as an example).
+
+| Localizing the learnt "background" concept vector  |
+| :-- |
+| ![image](./images/mask-bg-42.gif) ![image](./images/mask-bg-83.gif) ![image](./images/mask-bg-29.gif) |
+
+# [`./ffhq-edit-zoo.ipynb`](./ffhq-edit-zoo.ipynb)
+
+Quickly produce edits with annotated directions with pre-trained factors on FFHQ StyleGAN2.
+
+| Local image editing: "Big eyes"  |
+| :-- |
+| ![image](./images/qualitative.png) |
+
+## Setup
+
+Should you wish to run the notebooks, please consult this section below:
+
+### Install
+First, please install the dependencies with `pip install -r requirements.txt`, or alternatively with conda using `conda env create -f environment.yml`
+
+### Pre-trained models
+Should you wish to download the pre-trained models to run the notebooks, please first download them with:
+
+```bash
+wget -r -np -nH --cut-dirs=2 -R *index* http://eecs.qmul.ac.uk/~jo001/PandA-pretrained-models/
+```
+
+# citation
+
+If you find our work useful, please consider citing our paper:
+
+```bibtex
+@inproceedings{oldfield2023panda,
+    title={PandA: Unsupervised Learning of Parts and Appearances in the Feature Maps of GANs},
+    author={James Oldfield and Christos Tzelepis and Yannis Panagakis and Mihalis A. Nicolaou and Ioannis Patras},
+    booktitle={Int. Conf. Learn. Represent.},
+    year={2023}
+}
+```
+
+# contact
+
+**Please feel free to get in touch at**: `j.a.oldfield@qmul.ac.uk`
+
+---
+
+## credits
+
+- `./networks/genforce/` contains mostly code directly from [https://github.com/genforce/genforce](https://github.com/genforce/genforce).
+- `./networks/biggan/` contains mostly code directly from [https://github.com/huggingface/pytorch-pretrained-BigGAN](https://github.com/huggingface/pytorch-pretrained-BigGAN).
+- `./networks/stylegan3/` contains mostly code directly from [https://github.com/NVlabs/stylegan3](https://github.com/NVlabs/stylegan3).

requirements.txt

@@ -0,0 +1,14 @@
+boto3==1.21.8
+botocore==1.24.8
+bs4
+imageio==2.14.1
+jupyterlab
+matplotlib==3.5.1
+nltk==3.7
+numpy
+opencv-python==4.2.0.32
+Pillow==9.0.1
+requests==2.27.0
+tensorly==0.7.0
+torch==1.10.2
+torchvision==0.11.3

utils.py

@@ -0,0 +1,110 @@
+from matplotlib import pyplot as plt
+from matplotlib import gridspec
+import matplotlib.patches as mpatches
+import torch
+import numpy as np
+from PIL import Image
+
+
+def get_cols():
+    # list of perceptually distinct colours (for spatial factor plots)
+    return np.array([[255,0,0], [255,255,0], [0,234,255], [170,0,255], [255,127,0], [191,255,0], [0,149,255], [255,0,170], [255,212,0], [106,255,0], [0,64,255], [237,185,185], [185,215,237], [231,233,185], [220,185,237], [185,237,224], [143,35,35], [35,98,143], [143,106,35], [107,35,143], [79,143,35], [0,0,0], [115,115,115], [204,204,204]])
+
+
+def mapRange(value, inMin, inMax, outMin, outMax):
+    return outMin + (((value - inMin) / (inMax - inMin)) * (outMax - outMin))
+
+
+def plot_masks(Us, r, s, rs=256, save_path=None, title_factors=True):
+    """
+    Plots the parts factors with matplotlib for visualization
+
+    Parameters
+    ----------
+    Us : np.array
+        Learnt parts factor matrix.
+    r : int
+        Number of factors to show.
+    s : int
+        Dimensions of each part (h*w).
+    rs : int
+        Target size to downsize images to.
+    save_path : bool
+        Save figure?
+    title_factors : bool
+        Print matplotlib title on each part?
+    """
+
+    fig = plt.figure(constrained_layout=True, figsize=(20, 3))
+    spec = gridspec.GridSpec(ncols=r + 1, nrows=1, figure=fig)
+
+    for i in range(0, r):
+        fig.add_subplot(spec[i])
+
+        if title_factors:
+            plt.title(f'Part {i}')
+
+        part = Us[i].reshape([s, s])
+        part = mapRange(part, torch.min(part), torch.max(part), 0.0, 1.0) * 255
+        part = part.detach().cpu().numpy()
+        part = np.array(Image.fromarray(np.uint8(part)).convert('RGBA').resize((rs, rs), Image.NEAREST)) / 255
+
+        plt.axis('off')
+        plt.imshow(part, vmin=1, vmax=1, cmap='gray', alpha=1.00)
+
+    if save_path is not None:
+        plt.savefig(save_path)
+
+
+def plot_colours(image, Us, r, s, rs=128, save_path=None, alpha=1.0, seed=-1, legend=True):
+    """
+    Plots the parts factors over an image with matplotlib for visualization
+
+    Parameters
+    ----------
+    image : np.array
+        Image to visualize.
+    Us : np.array
+        Learnt parts factor matrix.
+    r : int
+        Number of factors to show.
+    s : int
+        Dimensions of each part (h*w).
+    rs : int
+        Target size to downsize images to.
+    alpha : float
+        Alpha value for the masks.
+    seed : int
+        Random seed when generating the colour palette (use -1 to use the provided "perceptually distinct" colour palette, but note this has a maximum of 30 colours or so).
+    legend : bool
+        Plot the legend, detailing the colour-coded parts key?
+    """
+
+    img = Image.fromarray(image).resize((rs, rs)).convert('RGBA')
+
+    # Use perceptually distinct colour list, or random seed (for e.g. if you have too many factors)
+    cols = get_cols()
+    if seed >= 0:
+        np.random.seed(seed)
+        cols = np.random.randint(0, 255, [r, 3])
+
+    plt.imshow(img, alpha=1.0)
+    plt.axis('off')
+
+    patches = []
+    for i in range(0, r):
+        mask = Us[i].detach().cpu().numpy().reshape([s, s])
+        mask = mapRange(mask, np.min(mask), np.max(mask), 0, 255)
+        mask = np.uint8(mask)
+        mask = np.array(Image.fromarray(mask).convert('L').resize((rs, rs)))
+        mask = (mask[:, :, None] / 255.) * np.array(np.concatenate([cols[i] / 255, [1]]))
+
+        patches += [mpatches.Patch(color=cols[i] / 255, label=f'Part {i}')]
+
+        plt.imshow(mask, vmin=0, vmax=1, alpha=alpha)
+
+    if legend:
+        plt.legend(title='Spatial factors', handles=patches, bbox_to_anchor=(1.01, 1.01), loc="upper left")
+
+    if save_path is not None:
+        plt.savefig(save_path, bbox_inches='tight', pad_inches=0.0)