Update app.py

app.py

@@ -27,11 +27,11 @@ def setup_environment():
     retinaface_path = "weights/detection_Resnet50_Final.pth"
     download_file(retinaface_url, retinaface_path)
 
-# Define a simplified CodeFormer architecture (instead of downloading codeformer_arch.py)
+# Define a simplified CodeFormer architecture (placeholder)
 class CodeFormer(torch.nn.Module):
     def __init__(self, dim_embd=512, codebook_size=1024, n_head=8, n_layer=9, connect_list=['32', '64', '128', '256']):
         super(CodeFormer, self).__init__()
-        # This is a simplified placeholder. In practice, you'd need the full architecture.
+        # Simplified placeholder (full architecture requires codeformer_arch.py)
         self.encoder = torch.nn.Sequential(
             torch.nn.Conv2d(3, dim_embd, kernel_size=3, stride=1, padding=1),
             torch.nn.ReLU(),
@@ -41,7 +41,6 @@ class CodeFormer(torch.nn.Module):
             torch.nn.ConvTranspose2d(dim_embd, 3, kernel_size=3, stride=1, padding=1),
             torch.nn.Sigmoid()
         )
-        # Note: This is a mock implementation. Full CodeFormer requires the actual codeformer_arch.py.
 
     def forward(self, x, w=0.5, adain=True):
         # Simplified forward pass (placeholder)
@@ -91,8 +90,8 @@ def enhance_image(image, fidelity_weight=0.5):
     face_helper = FaceRestoreHelper(upscale_factor=1, face_size=512, crop_ratio=(1, 1), det_model='retinaface_resnet50', save_ext='png', device='cpu')
     face_helper.clean_all()
     face_helper.read_image(img)
-    face_helper.get_face_landmarks_5(align=True)
-    face_helper.align_warp_face()
+    face_helper.get_face_landmarks_5()  # Removed align=True
+    face_helper.align_warp_face()      # Ensure alignment happens here
     
     # Enhance face with CodeFormer
     for cropped_face in face_helper.cropped_faces: