Update app.py

app.py

@@ -1,219 +1,227 @@
-# app.py
-# Updated: 2025-04-05 18:56:50 PM IST (Ludhiana, Punjab, India)
-# Fix: Removed .to(device) call after CodeFormer() initialization due to AttributeError.
-# Added check for .eval() method.
-
-import gradio as gr
+import os
+import sys
 import torch
-import cv2
+import gradio as gr
+from PIL import Image
 import numpy as np
-import os
-import time
-import warnings
-import traceback # Import traceback for detailed error printing
-
-print("--- Script Start ---")
-print(f"Current Time (IST): {time.strftime('%Y-%m-%d %H:%M:%S')}") # Using system time from HF Space
-
-# Suppress specific warnings or all warnings if needed
-warnings.filterwarnings("ignore")
-
-# --- Globals ---
-codeformer_net = None
-is_initialized = False
-init_error_message = ""
-
-# --- Attempt to Import CodeFormer ---
-try:
-    from codeformer import CodeFormer
-    print("Successfully imported CodeFormer.")
-except ImportError as e:
-    init_error_message = f"Error: CodeFormer package not found or import failed. Check requirements.txt installation. Details: {e}"
-    print(init_error_message)
-    # Keep script running so Gradio might load with an error message
-except Exception as e:
-    init_error_message = f"An unexpected error occurred during import: {e}\n{traceback.format_exc()}"
-    print(init_error_message)
-
-
-# --- Initialize Model (only if import succeeded) ---
-if 'CodeFormer' in globals(): # Check if import was successful
-    print("Attempting to initialize CodeFormer model (simplified init, no .to(device))...")
-    try:
-        # Define the target device (though we won't call .to() on the model object)
-        device = torch.device("cpu")
-        print(f"Target device check: {device}")
-
-        # *** MODIFIED LINE BELOW ***
-        # Initialize WITHOUT .to(device)
-        codeformer_net = CodeFormer()
-        print("CodeFormer() object created.")
-
-        # Check if the object has an eval method (common for models)
-        # If it doesn't, it might not be a standard torch.nn.Module
-        try:
-            codeformer_net.eval()
-            print("codeformer_net.eval() called successfully.")
-        except AttributeError:
-            print("Note: codeformer_net object does not have an eval() method. Assuming it's not needed or handled internally.")
-        except Exception as e_eval:
-             print(f"Warning: Calling codeformer_net.eval() failed: {e_eval}")
-
-
-        is_initialized = True
-        print("CodeFormer model object initialization appears successful (removed .to(device)).")
-
-    except FileNotFoundError as e:
-        # This error might occur if CodeFormer() tries to load weights internally and fails
-        init_error_message = f"Error: Could not find CodeFormer model weights during init. The package might have failed to download them automatically. Check logs. Details: {e}\n{traceback.format_exc()}"
-        print(init_error_message)
-        codeformer_net = None # Ensure model is None if init fails
-    except TypeError as e:
-         # Catch TypeErrors specifically from the init call
-         init_error_message = f"TypeError during CodeFormer init. The constructor signature might be wrong. Details: {e}\n{traceback.format_exc()}"
-         print(init_error_message)
-         codeformer_net = None
-    except Exception as e:
-        # Catch any other initialization errors
-        init_error_message = f"Error initializing CodeFormer model object: {e}\n{traceback.format_exc()}"
-        print(init_error_message)
-        codeformer_net = None # Ensure model is None if init fails
-else:
-    # If import failed, ensure message reflects that
-    if not init_error_message: # Safety net if import error wasn't captured somehow
-       init_error_message = "CodeFormer could not be imported. Cannot initialize model."
-    print("Skipping model initialization due to import failure.")
-
-
-# --- Processing Function ---
-def enhance_image(input_img, fidelity_weight, background_enhance, face_upsample):
-    """
-    Enhances the input image using CodeFormer.
-    Args:
-        input_img (np.ndarray): Input image from Gradio (RGB format).
-        fidelity_weight (float): Balances fidelity and quality (0 = best quality, 1 = best fidelity).
-        background_enhance (bool): Whether to enhance background using RealESRGAN.
-        face_upsample (bool): Whether to further upsample restored faces.
-    Returns:
-        np.ndarray: Enhanced image (RGB format) or None on error.
-        str: Status or processing time message.
-    """
-    print("--- enhance_image function called ---") # Log function entry
-
-    if not is_initialized or codeformer_net is None:
-        error_msg = f"ERROR: CodeFormer model is not available. Initialization failed. Check logs for details. Message: {init_error_message}"
-        print(error_msg)
-        # Return None for the image and the error message for the status textbox
-        return None, error_msg
-
-    if input_img is None:
-        print("Error: No input image provided.")
-        return None, "Error: No input image provided."
-
-    print(f"Processing image with fidelity: {fidelity_weight}, bg_enhance: {background_enhance}, face_upsample: {face_upsample}")
-    start_time = time.time()
-
-    try:
-        # 1. Convert RGB (from Gradio) to BGR (often expected by OpenCV/CodeFormer backend)
-        img_bgr = cv2.cvtColor(input_img, cv2.COLOR_RGB2BGR)
-        print("Input image converted from RGB to BGR.")
-
-        # 2. Select background upsampler based on checkbox
-        bg_upsampler = 'realesrgan' if background_enhance else None
-        print(f"Background upsampler selected: {bg_upsampler}")
-
-        # 3. Run CodeFormer enhancement
-        # Ensure parameters match the CodeFormer API you have installed.
-        # `w` corresponds to fidelity_weight.
-        print("Starting CodeFormer enhancement...")
-        # *** Assuming enhance method exists and works on CPU by default ***
-        with torch.no_grad(): # Keep torch.no_grad() as internal ops might use Torch
-            output_bgr, _, _ = codeformer_net.enhance(
-                img_bgr,
-                w=fidelity_weight,
-                adain=True, # AdaIN is commonly used
-                face_upsample=face_upsample,
-                bg_upsampler=bg_upsampler
-            )
-        print("CodeFormer enhancement finished.")
-
-        # 4. Convert BGR output back to RGB for Gradio display
-        output_rgb = cv2.cvtColor(output_bgr, cv2.COLOR_BGR2RGB)
-        print("Output image converted from BGR to RGB.")
-
-        end_time = time.time()
-        processing_time = end_time - start_time
-        time_msg = f"Processing finished in {processing_time:.2f} seconds (on CPU)."
-        print(time_msg)
-        return output_rgb, time_msg
-
-    except Exception as e:
-        # Catch errors during the enhancement process
-        error_details = f"Error during enhancement: {e}\n{traceback.format_exc()}"
-        print(error_details)
-        return None, error_details # Return error message to the status textbox
-
-# --- Gradio Interface Definition ---
-title = "CodeFormer Image Enhancement (CPU Demo)"
-# Dynamically update description based on model load status
-status_message = 'Model initialization appears successful.' if is_initialized else f'Model Load FAILED: {init_error_message}'
-description = f"""
-Upload an image to enhance its quality, particularly for faces, using CodeFormer.
-**Note:** This demo runs on a free Hugging Face CPU. Processing will be **SLOW** (expect seconds to minutes per image).
-Adjust the fidelity weight (0 = max quality enhancement, 1 = closer to original). Optionally enhance background and upsample faces.
-**Status:** {status_message} Check Logs for details if failed.
-"""
-article = "<p style='text-align: center'>CodeFormer CPU Demo | <a href='https://github.com/sczhou/CodeFormer' target='_blank'>Official Repo</a></p>"
-
-# Define examples (Make sure the 'examples' folder and files exist in your Space repo)
-example_list = []
-if os.path.exists("examples"):
-    example_list = [
-        ["examples/face1.png", 0.7, True, True],
-        ["examples/face2.png", 0.5, True, True],
-        ["examples/bg1.png", 0.8, True, False],
-    ]
-    print("Example files folder found.")
-else:
-    print("Note: 'examples' folder not found. Gradio examples will be empty.")
-
-
-print("Defining Gradio interface...")
-try:
-    iface = gr.Interface(
-        fn=enhance_image,
-        inputs=[
-            gr.Image(label="Upload Image", type="numpy"),
-            gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.7, label="Fidelity Weight (0 = Max Quality, 1 = Max Fidelity)"),
-            gr.Checkbox(label="Enhance Background (Uses RealESRGAN)", value=True),
-            gr.Checkbox(label="Upsample Restored Faces", value=True)
-        ],
-        outputs=[
-            gr.Image(label="Enhanced Image", type="numpy"),
-            gr.Textbox(label="Status / Processing Time") # Output for status messages and time
-        ],
-        title=title,
-        description=description,
-        article=article,
-        examples=example_list,
-        allow_flagging="never"
+import requests
+from io import BytesIO
+from basicsr.utils import imwrite
+from torchvision.transforms.functional import normalize
+
+# Clone the CodeFormer repository if it doesn't exist
+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
+
+# Add the CodeFormer directory to the system 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
+
+# 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 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)
+
+# Load CodeFormer model
+def load_codeformer_model():
+    # Force to use 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')(
+        dim_embd=512,
+        codebook_size=1024,
+        n_head=8,
+        n_layers=9,
+        connect_list=['32', '64', '128', '256']
+    ).to(device)
+    
+    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'])
+    else:
+        codeformer_net.load_state_dict(checkpoint['params'])
+    
+    codeformer_net.eval()
+    
+    # Setup face restoration 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
     )
-    print("Gradio interface defined successfully.")
-except Exception as e:
-    print(f"FATAL: Failed to define Gradio interface: {e}\n{traceback.format_exc()}")
-    # If interface definition fails, we can't launch.
-    raise RuntimeError("Could not create Gradio Interface.") from e
+    
+    return codeformer_net, face_helper
+
+# Process the image with CodeFormer
+def process_image(image, w=0.5, has_aligned=False):
+    device = torch.device('cpu')
+    codeformer_net, face_helper = load_codeformer_model()
+    
+    # Convert the input image to numpy array
+    if isinstance(image, Image.Image):
+        img = np.array(image)
+    else:
+        img = image
+    
+    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
+    if len(face_helper.cropped_faces) == 0:
+        return image, "No face detected. Please try another image."
+    
+    # CodeFormer inference
+    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
+                restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
+            del output
+            torch.cuda.empty_cache()
+        except Exception as error:
+            print(f'Error: {error}')
+            restored_face = cropped_face
+        
+        # Save the restored face
+        face_helper.add_restored_face(restored_face)
+    
+    # Get the 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])
+    
+    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
 
+# Main function
+def main():
+    app = create_gradio_interface()
+    app.launch(share=True)
 
-# --- Launch the App ---
 if __name__ == "__main__":
-    print("Attempting to launch Gradio app...")
-    try:
-        iface.launch()
-        print("Gradio app launched successfully.")
-        print("--- Script End (App Running) ---")
-    except Exception as e:
-        print(f"FATAL: Failed to launch Gradio interface: {e}\n{traceback.format_exc()}")
-        # Exit if launch fails. Logs should show the error.
-        exit(1)
-
+    main()