Update app.py

app.py

@@ -1,219 +1,17 @@
-import os
-import sys
-import torch
 import gradio as gr
-import numpy as np
-import requests
-from PIL import Image
+from gfpgan import GFPGANer
 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 repository and install dependencies
-if not os.path.exists('CodeFormer'):
-    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 repository to path
-sys.path.append('CodeFormer')
-
-# Create directories for model weights
-os.makedirs('CodeFormer/weights/CodeFormer', exist_ok=True)
-os.makedirs('CodeFormer/weights/facelib', exist_ok=True)
-
-# 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 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"
+restorer = GFPGANer(
+    model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.3.pth',
+    upscale=1,
+    arch='clean',
+    channel_multiplier=2,
+    bg_upsampler=None
 )
 
-# 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')
-    
-    # Initialize model
-    model = ARCH_REGISTRY.get('CodeFormer')(
-        dim_embd=512,
-        codebook_size=1024,
-        n_head=8,
-        n_layers=9,
-        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:
-        model.load_state_dict(checkpoint['params_ema'])
-    else:
-        model.load_state_dict(checkpoint['params'])
-    
-    model.eval()
-    
-    # Initialize face helper
-    face_helper = FaceRestoreHelper(
-        upscale_factor=1,
-        face_size=512,
-        crop_ratio=(1, 1),
-        det_model='retinaface_resnet50',
-        save_ext='png',
-        use_parse=True,
-        device=device
-    )
-    
-    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 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')
-    model, face_helper = load_model()
-    
-    # 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:
-        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)
-        face_helper.cropped_faces = [img]
-    else:
-        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()
-    
-    # Check if face was detected
-    if len(face_helper.cropped_faces) == 0:
-        return img, "No face detected in the image!"
-    
-    # Process each face
-    for idx, cropped_face in enumerate(face_helper.cropped_faces):
-        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 = 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
-        except Exception as e:
-            print(f'Error: {e}')
-            restored_face = cropped_face
-        
-        face_helper.add_restored_face(restored_face)
-    
-    # Get final result
-    if not has_aligned:
-        face_helper.get_inverse_affine(None)
-        restored_img = face_helper.paste_faces_to_input_image()
-    else:
-        restored_img = face_helper.restored_faces[0]
-    
-    # Return result as PIL image
-    return Image.fromarray(restored_img), "Face restoration complete!"
-
-# 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!"
-)
+def restore_image(img):
+    cropped_faces, restored_faces, restored_img = restorer.enhance(img, has_aligned=False, only_center_face=False)
+    return restored_img
 
-# Launch the app
-demo.launch()
+gr.Interface(fn=restore_image, inputs=gr.Image(type='numpy'), outputs=gr.Image(type='numpy')).launch()