simp

README.md

@@ -1,6 +1,6 @@
 ---
 title: PTI
-emoji: 🦀
+emoji: 🦀
 colorFrom: gray
 colorTo: pink
 sdk: gradio
@@ -34,4 +34,4 @@ Path to your main application file (which contains either `gradio` or `streamlit
 Path is relative to the root of the repository.
 
 `pinned`: _boolean_
-Whether the Space stays on top of your list.
+Whether the Space stays on top of your list.

configs/evaluation_config.py

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

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 = ''

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 = 300
+
+## Optimization
+pti_learning_rate = 3e-4
+first_inv_lr = 5e-3
+train_batch_size = 1
+use_last_w_pivots = False

configs/paths_config.py

@@ -0,0 +1,31 @@
+## Pretrained models paths
+e4e = './pretrained_models/e4e_ffhq_encode.pt'
+stylegan2_ada_ffhq = '/home/sayantan/PTI/pretrained_models/ffhq.pkl'
+style_clip_pretrained_mappers = ''
+ir_se50 = './pretrained_models/model_ir_se50.pth'
+dlib = './pretrained_models/align.dat'
+
+## Dirs for output files
+checkpoints_dir = './checkpoints'
+embedding_base_dir = './embeddings'
+styleclip_output_dir = './StyleCLIP_results'
+experiments_output_dir = './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 = 'rocky'
+
+## 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 = 'editings/interfacegan_directions/age.pt'
+interfacegan_smile = 'editings/interfacegan_directions/smile.pt'
+interfacegan_rotation = 'editings/interfacegan_directions/rotation.pt'
+ffhq_pca = 'editings/ganspace_pca/ffhq_pca.pt'

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

criteria/localitly_regulizer.py

@@ -0,0 +1,59 @@
+import torch
+import numpy as np
+import wandb
+from criteria import l2_loss
+from configs import hyperparameters
+from 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

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

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)

edit.py

@@ -0,0 +1,84 @@
+import wandb
+import click
+import os
+import sys
+import pickle
+import numpy as np
+from PIL import Image
+import glob
+import torch
+from configs import paths_config, hyperparameters, global_config
+from IPython.display import display
+import matplotlib.pyplot as plt
+from scripts.latent_editor_wrapper import LatentEditorWrapper
+
+
+image_dir_name = '/home/sayantan/processed_images'
+use_multi_id_training = False
+global_config.device = 'cuda'
+paths_config.e4e = '/home/sayantan/PTI/pretrained_models/e4e_ffhq_encode.pt'
+paths_config.input_data_id = image_dir_name
+paths_config.input_data_path = f'{image_dir_name}'
+paths_config.stylegan2_ada_ffhq = '/home/sayantan/PTI/pretrained_models/ffhq.pkl'
+paths_config.checkpoints_dir = '/home/sayantan/PTI/'
+paths_config.style_clip_pretrained_mappers = '/home/sayantan/PTI/pretrained_models'
+hyperparameters.use_locality_regularization = False
+hyperparameters.lpips_type = 'squeeze'
+
+model_id = "MYJJDFVGATAT"
+
+
+
+def display_alongside_source_image(images): 
+    res = np.concatenate([np.array(image) for image in images], axis=1)
+    return Image.fromarray(res)
+
+def load_generators(model_id, image_name):
+  with open(paths_config.stylegan2_ada_ffhq, 'rb') as f:
+    old_G = pickle.load(f)['G_ema'].cuda()
+    
+  with open(f'{paths_config.checkpoints_dir}/model_{model_id}_{image_name}.pt', 'rb') as f_new: 
+    new_G = torch.load(f_new).cuda()
+
+  return old_G, new_G
+
+def plot_syn_images(syn_images,text):
+  for img in syn_images: 
+      img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).detach().cpu().numpy()[0] 
+      plt.axis('off') 
+      resized_image = Image.fromarray(img,mode='RGB').resize((256,256))
+      display(resized_image)
+      #wandb.log({text: [wandb.Image(resized_image, caption="Label")]})
+      del img 
+      del resized_image 
+      torch.cuda.empty_cache()
+      
+def syn_images_wandb(img):
+  img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).detach().cpu().numpy()[0] 
+  plt.axis('off') 
+  resized_image = Image.fromarray(img,mode='RGB').resize((256,256)) 
+  return resized_image
+
+
+def edit(image_name):
+  generator_type = paths_config.multi_id_model_type if use_multi_id_training else image_name
+  old_G, new_G = load_generators(model_id, generator_type)
+  w_path_dir = f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}'
+  
+  embedding_dir = f'{w_path_dir}/{paths_config.pti_results_keyword}/{image_name}'
+  w_pivot = torch.load(f'{embedding_dir}/0.pt')
+  
+  old_image = old_G.synthesis(w_pivot, noise_mode='const', force_fp32 = True)
+  new_image = new_G.synthesis(w_pivot, noise_mode='const', force_fp32 = True)
+  
+  latent_editor = LatentEditorWrapper()
+  latents_after_edit = latent_editor.get_single_interface_gan_edits(w_pivot, [i for i in range(-5,5)])
+  
+  for direction, factor_and_edit in latents_after_edit.items():
+    for editkey in factor_and_edit.keys():
+      new_image = new_G.synthesis(factor_and_edit[editkey], noise_mode='const', force_fp32 = True)
+      image_pil = syn_images_wandb(new_image).save(f"/home/sayantan/PTI/{direction}/{editkey}/{image_name}.jpg")
+
+if __name__ == '__main__':
+    for image_name in [f.split(".")[0].split("_")[2] for f in sorted(glob.glob("*.pt"))]:
+        edit(image_name)

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

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

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)

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)

makedirs.py

@@ -0,0 +1,84 @@
+import wandb
+import click
+import os
+import sys
+import pickle
+import numpy as np
+from PIL import Image
+import torch
+from configs import paths_config, hyperparameters, global_config
+from IPython.display import display
+import matplotlib.pyplot as plt
+from scripts.latent_editor_wrapper import LatentEditorWrapper
+
+
+image_dir_name = '/home/sayantan/processed_images'
+use_multi_id_training = False
+global_config.device = 'cuda'
+paths_config.e4e = '/home/sayantan/PTI/pretrained_models/e4e_ffhq_encode.pt'
+paths_config.input_data_id = image_dir_name
+paths_config.input_data_path = f'{image_dir_name}'
+paths_config.stylegan2_ada_ffhq = '/home/sayantan/PTI/pretrained_models/ffhq.pkl'
+paths_config.checkpoints_dir = '/home/sayantan/PTI/'
+paths_config.style_clip_pretrained_mappers = '/home/sayantan/PTI/pretrained_models'
+hyperparameters.use_locality_regularization = False
+hyperparameters.lpips_type = 'squeeze'
+
+model_id = "MYJJDFVGATAT"
+
+
+
+def display_alongside_source_image(images): 
+    res = np.concatenate([np.array(image) for image in images], axis=1)
+    return Image.fromarray(res)
+
+def load_generators(model_id, image_name):
+  with open(paths_config.stylegan2_ada_ffhq, 'rb') as f:
+    old_G = pickle.load(f)['G_ema'].cuda()
+    
+  with open(f'{paths_config.checkpoints_dir}/model_{model_id}_{image_name}.pt', 'rb') as f_new: 
+    new_G = torch.load(f_new).cuda()
+
+  return old_G, new_G
+
+def plot_syn_images(syn_images,text):
+  for img in syn_images: 
+      img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).detach().cpu().numpy()[0] 
+      plt.axis('off') 
+      resized_image = Image.fromarray(img,mode='RGB').resize((256,256))
+      display(resized_image)
+      #wandb.log({text: [wandb.Image(resized_image, caption="Label")]})
+      del img 
+      del resized_image 
+      torch.cuda.empty_cache()
+      
+def syn_images_wandb(img):
+  img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).detach().cpu().numpy()[0] 
+  plt.axis('off') 
+  resized_image = Image.fromarray(img,mode='RGB').resize((256,256)) 
+  return resized_image
+
+@click.command()
+@click.pass_context
+@click.option('--image_name', prompt='image name', help='The name for image')
+
+def makedir(ctx: click.Context,image_name):
+  generator_type = paths_config.multi_id_model_type if use_multi_id_training else image_name
+  old_G, new_G = load_generators(model_id, generator_type)
+  w_path_dir = f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}'
+  
+  embedding_dir = f'{w_path_dir}/{paths_config.pti_results_keyword}/{image_name}'
+  w_pivot = torch.load(f'{embedding_dir}/0.pt')
+  
+  old_image = old_G.synthesis(w_pivot, noise_mode='const', force_fp32 = True)
+  new_image = new_G.synthesis(w_pivot, noise_mode='const', force_fp32 = True)
+  
+  latent_editor = LatentEditorWrapper()
+  latents_after_edit = latent_editor.get_single_interface_gan_edits(w_pivot, [i for i in range(-5,5)])
+  
+  for direction, factor_and_edit in latents_after_edit.items():
+    for editkey in factor_and_edit.keys():
+      os.makedirs(f"/home/sayantan/PTI/{direction}/{editkey}")
+
+if __name__ == '__main__':
+    makedir()

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

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

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()
+    
+    
+    
+    
+    
+    
+    

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)
+    
+    

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)
+    
+    
+    
+    
+    

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')
+    #%%
+
+    
+    

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
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+