fix

app.py

@@ -28,11 +28,18 @@ from PIL import Image
 import torch
 import torchvision.transforms as transforms
 from argparse import Namespace
+from e4e.utils.common import tensor2im
 from e4e.models.psp import pSp
 from e4e.models.encoders import psp_encoders
 from util import *
 from huggingface_hub import hf_hub_download
 
+transform = transforms.Compose([
+        transforms.Resize((256, 256)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+resize_dims = (256, 256)
+
 device= 'cpu'
 ffhq_model_path = hf_hub_download(repo_id="bankholdup/stylegan_petbreeder", filename="e4e_ffhq512.pt")
 
@@ -95,6 +102,15 @@ cat_decoder.eval()
 cat_decoder.to(device)
 
 
+def run_alignment(image_path):
+  import dlib
+  from e4e.utils.alignment import align_face
+  predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
+  aligned_image = align_face(filepath=image_path, predictor=predictor) 
+  print("Aligned image has shape: {}".format(aligned_image.size))
+  return aligned_image
+
+
 def gen_im(model_type='ffhq'):
   if model_type=='ffhq':
     imgs, _ = ffhq_decoder([ffhq_codes], input_is_latent=True, randomize_noise=False, return_latents=True)
@@ -109,15 +125,17 @@ def gen_im(model_type='ffhq'):
 
 def inference(img):  
     img.save('out.jpg')  
-    aligned_face = align_face('out.jpg')
+    #aligned_face = align_face('out.jpg')
+	input_image = run_alignment(image_path)
+	transformed_image = transform(input_image)
 	
-    ffhq_codes = ffhq_encoder(aligned_face.unsqueeze(0).to(device).float())
+    ffhq_codes = ffhq_encoder(transformed_image.unsqueeze(0).to(device).float())
     ffhq_codes = ffhq_codes + ffhq_latent_avg.repeat(ffhq_codes.shape[0], 1, 1)
 
-    cat_codes = cat_encoder(aligned_face.unsqueeze(0).to(device).float())
+    cat_codes = cat_encoder(transformed_image.unsqueeze(0).to(device).float())
     cat_codes = cat_codes + ffhq_latent_avg.repeat(cat_codes.shape[0], 1, 1)
 
-    dog_codes = dog_encoder(aligned_face.unsqueeze(0).to(device).float())
+    dog_codes = dog_encoder(transformed_image.unsqueeze(0).to(device).float())
     dog_codes = dog_codes + ffhq_latent_avg.repeat(dog_codes.shape[0], 1, 1)
         
     animal = "cat"