trysem ucalyptus commited on
Commit
4d9fdb5
0 Parent(s):

Duplicate from ucalyptus/PTI

Browse files

Co-authored-by: Sayantan Das <ucalyptus@users.noreply.huggingface.co>

This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. .gitattributes +27 -0
  2. .gitignore +5 -0
  3. README.md +38 -0
  4. app.py +26 -0
  5. configs/__init__.py +0 -0
  6. configs/__pycache__/__init__.cpython-36.pyc +0 -0
  7. configs/__pycache__/__init__.cpython-39.pyc +0 -0
  8. configs/__pycache__/global_config.cpython-36.pyc +0 -0
  9. configs/__pycache__/global_config.cpython-39.pyc +0 -0
  10. configs/__pycache__/hyperparameters.cpython-36.pyc +0 -0
  11. configs/__pycache__/hyperparameters.cpython-39.pyc +0 -0
  12. configs/__pycache__/paths_config.cpython-36.pyc +0 -0
  13. configs/__pycache__/paths_config.cpython-39.pyc +0 -0
  14. configs/evaluation_config.py +1 -0
  15. configs/global_config.py +12 -0
  16. configs/hyperparameters.py +28 -0
  17. configs/paths_config.py +31 -0
  18. criteria/__init__.py +0 -0
  19. criteria/__pycache__/__init__.cpython-36.pyc +0 -0
  20. criteria/__pycache__/__init__.cpython-39.pyc +0 -0
  21. criteria/__pycache__/l2_loss.cpython-36.pyc +0 -0
  22. criteria/__pycache__/l2_loss.cpython-39.pyc +0 -0
  23. criteria/__pycache__/localitly_regulizer.cpython-36.pyc +0 -0
  24. criteria/__pycache__/localitly_regulizer.cpython-39.pyc +0 -0
  25. criteria/l2_loss.py +8 -0
  26. criteria/localitly_regulizer.py +58 -0
  27. dnnlib/__init__.py +9 -0
  28. dnnlib/__pycache__/__init__.cpython-36.pyc +0 -0
  29. dnnlib/__pycache__/__init__.cpython-39.pyc +0 -0
  30. dnnlib/__pycache__/util.cpython-36.pyc +0 -0
  31. dnnlib/__pycache__/util.cpython-39.pyc +0 -0
  32. dnnlib/util.py +477 -0
  33. edit.py +83 -0
  34. editings/__pycache__/ganspace.cpython-39.pyc +0 -0
  35. editings/__pycache__/latent_editor.cpython-39.pyc +0 -0
  36. editings/ganspace.py +21 -0
  37. editings/ganspace_pca/ffhq_pca.pt +0 -0
  38. editings/interfacegan_directions/age.pt +0 -0
  39. editings/interfacegan_directions/rotation.pt +0 -0
  40. editings/interfacegan_directions/smile.pt +0 -0
  41. editings/latent_editor.py +23 -0
  42. evaluation/experiment_setting_creator.py +43 -0
  43. evaluation/qualitative_edit_comparison.py +156 -0
  44. makedirs.py +83 -0
  45. models/StyleCLIP/__init__.py +0 -0
  46. models/StyleCLIP/criteria/__init__.py +0 -0
  47. models/StyleCLIP/criteria/clip_loss.py +17 -0
  48. models/StyleCLIP/criteria/id_loss.py +39 -0
  49. models/StyleCLIP/global_directions/GUI.py +103 -0
  50. models/StyleCLIP/global_directions/GenerateImg.py +50 -0
.gitattributes ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bin.* filter=lfs diff=lfs merge=lfs -text
5
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.model filter=lfs diff=lfs merge=lfs -text
12
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
13
+ *.onnx filter=lfs diff=lfs merge=lfs -text
14
+ *.ot filter=lfs diff=lfs merge=lfs -text
15
+ *.parquet filter=lfs diff=lfs merge=lfs -text
16
+ *.pb filter=lfs diff=lfs merge=lfs -text
17
+ *.pt filter=lfs diff=lfs merge=lfs -text
18
+ *.pth filter=lfs diff=lfs merge=lfs -text
19
+ *.rar filter=lfs diff=lfs merge=lfs -text
20
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
21
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
22
+ *.tflite filter=lfs diff=lfs merge=lfs -text
23
+ *.tgz filter=lfs diff=lfs merge=lfs -text
24
+ *.xz filter=lfs diff=lfs merge=lfs -text
25
+ *.zip filter=lfs diff=lfs merge=lfs -text
26
+ *.zstandard filter=lfs diff=lfs merge=lfs -text
27
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
.gitignore ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ *.dat
2
+ *.pkl
3
+ *.zip
4
+ __pycache__/*
5
+ *.pyc
README.md ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ title: PTI
3
+ emoji: 🦀
4
+ colorFrom: gray
5
+ colorTo: pink
6
+ sdk: gradio
7
+ app_file: app.py
8
+ pinned: false
9
+ duplicated_from: ucalyptus/PTI
10
+ ---
11
+
12
+ # Configuration
13
+
14
+ `title`: _string_
15
+ Display title for the Space
16
+
17
+ `emoji`: _string_
18
+ Space emoji (emoji-only character allowed)
19
+
20
+ `colorFrom`: _string_
21
+ Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
22
+
23
+ `colorTo`: _string_
24
+ Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
25
+
26
+ `sdk`: _string_
27
+ Can be either `gradio`, `streamlit`, or `static`
28
+
29
+ `sdk_version` : _string_
30
+ Only applicable for `streamlit` SDK.
31
+ See [doc](https://hf.co/docs/hub/spaces) for more info on supported versions.
32
+
33
+ `app_file`: _string_
34
+ Path to your main application file (which contains either `gradio` or `streamlit` Python code, or `static` html code).
35
+ Path is relative to the root of the repository.
36
+
37
+ `pinned`: _boolean_
38
+ Whether the Space stays on top of your list.
app.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from PIL import Image
3
+ os.system("pip install gdown lpips gradio")
4
+ os.system("gdown --id 1HKmjg6iXsWr4aFPuU0gBXPGR83wqMzq7 -O align.dat")
5
+ os.system("wget https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/ffhq.pkl")
6
+ os.system("gdown https://github.com/ninja-build/ninja/releases/download/v1.10.2/ninja-linux.zip")
7
+ os.system("unzip ninja-linux.zip -d /usr/local/bin/")
8
+ os.system("sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force")
9
+ os.makedirs("embeddings/",exist_ok=True)
10
+ os.makedirs("checkpoints/",exist_ok=True)
11
+ os.makedirs("output/", exist_ok=True)
12
+ os.makedirs("images/",exist_ok=True)
13
+ import gradio as gr
14
+
15
+
16
+
17
+ def inference(img):
18
+ img.save("images/file.png")
19
+ os.system("python tune.py")
20
+ return Image.open("output/out.png")
21
+
22
+ title = "Pivotal Tuning for Latent-based editing of Real Images"
23
+ description = "Gradio Demo for Pivotal Tuning Inversion. To use it, simply upload your image. Ensure the image is 256 x 256 in size."
24
+ article = "<p style='text-align: center'><a href='https://github.com/danielroich/PTI' target='_blank'>Github Repo</a>"
25
+ gr.Interface(inference, [gr.inputs.Image(type="pil")], gr.outputs.Image(type="pil"),title=title,description=description,article=article,allow_flagging='never',allow_screenshot=False).launch(share=True)
26
+ os.system("rm output/out.png")
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 = 'e4e_ffhq_encode.pt'
3
+ stylegan2_ada_ffhq = 'ffhq.pkl'
4
+ style_clip_pretrained_mappers = ''
5
+ ir_se50 = 'model_ir_se50.pth'
6
+ dlib = '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,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import numpy as np
3
+ from criteria import l2_loss
4
+ from configs import hyperparameters
5
+ from configs import global_config
6
+
7
+
8
+ class Space_Regulizer:
9
+ def __init__(self, original_G, lpips_net):
10
+ self.original_G = original_G
11
+ self.morphing_regulizer_alpha = hyperparameters.regulizer_alpha
12
+ self.lpips_loss = lpips_net
13
+
14
+ def get_morphed_w_code(self, new_w_code, fixed_w):
15
+ interpolation_direction = new_w_code - fixed_w
16
+ interpolation_direction_norm = torch.norm(interpolation_direction, p=2)
17
+ direction_to_move = hyperparameters.regulizer_alpha * interpolation_direction / interpolation_direction_norm
18
+ result_w = fixed_w + direction_to_move
19
+ self.morphing_regulizer_alpha * fixed_w + (1 - self.morphing_regulizer_alpha) * new_w_code
20
+
21
+ return result_w
22
+
23
+ def get_image_from_ws(self, w_codes, G):
24
+ return torch.cat([G.synthesis(w_code, noise_mode='none', force_fp32=True) for w_code in w_codes])
25
+
26
+ def ball_holder_loss_lazy(self, new_G, num_of_sampled_latents, w_batch, use_wandb=False):
27
+ loss = 0.0
28
+
29
+ z_samples = np.random.randn(num_of_sampled_latents, self.original_G.z_dim)
30
+ w_samples = self.original_G.mapping(torch.from_numpy(z_samples).to(global_config.device), None,
31
+ truncation_psi=0.5)
32
+ territory_indicator_ws = [self.get_morphed_w_code(w_code.unsqueeze(0), w_batch) for w_code in w_samples]
33
+
34
+ for w_code in territory_indicator_ws:
35
+ new_img = new_G.synthesis(w_code, noise_mode='none', force_fp32=True)
36
+ with torch.no_grad():
37
+ old_img = self.original_G.synthesis(w_code, noise_mode='none', force_fp32=True)
38
+
39
+ if hyperparameters.regulizer_l2_lambda > 0:
40
+ l2_loss_val = l2_loss.l2_loss(old_img, new_img)
41
+ if use_wandb:
42
+ wandb.log({f'space_regulizer_l2_loss_val': l2_loss_val.detach().cpu()},
43
+ step=global_config.training_step)
44
+ loss += l2_loss_val * hyperparameters.regulizer_l2_lambda
45
+
46
+ if hyperparameters.regulizer_lpips_lambda > 0:
47
+ loss_lpips = self.lpips_loss(old_img, new_img)
48
+ loss_lpips = torch.mean(torch.squeeze(loss_lpips))
49
+ if use_wandb:
50
+ wandb.log({f'space_regulizer_lpips_loss_val': loss_lpips.detach().cpu()},
51
+ step=global_config.training_step)
52
+ loss += loss_lpips * hyperparameters.regulizer_lpips_lambda
53
+
54
+ return loss / len(territory_indicator_ws)
55
+
56
+ def space_regulizer_loss(self, new_G, w_batch, use_wandb):
57
+ ret_val = self.ball_holder_loss_lazy(new_G, hyperparameters.latent_ball_num_of_samples, w_batch, use_wandb)
58
+ 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,83 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import click
2
+ import os
3
+ import sys
4
+ import pickle
5
+ import numpy as np
6
+ from PIL import Image
7
+ import glob
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
+
62
+ def edit(image_name):
63
+ generator_type = paths_config.multi_id_model_type if use_multi_id_training else image_name
64
+ old_G, new_G = load_generators(model_id, generator_type)
65
+ w_path_dir = f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}'
66
+
67
+ embedding_dir = f'{w_path_dir}/{paths_config.pti_results_keyword}/{image_name}'
68
+ w_pivot = torch.load(f'{embedding_dir}/0.pt')
69
+
70
+ old_image = old_G.synthesis(w_pivot, noise_mode='const', force_fp32 = True)
71
+ new_image = new_G.synthesis(w_pivot, noise_mode='const', force_fp32 = True)
72
+
73
+ latent_editor = LatentEditorWrapper()
74
+ latents_after_edit = latent_editor.get_single_interface_gan_edits(w_pivot, [i for i in range(-5,5)])
75
+
76
+ for direction, factor_and_edit in latents_after_edit.items():
77
+ for editkey in factor_and_edit.keys():
78
+ new_image = new_G.synthesis(factor_and_edit[editkey], noise_mode='const', force_fp32 = True)
79
+ image_pil = syn_images_wandb(new_image).save(f"/home/sayantan/PTI/{direction}/{editkey}/{image_name}.jpg")
80
+
81
+ if __name__ == '__main__':
82
+ for image_name in [f.split(".")[0].split("_")[2] for f in sorted(glob.glob("*.pt"))]:
83
+ 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,83 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import click
2
+ import os
3
+ import sys
4
+ import pickle
5
+ import numpy as np
6
+ from PIL import Image
7
+ import torch
8
+ from configs import paths_config, hyperparameters, global_config
9
+ from IPython.display import display
10
+ import matplotlib.pyplot as plt
11
+ from scripts.latent_editor_wrapper import LatentEditorWrapper
12
+
13
+
14
+ image_dir_name = '/home/sayantan/processed_images'
15
+ use_multi_id_training = False
16
+ global_config.device = 'cuda'
17
+ paths_config.e4e = '/home/sayantan/PTI/pretrained_models/e4e_ffhq_encode.pt'
18
+ paths_config.input_data_id = image_dir_name
19
+ paths_config.input_data_path = f'{image_dir_name}'
20
+ paths_config.stylegan2_ada_ffhq = '/home/sayantan/PTI/pretrained_models/ffhq.pkl'
21
+ paths_config.checkpoints_dir = '/home/sayantan/PTI/'
22
+ paths_config.style_clip_pretrained_mappers = '/home/sayantan/PTI/pretrained_models'
23
+ hyperparameters.use_locality_regularization = False
24
+ hyperparameters.lpips_type = 'squeeze'
25
+
26
+ model_id = "MYJJDFVGATAT"
27
+
28
+
29
+
30
+ def display_alongside_source_image(images):
31
+ res = np.concatenate([np.array(image) for image in images], axis=1)
32
+ return Image.fromarray(res)
33
+
34
+ def load_generators(model_id, image_name):
35
+ with open(paths_config.stylegan2_ada_ffhq, 'rb') as f:
36
+ old_G = pickle.load(f)['G_ema'].cuda()
37
+
38
+ with open(f'{paths_config.checkpoints_dir}/model_{model_id}_{image_name}.pt', 'rb') as f_new:
39
+ new_G = torch.load(f_new).cuda()
40
+
41
+ return old_G, new_G
42
+
43
+ def plot_syn_images(syn_images,text):
44
+ for img in syn_images:
45
+ img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).detach().cpu().numpy()[0]
46
+ plt.axis('off')
47
+ resized_image = Image.fromarray(img,mode='RGB').resize((256,256))
48
+ display(resized_image)
49
+ #wandb.log({text: [wandb.Image(resized_image, caption="Label")]})
50
+ del img
51
+ del resized_image
52
+ torch.cuda.empty_cache()
53
+
54
+ def syn_images_wandb(img):
55
+ img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).detach().cpu().numpy()[0]
56
+ plt.axis('off')
57
+ resized_image = Image.fromarray(img,mode='RGB').resize((256,256))
58
+ return resized_image
59
+
60
+ @click.command()
61
+ @click.pass_context
62
+ @click.option('--image_name', prompt='image name', help='The name for image')
63
+
64
+ def makedir(ctx: click.Context,image_name):
65
+ generator_type = paths_config.multi_id_model_type if use_multi_id_training else image_name
66
+ old_G, new_G = load_generators(model_id, generator_type)
67
+ w_path_dir = f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}'
68
+
69
+ embedding_dir = f'{w_path_dir}/{paths_config.pti_results_keyword}/{image_name}'
70
+ w_pivot = torch.load(f'{embedding_dir}/0.pt')
71
+
72
+ old_image = old_G.synthesis(w_pivot, noise_mode='const', force_fp32 = True)
73
+ new_image = new_G.synthesis(w_pivot, noise_mode='const', force_fp32 = True)
74
+
75
+ latent_editor = LatentEditorWrapper()
76
+ latents_after_edit = latent_editor.get_single_interface_gan_edits(w_pivot, [i for i in range(-5,5)])
77
+
78
+ for direction, factor_and_edit in latents_after_edit.items():
79
+ for editkey in factor_and_edit.keys():
80
+ os.makedirs(f"/home/sayantan/PTI/{direction}/{editkey}")
81
+
82
+ if __name__ == '__main__':
83
+ 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
+