remove PTI from model

app.py

@@ -29,7 +29,7 @@ from ris.blend import blend_latents
 from ris.model import Generator as RIS_Generator
 
 #from models.pti.manipulator import Manipulator
-#from models.pti.wrapper import Generator_wrapper
+#from models.pti.wrapper import Generator as Generator_wrapper
 #from models.pti.e4e_projection import projection
 
 from metrics import FaceMetric
@@ -334,12 +334,6 @@ with gr.Blocks() as demo:
             else:
                 e4e_output = [None, None, None, None, e4e_metrics_text]
             output_imgs.extend(e4e_output)
-            if 'PTI' in inverter_bools:
-                pti_output = None, None, None, None
-                #manipulator.set_real_img_projection(src, inv_mode='w+', pti_mode='s')
-            else:
-                pti_output = None, None, None, None
-            output_imgs.extend(pti_output)
         return output_imgs
     submit_button.click(
         submit, 
@@ -349,9 +343,10 @@ with gr.Blocks() as demo:
             gd_bool, neutral_text, target_text, alpha, beta,
             ris_bool, ref_img
         ],
-        [output_hyperstyle_invert, output_hyperstyle_mapper, output_hyperstyle_gd, output_hyperstyle_ris, output_hypersyle_metrics,
-         output_e4e_invert, output_e4e_mapper, output_e4e_gd, output_e4e_ris, output_e4e_metrics,
-         output_pti_invert, output_pti_mapper, output_pti_gd, output_pti_ris]
+        [
+            output_hyperstyle_invert, output_hyperstyle_mapper, output_hyperstyle_gd, output_hyperstyle_ris, output_hypersyle_metrics,
+            output_e4e_invert, output_e4e_mapper, output_e4e_gd, output_e4e_ris, output_e4e_metrics,
+        ]
             )
 
 demo.launch()

configs.py

@@ -0,0 +1,86 @@
+import numpy as np
+
+
+class GENERATOR_CONFIGS:
+    """StyleGAN2-ada generator configuration
+    """
+    def __init__(self, resolution=1024):
+        channel_base = 32768 if resolution >= 1024 else 16384
+        self.G_kwargs = {
+            'class_name': 'training.networks.Generator',
+            'z_dim': 512,
+            'w_dim': 512,
+            'mapping_kwargs': {'num_layers': 8},
+            'synthesis_kwargs': {
+                'channel_base': channel_base,
+                'channel_max': 512,
+                'num_fp16_res': 4,
+                'conv_clamp': 256
+            }
+        }
+        self.common_kwargs = {'c_dim': 0, 'img_resolution': resolution, 'img_channels': 3}
+        self.w_idx_lst = [
+            0,1,        # 4
+            1,2,3,      # 8
+            3,4,5,      # 16
+            5,6,7,      # 32
+            7,8,9,      # 64
+            9,10,11,    # 128
+            11,12,13,   # 256
+            13,14,15,   # 512
+            15,16,17,   # 1024
+        ]
+        cutoff_idx = int(np.log2(1024/resolution) * (-3))
+        if resolution < 1024:
+            self.w_idx_lst = self.w_idx_lst[:cutoff_idx]
+
+
+class PATH_CONFIGS:
+    """Paths configuration
+    """
+    def __init__(self): 
+        self.e4e = 'pretrained/e4e_ffhq_encode.pt'
+        self.stylegan2_ada_ffhq = 'pretrained/ffhq.pkl'
+        self.ir_se50 = 'pretrained/model_ir_se50.pth'
+        self.dlib = 'pretrained/shape_predictor_68_face_landmarks.dat'
+
+class PTI_HPARAMS:
+    """Pivot-tuning-inversion related hyper-parameters
+    """
+    def __init__(self):
+        # Architectures
+        self.lpips_type = 'alex'
+        self.first_inv_type = 'w+'
+        self.optim_type = 'adam'
+
+        # Locality regularization
+        self.latent_ball_num_of_samples = 1
+        self.locality_regularization_interval = 1
+        self.use_locality_regularization = False
+        self.regulizer_l2_lambda = 0.1
+        self.regulizer_lpips_lambda = 0.1
+        self.regulizer_alpha = 30
+        
+        ## Loss
+        self.pt_l2_lambda = 1
+        self.pt_lpips_lambda = 1
+        
+        ## Steps
+        self.LPIPS_value_threshold = 0.06
+        self.max_pti_steps = 350
+        self.first_inv_steps = 450
+        self.max_images_to_invert = 30
+        
+        ## Optimization
+        self.pti_learning_rate = 3e-4
+        self.first_inv_lr = 5e-3
+        self.train_batch_size = 1
+
+
+class PTI_GLOBAL_CFGS:
+    def __init__(self):
+        self.training_step = 1
+        self.imgage_rec_result_log_snapshot = 100
+        self.pivotal_training_steps = 0
+        self.model_snapshot_interval = 400
+        self.run_name = ''

models/pti/__init__.py

@@ -0,0 +1,3 @@
+from . import dnnlib
+from . import torch_utils
+from . import training

models/pti/legacy.py

@@ -68,6 +68,8 @@ class _LegacyUnpickler(pickle.Unpickler):
     def find_class(self, module, name):
         if module == 'dnnlib.tflib.network' and name == 'Network':
             return _TFNetworkStub
+        if module.split('.')[0] in ['dlib_utils', 'dnnlib', 'torch_utils', 'training']:
+            module = 'models.pti.'+module
         return super().find_class(module, name)
 
 #----------------------------------------------------------------------------

pretrained_models/pti/ffhq.pkl

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

ris/manipulator.py

@@ -289,7 +289,7 @@ if __name__ == '__main__':
 
     device = torch.device('cuda:0')
     ckpt = args.ckpt
-    G = Generator(ckpt, device)
+    G = Generator_wrapper(ckpt, device)
 
     face_preprocess = args.face_preprocess
     dataset_name = args.dataset_name

ris/wrapper.py

@@ -7,15 +7,15 @@ import PIL.Image
 import torch
 from torchvision.transforms import transforms
 
-import dnnlib
+import models.pti.dnnlib as dnnlib
 import legacy
-from configs_gd import GENERATOR_CONFIGS
-from dlib_utils.face_alignment import image_align
-from dlib_utils.landmarks_detector import LandmarksDetector
-from torch_utils.misc import copy_params_and_buffers
+from models.pti.configs import GENERATOR_CONFIGS
+from models.pti.dlib_utils.face_alignment import image_align
+from models.pti.dlib_utils.landmarks_detector import LandmarksDetector
+from models.pti.torch_utils.misc import copy_params_and_buffers
 
-from pivot_tuning_inversion.utils.ImagesDataset import ImagesDataset, ImageLatentsDataset
-from pivot_tuning_inversion.training.coaches.multi_id_coach import MultiIDCoach
+from models.pti.pivot_tuning_inversion.utils.ImagesDataset import ImagesDataset, ImageLatentsDataset
+from models.pti.pivot_tuning_inversion.training.coaches.multi_id_coach import MultiIDCoach
 
 
 class FaceLandmarksDetector: