Update app.py

app.py

@@ -17,11 +17,46 @@ import lpips
 from model import *
 
 
-from e4e_projection import projection as e4e_projection
+#from e4e_projection import projection as e4e_projection
 
 from copy import deepcopy
 import imageio
 
+import os
+import sys
+import numpy as np
+from PIL import Image
+import torch
+import torchvision.transforms as transforms
+from argparse import Namespace
+from e4e.models.psp import pSp
+from util import *
+
+@ torch.no_grad()
+
+def projection(img, name, device='cuda'):
+ 
+    model_path = 'models/e4e_ffhq_encode.pt'
+    ckpt = torch.load(model_path, map_location='cpu')
+    opts = ckpt['opts']
+    opts['checkpoint_path'] = model_path
+    opts= Namespace(**opts)
+    net = pSp(opts, device).eval().to(device)
+    transform = transforms.Compose(
+        [
+            transforms.Resize(256),
+            transforms.CenterCrop(256),
+            transforms.ToTensor(),
+            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
+        ]
+    )
+    img = transform(img).unsqueeze(0).to(device)
+    images, w_plus = net(img, randomize_noise=False, return_latents=True)
+    result_file = {}
+    result_file['latent'] = w_plus[0]
+    torch.save(result_file, name)
+    return w_plus[0]
+
 os.makedirs('models', exist_ok=True)
 os.system("gdown https://drive.google.com/uc?id=1jtCg8HQ6RlTmLdnbT2PfW1FJ2AYkWqsK")
 os.system("cp e4e_ffhq_encode.pt models/e4e_ffhq_encode.pt")
@@ -113,7 +148,7 @@ generatorspider.load_state_dict(ckptspider["g"], strict=False)
 def inference(img, model):    
     aligned_face = align_face(img)
         
-    my_w = e4e_projection(aligned_face, "test.pt", device).unsqueeze(0)
+    my_w = projection(aligned_face, "test.pt", device).unsqueeze(0)
     if model == 'JoJo':
         with torch.no_grad():
             my_sample = generatorjojo(my_w, input_is_latent=True)