first

.gitattributes

@@ -1,3 +1,5 @@
 *.gif filter=lfs diff=lfs merge=lfs -text
 *.png filter=lfs diff=lfs merge=lfs -text
 *.jpg filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+. filter=lfs diff=lfs merge=lfs -text

PTI/.gitignore

@@ -0,0 +1,4 @@
+checkpoints
+__pycache__
+embeddings
+test

PTI/LICENSE

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

PTI/README.md

@@ -0,0 +1,230 @@
+# PTI: Pivotal Tuning for Latent-based editing of Real Images     (ACM TOG 2022)
+
+<!-- > Recently, a surge of advanced facial editing techniques have been proposed
+that leverage the generative power of a pre-trained StyleGAN. To successfully
+edit an image this way, one must first project (or invert) the image into
+the pre-trained generator’s domain. As it turns out, however, StyleGAN’s
+latent space induces an inherent tradeoff between distortion and editability,
+i.e. between maintaining the original appearance and convincingly altering
+some of its attributes. Practically, this means it is still challenging to
+apply ID-preserving facial latent-space editing to faces which are out of the
+generator’s domain. In this paper, we present an approach to bridge this
+gap. Our technique slightly alters the generator, so that an out-of-domain
+image is faithfully mapped into an in-domain latent code. The key idea is
+pivotal tuning — a brief training process that preserves the editing quality
+of an in-domain latent region, while changing its portrayed identity and
+appearance. In Pivotal Tuning Inversion (PTI), an initial inverted latent code
+serves as a pivot, around which the generator is fined-tuned. At the same
+time, a regularization term keeps nearby identities intact, to locally contain
+the effect. This surgical training process ends up altering appearance features
+that represent mostly identity, without affecting editing capabilities.
+To supplement this, we further show that pivotal tuning can also adjust the
+generator to accommodate a multitude of faces, while introducing negligible
+distortion on the rest of the domain. We validate our technique through
+inversion and editing metrics, and show preferable scores to state-of-the-art
+methods. We further qualitatively demonstrate our technique by applying
+advanced edits (such as pose, age, or expression) to numerous images of
+well-known and recognizable identities. Finally, we demonstrate resilience
+to harder cases, including heavy make-up, elaborate hairstyles and/or headwear,
+which otherwise could not have been successfully inverted and edited
+by state-of-the-art methods. -->
+
+<a href="https://arxiv.org/abs/2106.05744"><img src="https://img.shields.io/badge/arXiv-2008.00951-b31b1b.svg"></a>
+<a href="https://opensource.org/licenses/MIT"><img src="https://img.shields.io/badge/License-MIT-yellow.svg"></a>  
+Inference Notebook: <a href="https://colab.research.google.com/github/danielroich/PTI/blob/main/notebooks/inference_playground.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" height=20></a>  
+
+<p align="center">
+<img src="docs/teaser.jpg"/>  
+<br>
+Pivotal Tuning Inversion (PTI) enables employing off-the-shelf latent based
+semantic editing techniques on real images using StyleGAN. 
+PTI excels in identity preserving edits, portrayed through recognizable figures —
+Serena Williams and Robert Downey Jr. (top), and in handling faces which
+are clearly out-of-domain, e.g., due to heavy makeup (bottom).
+</br>
+</p>
+
+## Description   
+Official Implementation of our PTI paper + code for evaluation metrics. PTI introduces an optimization mechanizem for solving the StyleGAN inversion task.
+Providing near-perfect reconstruction results while maintaining the high editing abilitis of the native StyleGAN latent space W. For more details, see <a href="https://arxiv.org/abs/2106.05744"><img src="https://img.shields.io/badge/arXiv-2008.00951-b31b1b.svg"></a>
+
+## Recent Updates
+**2021.07.01**: Fixed files download phase in the inference notebook. Which might caused the notebook not to run smoothly.
+
+**2021.06.29**: Added support for CPU. In order to run PTI on CPU please change `device` parameter under `configs/global_config.py` to "cpu" instead of "cuda".
+
+**2021.06.25** : Adding mohawk edit using StyleCLIP+PTI in inference notebook.
+	      Updating documentation in inference notebook due to Google Drive rate limit reached.
+	      Currently, Google Drive does not allow to download the pretrined models using Colab automatically. Manual intervention might be needed.
+
+## Getting Started
+### Prerequisites
+- Linux or macOS
+- NVIDIA GPU + CUDA CuDNN (Not mandatory bur recommended)
+- Python 3
+
+### Installation
+- Dependencies:  
+	1. lpips
+	2. wandb
+	3. pytorch
+	4. torchvision
+	5. matplotlib
+	6. dlib
+- All dependencies can be installed using *pip install* and the package name
+
+## Pretrained Models
+Please download the pretrained models from the following links.
+
+### Auxiliary Models
+We provide various auxiliary models needed for PTI inversion task.  
+This includes the StyleGAN generator and pre-trained models used for loss computation.
+| Path | Description
+| :--- | :----------
+|[FFHQ StyleGAN](https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/ffhq.pkl) | StyleGAN2-ada model trained on FFHQ with 1024x1024 output resolution.
+|[Dlib alignment](https://drive.google.com/file/d/1HKmjg6iXsWr4aFPuU0gBXPGR83wqMzq7/view?usp=sharing) | Dlib alignment used for images preproccessing.
+|[FFHQ e4e encoder](https://drive.google.com/file/d/1ALC5CLA89Ouw40TwvxcwebhzWXM5YSCm/view?usp=sharing) | Pretrained e4e encoder. Used for StyleCLIP editing.
+
+Note: The StyleGAN model is used directly from the official [stylegan2-ada-pytorch implementation](https://github.com/NVlabs/stylegan2-ada-pytorch).
+For StyleCLIP pretrained mappers, please see [StyleCLIP's official routes](https://github.com/orpatashnik/StyleCLIP/blob/main/utils.py)
+
+
+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`. 
+
+
+## Inversion
+### Preparing your Data
+In order to invert a real image and edit it you should first align and crop it to the correct size. To do so you should perform *One* of the following steps: 
+1. Run `notebooks/align_data.ipynb` and change the "images_path" variable to the raw images path
+2. Run `utils/align_data.py` and change the "images_path" variable to the raw images path
+
+
+### Weights And Biases
+The project supports [Weights And Biases](https://wandb.ai/home) framework for experiment tracking. For the inversion task it enables visualization of the losses progression and the generator intermediate results during the initial inversion and the *Pivotal Tuning*(PT) procedure.
+
+The log frequency can be adjusted using the parameters defined at `configs/global_config.py` under the "Logs" subsection.
+
+There is no no need to have an account. However, in order to use the features provided by Weights and Biases you first have to register on their site.
+
+
+### Running PTI
+The main training script is `scripts/run_pti.py`. The script receives aligned and cropped images from paths configured in the "Input info" subscetion in
+ `configs/paths_config.py`. 
+Results are saved to directories found at "Dirs for output files" under `configs/paths_config.py`. This includes inversion latent codes and tuned generators. 
+The hyperparametrs for the inversion task can be found at  `configs/hyperparameters.py`. They are intilized to the default values used in the paper. 
+
+## Editing
+By default, we assume that all auxiliary edit directions are downloaded and saved to the directory `editings`. 
+However, you may use your own paths by changing the necessary values in `configs/path_configs.py` under "Edit directions" subsection.
+
+Example of editing code can be found at `scripts/latent_editor_wrapper.py`
+
+## Inference Notebooks
+To help visualize the results of PTI we provide a Jupyter notebook found in `notebooks/inference_playground.ipynb`.   
+The notebook will download the pretrained models and run inference on a sample image found online or 
+on images of your choosing. It is recommended to run this in [Google Colab](https://colab.research.google.com/github/danielroich/PTI/blob/main/notebooks/inference_playground.ipynb).
+
+The notebook demonstrates how to:
+- Invert an image using PTI
+- Visualise the inversion and use the PTI output
+- Edit the image after PTI using InterfaceGAN and StyleCLIP
+- Compare to other inversion methods
+
+## Evaluation
+Currently the repository supports qualitative evaluation for reconstruction of: PTI, SG2 (*W Space*), e4e, SG2Plus (*W+ Space*). 
+As well as editing using InterfaceGAN and GANSpace for the same inversion methods.
+To run the evaluation please see `evaluation/qualitative_edit_comparison.py`. Examples of the evaluation scripts are:
+
+<p align="center">
+<img src="docs/model_rec.jpg"/>  
+<br>
+Reconsturction comparison between different methods. The images order is: Original image, W+ inversion, e4e inversion, W inversion, PTI inversion
+</br>  
+</p>
+
+<p align="center">
+<img src="docs/stern_rotation.jpg"/>  
+<br>
+InterfaceGAN pose edit comparison between different methods. The images order is: Original, W+, e4e, W, PTI
+</br>  
+</p>
+
+<p align="center">
+<img src="docs/tyron_original.jpg" width="220" height="220"/>  
+<img src="docs/tyron_edit.jpg" width="220" height="220"/>
+<br>
+Image per edit or several edits without comparison
+</br>  
+</p>
+
+###  Coming Soon - Quantitative evaluation and StyleCLIP qualitative evaluation
+
+## Repository structure
+| Path | Description <img width=200>
+| :--- | :---
+| &boxvr;&nbsp; configs | Folder containing configs defining Hyperparameters, paths and logging
+| &boxvr;&nbsp; criteria | Folder containing various loss and regularization criterias for the optimization
+| &boxvr;&nbsp; dnnlib | Folder containing internal utils for StyleGAN2-ada
+| &boxvr;&nbsp; docs | Folder containing the latent space edit directions
+| &boxvr;&nbsp; editings | Folder containing images displayed in the README
+| &boxvr;&nbsp; environment | Folder containing Anaconda environment used in our experiments
+| &boxvr;&nbsp; licenses | Folder containing licenses of the open source projects used in this repository
+| &boxvr;&nbsp; models | Folder containing models used in different editing techniques and first phase inversion
+| &boxvr;&nbsp; notebooks | Folder with jupyter notebooks to demonstrate the usage of PTI end-to-end
+| &boxvr;&nbsp; scripts | Folder with running scripts for inversion, editing and metric computations
+| &boxvr;&nbsp; torch_utils | Folder containing internal utils for StyleGAN2-ada
+| &boxvr;&nbsp; training | Folder containing the core training logic of PTI
+| &boxvr;&nbsp; utils | Folder with various utility functions
+
+
+## Credits
+**StyleGAN2-ada model and implementation:**  
+https://github.com/NVlabs/stylegan2-ada-pytorch
+Copyright © 2021, NVIDIA Corporation.  
+Nvidia Source Code License https://nvlabs.github.io/stylegan2-ada-pytorch/license.html
+
+**LPIPS model and implementation:**  
+https://github.com/richzhang/PerceptualSimilarity  
+Copyright (c) 2020, Sou Uchida  
+License (BSD 2-Clause) https://github.com/richzhang/PerceptualSimilarity/blob/master/LICENSE
+
+**e4e model and implementation:**   
+https://github.com/omertov/encoder4editing
+Copyright (c) 2021 omertov  
+License (MIT) https://github.com/omertov/encoder4editing/blob/main/LICENSE
+
+**StyleCLIP model and implementation:**   
+https://github.com/orpatashnik/StyleCLIP
+Copyright (c) 2021 orpatashnik  
+License (MIT) https://github.com/orpatashnik/StyleCLIP/blob/main/LICENSE
+
+**InterfaceGAN implementation:**   
+https://github.com/genforce/interfacegan
+Copyright (c) 2020 genforce  
+License (MIT) https://github.com/genforce/interfacegan/blob/master/LICENSE
+
+**GANSpace implementation:**   
+https://github.com/harskish/ganspace
+Copyright (c) 2020 harkish  
+License (Apache License 2.0) https://github.com/harskish/ganspace/blob/master/LICENSE
+
+
+## Acknowledgments
+This repository structure is based on [encoder4editing](https://github.com/omertov/encoder4editing) and [ReStyle](https://github.com/yuval-alaluf/restyle-encoder) repositories
+
+## Contact
+For any inquiry please contact us at our email addresses: [email protected] or [email protected]
+
+
+## Citation
+If you use this code for your research, please cite:
+```
+@article{roich2021pivotal,
+  title={Pivotal Tuning for Latent-based Editing of Real Images},
+  author={Roich, Daniel and Mokady, Ron and Bermano, Amit H and Cohen-Or, Daniel},
+  publisher = {Association for Computing Machinery},
+  journal={ACM Trans. Graph.},
+  year={2021}
+}
+```

PTI/configs/evaluation_config.py

@@ -0,0 +1 @@
+evaluated_methods = ['e4e', 'SG2', 'SG2Plus']

PTI/configs/global_config.py

@@ -0,0 +1,12 @@
+# Device
+cuda_visible_devices = '0'
+device = 'cuda:0'
+
+# Logs
+training_step = 1
+image_rec_result_log_snapshot = 100
+pivotal_training_steps = 0
+model_snapshot_interval = 400
+
+# Run name to be updated during PTI
+run_name = ''

PTI/configs/hyperparameters.py

@@ -0,0 +1,28 @@
+## Architechture
+lpips_type = "alex"
+first_inv_type = "w"
+optim_type = "adam"
+
+## Locality regularization
+latent_ball_num_of_samples = 1
+locality_regularization_interval = 1
+use_locality_regularization = False
+regulizer_l2_lambda = 0.1
+regulizer_lpips_lambda = 0.1
+regulizer_alpha = 30
+
+## Loss
+pt_l2_lambda = 1
+pt_lpips_lambda = 1
+
+## Steps
+LPIPS_value_threshold = 0.06
+max_pti_steps = 350
+first_inv_steps = 450
+max_images_to_invert = 30
+
+## Optimization
+pti_learning_rate = 3e-4
+first_inv_lr = 5e-3
+train_batch_size = 1
+use_last_w_pivots = False

PTI/configs/paths_config.py

@@ -0,0 +1,31 @@
+## Pretrained models paths
+e4e = 'PTI/pretrained_models/e4e_ffhq_encode.pt'
+stylegan2_ada_ffhq = '../PTI/pretrained_models/ffhq.pkl'
+style_clip_pretrained_mappers = ''
+ir_se50 = 'PTI/pretrained_models/model_ir_se50.pth'
+dlib = 'PTI/pretrained_models/align.dat'
+
+## Dirs for output files
+checkpoints_dir = 'PTI/checkpoints'
+embedding_base_dir = 'PTI/embeddings'
+styleclip_output_dir = 'PTI/StyleCLIP_results'
+experiments_output_dir = 'PTI/output'
+
+## Input info
+### Input dir, where the images reside
+input_data_path = ''
+### Inversion identifier, used to keeping track of the inversion results. Both the latent code and the generator
+input_data_id = 'barcelona'
+
+## Keywords
+pti_results_keyword = 'PTI'
+e4e_results_keyword = 'e4e'
+sg2_results_keyword = 'SG2'
+sg2_plus_results_keyword = 'SG2_plus'
+multi_id_model_type = 'multi_id'
+
+## Edit directions
+interfacegan_age = 'PTI/editings/interfacegan_directions/age.pt'
+interfacegan_smile = 'PTI/editings/interfacegan_directions/smile.pt'
+interfacegan_rotation = 'PTI/editings/interfacegan_directions/rotation.pt'
+ffhq_pca = 'PTI/editings/ganspace_pca/ffhq_pca.pt'

PTI/criteria/l2_loss.py

@@ -0,0 +1,8 @@
+import torch
+
+l2_criterion = torch.nn.MSELoss(reduction='mean')
+
+
+def l2_loss(real_images, generated_images):
+    loss = l2_criterion(real_images, generated_images)
+    return loss

PTI/criteria/localitly_regulizer.py

@@ -0,0 +1,65 @@
+import torch
+import numpy as np
+from PTI.criteria import l2_loss
+from PTI.configs import hyperparameters
+from PTI.configs import global_config
+
+
+class Space_Regulizer:
+    def __init__(self, original_G, lpips_net):
+        self.original_G = original_G
+        self.morphing_regulizer_alpha = hyperparameters.regulizer_alpha
+        self.lpips_loss = lpips_net
+
+    def get_morphed_w_code(self, new_w_code, fixed_w):
+        interpolation_direction = new_w_code - fixed_w
+        interpolation_direction_norm = torch.norm(interpolation_direction, p=2)
+        direction_to_move = hyperparameters.regulizer_alpha * \
+            interpolation_direction / interpolation_direction_norm
+        result_w = fixed_w + direction_to_move
+        self.morphing_regulizer_alpha * fixed_w + \
+            (1 - self.morphing_regulizer_alpha) * new_w_code
+
+        return result_w
+
+    def get_image_from_ws(self, w_codes, G):
+        return torch.cat([G.synthesis(w_code, noise_mode='none', force_fp32=True) for w_code in w_codes])
+
+    def ball_holder_loss_lazy(self, new_G, num_of_sampled_latents, w_batch, use_wandb=False):
+        loss = 0.0
+
+        z_samples = np.random.randn(
+            num_of_sampled_latents, self.original_G.z_dim)
+        w_samples = self.original_G.mapping(torch.from_numpy(z_samples).to(global_config.device), None,
+                                            truncation_psi=0.5)
+        territory_indicator_ws = [self.get_morphed_w_code(
+            w_code.unsqueeze(0), w_batch) for w_code in w_samples]
+
+        for w_code in territory_indicator_ws:
+            new_img = new_G.synthesis(
+                w_code, noise_mode='none', force_fp32=True)
+            with torch.no_grad():
+                old_img = self.original_G.synthesis(
+                    w_code, noise_mode='none', force_fp32=True)
+
+            if hyperparameters.regulizer_l2_lambda > 0:
+                l2_loss_val = l2_loss.l2_loss(old_img, new_img)
+                if use_wandb:
+                    wandb.log({f'space_regulizer_l2_loss_val': l2_loss_val.detach().cpu()},
+                              step=global_config.training_step)
+                loss += l2_loss_val * hyperparameters.regulizer_l2_lambda
+
+            if hyperparameters.regulizer_lpips_lambda > 0:
+                loss_lpips = self.lpips_loss(old_img, new_img)
+                loss_lpips = torch.mean(torch.squeeze(loss_lpips))
+                if use_wandb:
+                    wandb.log({f'space_regulizer_lpips_loss_val': loss_lpips.detach().cpu()},
+                              step=global_config.training_step)
+                loss += loss_lpips * hyperparameters.regulizer_lpips_lambda
+
+        return loss / len(territory_indicator_ws)
+
+    def space_regulizer_loss(self, new_G, w_batch, use_wandb):
+        ret_val = self.ball_holder_loss_lazy(
+            new_G, hyperparameters.latent_ball_num_of_samples, w_batch, use_wandb)
+        return ret_val

PTI/dnnlib/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from .util import EasyDict, make_cache_dir_path

PTI/dnnlib/util.py

@@ -0,0 +1,477 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Miscellaneous utility classes and functions."""
+
+import ctypes
+import fnmatch
+import importlib
+import inspect
+import numpy as np
+import os
+import shutil
+import sys
+import types
+import io
+import pickle
+import re
+import requests
+import html
+import hashlib
+import glob
+import tempfile
+import urllib
+import urllib.request
+import uuid
+
+from distutils.util import strtobool
+from typing import Any, List, Tuple, Union
+
+
+# Util classes
+# ------------------------------------------------------------------------------------------
+
+
+class EasyDict(dict):
+    """Convenience class that behaves like a dict but allows access with the attribute syntax."""
+
+    def __getattr__(self, name: str) -> Any:
+        try:
+            return self[name]
+        except KeyError:
+            raise AttributeError(name)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        self[name] = value
+
+    def __delattr__(self, name: str) -> None:
+        del self[name]
+
+
+class Logger(object):
+    """Redirect stderr to stdout, optionally print stdout to a file, and optionally force flushing on both stdout and the file."""
+
+    def __init__(self, file_name: str = None, file_mode: str = "w", should_flush: bool = True):
+        self.file = None
+
+        if file_name is not None:
+            self.file = open(file_name, file_mode)
+
+        self.should_flush = should_flush
+        self.stdout = sys.stdout
+        self.stderr = sys.stderr
+
+        sys.stdout = self
+        sys.stderr = self
+
+    def __enter__(self) -> "Logger":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        self.close()
+
+    def write(self, text: Union[str, bytes]) -> None:
+        """Write text to stdout (and a file) and optionally flush."""
+        if isinstance(text, bytes):
+            text = text.decode()
+        if len(text) == 0: # workaround for a bug in VSCode debugger: sys.stdout.write(''); sys.stdout.flush() => crash
+            return
+
+        if self.file is not None:
+            self.file.write(text)
+
+        self.stdout.write(text)
+
+        if self.should_flush:
+            self.flush()
+
+    def flush(self) -> None:
+        """Flush written text to both stdout and a file, if open."""
+        if self.file is not None:
+            self.file.flush()
+
+        self.stdout.flush()
+
+    def close(self) -> None:
+        """Flush, close possible files, and remove stdout/stderr mirroring."""
+        self.flush()
+
+        # if using multiple loggers, prevent closing in wrong order
+        if sys.stdout is self:
+            sys.stdout = self.stdout
+        if sys.stderr is self:
+            sys.stderr = self.stderr
+
+        if self.file is not None:
+            self.file.close()
+            self.file = None
+
+
+# Cache directories
+# ------------------------------------------------------------------------------------------
+
+_dnnlib_cache_dir = None
+
+def set_cache_dir(path: str) -> None:
+    global _dnnlib_cache_dir
+    _dnnlib_cache_dir = path
+
+def make_cache_dir_path(*paths: str) -> str:
+    if _dnnlib_cache_dir is not None:
+        return os.path.join(_dnnlib_cache_dir, *paths)
+    if 'DNNLIB_CACHE_DIR' in os.environ:
+        return os.path.join(os.environ['DNNLIB_CACHE_DIR'], *paths)
+    if 'HOME' in os.environ:
+        return os.path.join(os.environ['HOME'], '.cache', 'dnnlib', *paths)
+    if 'USERPROFILE' in os.environ:
+        return os.path.join(os.environ['USERPROFILE'], '.cache', 'dnnlib', *paths)
+    return os.path.join(tempfile.gettempdir(), '.cache', 'dnnlib', *paths)
+
+# Small util functions
+# ------------------------------------------------------------------------------------------
+
+
+def format_time(seconds: Union[int, float]) -> str:
+    """Convert the seconds to human readable string with days, hours, minutes and seconds."""
+    s = int(np.rint(seconds))
+
+    if s < 60:
+        return "{0}s".format(s)
+    elif s < 60 * 60:
+        return "{0}m {1:02}s".format(s // 60, s % 60)
+    elif s < 24 * 60 * 60:
+        return "{0}h {1:02}m {2:02}s".format(s // (60 * 60), (s // 60) % 60, s % 60)
+    else:
+        return "{0}d {1:02}h {2:02}m".format(s // (24 * 60 * 60), (s // (60 * 60)) % 24, (s // 60) % 60)
+
+
+def ask_yes_no(question: str) -> bool:
+    """Ask the user the question until the user inputs a valid answer."""
+    while True:
+        try:
+            print("{0} [y/n]".format(question))
+            return strtobool(input().lower())
+        except ValueError:
+            pass
+
+
+def tuple_product(t: Tuple) -> Any:
+    """Calculate the product of the tuple elements."""
+    result = 1
+
+    for v in t:
+        result *= v
+
+    return result
+
+
+_str_to_ctype = {
+    "uint8": ctypes.c_ubyte,
+    "uint16": ctypes.c_uint16,
+    "uint32": ctypes.c_uint32,
+    "uint64": ctypes.c_uint64,
+    "int8": ctypes.c_byte,
+    "int16": ctypes.c_int16,
+    "int32": ctypes.c_int32,
+    "int64": ctypes.c_int64,
+    "float32": ctypes.c_float,
+    "float64": ctypes.c_double
+}
+
+
+def get_dtype_and_ctype(type_obj: Any) -> Tuple[np.dtype, Any]:
+    """Given a type name string (or an object having a __name__ attribute), return matching Numpy and ctypes types that have the same size in bytes."""
+    type_str = None
+
+    if isinstance(type_obj, str):
+        type_str = type_obj
+    elif hasattr(type_obj, "__name__"):
+        type_str = type_obj.__name__
+    elif hasattr(type_obj, "name"):
+        type_str = type_obj.name
+    else:
+        raise RuntimeError("Cannot infer type name from input")
+
+    assert type_str in _str_to_ctype.keys()
+
+    my_dtype = np.dtype(type_str)
+    my_ctype = _str_to_ctype[type_str]
+
+    assert my_dtype.itemsize == ctypes.sizeof(my_ctype)
+
+    return my_dtype, my_ctype
+
+
+def is_pickleable(obj: Any) -> bool:
+    try:
+        with io.BytesIO() as stream:
+            pickle.dump(obj, stream)
+        return True
+    except:
+        return False
+
+
+# Functionality to import modules/objects by name, and call functions by name
+# ------------------------------------------------------------------------------------------
+
+def get_module_from_obj_name(obj_name: str) -> Tuple[types.ModuleType, str]:
+    """Searches for the underlying module behind the name to some python object.
+    Returns the module and the object name (original name with module part removed)."""
+
+    # allow convenience shorthands, substitute them by full names
+    obj_name = re.sub("^np.", "numpy.", obj_name)
+    obj_name = re.sub("^tf.", "tensorflow.", obj_name)
+
+    # list alternatives for (module_name, local_obj_name)
+    parts = obj_name.split(".")
+    name_pairs = [(".".join(parts[:i]), ".".join(parts[i:])) for i in range(len(parts), 0, -1)]
+
+    # try each alternative in turn
+    for module_name, local_obj_name in name_pairs:
+        try:
+            module = importlib.import_module(module_name) # may raise ImportError
+            get_obj_from_module(module, local_obj_name) # may raise AttributeError
+            return module, local_obj_name
+        except:
+            pass
+
+    # maybe some of the modules themselves contain errors?
+    for module_name, _local_obj_name in name_pairs:
+        try:
+            importlib.import_module(module_name) # may raise ImportError
+        except ImportError:
+            if not str(sys.exc_info()[1]).startswith("No module named '" + module_name + "'"):
+                raise
+
+    # maybe the requested attribute is missing?
+    for module_name, local_obj_name in name_pairs:
+        try:
+            module = importlib.import_module(module_name) # may raise ImportError
+            get_obj_from_module(module, local_obj_name) # may raise AttributeError
+        except ImportError:
+            pass
+
+    # we are out of luck, but we have no idea why
+    raise ImportError(obj_name)
+
+
+def get_obj_from_module(module: types.ModuleType, obj_name: str) -> Any:
+    """Traverses the object name and returns the last (rightmost) python object."""
+    if obj_name == '':
+        return module
+    obj = module
+    for part in obj_name.split("."):
+        obj = getattr(obj, part)
+    return obj
+
+
+def get_obj_by_name(name: str) -> Any:
+    """Finds the python object with the given name."""
+    module, obj_name = get_module_from_obj_name(name)
+    return get_obj_from_module(module, obj_name)
+
+
+def call_func_by_name(*args, func_name: str = None, **kwargs) -> Any:
+    """Finds the python object with the given name and calls it as a function."""
+    assert func_name is not None
+    func_obj = get_obj_by_name(func_name)
+    assert callable(func_obj)
+    return func_obj(*args, **kwargs)
+
+
+def construct_class_by_name(*args, class_name: str = None, **kwargs) -> Any:
+    """Finds the python class with the given name and constructs it with the given arguments."""
+    return call_func_by_name(*args, func_name=class_name, **kwargs)
+
+
+def get_module_dir_by_obj_name(obj_name: str) -> str:
+    """Get the directory path of the module containing the given object name."""
+    module, _ = get_module_from_obj_name(obj_name)
+    return os.path.dirname(inspect.getfile(module))
+
+
+def is_top_level_function(obj: Any) -> bool:
+    """Determine whether the given object is a top-level function, i.e., defined at module scope using 'def'."""
+    return callable(obj) and obj.__name__ in sys.modules[obj.__module__].__dict__
+
+
+def get_top_level_function_name(obj: Any) -> str:
+    """Return the fully-qualified name of a top-level function."""
+    assert is_top_level_function(obj)
+    module = obj.__module__
+    if module == '__main__':
+        module = os.path.splitext(os.path.basename(sys.modules[module].__file__))[0]
+    return module + "." + obj.__name__
+
+
+# File system helpers
+# ------------------------------------------------------------------------------------------
+
+def list_dir_recursively_with_ignore(dir_path: str, ignores: List[str] = None, add_base_to_relative: bool = False) -> List[Tuple[str, str]]:
+    """List all files recursively in a given directory while ignoring given file and directory names.
+    Returns list of tuples containing both absolute and relative paths."""
+    assert os.path.isdir(dir_path)
+    base_name = os.path.basename(os.path.normpath(dir_path))
+
+    if ignores is None:
+        ignores = []
+
+    result = []
+
+    for root, dirs, files in os.walk(dir_path, topdown=True):
+        for ignore_ in ignores:
+            dirs_to_remove = [d for d in dirs if fnmatch.fnmatch(d, ignore_)]
+
+            # dirs need to be edited in-place
+            for d in dirs_to_remove:
+                dirs.remove(d)
+
+            files = [f for f in files if not fnmatch.fnmatch(f, ignore_)]
+
+        absolute_paths = [os.path.join(root, f) for f in files]
+        relative_paths = [os.path.relpath(p, dir_path) for p in absolute_paths]
+
+        if add_base_to_relative:
+            relative_paths = [os.path.join(base_name, p) for p in relative_paths]
+
+        assert len(absolute_paths) == len(relative_paths)
+        result += zip(absolute_paths, relative_paths)
+
+    return result
+
+
+def copy_files_and_create_dirs(files: List[Tuple[str, str]]) -> None:
+    """Takes in a list of tuples of (src, dst) paths and copies files.
+    Will create all necessary directories."""
+    for file in files:
+        target_dir_name = os.path.dirname(file[1])
+
+        # will create all intermediate-level directories
+        if not os.path.exists(target_dir_name):
+            os.makedirs(target_dir_name)
+
+        shutil.copyfile(file[0], file[1])
+
+
+# URL helpers
+# ------------------------------------------------------------------------------------------
+
+def is_url(obj: Any, allow_file_urls: bool = False) -> bool:
+    """Determine whether the given object is a valid URL string."""
+    if not isinstance(obj, str) or not "://" in obj:
+        return False
+    if allow_file_urls and obj.startswith('file://'):
+        return True
+    try:
+        res = requests.compat.urlparse(obj)
+        if not res.scheme or not res.netloc or not "." in res.netloc:
+            return False
+        res = requests.compat.urlparse(requests.compat.urljoin(obj, "/"))
+        if not res.scheme or not res.netloc or not "." in res.netloc:
+            return False
+    except:
+        return False
+    return True
+
+
+def open_url(url: str, cache_dir: str = None, num_attempts: int = 10, verbose: bool = True, return_filename: bool = False, cache: bool = True) -> Any:
+    """Download the given URL and return a binary-mode file object to access the data."""
+    assert num_attempts >= 1
+    assert not (return_filename and (not cache))
+
+    # Doesn't look like an URL scheme so interpret it as a local filename.
+    if not re.match('^[a-z]+://', url):
+        return url if return_filename else open(url, "rb")
+
+    # Handle file URLs.  This code handles unusual file:// patterns that
+    # arise on Windows:
+    #
+    # file:///c:/foo.txt
+    #
+    # which would translate to a local '/c:/foo.txt' filename that's
+    # invalid.  Drop the forward slash for such pathnames.
+    #
+    # If you touch this code path, you should test it on both Linux and
+    # Windows.
+    #
+    # Some internet resources suggest using urllib.request.url2pathname() but
+    # but that converts forward slashes to backslashes and this causes
+    # its own set of problems.
+    if url.startswith('file://'):
+        filename = urllib.parse.urlparse(url).path
+        if re.match(r'^/[a-zA-Z]:', filename):
+            filename = filename[1:]
+        return filename if return_filename else open(filename, "rb")
+
+    assert is_url(url)
+
+    # Lookup from cache.
+    if cache_dir is None:
+        cache_dir = make_cache_dir_path('downloads')
+
+    url_md5 = hashlib.md5(url.encode("utf-8")).hexdigest()
+    if cache:
+        cache_files = glob.glob(os.path.join(cache_dir, url_md5 + "_*"))
+        if len(cache_files) == 1:
+            filename = cache_files[0]
+            return filename if return_filename else open(filename, "rb")
+
+    # Download.
+    url_name = None
+    url_data = None
+    with requests.Session() as session:
+        if verbose:
+            print("Downloading %s ..." % url, end="", flush=True)
+        for attempts_left in reversed(range(num_attempts)):
+            try:
+                with session.get(url) as res:
+                    res.raise_for_status()
+                    if len(res.content) == 0:
+                        raise IOError("No data received")
+
+                    if len(res.content) < 8192:
+                        content_str = res.content.decode("utf-8")
+                        if "download_warning" in res.headers.get("Set-Cookie", ""):
+                            links = [html.unescape(link) for link in content_str.split('"') if "export=download" in link]
+                            if len(links) == 1:
+                                url = requests.compat.urljoin(url, links[0])
+                                raise IOError("Google Drive virus checker nag")
+                        if "Google Drive - Quota exceeded" in content_str:
+                            raise IOError("Google Drive download quota exceeded -- please try again later")
+
+                    match = re.search(r'filename="([^"]*)"', res.headers.get("Content-Disposition", ""))
+                    url_name = match[1] if match else url
+                    url_data = res.content
+                    if verbose:
+                        print(" done")
+                    break
+            except KeyboardInterrupt:
+                raise
+            except:
+                if not attempts_left:
+                    if verbose:
+                        print(" failed")
+                    raise
+                if verbose:
+                    print(".", end="", flush=True)
+
+    # Save to cache.
+    if cache:
+        safe_name = re.sub(r"[^0-9a-zA-Z-._]", "_", url_name)
+        cache_file = os.path.join(cache_dir, url_md5 + "_" + safe_name)
+        temp_file = os.path.join(cache_dir, "tmp_" + uuid.uuid4().hex + "_" + url_md5 + "_" + safe_name)
+        os.makedirs(cache_dir, exist_ok=True)
+        with open(temp_file, "wb") as f:
+            f.write(url_data)
+        os.replace(temp_file, cache_file) # atomic
+        if return_filename:
+            return cache_file
+
+    # Return data as file object.
+    assert not return_filename
+    return io.BytesIO(url_data)

PTI/editings/ganspace.py

@@ -0,0 +1,21 @@
+import torch
+
+
+def edit(latents, pca, edit_directions):
+    edit_latents = []
+    for latent in latents:
+        for pca_idx, start, end, strength in edit_directions:
+            delta = get_delta(pca, latent, pca_idx, strength)
+            delta_padded = torch.zeros(latent.shape).to('cuda')
+            delta_padded[start:end] += delta.repeat(end - start, 1)
+            edit_latents.append(latent + delta_padded)
+    return torch.stack(edit_latents)
+
+
+def get_delta(pca, latent, idx, strength):
+    w_centered = latent - pca['mean'].to('cuda')
+    lat_comp = pca['comp'].to('cuda')
+    lat_std = pca['std'].to('cuda')
+    w_coord = torch.sum(w_centered[0].reshape(-1)*lat_comp[idx].reshape(-1)) / lat_std[idx]
+    delta = (strength - w_coord)*lat_comp[idx]*lat_std[idx]
+    return delta

PTI/editings/latent_editor.py

@@ -0,0 +1,23 @@
+import torch
+
+from configs import paths_config
+from editings import ganspace
+from utils.data_utils import tensor2im
+
+
+class LatentEditor(object):
+
+    def apply_ganspace(self, latent, ganspace_pca, edit_directions):
+        edit_latents = ganspace.edit(latent, ganspace_pca, edit_directions)
+        return edit_latents
+
+    def apply_interfacegan(self, latent, direction, factor=1, factor_range=None):
+        edit_latents = []
+        if factor_range is not None:  # Apply a range of editing factors. for example, (-5, 5)
+            for f in range(*factor_range):
+                edit_latent = latent + f * direction
+                edit_latents.append(edit_latent)
+            edit_latents = torch.cat(edit_latents)
+        else:
+            edit_latents = latent + factor * direction
+        return edit_latents

PTI/evaluation/experiment_setting_creator.py

@@ -0,0 +1,43 @@
+import glob
+import os
+from configs import global_config, paths_config, hyperparameters
+from scripts.latent_creators.sg2_plus_latent_creator import SG2PlusLatentCreator
+from scripts.latent_creators.e4e_latent_creator import E4ELatentCreator
+from scripts.run_pti import run_PTI
+import pickle
+import torch
+from utils.models_utils import toogle_grad, load_old_G
+
+
+class ExperimentRunner:
+
+    def __init__(self, run_id=''):
+        self.images_paths = glob.glob(f'{paths_config.input_data_path}/*')
+        self.target_paths = glob.glob(f'{paths_config.input_data_path}/*')
+        self.run_id = run_id
+        self.sampled_ws = None
+
+        self.old_G = load_old_G()
+
+        toogle_grad(self.old_G, False)
+
+    def run_experiment(self, run_pt, create_other_latents, use_multi_id_training, use_wandb=False):
+        if run_pt:
+            self.run_id = run_PTI(self.run_id, use_wandb=use_wandb, use_multi_id_training=use_multi_id_training)
+        if create_other_latents:
+            sg2_plus_latent_creator = SG2PlusLatentCreator(use_wandb=use_wandb)
+            sg2_plus_latent_creator.create_latents()
+            e4e_latent_creator = E4ELatentCreator(use_wandb=use_wandb)
+            e4e_latent_creator.create_latents()
+
+        torch.cuda.empty_cache()
+
+        return self.run_id
+
+
+if __name__ == '__main__':
+    os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
+    os.environ['CUDA_VISIBLE_DEVICES'] = global_config.cuda_visible_devices
+
+    runner = ExperimentRunner()
+    runner.run_experiment(True, False, False)

PTI/evaluation/qualitative_edit_comparison.py

@@ -0,0 +1,156 @@
+import os
+from random import choice
+from string import ascii_uppercase
+from PIL import Image
+from tqdm import tqdm
+from scripts.latent_editor_wrapper import LatentEditorWrapper
+from evaluation.experiment_setting_creator import ExperimentRunner
+import torch
+from configs import paths_config, hyperparameters, evaluation_config
+from utils.log_utils import save_concat_image, save_single_image
+from utils.models_utils import load_tuned_G
+
+
+class EditComparison:
+
+    def __init__(self, save_single_images, save_concatenated_images, run_id):
+
+        self.run_id = run_id
+        self.experiment_creator = ExperimentRunner(run_id)
+        self.save_single_images = save_single_images
+        self.save_concatenated_images = save_concatenated_images
+        self.latent_editor = LatentEditorWrapper()
+
+    def save_reconstruction_images(self, image_latents, new_inv_image_latent, new_G, target_image):
+        if self.save_concatenated_images:
+            save_concat_image(self.concat_base_dir, image_latents, new_inv_image_latent, new_G,
+                              self.experiment_creator.old_G,
+                              'rec',
+                              target_image)
+
+        if self.save_single_images:
+            save_single_image(self.single_base_dir, new_inv_image_latent, new_G, 'rec')
+            target_image.save(f'{self.single_base_dir}/Original.jpg')
+
+    def create_output_dirs(self, full_image_name):
+        output_base_dir_path = f'{paths_config.experiments_output_dir}/{paths_config.input_data_id}/{self.run_id}/{full_image_name}'
+        os.makedirs(output_base_dir_path, exist_ok=True)
+
+        self.concat_base_dir = f'{output_base_dir_path}/concat_images'
+        self.single_base_dir = f'{output_base_dir_path}/single_images'
+
+        os.makedirs(self.concat_base_dir, exist_ok=True)
+        os.makedirs(self.single_base_dir, exist_ok=True)
+
+    def get_image_latent_codes(self, image_name):
+        image_latents = []
+        for method in evaluation_config.evaluated_methods:
+            if method == 'SG2':
+                image_latents.append(torch.load(
+                    f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}/'
+                    f'{paths_config.pti_results_keyword}/{image_name}/0.pt'))
+            else:
+                image_latents.append(torch.load(
+                    f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}/{method}/{image_name}/0.pt'))
+        new_inv_image_latent = torch.load(
+            f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}/{paths_config.pti_results_keyword}/{image_name}/0.pt')
+
+        return image_latents, new_inv_image_latent
+
+    def save_interfacegan_edits(self, image_latents, new_inv_image_latent, interfacegan_factors, new_G, target_image):
+        new_w_inv_edits = self.latent_editor.get_single_interface_gan_edits(new_inv_image_latent,
+                                                                            interfacegan_factors)
+
+        inv_edits = []
+        for latent in image_latents:
+            inv_edits.append(self.latent_editor.get_single_interface_gan_edits(latent, interfacegan_factors))
+
+        for direction, edits in new_w_inv_edits.items():
+            for factor, edit_tensor in edits.items():
+                if self.save_concatenated_images:
+                    save_concat_image(self.concat_base_dir, [edits[direction][factor] for edits in inv_edits],
+                                      new_w_inv_edits[direction][factor],
+                                      new_G,
+                                      self.experiment_creator.old_G,
+                                      f'{direction}_{factor}', target_image)
+                if self.save_single_images:
+                    save_single_image(self.single_base_dir, new_w_inv_edits[direction][factor], new_G,
+                                      f'{direction}_{factor}')
+
+    def save_ganspace_edits(self, image_latents, new_inv_image_latent, factors, new_G, target_image):
+        new_w_inv_edits = self.latent_editor.get_single_ganspace_edits(new_inv_image_latent, factors)
+        inv_edits = []
+        for latent in image_latents:
+            inv_edits.append(self.latent_editor.get_single_ganspace_edits(latent, factors))
+
+        for idx in range(len(new_w_inv_edits)):
+            if self.save_concatenated_images:
+                save_concat_image(self.concat_base_dir, [edit[idx] for edit in inv_edits], new_w_inv_edits[idx],
+                                  new_G,
+                                  self.experiment_creator.old_G,
+                                  f'ganspace_{idx}', target_image)
+            if self.save_single_images:
+                save_single_image(self.single_base_dir, new_w_inv_edits[idx], new_G,
+                                  f'ganspace_{idx}')
+
+    def run_experiment(self, run_pt, create_other_latents, use_multi_id_training, use_wandb=False):
+        images_counter = 0
+        new_G = None
+        interfacegan_factors = [val / 2 for val in range(-6, 7) if val != 0]
+        ganspace_factors = range(-20, 25, 5)
+        self.experiment_creator.run_experiment(run_pt, create_other_latents, use_multi_id_training, use_wandb)
+
+        if use_multi_id_training:
+            new_G = load_tuned_G(self.run_id, paths_config.multi_id_model_type)
+
+        for idx, image_path in tqdm(enumerate(self.experiment_creator.images_paths),
+                                    total=len(self.experiment_creator.images_paths)):
+
+            if images_counter >= hyperparameters.max_images_to_invert:
+                break
+
+            image_name = image_path.split('.')[0].split('/')[-1]
+            target_image = Image.open(self.experiment_creator.target_paths[idx])
+
+            if not use_multi_id_training:
+                new_G = load_tuned_G(self.run_id, image_name)
+
+            image_latents, new_inv_image_latent = self.get_image_latent_codes(image_name)
+
+            self.create_output_dirs(image_name)
+
+            self.save_reconstruction_images(image_latents, new_inv_image_latent, new_G, target_image)
+
+            self.save_interfacegan_edits(image_latents, new_inv_image_latent, interfacegan_factors, new_G, target_image)
+
+            self.save_ganspace_edits(image_latents, new_inv_image_latent, ganspace_factors, new_G, target_image)
+
+            target_image.close()
+            torch.cuda.empty_cache()
+            images_counter += 1
+
+
+def run_pti_and_full_edit(iid):
+    evaluation_config.evaluated_methods = ['SG2Plus', 'e4e', 'SG2']
+    edit_figure_creator = EditComparison(save_single_images=True, save_concatenated_images=True,
+                                         run_id=f'{paths_config.input_data_id}_pti_full_edit_{iid}')
+    edit_figure_creator.run_experiment(True, True, use_multi_id_training=False, use_wandb=False)
+
+
+def pti_no_comparison(iid):
+    evaluation_config.evaluated_methods = []
+    edit_figure_creator = EditComparison(save_single_images=True, save_concatenated_images=True,
+                                         run_id=f'{paths_config.input_data_id}_pti_no_comparison_{iid}')
+    edit_figure_creator.run_experiment(True, False, use_multi_id_training=False, use_wandb=False)
+
+
+def edits_for_existed_experiment(run_id):
+    evaluation_config.evaluated_methods = ['SG2Plus', 'e4e', 'SG2']
+    edit_figure_creator = EditComparison(save_single_images=True, save_concatenated_images=True,
+                                         run_id=run_id)
+    edit_figure_creator.run_experiment(False, True, use_multi_id_training=False, use_wandb=False)
+
+
+if __name__ == '__main__':
+    iid = ''.join(choice(ascii_uppercase) for i in range(7))
+    pti_no_comparison(iid)

PTI/models/StyleCLIP/criteria/clip_loss.py

@@ -0,0 +1,17 @@
+
+import torch
+import clip
+
+
+class CLIPLoss(torch.nn.Module):
+
+    def __init__(self, opts):
+        super(CLIPLoss, self).__init__()
+        self.model, self.preprocess = clip.load("ViT-B/32", device="cuda")
+        self.upsample = torch.nn.Upsample(scale_factor=7)
+        self.avg_pool = torch.nn.AvgPool2d(kernel_size=opts.stylegan_size // 32)
+
+    def forward(self, image, text):
+        image = self.avg_pool(self.upsample(image))
+        similarity = 1 - self.model(image, text)[0] / 100
+        return similarity

PTI/models/StyleCLIP/criteria/id_loss.py

@@ -0,0 +1,39 @@
+import torch
+from torch import nn
+
+from models.facial_recognition.model_irse import Backbone
+
+
+class IDLoss(nn.Module):
+    def __init__(self, opts):
+        super(IDLoss, self).__init__()
+        print('Loading ResNet ArcFace')
+        self.facenet = Backbone(input_size=112, num_layers=50, drop_ratio=0.6, mode='ir_se')
+        self.facenet.load_state_dict(torch.load(opts.ir_se50_weights))
+        self.pool = torch.nn.AdaptiveAvgPool2d((256, 256))
+        self.face_pool = torch.nn.AdaptiveAvgPool2d((112, 112))
+        self.facenet.eval()
+        self.opts = opts
+
+    def extract_feats(self, x):
+        if x.shape[2] != 256:
+            x = self.pool(x)
+        x = x[:, :, 35:223, 32:220]  # Crop interesting region
+        x = self.face_pool(x)
+        x_feats = self.facenet(x)
+        return x_feats
+
+    def forward(self, y_hat, y):
+        n_samples = y.shape[0]
+        y_feats = self.extract_feats(y)  # Otherwise use the feature from there
+        y_hat_feats = self.extract_feats(y_hat)
+        y_feats = y_feats.detach()
+        loss = 0
+        sim_improvement = 0
+        count = 0
+        for i in range(n_samples):
+            diff_target = y_hat_feats[i].dot(y_feats[i])
+            loss += 1 - diff_target
+            count += 1
+
+        return loss / count, sim_improvement / count

PTI/models/StyleCLIP/global_directions/GUI.py

@@ -0,0 +1,103 @@
+
+
+from tkinter import Tk,Frame ,Label,Button,messagebox,Canvas,Text,Scale
+from tkinter import  HORIZONTAL
+
+class View():
+    def __init__(self,master):
+        
+        self.width=600
+        self.height=600
+        
+        
+        self.root=master
+        self.root.geometry("600x600")
+        
+        self.left_frame=Frame(self.root,width=600)
+        self.left_frame.pack_propagate(0)
+        self.left_frame.pack(fill='both', side='left', expand='True')
+        
+        self.retrieval_frame=Frame(self.root,bg='snow3')
+        self.retrieval_frame.pack_propagate(0)
+        self.retrieval_frame.pack(fill='both', side='right', expand='True')
+        
+        self.bg_frame=Frame(self.left_frame,bg='snow3',height=600,width=600)
+        self.bg_frame.pack_propagate(0)
+        self.bg_frame.pack(fill='both', side='top', expand='True')
+        
+        self.command_frame=Frame(self.left_frame,bg='snow3')
+        self.command_frame.pack_propagate(0)
+        self.command_frame.pack(fill='both', side='bottom', expand='True')
+#        self.command_frame.grid(row=1, column=0,padx=0, pady=0)
+        
+        self.bg=Canvas(self.bg_frame,width=self.width,height=self.height, bg='gray')
+        self.bg.place(relx=0.5, rely=0.5, anchor='center')
+        
+        self.mani=Canvas(self.retrieval_frame,width=1024,height=1024, bg='gray') 
+        self.mani.grid(row=0, column=0,padx=0, pady=42)
+        
+        self.SetCommand()
+        
+        
+        
+    
+    def run(self):
+        self.root.mainloop()
+    
+    def helloCallBack(self):
+        category=self.set_category.get()
+        messagebox.showinfo( "Hello Python",category)
+    
+    def SetCommand(self):
+        
+        tmp = Label(self.command_frame, text="neutral", width=10 ,bg='snow3')
+        tmp.grid(row=1, column=0,padx=10, pady=10)
+        
+        tmp = Label(self.command_frame, text="a photo of a", width=10 ,bg='snow3')
+        tmp.grid(row=1, column=1,padx=10, pady=10)
+        
+        self.neutral = Text ( self.command_frame, height=2, width=30)
+        self.neutral.grid(row=1, column=2,padx=10, pady=10)
+        
+        
+        tmp = Label(self.command_frame, text="target", width=10 ,bg='snow3')
+        tmp.grid(row=2, column=0,padx=10, pady=10)
+        
+        tmp = Label(self.command_frame, text="a photo of a", width=10 ,bg='snow3')
+        tmp.grid(row=2, column=1,padx=10, pady=10)
+        
+        self.target = Text ( self.command_frame, height=2, width=30)
+        self.target.grid(row=2, column=2,padx=10, pady=10)
+        
+        tmp = Label(self.command_frame, text="strength", width=10 ,bg='snow3')
+        tmp.grid(row=3, column=0,padx=10, pady=10)
+        
+        self.alpha = Scale(self.command_frame, from_=-15, to=25, orient=HORIZONTAL,bg='snow3', length=250,resolution=0.01)
+        self.alpha.grid(row=3, column=2,padx=10, pady=10)
+        
+        
+        tmp = Label(self.command_frame, text="disentangle", width=10 ,bg='snow3')
+        tmp.grid(row=4, column=0,padx=10, pady=10)
+        
+        self.beta = Scale(self.command_frame, from_=0.08, to=0.4, orient=HORIZONTAL,bg='snow3', length=250,resolution=0.001)
+        self.beta.grid(row=4, column=2,padx=10, pady=10)
+        
+        self.reset = Button(self.command_frame, text='Reset') 
+        self.reset.grid(row=5, column=1,padx=10, pady=10)
+        
+        
+        self.set_init = Button(self.command_frame, text='Accept') 
+        self.set_init.grid(row=5, column=2,padx=10, pady=10)
+
+#%%
+if __name__ == "__main__":
+    master=Tk()
+    self=View(master)
+    self.run()
+    
+    
+    
+    
+    
+    
+    

PTI/models/StyleCLIP/global_directions/GenerateImg.py

@@ -0,0 +1,50 @@
+
+import os
+import numpy as np
+import argparse
+from manipulate import Manipulator
+
+from PIL import Image
+#%%
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Process some integers.')
+    
+    parser.add_argument('--dataset_name',type=str,default='ffhq',
+                    help='name of dataset, for example, ffhq')
+
+    args = parser.parse_args()
+    dataset_name=args.dataset_name
+    
+    if not os.path.isdir('./data/'+dataset_name):
+        os.system('mkdir ./data/'+dataset_name)
+    #%%
+    M=Manipulator(dataset_name=dataset_name)
+    np.set_printoptions(suppress=True)
+    print(M.dataset_name)
+    #%%
+    
+    M.img_index=0
+    M.num_images=50
+    M.alpha=[0]
+    M.step=1
+    lindex,bname=0,0
+    
+    M.manipulate_layers=[lindex]
+    codes,out=M.EditOneC(bname)
+    #%%
+    
+    for i in range(len(out)):
+        img=out[i,0]
+        img=Image.fromarray(img)
+        img.save('./data/'+dataset_name+'/'+str(i)+'.jpg')
+    #%%
+    w=np.load('./npy/'+dataset_name+'/W.npy')
+    
+    tmp=w[:M.num_images]
+    tmp=tmp[:,None,:]
+    tmp=np.tile(tmp,(1,M.Gs.components.synthesis.input_shape[1],1))
+    
+    np.save('./data/'+dataset_name+'/w_plus.npy',tmp)
+    
+    

PTI/models/StyleCLIP/global_directions/GetCode.py

@@ -0,0 +1,232 @@
+
+
+
+import os
+import pickle
+import numpy as np
+from dnnlib import tflib  
+import tensorflow as tf 
+
+import argparse
+
+def LoadModel(dataset_name):
+    # Initialize TensorFlow.
+    tflib.init_tf()
+    model_path='./model/'
+    model_name=dataset_name+'.pkl'
+    
+    tmp=os.path.join(model_path,model_name)
+    with open(tmp, 'rb') as f:
+        _, _, Gs = pickle.load(f)
+    return Gs
+
+def lerp(a,b,t):
+     return a + (b - a) * t
+
+#stylegan-ada
+def SelectName(layer_name,suffix):
+    if suffix==None:
+        tmp1='add:0' in layer_name 
+        tmp2='shape=(?,' in layer_name
+        tmp4='G_synthesis_1' in layer_name
+        tmp= tmp1 and tmp2 and tmp4  
+    else:
+        tmp1=('/Conv0_up'+suffix) in layer_name 
+        tmp2=('/Conv1'+suffix) in layer_name 
+        tmp3=('4x4/Conv'+suffix) in layer_name 
+        tmp4='G_synthesis_1' in layer_name
+        tmp5=('/ToRGB'+suffix) in layer_name
+        tmp= (tmp1 or tmp2 or tmp3 or tmp5) and tmp4 
+    return tmp
+
+
+def GetSNames(suffix):
+    #get style tensor name 
+    with tf.Session() as sess:
+        op = sess.graph.get_operations()
+    layers=[m.values() for m in op]
+    
+    
+    select_layers=[]
+    for layer in layers:
+        layer_name=str(layer)
+        if SelectName(layer_name,suffix):
+            select_layers.append(layer[0])
+    return select_layers
+
+def SelectName2(layer_name):
+    tmp1='mod_bias' in layer_name 
+    tmp2='mod_weight' in layer_name
+    tmp3='ToRGB' in layer_name 
+    
+    tmp= (tmp1 or tmp2) and (not tmp3) 
+    return tmp
+
+def GetKName(Gs):
+    
+    layers=[var for name, var in Gs.components.synthesis.vars.items()]
+    
+    select_layers=[]
+    for layer in layers:
+        layer_name=str(layer)
+        if SelectName2(layer_name):
+            select_layers.append(layer)
+    return select_layers
+
+def GetCode(Gs,random_state,num_img,num_once,dataset_name):
+    rnd = np.random.RandomState(random_state)  #5
+    
+    truncation_psi=0.7
+    truncation_cutoff=8
+    
+    dlatent_avg=Gs.get_var('dlatent_avg')
+    
+    dlatents=np.zeros((num_img,512),dtype='float32')
+    for i in range(int(num_img/num_once)):
+        src_latents =  rnd.randn(num_once, Gs.input_shape[1])
+        src_dlatents = Gs.components.mapping.run(src_latents, None) # [seed, layer, component]
+        
+        # Apply truncation trick.
+        if truncation_psi is not None and truncation_cutoff is not None:
+                layer_idx = np.arange(src_dlatents.shape[1])[np.newaxis, :, np.newaxis]
+                ones = np.ones(layer_idx.shape, dtype=np.float32)
+                coefs = np.where(layer_idx < truncation_cutoff, truncation_psi * ones, ones)
+                src_dlatents_np=lerp(dlatent_avg, src_dlatents, coefs)
+                src_dlatents=src_dlatents_np[:,0,:].astype('float32')
+                dlatents[(i*num_once):((i+1)*num_once),:]=src_dlatents
+    print('get all z and w')
+    
+    tmp='./npy/'+dataset_name+'/W'
+    np.save(tmp,dlatents)
+
+    
+def GetImg(Gs,num_img,num_once,dataset_name,save_name='images'):
+    print('Generate Image')
+    tmp='./npy/'+dataset_name+'/W.npy'
+    dlatents=np.load(tmp) 
+    fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
+    
+    all_images=[]
+    for i in range(int(num_img/num_once)):
+        print(i)
+        images=[]
+        for k in range(num_once):
+            tmp=dlatents[i*num_once+k]
+            tmp=tmp[None,None,:]
+            tmp=np.tile(tmp,(1,Gs.components.synthesis.input_shape[1],1))
+            image2= Gs.components.synthesis.run(tmp, randomize_noise=False, output_transform=fmt)
+            images.append(image2)
+            
+        images=np.concatenate(images)
+        
+        all_images.append(images)
+        
+    all_images=np.concatenate(all_images)
+    
+    tmp='./npy/'+dataset_name+'/'+save_name
+    np.save(tmp,all_images)
+
+def GetS(dataset_name,num_img):
+    print('Generate S')
+    tmp='./npy/'+dataset_name+'/W.npy'
+    dlatents=np.load(tmp)[:num_img]
+    
+    with tf.Session() as sess:
+        init = tf.global_variables_initializer()
+        sess.run(init)
+        
+        Gs=LoadModel(dataset_name)
+        Gs.print_layers()  #for ada
+        select_layers1=GetSNames(suffix=None)  #None,'/mul_1:0','/mod_weight/read:0','/MatMul:0'
+        dlatents=dlatents[:,None,:]
+        dlatents=np.tile(dlatents,(1,Gs.components.synthesis.input_shape[1],1))
+        
+        all_s = sess.run(
+            select_layers1,
+            feed_dict={'G_synthesis_1/dlatents_in:0': dlatents})
+    
+    layer_names=[layer.name for layer in select_layers1]
+    save_tmp=[layer_names,all_s]
+    return save_tmp
+
+    
+
+
+def convert_images_to_uint8(images, drange=[-1,1], nchw_to_nhwc=False):
+    """Convert a minibatch of images from float32 to uint8 with configurable dynamic range.
+    Can be used as an output transformation for Network.run().
+    """
+    if nchw_to_nhwc:
+        images = np.transpose(images, [0, 2, 3, 1])
+    
+    scale = 255 / (drange[1] - drange[0])
+    images = images * scale + (0.5 - drange[0] * scale)
+    
+    np.clip(images, 0, 255, out=images)
+    images=images.astype('uint8')
+    return images
+
+
+def GetCodeMS(dlatents):
+        m=[]
+        std=[]
+        for i in range(len(dlatents)):
+            tmp= dlatents[i] 
+            tmp_mean=tmp.mean(axis=0)
+            tmp_std=tmp.std(axis=0)
+            m.append(tmp_mean)
+            std.append(tmp_std)
+        return m,std
+
+
+
+#%%
+if __name__ == "__main__":
+    
+    
+    parser = argparse.ArgumentParser(description='Process some integers.')
+    
+    parser.add_argument('--dataset_name',type=str,default='ffhq',
+                    help='name of dataset, for example, ffhq')
+    parser.add_argument('--code_type',choices=['w','s','s_mean_std'],default='w')
+    
+    args = parser.parse_args()
+    random_state=5
+    num_img=100_000 
+    num_once=1_000
+    dataset_name=args.dataset_name
+    
+    if not os.path.isfile('./model/'+dataset_name+'.pkl'):
+        url='https://nvlabs-fi-cdn.nvidia.com/stylegan2/networks/'
+        name='stylegan2-'+dataset_name+'-config-f.pkl'
+        os.system('wget ' +url+name + '  -P  ./model/')
+        os.system('mv ./model/'+name+' ./model/'+dataset_name+'.pkl')
+    
+    if not os.path.isdir('./npy/'+dataset_name):
+        os.system('mkdir ./npy/'+dataset_name)
+    
+    if args.code_type=='w':
+        Gs=LoadModel(dataset_name=dataset_name)
+        GetCode(Gs,random_state,num_img,num_once,dataset_name)
+#        GetImg(Gs,num_img=num_img,num_once=num_once,dataset_name=dataset_name,save_name='images_100K') #no need 
+    elif args.code_type=='s':
+        save_name='S'
+        save_tmp=GetS(dataset_name,num_img=2_000)
+        tmp='./npy/'+dataset_name+'/'+save_name
+        with open(tmp, "wb") as fp:
+            pickle.dump(save_tmp, fp)
+        
+    elif args.code_type=='s_mean_std':
+        save_tmp=GetS(dataset_name,num_img=num_img)
+        dlatents=save_tmp[1]
+        m,std=GetCodeMS(dlatents)
+        save_tmp=[m,std]
+        save_name='S_mean_std'
+        tmp='./npy/'+dataset_name+'/'+save_name
+        with open(tmp, "wb") as fp:
+            pickle.dump(save_tmp, fp)
+    
+    
+    
+    
+    

PTI/models/StyleCLIP/global_directions/GetGUIData.py

@@ -0,0 +1,67 @@
+
+import os
+import numpy as np
+import argparse
+from manipulate import Manipulator
+import torch
+from PIL import Image
+#%%
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Process some integers.')
+    
+    parser.add_argument('--dataset_name',type=str,default='ffhq',
+                    help='name of dataset, for example, ffhq')
+
+    parser.add_argument('--real', action='store_true')
+
+    args = parser.parse_args()
+    dataset_name=args.dataset_name
+    
+    if not os.path.isdir('./data/'+dataset_name):
+        os.system('mkdir ./data/'+dataset_name)
+    #%%
+    M=Manipulator(dataset_name=dataset_name)
+    np.set_printoptions(suppress=True)
+    print(M.dataset_name)
+    #%%
+    #remove all .jpg
+    names=os.listdir('./data/'+dataset_name+'/')
+    for name in names:
+        if '.jpg' in name:
+            os.system('rm ./data/'+dataset_name+'/'+name)
+    
+    
+    #%%
+    if args.real:
+        latents=torch.load('./data/'+dataset_name+'/latents.pt')
+        w_plus=latents.cpu().detach().numpy()
+    else:
+        w=np.load('./npy/'+dataset_name+'/W.npy')
+        tmp=w[:50] #only use 50 images
+        tmp=tmp[:,None,:]
+        w_plus=np.tile(tmp,(1,M.Gs.components.synthesis.input_shape[1],1))
+    np.save('./data/'+dataset_name+'/w_plus.npy',w_plus)
+    
+    #%%
+    tmp=M.W2S(w_plus)
+    M.dlatents=tmp
+    
+    M.img_index=0
+    M.num_images=len(w_plus)
+    M.alpha=[0]
+    M.step=1
+    lindex,bname=0,0
+    
+    M.manipulate_layers=[lindex]
+    codes,out=M.EditOneC(bname)
+    #%%
+    
+    for i in range(len(out)):
+        img=out[i,0]
+        img=Image.fromarray(img)
+        img.save('./data/'+dataset_name+'/'+str(i)+'.jpg')
+    #%%
+
+    
+    

PTI/models/StyleCLIP/global_directions/Inference.py

@@ -0,0 +1,106 @@
+
+
+from manipulate import Manipulator
+import tensorflow as tf
+import numpy as np 
+import torch
+import clip
+from MapTS import GetBoundary,GetDt
+
+class StyleCLIP():
+    
+    def __init__(self,dataset_name='ffhq'):
+        print('load clip')
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model, preprocess = clip.load("ViT-B/32", device=device)
+        self.LoadData(dataset_name)
+        
+    def LoadData(self, dataset_name):
+        tf.keras.backend.clear_session()
+        M=Manipulator(dataset_name=dataset_name)
+        np.set_printoptions(suppress=True)
+        fs3=np.load('./npy/'+dataset_name+'/fs3.npy')
+        
+        self.M=M
+        self.fs3=fs3
+        
+        w_plus=np.load('./data/'+dataset_name+'/w_plus.npy')
+        self.M.dlatents=M.W2S(w_plus)
+        
+        if dataset_name=='ffhq':
+            self.c_threshold=20
+        else:
+            self.c_threshold=100
+        self.SetInitP()
+    
+    def SetInitP(self):
+        self.M.alpha=[3]
+        self.M.num_images=1
+        
+        self.target=''
+        self.neutral=''
+        self.GetDt2()
+        img_index=0
+        self.M.dlatent_tmp=[tmp[img_index:(img_index+1)] for tmp in self.M.dlatents]
+        
+        
+    def GetDt2(self):
+        classnames=[self.target,self.neutral]
+        dt=GetDt(classnames,self.model)
+        
+        self.dt=dt
+        num_cs=[]
+        betas=np.arange(0.1,0.3,0.01)
+        for i in range(len(betas)):
+            boundary_tmp2,num_c=GetBoundary(self.fs3,self.dt,self.M,threshold=betas[i])
+            print(betas[i])
+            num_cs.append(num_c)
+        
+        num_cs=np.array(num_cs)
+        select=num_cs>self.c_threshold
+        
+        if sum(select)==0:
+            self.beta=0.1
+        else:
+            self.beta=betas[select][-1]
+        
+    
+    def GetCode(self):
+        boundary_tmp2,num_c=GetBoundary(self.fs3,self.dt,self.M,threshold=self.beta)
+        codes=self.M.MSCode(self.M.dlatent_tmp,boundary_tmp2)
+        return codes
+    
+    def GetImg(self):
+        
+        codes=self.GetCode()
+        out=self.M.GenerateImg(codes)
+        img=out[0,0]
+        return img
+    
+
+
+
+#%%
+if __name__ == "__main__":
+    style_clip=StyleCLIP()
+    self=style_clip
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    

PTI/models/StyleCLIP/global_directions/MapTS.py

@@ -0,0 +1,394 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Thu Feb  4 17:36:31 2021
+
+@author: wuzongze
+"""
+
+import os
+#os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+#os.environ["CUDA_VISIBLE_DEVICES"] = "1" #(or "1" or "2")
+
+import sys 
+
+#sys.path=['', '/usr/local/tensorflow/avx-avx2-gpu/1.14.0/python3.7/site-packages', '/usr/local/matlab/2018b/lib/python3.7/site-packages', '/cs/labs/danix/wuzongze/pythonV/venv3.7/lib/python37.zip', '/cs/labs/danix/wuzongze/pythonV/venv3.7/lib/python3.7', '/cs/labs/danix/wuzongze/pythonV/venv3.7/lib/python3.7/lib-dynload', '/usr/lib/python3.7', '/cs/labs/danix/wuzongze/pythonV/venv3.7/lib/python3.7/site-packages', '/cs/labs/danix/wuzongze/pythonV/venv3.7/lib/python3.7/site-packages/copkmeans-1.5-py3.7.egg', '/cs/labs/danix/wuzongze/pythonV/venv3.7/lib/python3.7/site-packages/spherecluster-0.1.7-py3.7.egg', '/usr/lib/python3/dist-packages', '/usr/local/lib/python3.7/dist-packages', '/usr/lib/python3/dist-packages/IPython/extensions']
+
+import tensorflow as tf
+
+import numpy as np 
+import torch
+import clip
+from PIL import Image
+import pickle
+import copy
+import matplotlib.pyplot as plt
+
+def GetAlign(out,dt,model,preprocess):
+    imgs=out
+    imgs1=imgs.reshape([-1]+list(imgs.shape[2:]))
+    
+    tmp=[]
+    for i in range(len(imgs1)):
+        
+        img=Image.fromarray(imgs1[i])
+        image = preprocess(img).unsqueeze(0).to(device)
+        tmp.append(image)
+    
+    image=torch.cat(tmp)
+    
+    with torch.no_grad():
+        image_features = model.encode_image(image)
+        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
+    
+    image_features1=image_features.cpu().numpy()
+    
+    image_features1=image_features1.reshape(list(imgs.shape[:2])+[512])
+    
+    fd=image_features1[:,1:,:]-image_features1[:,:-1,:]
+    
+    fd1=fd.reshape([-1,512])
+    fd2=fd1/np.linalg.norm(fd1,axis=1)[:,None]
+    
+    tmp=np.dot(fd2,dt)
+    m=tmp.mean()
+    acc=np.sum(tmp>0)/len(tmp)
+    print(m,acc)
+    return m,acc
+
+
+def SplitS(ds_p,M,if_std):
+    all_ds=[]
+    start=0
+    for i in M.mindexs:
+        tmp=M.dlatents[i].shape[1]
+        end=start+tmp
+        tmp=ds_p[start:end]
+#        tmp=tmp*M.code_std[i]
+        
+        all_ds.append(tmp)
+        start=end
+    
+    all_ds2=[]
+    tmp_index=0
+    for i in range(len(M.s_names)):
+        if (not 'RGB' in M.s_names[i]) and (not len(all_ds[tmp_index])==0):
+            
+#            tmp=np.abs(all_ds[tmp_index]/M.code_std[i])
+#            print(i,tmp.mean())
+#            tmp=np.dot(M.latent_codes[i],all_ds[tmp_index])
+#            print(tmp)
+            if if_std:
+                tmp=all_ds[tmp_index]*M.code_std[i]
+            else:
+                tmp=all_ds[tmp_index]
+            
+            all_ds2.append(tmp)
+            tmp_index+=1
+        else:
+            tmp=np.zeros(len(M.dlatents[i][0]))
+            all_ds2.append(tmp)
+    return all_ds2
+
+
+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 {}.',
+]
+
+
+def zeroshot_classifier(classnames, templates,model):
+    with torch.no_grad():
+        zeroshot_weights = []
+        for classname in classnames:
+            texts = [template.format(classname) for template in templates] #format with class
+            texts = clip.tokenize(texts).cuda() #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).cuda()
+    return zeroshot_weights
+
+
+def GetDt(classnames,model):
+    text_features=zeroshot_classifier(classnames, imagenet_templates,model).t()
+    
+    dt=text_features[0]-text_features[1]
+    dt=dt.cpu().numpy()
+    
+#    t_m1=t_m/np.linalg.norm(t_m)
+#    dt=text_features.cpu().numpy()[0]-t_m1
+    print(np.linalg.norm(dt))
+    dt=dt/np.linalg.norm(dt)
+    return dt
+
+
+def GetBoundary(fs3,dt,M,threshold):
+    tmp=np.dot(fs3,dt)
+    
+    ds_imp=copy.copy(tmp)
+    select=np.abs(tmp)<threshold
+    num_c=np.sum(~select)
+
+
+    ds_imp[select]=0
+    tmp=np.abs(ds_imp).max()
+    ds_imp/=tmp
+    
+    boundary_tmp2=SplitS(ds_imp,M,if_std=True)
+    print('num of channels being manipulated:',num_c)
+    return boundary_tmp2,num_c
+
+def GetFs(file_path):
+    fs=np.load(file_path+'single_channel.npy')
+    tmp=np.linalg.norm(fs,axis=-1)
+    fs1=fs/tmp[:,:,:,None]
+    fs2=fs1[:,:,1,:]-fs1[:,:,0,:]  # 5*sigma - (-5)* sigma
+    fs3=fs2/np.linalg.norm(fs2,axis=-1)[:,:,None]
+    fs3=fs3.mean(axis=1)
+    fs3=fs3/np.linalg.norm(fs3,axis=-1)[:,None]
+    return fs3
+#%%
+
+if __name__ == "__main__":
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model, preprocess = clip.load("ViT-B/32", device=device)
+    #%%
+    sys.path.append('/cs/labs/danix/wuzongze/Gan_Manipulation/play')
+    from example_try import Manipulator4
+    
+    M=Manipulator4(dataset_name='ffhq',code_type='S')
+    np.set_printoptions(suppress=True)
+
+    #%%
+    
+    
+    file_path='/cs/labs/danix/wuzongze/Tansformer_Manipulation/CLIP/results/'+M.dataset_name+'/'
+    fs3=GetFs(file_path)
+    
+
+    
+    #%%
+    '''
+    text_features=zeroshot_classifier2(classnames, imagenet_templates) #.t()
+        
+    tmp=np.linalg.norm(text_features,axis=2)
+    text_features/=tmp[:,:,None]
+    dt=text_features[0]-text_features[1]
+    
+    tmp=np.linalg.norm(dt,axis=1)
+    dt/=tmp[:,None]
+    dt=dt.mean(axis=0)
+    '''
+    
+    #%%
+    '''
+    all_tmp=[]
+    tmp=torch.load('/cs/labs/danix/wuzongze/downloads/harris_latent.pt')
+    tmp=tmp.cpu().detach().numpy() #[:,:14,:]
+    all_tmp.append(tmp)
+    
+    tmp=torch.load('/cs/labs/danix/wuzongze/downloads/ariana_latent.pt')
+    tmp=tmp.cpu().detach().numpy() #[:,:14,:]
+    all_tmp.append(tmp)
+    
+    tmp=torch.load('/cs/labs/danix/wuzongze/downloads/federer.pt')
+    tmp=tmp.cpu().detach().numpy() #[:,:14,:]
+    all_tmp.append(tmp)
+    
+    all_tmp=np.array(all_tmp)[:,0]
+    
+    dlatent_tmp=M.W2S(all_tmp)
+    '''
+    '''
+    tmp=torch.load('/cs/labs/danix/wuzongze/downloads/all_cars.pt')
+    tmp=tmp.cpu().detach().numpy()[:300]
+    dlatent_tmp=M.W2S(tmp)
+    '''
+    '''
+    tmp=torch.load('/cs/labs/danix/wuzongze/downloads/faces.pt')
+    tmp=tmp.cpu().detach().numpy()[:100]
+    dlatent_tmp=M.W2S(tmp)
+    '''
+    #%%
+#    M.viz_size=1024
+    M.img_index=0
+    M.num_images=30
+    dlatent_tmp=[tmp[M.img_index:(M.img_index+M.num_images)] for tmp in M.dlatents]
+    #%%
+    
+    classnames=['face','face with glasses']
+    
+#    classnames=['car','classic car']
+#    classnames=['dog','happy dog']
+#    classnames=['bedroom','modern bedroom']
+    
+#    classnames=['church','church without watermark']
+#    classnames=['natural scene','natural scene without grass']
+    dt=GetDt(classnames,model)
+#    tmp=np.dot(fs3,dt)
+#    
+#    ds_imp=copy.copy(tmp)
+#    select=np.abs(tmp)<0.1
+#    num_c=np.sum(~select)
+#
+#
+#    ds_imp[select]=0
+#    tmp=np.abs(ds_imp).max()
+#    ds_imp/=tmp
+#    
+#    boundary_tmp2=SplitS(ds_imp,M,if_std=True)
+#    print('num of channels being manipulated:',num_c)
+    
+    boundary_tmp2=GetBoundary(fs3,dt,M,threshold=0.13)
+    
+    #%%
+    M.start_distance=-20
+    M.end_distance=20
+    M.step=7
+#    M.num_images=100
+    codes=M.MSCode(dlatent_tmp,boundary_tmp2)
+    out=M.GenerateImg(codes)
+    M.Vis2(str('tmp'),'filter2',out)
+    
+#    full=GetAlign(out,dt,model,preprocess)
+    
+    
+    #%%
+    boundary_tmp3=copy.copy(boundary_tmp2) #primary
+    boundary_tmp4=copy.copy(boundary_tmp2) #condition
+    #%%
+    boundary_tmp2=copy.copy(boundary_tmp3)
+    for i in range(len(boundary_tmp3)):
+        select=boundary_tmp4[i]==0
+        boundary_tmp2[i][~select]=0
+    
+    
+
+    
+    
+    
+    
+    #%%1
+    
+    
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

PTI/models/StyleCLIP/global_directions/PlayInteractively.py

@@ -0,0 +1,197 @@
+
+
+
+from tkinter import Tk 
+from PIL import Image, ImageTk
+from tkinter.filedialog import askopenfilename
+from GUI import View
+from Inference import StyleCLIP
+import argparse
+#%%
+
+
+class PlayInteractively():  #Controller
+    '''
+    followed Model View Controller Design Pattern
+    
+    controller, model, view
+    '''
+    def __init__(self,dataset_name='ffhq'):
+        
+        self.root = Tk()
+        self.view=View(self.root)
+        self.img_ratio=2
+        self.style_clip=StyleCLIP(dataset_name)
+        
+        self.view.neutral.bind("<Return>", self.text_n)
+        self.view.target.bind("<Return>", self.text_t)
+        self.view.alpha.bind('<ButtonRelease-1>', self.ChangeAlpha)
+        self.view.beta.bind('<ButtonRelease-1>', self.ChangeBeta)
+        self.view.set_init.bind('<ButtonPress-1>', self.SetInit) 
+        self.view.reset.bind('<ButtonPress-1>', self.Reset) 
+        self.view.bg.bind('<Double-1>', self.open_img)
+        
+        
+        self.drawn  = None
+        
+        self.view.target.delete(1.0, "end")
+        self.view.target.insert("end", self.style_clip.target)
+#        
+        self.view.neutral.delete(1.0, "end")
+        self.view.neutral.insert("end", self.style_clip.neutral)
+        
+    
+    def Reset(self,event):
+        self.style_clip.GetDt2()
+        self.style_clip.M.alpha=[0]
+        
+        self.view.beta.set(self.style_clip.beta)
+        self.view.alpha.set(0)
+        
+        img=self.style_clip.GetImg()
+        img=Image.fromarray(img)
+        img = ImageTk.PhotoImage(img)
+        self.addImage_m(img)
+        
+    
+    def SetInit(self,event):
+        codes=self.style_clip.GetCode()
+        self.style_clip.M.dlatent_tmp=[tmp[:,0] for tmp in codes]
+        print('set init')
+    
+    def ChangeAlpha(self,event):
+        tmp=self.view.alpha.get()
+        self.style_clip.M.alpha=[float(tmp)]
+        
+        img=self.style_clip.GetImg()
+        print('manipulate one')
+        img=Image.fromarray(img)
+        img = ImageTk.PhotoImage(img)
+        self.addImage_m(img)
+        
+    def ChangeBeta(self,event):
+        tmp=self.view.beta.get()
+        self.style_clip.beta=float(tmp)
+        
+        img=self.style_clip.GetImg()
+        print('manipulate one')
+        img=Image.fromarray(img)
+        img = ImageTk.PhotoImage(img)
+        self.addImage_m(img)
+
+    def ChangeDataset(self,event):
+        
+        dataset_name=self.view.set_category.get()
+        
+        self.style_clip.LoadData(dataset_name)
+        
+        self.view.target.delete(1.0, "end")
+        self.view.target.insert("end", self.style_clip.target)
+        
+        self.view.neutral.delete(1.0, "end")
+        self.view.neutral.insert("end", self.style_clip.neutral)
+    
+    def text_t(self,event):
+        tmp=self.view.target.get("1.0",'end')
+        tmp=tmp.replace('\n','')
+        
+        self.view.target.delete(1.0, "end")
+        self.view.target.insert("end", tmp)
+        
+        print('target',tmp,'###')
+        self.style_clip.target=tmp
+        self.style_clip.GetDt2()
+        self.view.beta.set(self.style_clip.beta)
+        self.view.alpha.set(3)
+        self.style_clip.M.alpha=[3]
+        
+        img=self.style_clip.GetImg()
+        print('manipulate one')
+        img=Image.fromarray(img)
+        img = ImageTk.PhotoImage(img)
+        self.addImage_m(img)
+        
+        
+    def text_n(self,event):
+        tmp=self.view.neutral.get("1.0",'end')
+        tmp=tmp.replace('\n','')
+        
+        self.view.neutral.delete(1.0, "end")
+        self.view.neutral.insert("end", tmp)
+        
+        print('neutral',tmp,'###')
+        self.style_clip.neutral=tmp
+        self.view.target.delete(1.0, "end")
+        self.view.target.insert("end", tmp)
+        
+        
+    def run(self):
+        self.root.mainloop()
+    
+    def addImage(self,img):
+        self.view.bg.create_image(self.view.width/2, self.view.height/2, image=img, anchor='center')
+        self.image=img #save a copy of image. if not the image will disappear
+        
+    def addImage_m(self,img):
+        self.view.mani.create_image(512, 512, image=img, anchor='center')
+        self.image2=img
+        
+    
+    def openfn(self):
+        filename = askopenfilename(title='open',initialdir='./data/'+self.style_clip.M.dataset_name+'/',filetypes=[("all image format", ".jpg"),("all image format", ".png")])
+        return filename
+    
+    def open_img(self,event):
+        x = self.openfn()
+        print(x)
+        
+        
+        img = Image.open(x)
+        img2 = img.resize(( 512,512), Image.ANTIALIAS)
+        img2 = ImageTk.PhotoImage(img2)
+        self.addImage(img2)
+        
+        img = ImageTk.PhotoImage(img)
+        self.addImage_m(img)
+        
+        img_index=x.split('/')[-1].split('.')[0]
+        img_index=int(img_index)
+        print(img_index)
+        self.style_clip.M.img_index=img_index
+        self.style_clip.M.dlatent_tmp=[tmp[img_index:(img_index+1)] for tmp in self.style_clip.M.dlatents]
+        
+        
+        self.style_clip.GetDt2()
+        self.view.beta.set(self.style_clip.beta)
+        self.view.alpha.set(3)
+        
+    #%%
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Process some integers.')
+    
+    parser.add_argument('--dataset_name',type=str,default='ffhq',
+                    help='name of dataset, for example, ffhq')
+    
+    args = parser.parse_args()
+    dataset_name=args.dataset_name
+    
+    self=PlayInteractively(dataset_name)
+    self.run()
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    

PTI/models/StyleCLIP/global_directions/SingleChannel.py

@@ -0,0 +1,109 @@
+
+
+
+import numpy as np 
+import torch
+import clip
+from PIL import Image
+import copy
+from manipulate import Manipulator
+import argparse
+
+def GetImgF(out,model,preprocess):
+    imgs=out
+    imgs1=imgs.reshape([-1]+list(imgs.shape[2:]))
+    
+    tmp=[]
+    for i in range(len(imgs1)):
+        
+        img=Image.fromarray(imgs1[i])
+        image = preprocess(img).unsqueeze(0).to(device)
+        tmp.append(image)
+    
+    image=torch.cat(tmp)
+    with torch.no_grad():
+        image_features = model.encode_image(image)
+    
+    image_features1=image_features.cpu().numpy()
+    image_features1=image_features1.reshape(list(imgs.shape[:2])+[512])
+    
+    return image_features1
+
+def GetFs(fs):
+    tmp=np.linalg.norm(fs,axis=-1)
+    fs1=fs/tmp[:,:,:,None]
+    fs2=fs1[:,:,1,:]-fs1[:,:,0,:]  # 5*sigma - (-5)* sigma
+    fs3=fs2/np.linalg.norm(fs2,axis=-1)[:,:,None]
+    fs3=fs3.mean(axis=1)
+    fs3=fs3/np.linalg.norm(fs3,axis=-1)[:,None]
+    return fs3
+
+#%%
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Process some integers.')
+    
+    parser.add_argument('--dataset_name',type=str,default='cat',
+                    help='name of dataset, for example, ffhq')
+    args = parser.parse_args()
+    dataset_name=args.dataset_name
+    
+    #%%
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model, preprocess = clip.load("ViT-B/32", device=device)
+    #%%
+    M=Manipulator(dataset_name=dataset_name)
+    np.set_printoptions(suppress=True)
+    print(M.dataset_name)
+    #%%
+    img_sindex=0
+    num_images=100
+    dlatents_o=[]
+    tmp=img_sindex*num_images
+    for i in range(len(M.dlatents)):
+        tmp1=M.dlatents[i][tmp:(tmp+num_images)]
+        dlatents_o.append(tmp1)
+    #%%
+    
+    all_f=[]
+    M.alpha=[-5,5] #ffhq 5
+    M.step=2
+    M.num_images=num_images
+    select=np.array(M.mindexs)<=16 #below or equal to 128 resolution 
+    mindexs2=np.array(M.mindexs)[select]
+    for lindex in mindexs2: #ignore ToRGB layers
+        print(lindex)
+        num_c=M.dlatents[lindex].shape[1]
+        for cindex in range(num_c):
+            
+            M.dlatents=copy.copy(dlatents_o)
+            M.dlatents[lindex][:,cindex]=M.code_mean[lindex][cindex]
+            
+            M.manipulate_layers=[lindex]
+            codes,out=M.EditOneC(cindex) 
+            image_features1=GetImgF(out,model,preprocess)
+            all_f.append(image_features1)
+    
+    all_f=np.array(all_f)
+    
+    fs3=GetFs(all_f)
+    
+    #%%
+    file_path='./npy/'+M.dataset_name+'/'
+    np.save(file_path+'fs3',fs3)
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    

PTI/models/StyleCLIP/global_directions/data/ffhq/w_plus.npy

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

PTI/models/StyleCLIP/global_directions/dnnlib/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from .util import EasyDict, make_cache_dir_path

PTI/models/StyleCLIP/global_directions/dnnlib/tflib/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from . import autosummary
+from . import network
+from . import optimizer
+from . import tfutil
+from . import custom_ops
+
+from .tfutil import *
+from .network import Network
+
+from .optimizer import Optimizer
+
+from .custom_ops import get_plugin

PTI/models/StyleCLIP/global_directions/dnnlib/tflib/autosummary.py

@@ -0,0 +1,193 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Helper for adding automatically tracked values to Tensorboard.
+
+Autosummary creates an identity op that internally keeps track of the input
+values and automatically shows up in TensorBoard. The reported value
+represents an average over input components. The average is accumulated
+constantly over time and flushed when save_summaries() is called.
+
+Notes:
+- The output tensor must be used as an input for something else in the
+  graph. Otherwise, the autosummary op will not get executed, and the average
+  value will not get accumulated.
+- It is perfectly fine to include autosummaries with the same name in
+  several places throughout the graph, even if they are executed concurrently.
+- It is ok to also pass in a python scalar or numpy array. In this case, it
+  is added to the average immediately.
+"""
+
+from collections import OrderedDict
+import numpy as np
+import tensorflow as tf
+from tensorboard import summary as summary_lib
+from tensorboard.plugins.custom_scalar import layout_pb2
+
+from . import tfutil
+from .tfutil import TfExpression
+from .tfutil import TfExpressionEx
+
+# Enable "Custom scalars" tab in TensorBoard for advanced formatting.
+# Disabled by default to reduce tfevents file size.
+enable_custom_scalars = False
+
+_dtype = tf.float64
+_vars = OrderedDict()  # name => [var, ...]
+_immediate = OrderedDict()  # name => update_op, update_value
+_finalized = False
+_merge_op = None
+
+
+def _create_var(name: str, value_expr: TfExpression) -> TfExpression:
+    """Internal helper for creating autosummary accumulators."""
+    assert not _finalized
+    name_id = name.replace("/", "_")
+    v = tf.cast(value_expr, _dtype)
+
+    if v.shape.is_fully_defined():
+        size = np.prod(v.shape.as_list())
+        size_expr = tf.constant(size, dtype=_dtype)
+    else:
+        size = None
+        size_expr = tf.reduce_prod(tf.cast(tf.shape(v), _dtype))
+
+    if size == 1:
+        if v.shape.ndims != 0:
+            v = tf.reshape(v, [])
+        v = [size_expr, v, tf.square(v)]
+    else:
+        v = [size_expr, tf.reduce_sum(v), tf.reduce_sum(tf.square(v))]
+    v = tf.cond(tf.is_finite(v[1]), lambda: tf.stack(v), lambda: tf.zeros(3, dtype=_dtype))
+
+    with tfutil.absolute_name_scope("Autosummary/" + name_id), tf.control_dependencies(None):
+        var = tf.Variable(tf.zeros(3, dtype=_dtype), trainable=False)  # [sum(1), sum(x), sum(x**2)]
+    update_op = tf.cond(tf.is_variable_initialized(var), lambda: tf.assign_add(var, v), lambda: tf.assign(var, v))
+
+    if name in _vars:
+        _vars[name].append(var)
+    else:
+        _vars[name] = [var]
+    return update_op
+
+
+def autosummary(name: str, value: TfExpressionEx, passthru: TfExpressionEx = None, condition: TfExpressionEx = True) -> TfExpressionEx:
+    """Create a new autosummary.
+
+    Args:
+        name:     Name to use in TensorBoard
+        value:    TensorFlow expression or python value to track
+        passthru: Optionally return this TF node without modifications but tack an autosummary update side-effect to this node.
+
+    Example use of the passthru mechanism:
+
+    n = autosummary('l2loss', loss, passthru=n)
+
+    This is a shorthand for the following code:
+
+    with tf.control_dependencies([autosummary('l2loss', loss)]):
+        n = tf.identity(n)
+    """
+    tfutil.assert_tf_initialized()
+    name_id = name.replace("/", "_")
+
+    if tfutil.is_tf_expression(value):
+        with tf.name_scope("summary_" + name_id), tf.device(value.device):
+            condition = tf.convert_to_tensor(condition, name='condition')
+            update_op = tf.cond(condition, lambda: tf.group(_create_var(name, value)), tf.no_op)
+            with tf.control_dependencies([update_op]):
+                return tf.identity(value if passthru is None else passthru)
+
+    else:  # python scalar or numpy array
+        assert not tfutil.is_tf_expression(passthru)
+        assert not tfutil.is_tf_expression(condition)
+        if condition:
+            if name not in _immediate:
+                with tfutil.absolute_name_scope("Autosummary/" + name_id), tf.device(None), tf.control_dependencies(None):
+                    update_value = tf.placeholder(_dtype)
+                    update_op = _create_var(name, update_value)
+                    _immediate[name] = update_op, update_value
+            update_op, update_value = _immediate[name]
+            tfutil.run(update_op, {update_value: value})
+        return value if passthru is None else passthru
+
+
+def finalize_autosummaries() -> None:
+    """Create the necessary ops to include autosummaries in TensorBoard report.
+    Note: This should be done only once per graph.
+    """
+    global _finalized
+    tfutil.assert_tf_initialized()
+
+    if _finalized:
+        return None
+
+    _finalized = True
+    tfutil.init_uninitialized_vars([var for vars_list in _vars.values() for var in vars_list])
+
+    # Create summary ops.
+    with tf.device(None), tf.control_dependencies(None):
+        for name, vars_list in _vars.items():
+            name_id = name.replace("/", "_")
+            with tfutil.absolute_name_scope("Autosummary/" + name_id):
+                moments = tf.add_n(vars_list)
+                moments /= moments[0]
+                with tf.control_dependencies([moments]):  # read before resetting
+                    reset_ops = [tf.assign(var, tf.zeros(3, dtype=_dtype)) for var in vars_list]
+                    with tf.name_scope(None), tf.control_dependencies(reset_ops):  # reset before reporting
+                        mean = moments[1]
+                        std = tf.sqrt(moments[2] - tf.square(moments[1]))
+                        tf.summary.scalar(name, mean)
+                        if enable_custom_scalars:
+                            tf.summary.scalar("xCustomScalars/" + name + "/margin_lo", mean - std)
+                            tf.summary.scalar("xCustomScalars/" + name + "/margin_hi", mean + std)
+
+    # Setup layout for custom scalars.
+    layout = None
+    if enable_custom_scalars:
+        cat_dict = OrderedDict()
+        for series_name in sorted(_vars.keys()):
+            p = series_name.split("/")
+            cat = p[0] if len(p) >= 2 else ""
+            chart = "/".join(p[1:-1]) if len(p) >= 3 else p[-1]
+            if cat not in cat_dict:
+                cat_dict[cat] = OrderedDict()
+            if chart not in cat_dict[cat]:
+                cat_dict[cat][chart] = []
+            cat_dict[cat][chart].append(series_name)
+        categories = []
+        for cat_name, chart_dict in cat_dict.items():
+            charts = []
+            for chart_name, series_names in chart_dict.items():
+                series = []
+                for series_name in series_names:
+                    series.append(layout_pb2.MarginChartContent.Series(
+                        value=series_name,
+                        lower="xCustomScalars/" + series_name + "/margin_lo",
+                        upper="xCustomScalars/" + series_name + "/margin_hi"))
+                margin = layout_pb2.MarginChartContent(series=series)
+                charts.append(layout_pb2.Chart(title=chart_name, margin=margin))
+            categories.append(layout_pb2.Category(title=cat_name, chart=charts))
+        layout = summary_lib.custom_scalar_pb(layout_pb2.Layout(category=categories))
+    return layout
+
+def save_summaries(file_writer, global_step=None):
+    """Call FileWriter.add_summary() with all summaries in the default graph,
+    automatically finalizing and merging them on the first call.
+    """
+    global _merge_op
+    tfutil.assert_tf_initialized()
+
+    if _merge_op is None:
+        layout = finalize_autosummaries()
+        if layout is not None:
+            file_writer.add_summary(layout)
+        with tf.device(None), tf.control_dependencies(None):
+            _merge_op = tf.summary.merge_all()
+
+    file_writer.add_summary(_merge_op.eval(), global_step)

PTI/models/StyleCLIP/global_directions/dnnlib/tflib/custom_ops.py

@@ -0,0 +1,181 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""TensorFlow custom ops builder.
+"""
+
+import glob
+import os
+import re
+import uuid
+import hashlib
+import tempfile
+import shutil
+import tensorflow as tf
+from tensorflow.python.client import device_lib # pylint: disable=no-name-in-module
+
+from .. import util
+
+#----------------------------------------------------------------------------
+# Global configs.
+
+cuda_cache_path = None
+cuda_cache_version_tag = 'v1'
+do_not_hash_included_headers = True # Speed up compilation by assuming that headers included by the CUDA code never change.
+verbose = True # Print status messages to stdout.
+
+#----------------------------------------------------------------------------
+# Internal helper funcs.
+
+def _find_compiler_bindir():
+    hostx64_paths = sorted(glob.glob('C:/Program Files (x86)/Microsoft Visual Studio/*/Professional/VC/Tools/MSVC/*/bin/Hostx64/x64'), reverse=True)
+    if hostx64_paths != []:
+        return hostx64_paths[0]
+    hostx64_paths = sorted(glob.glob('C:/Program Files (x86)/Microsoft Visual Studio/*/BuildTools/VC/Tools/MSVC/*/bin/Hostx64/x64'), reverse=True)
+    if hostx64_paths != []:
+        return hostx64_paths[0]
+    hostx64_paths = sorted(glob.glob('C:/Program Files (x86)/Microsoft Visual Studio/*/Community/VC/Tools/MSVC/*/bin/Hostx64/x64'), reverse=True)
+    if hostx64_paths != []:
+        return hostx64_paths[0]
+    vc_bin_dir = 'C:/Program Files (x86)/Microsoft Visual Studio 14.0/vc/bin'
+    if os.path.isdir(vc_bin_dir):
+        return vc_bin_dir
+    return None
+
+def _get_compute_cap(device):
+    caps_str = device.physical_device_desc
+    m = re.search('compute capability: (\\d+).(\\d+)', caps_str)
+    major = m.group(1)
+    minor = m.group(2)
+    return (major, minor)
+
+def _get_cuda_gpu_arch_string():
+    gpus = [x for x in device_lib.list_local_devices() if x.device_type == 'GPU']
+    if len(gpus) == 0:
+        raise RuntimeError('No GPU devices found')
+    (major, minor) = _get_compute_cap(gpus[0])
+    return 'sm_%s%s' % (major, minor)
+
+def _run_cmd(cmd):
+    with os.popen(cmd) as pipe:
+        output = pipe.read()
+        status = pipe.close()
+    if status is not None:
+        raise RuntimeError('NVCC returned an error. See below for full command line and output log:\n\n%s\n\n%s' % (cmd, output))
+
+def _prepare_nvcc_cli(opts):
+    cmd = 'nvcc ' + opts.strip()
+    cmd += ' --disable-warnings'
+    cmd += ' --include-path "%s"' % tf.sysconfig.get_include()
+    cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'protobuf_archive', 'src')
+    cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'com_google_absl')
+    cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'eigen_archive')
+
+    compiler_bindir = _find_compiler_bindir()
+    if compiler_bindir is None:
+        # Require that _find_compiler_bindir succeeds on Windows.  Allow
+        # nvcc to use whatever is the default on Linux.
+        if os.name == 'nt':
+            raise RuntimeError('Could not find MSVC/GCC/CLANG installation on this computer. Check compiler_bindir_search_path list in "%s".' % __file__)
+    else:
+        cmd += ' --compiler-bindir "%s"' % compiler_bindir
+    cmd += ' 2>&1'
+    return cmd
+
+#----------------------------------------------------------------------------
+# Main entry point.
+
+_plugin_cache = dict()
+
+def get_plugin(cuda_file, extra_nvcc_options=[]):
+    cuda_file_base = os.path.basename(cuda_file)
+    cuda_file_name, cuda_file_ext = os.path.splitext(cuda_file_base)
+
+    # Already in cache?
+    if cuda_file in _plugin_cache:
+        return _plugin_cache[cuda_file]
+
+    # Setup plugin.
+    if verbose:
+        print('Setting up TensorFlow plugin "%s": ' % cuda_file_base, end='', flush=True)
+    try:
+        # Hash CUDA source.
+        md5 = hashlib.md5()
+        with open(cuda_file, 'rb') as f:
+            md5.update(f.read())
+        md5.update(b'\n')
+
+        # Hash headers included by the CUDA code by running it through the preprocessor.
+        if not do_not_hash_included_headers:
+            if verbose:
+                print('Preprocessing... ', end='', flush=True)
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                tmp_file = os.path.join(tmp_dir, cuda_file_name + '_tmp' + cuda_file_ext)
+                _run_cmd(_prepare_nvcc_cli('"%s" --preprocess -o "%s" --keep --keep-dir "%s"' % (cuda_file, tmp_file, tmp_dir)))
+                with open(tmp_file, 'rb') as f:
+                    bad_file_str = ('"' + cuda_file.replace('\\', '/') + '"').encode('utf-8') # __FILE__ in error check macros
+                    good_file_str = ('"' + cuda_file_base + '"').encode('utf-8')
+                    for ln in f:
+                        if not ln.startswith(b'# ') and not ln.startswith(b'#line '): # ignore line number pragmas
+                            ln = ln.replace(bad_file_str, good_file_str)
+                            md5.update(ln)
+                    md5.update(b'\n')
+
+        # Select compiler configs.
+        compile_opts = ''
+        if os.name == 'nt':
+            compile_opts += '"%s"' % os.path.join(tf.sysconfig.get_lib(), 'python', '_pywrap_tensorflow_internal.lib')
+        elif os.name == 'posix':
+            compile_opts += f' --compiler-options \'-fPIC\''
+            compile_opts += f' --compiler-options \'{" ".join(tf.sysconfig.get_compile_flags())}\''
+            compile_opts += f' --linker-options \'{" ".join(tf.sysconfig.get_link_flags())}\''
+        else:
+            assert False # not Windows or Linux, w00t?
+        compile_opts += f' --gpu-architecture={_get_cuda_gpu_arch_string()}'
+        compile_opts += ' --use_fast_math'
+        for opt in extra_nvcc_options:
+            compile_opts += ' ' + opt
+        nvcc_cmd = _prepare_nvcc_cli(compile_opts)
+
+        # Hash build configuration.
+        md5.update(('nvcc_cmd: ' + nvcc_cmd).encode('utf-8') + b'\n')
+        md5.update(('tf.VERSION: ' + tf.VERSION).encode('utf-8') + b'\n')
+        md5.update(('cuda_cache_version_tag: ' + cuda_cache_version_tag).encode('utf-8') + b'\n')
+
+        # Compile if not already compiled.
+        cache_dir = util.make_cache_dir_path('tflib-cudacache') if cuda_cache_path is None else cuda_cache_path
+        bin_file_ext = '.dll' if os.name == 'nt' else '.so'
+        bin_file = os.path.join(cache_dir, cuda_file_name + '_' + md5.hexdigest() + bin_file_ext)
+        if not os.path.isfile(bin_file):
+            if verbose:
+                print('Compiling... ', end='', flush=True)
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                tmp_file = os.path.join(tmp_dir, cuda_file_name + '_tmp' + bin_file_ext)
+                _run_cmd(nvcc_cmd + ' "%s" --shared -o "%s" --keep --keep-dir "%s"' % (cuda_file, tmp_file, tmp_dir))
+                os.makedirs(cache_dir, exist_ok=True)
+                intermediate_file = os.path.join(cache_dir, cuda_file_name + '_' + uuid.uuid4().hex + '_tmp' + bin_file_ext)
+                shutil.copyfile(tmp_file, intermediate_file)
+                os.rename(intermediate_file, bin_file) # atomic
+
+        # Load.
+        if verbose:
+            print('Loading... ', end='', flush=True)
+        plugin = tf.load_op_library(bin_file)
+
+        # Add to cache.
+        _plugin_cache[cuda_file] = plugin
+        if verbose:
+            print('Done.', flush=True)
+        return plugin
+
+    except:
+        if verbose:
+            print('Failed!', flush=True)
+        raise
+
+#----------------------------------------------------------------------------

PTI/models/StyleCLIP/global_directions/dnnlib/tflib/network.py

@@ -0,0 +1,781 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Helper for managing networks."""
+
+import types
+import inspect
+import re
+import uuid
+import sys
+import copy
+import numpy as np
+import tensorflow as tf
+
+from collections import OrderedDict
+from typing import Any, List, Tuple, Union, Callable
+
+from . import tfutil
+from .. import util
+
+from .tfutil import TfExpression, TfExpressionEx
+
+# pylint: disable=protected-access
+# pylint: disable=attribute-defined-outside-init
+# pylint: disable=too-many-public-methods
+
+_import_handlers = []  # Custom import handlers for dealing with legacy data in pickle import.
+_import_module_src = dict()  # Source code for temporary modules created during pickle import.
+
+
+def import_handler(handler_func):
+    """Function decorator for declaring custom import handlers."""
+    _import_handlers.append(handler_func)
+    return handler_func
+
+
+class Network:
+    """Generic network abstraction.
+
+    Acts as a convenience wrapper for a parameterized network construction
+    function, providing several utility methods and convenient access to
+    the inputs/outputs/weights.
+
+    Network objects can be safely pickled and unpickled for long-term
+    archival purposes. The pickling works reliably as long as the underlying
+    network construction function is defined in a standalone Python module
+    that has no side effects or application-specific imports.
+
+    Args:
+        name: Network name. Used to select TensorFlow name and variable scopes. Defaults to build func name if None.
+        func_name: Fully qualified name of the underlying network construction function, or a top-level function object.
+        static_kwargs: Keyword arguments to be passed in to the network construction function.
+    """
+
+    def __init__(self, name: str = None, func_name: Any = None, **static_kwargs):
+        # Locate the user-specified build function.
+        assert isinstance(func_name, str) or util.is_top_level_function(func_name)
+        if util.is_top_level_function(func_name):
+            func_name = util.get_top_level_function_name(func_name)
+        module, func_name = util.get_module_from_obj_name(func_name)
+        func = util.get_obj_from_module(module, func_name)
+
+        # Dig up source code for the module containing the build function.
+        module_src = _import_module_src.get(module, None)
+        if module_src is None:
+            module_src = inspect.getsource(module)
+
+        # Initialize fields.
+        self._init_fields(name=(name or func_name), static_kwargs=static_kwargs, build_func=func, build_func_name=func_name, build_module_src=module_src)
+
+    def _init_fields(self, name: str, static_kwargs: dict, build_func: Callable, build_func_name: str, build_module_src: str) -> None:
+        tfutil.assert_tf_initialized()
+        assert isinstance(name, str)
+        assert len(name) >= 1
+        assert re.fullmatch(r"[A-Za-z0-9_.\\-]*", name)
+        assert isinstance(static_kwargs, dict)
+        assert util.is_pickleable(static_kwargs)
+        assert callable(build_func)
+        assert isinstance(build_func_name, str)
+        assert isinstance(build_module_src, str)
+
+        # Choose TensorFlow name scope.
+        with tf.name_scope(None):
+            scope = tf.get_default_graph().unique_name(name, mark_as_used=True)
+
+        # Query current TensorFlow device.
+        with tfutil.absolute_name_scope(scope), tf.control_dependencies(None):
+            device = tf.no_op(name="_QueryDevice").device
+
+        # Immutable state.
+        self._name                  = name
+        self._scope                 = scope
+        self._device                = device
+        self._static_kwargs         = util.EasyDict(copy.deepcopy(static_kwargs))
+        self._build_func            = build_func
+        self._build_func_name       = build_func_name
+        self._build_module_src      = build_module_src
+
+        # State before _init_graph().
+        self._var_inits             = dict()    # var_name => initial_value, set to None by _init_graph()
+        self._all_inits_known       = False     # Do we know for sure that _var_inits covers all the variables?
+        self._components            = None      # subnet_name => Network, None if the components are not known yet
+
+        # Initialized by _init_graph().
+        self._input_templates       = None
+        self._output_templates      = None
+        self._own_vars              = None
+
+        # Cached values initialized the respective methods.
+        self._input_shapes          = None
+        self._output_shapes         = None
+        self._input_names           = None
+        self._output_names          = None
+        self._vars                  = None
+        self._trainables            = None
+        self._var_global_to_local   = None
+        self._run_cache             = dict()
+
+    def _init_graph(self) -> None:
+        assert self._var_inits is not None
+        assert self._input_templates is None
+        assert self._output_templates is None
+        assert self._own_vars is None
+
+        # Initialize components.
+        if self._components is None:
+            self._components = util.EasyDict()
+
+        # Choose build func kwargs.
+        build_kwargs = dict(self.static_kwargs)
+        build_kwargs["is_template_graph"] = True
+        build_kwargs["components"] = self._components
+
+        # Override scope and device, and ignore surrounding control dependencies.
+        with tfutil.absolute_variable_scope(self.scope, reuse=False), tfutil.absolute_name_scope(self.scope), tf.device(self.device), tf.control_dependencies(None):
+            assert tf.get_variable_scope().name == self.scope
+            assert tf.get_default_graph().get_name_scope() == self.scope
+
+            # Create input templates.
+            self._input_templates = []
+            for param in inspect.signature(self._build_func).parameters.values():
+                if param.kind == param.POSITIONAL_OR_KEYWORD and param.default is param.empty:
+                    self._input_templates.append(tf.placeholder(tf.float32, name=param.name))
+
+            # Call build func.
+            out_expr = self._build_func(*self._input_templates, **build_kwargs)
+
+        # Collect output templates and variables.
+        assert tfutil.is_tf_expression(out_expr) or isinstance(out_expr, tuple)
+        self._output_templates = [out_expr] if tfutil.is_tf_expression(out_expr) else list(out_expr)
+        self._own_vars = OrderedDict((var.name[len(self.scope) + 1:].split(":")[0], var) for var in tf.global_variables(self.scope + "/"))
+
+        # Check for errors.
+        if len(self._input_templates) == 0:
+            raise ValueError("Network build func did not list any inputs.")
+        if len(self._output_templates) == 0:
+            raise ValueError("Network build func did not return any outputs.")
+        if any(not tfutil.is_tf_expression(t) for t in self._output_templates):
+            raise ValueError("Network outputs must be TensorFlow expressions.")
+        if any(t.shape.ndims is None for t in self._input_templates):
+            raise ValueError("Network input shapes not defined. Please call x.set_shape() for each input.")
+        if any(t.shape.ndims is None for t in self._output_templates):
+            raise ValueError("Network output shapes not defined. Please call x.set_shape() where applicable.")
+        if any(not isinstance(comp, Network) for comp in self._components.values()):
+            raise ValueError("Components of a Network must be Networks themselves.")
+        if len(self._components) != len(set(comp.name for comp in self._components.values())):
+            raise ValueError("Components of a Network must have unique names.")
+
+        # Initialize variables.
+        if len(self._var_inits):
+            tfutil.set_vars({self._get_vars()[name]: value for name, value in self._var_inits.items() if name in self._get_vars()})
+        remaining_inits = [var.initializer for name, var in self._own_vars.items() if name not in self._var_inits]
+        if self._all_inits_known:
+            assert len(remaining_inits) == 0
+        else:
+            tfutil.run(remaining_inits)
+        self._var_inits = None
+
+    @property
+    def name(self):
+        """User-specified name string."""
+        return self._name
+
+    @property
+    def scope(self):
+        """Unique TensorFlow scope containing template graph and variables, derived from the user-specified name."""
+        return self._scope
+
+    @property
+    def device(self):
+        """Name of the TensorFlow device that the weights of this network reside on. Determined by the current device at construction time."""
+        return self._device
+
+    @property
+    def static_kwargs(self):
+        """EasyDict of arguments passed to the user-supplied build func."""
+        return copy.deepcopy(self._static_kwargs)
+
+    @property
+    def components(self):
+        """EasyDict of sub-networks created by the build func."""
+        return copy.copy(self._get_components())
+
+    def _get_components(self):
+        if self._components is None:
+            self._init_graph()
+            assert self._components is not None
+        return self._components
+
+    @property
+    def input_shapes(self):
+        """List of input tensor shapes, including minibatch dimension."""
+        if self._input_shapes is None:
+            self._input_shapes = [t.shape.as_list() for t in self.input_templates]
+        return copy.deepcopy(self._input_shapes)
+
+    @property
+    def output_shapes(self):
+        """List of output tensor shapes, including minibatch dimension."""
+        if self._output_shapes is None:
+            self._output_shapes = [t.shape.as_list() for t in self.output_templates]
+        return copy.deepcopy(self._output_shapes)
+
+    @property
+    def input_shape(self):
+        """Short-hand for input_shapes[0]."""
+        return self.input_shapes[0]
+
+    @property
+    def output_shape(self):
+        """Short-hand for output_shapes[0]."""
+        return self.output_shapes[0]
+
+    @property
+    def num_inputs(self):
+        """Number of input tensors."""
+        return len(self.input_shapes)
+
+    @property
+    def num_outputs(self):
+        """Number of output tensors."""
+        return len(self.output_shapes)
+
+    @property
+    def input_names(self):
+        """Name string for each input."""
+        if self._input_names is None:
+            self._input_names = [t.name.split("/")[-1].split(":")[0] for t in self.input_templates]
+        return copy.copy(self._input_names)
+
+    @property
+    def output_names(self):
+        """Name string for each output."""
+        if self._output_names is None:
+            self._output_names = [t.name.split("/")[-1].split(":")[0] for t in self.output_templates]
+        return copy.copy(self._output_names)
+
+    @property
+    def input_templates(self):
+        """Input placeholders in the template graph."""
+        if self._input_templates is None:
+            self._init_graph()
+            assert self._input_templates is not None
+        return copy.copy(self._input_templates)
+
+    @property
+    def output_templates(self):
+        """Output tensors in the template graph."""
+        if self._output_templates is None:
+            self._init_graph()
+            assert self._output_templates is not None
+        return copy.copy(self._output_templates)
+
+    @property
+    def own_vars(self):
+        """Variables defined by this network (local_name => var), excluding sub-networks."""
+        return copy.copy(self._get_own_vars())
+
+    def _get_own_vars(self):
+        if self._own_vars is None:
+            self._init_graph()
+            assert self._own_vars is not None
+        return self._own_vars
+
+    @property
+    def vars(self):
+        """All variables (local_name => var)."""
+        return copy.copy(self._get_vars())
+
+    def _get_vars(self):
+        if self._vars is None:
+            self._vars = OrderedDict(self._get_own_vars())
+            for comp in self._get_components().values():
+                self._vars.update((comp.name + "/" + name, var) for name, var in comp._get_vars().items())
+        return self._vars
+
+    @property
+    def trainables(self):
+        """All trainable variables (local_name => var)."""
+        return copy.copy(self._get_trainables())
+
+    def _get_trainables(self):
+        if self._trainables is None:
+            self._trainables = OrderedDict((name, var) for name, var in self.vars.items() if var.trainable)
+        return self._trainables
+
+    @property
+    def var_global_to_local(self):
+        """Mapping from variable global names to local names."""
+        return copy.copy(self._get_var_global_to_local())
+
+    def _get_var_global_to_local(self):
+        if self._var_global_to_local is None:
+            self._var_global_to_local = OrderedDict((var.name.split(":")[0], name) for name, var in self.vars.items())
+        return self._var_global_to_local
+
+    def reset_own_vars(self) -> None:
+        """Re-initialize all variables of this network, excluding sub-networks."""
+        if self._var_inits is None or self._components is None:
+            tfutil.run([var.initializer for var in self._get_own_vars().values()])
+        else:
+            self._var_inits.clear()
+            self._all_inits_known = False
+
+    def reset_vars(self) -> None:
+        """Re-initialize all variables of this network, including sub-networks."""
+        if self._var_inits is None:
+            tfutil.run([var.initializer for var in self._get_vars().values()])
+        else:
+            self._var_inits.clear()
+            self._all_inits_known = False
+            if self._components is not None:
+                for comp in self._components.values():
+                    comp.reset_vars()
+
+    def reset_trainables(self) -> None:
+        """Re-initialize all trainable variables of this network, including sub-networks."""
+        tfutil.run([var.initializer for var in self._get_trainables().values()])
+
+    def get_output_for(self, *in_expr: TfExpression, return_as_list: bool = False, **dynamic_kwargs) -> Union[TfExpression, List[TfExpression]]:
+        """Construct TensorFlow expression(s) for the output(s) of this network, given the input expression(s).
+        The graph is placed on the current TensorFlow device."""
+        assert len(in_expr) == self.num_inputs
+        assert not all(expr is None for expr in in_expr)
+        self._get_vars()  # ensure that all variables have been created
+
+        # Choose build func kwargs.
+        build_kwargs = dict(self.static_kwargs)
+        build_kwargs.update(dynamic_kwargs)
+        build_kwargs["is_template_graph"] = False
+        build_kwargs["components"] = self._components
+
+        # Build TensorFlow graph to evaluate the network.
+        with tfutil.absolute_variable_scope(self.scope, reuse=True), tf.name_scope(self.name):
+            assert tf.get_variable_scope().name == self.scope
+            valid_inputs = [expr for expr in in_expr if expr is not None]
+            final_inputs = []
+            for expr, name, shape in zip(in_expr, self.input_names, self.input_shapes):
+                if expr is not None:
+                    expr = tf.identity(expr, name=name)
+                else:
+                    expr = tf.zeros([tf.shape(valid_inputs[0])[0]] + shape[1:], name=name)
+                final_inputs.append(expr)
+            out_expr = self._build_func(*final_inputs, **build_kwargs)
+
+        # Propagate input shapes back to the user-specified expressions.
+        for expr, final in zip(in_expr, final_inputs):
+            if isinstance(expr, tf.Tensor):
+                expr.set_shape(final.shape)
+
+        # Express outputs in the desired format.
+        assert tfutil.is_tf_expression(out_expr) or isinstance(out_expr, tuple)
+        if return_as_list:
+            out_expr = [out_expr] if tfutil.is_tf_expression(out_expr) else list(out_expr)
+        return out_expr
+
+    def get_var_local_name(self, var_or_global_name: Union[TfExpression, str]) -> str:
+        """Get the local name of a given variable, without any surrounding name scopes."""
+        assert tfutil.is_tf_expression(var_or_global_name) or isinstance(var_or_global_name, str)
+        global_name = var_or_global_name if isinstance(var_or_global_name, str) else var_or_global_name.name
+        return self._get_var_global_to_local()[global_name]
+
+    def find_var(self, var_or_local_name: Union[TfExpression, str]) -> TfExpression:
+        """Find variable by local or global name."""
+        assert tfutil.is_tf_expression(var_or_local_name) or isinstance(var_or_local_name, str)
+        return self._get_vars()[var_or_local_name] if isinstance(var_or_local_name, str) else var_or_local_name
+
+    def get_var(self, var_or_local_name: Union[TfExpression, str]) -> np.ndarray:
+        """Get the value of a given variable as NumPy array.
+        Note: This method is very inefficient -- prefer to use tflib.run(list_of_vars) whenever possible."""
+        return self.find_var(var_or_local_name).eval()
+
+    def set_var(self, var_or_local_name: Union[TfExpression, str], new_value: Union[int, float, np.ndarray]) -> None:
+        """Set the value of a given variable based on the given NumPy array.
+        Note: This method is very inefficient -- prefer to use tflib.set_vars() whenever possible."""
+        tfutil.set_vars({self.find_var(var_or_local_name): new_value})
+
+    def __getstate__(self) -> dict:
+        """Pickle export."""
+        state = dict()
+        state["version"]            = 5
+        state["name"]               = self.name
+        state["static_kwargs"]      = dict(self.static_kwargs)
+        state["components"]         = dict(self.components)
+        state["build_module_src"]   = self._build_module_src
+        state["build_func_name"]    = self._build_func_name
+        state["variables"]          = list(zip(self._get_own_vars().keys(), tfutil.run(list(self._get_own_vars().values()))))
+        state["input_shapes"]       = self.input_shapes
+        state["output_shapes"]      = self.output_shapes
+        state["input_names"]        = self.input_names
+        state["output_names"]       = self.output_names
+        return state
+
+    def __setstate__(self, state: dict) -> None:
+        """Pickle import."""
+
+        # Execute custom import handlers.
+        for handler in _import_handlers:
+            state = handler(state)
+
+        # Get basic fields.
+        assert state["version"] in [2, 3, 4, 5]
+        name = state["name"]
+        static_kwargs = state["static_kwargs"]
+        build_module_src = state["build_module_src"]
+        build_func_name = state["build_func_name"]
+
+        # Create temporary module from the imported source code.
+        module_name = "_tflib_network_import_" + uuid.uuid4().hex
+        module = types.ModuleType(module_name)
+        sys.modules[module_name] = module
+        _import_module_src[module] = build_module_src
+        exec(build_module_src, module.__dict__) # pylint: disable=exec-used
+        build_func = util.get_obj_from_module(module, build_func_name)
+
+        # Initialize fields.
+        self._init_fields(name=name, static_kwargs=static_kwargs, build_func=build_func, build_func_name=build_func_name, build_module_src=build_module_src)
+        self._var_inits.update(copy.deepcopy(state["variables"]))
+        self._all_inits_known   = True
+        self._components        = util.EasyDict(state.get("components", {}))
+        self._input_shapes      = copy.deepcopy(state.get("input_shapes", None))
+        self._output_shapes     = copy.deepcopy(state.get("output_shapes", None))
+        self._input_names       = copy.deepcopy(state.get("input_names", None))
+        self._output_names      = copy.deepcopy(state.get("output_names", None))
+
+    def clone(self, name: str = None, **new_static_kwargs) -> "Network":
+        """Create a clone of this network with its own copy of the variables."""
+        static_kwargs = dict(self.static_kwargs)
+        static_kwargs.update(new_static_kwargs)
+        net = object.__new__(Network)
+        net._init_fields(name=(name or self.name), static_kwargs=static_kwargs, build_func=self._build_func, build_func_name=self._build_func_name, build_module_src=self._build_module_src)
+        net.copy_vars_from(self)
+        return net
+
+    def copy_own_vars_from(self, src_net: "Network") -> None:
+        """Copy the values of all variables from the given network, excluding sub-networks."""
+
+        # Source has unknown variables or unknown components => init now.
+        if (src_net._var_inits is not None and not src_net._all_inits_known) or src_net._components is None:
+            src_net._get_vars()
+
+       # Both networks are inited => copy directly.
+        if src_net._var_inits is None and self._var_inits is None:
+            names = [name for name in self._get_own_vars().keys() if name in src_net._get_own_vars()]
+            tfutil.set_vars(tfutil.run({self._get_vars()[name]: src_net._get_vars()[name] for name in names}))
+            return
+
+        # Read from source.
+        if src_net._var_inits is None:
+            value_dict = tfutil.run(src_net._get_own_vars())
+        else:
+            value_dict = src_net._var_inits
+
+        # Write to destination.
+        if self._var_inits is None:
+            tfutil.set_vars({self._get_vars()[name]: value for name, value in value_dict.items() if name in self._get_vars()})
+        else:
+            self._var_inits.update(value_dict)
+
+    def copy_vars_from(self, src_net: "Network") -> None:
+        """Copy the values of all variables from the given network, including sub-networks."""
+
+        # Source has unknown variables or unknown components => init now.
+        if (src_net._var_inits is not None and not src_net._all_inits_known) or src_net._components is None:
+            src_net._get_vars()
+
+        # Source is inited, but destination components have not been created yet => set as initial values.
+        if src_net._var_inits is None and self._components is None:
+            self._var_inits.update(tfutil.run(src_net._get_vars()))
+            return
+
+        # Destination has unknown components => init now.
+        if self._components is None:
+            self._get_vars()
+
+        # Both networks are inited => copy directly.
+        if src_net._var_inits is None and self._var_inits is None:
+            names = [name for name in self._get_vars().keys() if name in src_net._get_vars()]
+            tfutil.set_vars(tfutil.run({self._get_vars()[name]: src_net._get_vars()[name] for name in names}))
+            return
+
+        # Copy recursively, component by component.
+        self.copy_own_vars_from(src_net)
+        for name, src_comp in src_net._components.items():
+            if name in self._components:
+                self._components[name].copy_vars_from(src_comp)
+
+    def copy_trainables_from(self, src_net: "Network") -> None:
+        """Copy the values of all trainable variables from the given network, including sub-networks."""
+        names = [name for name in self._get_trainables().keys() if name in src_net._get_trainables()]
+        tfutil.set_vars(tfutil.run({self._get_vars()[name]: src_net._get_vars()[name] for name in names}))
+
+    def convert(self, new_func_name: str, new_name: str = None, **new_static_kwargs) -> "Network":
+        """Create new network with the given parameters, and copy all variables from this network."""
+        if new_name is None:
+            new_name = self.name
+        static_kwargs = dict(self.static_kwargs)
+        static_kwargs.update(new_static_kwargs)
+        net = Network(name=new_name, func_name=new_func_name, **static_kwargs)
+        net.copy_vars_from(self)
+        return net
+
+    def setup_as_moving_average_of(self, src_net: "Network", beta: TfExpressionEx = 0.99, beta_nontrainable: TfExpressionEx = 0.0) -> tf.Operation:
+        """Construct a TensorFlow op that updates the variables of this network
+        to be slightly closer to those of the given network."""
+        with tfutil.absolute_name_scope(self.scope + "/_MovingAvg"):
+            ops = []
+            for name, var in self._get_vars().items():
+                if name in src_net._get_vars():
+                    cur_beta = beta if var.trainable else beta_nontrainable
+                    new_value = tfutil.lerp(src_net._get_vars()[name], var, cur_beta)
+                    ops.append(var.assign(new_value))
+            return tf.group(*ops)
+
+    def run(self,
+            *in_arrays: Tuple[Union[np.ndarray, None], ...],
+            input_transform: dict = None,
+            output_transform: dict = None,
+            return_as_list: bool = False,
+            print_progress: bool = False,
+            minibatch_size: int = None,
+            num_gpus: int = 1,
+            assume_frozen: bool = False,
+            **dynamic_kwargs) -> Union[np.ndarray, Tuple[np.ndarray, ...], List[np.ndarray]]:
+        """Run this network for the given NumPy array(s), and return the output(s) as NumPy array(s).
+
+        Args:
+            input_transform:    A dict specifying a custom transformation to be applied to the input tensor(s) before evaluating the network.
+                                The dict must contain a 'func' field that points to a top-level function. The function is called with the input
+                                TensorFlow expression(s) as positional arguments. Any remaining fields of the dict will be passed in as kwargs.
+            output_transform:   A dict specifying a custom transformation to be applied to the output tensor(s) after evaluating the network.
+                                The dict must contain a 'func' field that points to a top-level function. The function is called with the output
+                                TensorFlow expression(s) as positional arguments. Any remaining fields of the dict will be passed in as kwargs.
+            return_as_list:     True = return a list of NumPy arrays, False = return a single NumPy array, or a tuple if there are multiple outputs.
+            print_progress:     Print progress to the console? Useful for very large input arrays.
+            minibatch_size:     Maximum minibatch size to use, None = disable batching.
+            num_gpus:           Number of GPUs to use.
+            assume_frozen:      Improve multi-GPU performance by assuming that the trainable parameters will remain changed between calls.
+            dynamic_kwargs:     Additional keyword arguments to be passed into the network build function.
+        """
+        assert len(in_arrays) == self.num_inputs
+        assert not all(arr is None for arr in in_arrays)
+        assert input_transform is None or util.is_top_level_function(input_transform["func"])
+        assert output_transform is None or util.is_top_level_function(output_transform["func"])
+        output_transform, dynamic_kwargs = _handle_legacy_output_transforms(output_transform, dynamic_kwargs)
+        num_items = in_arrays[0].shape[0]
+        if minibatch_size is None:
+            minibatch_size = num_items
+
+        # Construct unique hash key from all arguments that affect the TensorFlow graph.
+        key = dict(input_transform=input_transform, output_transform=output_transform, num_gpus=num_gpus, assume_frozen=assume_frozen, dynamic_kwargs=dynamic_kwargs)
+        def unwind_key(obj):
+            if isinstance(obj, dict):
+                return [(key, unwind_key(value)) for key, value in sorted(obj.items())]
+            if callable(obj):
+                return util.get_top_level_function_name(obj)
+            return obj
+        key = repr(unwind_key(key))
+
+        # Build graph.
+        if key not in self._run_cache:
+            with tfutil.absolute_name_scope(self.scope + "/_Run"), tf.control_dependencies(None):
+                with tf.device("/cpu:0"):
+                    in_expr = [tf.placeholder(tf.float32, name=name) for name in self.input_names]
+                    in_split = list(zip(*[tf.split(x, num_gpus) for x in in_expr]))
+
+                out_split = []
+                for gpu in range(num_gpus):
+                    with tf.device(self.device if num_gpus == 1 else "/gpu:%d" % gpu):
+                        net_gpu = self.clone() if assume_frozen else self
+                        in_gpu = in_split[gpu]
+
+                        if input_transform is not None:
+                            in_kwargs = dict(input_transform)
+                            in_gpu = in_kwargs.pop("func")(*in_gpu, **in_kwargs)
+                            in_gpu = [in_gpu] if tfutil.is_tf_expression(in_gpu) else list(in_gpu)
+
+                        assert len(in_gpu) == self.num_inputs
+                        out_gpu = net_gpu.get_output_for(*in_gpu, return_as_list=True, **dynamic_kwargs)
+
+                        if output_transform is not None:
+                            out_kwargs = dict(output_transform)
+                            out_gpu = out_kwargs.pop("func")(*out_gpu, **out_kwargs)
+                            out_gpu = [out_gpu] if tfutil.is_tf_expression(out_gpu) else list(out_gpu)
+
+                        assert len(out_gpu) == self.num_outputs
+                        out_split.append(out_gpu)
+
+                with tf.device("/cpu:0"):
+                    out_expr = [tf.concat(outputs, axis=0) for outputs in zip(*out_split)]
+                    self._run_cache[key] = in_expr, out_expr
+
+        # Run minibatches.
+        in_expr, out_expr = self._run_cache[key]
+        out_arrays = [np.empty([num_items] + expr.shape.as_list()[1:], expr.dtype.name) for expr in out_expr]
+
+        for mb_begin in range(0, num_items, minibatch_size):
+            if print_progress:
+                print("\r%d / %d" % (mb_begin, num_items), end="")
+
+            mb_end = min(mb_begin + minibatch_size, num_items)
+            mb_num = mb_end - mb_begin
+            mb_in = [src[mb_begin : mb_end] if src is not None else np.zeros([mb_num] + shape[1:]) for src, shape in zip(in_arrays, self.input_shapes)]
+            mb_out = tf.get_default_session().run(out_expr, dict(zip(in_expr, mb_in)))
+
+            for dst, src in zip(out_arrays, mb_out):
+                dst[mb_begin: mb_end] = src
+
+        # Done.
+        if print_progress:
+            print("\r%d / %d" % (num_items, num_items))
+
+        if not return_as_list:
+            out_arrays = out_arrays[0] if len(out_arrays) == 1 else tuple(out_arrays)
+        return out_arrays
+
+    def list_ops(self) -> List[TfExpression]:
+        _ = self.output_templates  # ensure that the template graph has been created
+        include_prefix = self.scope + "/"
+        exclude_prefix = include_prefix + "_"
+        ops = tf.get_default_graph().get_operations()
+        ops = [op for op in ops if op.name.startswith(include_prefix)]
+        ops = [op for op in ops if not op.name.startswith(exclude_prefix)]
+        return ops
+
+    def list_layers(self) -> List[Tuple[str, TfExpression, List[TfExpression]]]:
+        """Returns a list of (layer_name, output_expr, trainable_vars) tuples corresponding to
+        individual layers of the network. Mainly intended to be used for reporting."""
+        layers = []
+
+        def recurse(scope, parent_ops, parent_vars, level):
+            if len(parent_ops) == 0 and len(parent_vars) == 0:
+                return
+
+            # Ignore specific patterns.
+            if any(p in scope for p in ["/Shape", "/strided_slice", "/Cast", "/concat", "/Assign"]):
+                return
+
+            # Filter ops and vars by scope.
+            global_prefix = scope + "/"
+            local_prefix = global_prefix[len(self.scope) + 1:]
+            cur_ops = [op for op in parent_ops if op.name.startswith(global_prefix) or op.name == global_prefix[:-1]]
+            cur_vars = [(name, var) for name, var in parent_vars if name.startswith(local_prefix) or name == local_prefix[:-1]]
+            if not cur_ops and not cur_vars:
+                return
+
+            # Filter out all ops related to variables.
+            for var in [op for op in cur_ops if op.type.startswith("Variable")]:
+                var_prefix = var.name + "/"
+                cur_ops = [op for op in cur_ops if not op.name.startswith(var_prefix)]
+
+            # Scope does not contain ops as immediate children => recurse deeper.
+            contains_direct_ops = any("/" not in op.name[len(global_prefix):] and op.type not in ["Identity", "Cast", "Transpose"] for op in cur_ops)
+            if (level == 0 or not contains_direct_ops) and (len(cur_ops) != 0 or len(cur_vars) != 0):
+                visited = set()
+                for rel_name in [op.name[len(global_prefix):] for op in cur_ops] + [name[len(local_prefix):] for name, _var in cur_vars]:
+                    token = rel_name.split("/")[0]
+                    if token not in visited:
+                        recurse(global_prefix + token, cur_ops, cur_vars, level + 1)
+                        visited.add(token)
+                return
+
+            # Report layer.
+            layer_name = scope[len(self.scope) + 1:]
+            layer_output = cur_ops[-1].outputs[0] if cur_ops else cur_vars[-1][1]
+            layer_trainables = [var for _name, var in cur_vars if var.trainable]
+            layers.append((layer_name, layer_output, layer_trainables))
+
+        recurse(self.scope, self.list_ops(), list(self._get_vars().items()), 0)
+        return layers
+
+    def print_layers(self, title: str = None, hide_layers_with_no_params: bool = False) -> None:
+        """Print a summary table of the network structure."""
+        rows = [[title if title is not None else self.name, "Params", "OutputShape", "WeightShape"]]
+        rows += [["---"] * 4]
+        total_params = 0
+
+        for layer_name, layer_output, layer_trainables in self.list_layers():
+            num_params = sum(int(np.prod(var.shape.as_list())) for var in layer_trainables)
+            weights = [var for var in layer_trainables if var.name.endswith("/weight:0")]
+            weights.sort(key=lambda x: len(x.name))
+            if len(weights) == 0 and len(layer_trainables) == 1:
+                weights = layer_trainables
+            total_params += num_params
+
+            if not hide_layers_with_no_params or num_params != 0:
+                num_params_str = str(num_params) if num_params > 0 else "-"
+                output_shape_str = str(layer_output.shape)
+                weight_shape_str = str(weights[0].shape) if len(weights) >= 1 else "-"
+                rows += [[layer_name, num_params_str, output_shape_str, weight_shape_str]]
+
+        rows += [["---"] * 4]
+        rows += [["Total", str(total_params), "", ""]]
+
+        widths = [max(len(cell) for cell in column) for column in zip(*rows)]
+        print()
+        for row in rows:
+            print("  ".join(cell + " " * (width - len(cell)) for cell, width in zip(row, widths)))
+        print()
+
+    def setup_weight_histograms(self, title: str = None) -> None:
+        """Construct summary ops to include histograms of all trainable parameters in TensorBoard."""
+        if title is None:
+            title = self.name
+
+        with tf.name_scope(None), tf.device(None), tf.control_dependencies(None):
+            for local_name, var in self._get_trainables().items():
+                if "/" in local_name:
+                    p = local_name.split("/")
+                    name = title + "_" + p[-1] + "/" + "_".join(p[:-1])
+                else:
+                    name = title + "_toplevel/" + local_name
+
+                tf.summary.histogram(name, var)
+
+#----------------------------------------------------------------------------
+# Backwards-compatible emulation of legacy output transformation in Network.run().
+
+_print_legacy_warning = True
+
+def _handle_legacy_output_transforms(output_transform, dynamic_kwargs):
+    global _print_legacy_warning
+    legacy_kwargs = ["out_mul", "out_add", "out_shrink", "out_dtype"]
+    if not any(kwarg in dynamic_kwargs for kwarg in legacy_kwargs):
+        return output_transform, dynamic_kwargs
+
+    if _print_legacy_warning:
+        _print_legacy_warning = False
+        print()
+        print("WARNING: Old-style output transformations in Network.run() are deprecated.")
+        print("Consider using 'output_transform=dict(func=tflib.convert_images_to_uint8)'")
+        print("instead of 'out_mul=127.5, out_add=127.5, out_dtype=np.uint8'.")
+        print()
+    assert output_transform is None
+
+    new_kwargs = dict(dynamic_kwargs)
+    new_transform = {kwarg: new_kwargs.pop(kwarg) for kwarg in legacy_kwargs if kwarg in dynamic_kwargs}
+    new_transform["func"] = _legacy_output_transform_func
+    return new_transform, new_kwargs
+
+def _legacy_output_transform_func(*expr, out_mul=1.0, out_add=0.0, out_shrink=1, out_dtype=None):
+    if out_mul != 1.0:
+        expr = [x * out_mul for x in expr]
+
+    if out_add != 0.0:
+        expr = [x + out_add for x in expr]
+
+    if out_shrink > 1:
+        ksize = [1, 1, out_shrink, out_shrink]
+        expr = [tf.nn.avg_pool(x, ksize=ksize, strides=ksize, padding="VALID", data_format="NCHW") for x in expr]
+
+    if out_dtype is not None:
+        if tf.as_dtype(out_dtype).is_integer:
+            expr = [tf.round(x) for x in expr]
+        expr = [tf.saturate_cast(x, out_dtype) for x in expr]
+    return expr

PTI/models/StyleCLIP/global_directions/dnnlib/tflib/ops/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+# empty