Update app.py

app.py

@@ -14,6 +14,7 @@ from torch import nn, autograd, optim
 from torch.nn import functional as F
 from tqdm import tqdm
 import lpips
+import time
 
 
 #from e4e_projection import projection as e4e_projection
@@ -35,6 +36,9 @@ from e4e.models.stylegan2.model import Generator
 from util import *
 from huggingface_hub import hf_hub_download
 
+import dlib
+from e4e.utils.alignment import align_face
+
 transform = transforms.Compose([
         transforms.Resize((256, 256)),
         transforms.ToTensor(),
@@ -102,12 +106,11 @@ cat_decoder.load_state_dict(cat_d_filt, strict=True)
 cat_decoder.eval()
 cat_decoder.to(device)
 
+dlib_path = hf_hub_download(repo_id="bankholdup/stylegan_petbreeder", filename="shape_predictor_68_face_landmarks.dat")
+predictor = dlib.shape_predictor(dlib_path)
+
 
 def run_alignment(image_path):
-  import dlib
-  from e4e.utils.alignment import align_face
-  dlib_path = hf_hub_download(repo_id="bankholdup/stylegan_petbreeder", filename="shape_predictor_68_face_landmarks.dat")
-  predictor = dlib.shape_predictor(dlib_path)
   aligned_image = align_face(filepath=image_path, predictor=predictor) 
   print("Aligned image has shape: {}".format(aligned_image.size))
   return aligned_image
@@ -124,10 +127,16 @@ def gen_im(ffhq_codes, dog_codes, cat_codes, model_type='ffhq'):
     imgs, _ = custom_decoder([custom_codes], input_is_latent=True, randomize_noise=False, return_latents=True)
   return tensor2im(imgs[0])
 
+def set_seed(rd):
+    np.random.seed(rd)
+    torch.manual_seed(rd)
 
 def inference(img):  
+    random_seed = round(time.time() * 1000)
+    set_seed(random_seed)
+    
     img.save('out.jpg')  
-    #aligned_face = align_face('out.jpg')
+
     input_image = run_alignment('out.jpg')
     transformed_image = transform(input_image)