app.py

# 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 torch
import cv2
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"
    )
    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


# --- 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)