add wrapper

.gitattributes

@@ -36,3 +36,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
  filter=lfs diff=lfs merge=lfs -text
 *.pkl* filter=lfs diff=lfs merge=lfs -text
  filter=lfs diff=lfs merge=lfs -text
+*.dat* filter=lfs diff=lfs merge=lfs -text
+*.pt* filter=lfs diff=lfs merge=lfs -text

find_direction.py

@@ -20,7 +20,7 @@ from PIL import Image
 from torch_utils import misc
 from torch_utils.ops import upfirdn2d
 import id_loss
-
+from copy import deepcopy
 
 def block_forward(self, x, img, ws, shapes, force_fp32=False, fused_modconv=None, **layer_kwargs):
         misc.assert_shape(ws, [None, self.num_conv + self.num_torgb, self.w_dim])
@@ -72,13 +72,14 @@ def unravel_index(index, shape):
     return tuple(reversed(out))
 
 def find_direction(
-    G,
+    GIn,
     text_prompt: str,
     truncation_psi: float = 0.7,
     noise_mode: str = "const",
     resolution: int = 256,
     identity_power: float = 0.5,
 ):
+    G = deepcopy(GIn)
     seeds=np.random.randint(0, 1000, 128)
 
     batch_size=1
@@ -160,8 +161,9 @@ def find_direction(
     styles_direction_grad_el2 = torch.zeros(1, 26, 512, device=device)
     styles_direction.requires_grad_()
 
-    global id_loss
-    id_loss = id_loss.IDLoss("a").to(device).eval()
+    global id_loss2
+    #id_loss = id_loss.IDLoss("a").to(device).eval()
+    id_loss2 = id_loss.IDLoss("a").to(device).eval()
 
     temp_photos = []
     grads = []
@@ -205,7 +207,7 @@ def find_direction(
                 x, img = block_forward(block, x, img, styles2[:, styles_idx:styles_idx+3, :], temp_shapes[k], noise_mode=noise_mode, force_fp32=True)
                 styles_idx += 3
 
-        identity_loss, _ = id_loss(img, img2)
+        identity_loss, _ = id_loss2(img, img2)
         identity_loss *= id_coeff
         img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255)
         img = (transf(img.permute(0, 3, 1, 2)) / 255).sub_(mean).div_(std)
@@ -238,7 +240,7 @@ def find_direction(
                 x, img = block_forward(block, x, img, styles2[:, styles_idx:styles_idx+3, :], temp_shapes[k], noise_mode=noise_mode, force_fp32=True)
                 styles_idx += 3
 
-        identity_loss, _ = id_loss(img, img2)
+        identity_loss, _ = id_loss2(img, img2)
         identity_loss *= id_coeff
         img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255)
         img = (transf(img.permute(0, 3, 1, 2)) / 255).sub_(mean).div_(std)

generator.py

@@ -21,6 +21,7 @@ from torch_utils.ops import conv2d_resample
 from torch_utils.ops import upfirdn2d
 from torch_utils.ops import bias_act
 from torch_utils.ops import fma
+from copy import deepcopy
 
 
 import click
@@ -125,13 +126,14 @@ def unravel_index(index, shape):
 
 
 def w_to_s(
-    G,
-    outdir: str,
-    projected_w: str,
+    GIn,
+    wsIn:np.ndarray,
+    outdir: str ="s_out",
     truncation_psi: float = 0.7,
     noise_mode: str = "const",
 ):
-
+    G=deepcopy(GIn)
+    
     # Use GPU if available
     if torch.cuda.is_available():
         device = torch.device("cuda")
@@ -144,8 +146,8 @@ def w_to_s(
     for i in G.parameters():
         i.requires_grad = True
 
-    ws = np.load(projected_w)['w']
-    ws = torch.tensor(ws, device=device)
+    # ws = np.load(projected_w)['w']
+    ws = torch.tensor(wsIn, device=device)
 
     block_ws = []
     with torch.autograd.profiler.record_function('split_ws'):
@@ -186,17 +188,18 @@ def w_to_s(
 
     styles = styles.detach()
     np.savez(f'{outdir}/input.npz', s=styles.cpu().numpy())
-
+    return styles.cpu().numpy()
 
 def generate_from_style(
-    G,
+    GIn,
+    styles: np.ndarray,
+    styles_direction: np.ndarray,
     outdir: str,
-    s_input: str,
-    text_prompt: str,
     change_power: int,
     truncation_psi: float = 0.7,
     noise_mode: str = "const",
 ):
+    G=deepcopy(GIn)
     # Use GPU if available
     if torch.cuda.is_available():
         device = torch.device("cuda")
@@ -225,16 +228,12 @@ def generate_from_style(
 
         temp_shapes.append(temp_shape)
 
-    if s_input is not None:
-        styles = np.load(s_input)['s']
-        styles_direction = np.load(f'{outdir}/direction_'+text_prompt.replace(" ", "_")+'.npz')['s']
-
-        styles_direction = torch.tensor(styles_direction, device=device)
-        styles = torch.tensor(styles, device=device)
+    styles_direction = torch.tensor(styles_direction, device=device)
+    styles = torch.tensor(styles, device=device)
 
     with torch.no_grad():
         imgs = []
-        grad_changes = [0, 0.25*change_power, 0.5*change_power, 0.75*change_power, change_power]
+        grad_changes = [change_power]
 
         for grad_change in grad_changes:
             styles += styles_direction*grad_change
@@ -256,5 +255,6 @@ def generate_from_style(
 
             styles -= styles_direction*grad_change
 
-        img_filepath = f'{outdir}/'+text_prompt.replace(" ", "_")+'_'+str(change_power)+'.jpeg'
-        PIL.Image.fromarray(np.concatenate(imgs, axis=1), 'RGB').save(img_filepath, quality=95)
+        output_image = PIL.Image.fromarray(np.concatenate(imgs, axis=1), 'RGB')
+        output_image.save(os.path.join(outdir, 'final_out.png'), quality=95)
+        return output_image

pretrained/e4e_ffhq_encode.pt

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

pretrained/shape_predictor_68_face_landmarks.dat

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

psp_wrapper.py

@@ -0,0 +1,54 @@
+from argparse import Namespace
+import sys
+
+sys.path.append(".")
+sys.path.append("..")
+sys.path.append("./encoder4editing")
+
+from PIL import Image
+import torch
+import torchvision.transforms as transforms
+import dlib
+from utils.alignment import align_face
+
+
+from utils.common import tensor2im
+from models.psp import pSp  # we use the pSp framework to load the e4e encoder.
+experiment_type = 'ffhq_encode'
+
+EXPERIMENT_DATA_ARGS = {
+    "ffhq_encode": {
+        "model_path": "encoder4editing/e4e_ffhq_encode.pt",
+        "image_path": "notebooks/images/input_img.jpg"
+    },
+}
+# Setup required image transformations
+EXPERIMENT_ARGS = EXPERIMENT_DATA_ARGS[experiment_type]
+EXPERIMENT_ARGS['transform'] = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+
+class psp_encoder:
+    def __init__(self, model_path: str, shape_predictor_path: str):
+        self.ckpt = torch.load(model_path, map_location='cpu')
+        self.opts = self.ckpt['opts']
+        # update the training options
+        self.opts['checkpoint_path'] = model_path
+        self.opts= Namespace(**self.opts)
+        self.net = pSp(self.opts)
+        self.net.eval()
+        self.net.cuda()
+        self.shape_predictor = dlib.shape_predictor(shape_predictor_path)
+
+    def get_w(self, image_path):
+        original_image = Image.open(image_path)
+        original_image = original_image.convert("RGB")
+        input_image = align_face(filepath=image_path, predictor=self.shape_predictor)
+        resize_dims = (256, 256)
+        input_image.resize(resize_dims)
+        img_transforms = EXPERIMENT_ARGS['transform']
+        transformed_image = img_transforms(input_image)
+        with torch.no_grad():
+            _, latents = self.net(transformed_image.unsqueeze(0).to("cuda").float(), randomize_noise=False, return_latents=True)
+            return latents.cpu().numpy()