Update app.py

app.py

@@ -2,76 +2,85 @@ import os
 import sys
 import torch
 import gradio as gr
-from PIL import Image
 import numpy as np
 import requests
-from io import BytesIO
-from basicsr.utils import imwrite
-from torchvision.transforms.functional import normalize
+from PIL import Image
+import cv2
+import subprocess
+
+# Function to run shell commands
+def run_command(command):
+    process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    stdout, stderr = process.communicate()
+    if process.returncode != 0:
+        print(f"Error executing command: {command}")
+        print(stderr.decode())
+    else:
+        print(stdout.decode())
+    return process.returncode
 
-# Clone the CodeFormer repository if it doesn't exist
+# Clone repository and install dependencies
 if not os.path.exists('CodeFormer'):
-    !git clone https://github.com/sczhou/CodeFormer.git
-    !pip install -r CodeFormer/requirements.txt
-    !pip install basicsr
-    !pip install facexlib
-    !pip install gradio>=3.25.0
-    !pip install realesrgan
-    !python CodeFormer/basicsr/setup.py develop
+    run_command("git clone https://github.com/sczhou/CodeFormer.git")
+    run_command("pip install -r CodeFormer/requirements.txt")
+    run_command("pip install basicsr facexlib realesrgan opencv-python")
+    run_command("python CodeFormer/basicsr/setup.py develop")
 
-# Add the CodeFormer directory to the system path
+# Add repository to path
 sys.path.append('CodeFormer')
 
-# Import necessary modules from CodeFormer
-from CodeFormer.basicsr.archs.codeformer_arch import CodeFormer
-from CodeFormer.basicsr.utils.registry import ARCH_REGISTRY
-from CodeFormer.facelib.utils.face_restoration_helper import FaceRestoreHelper
-from CodeFormer.facelib.detection.retinaface import retinaface
+# Create directories for model weights
+os.makedirs('CodeFormer/weights/CodeFormer', exist_ok=True)
+os.makedirs('CodeFormer/weights/facelib', exist_ok=True)
 
-# Function to download model weights
-def download_model_weights():
-    if not os.path.exists('CodeFormer/weights'):
-        os.makedirs('CodeFormer/weights/CodeFormer', exist_ok=True)
-        os.makedirs('CodeFormer/weights/facelib', exist_ok=True)
-    
-    # Download CodeFormer weights
-    codeformer_weight_path = 'CodeFormer/weights/CodeFormer/codeformer.pth'
-    if not os.path.exists(codeformer_weight_path):
-        print('Downloading CodeFormer weights...')
-        url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
-        r = requests.get(url, allow_redirects=True)
-        with open(codeformer_weight_path, 'wb') as f:
-            f.write(r.content)
-    
-    # Download detection model weights
-    detection_model_path = 'CodeFormer/weights/facelib/detection_Resnet50_Final.pth'
-    if not os.path.exists(detection_model_path):
-        print('Downloading face detection model weights...')
-        url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth'
-        r = requests.get(url, allow_redirects=True)
-        with open(detection_model_path, 'wb') as f:
-            f.write(r.content)
+# Download model weights
+def download_file(url, save_path):
+    if not os.path.exists(save_path):
+        print(f"Downloading {url} to {save_path}")
+        response = requests.get(url)
+        if response.status_code == 200:
+            with open(save_path, 'wb') as f:
+                f.write(response.content)
+            print(f"Downloaded {save_path}")
+        else:
+            print(f"Failed to download {url}, status code: {response.status_code}")
 
-    # Download parsing model weights
-    parsing_model_path = 'CodeFormer/weights/facelib/parsing_parsenet.pth'
-    if not os.path.exists(parsing_model_path):
-        print('Downloading face parsing model weights...')
-        url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth'
-        r = requests.get(url, allow_redirects=True)
-        with open(parsing_model_path, 'wb') as f:
-            f.write(r.content)
+# Download required model weights
+download_file(
+    "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth",
+    "CodeFormer/weights/CodeFormer/codeformer.pth"
+)
+download_file(
+    "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth",
+    "CodeFormer/weights/facelib/detection_Resnet50_Final.pth"
+)
+download_file(
+    "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth",
+    "CodeFormer/weights/facelib/parsing_parsenet.pth"
+)
 
-# Load CodeFormer model
-def load_codeformer_model():
-    # Force to use CPU
+# Import CodeFormer modules
+try:
+    from basicsr.archs.codeformer_arch import CodeFormer
+    from basicsr.utils.registry import ARCH_REGISTRY
+    from facelib.utils.face_restoration_helper import FaceRestoreHelper
+    from torchvision.transforms.functional import normalize
+except ImportError:
+    print("Error importing CodeFormer modules. Make sure all dependencies are installed correctly.")
+    # Try to install missing dependencies
+    run_command("cd CodeFormer && pip install -e .")
+    from basicsr.archs.codeformer_arch import CodeFormer
+    from basicsr.utils.registry import ARCH_REGISTRY
+    from facelib.utils.face_restoration_helper import FaceRestoreHelper
+    from torchvision.transforms.functional import normalize
+
+# Load model function
+def load_model():
+    print('Loading CodeFormer model on CPU...')
     device = torch.device('cpu')
-    print(f'Running on device: {device}')
-    
-    # Download model weights if they don't exist
-    download_model_weights()
     
-    # Load CodeFormer model
-    codeformer_net = ARCH_REGISTRY.get('CodeFormer')(
+    # Initialize model
+    model = ARCH_REGISTRY.get('CodeFormer')(
         dim_embd=512,
         codebook_size=1024,
         n_head=8,
@@ -79,17 +88,18 @@ def load_codeformer_model():
         connect_list=['32', '64', '128', '256']
     ).to(device)
     
+    # Load checkpoint
     ckpt_path = 'CodeFormer/weights/CodeFormer/codeformer.pth'
     checkpoint = torch.load(ckpt_path, map_location=device)
     
     if 'params_ema' in checkpoint:
-        codeformer_net.load_state_dict(checkpoint['params_ema'])
+        model.load_state_dict(checkpoint['params_ema'])
     else:
-        codeformer_net.load_state_dict(checkpoint['params'])
+        model.load_state_dict(checkpoint['params'])
     
-    codeformer_net.eval()
+    model.eval()
     
-    # Setup face restoration helper
+    # Initialize face helper
     face_helper = FaceRestoreHelper(
         upscale_factor=1,
         face_size=512,
@@ -100,128 +110,110 @@ def load_codeformer_model():
         device=device
     )
     
-    return codeformer_net, face_helper
+    return model, face_helper
+
+# Image conversion utilities
+def img2tensor(img, bgr2rgb=True, float32=True):
+    img = img.astype(np.float32) if float32 else img
+    if bgr2rgb:
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    img = torch.from_numpy(img.transpose(2, 0, 1))
+    return img
+
+def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)):
+    tensor = tensor.squeeze(0).float().detach().cpu().clamp_(*min_max)
+    tensor = (tensor - min_max[0]) / (min_max[1] - min_max[0])
+    
+    n_dim = tensor.dim()
+    if n_dim == 3:
+        img_np = tensor.numpy()
+        img_np = img_np.transpose(1, 2, 0)
+        if rgb2bgr:
+            img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+    elif n_dim == 2:
+        img_np = tensor.numpy()
+    else:
+        raise TypeError(f'Only support 3D and 2D tensor. But got {n_dim}D tensor.')
+    
+    if out_type == np.uint8:
+        img_np = (img_np * 255.0).round().astype(np.uint8)
+    
+    return img_np
 
-# Process the image with CodeFormer
-def process_image(image, w=0.5, has_aligned=False):
+# Process image function
+def process(image, w_value=0.5, has_aligned=False):
+    if image is None:
+        return None, "Please upload an image."
+    
     device = torch.device('cpu')
-    codeformer_net, face_helper = load_codeformer_model()
+    model, face_helper = load_model()
     
-    # Convert the input image to numpy array
+    # Convert PIL image to numpy array if needed
     if isinstance(image, Image.Image):
         img = np.array(image)
     else:
         img = image
-    
+
+    # Process aligned face or detect faces
     if has_aligned:
-        # The input image is already a cropped and aligned face
         face_helper.is_gray = len(img.shape) == 2 or (len(img.shape) == 3 and img.shape[2] == 1)
         if face_helper.is_gray:
             img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
-        # Prepare the face for processing
         face_helper.cropped_faces = [img]
     else:
-        # Detect and crop faces from the input image
         face_helper.clean_all()
         face_helper.read_image(img)
         face_helper.get_face_landmarks_5(only_center_face=False, resize=640, eye_dist_threshold=5)
         face_helper.align_warp_face()
     
-    # If no face is detected
+    # Check if face was detected
     if len(face_helper.cropped_faces) == 0:
-        return image, "No face detected. Please try another image."
+        return img, "No face detected in the image!"
     
-    # CodeFormer inference
+    # Process each face
     for idx, cropped_face in enumerate(face_helper.cropped_faces):
-        # Prepare the image for the model
         cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
         normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
         cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
         
         try:
             with torch.no_grad():
-                output = codeformer_net(cropped_face_t, w=w, adain=True)[0]
-                # Convert tensor to image
+                output = model(cropped_face_t, w=w_value, adain=True)[0]
                 restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
+            
+            # Free up memory
             del output
-            torch.cuda.empty_cache()
-        except Exception as error:
-            print(f'Error: {error}')
+        except Exception as e:
+            print(f'Error: {e}')
             restored_face = cropped_face
         
-        # Save the restored face
         face_helper.add_restored_face(restored_face)
     
-    # Get the final result
+    # Get final result
     if not has_aligned:
-        # Paste the restored faces back to the original image
         face_helper.get_inverse_affine(None)
         restored_img = face_helper.paste_faces_to_input_image()
-        restored_img = Image.fromarray(restored_img)
     else:
-        restored_img = Image.fromarray(face_helper.restored_faces[0])
+        restored_img = face_helper.restored_faces[0]
     
-    return restored_img, "Face successfully restored."
-
-# Helper functions for image conversion
-def img2tensor(img, bgr2rgb=True, float32=True):
-    img = img.astype(np.float32) if float32 else img
-    if bgr2rgb:
-        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    img = torch.from_numpy(img.transpose(2, 0, 1))
-    return img
-
-def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)):
-    tensor = tensor.squeeze(0).float().detach().cpu().clamp_(*min_max)
-    tensor = (tensor - min_max[0]) / (min_max[1] - min_max[0])
-    n_dim = tensor.dim()
-    if n_dim == 3:
-        img_np = tensor.numpy()
-        img_np = img_np.transpose(1, 2, 0)
-        if rgb2bgr:
-            img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
-    elif n_dim == 2:
-        img_np = tensor.numpy()
-    else:
-        raise TypeError(f'Only support 3D and 2D tensor. But got {n_dim}D tensor.')
-    if out_type == np.uint8:
-        img_np = (img_np * 255.0).round().astype(np.uint8)
-    return img_np
-
-# Create a Gradio interface for the app
-def create_gradio_interface():
-    with gr.Blocks(title="CodeFormer Face Restoration (CPU Version)") as app:
-        gr.Markdown("# CodeFormer Face Restoration (CPU Version)")
-        gr.Markdown("Upload a photo with faces to restore the quality. This model runs on CPU, so it might take a few minutes to process.")
-        
-        with gr.Row():
-            with gr.Column():
-                input_image = gr.Image(label="Input Image", type="pil")
-                w_slider = gr.Slider(0, 1, value=0.5, step=0.1, label="Fidelity Weight (0: more quality, 1: more identity)")
-                aligned_checkbox = gr.Checkbox(label="Input is an already aligned face", value=False)
-                process_button = gr.Button("Restore Face")
-            
-            with gr.Column():
-                output_image = gr.Image(label="Restored Image")
-                output_text = gr.Textbox(label="Status")
-        
-        process_button.click(
-            fn=process_image,
-            inputs=[input_image, w_slider, aligned_checkbox],
-            outputs=[output_image, output_text]
-        )
-        
-        gr.Markdown("Note: Lower fidelity weight (w) values create higher-quality results with more modifications, while higher values preserve more of the original identity.")
-        
-    return app
-
-# Import CV2 only when needed to avoid issues
-import cv2
+    # Return result as PIL image
+    return Image.fromarray(restored_img), "Face restoration complete!"
 
-# Main function
-def main():
-    app = create_gradio_interface()
-    app.launch(share=True)
+# Create Gradio interface
+demo = gr.Interface(
+    fn=process,
+    inputs=[
+        gr.Image(type="pil", label="Input Image"),
+        gr.Slider(0, 1, 0.5, step=0.01, label="Fidelity Weight (w) - Lower for quality, Higher for identity"),
+        gr.Checkbox(label="Is input an aligned face?", value=False)
+    ],
+    outputs=[
+        gr.Image(type="pil", label="Restored Image"),
+        gr.Textbox(label="Status")
+    ],
+    title="CodeFormer Face Restoration (CPU Version)",
+    description="Upload a photo with faces to restore the quality. Running on CPU, so please be patient!"
+)
 
-if __name__ == "__main__":
-    main()
+# Launch the app
+demo.launch()