Update app.py

app.py

@@ -6,7 +6,6 @@ import cv2
 import numpy as np
 from PIL import Image
 import urllib.request
-import tarfile
 
 # Function to download a file from a URL
 def download_file(url, dest):
@@ -22,39 +21,23 @@ def setup_environment():
     model_path = "weights/codeformer.pth"
     download_file(model_url, model_path)
 
-    # Download facexlib detection models (needed for face detection)
+    # Download facexlib detection models
     retinaface_url = "https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_Resnet50_Final.pth"
     retinaface_path = "weights/detection_Resnet50_Final.pth"
     download_file(retinaface_url, retinaface_path)
 
-# 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__()
-        # 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(),
-            torch.nn.Conv2d(dim_embd, dim_embd, kernel_size=3, stride=1, padding=1)
-        )
-        self.decoder = torch.nn.Sequential(
-            torch.nn.ConvTranspose2d(dim_embd, 3, kernel_size=3, stride=1, padding=1),
-            torch.nn.Sigmoid()
-        )
-
-    def forward(self, x, w=0.5, adain=True):
-        # Simplified forward pass (placeholder)
-        enc = self.encoder(x)
-        dec = self.decoder(enc)
-        return dec
+    # Download codeformer_arch.py
+    arch_url = "https://raw.githubusercontent.com/sczhou/CodeFormer/master/codeformer_arch.py"
+    download_file(arch_url, "codeformer_arch.py")
 
 # Load CodeFormer model
 def load_codeformer():
     setup_environment()
+    from codeformer_arch import CodeFormer
     model_path = "weights/codeformer.pth"
-    net = CodeFormer().to('cpu')
+    net = CodeFormer(dim_embd=512, codebook_size=1024, n_head=8, n_layer=9, connect_list=['32', '64', '128', '256']).to('cpu')
     checkpoint = torch.load(model_path, map_location='cpu')
-    net.load_state_dict(checkpoint, strict=False)  # strict=False due to simplified architecture
+    net.load_state_dict(checkpoint)
     net.eval()
     return net
 
@@ -90,8 +73,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()  # Removed align=True
-    face_helper.align_warp_face()      # Ensure alignment happens here
+    face_helper.get_face_landmarks_5()
+    face_helper.align_warp_face()
     
     # Enhance face with CodeFormer
     for cropped_face in face_helper.cropped_faces:
@@ -129,6 +112,5 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
-    # Ensure setup runs once
     setup_environment()
     demo.launch()