ucalyptus commited on
Commit
2d7efb8
1 Parent(s): c3699bb
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. README.md +2 -2
  2. configs/__init__.py +0 -0
  3. configs/__pycache__/__init__.cpython-36.pyc +0 -0
  4. configs/__pycache__/__init__.cpython-39.pyc +0 -0
  5. configs/__pycache__/global_config.cpython-36.pyc +0 -0
  6. configs/__pycache__/global_config.cpython-39.pyc +0 -0
  7. configs/__pycache__/hyperparameters.cpython-36.pyc +0 -0
  8. configs/__pycache__/hyperparameters.cpython-39.pyc +0 -0
  9. configs/__pycache__/paths_config.cpython-36.pyc +0 -0
  10. configs/__pycache__/paths_config.cpython-39.pyc +0 -0
  11. configs/evaluation_config.py +1 -0
  12. configs/global_config.py +12 -0
  13. configs/hyperparameters.py +28 -0
  14. configs/paths_config.py +31 -0
  15. criteria/__init__.py +0 -0
  16. criteria/__pycache__/__init__.cpython-36.pyc +0 -0
  17. criteria/__pycache__/__init__.cpython-39.pyc +0 -0
  18. criteria/__pycache__/l2_loss.cpython-36.pyc +0 -0
  19. criteria/__pycache__/l2_loss.cpython-39.pyc +0 -0
  20. criteria/__pycache__/localitly_regulizer.cpython-36.pyc +0 -0
  21. criteria/__pycache__/localitly_regulizer.cpython-39.pyc +0 -0
  22. criteria/l2_loss.py +8 -0
  23. criteria/localitly_regulizer.py +59 -0
  24. dnnlib/__init__.py +9 -0
  25. dnnlib/__pycache__/__init__.cpython-36.pyc +0 -0
  26. dnnlib/__pycache__/__init__.cpython-39.pyc +0 -0
  27. dnnlib/__pycache__/util.cpython-36.pyc +0 -0
  28. dnnlib/__pycache__/util.cpython-39.pyc +0 -0
  29. dnnlib/util.py +477 -0
  30. edit.py +84 -0
  31. editings/__pycache__/ganspace.cpython-39.pyc +0 -0
  32. editings/__pycache__/latent_editor.cpython-39.pyc +0 -0
  33. editings/ganspace.py +21 -0
  34. editings/ganspace_pca/ffhq_pca.pt +0 -0
  35. editings/interfacegan_directions/age.pt +0 -0
  36. editings/interfacegan_directions/rotation.pt +0 -0
  37. editings/interfacegan_directions/smile.pt +0 -0
  38. editings/latent_editor.py +23 -0
  39. evaluation/experiment_setting_creator.py +43 -0
  40. evaluation/qualitative_edit_comparison.py +156 -0
  41. makedirs.py +84 -0
  42. models/StyleCLIP/__init__.py +0 -0
  43. models/StyleCLIP/criteria/__init__.py +0 -0
  44. models/StyleCLIP/criteria/clip_loss.py +17 -0
  45. models/StyleCLIP/criteria/id_loss.py +39 -0
  46. models/StyleCLIP/global_directions/GUI.py +103 -0
  47. models/StyleCLIP/global_directions/GenerateImg.py +50 -0
  48. models/StyleCLIP/global_directions/GetCode.py +232 -0
  49. models/StyleCLIP/global_directions/GetGUIData.py +67 -0
  50. models/StyleCLIP/global_directions/Inference.py +106 -0
README.md CHANGED
@@ -1,6 +1,6 @@
1
  ---
2
  title: PTI
3
- emoji: 🦀
4
  colorFrom: gray
5
  colorTo: pink
6
  sdk: gradio
@@ -34,4 +34,4 @@ Path to your main application file (which contains either `gradio` or `streamlit
34
  Path is relative to the root of the repository.
35
 
36
  `pinned`: _boolean_
37
- Whether the Space stays on top of your list.
 
1
  ---
2
  title: PTI
3
+ emoji: 🦀
4
  colorFrom: gray
5
  colorTo: pink
6
  sdk: gradio
 
34
  Path is relative to the root of the repository.
35
 
36
  `pinned`: _boolean_
37
+ Whether the Space stays on top of your list.
configs/__init__.py ADDED
File without changes
configs/__pycache__/__init__.cpython-36.pyc ADDED
Binary file (121 Bytes). View file
 
configs/__pycache__/__init__.cpython-39.pyc ADDED
Binary file (129 Bytes). View file
 
configs/__pycache__/global_config.cpython-36.pyc ADDED
Binary file (331 Bytes). View file
 
configs/__pycache__/global_config.cpython-39.pyc ADDED
Binary file (339 Bytes). View file
 
configs/__pycache__/hyperparameters.cpython-36.pyc ADDED
Binary file (680 Bytes). View file
 
configs/__pycache__/hyperparameters.cpython-39.pyc ADDED
Binary file (693 Bytes). View file
 
configs/__pycache__/paths_config.cpython-36.pyc ADDED
Binary file (1.06 kB). View file
 
configs/__pycache__/paths_config.cpython-39.pyc ADDED
Binary file (1.04 kB). View file
 
configs/evaluation_config.py ADDED
@@ -0,0 +1 @@
 
 
1
+ evaluated_methods = ['e4e', 'SG2', 'SG2Plus']
configs/global_config.py ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ## Device
2
+ cuda_visible_devices = '0'
3
+ device = 'cuda:0'
4
+
5
+ ## Logs
6
+ training_step = 1
7
+ image_rec_result_log_snapshot = 100
8
+ pivotal_training_steps = 0
9
+ model_snapshot_interval = 400
10
+
11
+ ## Run name to be updated during PTI
12
+ run_name = ''
configs/hyperparameters.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ## Architechture
2
+ lpips_type = 'alex'
3
+ first_inv_type = 'w'
4
+ optim_type = 'adam'
5
+
6
+ ## Locality regularization
7
+ latent_ball_num_of_samples = 1
8
+ locality_regularization_interval = 1
9
+ use_locality_regularization = False
10
+ regulizer_l2_lambda = 0.1
11
+ regulizer_lpips_lambda = 0.1
12
+ regulizer_alpha = 30
13
+
14
+ ## Loss
15
+ pt_l2_lambda = 1
16
+ pt_lpips_lambda = 1
17
+
18
+ ## Steps
19
+ LPIPS_value_threshold = 0.06
20
+ max_pti_steps = 350
21
+ first_inv_steps = 450
22
+ max_images_to_invert = 300
23
+
24
+ ## Optimization
25
+ pti_learning_rate = 3e-4
26
+ first_inv_lr = 5e-3
27
+ train_batch_size = 1
28
+ use_last_w_pivots = False
configs/paths_config.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ## Pretrained models paths
2
+ e4e = './pretrained_models/e4e_ffhq_encode.pt'
3
+ stylegan2_ada_ffhq = '/home/sayantan/PTI/pretrained_models/ffhq.pkl'
4
+ style_clip_pretrained_mappers = ''
5
+ ir_se50 = './pretrained_models/model_ir_se50.pth'
6
+ dlib = './pretrained_models/align.dat'
7
+
8
+ ## Dirs for output files
9
+ checkpoints_dir = './checkpoints'
10
+ embedding_base_dir = './embeddings'
11
+ styleclip_output_dir = './StyleCLIP_results'
12
+ experiments_output_dir = './output'
13
+
14
+ ## Input info
15
+ ### Input dir, where the images reside
16
+ input_data_path = ''
17
+ ### Inversion identifier, used to keeping track of the inversion results. Both the latent code and the generator
18
+ input_data_id = 'rocky'
19
+
20
+ ## Keywords
21
+ pti_results_keyword = 'PTI'
22
+ e4e_results_keyword = 'e4e'
23
+ sg2_results_keyword = 'SG2'
24
+ sg2_plus_results_keyword = 'SG2_plus'
25
+ multi_id_model_type = 'multi_id'
26
+
27
+ ## Edit directions
28
+ interfacegan_age = 'editings/interfacegan_directions/age.pt'
29
+ interfacegan_smile = 'editings/interfacegan_directions/smile.pt'
30
+ interfacegan_rotation = 'editings/interfacegan_directions/rotation.pt'
31
+ ffhq_pca = 'editings/ganspace_pca/ffhq_pca.pt'
criteria/__init__.py ADDED
File without changes
criteria/__pycache__/__init__.cpython-36.pyc ADDED
Binary file (122 Bytes). View file
 
criteria/__pycache__/__init__.cpython-39.pyc ADDED
Binary file (130 Bytes). View file
 
criteria/__pycache__/l2_loss.cpython-36.pyc ADDED
Binary file (346 Bytes). View file
 
criteria/__pycache__/l2_loss.cpython-39.pyc ADDED
Binary file (358 Bytes). View file
 
criteria/__pycache__/localitly_regulizer.cpython-36.pyc ADDED
Binary file (2.86 kB). View file
 
criteria/__pycache__/localitly_regulizer.cpython-39.pyc ADDED
Binary file (2.9 kB). View file
 
criteria/l2_loss.py ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+ l2_criterion = torch.nn.MSELoss(reduction='mean')
4
+
5
+
6
+ def l2_loss(real_images, generated_images):
7
+ loss = l2_criterion(real_images, generated_images)
8
+ return loss
criteria/localitly_regulizer.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import numpy as np
3
+ import wandb
4
+ from criteria import l2_loss
5
+ from configs import hyperparameters
6
+ from configs import global_config
7
+
8
+
9
+ class Space_Regulizer:
10
+ def __init__(self, original_G, lpips_net):
11
+ self.original_G = original_G
12
+ self.morphing_regulizer_alpha = hyperparameters.regulizer_alpha
13
+ self.lpips_loss = lpips_net
14
+
15
+ def get_morphed_w_code(self, new_w_code, fixed_w):
16
+ interpolation_direction = new_w_code - fixed_w
17
+ interpolation_direction_norm = torch.norm(interpolation_direction, p=2)
18
+ direction_to_move = hyperparameters.regulizer_alpha * interpolation_direction / interpolation_direction_norm
19
+ result_w = fixed_w + direction_to_move
20
+ self.morphing_regulizer_alpha * fixed_w + (1 - self.morphing_regulizer_alpha) * new_w_code
21
+
22
+ return result_w
23
+
24
+ def get_image_from_ws(self, w_codes, G):
25
+ return torch.cat([G.synthesis(w_code, noise_mode='none', force_fp32=True) for w_code in w_codes])
26
+
27
+ def ball_holder_loss_lazy(self, new_G, num_of_sampled_latents, w_batch, use_wandb=False):
28
+ loss = 0.0
29
+
30
+ z_samples = np.random.randn(num_of_sampled_latents, self.original_G.z_dim)
31
+ w_samples = self.original_G.mapping(torch.from_numpy(z_samples).to(global_config.device), None,
32
+ truncation_psi=0.5)
33
+ territory_indicator_ws = [self.get_morphed_w_code(w_code.unsqueeze(0), w_batch) for w_code in w_samples]
34
+
35
+ for w_code in territory_indicator_ws:
36
+ new_img = new_G.synthesis(w_code, noise_mode='none', force_fp32=True)
37
+ with torch.no_grad():
38
+ old_img = self.original_G.synthesis(w_code, noise_mode='none', force_fp32=True)
39
+
40
+ if hyperparameters.regulizer_l2_lambda > 0:
41
+ l2_loss_val = l2_loss.l2_loss(old_img, new_img)
42
+ if use_wandb:
43
+ wandb.log({f'space_regulizer_l2_loss_val': l2_loss_val.detach().cpu()},
44
+ step=global_config.training_step)
45
+ loss += l2_loss_val * hyperparameters.regulizer_l2_lambda
46
+
47
+ if hyperparameters.regulizer_lpips_lambda > 0:
48
+ loss_lpips = self.lpips_loss(old_img, new_img)
49
+ loss_lpips = torch.mean(torch.squeeze(loss_lpips))
50
+ if use_wandb:
51
+ wandb.log({f'space_regulizer_lpips_loss_val': loss_lpips.detach().cpu()},
52
+ step=global_config.training_step)
53
+ loss += loss_lpips * hyperparameters.regulizer_lpips_lambda
54
+
55
+ return loss / len(territory_indicator_ws)
56
+
57
+ def space_regulizer_loss(self, new_G, w_batch, use_wandb):
58
+ ret_val = self.ball_holder_loss_lazy(new_G, hyperparameters.latent_ball_num_of_samples, w_batch, use_wandb)
59
+ return ret_val
dnnlib/__init__.py ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
2
+ #
3
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
4
+ # and proprietary rights in and to this software, related documentation
5
+ # and any modifications thereto. Any use, reproduction, disclosure or
6
+ # distribution of this software and related documentation without an express
7
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
8
+
9
+ from .util import EasyDict, make_cache_dir_path
dnnlib/__pycache__/__init__.cpython-36.pyc ADDED
Binary file (187 Bytes). View file
 
dnnlib/__pycache__/__init__.cpython-39.pyc ADDED
Binary file (195 Bytes). View file
 
dnnlib/__pycache__/util.cpython-36.pyc ADDED
Binary file (13.6 kB). View file
 
dnnlib/__pycache__/util.cpython-39.pyc ADDED
Binary file (13.7 kB). View file
 
dnnlib/util.py ADDED
@@ -0,0 +1,477 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
2
+ #
3
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
4
+ # and proprietary rights in and to this software, related documentation
5
+ # and any modifications thereto. Any use, reproduction, disclosure or
6
+ # distribution of this software and related documentation without an express
7
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
8
+
9
+ """Miscellaneous utility classes and functions."""
10
+
11
+ import ctypes
12
+ import fnmatch
13
+ import importlib
14
+ import inspect
15
+ import numpy as np
16
+ import os
17
+ import shutil
18
+ import sys
19
+ import types
20
+ import io
21
+ import pickle
22
+ import re
23
+ import requests
24
+ import html
25
+ import hashlib
26
+ import glob
27
+ import tempfile
28
+ import urllib
29
+ import urllib.request
30
+ import uuid
31
+
32
+ from distutils.util import strtobool
33
+ from typing import Any, List, Tuple, Union
34
+
35
+
36
+ # Util classes
37
+ # ------------------------------------------------------------------------------------------
38
+
39
+
40
+ class EasyDict(dict):
41
+ """Convenience class that behaves like a dict but allows access with the attribute syntax."""
42
+
43
+ def __getattr__(self, name: str) -> Any:
44
+ try:
45
+ return self[name]
46
+ except KeyError:
47
+ raise AttributeError(name)
48
+
49
+ def __setattr__(self, name: str, value: Any) -> None:
50
+ self[name] = value
51
+
52
+ def __delattr__(self, name: str) -> None:
53
+ del self[name]
54
+
55
+
56
+ class Logger(object):
57
+ """Redirect stderr to stdout, optionally print stdout to a file, and optionally force flushing on both stdout and the file."""
58
+
59
+ def __init__(self, file_name: str = None, file_mode: str = "w", should_flush: bool = True):
60
+ self.file = None
61
+
62
+ if file_name is not None:
63
+ self.file = open(file_name, file_mode)
64
+
65
+ self.should_flush = should_flush
66
+ self.stdout = sys.stdout
67
+ self.stderr = sys.stderr
68
+
69
+ sys.stdout = self
70
+ sys.stderr = self
71
+
72
+ def __enter__(self) -> "Logger":
73
+ return self
74
+
75
+ def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
76
+ self.close()
77
+
78
+ def write(self, text: Union[str, bytes]) -> None:
79
+ """Write text to stdout (and a file) and optionally flush."""
80
+ if isinstance(text, bytes):
81
+ text = text.decode()
82
+ if len(text) == 0: # workaround for a bug in VSCode debugger: sys.stdout.write(''); sys.stdout.flush() => crash
83
+ return
84
+
85
+ if self.file is not None:
86
+ self.file.write(text)
87
+
88
+ self.stdout.write(text)
89
+
90
+ if self.should_flush:
91
+ self.flush()
92
+
93
+ def flush(self) -> None:
94
+ """Flush written text to both stdout and a file, if open."""
95
+ if self.file is not None:
96
+ self.file.flush()
97
+
98
+ self.stdout.flush()
99
+
100
+ def close(self) -> None:
101
+ """Flush, close possible files, and remove stdout/stderr mirroring."""
102
+ self.flush()
103
+
104
+ # if using multiple loggers, prevent closing in wrong order
105
+ if sys.stdout is self:
106
+ sys.stdout = self.stdout
107
+ if sys.stderr is self:
108
+ sys.stderr = self.stderr
109
+
110
+ if self.file is not None:
111
+ self.file.close()
112
+ self.file = None
113
+
114
+
115
+ # Cache directories
116
+ # ------------------------------------------------------------------------------------------
117
+
118
+ _dnnlib_cache_dir = None
119
+
120
+ def set_cache_dir(path: str) -> None:
121
+ global _dnnlib_cache_dir
122
+ _dnnlib_cache_dir = path
123
+
124
+ def make_cache_dir_path(*paths: str) -> str:
125
+ if _dnnlib_cache_dir is not None:
126
+ return os.path.join(_dnnlib_cache_dir, *paths)
127
+ if 'DNNLIB_CACHE_DIR' in os.environ:
128
+ return os.path.join(os.environ['DNNLIB_CACHE_DIR'], *paths)
129
+ if 'HOME' in os.environ:
130
+ return os.path.join(os.environ['HOME'], '.cache', 'dnnlib', *paths)
131
+ if 'USERPROFILE' in os.environ:
132
+ return os.path.join(os.environ['USERPROFILE'], '.cache', 'dnnlib', *paths)
133
+ return os.path.join(tempfile.gettempdir(), '.cache', 'dnnlib', *paths)
134
+
135
+ # Small util functions
136
+ # ------------------------------------------------------------------------------------------
137
+
138
+
139
+ def format_time(seconds: Union[int, float]) -> str:
140
+ """Convert the seconds to human readable string with days, hours, minutes and seconds."""
141
+ s = int(np.rint(seconds))
142
+
143
+ if s < 60:
144
+ return "{0}s".format(s)
145
+ elif s < 60 * 60:
146
+ return "{0}m {1:02}s".format(s // 60, s % 60)
147
+ elif s < 24 * 60 * 60:
148
+ return "{0}h {1:02}m {2:02}s".format(s // (60 * 60), (s // 60) % 60, s % 60)
149
+ else:
150
+ return "{0}d {1:02}h {2:02}m".format(s // (24 * 60 * 60), (s // (60 * 60)) % 24, (s // 60) % 60)
151
+
152
+
153
+ def ask_yes_no(question: str) -> bool:
154
+ """Ask the user the question until the user inputs a valid answer."""
155
+ while True:
156
+ try:
157
+ print("{0} [y/n]".format(question))
158
+ return strtobool(input().lower())
159
+ except ValueError:
160
+ pass
161
+
162
+
163
+ def tuple_product(t: Tuple) -> Any:
164
+ """Calculate the product of the tuple elements."""
165
+ result = 1
166
+
167
+ for v in t:
168
+ result *= v
169
+
170
+ return result
171
+
172
+
173
+ _str_to_ctype = {
174
+ "uint8": ctypes.c_ubyte,
175
+ "uint16": ctypes.c_uint16,
176
+ "uint32": ctypes.c_uint32,
177
+ "uint64": ctypes.c_uint64,
178
+ "int8": ctypes.c_byte,
179
+ "int16": ctypes.c_int16,
180
+ "int32": ctypes.c_int32,
181
+ "int64": ctypes.c_int64,
182
+ "float32": ctypes.c_float,
183
+ "float64": ctypes.c_double
184
+ }
185
+
186
+
187
+ def get_dtype_and_ctype(type_obj: Any) -> Tuple[np.dtype, Any]:
188
+ """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."""
189
+ type_str = None
190
+
191
+ if isinstance(type_obj, str):
192
+ type_str = type_obj
193
+ elif hasattr(type_obj, "__name__"):
194
+ type_str = type_obj.__name__
195
+ elif hasattr(type_obj, "name"):
196
+ type_str = type_obj.name
197
+ else:
198
+ raise RuntimeError("Cannot infer type name from input")
199
+
200
+ assert type_str in _str_to_ctype.keys()
201
+
202
+ my_dtype = np.dtype(type_str)
203
+ my_ctype = _str_to_ctype[type_str]
204
+
205
+ assert my_dtype.itemsize == ctypes.sizeof(my_ctype)
206
+
207
+ return my_dtype, my_ctype
208
+
209
+
210
+ def is_pickleable(obj: Any) -> bool:
211
+ try:
212
+ with io.BytesIO() as stream:
213
+ pickle.dump(obj, stream)
214
+ return True
215
+ except:
216
+ return False
217
+
218
+
219
+ # Functionality to import modules/objects by name, and call functions by name
220
+ # ------------------------------------------------------------------------------------------
221
+
222
+ def get_module_from_obj_name(obj_name: str) -> Tuple[types.ModuleType, str]:
223
+ """Searches for the underlying module behind the name to some python object.
224
+ Returns the module and the object name (original name with module part removed)."""
225
+
226
+ # allow convenience shorthands, substitute them by full names
227
+ obj_name = re.sub("^np.", "numpy.", obj_name)
228
+ obj_name = re.sub("^tf.", "tensorflow.", obj_name)
229
+
230
+ # list alternatives for (module_name, local_obj_name)
231
+ parts = obj_name.split(".")
232
+ name_pairs = [(".".join(parts[:i]), ".".join(parts[i:])) for i in range(len(parts), 0, -1)]
233
+
234
+ # try each alternative in turn
235
+ for module_name, local_obj_name in name_pairs:
236
+ try:
237
+ module = importlib.import_module(module_name) # may raise ImportError
238
+ get_obj_from_module(module, local_obj_name) # may raise AttributeError
239
+ return module, local_obj_name
240
+ except:
241
+ pass
242
+
243
+ # maybe some of the modules themselves contain errors?
244
+ for module_name, _local_obj_name in name_pairs:
245
+ try:
246
+ importlib.import_module(module_name) # may raise ImportError
247
+ except ImportError:
248
+ if not str(sys.exc_info()[1]).startswith("No module named '" + module_name + "'"):
249
+ raise
250
+
251
+ # maybe the requested attribute is missing?
252
+ for module_name, local_obj_name in name_pairs:
253
+ try:
254
+ module = importlib.import_module(module_name) # may raise ImportError
255
+ get_obj_from_module(module, local_obj_name) # may raise AttributeError
256
+ except ImportError:
257
+ pass
258
+
259
+ # we are out of luck, but we have no idea why
260
+ raise ImportError(obj_name)
261
+
262
+
263
+ def get_obj_from_module(module: types.ModuleType, obj_name: str) -> Any:
264
+ """Traverses the object name and returns the last (rightmost) python object."""
265
+ if obj_name == '':
266
+ return module
267
+ obj = module
268
+ for part in obj_name.split("."):
269
+ obj = getattr(obj, part)
270
+ return obj
271
+
272
+
273
+ def get_obj_by_name(name: str) -> Any:
274
+ """Finds the python object with the given name."""
275
+ module, obj_name = get_module_from_obj_name(name)
276
+ return get_obj_from_module(module, obj_name)
277
+
278
+
279
+ def call_func_by_name(*args, func_name: str = None, **kwargs) -> Any:
280
+ """Finds the python object with the given name and calls it as a function."""
281
+ assert func_name is not None
282
+ func_obj = get_obj_by_name(func_name)
283
+ assert callable(func_obj)
284
+ return func_obj(*args, **kwargs)
285
+
286
+
287
+ def construct_class_by_name(*args, class_name: str = None, **kwargs) -> Any:
288
+ """Finds the python class with the given name and constructs it with the given arguments."""
289
+ return call_func_by_name(*args, func_name=class_name, **kwargs)
290
+
291
+
292
+ def get_module_dir_by_obj_name(obj_name: str) -> str:
293
+ """Get the directory path of the module containing the given object name."""
294
+ module, _ = get_module_from_obj_name(obj_name)
295
+ return os.path.dirname(inspect.getfile(module))
296
+
297
+
298
+ def is_top_level_function(obj: Any) -> bool:
299
+ """Determine whether the given object is a top-level function, i.e., defined at module scope using 'def'."""
300
+ return callable(obj) and obj.__name__ in sys.modules[obj.__module__].__dict__
301
+
302
+
303
+ def get_top_level_function_name(obj: Any) -> str:
304
+ """Return the fully-qualified name of a top-level function."""
305
+ assert is_top_level_function(obj)
306
+ module = obj.__module__
307
+ if module == '__main__':
308
+ module = os.path.splitext(os.path.basename(sys.modules[module].__file__))[0]
309
+ return module + "." + obj.__name__
310
+
311
+
312
+ # File system helpers
313
+ # ------------------------------------------------------------------------------------------
314
+
315
+ def list_dir_recursively_with_ignore(dir_path: str, ignores: List[str] = None, add_base_to_relative: bool = False) -> List[Tuple[str, str]]:
316
+ """List all files recursively in a given directory while ignoring given file and directory names.
317
+ Returns list of tuples containing both absolute and relative paths."""
318
+ assert os.path.isdir(dir_path)
319
+ base_name = os.path.basename(os.path.normpath(dir_path))
320
+
321
+ if ignores is None:
322
+ ignores = []
323
+
324
+ result = []
325
+
326
+ for root, dirs, files in os.walk(dir_path, topdown=True):
327
+ for ignore_ in ignores:
328
+ dirs_to_remove = [d for d in dirs if fnmatch.fnmatch(d, ignore_)]
329
+
330
+ # dirs need to be edited in-place
331
+ for d in dirs_to_remove:
332
+ dirs.remove(d)
333
+
334
+ files = [f for f in files if not fnmatch.fnmatch(f, ignore_)]
335
+
336
+ absolute_paths = [os.path.join(root, f) for f in files]
337
+ relative_paths = [os.path.relpath(p, dir_path) for p in absolute_paths]
338
+
339
+ if add_base_to_relative:
340
+ relative_paths = [os.path.join(base_name, p) for p in relative_paths]
341
+
342
+ assert len(absolute_paths) == len(relative_paths)
343
+ result += zip(absolute_paths, relative_paths)
344
+
345
+ return result
346
+
347
+
348
+ def copy_files_and_create_dirs(files: List[Tuple[str, str]]) -> None:
349
+ """Takes in a list of tuples of (src, dst) paths and copies files.
350
+ Will create all necessary directories."""
351
+ for file in files:
352
+ target_dir_name = os.path.dirname(file[1])
353
+
354
+ # will create all intermediate-level directories
355
+ if not os.path.exists(target_dir_name):
356
+ os.makedirs(target_dir_name)
357
+
358
+ shutil.copyfile(file[0], file[1])
359
+
360
+
361
+ # URL helpers
362
+ # ------------------------------------------------------------------------------------------
363
+
364
+ def is_url(obj: Any, allow_file_urls: bool = False) -> bool:
365
+ """Determine whether the given object is a valid URL string."""
366
+ if not isinstance(obj, str) or not "://" in obj:
367
+ return False
368
+ if allow_file_urls and obj.startswith('file://'):
369
+ return True
370
+ try:
371
+ res = requests.compat.urlparse(obj)
372
+ if not res.scheme or not res.netloc or not "." in res.netloc:
373
+ return False
374
+ res = requests.compat.urlparse(requests.compat.urljoin(obj, "/"))
375
+ if not res.scheme or not res.netloc or not "." in res.netloc:
376
+ return False
377
+ except:
378
+ return False
379
+ return True
380
+
381
+
382
+ def open_url(url: str, cache_dir: str = None, num_attempts: int = 10, verbose: bool = True, return_filename: bool = False, cache: bool = True) -> Any:
383
+ """Download the given URL and return a binary-mode file object to access the data."""
384
+ assert num_attempts >= 1
385
+ assert not (return_filename and (not cache))
386
+
387
+ # Doesn't look like an URL scheme so interpret it as a local filename.
388
+ if not re.match('^[a-z]+://', url):
389
+ return url if return_filename else open(url, "rb")
390
+
391
+ # Handle file URLs. This code handles unusual file:// patterns that
392
+ # arise on Windows:
393
+ #
394
+ # file:///c:/foo.txt
395
+ #
396
+ # which would translate to a local '/c:/foo.txt' filename that's
397
+ # invalid. Drop the forward slash for such pathnames.
398
+ #
399
+ # If you touch this code path, you should test it on both Linux and
400
+ # Windows.
401
+ #
402
+ # Some internet resources suggest using urllib.request.url2pathname() but
403
+ # but that converts forward slashes to backslashes and this causes
404
+ # its own set of problems.
405
+ if url.startswith('file://'):
406
+ filename = urllib.parse.urlparse(url).path
407
+ if re.match(r'^/[a-zA-Z]:', filename):
408
+ filename = filename[1:]
409
+ return filename if return_filename else open(filename, "rb")
410
+
411
+ assert is_url(url)
412
+
413
+ # Lookup from cache.
414
+ if cache_dir is None:
415
+ cache_dir = make_cache_dir_path('downloads')
416
+
417
+ url_md5 = hashlib.md5(url.encode("utf-8")).hexdigest()
418
+ if cache:
419
+ cache_files = glob.glob(os.path.join(cache_dir, url_md5 + "_*"))
420
+ if len(cache_files) == 1:
421
+ filename = cache_files[0]
422
+ return filename if return_filename else open(filename, "rb")
423
+
424
+ # Download.
425
+ url_name = None
426
+ url_data = None
427
+ with requests.Session() as session:
428
+ if verbose:
429
+ print("Downloading %s ..." % url, end="", flush=True)
430
+ for attempts_left in reversed(range(num_attempts)):
431
+ try:
432
+ with session.get(url) as res:
433
+ res.raise_for_status()
434
+ if len(res.content) == 0:
435
+ raise IOError("No data received")
436
+
437
+ if len(res.content) < 8192:
438
+ content_str = res.content.decode("utf-8")
439
+ if "download_warning" in res.headers.get("Set-Cookie", ""):
440
+ links = [html.unescape(link) for link in content_str.split('"') if "export=download" in link]
441
+ if len(links) == 1:
442
+ url = requests.compat.urljoin(url, links[0])
443
+ raise IOError("Google Drive virus checker nag")
444
+ if "Google Drive - Quota exceeded" in content_str:
445
+ raise IOError("Google Drive download quota exceeded -- please try again later")
446
+
447
+ match = re.search(r'filename="([^"]*)"', res.headers.get("Content-Disposition", ""))
448
+ url_name = match[1] if match else url
449
+ url_data = res.content
450
+ if verbose:
451
+ print(" done")
452
+ break
453
+ except KeyboardInterrupt:
454
+ raise
455
+ except:
456
+ if not attempts_left:
457
+ if verbose:
458
+ print(" failed")
459
+ raise
460
+ if verbose:
461
+ print(".", end="", flush=True)
462
+
463
+ # Save to cache.
464
+ if cache:
465
+ safe_name = re.sub(r"[^0-9a-zA-Z-._]", "_", url_name)
466
+ cache_file = os.path.join(cache_dir, url_md5 + "_" + safe_name)
467
+ temp_file = os.path.join(cache_dir, "tmp_" + uuid.uuid4().hex + "_" + url_md5 + "_" + safe_name)
468
+ os.makedirs(cache_dir, exist_ok=True)
469
+ with open(temp_file, "wb") as f:
470
+ f.write(url_data)
471
+ os.replace(temp_file, cache_file) # atomic
472
+ if return_filename:
473
+ return cache_file
474
+
475
+ # Return data as file object.
476
+ assert not return_filename
477
+ return io.BytesIO(url_data)
edit.py ADDED
@@ -0,0 +1,84 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import wandb
2
+ import click
3
+ import os
4
+ import sys
5
+ import pickle
6
+ import numpy as np
7
+ from PIL import Image
8
+ import glob
9
+ import torch
10
+ from configs import paths_config, hyperparameters, global_config
11
+ from IPython.display import display
12
+ import matplotlib.pyplot as plt
13
+ from scripts.latent_editor_wrapper import LatentEditorWrapper
14
+
15
+
16
+ image_dir_name = '/home/sayantan/processed_images'
17
+ use_multi_id_training = False
18
+ global_config.device = 'cuda'
19
+ paths_config.e4e = '/home/sayantan/PTI/pretrained_models/e4e_ffhq_encode.pt'
20
+ paths_config.input_data_id = image_dir_name
21
+ paths_config.input_data_path = f'{image_dir_name}'
22
+ paths_config.stylegan2_ada_ffhq = '/home/sayantan/PTI/pretrained_models/ffhq.pkl'
23
+ paths_config.checkpoints_dir = '/home/sayantan/PTI/'
24
+ paths_config.style_clip_pretrained_mappers = '/home/sayantan/PTI/pretrained_models'
25
+ hyperparameters.use_locality_regularization = False
26
+ hyperparameters.lpips_type = 'squeeze'
27
+
28
+ model_id = "MYJJDFVGATAT"
29
+
30
+
31
+
32
+ def display_alongside_source_image(images):
33
+ res = np.concatenate([np.array(image) for image in images], axis=1)
34
+ return Image.fromarray(res)
35
+
36
+ def load_generators(model_id, image_name):
37
+ with open(paths_config.stylegan2_ada_ffhq, 'rb') as f:
38
+ old_G = pickle.load(f)['G_ema'].cuda()
39
+
40
+ with open(f'{paths_config.checkpoints_dir}/model_{model_id}_{image_name}.pt', 'rb') as f_new:
41
+ new_G = torch.load(f_new).cuda()
42
+
43
+ return old_G, new_G
44
+
45
+ def plot_syn_images(syn_images,text):
46
+ for img in syn_images:
47
+ img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).detach().cpu().numpy()[0]
48
+ plt.axis('off')
49
+ resized_image = Image.fromarray(img,mode='RGB').resize((256,256))
50
+ display(resized_image)
51
+ #wandb.log({text: [wandb.Image(resized_image, caption="Label")]})
52
+ del img
53
+ del resized_image
54
+ torch.cuda.empty_cache()
55
+
56
+ def syn_images_wandb(img):
57
+ img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).detach().cpu().numpy()[0]
58
+ plt.axis('off')
59
+ resized_image = Image.fromarray(img,mode='RGB').resize((256,256))
60
+ return resized_image
61
+
62
+
63
+ def edit(image_name):
64
+ generator_type = paths_config.multi_id_model_type if use_multi_id_training else image_name
65
+ old_G, new_G = load_generators(model_id, generator_type)
66
+ w_path_dir = f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}'
67
+
68
+ embedding_dir = f'{w_path_dir}/{paths_config.pti_results_keyword}/{image_name}'
69
+ w_pivot = torch.load(f'{embedding_dir}/0.pt')
70
+
71
+ old_image = old_G.synthesis(w_pivot, noise_mode='const', force_fp32 = True)
72
+ new_image = new_G.synthesis(w_pivot, noise_mode='const', force_fp32 = True)
73
+
74
+ latent_editor = LatentEditorWrapper()
75
+ latents_after_edit = latent_editor.get_single_interface_gan_edits(w_pivot, [i for i in range(-5,5)])
76
+
77
+ for direction, factor_and_edit in latents_after_edit.items():
78
+ for editkey in factor_and_edit.keys():
79
+ new_image = new_G.synthesis(factor_and_edit[editkey], noise_mode='const', force_fp32 = True)
80
+ image_pil = syn_images_wandb(new_image).save(f"/home/sayantan/PTI/{direction}/{editkey}/{image_name}.jpg")
81
+
82
+ if __name__ == '__main__':
83
+ for image_name in [f.split(".")[0].split("_")[2] for f in sorted(glob.glob("*.pt"))]:
84
+ edit(image_name)
editings/__pycache__/ganspace.cpython-39.pyc ADDED
Binary file (878 Bytes). View file
 
editings/__pycache__/latent_editor.cpython-39.pyc ADDED
Binary file (994 Bytes). View file
 
editings/ganspace.py ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+
4
+ def edit(latents, pca, edit_directions):
5
+ edit_latents = []
6
+ for latent in latents:
7
+ for pca_idx, start, end, strength in edit_directions:
8
+ delta = get_delta(pca, latent, pca_idx, strength)
9
+ delta_padded = torch.zeros(latent.shape).to('cuda')
10
+ delta_padded[start:end] += delta.repeat(end - start, 1)
11
+ edit_latents.append(latent + delta_padded)
12
+ return torch.stack(edit_latents)
13
+
14
+
15
+ def get_delta(pca, latent, idx, strength):
16
+ w_centered = latent - pca['mean'].to('cuda')
17
+ lat_comp = pca['comp'].to('cuda')
18
+ lat_std = pca['std'].to('cuda')
19
+ w_coord = torch.sum(w_centered[0].reshape(-1)*lat_comp[idx].reshape(-1)) / lat_std[idx]
20
+ delta = (strength - w_coord)*lat_comp[idx]*lat_std[idx]
21
+ return delta
editings/ganspace_pca/ffhq_pca.pt ADDED
Binary file (168 kB). View file
 
editings/interfacegan_directions/age.pt ADDED
Binary file (2.81 kB). View file
 
editings/interfacegan_directions/rotation.pt ADDED
Binary file (2.81 kB). View file
 
editings/interfacegan_directions/smile.pt ADDED
Binary file (2.81 kB). View file
 
editings/latent_editor.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+ from configs import paths_config
4
+ from editings import ganspace
5
+ from utils.data_utils import tensor2im
6
+
7
+
8
+ class LatentEditor(object):
9
+
10
+ def apply_ganspace(self, latent, ganspace_pca, edit_directions):
11
+ edit_latents = ganspace.edit(latent, ganspace_pca, edit_directions)
12
+ return edit_latents
13
+
14
+ def apply_interfacegan(self, latent, direction, factor=1, factor_range=None):
15
+ edit_latents = []
16
+ if factor_range is not None: # Apply a range of editing factors. for example, (-5, 5)
17
+ for f in range(*factor_range):
18
+ edit_latent = latent + f * direction
19
+ edit_latents.append(edit_latent)
20
+ edit_latents = torch.cat(edit_latents)
21
+ else:
22
+ edit_latents = latent + factor * direction
23
+ return edit_latents
evaluation/experiment_setting_creator.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import glob
2
+ import os
3
+ from configs import global_config, paths_config, hyperparameters
4
+ from scripts.latent_creators.sg2_plus_latent_creator import SG2PlusLatentCreator
5
+ from scripts.latent_creators.e4e_latent_creator import E4ELatentCreator
6
+ from scripts.run_pti import run_PTI
7
+ import pickle
8
+ import torch
9
+ from utils.models_utils import toogle_grad, load_old_G
10
+
11
+
12
+ class ExperimentRunner:
13
+
14
+ def __init__(self, run_id=''):
15
+ self.images_paths = glob.glob(f'{paths_config.input_data_path}/*')
16
+ self.target_paths = glob.glob(f'{paths_config.input_data_path}/*')
17
+ self.run_id = run_id
18
+ self.sampled_ws = None
19
+
20
+ self.old_G = load_old_G()
21
+
22
+ toogle_grad(self.old_G, False)
23
+
24
+ def run_experiment(self, run_pt, create_other_latents, use_multi_id_training, use_wandb=False):
25
+ if run_pt:
26
+ self.run_id = run_PTI(self.run_id, use_wandb=use_wandb, use_multi_id_training=use_multi_id_training)
27
+ if create_other_latents:
28
+ sg2_plus_latent_creator = SG2PlusLatentCreator(use_wandb=use_wandb)
29
+ sg2_plus_latent_creator.create_latents()
30
+ e4e_latent_creator = E4ELatentCreator(use_wandb=use_wandb)
31
+ e4e_latent_creator.create_latents()
32
+
33
+ torch.cuda.empty_cache()
34
+
35
+ return self.run_id
36
+
37
+
38
+ if __name__ == '__main__':
39
+ os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
40
+ os.environ['CUDA_VISIBLE_DEVICES'] = global_config.cuda_visible_devices
41
+
42
+ runner = ExperimentRunner()
43
+ runner.run_experiment(True, False, False)
evaluation/qualitative_edit_comparison.py ADDED
@@ -0,0 +1,156 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from random import choice
3
+ from string import ascii_uppercase
4
+ from PIL import Image
5
+ from tqdm import tqdm
6
+ from scripts.latent_editor_wrapper import LatentEditorWrapper
7
+ from evaluation.experiment_setting_creator import ExperimentRunner
8
+ import torch
9
+ from configs import paths_config, hyperparameters, evaluation_config
10
+ from utils.log_utils import save_concat_image, save_single_image
11
+ from utils.models_utils import load_tuned_G
12
+
13
+
14
+ class EditComparison:
15
+
16
+ def __init__(self, save_single_images, save_concatenated_images, run_id):
17
+
18
+ self.run_id = run_id
19
+ self.experiment_creator = ExperimentRunner(run_id)
20
+ self.save_single_images = save_single_images
21
+ self.save_concatenated_images = save_concatenated_images
22
+ self.latent_editor = LatentEditorWrapper()
23
+
24
+ def save_reconstruction_images(self, image_latents, new_inv_image_latent, new_G, target_image):
25
+ if self.save_concatenated_images:
26
+ save_concat_image(self.concat_base_dir, image_latents, new_inv_image_latent, new_G,
27
+ self.experiment_creator.old_G,
28
+ 'rec',
29
+ target_image)
30
+
31
+ if self.save_single_images:
32
+ save_single_image(self.single_base_dir, new_inv_image_latent, new_G, 'rec')
33
+ target_image.save(f'{self.single_base_dir}/Original.jpg')
34
+
35
+ def create_output_dirs(self, full_image_name):
36
+ output_base_dir_path = f'{paths_config.experiments_output_dir}/{paths_config.input_data_id}/{self.run_id}/{full_image_name}'
37
+ os.makedirs(output_base_dir_path, exist_ok=True)
38
+
39
+ self.concat_base_dir = f'{output_base_dir_path}/concat_images'
40
+ self.single_base_dir = f'{output_base_dir_path}/single_images'
41
+
42
+ os.makedirs(self.concat_base_dir, exist_ok=True)
43
+ os.makedirs(self.single_base_dir, exist_ok=True)
44
+
45
+ def get_image_latent_codes(self, image_name):
46
+ image_latents = []
47
+ for method in evaluation_config.evaluated_methods:
48
+ if method == 'SG2':
49
+ image_latents.append(torch.load(
50
+ f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}/'
51
+ f'{paths_config.pti_results_keyword}/{image_name}/0.pt'))
52
+ else:
53
+ image_latents.append(torch.load(
54
+ f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}/{method}/{image_name}/0.pt'))
55
+ new_inv_image_latent = torch.load(
56
+ f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}/{paths_config.pti_results_keyword}/{image_name}/0.pt')
57
+
58
+ return image_latents, new_inv_image_latent
59
+
60
+ def save_interfacegan_edits(self, image_latents, new_inv_image_latent, interfacegan_factors, new_G, target_image):
61
+ new_w_inv_edits = self.latent_editor.get_single_interface_gan_edits(new_inv_image_latent,
62
+ interfacegan_factors)
63
+
64
+ inv_edits = []
65
+ for latent in image_latents:
66
+ inv_edits.append(self.latent_editor.get_single_interface_gan_edits(latent, interfacegan_factors))
67
+
68
+ for direction, edits in new_w_inv_edits.items():
69
+ for factor, edit_tensor in edits.items():
70
+ if self.save_concatenated_images:
71
+ save_concat_image(self.concat_base_dir, [edits[direction][factor] for edits in inv_edits],
72
+ new_w_inv_edits[direction][factor],
73
+ new_G,
74
+ self.experiment_creator.old_G,
75
+ f'{direction}_{factor}', target_image)
76
+ if self.save_single_images:
77
+ save_single_image(self.single_base_dir, new_w_inv_edits[direction][factor], new_G,
78
+ f'{direction}_{factor}')
79
+
80
+ def save_ganspace_edits(self, image_latents, new_inv_image_latent, factors, new_G, target_image):
81
+ new_w_inv_edits = self.latent_editor.get_single_ganspace_edits(new_inv_image_latent, factors)
82
+ inv_edits = []
83
+ for latent in image_latents:
84
+ inv_edits.append(self.latent_editor.get_single_ganspace_edits(latent, factors))
85
+
86
+ for idx in range(len(new_w_inv_edits)):
87
+ if self.save_concatenated_images:
88
+ save_concat_image(self.concat_base_dir, [edit[idx] for edit in inv_edits], new_w_inv_edits[idx],
89
+ new_G,
90
+ self.experiment_creator.old_G,
91
+ f'ganspace_{idx}', target_image)
92
+ if self.save_single_images:
93
+ save_single_image(self.single_base_dir, new_w_inv_edits[idx], new_G,
94
+ f'ganspace_{idx}')
95
+
96
+ def run_experiment(self, run_pt, create_other_latents, use_multi_id_training, use_wandb=False):
97
+ images_counter = 0
98
+ new_G = None
99
+ interfacegan_factors = [val / 2 for val in range(-6, 7) if val != 0]
100
+ ganspace_factors = range(-20, 25, 5)
101
+ self.experiment_creator.run_experiment(run_pt, create_other_latents, use_multi_id_training, use_wandb)
102
+
103
+ if use_multi_id_training:
104
+ new_G = load_tuned_G(self.run_id, paths_config.multi_id_model_type)
105
+
106
+ for idx, image_path in tqdm(enumerate(self.experiment_creator.images_paths),
107
+ total=len(self.experiment_creator.images_paths)):
108
+
109
+ if images_counter >= hyperparameters.max_images_to_invert:
110
+ break
111
+
112
+ image_name = image_path.split('.')[0].split('/')[-1]
113
+ target_image = Image.open(self.experiment_creator.target_paths[idx])
114
+
115
+ if not use_multi_id_training:
116
+ new_G = load_tuned_G(self.run_id, image_name)
117
+
118
+ image_latents, new_inv_image_latent = self.get_image_latent_codes(image_name)
119
+
120
+ self.create_output_dirs(image_name)
121
+
122
+ self.save_reconstruction_images(image_latents, new_inv_image_latent, new_G, target_image)
123
+
124
+ self.save_interfacegan_edits(image_latents, new_inv_image_latent, interfacegan_factors, new_G, target_image)
125
+
126
+ self.save_ganspace_edits(image_latents, new_inv_image_latent, ganspace_factors, new_G, target_image)
127
+
128
+ target_image.close()
129
+ torch.cuda.empty_cache()
130
+ images_counter += 1
131
+
132
+
133
+ def run_pti_and_full_edit(iid):
134
+ evaluation_config.evaluated_methods = ['SG2Plus', 'e4e', 'SG2']
135
+ edit_figure_creator = EditComparison(save_single_images=True, save_concatenated_images=True,
136
+ run_id=f'{paths_config.input_data_id}_pti_full_edit_{iid}')
137
+ edit_figure_creator.run_experiment(True, True, use_multi_id_training=False, use_wandb=False)
138
+
139
+
140
+ def pti_no_comparison(iid):
141
+ evaluation_config.evaluated_methods = []
142
+ edit_figure_creator = EditComparison(save_single_images=True, save_concatenated_images=True,
143
+ run_id=f'{paths_config.input_data_id}_pti_no_comparison_{iid}')
144
+ edit_figure_creator.run_experiment(True, False, use_multi_id_training=False, use_wandb=False)
145
+
146
+
147
+ def edits_for_existed_experiment(run_id):
148
+ evaluation_config.evaluated_methods = ['SG2Plus', 'e4e', 'SG2']
149
+ edit_figure_creator = EditComparison(save_single_images=True, save_concatenated_images=True,
150
+ run_id=run_id)
151
+ edit_figure_creator.run_experiment(False, True, use_multi_id_training=False, use_wandb=False)
152
+
153
+
154
+ if __name__ == '__main__':
155
+ iid = ''.join(choice(ascii_uppercase) for i in range(7))
156
+ pti_no_comparison(iid)
makedirs.py ADDED
@@ -0,0 +1,84 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import wandb
2
+ import click
3
+ import os
4
+ import sys
5
+ import pickle
6
+ import numpy as np
7
+ from PIL import Image
8
+ import torch
9
+ from configs import paths_config, hyperparameters, global_config
10
+ from IPython.display import display
11
+ import matplotlib.pyplot as plt
12
+ from scripts.latent_editor_wrapper import LatentEditorWrapper
13
+
14
+
15
+ image_dir_name = '/home/sayantan/processed_images'
16
+ use_multi_id_training = False
17
+ global_config.device = 'cuda'
18
+ paths_config.e4e = '/home/sayantan/PTI/pretrained_models/e4e_ffhq_encode.pt'
19
+ paths_config.input_data_id = image_dir_name
20
+ paths_config.input_data_path = f'{image_dir_name}'
21
+ paths_config.stylegan2_ada_ffhq = '/home/sayantan/PTI/pretrained_models/ffhq.pkl'
22
+ paths_config.checkpoints_dir = '/home/sayantan/PTI/'
23
+ paths_config.style_clip_pretrained_mappers = '/home/sayantan/PTI/pretrained_models'
24
+ hyperparameters.use_locality_regularization = False
25
+ hyperparameters.lpips_type = 'squeeze'
26
+
27
+ model_id = "MYJJDFVGATAT"
28
+
29
+
30
+
31
+ def display_alongside_source_image(images):
32
+ res = np.concatenate([np.array(image) for image in images], axis=1)
33
+ return Image.fromarray(res)
34
+
35
+ def load_generators(model_id, image_name):
36
+ with open(paths_config.stylegan2_ada_ffhq, 'rb') as f:
37
+ old_G = pickle.load(f)['G_ema'].cuda()
38
+
39
+ with open(f'{paths_config.checkpoints_dir}/model_{model_id}_{image_name}.pt', 'rb') as f_new:
40
+ new_G = torch.load(f_new).cuda()
41
+
42
+ return old_G, new_G
43
+
44
+ def plot_syn_images(syn_images,text):
45
+ for img in syn_images:
46
+ img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).detach().cpu().numpy()[0]
47
+ plt.axis('off')
48
+ resized_image = Image.fromarray(img,mode='RGB').resize((256,256))
49
+ display(resized_image)
50
+ #wandb.log({text: [wandb.Image(resized_image, caption="Label")]})
51
+ del img
52
+ del resized_image
53
+ torch.cuda.empty_cache()
54
+
55
+ def syn_images_wandb(img):
56
+ img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).detach().cpu().numpy()[0]
57
+ plt.axis('off')
58
+ resized_image = Image.fromarray(img,mode='RGB').resize((256,256))
59
+ return resized_image
60
+
61
+ @click.command()
62
+ @click.pass_context
63
+ @click.option('--image_name', prompt='image name', help='The name for image')
64
+
65
+ def makedir(ctx: click.Context,image_name):
66
+ generator_type = paths_config.multi_id_model_type if use_multi_id_training else image_name
67
+ old_G, new_G = load_generators(model_id, generator_type)
68
+ w_path_dir = f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}'
69
+
70
+ embedding_dir = f'{w_path_dir}/{paths_config.pti_results_keyword}/{image_name}'
71
+ w_pivot = torch.load(f'{embedding_dir}/0.pt')
72
+
73
+ old_image = old_G.synthesis(w_pivot, noise_mode='const', force_fp32 = True)
74
+ new_image = new_G.synthesis(w_pivot, noise_mode='const', force_fp32 = True)
75
+
76
+ latent_editor = LatentEditorWrapper()
77
+ latents_after_edit = latent_editor.get_single_interface_gan_edits(w_pivot, [i for i in range(-5,5)])
78
+
79
+ for direction, factor_and_edit in latents_after_edit.items():
80
+ for editkey in factor_and_edit.keys():
81
+ os.makedirs(f"/home/sayantan/PTI/{direction}/{editkey}")
82
+
83
+ if __name__ == '__main__':
84
+ makedir()
models/StyleCLIP/__init__.py ADDED
File without changes
models/StyleCLIP/criteria/__init__.py ADDED
File without changes
models/StyleCLIP/criteria/clip_loss.py ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import torch
3
+ import clip
4
+
5
+
6
+ class CLIPLoss(torch.nn.Module):
7
+
8
+ def __init__(self, opts):
9
+ super(CLIPLoss, self).__init__()
10
+ self.model, self.preprocess = clip.load("ViT-B/32", device="cuda")
11
+ self.upsample = torch.nn.Upsample(scale_factor=7)
12
+ self.avg_pool = torch.nn.AvgPool2d(kernel_size=opts.stylegan_size // 32)
13
+
14
+ def forward(self, image, text):
15
+ image = self.avg_pool(self.upsample(image))
16
+ similarity = 1 - self.model(image, text)[0] / 100
17
+ return similarity
models/StyleCLIP/criteria/id_loss.py ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from torch import nn
3
+
4
+ from models.facial_recognition.model_irse import Backbone
5
+
6
+
7
+ class IDLoss(nn.Module):
8
+ def __init__(self, opts):
9
+ super(IDLoss, self).__init__()
10
+ print('Loading ResNet ArcFace')
11
+ self.facenet = Backbone(input_size=112, num_layers=50, drop_ratio=0.6, mode='ir_se')
12
+ self.facenet.load_state_dict(torch.load(opts.ir_se50_weights))
13
+ self.pool = torch.nn.AdaptiveAvgPool2d((256, 256))
14
+ self.face_pool = torch.nn.AdaptiveAvgPool2d((112, 112))
15
+ self.facenet.eval()
16
+ self.opts = opts
17
+
18
+ def extract_feats(self, x):
19
+ if x.shape[2] != 256:
20
+ x = self.pool(x)
21
+ x = x[:, :, 35:223, 32:220] # Crop interesting region
22
+ x = self.face_pool(x)
23
+ x_feats = self.facenet(x)
24
+ return x_feats
25
+
26
+ def forward(self, y_hat, y):
27
+ n_samples = y.shape[0]
28
+ y_feats = self.extract_feats(y) # Otherwise use the feature from there
29
+ y_hat_feats = self.extract_feats(y_hat)
30
+ y_feats = y_feats.detach()
31
+ loss = 0
32
+ sim_improvement = 0
33
+ count = 0
34
+ for i in range(n_samples):
35
+ diff_target = y_hat_feats[i].dot(y_feats[i])
36
+ loss += 1 - diff_target
37
+ count += 1
38
+
39
+ return loss / count, sim_improvement / count
models/StyleCLIP/global_directions/GUI.py ADDED
@@ -0,0 +1,103 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+
3
+ from tkinter import Tk,Frame ,Label,Button,messagebox,Canvas,Text,Scale
4
+ from tkinter import HORIZONTAL
5
+
6
+ class View():
7
+ def __init__(self,master):
8
+
9
+ self.width=600
10
+ self.height=600
11
+
12
+
13
+ self.root=master
14
+ self.root.geometry("600x600")
15
+
16
+ self.left_frame=Frame(self.root,width=600)
17
+ self.left_frame.pack_propagate(0)
18
+ self.left_frame.pack(fill='both', side='left', expand='True')
19
+
20
+ self.retrieval_frame=Frame(self.root,bg='snow3')
21
+ self.retrieval_frame.pack_propagate(0)
22
+ self.retrieval_frame.pack(fill='both', side='right', expand='True')
23
+
24
+ self.bg_frame=Frame(self.left_frame,bg='snow3',height=600,width=600)
25
+ self.bg_frame.pack_propagate(0)
26
+ self.bg_frame.pack(fill='both', side='top', expand='True')
27
+
28
+ self.command_frame=Frame(self.left_frame,bg='snow3')
29
+ self.command_frame.pack_propagate(0)
30
+ self.command_frame.pack(fill='both', side='bottom', expand='True')
31
+ # self.command_frame.grid(row=1, column=0,padx=0, pady=0)
32
+
33
+ self.bg=Canvas(self.bg_frame,width=self.width,height=self.height, bg='gray')
34
+ self.bg.place(relx=0.5, rely=0.5, anchor='center')
35
+
36
+ self.mani=Canvas(self.retrieval_frame,width=1024,height=1024, bg='gray')
37
+ self.mani.grid(row=0, column=0,padx=0, pady=42)
38
+
39
+ self.SetCommand()
40
+
41
+
42
+
43
+
44
+ def run(self):
45
+ self.root.mainloop()
46
+
47
+ def helloCallBack(self):
48
+ category=self.set_category.get()
49
+ messagebox.showinfo( "Hello Python",category)
50
+
51
+ def SetCommand(self):
52
+
53
+ tmp = Label(self.command_frame, text="neutral", width=10 ,bg='snow3')
54
+ tmp.grid(row=1, column=0,padx=10, pady=10)
55
+
56
+ tmp = Label(self.command_frame, text="a photo of a", width=10 ,bg='snow3')
57
+ tmp.grid(row=1, column=1,padx=10, pady=10)
58
+
59
+ self.neutral = Text ( self.command_frame, height=2, width=30)
60
+ self.neutral.grid(row=1, column=2,padx=10, pady=10)
61
+
62
+
63
+ tmp = Label(self.command_frame, text="target", width=10 ,bg='snow3')
64
+ tmp.grid(row=2, column=0,padx=10, pady=10)
65
+
66
+ tmp = Label(self.command_frame, text="a photo of a", width=10 ,bg='snow3')
67
+ tmp.grid(row=2, column=1,padx=10, pady=10)
68
+
69
+ self.target = Text ( self.command_frame, height=2, width=30)
70
+ self.target.grid(row=2, column=2,padx=10, pady=10)
71
+
72
+ tmp = Label(self.command_frame, text="strength", width=10 ,bg='snow3')
73
+ tmp.grid(row=3, column=0,padx=10, pady=10)
74
+
75
+ self.alpha = Scale(self.command_frame, from_=-15, to=25, orient=HORIZONTAL,bg='snow3', length=250,resolution=0.01)
76
+ self.alpha.grid(row=3, column=2,padx=10, pady=10)
77
+
78
+
79
+ tmp = Label(self.command_frame, text="disentangle", width=10 ,bg='snow3')
80
+ tmp.grid(row=4, column=0,padx=10, pady=10)
81
+
82
+ self.beta = Scale(self.command_frame, from_=0.08, to=0.4, orient=HORIZONTAL,bg='snow3', length=250,resolution=0.001)
83
+ self.beta.grid(row=4, column=2,padx=10, pady=10)
84
+
85
+ self.reset = Button(self.command_frame, text='Reset')
86
+ self.reset.grid(row=5, column=1,padx=10, pady=10)
87
+
88
+
89
+ self.set_init = Button(self.command_frame, text='Accept')
90
+ self.set_init.grid(row=5, column=2,padx=10, pady=10)
91
+
92
+ #%%
93
+ if __name__ == "__main__":
94
+ master=Tk()
95
+ self=View(master)
96
+ self.run()
97
+
98
+
99
+
100
+
101
+
102
+
103
+
models/StyleCLIP/global_directions/GenerateImg.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import os
3
+ import numpy as np
4
+ import argparse
5
+ from manipulate import Manipulator
6
+
7
+ from PIL import Image
8
+ #%%
9
+
10
+ if __name__ == "__main__":
11
+ parser = argparse.ArgumentParser(description='Process some integers.')
12
+
13
+ parser.add_argument('--dataset_name',type=str,default='ffhq',
14
+ help='name of dataset, for example, ffhq')
15
+
16
+ args = parser.parse_args()
17
+ dataset_name=args.dataset_name
18
+
19
+ if not os.path.isdir('./data/'+dataset_name):
20
+ os.system('mkdir ./data/'+dataset_name)
21
+ #%%
22
+ M=Manipulator(dataset_name=dataset_name)
23
+ np.set_printoptions(suppress=True)
24
+ print(M.dataset_name)
25
+ #%%
26
+
27
+ M.img_index=0
28
+ M.num_images=50
29
+ M.alpha=[0]
30
+ M.step=1
31
+ lindex,bname=0,0
32
+
33
+ M.manipulate_layers=[lindex]
34
+ codes,out=M.EditOneC(bname)
35
+ #%%
36
+
37
+ for i in range(len(out)):
38
+ img=out[i,0]
39
+ img=Image.fromarray(img)
40
+ img.save('./data/'+dataset_name+'/'+str(i)+'.jpg')
41
+ #%%
42
+ w=np.load('./npy/'+dataset_name+'/W.npy')
43
+
44
+ tmp=w[:M.num_images]
45
+ tmp=tmp[:,None,:]
46
+ tmp=np.tile(tmp,(1,M.Gs.components.synthesis.input_shape[1],1))
47
+
48
+ np.save('./data/'+dataset_name+'/w_plus.npy',tmp)
49
+
50
+
models/StyleCLIP/global_directions/GetCode.py ADDED
@@ -0,0 +1,232 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+
3
+
4
+ import os
5
+ import pickle
6
+ import numpy as np
7
+ from dnnlib import tflib
8
+ import tensorflow as tf
9
+
10
+ import argparse
11
+
12
+ def LoadModel(dataset_name):
13
+ # Initialize TensorFlow.
14
+ tflib.init_tf()
15
+ model_path='./model/'
16
+ model_name=dataset_name+'.pkl'
17
+
18
+ tmp=os.path.join(model_path,model_name)
19
+ with open(tmp, 'rb') as f:
20
+ _, _, Gs = pickle.load(f)
21
+ return Gs
22
+
23
+ def lerp(a,b,t):
24
+ return a + (b - a) * t
25
+
26
+ #stylegan-ada
27
+ def SelectName(layer_name,suffix):
28
+ if suffix==None:
29
+ tmp1='add:0' in layer_name
30
+ tmp2='shape=(?,' in layer_name
31
+ tmp4='G_synthesis_1' in layer_name
32
+ tmp= tmp1 and tmp2 and tmp4
33
+ else:
34
+ tmp1=('/Conv0_up'+suffix) in layer_name
35
+ tmp2=('/Conv1'+suffix) in layer_name
36
+ tmp3=('4x4/Conv'+suffix) in layer_name
37
+ tmp4='G_synthesis_1' in layer_name
38
+ tmp5=('/ToRGB'+suffix) in layer_name
39
+ tmp= (tmp1 or tmp2 or tmp3 or tmp5) and tmp4
40
+ return tmp
41
+
42
+
43
+ def GetSNames(suffix):
44
+ #get style tensor name
45
+ with tf.Session() as sess:
46
+ op = sess.graph.get_operations()
47
+ layers=[m.values() for m in op]
48
+
49
+
50
+ select_layers=[]
51
+ for layer in layers:
52
+ layer_name=str(layer)
53
+ if SelectName(layer_name,suffix):
54
+ select_layers.append(layer[0])
55
+ return select_layers
56
+
57
+ def SelectName2(layer_name):
58
+ tmp1='mod_bias' in layer_name
59
+ tmp2='mod_weight' in layer_name
60
+ tmp3='ToRGB' in layer_name
61
+
62
+ tmp= (tmp1 or tmp2) and (not tmp3)
63
+ return tmp
64
+
65
+ def GetKName(Gs):
66
+
67
+ layers=[var for name, var in Gs.components.synthesis.vars.items()]
68
+
69
+ select_layers=[]
70
+ for layer in layers:
71
+ layer_name=str(layer)
72
+ if SelectName2(layer_name):
73
+ select_layers.append(layer)
74
+ return select_layers
75
+
76
+ def GetCode(Gs,random_state,num_img,num_once,dataset_name):
77
+ rnd = np.random.RandomState(random_state) #5
78
+
79
+ truncation_psi=0.7
80
+ truncation_cutoff=8
81
+
82
+ dlatent_avg=Gs.get_var('dlatent_avg')
83
+
84
+ dlatents=np.zeros((num_img,512),dtype='float32')
85
+ for i in range(int(num_img/num_once)):
86
+ src_latents = rnd.randn(num_once, Gs.input_shape[1])
87
+ src_dlatents = Gs.components.mapping.run(src_latents, None) # [seed, layer, component]
88
+
89
+ # Apply truncation trick.
90
+ if truncation_psi is not None and truncation_cutoff is not None:
91
+ layer_idx = np.arange(src_dlatents.shape[1])[np.newaxis, :, np.newaxis]
92
+ ones = np.ones(layer_idx.shape, dtype=np.float32)
93
+ coefs = np.where(layer_idx < truncation_cutoff, truncation_psi * ones, ones)
94
+ src_dlatents_np=lerp(dlatent_avg, src_dlatents, coefs)
95
+ src_dlatents=src_dlatents_np[:,0,:].astype('float32')
96
+ dlatents[(i*num_once):((i+1)*num_once),:]=src_dlatents
97
+ print('get all z and w')
98
+
99
+ tmp='./npy/'+dataset_name+'/W'
100
+ np.save(tmp,dlatents)
101
+
102
+
103
+ def GetImg(Gs,num_img,num_once,dataset_name,save_name='images'):
104
+ print('Generate Image')
105
+ tmp='./npy/'+dataset_name+'/W.npy'
106
+ dlatents=np.load(tmp)
107
+ fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
108
+
109
+ all_images=[]
110
+ for i in range(int(num_img/num_once)):
111
+ print(i)
112
+ images=[]
113
+ for k in range(num_once):
114
+ tmp=dlatents[i*num_once+k]
115
+ tmp=tmp[None,None,:]
116
+ tmp=np.tile(tmp,(1,Gs.components.synthesis.input_shape[1],1))
117
+ image2= Gs.components.synthesis.run(tmp, randomize_noise=False, output_transform=fmt)
118
+ images.append(image2)
119
+
120
+ images=np.concatenate(images)
121
+
122
+ all_images.append(images)
123
+
124
+ all_images=np.concatenate(all_images)
125
+
126
+ tmp='./npy/'+dataset_name+'/'+save_name
127
+ np.save(tmp,all_images)
128
+
129
+ def GetS(dataset_name,num_img):
130
+ print('Generate S')
131
+ tmp='./npy/'+dataset_name+'/W.npy'
132
+ dlatents=np.load(tmp)[:num_img]
133
+
134
+ with tf.Session() as sess:
135
+ init = tf.global_variables_initializer()
136
+ sess.run(init)
137
+
138
+ Gs=LoadModel(dataset_name)
139
+ Gs.print_layers() #for ada
140
+ select_layers1=GetSNames(suffix=None) #None,'/mul_1:0','/mod_weight/read:0','/MatMul:0'
141
+ dlatents=dlatents[:,None,:]
142
+ dlatents=np.tile(dlatents,(1,Gs.components.synthesis.input_shape[1],1))
143
+
144
+ all_s = sess.run(
145
+ select_layers1,
146
+ feed_dict={'G_synthesis_1/dlatents_in:0': dlatents})
147
+
148
+ layer_names=[layer.name for layer in select_layers1]
149
+ save_tmp=[layer_names,all_s]
150
+ return save_tmp
151
+
152
+
153
+
154
+
155
+ def convert_images_to_uint8(images, drange=[-1,1], nchw_to_nhwc=False):
156
+ """Convert a minibatch of images from float32 to uint8 with configurable dynamic range.
157
+ Can be used as an output transformation for Network.run().
158
+ """
159
+ if nchw_to_nhwc:
160
+ images = np.transpose(images, [0, 2, 3, 1])
161
+
162
+ scale = 255 / (drange[1] - drange[0])
163
+ images = images * scale + (0.5 - drange[0] * scale)
164
+
165
+ np.clip(images, 0, 255, out=images)
166
+ images=images.astype('uint8')
167
+ return images
168
+
169
+
170
+ def GetCodeMS(dlatents):
171
+ m=[]
172
+ std=[]
173
+ for i in range(len(dlatents)):
174
+ tmp= dlatents[i]
175
+ tmp_mean=tmp.mean(axis=0)
176
+ tmp_std=tmp.std(axis=0)
177
+ m.append(tmp_mean)
178
+ std.append(tmp_std)
179
+ return m,std
180
+
181
+
182
+
183
+ #%%
184
+ if __name__ == "__main__":
185
+
186
+
187
+ parser = argparse.ArgumentParser(description='Process some integers.')
188
+
189
+ parser.add_argument('--dataset_name',type=str,default='ffhq',
190
+ help='name of dataset, for example, ffhq')
191
+ parser.add_argument('--code_type',choices=['w','s','s_mean_std'],default='w')
192
+
193
+ args = parser.parse_args()
194
+ random_state=5
195
+ num_img=100_000
196
+ num_once=1_000
197
+ dataset_name=args.dataset_name
198
+
199
+ if not os.path.isfile('./model/'+dataset_name+'.pkl'):
200
+ url='https://nvlabs-fi-cdn.nvidia.com/stylegan2/networks/'
201
+ name='stylegan2-'+dataset_name+'-config-f.pkl'
202
+ os.system('wget ' +url+name + ' -P ./model/')
203
+ os.system('mv ./model/'+name+' ./model/'+dataset_name+'.pkl')
204
+
205
+ if not os.path.isdir('./npy/'+dataset_name):
206
+ os.system('mkdir ./npy/'+dataset_name)
207
+
208
+ if args.code_type=='w':
209
+ Gs=LoadModel(dataset_name=dataset_name)
210
+ GetCode(Gs,random_state,num_img,num_once,dataset_name)
211
+ # GetImg(Gs,num_img=num_img,num_once=num_once,dataset_name=dataset_name,save_name='images_100K') #no need
212
+ elif args.code_type=='s':
213
+ save_name='S'
214
+ save_tmp=GetS(dataset_name,num_img=2_000)
215
+ tmp='./npy/'+dataset_name+'/'+save_name
216
+ with open(tmp, "wb") as fp:
217
+ pickle.dump(save_tmp, fp)
218
+
219
+ elif args.code_type=='s_mean_std':
220
+ save_tmp=GetS(dataset_name,num_img=num_img)
221
+ dlatents=save_tmp[1]
222
+ m,std=GetCodeMS(dlatents)
223
+ save_tmp=[m,std]
224
+ save_name='S_mean_std'
225
+ tmp='./npy/'+dataset_name+'/'+save_name
226
+ with open(tmp, "wb") as fp:
227
+ pickle.dump(save_tmp, fp)
228
+
229
+
230
+
231
+
232
+
models/StyleCLIP/global_directions/GetGUIData.py ADDED
@@ -0,0 +1,67 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import os
3
+ import numpy as np
4
+ import argparse
5
+ from manipulate import Manipulator
6
+ import torch
7
+ from PIL import Image
8
+ #%%
9
+
10
+ if __name__ == "__main__":
11
+ parser = argparse.ArgumentParser(description='Process some integers.')
12
+
13
+ parser.add_argument('--dataset_name',type=str,default='ffhq',
14
+ help='name of dataset, for example, ffhq')
15
+
16
+ parser.add_argument('--real', action='store_true')
17
+
18
+ args = parser.parse_args()
19
+ dataset_name=args.dataset_name
20
+
21
+ if not os.path.isdir('./data/'+dataset_name):
22
+ os.system('mkdir ./data/'+dataset_name)
23
+ #%%
24
+ M=Manipulator(dataset_name=dataset_name)
25
+ np.set_printoptions(suppress=True)
26
+ print(M.dataset_name)
27
+ #%%
28
+ #remove all .jpg
29
+ names=os.listdir('./data/'+dataset_name+'/')
30
+ for name in names:
31
+ if '.jpg' in name:
32
+ os.system('rm ./data/'+dataset_name+'/'+name)
33
+
34
+
35
+ #%%
36
+ if args.real:
37
+ latents=torch.load('./data/'+dataset_name+'/latents.pt')
38
+ w_plus=latents.cpu().detach().numpy()
39
+ else:
40
+ w=np.load('./npy/'+dataset_name+'/W.npy')
41
+ tmp=w[:50] #only use 50 images
42
+ tmp=tmp[:,None,:]
43
+ w_plus=np.tile(tmp,(1,M.Gs.components.synthesis.input_shape[1],1))
44
+ np.save('./data/'+dataset_name+'/w_plus.npy',w_plus)
45
+
46
+ #%%
47
+ tmp=M.W2S(w_plus)
48
+ M.dlatents=tmp
49
+
50
+ M.img_index=0
51
+ M.num_images=len(w_plus)
52
+ M.alpha=[0]
53
+ M.step=1
54
+ lindex,bname=0,0
55
+
56
+ M.manipulate_layers=[lindex]
57
+ codes,out=M.EditOneC(bname)
58
+ #%%
59
+
60
+ for i in range(len(out)):
61
+ img=out[i,0]
62
+ img=Image.fromarray(img)
63
+ img.save('./data/'+dataset_name+'/'+str(i)+'.jpg')
64
+ #%%
65
+
66
+
67
+
models/StyleCLIP/global_directions/Inference.py ADDED
@@ -0,0 +1,106 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+
3
+ from manipulate import Manipulator
4
+ import tensorflow as tf
5
+ import numpy as np
6
+ import torch
7
+ import clip
8
+ from MapTS import GetBoundary,GetDt
9
+
10
+ class StyleCLIP():
11
+
12
+ def __init__(self,dataset_name='ffhq'):
13
+ print('load clip')
14
+ device = "cuda" if torch.cuda.is_available() else "cpu"
15
+ self.model, preprocess = clip.load("ViT-B/32", device=device)
16
+ self.LoadData(dataset_name)
17
+
18
+ def LoadData(self, dataset_name):
19
+ tf.keras.backend.clear_session()
20
+ M=Manipulator(dataset_name=dataset_name)
21
+ np.set_printoptions(suppress=True)
22
+ fs3=np.load('./npy/'+dataset_name+'/fs3.npy')
23
+
24
+ self.M=M
25
+ self.fs3=fs3
26
+
27
+ w_plus=np.load('./data/'+dataset_name+'/w_plus.npy')
28
+ self.M.dlatents=M.W2S(w_plus)
29
+
30
+ if dataset_name=='ffhq':
31
+ self.c_threshold=20
32
+ else:
33
+ self.c_threshold=100
34
+ self.SetInitP()
35
+
36
+ def SetInitP(self):
37
+ self.M.alpha=[3]
38
+ self.M.num_images=1
39
+
40
+ self.target=''
41
+ self.neutral=''
42
+ self.GetDt2()
43
+ img_index=0
44
+ self.M.dlatent_tmp=[tmp[img_index:(img_index+1)] for tmp in self.M.dlatents]
45
+
46
+
47
+ def GetDt2(self):
48
+ classnames=[self.target,self.neutral]
49
+ dt=GetDt(classnames,self.model)
50
+
51
+ self.dt=dt
52
+ num_cs=[]
53
+ betas=np.arange(0.1,0.3,0.01)
54
+ for i in range(len(betas)):
55
+ boundary_tmp2,num_c=GetBoundary(self.fs3,self.dt,self.M,threshold=betas[i])
56
+ print(betas[i])
57
+ num_cs.append(num_c)
58
+
59
+ num_cs=np.array(num_cs)
60
+ select=num_cs>self.c_threshold
61
+
62
+ if sum(select)==0:
63
+ self.beta=0.1
64
+ else:
65
+ self.beta=betas[select][-1]
66
+
67
+
68
+ def GetCode(self):
69
+ boundary_tmp2,num_c=GetBoundary(self.fs3,self.dt,self.M,threshold=self.beta)
70
+ codes=self.M.MSCode(self.M.dlatent_tmp,boundary_tmp2)
71
+ return codes
72
+
73
+ def GetImg(self):
74
+
75
+ codes=self.GetCode()
76
+ out=self.M.GenerateImg(codes)
77
+ img=out[0,0]
78
+ return img
79
+
80
+
81
+
82
+
83
+ #%%
84
+ if __name__ == "__main__":
85
+ style_clip=StyleCLIP()
86
+ self=style_clip
87
+
88
+
89
+
90
+
91
+
92
+
93
+
94
+
95
+
96
+
97
+
98
+
99
+
100
+
101
+
102
+
103
+
104
+
105
+
106
+