use facexlib

app.py

@@ -3,7 +3,11 @@ from PIL import Image
 import torch
 import gradio as gr
 os.system("pip install gradio==2.5.3")
-os.system("pip install autocrop")
+
+os.system("pip install facexlib")
+
+from facexlib.utils.face_restoration_helper import FaceRestoreHelper
+#os.system("pip install autocrop")
 #os.system("pip install dlib")
 from autocrop import Cropper
 import torch
@@ -34,7 +38,15 @@ os.makedirs('models', exist_ok=True)
 #os.system("bzip2 -dk shape_predictor_68_face_landmarks.dat.bz2")
 #os.system("mv shape_predictor_68_face_landmarks.dat models/dlibshape_predictor_68_face_landmarks.dat")
 
-cropper = Cropper(face_percent=80)
+#cropper = Cropper(face_percent=80)
+
+face_helper = FaceRestoreHelper(
+        upscale_factor=0,
+        face_size=512,
+        crop_ratio=(1, 1),
+        det_model='retinaface_resnet50',
+        save_ext='png',
+        device='cpu')
 
 device = 'cpu' 
 
@@ -85,15 +97,21 @@ generatorjinx.load_state_dict(ckptjinx["g"], strict=False)
 
 
 def inference(img, model):    
+    face_helper.clean_all()
     #aligned_face = align_face(img)
-    cropped_array = cropper.crop(img[:,:,::-1])
+    #cropped_array = cropper.crop(img[:,:,::-1])
     
-    if cropped_array.any():
-        aligned_face = Image.fromarray(cropped_array)
-    else:
-        aligned_face = Image.fromarray(img[:,:,::-1])
-
-    my_w = e4e_projection(aligned_face, "test.pt", device).unsqueeze(0)
+    #if cropped_array.any():
+        #aligned_face = Image.fromarray(cropped_array)
+    #else:
+        #aligned_face = Image.fromarray(img[:,:,::-1])
+    
+    face_helper.read_image(img)
+    face_helper.get_face_landmarks_5(only_center_face=False, eye_dist_threshold=5)
+    face_helper.align_warp_face(save_cropped_path=".")
+    pilimg = Image.open("./_00.png")
+    
+    my_w = e4e_projection(pilimg, "test.pt", device).unsqueeze(0)
     if model == 'JoJo':
         with torch.no_grad():
             my_sample = generatorjojo(my_w, input_is_latent=True)