Update app.py

app.py

@@ -32,20 +32,17 @@ os.chdir("..")
 shutil.rmtree("HairFastGAN")
 
 from hair_swap import HairFast, get_parser
+from utils.shape_predictor import align_face
 
 hair_fast = HairFast(get_parser().parse_args([]))
 
 def resize(image_path):
     img = Image.open(image_path)
-    square_size = 1024
-    
-    left = (img.width - square_size) / 2
-    top = (img.height - square_size) / 2
-    right = (img.width + square_size) / 2
-    bottom = (img.height + square_size) / 2
+
+    if img.size != (1024, 1024):
+        img = align_face(img, return_tensors=False)[0]
     
-    img_cropped = img.crop((left, top, right, bottom))
-    return img_cropped
+    return img
 
 @spaces.GPU
 def swap_hair(source, target_1, target_2, progress=gr.Progress(track_tqdm=True)):