diff --git a/README.md b/README.md
index ae6ed8cffd5b13622576c4b05270e521ccd631a0..fcd4d29957d71326e7fe3b2436f3e1afca19065f 100644
--- a/README.md
+++ b/README.md
@@ -1,8 +1,8 @@
---
title: StyleGAN NADA
-emoji: 💩
-colorFrom: red
-colorTo: pink
+emoji: 🌖
+colorFrom: blue
+colorTo: gray
sdk: gradio
sdk_version: 3.0.2
app_file: app.py
diff --git a/app.py b/app.py
new file mode 100644
index 0000000000000000000000000000000000000000..8d160e28b9ea3648c1d847a6c47086ab3424e904
--- /dev/null
+++ b/app.py
@@ -0,0 +1,387 @@
+import os
+import random
+
+import torch
+import gradio as gr
+
+from e4e.models.psp import pSp
+from util import *
+from huggingface_hub import hf_hub_download
+
+import tempfile
+from argparse import Namespace
+import shutil
+
+import dlib
+import numpy as np
+import torchvision.transforms as transforms
+from torchvision import utils
+
+from model.sg2_model import Generator
+from generate_videos import generate_frames, video_from_interpolations, project_code_by_edit_name
+from styleclip.styleclip_global import project_code_with_styleclip, style_tensor_to_style_dict
+
+import clip
+
+model_dir = "models"
+os.makedirs(model_dir, exist_ok=True)
+
+model_repos = {"e4e": ("akhaliq/JoJoGAN_e4e_ffhq_encode", "e4e_ffhq_encode.pt"),
+ "dlib": ("akhaliq/jojogan_dlib", "shape_predictor_68_face_landmarks.dat"),
+ "sc_fs3": ("rinong/stylegan-nada-models", "fs3.npy"),
+ "base": ("akhaliq/jojogan-stylegan2-ffhq-config-f", "stylegan2-ffhq-config-f.pt"),
+ "sketch": ("rinong/stylegan-nada-models", "sketch.pt"),
+ "joker": ("rinong/stylegan-nada-models", "joker.pt"),
+ "pixar": ("rinong/stylegan-nada-models", "pixar.pt"),
+ "botero": ("rinong/stylegan-nada-models", "botero.pt"),
+ "white_walker": ("rinong/stylegan-nada-models", "white_walker.pt"),
+ "zuckerberg": ("rinong/stylegan-nada-models", "zuckerberg.pt"),
+ "simpson": ("rinong/stylegan-nada-models", "simpson.pt"),
+ "ssj": ("rinong/stylegan-nada-models", "ssj.pt"),
+ "cubism": ("rinong/stylegan-nada-models", "cubism.pt"),
+ "disney_princess": ("rinong/stylegan-nada-models", "disney_princess.pt"),
+ "edvard_munch": ("rinong/stylegan-nada-models", "edvard_munch.pt"),
+ "van_gogh": ("rinong/stylegan-nada-models", "van_gogh.pt"),
+ "oil": ("rinong/stylegan-nada-models", "oil.pt"),
+ "rick_morty": ("rinong/stylegan-nada-models", "rick_morty.pt"),
+ "anime": ("rinong/stylegan-nada-models", "anime.pt"),
+ "shrek": ("rinong/stylegan-nada-models", "shrek.pt"),
+ "thanos": ("rinong/stylegan-nada-models", "thanos.pt"),
+ "ukiyoe": ("rinong/stylegan-nada-models", "ukiyoe.pt"),
+ "groot": ("rinong/stylegan-nada-models", "groot.pt"),
+ "witcher": ("rinong/stylegan-nada-models", "witcher.pt"),
+ "grafitti_on_wall": ("rinong/stylegan-nada-models", "grafitti_on_wall.pt"),
+ "modernism": ("rinong/stylegan-nada-models", "modernism.pt"),
+ "marble": ("rinong/stylegan-nada-models", "marble.pt"),
+ "vintage_comics": ("rinong/stylegan-nada-models", "vintage_comics.pt"),
+ "crochet": ("rinong/stylegan-nada-models", "crochet.pt"),
+ "modigliani": ("rinong/stylegan-nada-models", "modigliani.pt"),
+ "ghibli": ("rinong/stylegan-nada-models", "ghibli.pt"),
+ "elf": ("rinong/stylegan-nada-models", "elf.pt"),
+ "zombie": ("rinong/stylegan-nada-models", "zombie.pt"),
+ "werewolf": ("rinong/stylegan-nada-models", "werewolf.pt"),
+ "plastic_puppet": ("rinong/stylegan-nada-models", "plastic_puppet.pt"),
+ "mona_lisa": ("rinong/stylegan-nada-models", "mona_lisa.pt"),
+ }
+
+def get_models():
+ os.makedirs(model_dir, exist_ok=True)
+
+ model_paths = {}
+
+ for model_name, repo_details in model_repos.items():
+ download_path = hf_hub_download(repo_id=repo_details[0], filename=repo_details[1])
+ model_paths[model_name] = download_path
+
+ return model_paths
+
+model_paths = get_models()
+
+class ImageEditor(object):
+ def __init__(self):
+ self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+ latent_size = 512
+ n_mlp = 8
+ channel_mult = 2
+ model_size = 1024
+
+ self.generators = {}
+
+ self.model_list = [name for name in model_paths.keys() if name not in ["e4e", "dlib", "sc_fs3"]]
+
+ for model in self.model_list:
+ g_ema = Generator(
+ model_size, latent_size, n_mlp, channel_multiplier=channel_mult
+ ).to(self.device)
+
+ checkpoint = torch.load(model_paths[model], map_location=self.device)
+
+ g_ema.load_state_dict(checkpoint['g_ema'])
+
+ self.generators[model] = g_ema
+
+ self.experiment_args = {"model_path": model_paths["e4e"]}
+ self.experiment_args["transform"] = transforms.Compose(
+ [
+ transforms.Resize((256, 256)),
+ transforms.ToTensor(),
+ transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
+ ]
+ )
+ self.resize_dims = (256, 256)
+
+ model_path = self.experiment_args["model_path"]
+
+ ckpt = torch.load(model_path, map_location="cpu")
+ opts = ckpt["opts"]
+
+ opts["checkpoint_path"] = model_path
+ opts = Namespace(**opts)
+
+ self.e4e_net = pSp(opts, self.device)
+ self.e4e_net.eval()
+
+ self.shape_predictor = dlib.shape_predictor(
+ model_paths["dlib"]
+ )
+
+ self.styleclip_fs3 = torch.from_numpy(np.load(model_paths["sc_fs3"])).to(self.device)
+
+ self.clip_model, _ = clip.load("ViT-B/32", device=self.device)
+
+ print("setup complete")
+
+ def get_style_list(self):
+ style_list = []
+
+ for key in self.generators:
+ style_list.append(key)
+
+ return style_list
+
+ def invert_image(self, input_image):
+ input_image = self.run_alignment(str(input_image))
+
+ input_image = input_image.resize(self.resize_dims)
+
+ img_transforms = self.experiment_args["transform"]
+ transformed_image = img_transforms(input_image)
+
+ with torch.no_grad():
+ images, latents = self.run_on_batch(transformed_image.unsqueeze(0))
+ result_image, latent = images[0], latents[0]
+
+ inverted_latent = latent.unsqueeze(0).unsqueeze(1)
+
+ return inverted_latent
+
+ def get_generators_for_styles(self, output_styles, loop_styles=False):
+
+ if "base" in output_styles: # always start with base if chosen
+ output_styles.insert(0, output_styles.pop(output_styles.index("base")))
+ if loop_styles:
+ output_styles.append(output_styles[0])
+
+ return [self.generators[style] for style in output_styles]
+
+ def _pack_edits(func):
+ def inner(self,
+ edit_type_choice,
+ pose_slider,
+ smile_slider,
+ gender_slider,
+ age_slider,
+ hair_slider,
+ src_text_styleclip,
+ tar_text_styleclip,
+ alpha_styleclip,
+ beta_styleclip,
+ *args):
+
+ edit_choices = {"edit_type": edit_type_choice,
+ "pose": pose_slider,
+ "smile": smile_slider,
+ "gender": gender_slider,
+ "age": age_slider,
+ "hair_length": hair_slider,
+ "src_text": src_text_styleclip,
+ "tar_text": tar_text_styleclip,
+ "alpha": alpha_styleclip,
+ "beta": beta_styleclip}
+
+
+ return func(self, *args, edit_choices)
+
+ return inner
+
+ def get_target_latents(self, source_latent, edit_choices, generators):
+
+ target_latents = []
+
+ if edit_choices["edit_type"] == "InterFaceGAN":
+ np_source_latent = source_latent.squeeze(0).cpu().detach().numpy()
+
+ for attribute_name in ["pose", "smile", "gender", "age", "hair_length"]:
+ strength = edit_choices[attribute_name]
+ if strength != 0.0:
+ projected_code_np = project_code_by_edit_name(np_source_latent, attribute_name, strength)
+ target_latents.append(torch.from_numpy(projected_code_np).float().to(self.device))
+
+ elif edit_choices["edit_type"] == "StyleCLIP":
+ if edit_choices["alpha"] != 0.0:
+ source_s_dict = generators[0].get_s_code(source_latent, input_is_latent=True)[0]
+ target_latents.append(project_code_with_styleclip(source_s_dict,
+ edit_choices["src_text"],
+ edit_choices["tar_text"],
+ edit_choices["alpha"],
+ edit_choices["beta"],
+ generators[0],
+ self.styleclip_fs3,
+ self.clip_model))
+
+ # if edit type is none or if all sliders were set to 0
+ if not target_latents:
+ target_latents = [source_latent.squeeze(0), ] * max((len(generators) - 1), 1)
+
+ return target_latents
+
+ @_pack_edits
+ def edit_image(self, input, output_styles, edit_choices):
+ return self.predict(input, output_styles, edit_choices=edit_choices)
+
+ @_pack_edits
+ def edit_video(self, input, output_styles, loop_styles, edit_choices):
+ return self.predict(input, output_styles, generate_video=True, loop_styles=loop_styles, edit_choices=edit_choices)
+
+ def predict(
+ self,
+ input, # Input image path
+ output_styles, # Style checkbox options.
+ generate_video = False, # Generate a video instead of an output image
+ loop_styles = False, # Loop back to the initial style
+ edit_choices = None, # Optional dictionary with edit choice arguments
+ ):
+
+ if edit_choices is None:
+ edit_choices = {"edit_type": "None"}
+
+ # @title Align image
+ out_dir = tempfile.mkdtemp()
+
+ inverted_latent = self.invert_image(input)
+ generators = self.get_generators_for_styles(output_styles, loop_styles)
+
+ target_latents = self.get_target_latents(inverted_latent, edit_choices, generators)
+
+ if not generate_video:
+ output_paths = []
+
+ with torch.no_grad():
+ for g_ema in generators:
+ latent_for_gen = random.choice(target_latents)
+
+ if edit_choices["edit_type"] == "StyleCLIP":
+ latent_for_gen = style_tensor_to_style_dict(latent_for_gen, g_ema)
+ img, _ = g_ema(latent_for_gen, input_is_s_code=True, input_is_latent=True, truncation=1, randomize_noise=False)
+ else:
+ img, _ = g_ema([latent_for_gen], input_is_latent=True, truncation=1, randomize_noise=False)
+
+ output_path = os.path.join(out_dir, f"out_{len(output_paths)}.jpg")
+ utils.save_image(img, output_path, nrow=1, normalize=True, range=(-1, 1))
+
+ output_paths.append(output_path)
+
+ return output_paths
+
+ return self.generate_vid(generators, inverted_latent, target_latents, out_dir)
+
+ def generate_vid(self, generators, source_latent, target_latents, out_dir):
+
+ fps = 24
+
+ with tempfile.TemporaryDirectory() as dirpath:
+ generate_frames(source_latent, target_latents, generators, dirpath)
+ video_from_interpolations(fps, dirpath)
+
+ gen_path = os.path.join(dirpath, "out.mp4")
+ out_path = os.path.join(out_dir, "out.mp4")
+
+ shutil.copy2(gen_path, out_path)
+
+ return out_path
+
+ def run_alignment(self, image_path):
+ aligned_image = align_face(filepath=image_path, predictor=self.shape_predictor)
+ print("Aligned image has shape: {}".format(aligned_image.size))
+ return aligned_image
+
+ def run_on_batch(self, inputs):
+ images, latents = self.e4e_net(
+ inputs.to(self.device).float(), randomize_noise=False, return_latents=True
+ )
+ return images, latents
+
+editor = ImageEditor()
+
+blocks = gr.Blocks()
+
+with blocks:
+ gr.Markdown("StyleGAN-NADA
")
+ gr.Markdown(
+ "Inference demo for StyleGAN-NADA: CLIP-Guided Domain Adaptation of Image Generators (SIGGRAPH 2022).
"
+ )
+ gr.Markdown(
+ "Usage
Upload an image of your face, pick your desired output styles, and apply StyleGAN-based editing.
"
+ "Choose the edit image tab to create static images in all chosen styles. Choose the video tab in order to interpolate between all chosen styles
(To make it easier on the servers, we've limited video length. If you add too many styles (we recommend no more than 3!), they'll pass in the blink of an eye! 🤗)
"
+ )
+ gr.Markdown(
+ "For more information about the paper and code for training your own models (with text or images), please visit our project page or the official repository."
+ )
+
+ gr.Markdown("A note on social impact
This model relies on StyleGAN and CLIP, both of which are prone to biases inherited from their training data and their architecture. These may include (but are not limited to) poor representation of minorities or the perpetution of societal biases, such as gender norms. In particular, StyleGAN editing may induce undesired changes in skin tones. Moreover, generative models can, and have been used to create deep fake imagery which may assist in the spread of propaganda. However,
tools are available for identifying StyleGAN generated imagery, and any 'realistic' results produced by this model should be easily identifiable through such tools.
StyleGAN-NADA: CLIP-Guided Domain Adaptation of Image Generators | Project Page | Code
"
+ gr.Markdown(article)
+
+blocks.launch(enable_queue=True)
\ No newline at end of file
diff --git a/e4e/LICENSE b/e4e/LICENSE
new file mode 100644
index 0000000000000000000000000000000000000000..88ba9d421ea8acea9b4e3937535e72c282b3d4e6
--- /dev/null
+++ b/e4e/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 omertov
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/e4e/README.md b/e4e/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..14b6bc701b2bad3c2fc7b1d9b36f1892681ded5f
--- /dev/null
+++ b/e4e/README.md
@@ -0,0 +1,142 @@
+# Designing an Encoder for StyleGAN Image Manipulation
+ 
+ 
+ [](http://colab.research.google.com/github/omertov/encoder4editing/blob/main/notebooks/inference_playground.ipynb)
+
+> Recently, there has been a surge of diverse methods for performing image editing by employing pre-trained unconditional generators. Applying these methods on real images, however, remains a challenge, as it necessarily requires the inversion of the images into their latent space. To successfully invert a real image, one needs to find a latent code that reconstructs the input image accurately, and more importantly, allows for its meaningful manipulation. In this paper, we carefully study the latent space of StyleGAN, the state-of-the-art unconditional generator. We identify and analyze the existence of a distortion-editability tradeoff and a distortion-perception tradeoff within the StyleGAN latent space. We then suggest two principles for designing encoders in a manner that allows one to control the proximity of the inversions to regions that StyleGAN was originally trained on. We present an encoder based on our two principles that is specifically designed for facilitating editing on real images by balancing these tradeoffs. By evaluating its performance qualitatively and quantitatively on numerous challenging domains, including cars and horses, we show that our inversion method, followed by common editing techniques, achieves superior real-image editing quality, with only a small reconstruction accuracy drop.
+
+
+
+
+
+## Description
+Official Implementation of "Designing an Encoder for StyleGAN Image Manipulation" paper for both training and evaluation.
+The e4e encoder is specifically designed to complement existing image manipulation techniques performed over StyleGAN's latent space.
+
+## Recent Updates
+`2021.03.25`: Add pose editing direction.
+
+## Getting Started
+### Prerequisites
+- Linux or macOS
+- NVIDIA GPU + CUDA CuDNN (CPU may be possible with some modifications, but is not inherently supported)
+- Python 3
+
+### Installation
+- Clone the repository:
+```
+git clone https://github.com/omertov/encoder4editing.git
+cd encoder4editing
+```
+- Dependencies:
+We recommend running this repository using [Anaconda](https://docs.anaconda.com/anaconda/install/).
+All dependencies for defining the environment are provided in `environment/e4e_env.yaml`.
+
+### Inference Notebook
+We provide a Jupyter notebook found in `notebooks/inference_playground.ipynb` that allows one to encode and perform several editings on real images using StyleGAN.
+
+### Pretrained Models
+Please download the pre-trained models from the following links. Each e4e model contains the entire pSp framework architecture, including the encoder and decoder weights.
+| Path | Description
+| :--- | :----------
+|[FFHQ Inversion](https://drive.google.com/file/d/1cUv_reLE6k3604or78EranS7XzuVMWeO/view?usp=sharing) | FFHQ e4e encoder.
+|[Cars Inversion](https://drive.google.com/file/d/17faPqBce2m1AQeLCLHUVXaDfxMRU2QcV/view?usp=sharing) | Cars e4e encoder.
+|[Horse Inversion](https://drive.google.com/file/d/1TkLLnuX86B_BMo2ocYD0kX9kWh53rUVX/view?usp=sharing) | Horse e4e encoder.
+|[Church Inversion](https://drive.google.com/file/d/1-L0ZdnQLwtdy6-A_Ccgq5uNJGTqE7qBa/view?usp=sharing) | Church e4e encoder.
+
+If you wish to use one of the pretrained models for training or inference, you may do so using the flag `--checkpoint_path`.
+
+In addition, we provide various auxiliary models needed for training your own e4e model from scratch.
+| Path | Description
+| :--- | :----------
+|[FFHQ StyleGAN](https://drive.google.com/file/d/1EM87UquaoQmk17Q8d5kYIAHqu0dkYqdT/view?usp=sharing) | StyleGAN model pretrained on FFHQ taken from [rosinality](https://github.com/rosinality/stylegan2-pytorch) with 1024x1024 output resolution.
+|[IR-SE50 Model](https://drive.google.com/file/d/1KW7bjndL3QG3sxBbZxreGHigcCCpsDgn/view?usp=sharing) | Pretrained IR-SE50 model taken from [TreB1eN](https://github.com/TreB1eN/InsightFace_Pytorch) for use in our ID loss during training.
+|[MOCOv2 Model](https://drive.google.com/file/d/18rLcNGdteX5LwT7sv_F7HWr12HpVEzVe/view?usp=sharing) | Pretrained ResNet-50 model trained using MOCOv2 for use in our simmilarity loss for domains other then human faces during training.
+
+By default, we assume that all auxiliary models are downloaded and saved to the directory `pretrained_models`. However, you may use your own paths by changing the necessary values in `configs/path_configs.py`.
+
+## Training
+To train the e4e encoder, make sure the paths to the required models, as well as training and testing data is configured in `configs/path_configs.py` and `configs/data_configs.py`.
+#### **Training the e4e Encoder**
+```
+python scripts/train.py \
+--dataset_type cars_encode \
+--exp_dir new/experiment/directory \
+--start_from_latent_avg \
+--use_w_pool \
+--w_discriminator_lambda 0.1 \
+--progressive_start 20000 \
+--id_lambda 0.5 \
+--val_interval 10000 \
+--max_steps 200000 \
+--stylegan_size 512 \
+--stylegan_weights path/to/pretrained/stylegan.pt \
+--workers 8 \
+--batch_size 8 \
+--test_batch_size 4 \
+--test_workers 4
+```
+
+#### Training on your own dataset
+In order to train the e4e encoder on a custom dataset, perform the following adjustments:
+1. Insert the paths to your train and test data into the `dataset_paths` variable defined in `configs/paths_config.py`:
+```
+dataset_paths = {
+ 'my_train_data': '/path/to/train/images/directory',
+ 'my_test_data': '/path/to/test/images/directory'
+}
+```
+2. Configure a new dataset under the DATASETS variable defined in `configs/data_configs.py`:
+```
+DATASETS = {
+ 'my_data_encode': {
+ 'transforms': transforms_config.EncodeTransforms,
+ 'train_source_root': dataset_paths['my_train_data'],
+ 'train_target_root': dataset_paths['my_train_data'],
+ 'test_source_root': dataset_paths['my_test_data'],
+ 'test_target_root': dataset_paths['my_test_data']
+ }
+}
+```
+Refer to `configs/transforms_config.py` for the transformations applied to the train and test images during training.
+
+3. Finally, run a training session with `--dataset_type my_data_encode`.
+
+## Inference
+Having trained your model, you can use `scripts/inference.py` to apply the model on a set of images.
+For example,
+```
+python scripts/inference.py \
+--images_dir=/path/to/images/directory \
+--save_dir=/path/to/saving/directory \
+path/to/checkpoint.pt
+```
+
+## Latent Editing Consistency (LEC)
+As described in the paper, we suggest a new metric, Latent Editing Consistency (LEC), for evaluating the encoder's
+performance.
+We provide an example for calculating the metric over the FFHQ StyleGAN using the aging editing direction in
+`metrics/LEC.py`.
+
+To run the example:
+```
+cd metrics
+python LEC.py \
+--images_dir=/path/to/images/directory \
+path/to/checkpoint.pt
+```
+
+## Acknowledgments
+This code borrows heavily from [pixel2style2pixel](https://github.com/eladrich/pixel2style2pixel)
+
+## Citation
+If you use this code for your research, please cite our paper Designing an Encoder for StyleGAN Image Manipulation:
+
+```
+@article{tov2021designing,
+ title={Designing an Encoder for StyleGAN Image Manipulation},
+ author={Tov, Omer and Alaluf, Yuval and Nitzan, Yotam and Patashnik, Or and Cohen-Or, Daniel},
+ journal={arXiv preprint arXiv:2102.02766},
+ year={2021}
+}
+```
diff --git a/e4e/configs/__init__.py b/e4e/configs/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/e4e/configs/data_configs.py b/e4e/configs/data_configs.py
new file mode 100644
index 0000000000000000000000000000000000000000..deccb0b1c266ad4b6abaef53d67ec1ed0ddbd462
--- /dev/null
+++ b/e4e/configs/data_configs.py
@@ -0,0 +1,41 @@
+from configs import transforms_config
+from configs.paths_config import dataset_paths
+
+
+DATASETS = {
+ 'ffhq_encode': {
+ 'transforms': transforms_config.EncodeTransforms,
+ 'train_source_root': dataset_paths['ffhq'],
+ 'train_target_root': dataset_paths['ffhq'],
+ 'test_source_root': dataset_paths['celeba_test'],
+ 'test_target_root': dataset_paths['celeba_test'],
+ },
+ 'cars_encode': {
+ 'transforms': transforms_config.CarsEncodeTransforms,
+ 'train_source_root': dataset_paths['cars_train'],
+ 'train_target_root': dataset_paths['cars_train'],
+ 'test_source_root': dataset_paths['cars_test'],
+ 'test_target_root': dataset_paths['cars_test'],
+ },
+ 'horse_encode': {
+ 'transforms': transforms_config.EncodeTransforms,
+ 'train_source_root': dataset_paths['horse_train'],
+ 'train_target_root': dataset_paths['horse_train'],
+ 'test_source_root': dataset_paths['horse_test'],
+ 'test_target_root': dataset_paths['horse_test'],
+ },
+ 'church_encode': {
+ 'transforms': transforms_config.EncodeTransforms,
+ 'train_source_root': dataset_paths['church_train'],
+ 'train_target_root': dataset_paths['church_train'],
+ 'test_source_root': dataset_paths['church_test'],
+ 'test_target_root': dataset_paths['church_test'],
+ },
+ 'cats_encode': {
+ 'transforms': transforms_config.EncodeTransforms,
+ 'train_source_root': dataset_paths['cats_train'],
+ 'train_target_root': dataset_paths['cats_train'],
+ 'test_source_root': dataset_paths['cats_test'],
+ 'test_target_root': dataset_paths['cats_test'],
+ }
+}
diff --git a/e4e/configs/paths_config.py b/e4e/configs/paths_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..4604f6063b8125364a52a492de52fcc54004f373
--- /dev/null
+++ b/e4e/configs/paths_config.py
@@ -0,0 +1,28 @@
+dataset_paths = {
+ # Face Datasets (In the paper: FFHQ - train, CelebAHQ - test)
+ 'ffhq': '',
+ 'celeba_test': '',
+
+ # Cars Dataset (In the paper: Stanford cars)
+ 'cars_train': '',
+ 'cars_test': '',
+
+ # Horse Dataset (In the paper: LSUN Horse)
+ 'horse_train': '',
+ 'horse_test': '',
+
+ # Church Dataset (In the paper: LSUN Church)
+ 'church_train': '',
+ 'church_test': '',
+
+ # Cats Dataset (In the paper: LSUN Cat)
+ 'cats_train': '',
+ 'cats_test': ''
+}
+
+model_paths = {
+ 'stylegan_ffhq': 'pretrained_models/stylegan2-ffhq-config-f.pt',
+ 'ir_se50': 'pretrained_models/model_ir_se50.pth',
+ 'shape_predictor': 'pretrained_models/shape_predictor_68_face_landmarks.dat',
+ 'moco': 'pretrained_models/moco_v2_800ep_pretrain.pth'
+}
diff --git a/e4e/configs/transforms_config.py b/e4e/configs/transforms_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..ac12b5d5ba0571f21715e0f6b24b9c1ebe84bf72
--- /dev/null
+++ b/e4e/configs/transforms_config.py
@@ -0,0 +1,62 @@
+from abc import abstractmethod
+import torchvision.transforms as transforms
+
+
+class TransformsConfig(object):
+
+ def __init__(self, opts):
+ self.opts = opts
+
+ @abstractmethod
+ def get_transforms(self):
+ pass
+
+
+class EncodeTransforms(TransformsConfig):
+
+ def __init__(self, opts):
+ super(EncodeTransforms, self).__init__(opts)
+
+ def get_transforms(self):
+ transforms_dict = {
+ 'transform_gt_train': transforms.Compose([
+ transforms.Resize((256, 256)),
+ transforms.RandomHorizontalFlip(0.5),
+ transforms.ToTensor(),
+ transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+ 'transform_source': None,
+ 'transform_test': transforms.Compose([
+ transforms.Resize((256, 256)),
+ transforms.ToTensor(),
+ transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+ 'transform_inference': transforms.Compose([
+ transforms.Resize((256, 256)),
+ transforms.ToTensor(),
+ transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+ }
+ return transforms_dict
+
+
+class CarsEncodeTransforms(TransformsConfig):
+
+ def __init__(self, opts):
+ super(CarsEncodeTransforms, self).__init__(opts)
+
+ def get_transforms(self):
+ transforms_dict = {
+ 'transform_gt_train': transforms.Compose([
+ transforms.Resize((192, 256)),
+ transforms.RandomHorizontalFlip(0.5),
+ transforms.ToTensor(),
+ transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+ 'transform_source': None,
+ 'transform_test': transforms.Compose([
+ transforms.Resize((192, 256)),
+ transforms.ToTensor(),
+ transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+ 'transform_inference': transforms.Compose([
+ transforms.Resize((192, 256)),
+ transforms.ToTensor(),
+ transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+ }
+ return transforms_dict
diff --git a/e4e/models/__init__.py b/e4e/models/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/e4e/models/discriminator.py b/e4e/models/discriminator.py
new file mode 100644
index 0000000000000000000000000000000000000000..16bf3722c7f2e35cdc9bd177a33ed0975e67200d
--- /dev/null
+++ b/e4e/models/discriminator.py
@@ -0,0 +1,20 @@
+from torch import nn
+
+
+class LatentCodesDiscriminator(nn.Module):
+ def __init__(self, style_dim, n_mlp):
+ super().__init__()
+
+ self.style_dim = style_dim
+
+ layers = []
+ for i in range(n_mlp-1):
+ layers.append(
+ nn.Linear(style_dim, style_dim)
+ )
+ layers.append(nn.LeakyReLU(0.2))
+ layers.append(nn.Linear(512, 1))
+ self.mlp = nn.Sequential(*layers)
+
+ def forward(self, w):
+ return self.mlp(w)
diff --git a/e4e/models/encoders/__init__.py b/e4e/models/encoders/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/e4e/models/encoders/helpers.py b/e4e/models/encoders/helpers.py
new file mode 100644
index 0000000000000000000000000000000000000000..c4a58b34ea5ca6912fe53c63dede0a8696f5c024
--- /dev/null
+++ b/e4e/models/encoders/helpers.py
@@ -0,0 +1,140 @@
+from collections import namedtuple
+import torch
+import torch.nn.functional as F
+from torch.nn import Conv2d, BatchNorm2d, PReLU, ReLU, Sigmoid, MaxPool2d, AdaptiveAvgPool2d, Sequential, Module
+
+"""
+ArcFace implementation from [TreB1eN](https://github.com/TreB1eN/InsightFace_Pytorch)
+"""
+
+
+class Flatten(Module):
+ def forward(self, input):
+ return input.view(input.size(0), -1)
+
+
+def l2_norm(input, axis=1):
+ norm = torch.norm(input, 2, axis, True)
+ output = torch.div(input, norm)
+ return output
+
+
+class Bottleneck(namedtuple('Block', ['in_channel', 'depth', 'stride'])):
+ """ A named tuple describing a ResNet block. """
+
+
+def get_block(in_channel, depth, num_units, stride=2):
+ return [Bottleneck(in_channel, depth, stride)] + [Bottleneck(depth, depth, 1) for i in range(num_units - 1)]
+
+
+def get_blocks(num_layers):
+ if num_layers == 50:
+ blocks = [
+ get_block(in_channel=64, depth=64, num_units=3),
+ get_block(in_channel=64, depth=128, num_units=4),
+ get_block(in_channel=128, depth=256, num_units=14),
+ get_block(in_channel=256, depth=512, num_units=3)
+ ]
+ elif num_layers == 100:
+ blocks = [
+ get_block(in_channel=64, depth=64, num_units=3),
+ get_block(in_channel=64, depth=128, num_units=13),
+ get_block(in_channel=128, depth=256, num_units=30),
+ get_block(in_channel=256, depth=512, num_units=3)
+ ]
+ elif num_layers == 152:
+ blocks = [
+ get_block(in_channel=64, depth=64, num_units=3),
+ get_block(in_channel=64, depth=128, num_units=8),
+ get_block(in_channel=128, depth=256, num_units=36),
+ get_block(in_channel=256, depth=512, num_units=3)
+ ]
+ else:
+ raise ValueError("Invalid number of layers: {}. Must be one of [50, 100, 152]".format(num_layers))
+ return blocks
+
+
+class SEModule(Module):
+ def __init__(self, channels, reduction):
+ super(SEModule, self).__init__()
+ self.avg_pool = AdaptiveAvgPool2d(1)
+ self.fc1 = Conv2d(channels, channels // reduction, kernel_size=1, padding=0, bias=False)
+ self.relu = ReLU(inplace=True)
+ self.fc2 = Conv2d(channels // reduction, channels, kernel_size=1, padding=0, bias=False)
+ self.sigmoid = Sigmoid()
+
+ def forward(self, x):
+ module_input = x
+ x = self.avg_pool(x)
+ x = self.fc1(x)
+ x = self.relu(x)
+ x = self.fc2(x)
+ x = self.sigmoid(x)
+ return module_input * x
+
+
+class bottleneck_IR(Module):
+ def __init__(self, in_channel, depth, stride):
+ super(bottleneck_IR, self).__init__()
+ if in_channel == depth:
+ self.shortcut_layer = MaxPool2d(1, stride)
+ else:
+ self.shortcut_layer = Sequential(
+ Conv2d(in_channel, depth, (1, 1), stride, bias=False),
+ BatchNorm2d(depth)
+ )
+ self.res_layer = Sequential(
+ BatchNorm2d(in_channel),
+ Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False), PReLU(depth),
+ Conv2d(depth, depth, (3, 3), stride, 1, bias=False), BatchNorm2d(depth)
+ )
+
+ def forward(self, x):
+ shortcut = self.shortcut_layer(x)
+ res = self.res_layer(x)
+ return res + shortcut
+
+
+class bottleneck_IR_SE(Module):
+ def __init__(self, in_channel, depth, stride):
+ super(bottleneck_IR_SE, self).__init__()
+ if in_channel == depth:
+ self.shortcut_layer = MaxPool2d(1, stride)
+ else:
+ self.shortcut_layer = Sequential(
+ Conv2d(in_channel, depth, (1, 1), stride, bias=False),
+ BatchNorm2d(depth)
+ )
+ self.res_layer = Sequential(
+ BatchNorm2d(in_channel),
+ Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False),
+ PReLU(depth),
+ Conv2d(depth, depth, (3, 3), stride, 1, bias=False),
+ BatchNorm2d(depth),
+ SEModule(depth, 16)
+ )
+
+ def forward(self, x):
+ shortcut = self.shortcut_layer(x)
+ res = self.res_layer(x)
+ return res + shortcut
+
+
+def _upsample_add(x, y):
+ """Upsample and add two feature maps.
+ Args:
+ x: (Variable) top feature map to be upsampled.
+ y: (Variable) lateral feature map.
+ Returns:
+ (Variable) added feature map.
+ Note in PyTorch, when input size is odd, the upsampled feature map
+ with `F.upsample(..., scale_factor=2, mode='nearest')`
+ maybe not equal to the lateral feature map size.
+ e.g.
+ original input size: [N,_,15,15] ->
+ conv2d feature map size: [N,_,8,8] ->
+ upsampled feature map size: [N,_,16,16]
+ So we choose bilinear upsample which supports arbitrary output sizes.
+ """
+ _, _, H, W = y.size()
+ return F.interpolate(x, size=(H, W), mode='bilinear', align_corners=True) + y
diff --git a/e4e/models/encoders/model_irse.py b/e4e/models/encoders/model_irse.py
new file mode 100644
index 0000000000000000000000000000000000000000..6a94d67542f961ff6533f0335cf4cb0fa54024fb
--- /dev/null
+++ b/e4e/models/encoders/model_irse.py
@@ -0,0 +1,84 @@
+from torch.nn import Linear, Conv2d, BatchNorm1d, BatchNorm2d, PReLU, Dropout, Sequential, Module
+from e4e.models.encoders.helpers import get_blocks, Flatten, bottleneck_IR, bottleneck_IR_SE, l2_norm
+
+"""
+Modified Backbone implementation from [TreB1eN](https://github.com/TreB1eN/InsightFace_Pytorch)
+"""
+
+
+class Backbone(Module):
+ def __init__(self, input_size, num_layers, mode='ir', drop_ratio=0.4, affine=True):
+ super(Backbone, self).__init__()
+ assert input_size in [112, 224], "input_size should be 112 or 224"
+ assert num_layers in [50, 100, 152], "num_layers should be 50, 100 or 152"
+ assert mode in ['ir', 'ir_se'], "mode should be ir or ir_se"
+ blocks = get_blocks(num_layers)
+ if mode == 'ir':
+ unit_module = bottleneck_IR
+ elif mode == 'ir_se':
+ unit_module = bottleneck_IR_SE
+ self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1, bias=False),
+ BatchNorm2d(64),
+ PReLU(64))
+ if input_size == 112:
+ self.output_layer = Sequential(BatchNorm2d(512),
+ Dropout(drop_ratio),
+ Flatten(),
+ Linear(512 * 7 * 7, 512),
+ BatchNorm1d(512, affine=affine))
+ else:
+ self.output_layer = Sequential(BatchNorm2d(512),
+ Dropout(drop_ratio),
+ Flatten(),
+ Linear(512 * 14 * 14, 512),
+ BatchNorm1d(512, affine=affine))
+
+ modules = []
+ for block in blocks:
+ for bottleneck in block:
+ modules.append(unit_module(bottleneck.in_channel,
+ bottleneck.depth,
+ bottleneck.stride))
+ self.body = Sequential(*modules)
+
+ def forward(self, x):
+ x = self.input_layer(x)
+ x = self.body(x)
+ x = self.output_layer(x)
+ return l2_norm(x)
+
+
+def IR_50(input_size):
+ """Constructs a ir-50 model."""
+ model = Backbone(input_size, num_layers=50, mode='ir', drop_ratio=0.4, affine=False)
+ return model
+
+
+def IR_101(input_size):
+ """Constructs a ir-101 model."""
+ model = Backbone(input_size, num_layers=100, mode='ir', drop_ratio=0.4, affine=False)
+ return model
+
+
+def IR_152(input_size):
+ """Constructs a ir-152 model."""
+ model = Backbone(input_size, num_layers=152, mode='ir', drop_ratio=0.4, affine=False)
+ return model
+
+
+def IR_SE_50(input_size):
+ """Constructs a ir_se-50 model."""
+ model = Backbone(input_size, num_layers=50, mode='ir_se', drop_ratio=0.4, affine=False)
+ return model
+
+
+def IR_SE_101(input_size):
+ """Constructs a ir_se-101 model."""
+ model = Backbone(input_size, num_layers=100, mode='ir_se', drop_ratio=0.4, affine=False)
+ return model
+
+
+def IR_SE_152(input_size):
+ """Constructs a ir_se-152 model."""
+ model = Backbone(input_size, num_layers=152, mode='ir_se', drop_ratio=0.4, affine=False)
+ return model
diff --git a/e4e/models/encoders/psp_encoders.py b/e4e/models/encoders/psp_encoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..dc49acd11f062cbd29f839ee3c04bce7fa84f479
--- /dev/null
+++ b/e4e/models/encoders/psp_encoders.py
@@ -0,0 +1,200 @@
+from enum import Enum
+import math
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import Conv2d, BatchNorm2d, PReLU, Sequential, Module
+
+from e4e.models.encoders.helpers import get_blocks, bottleneck_IR, bottleneck_IR_SE, _upsample_add
+from e4e.models.stylegan2.model import EqualLinear
+
+
+class ProgressiveStage(Enum):
+ WTraining = 0
+ Delta1Training = 1
+ Delta2Training = 2
+ Delta3Training = 3
+ Delta4Training = 4
+ Delta5Training = 5
+ Delta6Training = 6
+ Delta7Training = 7
+ Delta8Training = 8
+ Delta9Training = 9
+ Delta10Training = 10
+ Delta11Training = 11
+ Delta12Training = 12
+ Delta13Training = 13
+ Delta14Training = 14
+ Delta15Training = 15
+ Delta16Training = 16
+ Delta17Training = 17
+ Inference = 18
+
+
+class GradualStyleBlock(Module):
+ def __init__(self, in_c, out_c, spatial):
+ super(GradualStyleBlock, self).__init__()
+ self.out_c = out_c
+ self.spatial = spatial
+ num_pools = int(np.log2(spatial))
+ modules = []
+ modules += [Conv2d(in_c, out_c, kernel_size=3, stride=2, padding=1),
+ nn.LeakyReLU()]
+ for i in range(num_pools - 1):
+ modules += [
+ Conv2d(out_c, out_c, kernel_size=3, stride=2, padding=1),
+ nn.LeakyReLU()
+ ]
+ self.convs = nn.Sequential(*modules)
+ self.linear = EqualLinear(out_c, out_c, lr_mul=1)
+
+ def forward(self, x):
+ x = self.convs(x)
+ x = x.view(-1, self.out_c)
+ x = self.linear(x)
+ return x
+
+
+class GradualStyleEncoder(Module):
+ def __init__(self, num_layers, mode='ir', opts=None):
+ super(GradualStyleEncoder, self).__init__()
+ assert num_layers in [50, 100, 152], 'num_layers should be 50,100, or 152'
+ assert mode in ['ir', 'ir_se'], 'mode should be ir or ir_se'
+ blocks = get_blocks(num_layers)
+ if mode == 'ir':
+ unit_module = bottleneck_IR
+ elif mode == 'ir_se':
+ unit_module = bottleneck_IR_SE
+ self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1, bias=False),
+ BatchNorm2d(64),
+ PReLU(64))
+ modules = []
+ for block in blocks:
+ for bottleneck in block:
+ modules.append(unit_module(bottleneck.in_channel,
+ bottleneck.depth,
+ bottleneck.stride))
+ self.body = Sequential(*modules)
+
+ self.styles = nn.ModuleList()
+ log_size = int(math.log(opts.stylegan_size, 2))
+ self.style_count = 2 * log_size - 2
+ self.coarse_ind = 3
+ self.middle_ind = 7
+ for i in range(self.style_count):
+ if i < self.coarse_ind:
+ style = GradualStyleBlock(512, 512, 16)
+ elif i < self.middle_ind:
+ style = GradualStyleBlock(512, 512, 32)
+ else:
+ style = GradualStyleBlock(512, 512, 64)
+ self.styles.append(style)
+ self.latlayer1 = nn.Conv2d(256, 512, kernel_size=1, stride=1, padding=0)
+ self.latlayer2 = nn.Conv2d(128, 512, kernel_size=1, stride=1, padding=0)
+
+ def forward(self, x):
+ x = self.input_layer(x)
+
+ latents = []
+ modulelist = list(self.body._modules.values())
+ for i, l in enumerate(modulelist):
+ x = l(x)
+ if i == 6:
+ c1 = x
+ elif i == 20:
+ c2 = x
+ elif i == 23:
+ c3 = x
+
+ for j in range(self.coarse_ind):
+ latents.append(self.styles[j](c3))
+
+ p2 = _upsample_add(c3, self.latlayer1(c2))
+ for j in range(self.coarse_ind, self.middle_ind):
+ latents.append(self.styles[j](p2))
+
+ p1 = _upsample_add(p2, self.latlayer2(c1))
+ for j in range(self.middle_ind, self.style_count):
+ latents.append(self.styles[j](p1))
+
+ out = torch.stack(latents, dim=1)
+ return out
+
+
+class Encoder4Editing(Module):
+ def __init__(self, num_layers, mode='ir', opts=None):
+ super(Encoder4Editing, self).__init__()
+ assert num_layers in [50, 100, 152], 'num_layers should be 50,100, or 152'
+ assert mode in ['ir', 'ir_se'], 'mode should be ir or ir_se'
+ blocks = get_blocks(num_layers)
+ if mode == 'ir':
+ unit_module = bottleneck_IR
+ elif mode == 'ir_se':
+ unit_module = bottleneck_IR_SE
+ self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1, bias=False),
+ BatchNorm2d(64),
+ PReLU(64))
+ modules = []
+ for block in blocks:
+ for bottleneck in block:
+ modules.append(unit_module(bottleneck.in_channel,
+ bottleneck.depth,
+ bottleneck.stride))
+ self.body = Sequential(*modules)
+
+ self.styles = nn.ModuleList()
+ log_size = int(math.log(opts.stylegan_size, 2))
+ self.style_count = 2 * log_size - 2
+ self.coarse_ind = 3
+ self.middle_ind = 7
+
+ for i in range(self.style_count):
+ if i < self.coarse_ind:
+ style = GradualStyleBlock(512, 512, 16)
+ elif i < self.middle_ind:
+ style = GradualStyleBlock(512, 512, 32)
+ else:
+ style = GradualStyleBlock(512, 512, 64)
+ self.styles.append(style)
+
+ self.latlayer1 = nn.Conv2d(256, 512, kernel_size=1, stride=1, padding=0)
+ self.latlayer2 = nn.Conv2d(128, 512, kernel_size=1, stride=1, padding=0)
+
+ self.progressive_stage = ProgressiveStage.Inference
+
+ def get_deltas_starting_dimensions(self):
+ ''' Get a list of the initial dimension of every delta from which it is applied '''
+ return list(range(self.style_count)) # Each dimension has a delta applied to it
+
+ def set_progressive_stage(self, new_stage: ProgressiveStage):
+ self.progressive_stage = new_stage
+ print('Changed progressive stage to: ', new_stage)
+
+ def forward(self, x):
+ x = self.input_layer(x)
+
+ modulelist = list(self.body._modules.values())
+ for i, l in enumerate(modulelist):
+ x = l(x)
+ if i == 6:
+ c1 = x
+ elif i == 20:
+ c2 = x
+ elif i == 23:
+ c3 = x
+
+ # Infer main W and duplicate it
+ w0 = self.styles[0](c3)
+ w = w0.repeat(self.style_count, 1, 1).permute(1, 0, 2)
+ stage = self.progressive_stage.value
+ features = c3
+ for i in range(1, min(stage + 1, self.style_count)): # Infer additional deltas
+ if i == self.coarse_ind:
+ p2 = _upsample_add(c3, self.latlayer1(c2)) # FPN's middle features
+ features = p2
+ elif i == self.middle_ind:
+ p1 = _upsample_add(p2, self.latlayer2(c1)) # FPN's fine features
+ features = p1
+ delta_i = self.styles[i](features)
+ w[:, i] += delta_i
+ return w
diff --git a/e4e/models/latent_codes_pool.py b/e4e/models/latent_codes_pool.py
new file mode 100644
index 0000000000000000000000000000000000000000..0281d4b5e80f8eb26e824fa35b4f908dcb6634e6
--- /dev/null
+++ b/e4e/models/latent_codes_pool.py
@@ -0,0 +1,55 @@
+import random
+import torch
+
+
+class LatentCodesPool:
+ """This class implements latent codes buffer that stores previously generated w latent codes.
+ This buffer enables us to update discriminators using a history of generated w's
+ rather than the ones produced by the latest encoder.
+ """
+
+ def __init__(self, pool_size):
+ """Initialize the ImagePool class
+ Parameters:
+ pool_size (int) -- the size of image buffer, if pool_size=0, no buffer will be created
+ """
+ self.pool_size = pool_size
+ if self.pool_size > 0: # create an empty pool
+ self.num_ws = 0
+ self.ws = []
+
+ def query(self, ws):
+ """Return w's from the pool.
+ Parameters:
+ ws: the latest generated w's from the generator
+ Returns w's from the buffer.
+ By 50/100, the buffer will return input w's.
+ By 50/100, the buffer will return w's previously stored in the buffer,
+ and insert the current w's to the buffer.
+ """
+ if self.pool_size == 0: # if the buffer size is 0, do nothing
+ return ws
+ return_ws = []
+ for w in ws: # ws.shape: (batch, 512) or (batch, n_latent, 512)
+ # w = torch.unsqueeze(image.data, 0)
+ if w.ndim == 2:
+ i = random.randint(0, len(w) - 1) # apply a random latent index as a candidate
+ w = w[i]
+ self.handle_w(w, return_ws)
+ return_ws = torch.stack(return_ws, 0) # collect all the images and return
+ return return_ws
+
+ def handle_w(self, w, return_ws):
+ if self.num_ws < self.pool_size: # if the buffer is not full; keep inserting current codes to the buffer
+ self.num_ws = self.num_ws + 1
+ self.ws.append(w)
+ return_ws.append(w)
+ else:
+ p = random.uniform(0, 1)
+ if p > 0.5: # by 50% chance, the buffer will return a previously stored latent code, and insert the current code into the buffer
+ random_id = random.randint(0, self.pool_size - 1) # randint is inclusive
+ tmp = self.ws[random_id].clone()
+ self.ws[random_id] = w
+ return_ws.append(tmp)
+ else: # by another 50% chance, the buffer will return the current image
+ return_ws.append(w)
diff --git a/e4e/models/psp.py b/e4e/models/psp.py
new file mode 100644
index 0000000000000000000000000000000000000000..36c0b2b7b3fdd28bc32272d0d8fcff24e4848355
--- /dev/null
+++ b/e4e/models/psp.py
@@ -0,0 +1,99 @@
+import matplotlib
+
+matplotlib.use('Agg')
+import torch
+from torch import nn
+from e4e.models.encoders import psp_encoders
+from e4e.models.stylegan2.model import Generator
+from e4e.configs.paths_config import model_paths
+
+
+def get_keys(d, name):
+ if 'state_dict' in d:
+ d = d['state_dict']
+ d_filt = {k[len(name) + 1:]: v for k, v in d.items() if k[:len(name)] == name}
+ return d_filt
+
+
+class pSp(nn.Module):
+
+ def __init__(self, opts, device):
+ super(pSp, self).__init__()
+ self.opts = opts
+ self.device = device
+ # Define architecture
+ self.encoder = self.set_encoder()
+ self.decoder = Generator(opts.stylegan_size, 512, 8, channel_multiplier=2)
+ self.face_pool = torch.nn.AdaptiveAvgPool2d((256, 256))
+ # Load weights if needed
+ self.load_weights()
+
+ def set_encoder(self):
+ if self.opts.encoder_type == 'GradualStyleEncoder':
+ encoder = psp_encoders.GradualStyleEncoder(50, 'ir_se', self.opts)
+ elif self.opts.encoder_type == 'Encoder4Editing':
+ encoder = psp_encoders.Encoder4Editing(50, 'ir_se', self.opts)
+ else:
+ raise Exception('{} is not a valid encoders'.format(self.opts.encoder_type))
+ return encoder
+
+ def load_weights(self):
+ if self.opts.checkpoint_path is not None:
+ print('Loading e4e over the pSp framework from checkpoint: {}'.format(self.opts.checkpoint_path))
+ ckpt = torch.load(self.opts.checkpoint_path, map_location='cpu')
+ self.encoder.load_state_dict(get_keys(ckpt, 'encoder'), strict=True)
+ self.decoder.load_state_dict(get_keys(ckpt, 'decoder'), strict=True)
+ self.__load_latent_avg(ckpt)
+ else:
+ print('Loading encoders weights from irse50!')
+ encoder_ckpt = torch.load(model_paths['ir_se50'])
+ self.encoder.load_state_dict(encoder_ckpt, strict=False)
+ print('Loading decoder weights from pretrained!')
+ ckpt = torch.load(self.opts.stylegan_weights)
+ self.decoder.load_state_dict(ckpt['g_ema'], strict=False)
+ self.__load_latent_avg(ckpt, repeat=self.encoder.style_count)
+
+ def forward(self, x, resize=True, latent_mask=None, input_code=False, randomize_noise=True,
+ inject_latent=None, return_latents=False, alpha=None):
+ if input_code:
+ codes = x
+ else:
+ codes = self.encoder(x)
+ # normalize with respect to the center of an average face
+ if self.opts.start_from_latent_avg:
+ if codes.ndim == 2:
+ codes = codes + self.latent_avg.repeat(codes.shape[0], 1, 1)[:, 0, :]
+ else:
+ codes = codes + self.latent_avg.repeat(codes.shape[0], 1, 1)
+
+ if latent_mask is not None:
+ for i in latent_mask:
+ if inject_latent is not None:
+ if alpha is not None:
+ codes[:, i] = alpha * inject_latent[:, i] + (1 - alpha) * codes[:, i]
+ else:
+ codes[:, i] = inject_latent[:, i]
+ else:
+ codes[:, i] = 0
+
+ input_is_latent = not input_code
+ images, result_latent = self.decoder([codes],
+ input_is_latent=input_is_latent,
+ randomize_noise=randomize_noise,
+ return_latents=return_latents)
+
+ if resize:
+ images = self.face_pool(images)
+
+ if return_latents:
+ return images, result_latent
+ else:
+ return images
+
+ def __load_latent_avg(self, ckpt, repeat=None):
+ if 'latent_avg' in ckpt:
+ self.latent_avg = ckpt['latent_avg'].to(self.device)
+ if repeat is not None:
+ self.latent_avg = self.latent_avg.repeat(repeat, 1)
+ else:
+ self.latent_avg = None
diff --git a/e4e/models/stylegan2/__init__.py b/e4e/models/stylegan2/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/e4e/models/stylegan2/model.py b/e4e/models/stylegan2/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..fcb12af85669ab6fd7f79cb14ddbdf80b2fbd83d
--- /dev/null
+++ b/e4e/models/stylegan2/model.py
@@ -0,0 +1,678 @@
+import math
+import random
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+if torch.cuda.is_available():
+ from op.fused_act import FusedLeakyReLU, fused_leaky_relu
+ from op.upfirdn2d import upfirdn2d
+else:
+ from op.fused_act_cpu import FusedLeakyReLU, fused_leaky_relu
+ from op.upfirdn2d_cpu import upfirdn2d
+
+
+class PixelNorm(nn.Module):
+ def __init__(self):
+ super().__init__()
+
+ def forward(self, input):
+ return input * torch.rsqrt(torch.mean(input ** 2, dim=1, keepdim=True) + 1e-8)
+
+
+def make_kernel(k):
+ k = torch.tensor(k, dtype=torch.float32)
+
+ if k.ndim == 1:
+ k = k[None, :] * k[:, None]
+
+ k /= k.sum()
+
+ return k
+
+
+class Upsample(nn.Module):
+ def __init__(self, kernel, factor=2):
+ super().__init__()
+
+ self.factor = factor
+ kernel = make_kernel(kernel) * (factor ** 2)
+ self.register_buffer('kernel', kernel)
+
+ p = kernel.shape[0] - factor
+
+ pad0 = (p + 1) // 2 + factor - 1
+ pad1 = p // 2
+
+ self.pad = (pad0, pad1)
+
+ def forward(self, input):
+ out = upfirdn2d(input, self.kernel, up=self.factor, down=1, pad=self.pad)
+
+ return out
+
+
+class Downsample(nn.Module):
+ def __init__(self, kernel, factor=2):
+ super().__init__()
+
+ self.factor = factor
+ kernel = make_kernel(kernel)
+ self.register_buffer('kernel', kernel)
+
+ p = kernel.shape[0] - factor
+
+ pad0 = (p + 1) // 2
+ pad1 = p // 2
+
+ self.pad = (pad0, pad1)
+
+ def forward(self, input):
+ out = upfirdn2d(input, self.kernel, up=1, down=self.factor, pad=self.pad)
+
+ return out
+
+
+class Blur(nn.Module):
+ def __init__(self, kernel, pad, upsample_factor=1):
+ super().__init__()
+
+ kernel = make_kernel(kernel)
+
+ if upsample_factor > 1:
+ kernel = kernel * (upsample_factor ** 2)
+
+ self.register_buffer('kernel', kernel)
+
+ self.pad = pad
+
+ def forward(self, input):
+ out = upfirdn2d(input, self.kernel, pad=self.pad)
+
+ return out
+
+
+class EqualConv2d(nn.Module):
+ def __init__(
+ self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True
+ ):
+ super().__init__()
+
+ self.weight = nn.Parameter(
+ torch.randn(out_channel, in_channel, kernel_size, kernel_size)
+ )
+ self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2)
+
+ self.stride = stride
+ self.padding = padding
+
+ if bias:
+ self.bias = nn.Parameter(torch.zeros(out_channel))
+
+ else:
+ self.bias = None
+
+ def forward(self, input):
+ out = F.conv2d(
+ input,
+ self.weight * self.scale,
+ bias=self.bias,
+ stride=self.stride,
+ padding=self.padding,
+ )
+
+ return out
+
+ def __repr__(self):
+ return (
+ f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},'
+ f' {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})'
+ )
+
+
+class EqualLinear(nn.Module):
+ def __init__(
+ self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None
+ ):
+ super().__init__()
+
+ self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul))
+
+ if bias:
+ self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init))
+
+ else:
+ self.bias = None
+
+ self.activation = activation
+
+ self.scale = (1 / math.sqrt(in_dim)) * lr_mul
+ self.lr_mul = lr_mul
+
+ def forward(self, input):
+ if self.activation:
+ out = F.linear(input, self.weight * self.scale)
+ out = fused_leaky_relu(out, self.bias * self.lr_mul)
+
+ else:
+ out = F.linear(
+ input, self.weight * self.scale, bias=self.bias * self.lr_mul
+ )
+
+ return out
+
+ def __repr__(self):
+ return (
+ f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})'
+ )
+
+
+class ScaledLeakyReLU(nn.Module):
+ def __init__(self, negative_slope=0.2):
+ super().__init__()
+
+ self.negative_slope = negative_slope
+
+ def forward(self, input):
+ out = F.leaky_relu(input, negative_slope=self.negative_slope)
+
+ return out * math.sqrt(2)
+
+
+class ModulatedConv2d(nn.Module):
+ def __init__(
+ self,
+ in_channel,
+ out_channel,
+ kernel_size,
+ style_dim,
+ demodulate=True,
+ upsample=False,
+ downsample=False,
+ blur_kernel=[1, 3, 3, 1],
+ ):
+ super().__init__()
+
+ self.eps = 1e-8
+ self.kernel_size = kernel_size
+ self.in_channel = in_channel
+ self.out_channel = out_channel
+ self.upsample = upsample
+ self.downsample = downsample
+
+ if upsample:
+ factor = 2
+ p = (len(blur_kernel) - factor) - (kernel_size - 1)
+ pad0 = (p + 1) // 2 + factor - 1
+ pad1 = p // 2 + 1
+
+ self.blur = Blur(blur_kernel, pad=(pad0, pad1), upsample_factor=factor)
+
+ if downsample:
+ factor = 2
+ p = (len(blur_kernel) - factor) + (kernel_size - 1)
+ pad0 = (p + 1) // 2
+ pad1 = p // 2
+
+ self.blur = Blur(blur_kernel, pad=(pad0, pad1))
+
+ fan_in = in_channel * kernel_size ** 2
+ self.scale = 1 / math.sqrt(fan_in)
+ self.padding = kernel_size // 2
+
+ self.weight = nn.Parameter(
+ torch.randn(1, out_channel, in_channel, kernel_size, kernel_size)
+ )
+
+ self.modulation = EqualLinear(style_dim, in_channel, bias_init=1)
+
+ self.demodulate = demodulate
+
+ def __repr__(self):
+ return (
+ f'{self.__class__.__name__}({self.in_channel}, {self.out_channel}, {self.kernel_size}, '
+ f'upsample={self.upsample}, downsample={self.downsample})'
+ )
+
+ def forward(self, input, style):
+ batch, in_channel, height, width = input.shape
+
+ style = self.modulation(style).view(batch, 1, in_channel, 1, 1)
+ weight = self.scale * self.weight * style
+
+ if self.demodulate:
+ demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + 1e-8)
+ weight = weight * demod.view(batch, self.out_channel, 1, 1, 1)
+
+ weight = weight.view(
+ batch * self.out_channel, in_channel, self.kernel_size, self.kernel_size
+ )
+
+ if self.upsample:
+ input = input.view(1, batch * in_channel, height, width)
+ weight = weight.view(
+ batch, self.out_channel, in_channel, self.kernel_size, self.kernel_size
+ )
+ weight = weight.transpose(1, 2).reshape(
+ batch * in_channel, self.out_channel, self.kernel_size, self.kernel_size
+ )
+ out = F.conv_transpose2d(input, weight, padding=0, stride=2, groups=batch)
+ _, _, height, width = out.shape
+ out = out.view(batch, self.out_channel, height, width)
+ out = self.blur(out)
+
+ elif self.downsample:
+ input = self.blur(input)
+ _, _, height, width = input.shape
+ input = input.view(1, batch * in_channel, height, width)
+ out = F.conv2d(input, weight, padding=0, stride=2, groups=batch)
+ _, _, height, width = out.shape
+ out = out.view(batch, self.out_channel, height, width)
+
+ else:
+ input = input.view(1, batch * in_channel, height, width)
+ out = F.conv2d(input, weight, padding=self.padding, groups=batch)
+ _, _, height, width = out.shape
+ out = out.view(batch, self.out_channel, height, width)
+
+ return out
+
+
+class NoiseInjection(nn.Module):
+ def __init__(self):
+ super().__init__()
+
+ self.weight = nn.Parameter(torch.zeros(1))
+
+ def forward(self, image, noise=None):
+ if noise is None:
+ batch, _, height, width = image.shape
+ noise = image.new_empty(batch, 1, height, width).normal_()
+
+ return image + self.weight * noise
+
+
+class ConstantInput(nn.Module):
+ def __init__(self, channel, size=4):
+ super().__init__()
+
+ self.input = nn.Parameter(torch.randn(1, channel, size, size))
+
+ def forward(self, input):
+ batch = input.shape[0]
+ out = self.input.repeat(batch, 1, 1, 1)
+
+ return out
+
+
+class StyledConv(nn.Module):
+ def __init__(
+ self,
+ in_channel,
+ out_channel,
+ kernel_size,
+ style_dim,
+ upsample=False,
+ blur_kernel=[1, 3, 3, 1],
+ demodulate=True,
+ ):
+ super().__init__()
+
+ self.conv = ModulatedConv2d(
+ in_channel,
+ out_channel,
+ kernel_size,
+ style_dim,
+ upsample=upsample,
+ blur_kernel=blur_kernel,
+ demodulate=demodulate,
+ )
+
+ self.noise = NoiseInjection()
+ # self.bias = nn.Parameter(torch.zeros(1, out_channel, 1, 1))
+ # self.activate = ScaledLeakyReLU(0.2)
+ self.activate = FusedLeakyReLU(out_channel)
+
+ def forward(self, input, style, noise=None):
+ out = self.conv(input, style)
+ out = self.noise(out, noise=noise)
+ # out = out + self.bias
+ out = self.activate(out)
+
+ return out
+
+
+class ToRGB(nn.Module):
+ def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[1, 3, 3, 1]):
+ super().__init__()
+
+ if upsample:
+ self.upsample = Upsample(blur_kernel)
+
+ self.conv = ModulatedConv2d(in_channel, 3, 1, style_dim, demodulate=False)
+ self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1))
+
+ def forward(self, input, style, skip=None):
+ out = self.conv(input, style)
+ out = out + self.bias
+
+ if skip is not None:
+ skip = self.upsample(skip)
+
+ out = out + skip
+
+ return out
+
+
+class Generator(nn.Module):
+ def __init__(
+ self,
+ size,
+ style_dim,
+ n_mlp,
+ channel_multiplier=2,
+ blur_kernel=[1, 3, 3, 1],
+ lr_mlp=0.01,
+ ):
+ super().__init__()
+
+ self.size = size
+
+ self.style_dim = style_dim
+
+ layers = [PixelNorm()]
+
+ for i in range(n_mlp):
+ layers.append(
+ EqualLinear(
+ style_dim, style_dim, lr_mul=lr_mlp, activation='fused_lrelu'
+ )
+ )
+
+ self.style = nn.Sequential(*layers)
+
+ self.channels = {
+ 4: 512,
+ 8: 512,
+ 16: 512,
+ 32: 512,
+ 64: 256 * channel_multiplier,
+ 128: 128 * channel_multiplier,
+ 256: 64 * channel_multiplier,
+ 512: 32 * channel_multiplier,
+ 1024: 16 * channel_multiplier,
+ }
+
+ self.input = ConstantInput(self.channels[4])
+ self.conv1 = StyledConv(
+ self.channels[4], self.channels[4], 3, style_dim, blur_kernel=blur_kernel
+ )
+ self.to_rgb1 = ToRGB(self.channels[4], style_dim, upsample=False)
+
+ self.log_size = int(math.log(size, 2))
+ self.num_layers = (self.log_size - 2) * 2 + 1
+
+ self.convs = nn.ModuleList()
+ self.upsamples = nn.ModuleList()
+ self.to_rgbs = nn.ModuleList()
+ self.noises = nn.Module()
+
+ in_channel = self.channels[4]
+
+ for layer_idx in range(self.num_layers):
+ res = (layer_idx + 5) // 2
+ shape = [1, 1, 2 ** res, 2 ** res]
+ self.noises.register_buffer(f'noise_{layer_idx}', torch.randn(*shape))
+
+ for i in range(3, self.log_size + 1):
+ out_channel = self.channels[2 ** i]
+
+ self.convs.append(
+ StyledConv(
+ in_channel,
+ out_channel,
+ 3,
+ style_dim,
+ upsample=True,
+ blur_kernel=blur_kernel,
+ )
+ )
+
+ self.convs.append(
+ StyledConv(
+ out_channel, out_channel, 3, style_dim, blur_kernel=blur_kernel
+ )
+ )
+
+ self.to_rgbs.append(ToRGB(out_channel, style_dim))
+
+ in_channel = out_channel
+
+ self.n_latent = self.log_size * 2 - 2
+
+ def make_noise(self):
+ device = self.input.input.device
+
+ noises = [torch.randn(1, 1, 2 ** 2, 2 ** 2, device=device)]
+
+ for i in range(3, self.log_size + 1):
+ for _ in range(2):
+ noises.append(torch.randn(1, 1, 2 ** i, 2 ** i, device=device))
+
+ return noises
+
+ def mean_latent(self, n_latent):
+ latent_in = torch.randn(
+ n_latent, self.style_dim, device=self.input.input.device
+ )
+ latent = self.style(latent_in).mean(0, keepdim=True)
+
+ return latent
+
+ def get_latent(self, input):
+ return self.style(input)
+
+ def forward(
+ self,
+ styles,
+ return_latents=False,
+ return_features=False,
+ inject_index=None,
+ truncation=1,
+ truncation_latent=None,
+ input_is_latent=False,
+ noise=None,
+ randomize_noise=True,
+ ):
+ if not input_is_latent:
+ styles = [self.style(s) for s in styles]
+
+ if noise is None:
+ if randomize_noise:
+ noise = [None] * self.num_layers
+ else:
+ noise = [
+ getattr(self.noises, f'noise_{i}') for i in range(self.num_layers)
+ ]
+
+ if truncation < 1:
+ style_t = []
+
+ for style in styles:
+ style_t.append(
+ truncation_latent + truncation * (style - truncation_latent)
+ )
+
+ styles = style_t
+
+ if len(styles) < 2:
+ inject_index = self.n_latent
+
+ if styles[0].ndim < 3:
+ latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
+ else:
+ latent = styles[0]
+
+ else:
+ if inject_index is None:
+ inject_index = random.randint(1, self.n_latent - 1)
+
+ latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
+ latent2 = styles[1].unsqueeze(1).repeat(1, self.n_latent - inject_index, 1)
+
+ latent = torch.cat([latent, latent2], 1)
+
+ out = self.input(latent)
+ out = self.conv1(out, latent[:, 0], noise=noise[0])
+
+ skip = self.to_rgb1(out, latent[:, 1])
+
+ i = 1
+ for conv1, conv2, noise1, noise2, to_rgb in zip(
+ self.convs[::2], self.convs[1::2], noise[1::2], noise[2::2], self.to_rgbs
+ ):
+ out = conv1(out, latent[:, i], noise=noise1)
+ out = conv2(out, latent[:, i + 1], noise=noise2)
+ skip = to_rgb(out, latent[:, i + 2], skip)
+
+ i += 2
+
+ image = skip
+
+ if return_latents:
+ return image, latent
+ elif return_features:
+ return image, out
+ else:
+ return image, None
+
+
+class ConvLayer(nn.Sequential):
+ def __init__(
+ self,
+ in_channel,
+ out_channel,
+ kernel_size,
+ downsample=False,
+ blur_kernel=[1, 3, 3, 1],
+ bias=True,
+ activate=True,
+ ):
+ layers = []
+
+ if downsample:
+ factor = 2
+ p = (len(blur_kernel) - factor) + (kernel_size - 1)
+ pad0 = (p + 1) // 2
+ pad1 = p // 2
+
+ layers.append(Blur(blur_kernel, pad=(pad0, pad1)))
+
+ stride = 2
+ self.padding = 0
+
+ else:
+ stride = 1
+ self.padding = kernel_size // 2
+
+ layers.append(
+ EqualConv2d(
+ in_channel,
+ out_channel,
+ kernel_size,
+ padding=self.padding,
+ stride=stride,
+ bias=bias and not activate,
+ )
+ )
+
+ if activate:
+ if bias:
+ layers.append(FusedLeakyReLU(out_channel))
+
+ else:
+ layers.append(ScaledLeakyReLU(0.2))
+
+ super().__init__(*layers)
+
+
+class ResBlock(nn.Module):
+ def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]):
+ super().__init__()
+
+ self.conv1 = ConvLayer(in_channel, in_channel, 3)
+ self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True)
+
+ self.skip = ConvLayer(
+ in_channel, out_channel, 1, downsample=True, activate=False, bias=False
+ )
+
+ def forward(self, input):
+ out = self.conv1(input)
+ out = self.conv2(out)
+
+ skip = self.skip(input)
+ out = (out + skip) / math.sqrt(2)
+
+ return out
+
+
+class Discriminator(nn.Module):
+ def __init__(self, size, channel_multiplier=2, blur_kernel=[1, 3, 3, 1]):
+ super().__init__()
+
+ channels = {
+ 4: 512,
+ 8: 512,
+ 16: 512,
+ 32: 512,
+ 64: 256 * channel_multiplier,
+ 128: 128 * channel_multiplier,
+ 256: 64 * channel_multiplier,
+ 512: 32 * channel_multiplier,
+ 1024: 16 * channel_multiplier,
+ }
+
+ convs = [ConvLayer(3, channels[size], 1)]
+
+ log_size = int(math.log(size, 2))
+
+ in_channel = channels[size]
+
+ for i in range(log_size, 2, -1):
+ out_channel = channels[2 ** (i - 1)]
+
+ convs.append(ResBlock(in_channel, out_channel, blur_kernel))
+
+ in_channel = out_channel
+
+ self.convs = nn.Sequential(*convs)
+
+ self.stddev_group = 4
+ self.stddev_feat = 1
+
+ self.final_conv = ConvLayer(in_channel + 1, channels[4], 3)
+ self.final_linear = nn.Sequential(
+ EqualLinear(channels[4] * 4 * 4, channels[4], activation='fused_lrelu'),
+ EqualLinear(channels[4], 1),
+ )
+
+ def forward(self, input):
+ out = self.convs(input)
+
+ batch, channel, height, width = out.shape
+ group = min(batch, self.stddev_group)
+ stddev = out.view(
+ group, -1, self.stddev_feat, channel // self.stddev_feat, height, width
+ )
+ stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-8)
+ stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
+ stddev = stddev.repeat(group, 1, height, width)
+ out = torch.cat([out, stddev], 1)
+
+ out = self.final_conv(out)
+
+ out = out.view(batch, -1)
+ out = self.final_linear(out)
+
+ return out
diff --git a/e4e/models/stylegan2/op/__init__.py b/e4e/models/stylegan2/op/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/e4e/models/stylegan2/op/fused_act.py b/e4e/models/stylegan2/op/fused_act.py
new file mode 100644
index 0000000000000000000000000000000000000000..973a84fffde53668d31397da5fb993bbc95f7be0
--- /dev/null
+++ b/e4e/models/stylegan2/op/fused_act.py
@@ -0,0 +1,85 @@
+import os
+
+import torch
+from torch import nn
+from torch.autograd import Function
+from torch.utils.cpp_extension import load
+
+module_path = os.path.dirname(__file__)
+fused = load(
+ 'fused',
+ sources=[
+ os.path.join(module_path, 'fused_bias_act.cpp'),
+ os.path.join(module_path, 'fused_bias_act_kernel.cu'),
+ ],
+)
+
+
+class FusedLeakyReLUFunctionBackward(Function):
+ @staticmethod
+ def forward(ctx, grad_output, out, negative_slope, scale):
+ ctx.save_for_backward(out)
+ ctx.negative_slope = negative_slope
+ ctx.scale = scale
+
+ empty = grad_output.new_empty(0)
+
+ grad_input = fused.fused_bias_act(
+ grad_output, empty, out, 3, 1, negative_slope, scale
+ )
+
+ dim = [0]
+
+ if grad_input.ndim > 2:
+ dim += list(range(2, grad_input.ndim))
+
+ grad_bias = grad_input.sum(dim).detach()
+
+ return grad_input, grad_bias
+
+ @staticmethod
+ def backward(ctx, gradgrad_input, gradgrad_bias):
+ out, = ctx.saved_tensors
+ gradgrad_out = fused.fused_bias_act(
+ gradgrad_input, gradgrad_bias, out, 3, 1, ctx.negative_slope, ctx.scale
+ )
+
+ return gradgrad_out, None, None, None
+
+
+class FusedLeakyReLUFunction(Function):
+ @staticmethod
+ def forward(ctx, input, bias, negative_slope, scale):
+ empty = input.new_empty(0)
+ out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale)
+ ctx.save_for_backward(out)
+ ctx.negative_slope = negative_slope
+ ctx.scale = scale
+
+ return out
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ out, = ctx.saved_tensors
+
+ grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply(
+ grad_output, out, ctx.negative_slope, ctx.scale
+ )
+
+ return grad_input, grad_bias, None, None
+
+
+class FusedLeakyReLU(nn.Module):
+ def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5):
+ super().__init__()
+
+ self.bias = nn.Parameter(torch.zeros(channel))
+ self.negative_slope = negative_slope
+ self.scale = scale
+
+ def forward(self, input):
+ return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale)
+
+
+def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5):
+ return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale)
diff --git a/e4e/models/stylegan2/op/fused_bias_act.cpp b/e4e/models/stylegan2/op/fused_bias_act.cpp
new file mode 100644
index 0000000000000000000000000000000000000000..02be898f970bcc8ea297867fcaa4e71b24b3d949
--- /dev/null
+++ b/e4e/models/stylegan2/op/fused_bias_act.cpp
@@ -0,0 +1,21 @@
+#include
+
+
+torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,
+ int act, int grad, float alpha, float scale);
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+torch::Tensor fused_bias_act(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,
+ int act, int grad, float alpha, float scale) {
+ CHECK_CUDA(input);
+ CHECK_CUDA(bias);
+
+ return fused_bias_act_op(input, bias, refer, act, grad, alpha, scale);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+ m.def("fused_bias_act", &fused_bias_act, "fused bias act (CUDA)");
+}
\ No newline at end of file
diff --git a/e4e/models/stylegan2/op/fused_bias_act_kernel.cu b/e4e/models/stylegan2/op/fused_bias_act_kernel.cu
new file mode 100644
index 0000000000000000000000000000000000000000..c9fa56fea7ede7072dc8925cfb0148f136eb85b8
--- /dev/null
+++ b/e4e/models/stylegan2/op/fused_bias_act_kernel.cu
@@ -0,0 +1,99 @@
+// Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+//
+// This work is made available under the Nvidia Source Code License-NC.
+// To view a copy of this license, visit
+// https://nvlabs.github.io/stylegan2/license.html
+
+#include
+
+#include
+#include
+#include
+#include
+
+#include
+#include
+
+
+template
+static __global__ void fused_bias_act_kernel(scalar_t* out, const scalar_t* p_x, const scalar_t* p_b, const scalar_t* p_ref,
+ int act, int grad, scalar_t alpha, scalar_t scale, int loop_x, int size_x, int step_b, int size_b, int use_bias, int use_ref) {
+ int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x;
+
+ scalar_t zero = 0.0;
+
+ for (int loop_idx = 0; loop_idx < loop_x && xi < size_x; loop_idx++, xi += blockDim.x) {
+ scalar_t x = p_x[xi];
+
+ if (use_bias) {
+ x += p_b[(xi / step_b) % size_b];
+ }
+
+ scalar_t ref = use_ref ? p_ref[xi] : zero;
+
+ scalar_t y;
+
+ switch (act * 10 + grad) {
+ default:
+ case 10: y = x; break;
+ case 11: y = x; break;
+ case 12: y = 0.0; break;
+
+ case 30: y = (x > 0.0) ? x : x * alpha; break;
+ case 31: y = (ref > 0.0) ? x : x * alpha; break;
+ case 32: y = 0.0; break;
+ }
+
+ out[xi] = y * scale;
+ }
+}
+
+
+torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,
+ int act, int grad, float alpha, float scale) {
+ int curDevice = -1;
+ cudaGetDevice(&curDevice);
+ cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice);
+
+ auto x = input.contiguous();
+ auto b = bias.contiguous();
+ auto ref = refer.contiguous();
+
+ int use_bias = b.numel() ? 1 : 0;
+ int use_ref = ref.numel() ? 1 : 0;
+
+ int size_x = x.numel();
+ int size_b = b.numel();
+ int step_b = 1;
+
+ for (int i = 1 + 1; i < x.dim(); i++) {
+ step_b *= x.size(i);
+ }
+
+ int loop_x = 4;
+ int block_size = 4 * 32;
+ int grid_size = (size_x - 1) / (loop_x * block_size) + 1;
+
+ auto y = torch::empty_like(x);
+
+ AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "fused_bias_act_kernel", [&] {
+ fused_bias_act_kernel<<>>(
+ y.data_ptr(),
+ x.data_ptr(),
+ b.data_ptr(),
+ ref.data_ptr(),
+ act,
+ grad,
+ alpha,
+ scale,
+ loop_x,
+ size_x,
+ step_b,
+ size_b,
+ use_bias,
+ use_ref
+ );
+ });
+
+ return y;
+}
\ No newline at end of file
diff --git a/e4e/models/stylegan2/op/upfirdn2d.cpp b/e4e/models/stylegan2/op/upfirdn2d.cpp
new file mode 100644
index 0000000000000000000000000000000000000000..d2e633dc896433c205e18bc3e455539192ff968e
--- /dev/null
+++ b/e4e/models/stylegan2/op/upfirdn2d.cpp
@@ -0,0 +1,23 @@
+#include
+
+
+torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel,
+ int up_x, int up_y, int down_x, int down_y,
+ int pad_x0, int pad_x1, int pad_y0, int pad_y1);
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+torch::Tensor upfirdn2d(const torch::Tensor& input, const torch::Tensor& kernel,
+ int up_x, int up_y, int down_x, int down_y,
+ int pad_x0, int pad_x1, int pad_y0, int pad_y1) {
+ CHECK_CUDA(input);
+ CHECK_CUDA(kernel);
+
+ return upfirdn2d_op(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+ m.def("upfirdn2d", &upfirdn2d, "upfirdn2d (CUDA)");
+}
\ No newline at end of file
diff --git a/e4e/models/stylegan2/op/upfirdn2d.py b/e4e/models/stylegan2/op/upfirdn2d.py
new file mode 100644
index 0000000000000000000000000000000000000000..7bc5a1e331c2bbb1893ac748cfd0f144ff0651b4
--- /dev/null
+++ b/e4e/models/stylegan2/op/upfirdn2d.py
@@ -0,0 +1,184 @@
+import os
+
+import torch
+from torch.autograd import Function
+from torch.utils.cpp_extension import load
+
+module_path = os.path.dirname(__file__)
+upfirdn2d_op = load(
+ 'upfirdn2d',
+ sources=[
+ os.path.join(module_path, 'upfirdn2d.cpp'),
+ os.path.join(module_path, 'upfirdn2d_kernel.cu'),
+ ],
+)
+
+
+class UpFirDn2dBackward(Function):
+ @staticmethod
+ def forward(
+ ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, in_size, out_size
+ ):
+ up_x, up_y = up
+ down_x, down_y = down
+ g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1 = g_pad
+
+ grad_output = grad_output.reshape(-1, out_size[0], out_size[1], 1)
+
+ grad_input = upfirdn2d_op.upfirdn2d(
+ grad_output,
+ grad_kernel,
+ down_x,
+ down_y,
+ up_x,
+ up_y,
+ g_pad_x0,
+ g_pad_x1,
+ g_pad_y0,
+ g_pad_y1,
+ )
+ grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], in_size[3])
+
+ ctx.save_for_backward(kernel)
+
+ pad_x0, pad_x1, pad_y0, pad_y1 = pad
+
+ ctx.up_x = up_x
+ ctx.up_y = up_y
+ ctx.down_x = down_x
+ ctx.down_y = down_y
+ ctx.pad_x0 = pad_x0
+ ctx.pad_x1 = pad_x1
+ ctx.pad_y0 = pad_y0
+ ctx.pad_y1 = pad_y1
+ ctx.in_size = in_size
+ ctx.out_size = out_size
+
+ return grad_input
+
+ @staticmethod
+ def backward(ctx, gradgrad_input):
+ kernel, = ctx.saved_tensors
+
+ gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], ctx.in_size[3], 1)
+
+ gradgrad_out = upfirdn2d_op.upfirdn2d(
+ gradgrad_input,
+ kernel,
+ ctx.up_x,
+ ctx.up_y,
+ ctx.down_x,
+ ctx.down_y,
+ ctx.pad_x0,
+ ctx.pad_x1,
+ ctx.pad_y0,
+ ctx.pad_y1,
+ )
+ # gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.out_size[0], ctx.out_size[1], ctx.in_size[3])
+ gradgrad_out = gradgrad_out.view(
+ ctx.in_size[0], ctx.in_size[1], ctx.out_size[0], ctx.out_size[1]
+ )
+
+ return gradgrad_out, None, None, None, None, None, None, None, None
+
+
+class UpFirDn2d(Function):
+ @staticmethod
+ def forward(ctx, input, kernel, up, down, pad):
+ up_x, up_y = up
+ down_x, down_y = down
+ pad_x0, pad_x1, pad_y0, pad_y1 = pad
+
+ kernel_h, kernel_w = kernel.shape
+ batch, channel, in_h, in_w = input.shape
+ ctx.in_size = input.shape
+
+ input = input.reshape(-1, in_h, in_w, 1)
+
+ ctx.save_for_backward(kernel, torch.flip(kernel, [0, 1]))
+
+ out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1
+ out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1
+ ctx.out_size = (out_h, out_w)
+
+ ctx.up = (up_x, up_y)
+ ctx.down = (down_x, down_y)
+ ctx.pad = (pad_x0, pad_x1, pad_y0, pad_y1)
+
+ g_pad_x0 = kernel_w - pad_x0 - 1
+ g_pad_y0 = kernel_h - pad_y0 - 1
+ g_pad_x1 = in_w * up_x - out_w * down_x + pad_x0 - up_x + 1
+ g_pad_y1 = in_h * up_y - out_h * down_y + pad_y0 - up_y + 1
+
+ ctx.g_pad = (g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1)
+
+ out = upfirdn2d_op.upfirdn2d(
+ input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
+ )
+ # out = out.view(major, out_h, out_w, minor)
+ out = out.view(-1, channel, out_h, out_w)
+
+ return out
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ kernel, grad_kernel = ctx.saved_tensors
+
+ grad_input = UpFirDn2dBackward.apply(
+ grad_output,
+ kernel,
+ grad_kernel,
+ ctx.up,
+ ctx.down,
+ ctx.pad,
+ ctx.g_pad,
+ ctx.in_size,
+ ctx.out_size,
+ )
+
+ return grad_input, None, None, None, None
+
+
+def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
+ out = UpFirDn2d.apply(
+ input, kernel, (up, up), (down, down), (pad[0], pad[1], pad[0], pad[1])
+ )
+
+ return out
+
+
+def upfirdn2d_native(
+ input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
+):
+ _, in_h, in_w, minor = input.shape
+ kernel_h, kernel_w = kernel.shape
+
+ out = input.view(-1, in_h, 1, in_w, 1, minor)
+ out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1])
+ out = out.view(-1, in_h * up_y, in_w * up_x, minor)
+
+ out = F.pad(
+ out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)]
+ )
+ out = out[
+ :,
+ max(-pad_y0, 0): out.shape[1] - max(-pad_y1, 0),
+ max(-pad_x0, 0): out.shape[2] - max(-pad_x1, 0),
+ :,
+ ]
+
+ out = out.permute(0, 3, 1, 2)
+ out = out.reshape(
+ [-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1]
+ )
+ w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
+ out = F.conv2d(out, w)
+ out = out.reshape(
+ -1,
+ minor,
+ in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
+ in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1,
+ )
+ out = out.permute(0, 2, 3, 1)
+
+ return out[:, ::down_y, ::down_x, :]
diff --git a/e4e/models/stylegan2/op/upfirdn2d_kernel.cu b/e4e/models/stylegan2/op/upfirdn2d_kernel.cu
new file mode 100644
index 0000000000000000000000000000000000000000..2a710aa6adc3d43ac93136a1814e3c39970e1c7e
--- /dev/null
+++ b/e4e/models/stylegan2/op/upfirdn2d_kernel.cu
@@ -0,0 +1,272 @@
+// Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+//
+// This work is made available under the Nvidia Source Code License-NC.
+// To view a copy of this license, visit
+// https://nvlabs.github.io/stylegan2/license.html
+
+#include
+
+#include
+#include
+#include
+#include
+
+#include
+#include
+
+
+static __host__ __device__ __forceinline__ int floor_div(int a, int b) {
+ int c = a / b;
+
+ if (c * b > a) {
+ c--;
+ }
+
+ return c;
+}
+
+
+struct UpFirDn2DKernelParams {
+ int up_x;
+ int up_y;
+ int down_x;
+ int down_y;
+ int pad_x0;
+ int pad_x1;
+ int pad_y0;
+ int pad_y1;
+
+ int major_dim;
+ int in_h;
+ int in_w;
+ int minor_dim;
+ int kernel_h;
+ int kernel_w;
+ int out_h;
+ int out_w;
+ int loop_major;
+ int loop_x;
+};
+
+
+template
+__global__ void upfirdn2d_kernel(scalar_t* out, const scalar_t* input, const scalar_t* kernel, const UpFirDn2DKernelParams p) {
+ const int tile_in_h = ((tile_out_h - 1) * down_y + kernel_h - 1) / up_y + 1;
+ const int tile_in_w = ((tile_out_w - 1) * down_x + kernel_w - 1) / up_x + 1;
+
+ __shared__ volatile float sk[kernel_h][kernel_w];
+ __shared__ volatile float sx[tile_in_h][tile_in_w];
+
+ int minor_idx = blockIdx.x;
+ int tile_out_y = minor_idx / p.minor_dim;
+ minor_idx -= tile_out_y * p.minor_dim;
+ tile_out_y *= tile_out_h;
+ int tile_out_x_base = blockIdx.y * p.loop_x * tile_out_w;
+ int major_idx_base = blockIdx.z * p.loop_major;
+
+ if (tile_out_x_base >= p.out_w | tile_out_y >= p.out_h | major_idx_base >= p.major_dim) {
+ return;
+ }
+
+ for (int tap_idx = threadIdx.x; tap_idx < kernel_h * kernel_w; tap_idx += blockDim.x) {
+ int ky = tap_idx / kernel_w;
+ int kx = tap_idx - ky * kernel_w;
+ scalar_t v = 0.0;
+
+ if (kx < p.kernel_w & ky < p.kernel_h) {
+ v = kernel[(p.kernel_h - 1 - ky) * p.kernel_w + (p.kernel_w - 1 - kx)];
+ }
+
+ sk[ky][kx] = v;
+ }
+
+ for (int loop_major = 0, major_idx = major_idx_base; loop_major < p.loop_major & major_idx < p.major_dim; loop_major++, major_idx++) {
+ for (int loop_x = 0, tile_out_x = tile_out_x_base; loop_x < p.loop_x & tile_out_x < p.out_w; loop_x++, tile_out_x += tile_out_w) {
+ int tile_mid_x = tile_out_x * down_x + up_x - 1 - p.pad_x0;
+ int tile_mid_y = tile_out_y * down_y + up_y - 1 - p.pad_y0;
+ int tile_in_x = floor_div(tile_mid_x, up_x);
+ int tile_in_y = floor_div(tile_mid_y, up_y);
+
+ __syncthreads();
+
+ for (int in_idx = threadIdx.x; in_idx < tile_in_h * tile_in_w; in_idx += blockDim.x) {
+ int rel_in_y = in_idx / tile_in_w;
+ int rel_in_x = in_idx - rel_in_y * tile_in_w;
+ int in_x = rel_in_x + tile_in_x;
+ int in_y = rel_in_y + tile_in_y;
+
+ scalar_t v = 0.0;
+
+ if (in_x >= 0 & in_y >= 0 & in_x < p.in_w & in_y < p.in_h) {
+ v = input[((major_idx * p.in_h + in_y) * p.in_w + in_x) * p.minor_dim + minor_idx];
+ }
+
+ sx[rel_in_y][rel_in_x] = v;
+ }
+
+ __syncthreads();
+ for (int out_idx = threadIdx.x; out_idx < tile_out_h * tile_out_w; out_idx += blockDim.x) {
+ int rel_out_y = out_idx / tile_out_w;
+ int rel_out_x = out_idx - rel_out_y * tile_out_w;
+ int out_x = rel_out_x + tile_out_x;
+ int out_y = rel_out_y + tile_out_y;
+
+ int mid_x = tile_mid_x + rel_out_x * down_x;
+ int mid_y = tile_mid_y + rel_out_y * down_y;
+ int in_x = floor_div(mid_x, up_x);
+ int in_y = floor_div(mid_y, up_y);
+ int rel_in_x = in_x - tile_in_x;
+ int rel_in_y = in_y - tile_in_y;
+ int kernel_x = (in_x + 1) * up_x - mid_x - 1;
+ int kernel_y = (in_y + 1) * up_y - mid_y - 1;
+
+ scalar_t v = 0.0;
+
+ #pragma unroll
+ for (int y = 0; y < kernel_h / up_y; y++)
+ #pragma unroll
+ for (int x = 0; x < kernel_w / up_x; x++)
+ v += sx[rel_in_y + y][rel_in_x + x] * sk[kernel_y + y * up_y][kernel_x + x * up_x];
+
+ if (out_x < p.out_w & out_y < p.out_h) {
+ out[((major_idx * p.out_h + out_y) * p.out_w + out_x) * p.minor_dim + minor_idx] = v;
+ }
+ }
+ }
+ }
+}
+
+
+torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel,
+ int up_x, int up_y, int down_x, int down_y,
+ int pad_x0, int pad_x1, int pad_y0, int pad_y1) {
+ int curDevice = -1;
+ cudaGetDevice(&curDevice);
+ cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice);
+
+ UpFirDn2DKernelParams p;
+
+ auto x = input.contiguous();
+ auto k = kernel.contiguous();
+
+ p.major_dim = x.size(0);
+ p.in_h = x.size(1);
+ p.in_w = x.size(2);
+ p.minor_dim = x.size(3);
+ p.kernel_h = k.size(0);
+ p.kernel_w = k.size(1);
+ p.up_x = up_x;
+ p.up_y = up_y;
+ p.down_x = down_x;
+ p.down_y = down_y;
+ p.pad_x0 = pad_x0;
+ p.pad_x1 = pad_x1;
+ p.pad_y0 = pad_y0;
+ p.pad_y1 = pad_y1;
+
+ p.out_h = (p.in_h * p.up_y + p.pad_y0 + p.pad_y1 - p.kernel_h + p.down_y) / p.down_y;
+ p.out_w = (p.in_w * p.up_x + p.pad_x0 + p.pad_x1 - p.kernel_w + p.down_x) / p.down_x;
+
+ auto out = at::empty({p.major_dim, p.out_h, p.out_w, p.minor_dim}, x.options());
+
+ int mode = -1;
+
+ int tile_out_h;
+ int tile_out_w;
+
+ if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 4 && p.kernel_w <= 4) {
+ mode = 1;
+ tile_out_h = 16;
+ tile_out_w = 64;
+ }
+
+ if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 3 && p.kernel_w <= 3) {
+ mode = 2;
+ tile_out_h = 16;
+ tile_out_w = 64;
+ }
+
+ if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 4 && p.kernel_w <= 4) {
+ mode = 3;
+ tile_out_h = 16;
+ tile_out_w = 64;
+ }
+
+ if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 2 && p.kernel_w <= 2) {
+ mode = 4;
+ tile_out_h = 16;
+ tile_out_w = 64;
+ }
+
+ if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 && p.kernel_h <= 4 && p.kernel_w <= 4) {
+ mode = 5;
+ tile_out_h = 8;
+ tile_out_w = 32;
+ }
+
+ if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 && p.kernel_h <= 2 && p.kernel_w <= 2) {
+ mode = 6;
+ tile_out_h = 8;
+ tile_out_w = 32;
+ }
+
+ dim3 block_size;
+ dim3 grid_size;
+
+ if (tile_out_h > 0 && tile_out_w) {
+ p.loop_major = (p.major_dim - 1) / 16384 + 1;
+ p.loop_x = 1;
+ block_size = dim3(32 * 8, 1, 1);
+ grid_size = dim3(((p.out_h - 1) / tile_out_h + 1) * p.minor_dim,
+ (p.out_w - 1) / (p.loop_x * tile_out_w) + 1,
+ (p.major_dim - 1) / p.loop_major + 1);
+ }
+
+ AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "upfirdn2d_cuda", [&] {
+ switch (mode) {
+ case 1:
+ upfirdn2d_kernel<<>>(
+ out.data_ptr(), x.data_ptr(), k.data_ptr(), p
+ );
+
+ break;
+
+ case 2:
+ upfirdn2d_kernel<<>>(
+ out.data_ptr(), x.data_ptr(), k.data_ptr(), p
+ );
+
+ break;
+
+ case 3:
+ upfirdn2d_kernel<<>>(
+ out.data_ptr(), x.data_ptr(), k.data_ptr(), p
+ );
+
+ break;
+
+ case 4:
+ upfirdn2d_kernel<<>>(
+ out.data_ptr(), x.data_ptr(), k.data_ptr(), p
+ );
+
+ break;
+
+ case 5:
+ upfirdn2d_kernel<<>>(
+ out.data_ptr(), x.data_ptr(), k.data_ptr(), p
+ );
+
+ break;
+
+ case 6:
+ upfirdn2d_kernel<<>>(
+ out.data_ptr(), x.data_ptr(), k.data_ptr(), p
+ );
+
+ break;
+ }
+ });
+
+ return out;
+}
\ No newline at end of file
diff --git a/e4e/options/__init__.py b/e4e/options/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/e4e/options/train_options.py b/e4e/options/train_options.py
new file mode 100644
index 0000000000000000000000000000000000000000..583ea1423fdc9a649cd7044d74d554bf0ac2bf51
--- /dev/null
+++ b/e4e/options/train_options.py
@@ -0,0 +1,84 @@
+from argparse import ArgumentParser
+from configs.paths_config import model_paths
+
+
+class TrainOptions:
+
+ def __init__(self):
+ self.parser = ArgumentParser()
+ self.initialize()
+
+ def initialize(self):
+ self.parser.add_argument('--exp_dir', type=str, help='Path to experiment output directory')
+ self.parser.add_argument('--dataset_type', default='ffhq_encode', type=str,
+ help='Type of dataset/experiment to run')
+ self.parser.add_argument('--encoder_type', default='Encoder4Editing', type=str, help='Which encoder to use')
+
+ self.parser.add_argument('--batch_size', default=4, type=int, help='Batch size for training')
+ self.parser.add_argument('--test_batch_size', default=2, type=int, help='Batch size for testing and inference')
+ self.parser.add_argument('--workers', default=4, type=int, help='Number of train dataloader workers')
+ self.parser.add_argument('--test_workers', default=2, type=int,
+ help='Number of test/inference dataloader workers')
+
+ self.parser.add_argument('--learning_rate', default=0.0001, type=float, help='Optimizer learning rate')
+ self.parser.add_argument('--optim_name', default='ranger', type=str, help='Which optimizer to use')
+ self.parser.add_argument('--train_decoder', default=False, type=bool, help='Whether to train the decoder model')
+ self.parser.add_argument('--start_from_latent_avg', action='store_true',
+ help='Whether to add average latent vector to generate codes from encoder.')
+ self.parser.add_argument('--lpips_type', default='alex', type=str, help='LPIPS backbone')
+
+ self.parser.add_argument('--lpips_lambda', default=0.8, type=float, help='LPIPS loss multiplier factor')
+ self.parser.add_argument('--id_lambda', default=0.1, type=float, help='ID loss multiplier factor')
+ self.parser.add_argument('--l2_lambda', default=1.0, type=float, help='L2 loss multiplier factor')
+
+ self.parser.add_argument('--stylegan_weights', default=model_paths['stylegan_ffhq'], type=str,
+ help='Path to StyleGAN model weights')
+ self.parser.add_argument('--stylegan_size', default=1024, type=int,
+ help='size of pretrained StyleGAN Generator')
+ self.parser.add_argument('--checkpoint_path', default=None, type=str, help='Path to pSp model checkpoint')
+
+ self.parser.add_argument('--max_steps', default=500000, type=int, help='Maximum number of training steps')
+ self.parser.add_argument('--image_interval', default=100, type=int,
+ help='Interval for logging train images during training')
+ self.parser.add_argument('--board_interval', default=50, type=int,
+ help='Interval for logging metrics to tensorboard')
+ self.parser.add_argument('--val_interval', default=1000, type=int, help='Validation interval')
+ self.parser.add_argument('--save_interval', default=None, type=int, help='Model checkpoint interval')
+
+ # Discriminator flags
+ self.parser.add_argument('--w_discriminator_lambda', default=0, type=float, help='Dw loss multiplier')
+ self.parser.add_argument('--w_discriminator_lr', default=2e-5, type=float, help='Dw learning rate')
+ self.parser.add_argument("--r1", type=float, default=10, help="weight of the r1 regularization")
+ self.parser.add_argument("--d_reg_every", type=int, default=16,
+ help="interval for applying r1 regularization")
+ self.parser.add_argument('--use_w_pool', action='store_true',
+ help='Whether to store a latnet codes pool for the discriminator\'s training')
+ self.parser.add_argument("--w_pool_size", type=int, default=50,
+ help="W\'s pool size, depends on --use_w_pool")
+
+ # e4e specific
+ self.parser.add_argument('--delta_norm', type=int, default=2, help="norm type of the deltas")
+ self.parser.add_argument('--delta_norm_lambda', type=float, default=2e-4, help="lambda for delta norm loss")
+
+ # Progressive training
+ self.parser.add_argument('--progressive_steps', nargs='+', type=int, default=None,
+ help="The training steps of training new deltas. steps[i] starts the delta_i training")
+ self.parser.add_argument('--progressive_start', type=int, default=None,
+ help="The training step to start training the deltas, overrides progressive_steps")
+ self.parser.add_argument('--progressive_step_every', type=int, default=2_000,
+ help="Amount of training steps for each progressive step")
+
+ # Save additional training info to enable future training continuation from produced checkpoints
+ self.parser.add_argument('--save_training_data', action='store_true',
+ help='Save intermediate training data to resume training from the checkpoint')
+ self.parser.add_argument('--sub_exp_dir', default=None, type=str, help='Name of sub experiment directory')
+ self.parser.add_argument('--keep_optimizer', action='store_true',
+ help='Whether to continue from the checkpoint\'s optimizer')
+ self.parser.add_argument('--resume_training_from_ckpt', default=None, type=str,
+ help='Path to training checkpoint, works when --save_training_data was set to True')
+ self.parser.add_argument('--update_param_list', nargs='+', type=str, default=None,
+ help="Name of training parameters to update the loaded training checkpoint")
+
+ def parse(self):
+ opts = self.parser.parse_args()
+ return opts
diff --git a/e4e/scripts/calc_losses_on_images.py b/e4e/scripts/calc_losses_on_images.py
new file mode 100644
index 0000000000000000000000000000000000000000..32b6bcee854da7ae357daf82bd986f30db9fb72c
--- /dev/null
+++ b/e4e/scripts/calc_losses_on_images.py
@@ -0,0 +1,87 @@
+from argparse import ArgumentParser
+import os
+import json
+import sys
+from tqdm import tqdm
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+import torchvision.transforms as transforms
+
+sys.path.append(".")
+sys.path.append("..")
+
+from criteria.lpips.lpips import LPIPS
+from datasets.gt_res_dataset import GTResDataset
+
+
+def parse_args():
+ parser = ArgumentParser(add_help=False)
+ parser.add_argument('--mode', type=str, default='lpips', choices=['lpips', 'l2'])
+ parser.add_argument('--data_path', type=str, default='results')
+ parser.add_argument('--gt_path', type=str, default='gt_images')
+ parser.add_argument('--workers', type=int, default=4)
+ parser.add_argument('--batch_size', type=int, default=4)
+ parser.add_argument('--is_cars', action='store_true')
+ args = parser.parse_args()
+ return args
+
+
+def run(args):
+ resize_dims = (256, 256)
+ if args.is_cars:
+ resize_dims = (192, 256)
+ transform = transforms.Compose([transforms.Resize(resize_dims),
+ transforms.ToTensor(),
+ transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+
+ print('Loading dataset')
+ dataset = GTResDataset(root_path=args.data_path,
+ gt_dir=args.gt_path,
+ transform=transform)
+
+ dataloader = DataLoader(dataset,
+ batch_size=args.batch_size,
+ shuffle=False,
+ num_workers=int(args.workers),
+ drop_last=True)
+
+ if args.mode == 'lpips':
+ loss_func = LPIPS(net_type='alex')
+ elif args.mode == 'l2':
+ loss_func = torch.nn.MSELoss()
+ else:
+ raise Exception('Not a valid mode!')
+ loss_func.cuda()
+
+ global_i = 0
+ scores_dict = {}
+ all_scores = []
+ for result_batch, gt_batch in tqdm(dataloader):
+ for i in range(args.batch_size):
+ loss = float(loss_func(result_batch[i:i + 1].cuda(), gt_batch[i:i + 1].cuda()))
+ all_scores.append(loss)
+ im_path = dataset.pairs[global_i][0]
+ scores_dict[os.path.basename(im_path)] = loss
+ global_i += 1
+
+ all_scores = list(scores_dict.values())
+ mean = np.mean(all_scores)
+ std = np.std(all_scores)
+ result_str = 'Average loss is {:.2f}+-{:.2f}'.format(mean, std)
+ print('Finished with ', args.data_path)
+ print(result_str)
+
+ out_path = os.path.join(os.path.dirname(args.data_path), 'inference_metrics')
+ if not os.path.exists(out_path):
+ os.makedirs(out_path)
+
+ with open(os.path.join(out_path, 'stat_{}.txt'.format(args.mode)), 'w') as f:
+ f.write(result_str)
+ with open(os.path.join(out_path, 'scores_{}.json'.format(args.mode)), 'w') as f:
+ json.dump(scores_dict, f)
+
+
+if __name__ == '__main__':
+ args = parse_args()
+ run(args)
diff --git a/e4e/scripts/inference.py b/e4e/scripts/inference.py
new file mode 100644
index 0000000000000000000000000000000000000000..185b9b34db85dcd97b9793bd5dbfc9d1ca046549
--- /dev/null
+++ b/e4e/scripts/inference.py
@@ -0,0 +1,133 @@
+import argparse
+
+import torch
+import numpy as np
+import sys
+import os
+import dlib
+
+sys.path.append(".")
+sys.path.append("..")
+
+from configs import data_configs, paths_config
+from datasets.inference_dataset import InferenceDataset
+from torch.utils.data import DataLoader
+from utils.model_utils import setup_model
+from utils.common import tensor2im
+from utils.alignment import align_face
+from PIL import Image
+
+
+def main(args):
+ net, opts = setup_model(args.ckpt, device)
+ is_cars = 'cars_' in opts.dataset_type
+ generator = net.decoder
+ generator.eval()
+ args, data_loader = setup_data_loader(args, opts)
+
+ # Check if latents exist
+ latents_file_path = os.path.join(args.save_dir, 'latents.pt')
+ if os.path.exists(latents_file_path):
+ latent_codes = torch.load(latents_file_path).to(device)
+ else:
+ latent_codes = get_all_latents(net, data_loader, args.n_sample, is_cars=is_cars)
+ torch.save(latent_codes, latents_file_path)
+
+ if not args.latents_only:
+ generate_inversions(args, generator, latent_codes, is_cars=is_cars)
+
+
+def setup_data_loader(args, opts):
+ dataset_args = data_configs.DATASETS[opts.dataset_type]
+ transforms_dict = dataset_args['transforms'](opts).get_transforms()
+ images_path = args.images_dir if args.images_dir is not None else dataset_args['test_source_root']
+ print(f"images path: {images_path}")
+ align_function = None
+ if args.align:
+ align_function = run_alignment
+ test_dataset = InferenceDataset(root=images_path,
+ transform=transforms_dict['transform_test'],
+ preprocess=align_function,
+ opts=opts)
+
+ data_loader = DataLoader(test_dataset,
+ batch_size=args.batch,
+ shuffle=False,
+ num_workers=2,
+ drop_last=True)
+
+ print(f'dataset length: {len(test_dataset)}')
+
+ if args.n_sample is None:
+ args.n_sample = len(test_dataset)
+ return args, data_loader
+
+
+def get_latents(net, x, is_cars=False):
+ codes = net.encoder(x)
+ if net.opts.start_from_latent_avg:
+ if codes.ndim == 2:
+ codes = codes + net.latent_avg.repeat(codes.shape[0], 1, 1)[:, 0, :]
+ else:
+ codes = codes + net.latent_avg.repeat(codes.shape[0], 1, 1)
+ if codes.shape[1] == 18 and is_cars:
+ codes = codes[:, :16, :]
+ return codes
+
+
+def get_all_latents(net, data_loader, n_images=None, is_cars=False):
+ all_latents = []
+ i = 0
+ with torch.no_grad():
+ for batch in data_loader:
+ if n_images is not None and i > n_images:
+ break
+ x = batch
+ inputs = x.to(device).float()
+ latents = get_latents(net, inputs, is_cars)
+ all_latents.append(latents)
+ i += len(latents)
+ return torch.cat(all_latents)
+
+
+def save_image(img, save_dir, idx):
+ result = tensor2im(img)
+ im_save_path = os.path.join(save_dir, f"{idx:05d}.jpg")
+ Image.fromarray(np.array(result)).save(im_save_path)
+
+
+@torch.no_grad()
+def generate_inversions(args, g, latent_codes, is_cars):
+ print('Saving inversion images')
+ inversions_directory_path = os.path.join(args.save_dir, 'inversions')
+ os.makedirs(inversions_directory_path, exist_ok=True)
+ for i in range(args.n_sample):
+ imgs, _ = g([latent_codes[i].unsqueeze(0)], input_is_latent=True, randomize_noise=False, return_latents=True)
+ if is_cars:
+ imgs = imgs[:, :, 64:448, :]
+ save_image(imgs[0], inversions_directory_path, i + 1)
+
+
+def run_alignment(image_path):
+ predictor = dlib.shape_predictor(paths_config.model_paths['shape_predictor'])
+ aligned_image = align_face(filepath=image_path, predictor=predictor)
+ print("Aligned image has shape: {}".format(aligned_image.size))
+ return aligned_image
+
+
+if __name__ == "__main__":
+ device = "cuda"
+
+ parser = argparse.ArgumentParser(description="Inference")
+ parser.add_argument("--images_dir", type=str, default=None,
+ help="The directory of the images to be inverted")
+ parser.add_argument("--save_dir", type=str, default=None,
+ help="The directory to save the latent codes and inversion images. (default: images_dir")
+ parser.add_argument("--batch", type=int, default=1, help="batch size for the generator")
+ parser.add_argument("--n_sample", type=int, default=None, help="number of the samples to infer.")
+ parser.add_argument("--latents_only", action="store_true", help="infer only the latent codes of the directory")
+ parser.add_argument("--align", action="store_true", help="align face images before inference")
+ parser.add_argument("ckpt", metavar="CHECKPOINT", help="path to generator checkpoint")
+
+ args = parser.parse_args()
+ main(args)
diff --git a/e4e/scripts/train.py b/e4e/scripts/train.py
new file mode 100644
index 0000000000000000000000000000000000000000..d885cfde49a0b21140e663e475918698d5e51ee3
--- /dev/null
+++ b/e4e/scripts/train.py
@@ -0,0 +1,88 @@
+"""
+This file runs the main training/val loop
+"""
+import os
+import json
+import math
+import sys
+import pprint
+import torch
+from argparse import Namespace
+
+sys.path.append(".")
+sys.path.append("..")
+
+from options.train_options import TrainOptions
+from training.coach import Coach
+
+
+def main():
+ opts = TrainOptions().parse()
+ previous_train_ckpt = None
+ if opts.resume_training_from_ckpt:
+ opts, previous_train_ckpt = load_train_checkpoint(opts)
+ else:
+ setup_progressive_steps(opts)
+ create_initial_experiment_dir(opts)
+
+ coach = Coach(opts, previous_train_ckpt)
+ coach.train()
+
+
+def load_train_checkpoint(opts):
+ train_ckpt_path = opts.resume_training_from_ckpt
+ previous_train_ckpt = torch.load(opts.resume_training_from_ckpt, map_location='cpu')
+ new_opts_dict = vars(opts)
+ opts = previous_train_ckpt['opts']
+ opts['resume_training_from_ckpt'] = train_ckpt_path
+ update_new_configs(opts, new_opts_dict)
+ pprint.pprint(opts)
+ opts = Namespace(**opts)
+ if opts.sub_exp_dir is not None:
+ sub_exp_dir = opts.sub_exp_dir
+ opts.exp_dir = os.path.join(opts.exp_dir, sub_exp_dir)
+ create_initial_experiment_dir(opts)
+ return opts, previous_train_ckpt
+
+
+def setup_progressive_steps(opts):
+ log_size = int(math.log(opts.stylegan_size, 2))
+ num_style_layers = 2*log_size - 2
+ num_deltas = num_style_layers - 1
+ if opts.progressive_start is not None: # If progressive delta training
+ opts.progressive_steps = [0]
+ next_progressive_step = opts.progressive_start
+ for i in range(num_deltas):
+ opts.progressive_steps.append(next_progressive_step)
+ next_progressive_step += opts.progressive_step_every
+
+ assert opts.progressive_steps is None or is_valid_progressive_steps(opts, num_style_layers), \
+ "Invalid progressive training input"
+
+
+def is_valid_progressive_steps(opts, num_style_layers):
+ return len(opts.progressive_steps) == num_style_layers and opts.progressive_steps[0] == 0
+
+
+def create_initial_experiment_dir(opts):
+ if os.path.exists(opts.exp_dir):
+ raise Exception('Oops... {} already exists'.format(opts.exp_dir))
+ os.makedirs(opts.exp_dir)
+
+ opts_dict = vars(opts)
+ pprint.pprint(opts_dict)
+ with open(os.path.join(opts.exp_dir, 'opt.json'), 'w') as f:
+ json.dump(opts_dict, f, indent=4, sort_keys=True)
+
+
+def update_new_configs(ckpt_opts, new_opts):
+ for k, v in new_opts.items():
+ if k not in ckpt_opts:
+ ckpt_opts[k] = v
+ if new_opts['update_param_list']:
+ for param in new_opts['update_param_list']:
+ ckpt_opts[param] = new_opts[param]
+
+
+if __name__ == '__main__':
+ main()
diff --git a/e4e/utils/__init__.py b/e4e/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/e4e/utils/alignment.py b/e4e/utils/alignment.py
new file mode 100644
index 0000000000000000000000000000000000000000..a02798f0f7c9fdcc319f7884a491b9e6580cc8aa
--- /dev/null
+++ b/e4e/utils/alignment.py
@@ -0,0 +1,115 @@
+import numpy as np
+import PIL
+import PIL.Image
+import scipy
+import scipy.ndimage
+import dlib
+
+
+def get_landmark(filepath, predictor):
+ """get landmark with dlib
+ :return: np.array shape=(68, 2)
+ """
+ detector = dlib.get_frontal_face_detector()
+
+ img = dlib.load_rgb_image(filepath)
+ dets = detector(img, 1)
+
+ for k, d in enumerate(dets):
+ shape = predictor(img, d)
+
+ t = list(shape.parts())
+ a = []
+ for tt in t:
+ a.append([tt.x, tt.y])
+ lm = np.array(a)
+ return lm
+
+
+def align_face(filepath, predictor):
+ """
+ :param filepath: str
+ :return: PIL Image
+ """
+
+ lm = get_landmark(filepath, predictor)
+
+ lm_chin = lm[0: 17] # left-right
+ lm_eyebrow_left = lm[17: 22] # left-right
+ lm_eyebrow_right = lm[22: 27] # left-right
+ lm_nose = lm[27: 31] # top-down
+ lm_nostrils = lm[31: 36] # top-down
+ lm_eye_left = lm[36: 42] # left-clockwise
+ lm_eye_right = lm[42: 48] # left-clockwise
+ lm_mouth_outer = lm[48: 60] # left-clockwise
+ lm_mouth_inner = lm[60: 68] # left-clockwise
+
+ # Calculate auxiliary vectors.
+ eye_left = np.mean(lm_eye_left, axis=0)
+ eye_right = np.mean(lm_eye_right, axis=0)
+ eye_avg = (eye_left + eye_right) * 0.5
+ eye_to_eye = eye_right - eye_left
+ mouth_left = lm_mouth_outer[0]
+ mouth_right = lm_mouth_outer[6]
+ mouth_avg = (mouth_left + mouth_right) * 0.5
+ eye_to_mouth = mouth_avg - eye_avg
+
+ # Choose oriented crop rectangle.
+ x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1]
+ x /= np.hypot(*x)
+ x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8)
+ y = np.flipud(x) * [-1, 1]
+ c = eye_avg + eye_to_mouth * 0.1
+ quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])
+ qsize = np.hypot(*x) * 2
+
+ # read image
+ img = PIL.Image.open(filepath)
+
+ output_size = 256
+ transform_size = 256
+ enable_padding = True
+
+ # Shrink.
+ shrink = int(np.floor(qsize / output_size * 0.5))
+ if shrink > 1:
+ rsize = (int(np.rint(float(img.size[0]) / shrink)), int(np.rint(float(img.size[1]) / shrink)))
+ img = img.resize(rsize, PIL.Image.ANTIALIAS)
+ quad /= shrink
+ qsize /= shrink
+
+ # Crop.
+ border = max(int(np.rint(qsize * 0.1)), 3)
+ crop = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),
+ int(np.ceil(max(quad[:, 1]))))
+ crop = (max(crop[0] - border, 0), max(crop[1] - border, 0), min(crop[2] + border, img.size[0]),
+ min(crop[3] + border, img.size[1]))
+ if crop[2] - crop[0] < img.size[0] or crop[3] - crop[1] < img.size[1]:
+ img = img.crop(crop)
+ quad -= crop[0:2]
+
+ # Pad.
+ pad = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),
+ int(np.ceil(max(quad[:, 1]))))
+ pad = (max(-pad[0] + border, 0), max(-pad[1] + border, 0), max(pad[2] - img.size[0] + border, 0),
+ max(pad[3] - img.size[1] + border, 0))
+ if enable_padding and max(pad) > border - 4:
+ pad = np.maximum(pad, int(np.rint(qsize * 0.3)))
+ img = np.pad(np.float32(img), ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect')
+ h, w, _ = img.shape
+ y, x, _ = np.ogrid[:h, :w, :1]
+ mask = np.maximum(1.0 - np.minimum(np.float32(x) / pad[0], np.float32(w - 1 - x) / pad[2]),
+ 1.0 - np.minimum(np.float32(y) / pad[1], np.float32(h - 1 - y) / pad[3]))
+ blur = qsize * 0.02
+ img += (scipy.ndimage.gaussian_filter(img, [blur, blur, 0]) - img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0)
+ img += (np.median(img, axis=(0, 1)) - img) * np.clip(mask, 0.0, 1.0)
+ img = PIL.Image.fromarray(np.uint8(np.clip(np.rint(img), 0, 255)), 'RGB')
+ quad += pad[:2]
+
+ # Transform.
+ img = img.transform((transform_size, transform_size), PIL.Image.QUAD, (quad + 0.5).flatten(), PIL.Image.BILINEAR)
+ if output_size < transform_size:
+ img = img.resize((output_size, output_size), PIL.Image.ANTIALIAS)
+
+ # Return aligned image.
+ return img
diff --git a/e4e/utils/common.py b/e4e/utils/common.py
new file mode 100644
index 0000000000000000000000000000000000000000..b19e18ddcb78b06678fa18e4a76da44fc511b789
--- /dev/null
+++ b/e4e/utils/common.py
@@ -0,0 +1,55 @@
+from PIL import Image
+import matplotlib.pyplot as plt
+
+
+# Log images
+def log_input_image(x, opts):
+ return tensor2im(x)
+
+
+def tensor2im(var):
+ # var shape: (3, H, W)
+ var = var.cpu().detach().transpose(0, 2).transpose(0, 1).numpy()
+ var = ((var + 1) / 2)
+ var[var < 0] = 0
+ var[var > 1] = 1
+ var = var * 255
+ return Image.fromarray(var.astype('uint8'))
+
+
+def vis_faces(log_hooks):
+ display_count = len(log_hooks)
+ fig = plt.figure(figsize=(8, 4 * display_count))
+ gs = fig.add_gridspec(display_count, 3)
+ for i in range(display_count):
+ hooks_dict = log_hooks[i]
+ fig.add_subplot(gs[i, 0])
+ if 'diff_input' in hooks_dict:
+ vis_faces_with_id(hooks_dict, fig, gs, i)
+ else:
+ vis_faces_no_id(hooks_dict, fig, gs, i)
+ plt.tight_layout()
+ return fig
+
+
+def vis_faces_with_id(hooks_dict, fig, gs, i):
+ plt.imshow(hooks_dict['input_face'])
+ plt.title('Input\nOut Sim={:.2f}'.format(float(hooks_dict['diff_input'])))
+ fig.add_subplot(gs[i, 1])
+ plt.imshow(hooks_dict['target_face'])
+ plt.title('Target\nIn={:.2f}, Out={:.2f}'.format(float(hooks_dict['diff_views']),
+ float(hooks_dict['diff_target'])))
+ fig.add_subplot(gs[i, 2])
+ plt.imshow(hooks_dict['output_face'])
+ plt.title('Output\n Target Sim={:.2f}'.format(float(hooks_dict['diff_target'])))
+
+
+def vis_faces_no_id(hooks_dict, fig, gs, i):
+ plt.imshow(hooks_dict['input_face'], cmap="gray")
+ plt.title('Input')
+ fig.add_subplot(gs[i, 1])
+ plt.imshow(hooks_dict['target_face'])
+ plt.title('Target')
+ fig.add_subplot(gs[i, 2])
+ plt.imshow(hooks_dict['output_face'])
+ plt.title('Output')
diff --git a/e4e/utils/data_utils.py b/e4e/utils/data_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..f1ba79f4a2d5cc2b97dce76d87bf6e7cdebbc257
--- /dev/null
+++ b/e4e/utils/data_utils.py
@@ -0,0 +1,25 @@
+"""
+Code adopted from pix2pixHD:
+https://github.com/NVIDIA/pix2pixHD/blob/master/data/image_folder.py
+"""
+import os
+
+IMG_EXTENSIONS = [
+ '.jpg', '.JPG', '.jpeg', '.JPEG',
+ '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tiff'
+]
+
+
+def is_image_file(filename):
+ return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+
+def make_dataset(dir):
+ images = []
+ assert os.path.isdir(dir), '%s is not a valid directory' % dir
+ for root, _, fnames in sorted(os.walk(dir)):
+ for fname in fnames:
+ if is_image_file(fname):
+ path = os.path.join(root, fname)
+ images.append(path)
+ return images
diff --git a/e4e/utils/model_utils.py b/e4e/utils/model_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..e51e95578f72b3218d6d832e3b604193cb68c1d7
--- /dev/null
+++ b/e4e/utils/model_utils.py
@@ -0,0 +1,35 @@
+import torch
+import argparse
+from models.psp import pSp
+from models.encoders.psp_encoders import Encoder4Editing
+
+
+def setup_model(checkpoint_path, device='cuda'):
+ ckpt = torch.load(checkpoint_path, map_location='cpu')
+ opts = ckpt['opts']
+
+ opts['checkpoint_path'] = checkpoint_path
+ opts['device'] = device
+ opts = argparse.Namespace(**opts)
+
+ net = pSp(opts)
+ net.eval()
+ net = net.to(device)
+ return net, opts
+
+
+def load_e4e_standalone(checkpoint_path, device='cuda'):
+ ckpt = torch.load(checkpoint_path, map_location='cpu')
+ opts = argparse.Namespace(**ckpt['opts'])
+ e4e = Encoder4Editing(50, 'ir_se', opts)
+ e4e_dict = {k.replace('encoder.', ''): v for k, v in ckpt['state_dict'].items() if k.startswith('encoder.')}
+ e4e.load_state_dict(e4e_dict)
+ e4e.eval()
+ e4e = e4e.to(device)
+ latent_avg = ckpt['latent_avg'].to(device)
+
+ def add_latent_avg(model, inputs, outputs):
+ return outputs + latent_avg.repeat(outputs.shape[0], 1, 1)
+
+ e4e.register_forward_hook(add_latent_avg)
+ return e4e
diff --git a/e4e/utils/train_utils.py b/e4e/utils/train_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..0c55177f7442010bc1fcc64de3d142585c22adc0
--- /dev/null
+++ b/e4e/utils/train_utils.py
@@ -0,0 +1,13 @@
+
+def aggregate_loss_dict(agg_loss_dict):
+ mean_vals = {}
+ for output in agg_loss_dict:
+ for key in output:
+ mean_vals[key] = mean_vals.setdefault(key, []) + [output[key]]
+ for key in mean_vals:
+ if len(mean_vals[key]) > 0:
+ mean_vals[key] = sum(mean_vals[key]) / len(mean_vals[key])
+ else:
+ print('{} has no value'.format(key))
+ mean_vals[key] = 0
+ return mean_vals
diff --git a/editing/interfacegan_boundaries/age.pt b/editing/interfacegan_boundaries/age.pt
new file mode 100644
index 0000000000000000000000000000000000000000..97024345fd8004681b65c8178653144326f3a631
Binary files /dev/null and b/editing/interfacegan_boundaries/age.pt differ
diff --git a/editing/interfacegan_boundaries/beard.pt b/editing/interfacegan_boundaries/beard.pt
new file mode 100644
index 0000000000000000000000000000000000000000..453f6686fee5f81cd1ddcaccff320b69cd15e55c
Binary files /dev/null and b/editing/interfacegan_boundaries/beard.pt differ
diff --git a/editing/interfacegan_boundaries/gender.pt b/editing/interfacegan_boundaries/gender.pt
new file mode 100644
index 0000000000000000000000000000000000000000..411a54bf8fe6fa13fb50c2b78505b3aabab9dde0
Binary files /dev/null and b/editing/interfacegan_boundaries/gender.pt differ
diff --git a/editing/interfacegan_boundaries/hair_length.pt b/editing/interfacegan_boundaries/hair_length.pt
new file mode 100644
index 0000000000000000000000000000000000000000..7721493285bafb398e34335a6a15855a7ba6f1fc
Binary files /dev/null and b/editing/interfacegan_boundaries/hair_length.pt differ
diff --git a/editing/interfacegan_boundaries/pose.pt b/editing/interfacegan_boundaries/pose.pt
new file mode 100644
index 0000000000000000000000000000000000000000..aa828a41df28f674e3a190520fa76eb037cd8d1d
Binary files /dev/null and b/editing/interfacegan_boundaries/pose.pt differ
diff --git a/editing/interfacegan_boundaries/smile.pt b/editing/interfacegan_boundaries/smile.pt
new file mode 100644
index 0000000000000000000000000000000000000000..357c9361a5f53f18c289459879ad4acfd194cb4a
Binary files /dev/null and b/editing/interfacegan_boundaries/smile.pt differ
diff --git a/generate_videos.py b/generate_videos.py
new file mode 100644
index 0000000000000000000000000000000000000000..0fb1ec6d9cef5ad3d36c885f989cd7f7346459a0
--- /dev/null
+++ b/generate_videos.py
@@ -0,0 +1,129 @@
+import os
+import argparse
+
+import torch
+from torchvision import utils
+
+from model.sg2_model import Generator
+from tqdm import tqdm
+from pathlib import Path
+
+import numpy as np
+
+import subprocess
+import shutil
+import copy
+
+from styleclip.styleclip_global import style_tensor_to_style_dict, style_dict_to_style_tensor
+
+VALID_EDITS = ["pose", "age", "smile", "gender", "hair_length", "beard"]
+
+SUGGESTED_DISTANCES = {
+ "pose": 3.0,
+ "smile": 2.0,
+ "age": 4.0,
+ "gender": 3.0,
+ "hair_length": -4.0,
+ "beard": 2.0
+ }
+
+def project_code(latent_code, boundary, distance=3.0):
+
+ if len(boundary) == 2:
+ boundary = boundary.reshape(1, 1, -1)
+
+ return latent_code + distance * boundary
+
+def project_code_by_edit_name(latent_code, name, strength):
+ boundary_dir = Path(os.path.abspath(__file__)).parents[0].joinpath("editing", "interfacegan_boundaries")
+
+ distance = SUGGESTED_DISTANCES[name] * strength
+ boundary = torch.load(os.path.join(boundary_dir, f'{name}.pt'), map_location="cpu").numpy()
+
+ return project_code(latent_code, boundary, distance)
+
+def generate_frames(source_latent, target_latents, g_ema_list, output_dir):
+
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+
+ code_is_s = target_latents[0].size()[1] == 9088
+
+ if code_is_s:
+ source_s_dict = g_ema_list[0].get_s_code(source_latent, input_is_latent=True)[0]
+ np_latent = style_dict_to_style_tensor(source_s_dict, g_ema_list[0]).cpu().detach().numpy()
+ else:
+ np_latent = source_latent.squeeze(0).cpu().detach().numpy()
+
+ np_target_latents = [target_latent.cpu().detach().numpy() for target_latent in target_latents]
+
+ num_alphas = 20 if code_is_s else min(10, 30 // len(target_latents))
+
+ alphas = np.linspace(0, 1, num=num_alphas)
+
+ latents = interpolate_with_target_latents(np_latent, np_target_latents, alphas)
+
+ segments = len(g_ema_list) - 1
+
+ if segments:
+ segment_length = len(latents) / segments
+
+ g_ema = copy.deepcopy(g_ema_list[0])
+
+ src_pars = dict(g_ema.named_parameters())
+ mix_pars = [dict(model.named_parameters()) for model in g_ema_list]
+ else:
+ g_ema = g_ema_list[0]
+
+ print("Generating frames for video...")
+ for idx, latent in tqdm(enumerate(latents), total=len(latents)):
+
+ if segments:
+ mix_alpha = (idx % segment_length) * 1.0 / segment_length
+ segment_id = int(idx // segment_length)
+
+ for k in src_pars.keys():
+ src_pars[k].data.copy_(mix_pars[segment_id][k] * (1 - mix_alpha) + mix_pars[segment_id + 1][k] * mix_alpha)
+
+ if idx == 0 or segments or latent is not latents[idx - 1]:
+ latent_tensor = torch.from_numpy(latent).float().to(device)
+
+ with torch.no_grad():
+ if code_is_s:
+ latent_for_gen = style_tensor_to_style_dict(latent_tensor, g_ema)
+ img, _ = g_ema(latent_for_gen, input_is_s_code=True, input_is_latent=True, truncation=1, randomize_noise=False)
+ else:
+ img, _ = g_ema([latent_tensor], input_is_latent=True, truncation=1, randomize_noise=False)
+
+ utils.save_image(img, f"{output_dir}/{str(idx).zfill(3)}.jpg", nrow=1, normalize=True, scale_each=True, range=(-1, 1))
+
+def interpolate_forward_backward(source_latent, target_latent, alphas):
+ latents_forward = [a * target_latent + (1-a) * source_latent for a in alphas] # interpolate from source to target
+ latents_backward = latents_forward[::-1] # interpolate from target to source
+ return latents_forward + [target_latent] * len(alphas) + latents_backward # forward + short delay at target + return
+
+def interpolate_with_target_latents(source_latent, target_latents, alphas):
+ # interpolate latent codes with all targets
+
+ print("Interpolating latent codes...")
+
+ latents = []
+ for target_latent in target_latents:
+ latents.extend(interpolate_forward_backward(source_latent, target_latent, alphas))
+
+ return latents
+
+def video_from_interpolations(fps, output_dir):
+
+ # combine frames to a video
+ command = ["ffmpeg",
+ "-r", f"{fps}",
+ "-i", f"{output_dir}/%03d.jpg",
+ "-c:v", "libx264",
+ "-vf", f"fps={fps}",
+ "-pix_fmt", "yuv420p",
+ f"{output_dir}/out.mp4"]
+
+ subprocess.call(command)
+
+
+
diff --git a/model/sg2_model.py b/model/sg2_model.py
new file mode 100755
index 0000000000000000000000000000000000000000..76b6f2bf128a5771c69dd7ce0e28362f1d82e30c
--- /dev/null
+++ b/model/sg2_model.py
@@ -0,0 +1,817 @@
+import math
+import random
+import functools
+import operator
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.autograd import Function
+
+from op import conv2d_gradfix
+
+if torch.cuda.is_available():
+ from op.fused_act import FusedLeakyReLU, fused_leaky_relu
+ from op.upfirdn2d import upfirdn2d
+else:
+ from op.fused_act_cpu import FusedLeakyReLU, fused_leaky_relu
+ from op.upfirdn2d_cpu import upfirdn2d
+
+
+class PixelNorm(nn.Module):
+ def __init__(self):
+ super().__init__()
+
+ def forward(self, input):
+ return input * torch.rsqrt(torch.mean(input ** 2, dim=1, keepdim=True) + 1e-8)
+
+
+def make_kernel(k):
+ k = torch.tensor(k, dtype=torch.float32)
+
+ if k.ndim == 1:
+ k = k[None, :] * k[:, None]
+
+ k /= k.sum()
+
+ return k
+
+
+class Upsample(nn.Module):
+ def __init__(self, kernel, factor=2):
+ super().__init__()
+
+ self.factor = factor
+ kernel = make_kernel(kernel) * (factor ** 2)
+ self.register_buffer("kernel", kernel)
+
+ p = kernel.shape[0] - factor
+
+ pad0 = (p + 1) // 2 + factor - 1
+ pad1 = p // 2
+
+ self.pad = (pad0, pad1)
+
+ def forward(self, input):
+ out = upfirdn2d(input, self.kernel, up=self.factor, down=1, pad=self.pad)
+
+ return out
+
+
+class Downsample(nn.Module):
+ def __init__(self, kernel, factor=2):
+ super().__init__()
+
+ self.factor = factor
+ kernel = make_kernel(kernel)
+ self.register_buffer("kernel", kernel)
+
+ p = kernel.shape[0] - factor
+
+ pad0 = (p + 1) // 2
+ pad1 = p // 2
+
+ self.pad = (pad0, pad1)
+
+ def forward(self, input):
+ out = upfirdn2d(input, self.kernel, up=1, down=self.factor, pad=self.pad)
+
+ return out
+
+
+class Blur(nn.Module):
+ def __init__(self, kernel, pad, upsample_factor=1):
+ super().__init__()
+
+ kernel = make_kernel(kernel)
+
+ if upsample_factor > 1:
+ kernel = kernel * (upsample_factor ** 2)
+
+ self.register_buffer("kernel", kernel)
+
+ self.pad = pad
+
+ def forward(self, input):
+ out = upfirdn2d(input, self.kernel, pad=self.pad)
+
+ return out
+
+
+class EqualConv2d(nn.Module):
+ def __init__(
+ self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True
+ ):
+ super().__init__()
+
+ self.weight = nn.Parameter(
+ torch.randn(out_channel, in_channel, kernel_size, kernel_size)
+ )
+ self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2)
+
+ self.stride = stride
+ self.padding = padding
+
+ if bias:
+ self.bias = nn.Parameter(torch.zeros(out_channel))
+
+ else:
+ self.bias = None
+
+ def forward(self, input):
+ out = conv2d_gradfix.conv2d(
+ input,
+ self.weight * self.scale,
+ bias=self.bias,
+ stride=self.stride,
+ padding=self.padding,
+ )
+
+ return out
+
+ def __repr__(self):
+ return (
+ f"{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},"
+ f" {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})"
+ )
+
+
+class EqualLinear(nn.Module):
+ def __init__(
+ self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None
+ ):
+ super().__init__()
+
+ self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul))
+
+ if bias:
+ self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init))
+
+ else:
+ self.bias = None
+
+ self.activation = activation
+
+ self.scale = (1 / math.sqrt(in_dim)) * lr_mul
+ self.lr_mul = lr_mul
+
+ def forward(self, input):
+ if self.activation:
+ out = F.linear(input, self.weight * self.scale)
+ out = fused_leaky_relu(out, self.bias * self.lr_mul)
+
+ else:
+ out = F.linear(
+ input, self.weight * self.scale, bias=self.bias * self.lr_mul
+ )
+
+ return out
+
+ def __repr__(self):
+ return (
+ f"{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})"
+ )
+
+
+class ModulatedConv2d(nn.Module):
+ def __init__(
+ self,
+ in_channel,
+ out_channel,
+ kernel_size,
+ style_dim,
+ demodulate=True,
+ upsample=False,
+ downsample=False,
+ blur_kernel=[1, 3, 3, 1],
+ fused=True,
+ ):
+ super().__init__()
+
+ self.eps = 1e-8
+ self.kernel_size = kernel_size
+ self.in_channel = in_channel
+ self.out_channel = out_channel
+ self.upsample = upsample
+ self.downsample = downsample
+
+ if upsample:
+ factor = 2
+ p = (len(blur_kernel) - factor) - (kernel_size - 1)
+ pad0 = (p + 1) // 2 + factor - 1
+ pad1 = p // 2 + 1
+
+ self.blur = Blur(blur_kernel, pad=(pad0, pad1), upsample_factor=factor)
+
+ if downsample:
+ factor = 2
+ p = (len(blur_kernel) - factor) + (kernel_size - 1)
+ pad0 = (p + 1) // 2
+ pad1 = p // 2
+
+ self.blur = Blur(blur_kernel, pad=(pad0, pad1))
+
+ fan_in = in_channel * kernel_size ** 2
+ self.scale = 1 / math.sqrt(fan_in)
+ self.padding = kernel_size // 2
+
+ self.weight = nn.Parameter(
+ torch.randn(1, out_channel, in_channel, kernel_size, kernel_size)
+ )
+
+ self.modulation = EqualLinear(style_dim, in_channel, bias_init=1)
+
+ self.demodulate = demodulate
+ self.fused = fused
+
+ def __repr__(self):
+ return (
+ f"{self.__class__.__name__}({self.in_channel}, {self.out_channel}, {self.kernel_size}, "
+ f"upsample={self.upsample}, downsample={self.downsample})"
+ )
+
+ def forward(self, input, style, is_s_code=False):
+ batch, in_channel, height, width = input.shape
+
+ if not self.fused:
+ weight = self.scale * self.weight.squeeze(0)
+
+ if is_s_code:
+ style = style[self.modulation]
+ else:
+ style = self.modulation(style)
+
+ if self.demodulate:
+ w = weight.unsqueeze(0) * style.view(batch, 1, in_channel, 1, 1)
+ dcoefs = (w.square().sum((2, 3, 4)) + 1e-8).rsqrt()
+
+ input = input * style.reshape(batch, in_channel, 1, 1)
+
+ if self.upsample:
+ weight = weight.transpose(0, 1)
+ out = conv2d_gradfix.conv_transpose2d(
+ input, weight, padding=0, stride=2
+ )
+ out = self.blur(out)
+
+ elif self.downsample:
+ input = self.blur(input)
+ out = conv2d_gradfix.conv2d(input, weight, padding=0, stride=2)
+
+ else:
+ out = conv2d_gradfix.conv2d(input, weight, padding=self.padding)
+
+ if self.demodulate:
+ out = out * dcoefs.view(batch, -1, 1, 1)
+
+ return out
+
+ if is_s_code:
+ style = style[self.modulation]
+ else:
+ style = self.modulation(style)
+
+ style = style.view(batch, 1, in_channel, 1, 1)
+ weight = self.scale * self.weight * style
+
+ if self.demodulate:
+ demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + 1e-8)
+ weight = weight * demod.view(batch, self.out_channel, 1, 1, 1)
+
+ weight = weight.view(
+ batch * self.out_channel, in_channel, self.kernel_size, self.kernel_size
+ )
+
+ if self.upsample:
+ input = input.view(1, batch * in_channel, height, width)
+ weight = weight.view(
+ batch, self.out_channel, in_channel, self.kernel_size, self.kernel_size
+ )
+ weight = weight.transpose(1, 2).reshape(
+ batch * in_channel, self.out_channel, self.kernel_size, self.kernel_size
+ )
+ out = conv2d_gradfix.conv_transpose2d(
+ input, weight, padding=0, stride=2, groups=batch
+ )
+ _, _, height, width = out.shape
+ out = out.view(batch, self.out_channel, height, width)
+ out = self.blur(out)
+
+ elif self.downsample:
+ input = self.blur(input)
+ _, _, height, width = input.shape
+ input = input.view(1, batch * in_channel, height, width)
+ out = conv2d_gradfix.conv2d(
+ input, weight, padding=0, stride=2, groups=batch
+ )
+ _, _, height, width = out.shape
+ out = out.view(batch, self.out_channel, height, width)
+
+ else:
+ input = input.view(1, batch * in_channel, height, width)
+ out = conv2d_gradfix.conv2d(
+ input, weight, padding=self.padding, groups=batch
+ )
+ _, _, height, width = out.shape
+ out = out.view(batch, self.out_channel, height, width)
+
+ return out
+
+
+class NoiseInjection(nn.Module):
+ def __init__(self):
+ super().__init__()
+
+ self.weight = nn.Parameter(torch.zeros(1))
+
+ def forward(self, image, noise=None):
+ if noise is None:
+ batch, _, height, width = image.shape
+ noise = image.new_empty(batch, 1, height, width).normal_()
+
+ return image + self.weight * noise
+
+
+class ConstantInput(nn.Module):
+ def __init__(self, channel, size=4):
+ super().__init__()
+
+ self.input = nn.Parameter(torch.randn(1, channel, size, size))
+
+ def forward(self, input, is_s_code=False):
+ if not is_s_code:
+ batch = input.shape[0]
+ else:
+ batch = next(iter(input.values())).shape[0]
+
+ out = self.input.repeat(batch, 1, 1, 1)
+
+ return out
+
+
+class StyledConv(nn.Module):
+ def __init__(
+ self,
+ in_channel,
+ out_channel,
+ kernel_size,
+ style_dim,
+ upsample=False,
+ blur_kernel=[1, 3, 3, 1],
+ demodulate=True,
+ ):
+ super().__init__()
+
+ self.conv = ModulatedConv2d(
+ in_channel,
+ out_channel,
+ kernel_size,
+ style_dim,
+ upsample=upsample,
+ blur_kernel=blur_kernel,
+ demodulate=demodulate,
+ )
+
+ self.noise = NoiseInjection()
+ # self.bias = nn.Parameter(torch.zeros(1, out_channel, 1, 1))
+ # self.activate = ScaledLeakyReLU(0.2)
+ self.activate = FusedLeakyReLU(out_channel)
+
+ def forward(self, input, style, noise=None, is_s_code=False):
+ out = self.conv(input, style, is_s_code=is_s_code)
+ out = self.noise(out, noise=noise)
+ # out = out + self.bias
+ out = self.activate(out)
+
+ return out
+
+
+class ToRGB(nn.Module):
+ def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[1, 3, 3, 1]):
+ super().__init__()
+
+ if upsample:
+ self.upsample = Upsample(blur_kernel)
+
+ self.conv = ModulatedConv2d(in_channel, 3, 1, style_dim, demodulate=False)
+ self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1))
+
+ def forward(self, input, style, skip=None, is_s_code=False):
+ out = self.conv(input, style, is_s_code=is_s_code)
+ out = out + self.bias
+
+ if skip is not None:
+ skip = self.upsample(skip)
+
+ out = out + skip
+
+ return out
+
+
+class Generator(nn.Module):
+ def __init__(
+ self,
+ size,
+ style_dim,
+ n_mlp,
+ channel_multiplier=2,
+ blur_kernel=[1, 3, 3, 1],
+ lr_mlp=0.01,
+ ):
+ super().__init__()
+
+ self.size = size
+
+ self.style_dim = style_dim
+
+ layers = [PixelNorm()]
+
+ for i in range(n_mlp):
+ layers.append(
+ EqualLinear(
+ style_dim, style_dim, lr_mul=lr_mlp, activation="fused_lrelu"
+ )
+ )
+
+ self.style = nn.Sequential(*layers)
+
+ self.channels = {
+ 4: 512,
+ 8: 512,
+ 16: 512,
+ 32: 512,
+ 64: 256 * channel_multiplier,
+ 128: 128 * channel_multiplier,
+ 256: 64 * channel_multiplier,
+ 512: 32 * channel_multiplier,
+ 1024: 16 * channel_multiplier,
+ }
+
+ self.input = ConstantInput(self.channels[4])
+ self.conv1 = StyledConv(
+ self.channels[4], self.channels[4], 3, style_dim, blur_kernel=blur_kernel
+ )
+ self.to_rgb1 = ToRGB(self.channels[4], style_dim, upsample=False)
+
+ self.log_size = int(math.log(size, 2))
+ self.num_layers = (self.log_size - 2) * 2 + 1
+
+ self.convs = nn.ModuleList()
+ self.upsamples = nn.ModuleList()
+ self.to_rgbs = nn.ModuleList()
+ self.noises = nn.Module()
+
+ in_channel = self.channels[4]
+
+ for layer_idx in range(self.num_layers):
+ res = (layer_idx + 5) // 2
+ shape = [1, 1, 2 ** res, 2 ** res]
+ self.noises.register_buffer(f"noise_{layer_idx}", torch.randn(*shape))
+
+ for i in range(3, self.log_size + 1):
+ out_channel = self.channels[2 ** i]
+
+ self.convs.append(
+ StyledConv(
+ in_channel,
+ out_channel,
+ 3,
+ style_dim,
+ upsample=True,
+ blur_kernel=blur_kernel,
+ )
+ )
+
+ self.convs.append(
+ StyledConv(
+ out_channel, out_channel, 3, style_dim, blur_kernel=blur_kernel
+ )
+ )
+
+ self.to_rgbs.append(ToRGB(out_channel, style_dim))
+
+ in_channel = out_channel
+
+ self.n_latent = self.log_size * 2 - 2
+
+
+ self.modulation_layers = [self.conv1.conv.modulation, self.to_rgb1.conv.modulation] + \
+ [layer.conv.modulation for layer in self.convs] + \
+ [layer.conv.modulation for layer in self.to_rgbs]
+
+ def make_noise(self):
+ device = self.input.input.device
+
+ noises = [torch.randn(1, 1, 2 ** 2, 2 ** 2, device=device)]
+
+ for i in range(3, self.log_size + 1):
+ for _ in range(2):
+ noises.append(torch.randn(1, 1, 2 ** i, 2 ** i, device=device))
+
+ return noises
+
+ def mean_latent(self, n_latent):
+ latent_in = torch.randn(
+ n_latent, self.style_dim, device=self.input.input.device
+ )
+ latent = self.style(latent_in).mean(0, keepdim=True)
+
+ return latent
+
+ def get_latent(self, input):
+ return self.style(input)
+
+ def get_s_code(self, styles, input_is_latent):
+
+ if not input_is_latent:
+ styles = [self.style(s) for s in styles]
+
+ s_codes = [{# const block
+ self.modulation_layers[0]: self.modulation_layers[0](style[:, 0]), #s0
+ self.modulation_layers[1]: self.modulation_layers[1](style[:, 1]), #s1
+ # conv layers
+ self.modulation_layers[2]: self.modulation_layers[2](style[:, 1]), #s2
+ self.modulation_layers[3]: self.modulation_layers[3](style[:, 2]), #s3
+ self.modulation_layers[4]: self.modulation_layers[4](style[:, 3]), #s5
+ self.modulation_layers[5]: self.modulation_layers[5](style[:, 4]), #s6
+ self.modulation_layers[6]: self.modulation_layers[6](style[:, 5]), #s8
+ self.modulation_layers[7]: self.modulation_layers[7](style[:, 6]), #s9
+ self.modulation_layers[8]: self.modulation_layers[8](style[:, 7]), #s11
+ self.modulation_layers[9]: self.modulation_layers[9](style[:, 8]), #s12
+ self.modulation_layers[10]: self.modulation_layers[10](style[:, 9]), #s14
+ self.modulation_layers[11]: self.modulation_layers[11](style[:, 10]), #s15
+ self.modulation_layers[12]: self.modulation_layers[12](style[:, 11]), #s17
+ self.modulation_layers[13]: self.modulation_layers[13](style[:, 12]), #s18
+ self.modulation_layers[14]: self.modulation_layers[14](style[:, 13]), #s20
+ self.modulation_layers[15]: self.modulation_layers[15](style[:, 14]), #s21
+ self.modulation_layers[16]: self.modulation_layers[16](style[:, 15]), #s23
+ self.modulation_layers[17]: self.modulation_layers[17](style[:, 16]), #s24
+ # toRGB layers
+ self.modulation_layers[18]: self.modulation_layers[18](style[:, 3]), #s4
+ self.modulation_layers[19]: self.modulation_layers[19](style[:, 5]), #s7
+ self.modulation_layers[20]: self.modulation_layers[20](style[:, 7]), #s10
+ self.modulation_layers[21]: self.modulation_layers[21](style[:, 9]), #s13
+ self.modulation_layers[22]: self.modulation_layers[22](style[:, 11]), #s16
+ self.modulation_layers[23]: self.modulation_layers[23](style[:, 13]), #s19
+ self.modulation_layers[24]: self.modulation_layers[24](style[:, 15]), #s22
+ self.modulation_layers[25]: self.modulation_layers[25](style[:, 17]), #s25
+ } for style in styles]
+
+ return s_codes
+
+
+ def forward(
+ self,
+ styles,
+ return_latents=False,
+ inject_index=None,
+ truncation=1,
+ truncation_latent=None,
+ input_is_latent=False,
+ input_is_s_code=False,
+ noise=None,
+ randomize_noise=True,
+ ):
+ if not input_is_s_code:
+ return self.forward_with_w(styles, return_latents, inject_index, truncation, truncation_latent, input_is_latent, noise, randomize_noise)
+
+ return self.forward_with_s(styles, return_latents, noise, randomize_noise)
+
+ def forward_with_w(
+ self,
+ styles,
+ return_latents=False,
+ inject_index=None,
+ truncation=1,
+ truncation_latent=None,
+ input_is_latent=False,
+ noise=None,
+ randomize_noise=True,
+ ):
+ if not input_is_latent:
+ styles = [self.style(s) for s in styles]
+
+ if noise is None:
+ if randomize_noise:
+ noise = [None] * self.num_layers
+ else:
+ noise = [
+ getattr(self.noises, f"noise_{i}") for i in range(self.num_layers)
+ ]
+
+ if truncation < 1:
+ style_t = []
+
+ for style in styles:
+ style_t.append(
+ truncation_latent + truncation * (style - truncation_latent)
+ )
+
+ styles = style_t
+
+ if len(styles) < 2:
+ inject_index = self.n_latent
+
+ if styles[0].ndim < 3:
+ latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
+
+ else:
+ latent = styles[0]
+
+ else:
+ if inject_index is None:
+ inject_index = random.randint(1, self.n_latent - 1)
+
+ latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
+ latent2 = styles[1].unsqueeze(1).repeat(1, self.n_latent - inject_index, 1)
+
+ latent = torch.cat([latent, latent2], 1)
+
+ out = self.input(latent)
+ out = self.conv1(out, latent[:, 0], noise=noise[0])
+
+ skip = self.to_rgb1(out, latent[:, 1])
+
+ i = 1
+ for conv1, conv2, noise1, noise2, to_rgb in zip(
+ self.convs[::2], self.convs[1::2], noise[1::2], noise[2::2], self.to_rgbs
+ ):
+ out = conv1(out, latent[:, i], noise=noise1)
+ out = conv2(out, latent[:, i + 1], noise=noise2)
+ skip = to_rgb(out, latent[:, i + 2], skip)
+
+ i += 2
+
+ image = skip
+
+ if return_latents:
+ return image, latent
+
+ else:
+ return image, None
+
+ def forward_with_s(
+ self,
+ styles,
+ return_latents=False,
+ noise=None,
+ randomize_noise=True,
+ ):
+
+ if noise is None:
+ if randomize_noise:
+ noise = [None] * self.num_layers
+ else:
+ noise = [
+ getattr(self.noises, f"noise_{i}") for i in range(self.num_layers)
+ ]
+
+ out = self.input(styles, is_s_code=True)
+ out = self.conv1(out, styles, is_s_code=True, noise=noise[0])
+
+ skip = self.to_rgb1(out, styles, is_s_code=True)
+
+ i = 1
+ for conv1, conv2, noise1, noise2, to_rgb in zip(
+ self.convs[::2], self.convs[1::2], noise[1::2], noise[2::2], self.to_rgbs
+ ):
+ out = conv1(out, styles, is_s_code=True, noise=noise1)
+ out = conv2(out, styles, is_s_code=True, noise=noise2)
+ skip = to_rgb(out, styles, skip, is_s_code=True)
+
+ i += 2
+
+ image = skip
+
+ if return_latents:
+ return image, styles
+
+ else:
+ return image, None
+
+class ConvLayer(nn.Sequential):
+ def __init__(
+ self,
+ in_channel,
+ out_channel,
+ kernel_size,
+ downsample=False,
+ blur_kernel=[1, 3, 3, 1],
+ bias=True,
+ activate=True,
+ ):
+ layers = []
+
+ if downsample:
+ factor = 2
+ p = (len(blur_kernel) - factor) + (kernel_size - 1)
+ pad0 = (p + 1) // 2
+ pad1 = p // 2
+
+ layers.append(Blur(blur_kernel, pad=(pad0, pad1)))
+
+ stride = 2
+ self.padding = 0
+
+ else:
+ stride = 1
+ self.padding = kernel_size // 2
+
+ layers.append(
+ EqualConv2d(
+ in_channel,
+ out_channel,
+ kernel_size,
+ padding=self.padding,
+ stride=stride,
+ bias=bias and not activate,
+ )
+ )
+
+ if activate:
+ layers.append(FusedLeakyReLU(out_channel, bias=bias))
+
+ super().__init__(*layers)
+
+
+class ResBlock(nn.Module):
+ def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]):
+ super().__init__()
+
+ self.conv1 = ConvLayer(in_channel, in_channel, 3)
+ self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True)
+
+ self.skip = ConvLayer(
+ in_channel, out_channel, 1, downsample=True, activate=False, bias=False
+ )
+
+ def forward(self, input):
+ out = self.conv1(input)
+ out = self.conv2(out)
+
+ skip = self.skip(input)
+ out = (out + skip) / math.sqrt(2)
+
+ return out
+
+
+class Discriminator(nn.Module):
+ def __init__(self, size, channel_multiplier=2, blur_kernel=[1, 3, 3, 1]):
+ super().__init__()
+
+ channels = {
+ 4: 512,
+ 8: 512,
+ 16: 512,
+ 32: 512,
+ 64: 256 * channel_multiplier,
+ 128: 128 * channel_multiplier,
+ 256: 64 * channel_multiplier,
+ 512: 32 * channel_multiplier,
+ 1024: 16 * channel_multiplier,
+ }
+
+ convs = [ConvLayer(3, channels[size], 1)]
+
+ log_size = int(math.log(size, 2))
+
+ in_channel = channels[size]
+
+ for i in range(log_size, 2, -1):
+ out_channel = channels[2 ** (i - 1)]
+
+ convs.append(ResBlock(in_channel, out_channel, blur_kernel))
+
+ in_channel = out_channel
+
+ self.convs = nn.Sequential(*convs)
+
+ self.stddev_group = 4
+ self.stddev_feat = 1
+
+ self.final_conv = ConvLayer(in_channel + 1, channels[4], 3)
+ self.final_linear = nn.Sequential(
+ EqualLinear(channels[4] * 4 * 4, channels[4], activation="fused_lrelu"),
+ EqualLinear(channels[4], 1),
+ )
+
+ def forward(self, input):
+ out = self.convs(input)
+
+ batch, channel, height, width = out.shape
+ group = min(batch, self.stddev_group)
+ stddev = out.view(
+ group, -1, self.stddev_feat, channel // self.stddev_feat, height, width
+ )
+ stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-8)
+ stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
+ stddev = stddev.repeat(group, 1, height, width)
+ out = torch.cat([out, stddev], 1)
+
+ out = self.final_conv(out)
+
+ out = out.view(batch, -1)
+ out = self.final_linear(out)
+
+ return out
+
diff --git a/op/__init__.py b/op/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/op/__pycache__/__init__.cpython-37.pyc b/op/__pycache__/__init__.cpython-37.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e1cd6edc2e0bced60aed726840d54f9a71a9c815
Binary files /dev/null and b/op/__pycache__/__init__.cpython-37.pyc differ
diff --git a/op/__pycache__/__init__.cpython-38.pyc b/op/__pycache__/__init__.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5ea25e5af3781755c8cbf61e3e017138ae3eafda
Binary files /dev/null and b/op/__pycache__/__init__.cpython-38.pyc differ
diff --git a/op/__pycache__/conv2d_gradfix.cpython-37.pyc b/op/__pycache__/conv2d_gradfix.cpython-37.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cfa25b42699d412fcf4fe261958c7fc0f18be2a9
Binary files /dev/null and b/op/__pycache__/conv2d_gradfix.cpython-37.pyc differ
diff --git a/op/__pycache__/conv2d_gradfix.cpython-38.pyc b/op/__pycache__/conv2d_gradfix.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..348760e25c6d479ce605ad12d04b2a61d362eaea
Binary files /dev/null and b/op/__pycache__/conv2d_gradfix.cpython-38.pyc differ
diff --git a/op/__pycache__/fused_act.cpython-37.pyc b/op/__pycache__/fused_act.cpython-37.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c26c07421b3965c3c2b48074706048dcb0fc5455
Binary files /dev/null and b/op/__pycache__/fused_act.cpython-37.pyc differ
diff --git a/op/__pycache__/fused_act.cpython-38.pyc b/op/__pycache__/fused_act.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a68ecfafa282521d70a503e571160232a4381661
Binary files /dev/null and b/op/__pycache__/fused_act.cpython-38.pyc differ
diff --git a/op/__pycache__/upfirdn2d.cpython-37.pyc b/op/__pycache__/upfirdn2d.cpython-37.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d3351be21e5d0adcbd4f16cda557d6212ada468e
Binary files /dev/null and b/op/__pycache__/upfirdn2d.cpython-37.pyc differ
diff --git a/op/__pycache__/upfirdn2d.cpython-38.pyc b/op/__pycache__/upfirdn2d.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c7b3c3cf96b0990ef7d56093cd85507b1e53a963
Binary files /dev/null and b/op/__pycache__/upfirdn2d.cpython-38.pyc differ
diff --git a/op/conv2d_gradfix.py b/op/conv2d_gradfix.py
new file mode 100644
index 0000000000000000000000000000000000000000..bb2f94bbcb8132299fd4d538972d32bd7ff6e7d6
--- /dev/null
+++ b/op/conv2d_gradfix.py
@@ -0,0 +1,227 @@
+import contextlib
+import warnings
+
+import torch
+from torch import autograd
+from torch.nn import functional as F
+
+enabled = True
+weight_gradients_disabled = False
+
+
+@contextlib.contextmanager
+def no_weight_gradients():
+ global weight_gradients_disabled
+
+ old = weight_gradients_disabled
+ weight_gradients_disabled = True
+ yield
+ weight_gradients_disabled = old
+
+
+def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
+ if could_use_op(input):
+ return conv2d_gradfix(
+ transpose=False,
+ weight_shape=weight.shape,
+ stride=stride,
+ padding=padding,
+ output_padding=0,
+ dilation=dilation,
+ groups=groups,
+ ).apply(input, weight, bias)
+
+ return F.conv2d(
+ input=input,
+ weight=weight,
+ bias=bias,
+ stride=stride,
+ padding=padding,
+ dilation=dilation,
+ groups=groups,
+ )
+
+
+def conv_transpose2d(
+ input,
+ weight,
+ bias=None,
+ stride=1,
+ padding=0,
+ output_padding=0,
+ groups=1,
+ dilation=1,
+):
+ if could_use_op(input):
+ return conv2d_gradfix(
+ transpose=True,
+ weight_shape=weight.shape,
+ stride=stride,
+ padding=padding,
+ output_padding=output_padding,
+ groups=groups,
+ dilation=dilation,
+ ).apply(input, weight, bias)
+
+ return F.conv_transpose2d(
+ input=input,
+ weight=weight,
+ bias=bias,
+ stride=stride,
+ padding=padding,
+ output_padding=output_padding,
+ dilation=dilation,
+ groups=groups,
+ )
+
+
+def could_use_op(input):
+ if (not enabled) or (not torch.backends.cudnn.enabled):
+ return False
+
+ if input.device.type != "cuda":
+ return False
+
+ if any(torch.__version__.startswith(x) for x in ["1.7.", "1.8."]):
+ return True
+
+ warnings.warn(
+ f"conv2d_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.conv2d()."
+ )
+
+ return False
+
+
+def ensure_tuple(xs, ndim):
+ xs = tuple(xs) if isinstance(xs, (tuple, list)) else (xs,) * ndim
+
+ return xs
+
+
+conv2d_gradfix_cache = dict()
+
+
+def conv2d_gradfix(
+ transpose, weight_shape, stride, padding, output_padding, dilation, groups
+):
+ ndim = 2
+ weight_shape = tuple(weight_shape)
+ stride = ensure_tuple(stride, ndim)
+ padding = ensure_tuple(padding, ndim)
+ output_padding = ensure_tuple(output_padding, ndim)
+ dilation = ensure_tuple(dilation, ndim)
+
+ key = (transpose, weight_shape, stride, padding, output_padding, dilation, groups)
+ if key in conv2d_gradfix_cache:
+ return conv2d_gradfix_cache[key]
+
+ common_kwargs = dict(
+ stride=stride, padding=padding, dilation=dilation, groups=groups
+ )
+
+ def calc_output_padding(input_shape, output_shape):
+ if transpose:
+ return [0, 0]
+
+ return [
+ input_shape[i + 2]
+ - (output_shape[i + 2] - 1) * stride[i]
+ - (1 - 2 * padding[i])
+ - dilation[i] * (weight_shape[i + 2] - 1)
+ for i in range(ndim)
+ ]
+
+ class Conv2d(autograd.Function):
+ @staticmethod
+ def forward(ctx, input, weight, bias):
+ if not transpose:
+ out = F.conv2d(input=input, weight=weight, bias=bias, **common_kwargs)
+
+ else:
+ out = F.conv_transpose2d(
+ input=input,
+ weight=weight,
+ bias=bias,
+ output_padding=output_padding,
+ **common_kwargs,
+ )
+
+ ctx.save_for_backward(input, weight)
+
+ return out
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ input, weight = ctx.saved_tensors
+ grad_input, grad_weight, grad_bias = None, None, None
+
+ if ctx.needs_input_grad[0]:
+ p = calc_output_padding(
+ input_shape=input.shape, output_shape=grad_output.shape
+ )
+ grad_input = conv2d_gradfix(
+ transpose=(not transpose),
+ weight_shape=weight_shape,
+ output_padding=p,
+ **common_kwargs,
+ ).apply(grad_output, weight, None)
+
+ if ctx.needs_input_grad[1] and not weight_gradients_disabled:
+ grad_weight = Conv2dGradWeight.apply(grad_output, input)
+
+ if ctx.needs_input_grad[2]:
+ grad_bias = grad_output.sum((0, 2, 3))
+
+ return grad_input, grad_weight, grad_bias
+
+ class Conv2dGradWeight(autograd.Function):
+ @staticmethod
+ def forward(ctx, grad_output, input):
+ op = torch._C._jit_get_operation(
+ "aten::cudnn_convolution_backward_weight"
+ if not transpose
+ else "aten::cudnn_convolution_transpose_backward_weight"
+ )
+ flags = [
+ torch.backends.cudnn.benchmark,
+ torch.backends.cudnn.deterministic,
+ torch.backends.cudnn.allow_tf32,
+ ]
+ grad_weight = op(
+ weight_shape,
+ grad_output,
+ input,
+ padding,
+ stride,
+ dilation,
+ groups,
+ *flags,
+ )
+ ctx.save_for_backward(grad_output, input)
+
+ return grad_weight
+
+ @staticmethod
+ def backward(ctx, grad_grad_weight):
+ grad_output, input = ctx.saved_tensors
+ grad_grad_output, grad_grad_input = None, None
+
+ if ctx.needs_input_grad[0]:
+ grad_grad_output = Conv2d.apply(input, grad_grad_weight, None)
+
+ if ctx.needs_input_grad[1]:
+ p = calc_output_padding(
+ input_shape=input.shape, output_shape=grad_output.shape
+ )
+ grad_grad_input = conv2d_gradfix(
+ transpose=(not transpose),
+ weight_shape=weight_shape,
+ output_padding=p,
+ **common_kwargs,
+ ).apply(grad_output, grad_grad_weight, None)
+
+ return grad_grad_output, grad_grad_input
+
+ conv2d_gradfix_cache[key] = Conv2d
+
+ return Conv2d
diff --git a/op/fused_act.py b/op/fused_act.py
new file mode 100755
index 0000000000000000000000000000000000000000..8459d510d7b79684779dfe47f5b46d81c94b4a4d
--- /dev/null
+++ b/op/fused_act.py
@@ -0,0 +1,86 @@
+import os
+
+import torch
+from torch import nn
+from torch.autograd import Function
+from torch.utils.cpp_extension import load
+
+
+module_path = os.path.dirname(__file__)
+fused = load(
+ 'fused',
+ sources=[
+ os.path.join(module_path, 'fused_bias_act.cpp'),
+ os.path.join(module_path, 'fused_bias_act_kernel.cu'),
+ ],
+)
+
+
+class FusedLeakyReLUFunctionBackward(Function):
+ @staticmethod
+ def forward(ctx, grad_output, out, negative_slope, scale):
+ ctx.save_for_backward(out)
+ ctx.negative_slope = negative_slope
+ ctx.scale = scale
+
+ empty = grad_output.new_empty(0)
+
+ grad_input = fused.fused_bias_act(
+ grad_output, empty, out, 3, 1, negative_slope, scale
+ )
+
+ dim = [0]
+
+ if grad_input.ndim > 2:
+ dim += list(range(2, grad_input.ndim))
+
+ grad_bias = grad_input.sum(dim).detach()
+
+ return grad_input, grad_bias
+
+ @staticmethod
+ def backward(ctx, gradgrad_input, gradgrad_bias):
+ out, = ctx.saved_tensors
+ gradgrad_out = fused.fused_bias_act(
+ gradgrad_input, gradgrad_bias, out, 3, 1, ctx.negative_slope, ctx.scale
+ )
+
+ return gradgrad_out, None, None, None
+
+
+class FusedLeakyReLUFunction(Function):
+ @staticmethod
+ def forward(ctx, input, bias, negative_slope, scale):
+ empty = input.new_empty(0)
+ out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale)
+ ctx.save_for_backward(out)
+ ctx.negative_slope = negative_slope
+ ctx.scale = scale
+
+ return out
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ out, = ctx.saved_tensors
+
+ grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply(
+ grad_output, out, ctx.negative_slope, ctx.scale
+ )
+
+ return grad_input, grad_bias, None, None
+
+
+class FusedLeakyReLU(nn.Module):
+ def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5):
+ super().__init__()
+
+ self.bias = nn.Parameter(torch.zeros(channel))
+ self.negative_slope = negative_slope
+ self.scale = scale
+
+ def forward(self, input):
+ return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale)
+
+
+def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5):
+ return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale)
diff --git a/op/fused_act_cpu.py b/op/fused_act_cpu.py
new file mode 100755
index 0000000000000000000000000000000000000000..f997dafdd53aa9f4bbe07af6746c67a2c6dcb4c7
--- /dev/null
+++ b/op/fused_act_cpu.py
@@ -0,0 +1,41 @@
+import os
+
+import torch
+from torch import nn
+from torch.autograd import Function
+from torch.nn import functional as F
+
+
+module_path = os.path.dirname(__file__)
+
+
+class FusedLeakyReLU(nn.Module):
+ def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5):
+ super().__init__()
+
+ self.bias = nn.Parameter(torch.zeros(channel))
+ self.negative_slope = negative_slope
+ self.scale = scale
+
+ def forward(self, input):
+ return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale)
+
+def fused_leaky_relu(input, bias=None, negative_slope=0.2, scale=2 ** 0.5):
+ if input.device.type == "cpu":
+ if bias is not None:
+ rest_dim = [1] * (input.ndim - bias.ndim - 1)
+ return (
+ F.leaky_relu(
+ input + bias.view(1, bias.shape[0], *rest_dim), negative_slope=0.2
+ )
+ * scale
+ )
+
+ else:
+ return F.leaky_relu(input, negative_slope=0.2) * scale
+
+ else:
+ return FusedLeakyReLUFunction.apply(
+ input.contiguous(), bias, negative_slope, scale
+ )
+
diff --git a/op/fused_bias_act.cpp b/op/fused_bias_act.cpp
new file mode 100755
index 0000000000000000000000000000000000000000..02be898f970bcc8ea297867fcaa4e71b24b3d949
--- /dev/null
+++ b/op/fused_bias_act.cpp
@@ -0,0 +1,21 @@
+#include
+
+
+torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,
+ int act, int grad, float alpha, float scale);
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+torch::Tensor fused_bias_act(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,
+ int act, int grad, float alpha, float scale) {
+ CHECK_CUDA(input);
+ CHECK_CUDA(bias);
+
+ return fused_bias_act_op(input, bias, refer, act, grad, alpha, scale);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+ m.def("fused_bias_act", &fused_bias_act, "fused bias act (CUDA)");
+}
\ No newline at end of file
diff --git a/op/fused_bias_act_kernel.cu b/op/fused_bias_act_kernel.cu
new file mode 100755
index 0000000000000000000000000000000000000000..c9fa56fea7ede7072dc8925cfb0148f136eb85b8
--- /dev/null
+++ b/op/fused_bias_act_kernel.cu
@@ -0,0 +1,99 @@
+// Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+//
+// This work is made available under the Nvidia Source Code License-NC.
+// To view a copy of this license, visit
+// https://nvlabs.github.io/stylegan2/license.html
+
+#include
+
+#include
+#include
+#include
+#include
+
+#include
+#include
+
+
+template
+static __global__ void fused_bias_act_kernel(scalar_t* out, const scalar_t* p_x, const scalar_t* p_b, const scalar_t* p_ref,
+ int act, int grad, scalar_t alpha, scalar_t scale, int loop_x, int size_x, int step_b, int size_b, int use_bias, int use_ref) {
+ int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x;
+
+ scalar_t zero = 0.0;
+
+ for (int loop_idx = 0; loop_idx < loop_x && xi < size_x; loop_idx++, xi += blockDim.x) {
+ scalar_t x = p_x[xi];
+
+ if (use_bias) {
+ x += p_b[(xi / step_b) % size_b];
+ }
+
+ scalar_t ref = use_ref ? p_ref[xi] : zero;
+
+ scalar_t y;
+
+ switch (act * 10 + grad) {
+ default:
+ case 10: y = x; break;
+ case 11: y = x; break;
+ case 12: y = 0.0; break;
+
+ case 30: y = (x > 0.0) ? x : x * alpha; break;
+ case 31: y = (ref > 0.0) ? x : x * alpha; break;
+ case 32: y = 0.0; break;
+ }
+
+ out[xi] = y * scale;
+ }
+}
+
+
+torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,
+ int act, int grad, float alpha, float scale) {
+ int curDevice = -1;
+ cudaGetDevice(&curDevice);
+ cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice);
+
+ auto x = input.contiguous();
+ auto b = bias.contiguous();
+ auto ref = refer.contiguous();
+
+ int use_bias = b.numel() ? 1 : 0;
+ int use_ref = ref.numel() ? 1 : 0;
+
+ int size_x = x.numel();
+ int size_b = b.numel();
+ int step_b = 1;
+
+ for (int i = 1 + 1; i < x.dim(); i++) {
+ step_b *= x.size(i);
+ }
+
+ int loop_x = 4;
+ int block_size = 4 * 32;
+ int grid_size = (size_x - 1) / (loop_x * block_size) + 1;
+
+ auto y = torch::empty_like(x);
+
+ AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "fused_bias_act_kernel", [&] {
+ fused_bias_act_kernel<<>>(
+ y.data_ptr(),
+ x.data_ptr(),
+ b.data_ptr(),
+ ref.data_ptr(),
+ act,
+ grad,
+ alpha,
+ scale,
+ loop_x,
+ size_x,
+ step_b,
+ size_b,
+ use_bias,
+ use_ref
+ );
+ });
+
+ return y;
+}
\ No newline at end of file
diff --git a/op/upfirdn2d.cpp b/op/upfirdn2d.cpp
new file mode 100755
index 0000000000000000000000000000000000000000..d2e633dc896433c205e18bc3e455539192ff968e
--- /dev/null
+++ b/op/upfirdn2d.cpp
@@ -0,0 +1,23 @@
+#include
+
+
+torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel,
+ int up_x, int up_y, int down_x, int down_y,
+ int pad_x0, int pad_x1, int pad_y0, int pad_y1);
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+torch::Tensor upfirdn2d(const torch::Tensor& input, const torch::Tensor& kernel,
+ int up_x, int up_y, int down_x, int down_y,
+ int pad_x0, int pad_x1, int pad_y0, int pad_y1) {
+ CHECK_CUDA(input);
+ CHECK_CUDA(kernel);
+
+ return upfirdn2d_op(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+ m.def("upfirdn2d", &upfirdn2d, "upfirdn2d (CUDA)");
+}
\ No newline at end of file
diff --git a/op/upfirdn2d.py b/op/upfirdn2d.py
new file mode 100755
index 0000000000000000000000000000000000000000..f1bbf96777f2c7267c1fef1733972014684ea22b
--- /dev/null
+++ b/op/upfirdn2d.py
@@ -0,0 +1,187 @@
+import os
+
+import torch
+from torch.autograd import Function
+from torch.utils.cpp_extension import load
+
+
+module_path = os.path.dirname(__file__)
+upfirdn2d_op = load(
+ 'upfirdn2d',
+ sources=[
+ os.path.join(module_path, 'upfirdn2d.cpp'),
+ os.path.join(module_path, 'upfirdn2d_kernel.cu'),
+ ],
+)
+
+
+class UpFirDn2dBackward(Function):
+ @staticmethod
+ def forward(
+ ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, in_size, out_size
+ ):
+
+ up_x, up_y = up
+ down_x, down_y = down
+ g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1 = g_pad
+
+ grad_output = grad_output.reshape(-1, out_size[0], out_size[1], 1)
+
+ grad_input = upfirdn2d_op.upfirdn2d(
+ grad_output,
+ grad_kernel,
+ down_x,
+ down_y,
+ up_x,
+ up_y,
+ g_pad_x0,
+ g_pad_x1,
+ g_pad_y0,
+ g_pad_y1,
+ )
+ grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], in_size[3])
+
+ ctx.save_for_backward(kernel)
+
+ pad_x0, pad_x1, pad_y0, pad_y1 = pad
+
+ ctx.up_x = up_x
+ ctx.up_y = up_y
+ ctx.down_x = down_x
+ ctx.down_y = down_y
+ ctx.pad_x0 = pad_x0
+ ctx.pad_x1 = pad_x1
+ ctx.pad_y0 = pad_y0
+ ctx.pad_y1 = pad_y1
+ ctx.in_size = in_size
+ ctx.out_size = out_size
+
+ return grad_input
+
+ @staticmethod
+ def backward(ctx, gradgrad_input):
+ kernel, = ctx.saved_tensors
+
+ gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], ctx.in_size[3], 1)
+
+ gradgrad_out = upfirdn2d_op.upfirdn2d(
+ gradgrad_input,
+ kernel,
+ ctx.up_x,
+ ctx.up_y,
+ ctx.down_x,
+ ctx.down_y,
+ ctx.pad_x0,
+ ctx.pad_x1,
+ ctx.pad_y0,
+ ctx.pad_y1,
+ )
+ # gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.out_size[0], ctx.out_size[1], ctx.in_size[3])
+ gradgrad_out = gradgrad_out.view(
+ ctx.in_size[0], ctx.in_size[1], ctx.out_size[0], ctx.out_size[1]
+ )
+
+ return gradgrad_out, None, None, None, None, None, None, None, None
+
+
+class UpFirDn2d(Function):
+ @staticmethod
+ def forward(ctx, input, kernel, up, down, pad):
+ up_x, up_y = up
+ down_x, down_y = down
+ pad_x0, pad_x1, pad_y0, pad_y1 = pad
+
+ kernel_h, kernel_w = kernel.shape
+ batch, channel, in_h, in_w = input.shape
+ ctx.in_size = input.shape
+
+ input = input.reshape(-1, in_h, in_w, 1)
+
+ ctx.save_for_backward(kernel, torch.flip(kernel, [0, 1]))
+
+ out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1
+ out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1
+ ctx.out_size = (out_h, out_w)
+
+ ctx.up = (up_x, up_y)
+ ctx.down = (down_x, down_y)
+ ctx.pad = (pad_x0, pad_x1, pad_y0, pad_y1)
+
+ g_pad_x0 = kernel_w - pad_x0 - 1
+ g_pad_y0 = kernel_h - pad_y0 - 1
+ g_pad_x1 = in_w * up_x - out_w * down_x + pad_x0 - up_x + 1
+ g_pad_y1 = in_h * up_y - out_h * down_y + pad_y0 - up_y + 1
+
+ ctx.g_pad = (g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1)
+
+ out = upfirdn2d_op.upfirdn2d(
+ input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
+ )
+ # out = out.view(major, out_h, out_w, minor)
+ out = out.view(-1, channel, out_h, out_w)
+
+ return out
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ kernel, grad_kernel = ctx.saved_tensors
+
+ grad_input = UpFirDn2dBackward.apply(
+ grad_output,
+ kernel,
+ grad_kernel,
+ ctx.up,
+ ctx.down,
+ ctx.pad,
+ ctx.g_pad,
+ ctx.in_size,
+ ctx.out_size,
+ )
+
+ return grad_input, None, None, None, None
+
+
+def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
+ out = UpFirDn2d.apply(
+ input, kernel, (up, up), (down, down), (pad[0], pad[1], pad[0], pad[1])
+ )
+
+ return out
+
+
+def upfirdn2d_native(
+ input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
+):
+ _, in_h, in_w, minor = input.shape
+ kernel_h, kernel_w = kernel.shape
+
+ out = input.view(-1, in_h, 1, in_w, 1, minor)
+ out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1])
+ out = out.view(-1, in_h * up_y, in_w * up_x, minor)
+
+ out = F.pad(
+ out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)]
+ )
+ out = out[
+ :,
+ max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0),
+ max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0),
+ :,
+ ]
+
+ out = out.permute(0, 3, 1, 2)
+ out = out.reshape(
+ [-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1]
+ )
+ w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
+ out = F.conv2d(out, w)
+ out = out.reshape(
+ -1,
+ minor,
+ in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
+ in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1,
+ )
+ out = out.permute(0, 2, 3, 1)
+
+ return out[:, ::down_y, ::down_x, :]
+
diff --git a/op/upfirdn2d_cpu.py b/op/upfirdn2d_cpu.py
new file mode 100755
index 0000000000000000000000000000000000000000..a0f820b4c81e03598589b1ea6b95cf9bef9b04f8
--- /dev/null
+++ b/op/upfirdn2d_cpu.py
@@ -0,0 +1,60 @@
+import os
+
+import torch
+from torch.autograd import Function
+from torch.nn import functional as F
+
+
+
+module_path = os.path.dirname(__file__)
+
+def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
+ out = upfirdn2d_native(
+ input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1]
+ )
+
+ return out
+
+
+def upfirdn2d_native(
+ input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
+):
+ _, channel, in_h, in_w = input.shape
+ input = input.reshape(-1, in_h, in_w, 1)
+
+ _, in_h, in_w, minor = input.shape
+ kernel_h, kernel_w = kernel.shape
+
+ out = input.view(-1, in_h, 1, in_w, 1, minor)
+ out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1])
+ out = out.view(-1, in_h * up_y, in_w * up_x, minor)
+
+ out = F.pad(
+ out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)]
+ )
+ out = out[
+ :,
+ max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0),
+ max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0),
+ :,
+ ]
+
+ out = out.permute(0, 3, 1, 2)
+ out = out.reshape(
+ [-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1]
+ )
+ w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
+ out = F.conv2d(out, w)
+ out = out.reshape(
+ -1,
+ minor,
+ in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
+ in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1,
+ )
+ out = out.permute(0, 2, 3, 1)
+ out = out[:, ::down_y, ::down_x, :]
+
+ out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h + down_y) // down_y
+ out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w + down_x) // down_x
+
+ return out.view(-1, channel, out_h, out_w)
diff --git a/op/upfirdn2d_kernel.cu b/op/upfirdn2d_kernel.cu
new file mode 100755
index 0000000000000000000000000000000000000000..2a710aa6adc3d43ac93136a1814e3c39970e1c7e
--- /dev/null
+++ b/op/upfirdn2d_kernel.cu
@@ -0,0 +1,272 @@
+// Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+//
+// This work is made available under the Nvidia Source Code License-NC.
+// To view a copy of this license, visit
+// https://nvlabs.github.io/stylegan2/license.html
+
+#include
+
+#include
+#include
+#include
+#include
+
+#include
+#include
+
+
+static __host__ __device__ __forceinline__ int floor_div(int a, int b) {
+ int c = a / b;
+
+ if (c * b > a) {
+ c--;
+ }
+
+ return c;
+}
+
+
+struct UpFirDn2DKernelParams {
+ int up_x;
+ int up_y;
+ int down_x;
+ int down_y;
+ int pad_x0;
+ int pad_x1;
+ int pad_y0;
+ int pad_y1;
+
+ int major_dim;
+ int in_h;
+ int in_w;
+ int minor_dim;
+ int kernel_h;
+ int kernel_w;
+ int out_h;
+ int out_w;
+ int loop_major;
+ int loop_x;
+};
+
+
+template
+__global__ void upfirdn2d_kernel(scalar_t* out, const scalar_t* input, const scalar_t* kernel, const UpFirDn2DKernelParams p) {
+ const int tile_in_h = ((tile_out_h - 1) * down_y + kernel_h - 1) / up_y + 1;
+ const int tile_in_w = ((tile_out_w - 1) * down_x + kernel_w - 1) / up_x + 1;
+
+ __shared__ volatile float sk[kernel_h][kernel_w];
+ __shared__ volatile float sx[tile_in_h][tile_in_w];
+
+ int minor_idx = blockIdx.x;
+ int tile_out_y = minor_idx / p.minor_dim;
+ minor_idx -= tile_out_y * p.minor_dim;
+ tile_out_y *= tile_out_h;
+ int tile_out_x_base = blockIdx.y * p.loop_x * tile_out_w;
+ int major_idx_base = blockIdx.z * p.loop_major;
+
+ if (tile_out_x_base >= p.out_w | tile_out_y >= p.out_h | major_idx_base >= p.major_dim) {
+ return;
+ }
+
+ for (int tap_idx = threadIdx.x; tap_idx < kernel_h * kernel_w; tap_idx += blockDim.x) {
+ int ky = tap_idx / kernel_w;
+ int kx = tap_idx - ky * kernel_w;
+ scalar_t v = 0.0;
+
+ if (kx < p.kernel_w & ky < p.kernel_h) {
+ v = kernel[(p.kernel_h - 1 - ky) * p.kernel_w + (p.kernel_w - 1 - kx)];
+ }
+
+ sk[ky][kx] = v;
+ }
+
+ for (int loop_major = 0, major_idx = major_idx_base; loop_major < p.loop_major & major_idx < p.major_dim; loop_major++, major_idx++) {
+ for (int loop_x = 0, tile_out_x = tile_out_x_base; loop_x < p.loop_x & tile_out_x < p.out_w; loop_x++, tile_out_x += tile_out_w) {
+ int tile_mid_x = tile_out_x * down_x + up_x - 1 - p.pad_x0;
+ int tile_mid_y = tile_out_y * down_y + up_y - 1 - p.pad_y0;
+ int tile_in_x = floor_div(tile_mid_x, up_x);
+ int tile_in_y = floor_div(tile_mid_y, up_y);
+
+ __syncthreads();
+
+ for (int in_idx = threadIdx.x; in_idx < tile_in_h * tile_in_w; in_idx += blockDim.x) {
+ int rel_in_y = in_idx / tile_in_w;
+ int rel_in_x = in_idx - rel_in_y * tile_in_w;
+ int in_x = rel_in_x + tile_in_x;
+ int in_y = rel_in_y + tile_in_y;
+
+ scalar_t v = 0.0;
+
+ if (in_x >= 0 & in_y >= 0 & in_x < p.in_w & in_y < p.in_h) {
+ v = input[((major_idx * p.in_h + in_y) * p.in_w + in_x) * p.minor_dim + minor_idx];
+ }
+
+ sx[rel_in_y][rel_in_x] = v;
+ }
+
+ __syncthreads();
+ for (int out_idx = threadIdx.x; out_idx < tile_out_h * tile_out_w; out_idx += blockDim.x) {
+ int rel_out_y = out_idx / tile_out_w;
+ int rel_out_x = out_idx - rel_out_y * tile_out_w;
+ int out_x = rel_out_x + tile_out_x;
+ int out_y = rel_out_y + tile_out_y;
+
+ int mid_x = tile_mid_x + rel_out_x * down_x;
+ int mid_y = tile_mid_y + rel_out_y * down_y;
+ int in_x = floor_div(mid_x, up_x);
+ int in_y = floor_div(mid_y, up_y);
+ int rel_in_x = in_x - tile_in_x;
+ int rel_in_y = in_y - tile_in_y;
+ int kernel_x = (in_x + 1) * up_x - mid_x - 1;
+ int kernel_y = (in_y + 1) * up_y - mid_y - 1;
+
+ scalar_t v = 0.0;
+
+ #pragma unroll
+ for (int y = 0; y < kernel_h / up_y; y++)
+ #pragma unroll
+ for (int x = 0; x < kernel_w / up_x; x++)
+ v += sx[rel_in_y + y][rel_in_x + x] * sk[kernel_y + y * up_y][kernel_x + x * up_x];
+
+ if (out_x < p.out_w & out_y < p.out_h) {
+ out[((major_idx * p.out_h + out_y) * p.out_w + out_x) * p.minor_dim + minor_idx] = v;
+ }
+ }
+ }
+ }
+}
+
+
+torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel,
+ int up_x, int up_y, int down_x, int down_y,
+ int pad_x0, int pad_x1, int pad_y0, int pad_y1) {
+ int curDevice = -1;
+ cudaGetDevice(&curDevice);
+ cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice);
+
+ UpFirDn2DKernelParams p;
+
+ auto x = input.contiguous();
+ auto k = kernel.contiguous();
+
+ p.major_dim = x.size(0);
+ p.in_h = x.size(1);
+ p.in_w = x.size(2);
+ p.minor_dim = x.size(3);
+ p.kernel_h = k.size(0);
+ p.kernel_w = k.size(1);
+ p.up_x = up_x;
+ p.up_y = up_y;
+ p.down_x = down_x;
+ p.down_y = down_y;
+ p.pad_x0 = pad_x0;
+ p.pad_x1 = pad_x1;
+ p.pad_y0 = pad_y0;
+ p.pad_y1 = pad_y1;
+
+ p.out_h = (p.in_h * p.up_y + p.pad_y0 + p.pad_y1 - p.kernel_h + p.down_y) / p.down_y;
+ p.out_w = (p.in_w * p.up_x + p.pad_x0 + p.pad_x1 - p.kernel_w + p.down_x) / p.down_x;
+
+ auto out = at::empty({p.major_dim, p.out_h, p.out_w, p.minor_dim}, x.options());
+
+ int mode = -1;
+
+ int tile_out_h;
+ int tile_out_w;
+
+ if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 4 && p.kernel_w <= 4) {
+ mode = 1;
+ tile_out_h = 16;
+ tile_out_w = 64;
+ }
+
+ if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 3 && p.kernel_w <= 3) {
+ mode = 2;
+ tile_out_h = 16;
+ tile_out_w = 64;
+ }
+
+ if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 4 && p.kernel_w <= 4) {
+ mode = 3;
+ tile_out_h = 16;
+ tile_out_w = 64;
+ }
+
+ if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 2 && p.kernel_w <= 2) {
+ mode = 4;
+ tile_out_h = 16;
+ tile_out_w = 64;
+ }
+
+ if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 && p.kernel_h <= 4 && p.kernel_w <= 4) {
+ mode = 5;
+ tile_out_h = 8;
+ tile_out_w = 32;
+ }
+
+ if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 && p.kernel_h <= 2 && p.kernel_w <= 2) {
+ mode = 6;
+ tile_out_h = 8;
+ tile_out_w = 32;
+ }
+
+ dim3 block_size;
+ dim3 grid_size;
+
+ if (tile_out_h > 0 && tile_out_w) {
+ p.loop_major = (p.major_dim - 1) / 16384 + 1;
+ p.loop_x = 1;
+ block_size = dim3(32 * 8, 1, 1);
+ grid_size = dim3(((p.out_h - 1) / tile_out_h + 1) * p.minor_dim,
+ (p.out_w - 1) / (p.loop_x * tile_out_w) + 1,
+ (p.major_dim - 1) / p.loop_major + 1);
+ }
+
+ AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "upfirdn2d_cuda", [&] {
+ switch (mode) {
+ case 1:
+ upfirdn2d_kernel<<>>(
+ out.data_ptr(), x.data_ptr(), k.data_ptr(), p
+ );
+
+ break;
+
+ case 2:
+ upfirdn2d_kernel<<>>(
+ out.data_ptr(), x.data_ptr(), k.data_ptr(), p
+ );
+
+ break;
+
+ case 3:
+ upfirdn2d_kernel<<>>(
+ out.data_ptr(), x.data_ptr(), k.data_ptr(), p
+ );
+
+ break;
+
+ case 4:
+ upfirdn2d_kernel<<>>(
+ out.data_ptr(), x.data_ptr(), k.data_ptr(), p
+ );
+
+ break;
+
+ case 5:
+ upfirdn2d_kernel<<>>(
+ out.data_ptr(), x.data_ptr(), k.data_ptr(), p
+ );
+
+ break;
+
+ case 6:
+ upfirdn2d_kernel<<>>(
+ out.data_ptr(), x.data_ptr(), k.data_ptr(), p
+ );
+
+ break;
+ }
+ });
+
+ return out;
+}
\ No newline at end of file
diff --git a/packages.txt b/packages.txt
new file mode 100644
index 0000000000000000000000000000000000000000..ace9fff330d5ce42b11d0e450d7cde99ccecfa77
--- /dev/null
+++ b/packages.txt
@@ -0,0 +1,4 @@
+ffmpeg
+libsm6
+libxext6
+cmake
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..6ef55b4ad5e0677f9545ce5e2fc7656b32913bb3
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,13 @@
+tqdm
+gdown
+scikit-learn==0.22
+scipy
+lpips
+opencv-python-headless
+torch
+torchvision
+imageio
+dlib
+ftfy
+regex
+git+https://github.com/openai/CLIP.git
\ No newline at end of file
diff --git a/styleclip/styleclip_global.py b/styleclip/styleclip_global.py
new file mode 100644
index 0000000000000000000000000000000000000000..96fa7569ebd51a5e6c2deddb57ccceb4f4376904
--- /dev/null
+++ b/styleclip/styleclip_global.py
@@ -0,0 +1,181 @@
+'''
+Code adapted from Stitch it in Time by Tzaban et al.
+https://github.com/rotemtzaban/STIT
+'''
+
+
+import numpy as np
+import torch
+from tqdm import tqdm
+from pathlib import Path
+import os
+
+import clip
+
+imagenet_templates = [
+ 'a bad photo of a {}.',
+ 'a photo of many {}.',
+ 'a sculpture of a {}.',
+ 'a photo of the hard to see {}.',
+ 'a low resolution photo of the {}.',
+ 'a rendering of a {}.',
+ 'graffiti of a {}.',
+ 'a bad photo of the {}.',
+ 'a cropped photo of the {}.',
+ 'a tattoo of a {}.',
+ 'the embroidered {}.',
+ 'a photo of a hard to see {}.',
+ 'a bright photo of a {}.',
+ 'a photo of a clean {}.',
+ 'a photo of a dirty {}.',
+ 'a dark photo of the {}.',
+ 'a drawing of a {}.',
+ 'a photo of my {}.',
+ 'the plastic {}.',
+ 'a photo of the cool {}.',
+ 'a close-up photo of a {}.',
+ 'a black and white photo of the {}.',
+ 'a painting of the {}.',
+ 'a painting of a {}.',
+ 'a pixelated photo of the {}.',
+ 'a sculpture of the {}.',
+ 'a bright photo of the {}.',
+ 'a cropped photo of a {}.',
+ 'a plastic {}.',
+ 'a photo of the dirty {}.',
+ 'a jpeg corrupted photo of a {}.',
+ 'a blurry photo of the {}.',
+ 'a photo of the {}.',
+ 'a good photo of the {}.',
+ 'a rendering of the {}.',
+ 'a {} in a video game.',
+ 'a photo of one {}.',
+ 'a doodle of a {}.',
+ 'a close-up photo of the {}.',
+ 'a photo of a {}.',
+ 'the origami {}.',
+ 'the {} in a video game.',
+ 'a sketch of a {}.',
+ 'a doodle of the {}.',
+ 'a origami {}.',
+ 'a low resolution photo of a {}.',
+ 'the toy {}.',
+ 'a rendition of the {}.',
+ 'a photo of the clean {}.',
+ 'a photo of a large {}.',
+ 'a rendition of a {}.',
+ 'a photo of a nice {}.',
+ 'a photo of a weird {}.',
+ 'a blurry photo of a {}.',
+ 'a cartoon {}.',
+ 'art of a {}.',
+ 'a sketch of the {}.',
+ 'a embroidered {}.',
+ 'a pixelated photo of a {}.',
+ 'itap of the {}.',
+ 'a jpeg corrupted photo of the {}.',
+ 'a good photo of a {}.',
+ 'a plushie {}.',
+ 'a photo of the nice {}.',
+ 'a photo of the small {}.',
+ 'a photo of the weird {}.',
+ 'the cartoon {}.',
+ 'art of the {}.',
+ 'a drawing of the {}.',
+ 'a photo of the large {}.',
+ 'a black and white photo of a {}.',
+ 'the plushie {}.',
+ 'a dark photo of a {}.',
+ 'itap of a {}.',
+ 'graffiti of the {}.',
+ 'a toy {}.',
+ 'itap of my {}.',
+ 'a photo of a cool {}.',
+ 'a photo of a small {}.',
+ 'a tattoo of the {}.',
+]
+
+CONV_CODE_INDICES = [(0, 512), (1024, 1536), (1536, 2048), (2560, 3072), (3072, 3584), (4096, 4608), (4608, 5120), (5632, 6144), (6144, 6656), (7168, 7680), (7680, 7936), (8192, 8448), (8448, 8576), (8704, 8832), (8832, 8896), (8960, 9024), (9024, 9056)]
+FFHQ_CODE_INDICES = [(0, 512), (512, 1024), (1024, 1536), (1536, 2048), (2560, 3072), (3072, 3584), (4096, 4608), (4608, 5120), (5632, 6144), (6144, 6656), (7168, 7680), (7680, 7936), (8192, 8448), (8448, 8576), (8704, 8832), (8832, 8896), (8960, 9024), (9024, 9056)] + \
+ [(2048, 2560), (3584, 4096), (5120, 5632), (6656, 7168), (7936, 8192), (8576, 8704), (8896, 8960), (9056, 9088)]
+
+def zeroshot_classifier(model, classnames, templates, device):
+
+ with torch.no_grad():
+ zeroshot_weights = []
+ for classname in tqdm(classnames):
+ texts = [template.format(classname) for template in templates] # format with class
+ texts = clip.tokenize(texts).to(device) # tokenize
+ class_embeddings = model.encode_text(texts) # embed with text encoder
+ class_embeddings /= class_embeddings.norm(dim=-1, keepdim=True)
+ class_embedding = class_embeddings.mean(dim=0)
+ class_embedding /= class_embedding.norm()
+ zeroshot_weights.append(class_embedding)
+ zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(device)
+ return zeroshot_weights
+
+def expand_to_full_dim(partial_tensor):
+ full_dim_tensor = torch.zeros(size=(1, 9088))
+
+ start_idx = 0
+ for conv_start, conv_end in CONV_CODE_INDICES:
+ length = conv_end - conv_start
+ full_dim_tensor[:, conv_start:conv_end] = partial_tensor[start_idx:start_idx + length]
+ start_idx += length
+
+ return full_dim_tensor
+
+def get_direction(neutral_class, target_class, beta, di, clip_model=None):
+
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+
+ if clip_model is None:
+ clip_model, _ = clip.load("ViT-B/32", device=device)
+
+ class_names = [neutral_class, target_class]
+ class_weights = zeroshot_classifier(clip_model, class_names, imagenet_templates, device)
+
+ dt = class_weights[:, 1] - class_weights[:, 0]
+ dt = dt / dt.norm()
+
+ dt = dt.float()
+ di = di.float()
+
+ relevance = di @ dt
+ mask = relevance.abs() > beta
+ direction = relevance * mask
+ direction_max = direction.abs().max()
+ if direction_max > 0:
+ direction = direction / direction_max
+ else:
+ raise ValueError(f'Beta value {beta} is too high for mapping from {neutral_class} to {target_class},'
+ f' try setting it to a lower value')
+ return direction
+
+def style_tensor_to_style_dict(style_tensor, refernce_generator):
+ style_layers = refernce_generator.modulation_layers
+
+ style_dict = {}
+ for layer_idx, layer in enumerate(style_layers):
+ style_dict[layer] = style_tensor[:, FFHQ_CODE_INDICES[layer_idx][0]:FFHQ_CODE_INDICES[layer_idx][1]]
+
+ return style_dict
+
+def style_dict_to_style_tensor(style_dict, reference_generator):
+ style_layers = reference_generator.modulation_layers
+
+ style_tensor = torch.zeros(size=(1, 9088))
+ for layer in style_dict:
+ layer_idx = style_layers.index(layer)
+ style_tensor[:, FFHQ_CODE_INDICES[layer_idx][0]:FFHQ_CODE_INDICES[layer_idx][1]] = style_dict[layer]
+
+ return style_tensor
+
+def project_code_with_styleclip(source_latent, source_class, target_class, alpha, beta, reference_generator, di, clip_model=None):
+ edit_direction = get_direction(source_class, target_class, beta, di, clip_model)
+
+ edit_full_dim = expand_to_full_dim(edit_direction)
+
+ source_s = style_dict_to_style_tensor(source_latent, reference_generator)
+
+ return source_s + alpha * edit_full_dim
\ No newline at end of file
diff --git a/util.py b/util.py
new file mode 100644
index 0000000000000000000000000000000000000000..083b56170f5feb72eccfebd38a53aed70db32064
--- /dev/null
+++ b/util.py
@@ -0,0 +1,136 @@
+from matplotlib import pyplot as plt
+import torch
+import torch.nn.functional as F
+import os
+import dlib
+from PIL import Image
+import numpy as np
+import scipy
+import scipy.ndimage
+import torchvision.transforms as transforms
+
+def display_image(image, size=None, mode='nearest', unnorm=False, title=''):
+ # image is [3,h,w] or [1,3,h,w] tensor [0,1]
+ if not isinstance(image, torch.Tensor):
+ image = transforms.ToTensor()(image).unsqueeze(0)
+ if image.is_cuda:
+ image = image.cpu()
+ if size is not None and image.size(-1) != size:
+ image = F.interpolate(image, size=(size,size), mode=mode)
+ if image.dim() == 4:
+ image = image[0]
+ image = image.permute(1, 2, 0).detach().numpy()
+ plt.figure()
+ plt.title(title)
+ plt.axis('off')
+ plt.imshow(image)
+
+def get_landmark(filepath, predictor):
+ """get landmark with dlib
+ :return: np.array shape=(68, 2)
+ """
+ detector = dlib.get_frontal_face_detector()
+
+ img = dlib.load_rgb_image(filepath)
+ dets = detector(img, 1)
+ assert len(dets) > 0, "Face not detected, try another face image"
+
+ for k, d in enumerate(dets):
+ shape = predictor(img, d)
+
+ t = list(shape.parts())
+ a = []
+ for tt in t:
+ a.append([tt.x, tt.y])
+ lm = np.array(a)
+ return lm
+
+def align_face(filepath, predictor, output_size=256, transform_size=1024, enable_padding=True):
+
+ """
+ :param filepath: str
+ :return: PIL Image
+ """
+ lm = get_landmark(filepath, predictor)
+
+ lm_chin = lm[0: 17] # left-right
+ lm_eyebrow_left = lm[17: 22] # left-right
+ lm_eyebrow_right = lm[22: 27] # left-right
+ lm_nose = lm[27: 31] # top-down
+ lm_nostrils = lm[31: 36] # top-down
+ lm_eye_left = lm[36: 42] # left-clockwise
+ lm_eye_right = lm[42: 48] # left-clockwise
+ lm_mouth_outer = lm[48: 60] # left-clockwise
+ lm_mouth_inner = lm[60: 68] # left-clockwise
+
+ # Calculate auxiliary vectors.
+ eye_left = np.mean(lm_eye_left, axis=0)
+ eye_right = np.mean(lm_eye_right, axis=0)
+ eye_avg = (eye_left + eye_right) * 0.5
+ eye_to_eye = eye_right - eye_left
+ mouth_left = lm_mouth_outer[0]
+ mouth_right = lm_mouth_outer[6]
+ mouth_avg = (mouth_left + mouth_right) * 0.5
+ eye_to_mouth = mouth_avg - eye_avg
+
+ # Choose oriented crop rectangle.
+ x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1]
+ x /= np.hypot(*x)
+ x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8)
+ y = np.flipud(x) * [-1, 1]
+ c = eye_avg + eye_to_mouth * 0.1
+ quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])
+ qsize = np.hypot(*x) * 2
+
+ # read image
+ img = Image.open(filepath)
+
+ transform_size = output_size
+ enable_padding = True
+
+ # Shrink.
+ shrink = int(np.floor(qsize / output_size * 0.5))
+ if shrink > 1:
+ rsize = (int(np.rint(float(img.size[0]) / shrink)), int(np.rint(float(img.size[1]) / shrink)))
+ img = img.resize(rsize, Image.ANTIALIAS)
+ quad /= shrink
+ qsize /= shrink
+
+ # Crop.
+ border = max(int(np.rint(qsize * 0.1)), 3)
+ crop = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),
+ int(np.ceil(max(quad[:, 1]))))
+ crop = (max(crop[0] - border, 0), max(crop[1] - border, 0), min(crop[2] + border, img.size[0]),
+ min(crop[3] + border, img.size[1]))
+ if crop[2] - crop[0] < img.size[0] or crop[3] - crop[1] < img.size[1]:
+ img = img.crop(crop)
+ quad -= crop[0:2]
+
+ # Pad.
+ pad = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),
+ int(np.ceil(max(quad[:, 1]))))
+ pad = (max(-pad[0] + border, 0), max(-pad[1] + border, 0), max(pad[2] - img.size[0] + border, 0),
+ max(pad[3] - img.size[1] + border, 0))
+ if enable_padding and max(pad) > border - 4:
+ pad = np.maximum(pad, int(np.rint(qsize * 0.3)))
+ img = np.pad(np.float32(img), ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect')
+ h, w, _ = img.shape
+ y, x, _ = np.ogrid[:h, :w, :1]
+ mask = np.maximum(1.0 - np.minimum(np.float32(x) / pad[0], np.float32(w - 1 - x) / pad[2]),
+ 1.0 - np.minimum(np.float32(y) / pad[1], np.float32(h - 1 - y) / pad[3]))
+ blur = qsize * 0.02
+ img += (scipy.ndimage.gaussian_filter(img, [blur, blur, 0]) - img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0)
+ img += (np.median(img, axis=(0, 1)) - img) * np.clip(mask, 0.0, 1.0)
+ img = Image.fromarray(np.uint8(np.clip(np.rint(img), 0, 255)), 'RGB')
+ quad += pad[:2]
+
+ # Transform.
+ img = img.transform((transform_size, transform_size), Image.QUAD, (quad + 0.5).flatten(), Image.BILINEAR)
+ if output_size < transform_size:
+ img = img.resize((output_size, output_size), Image.ANTIALIAS)
+
+ # Return aligned image.
+ return img
+
+def strip_path_extension(path):
+ return os.path.splitext(path)[0]
\ No newline at end of file