handler.py

import io
import os
import cv2
import torch
import base64
import logging
import requests
import numpy as np
from PIL import Image
from gfpgan import GFPGANer

# ======================================================
#  LOGGING CONFIGURATION
# ======================================================
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
logger.debug("📦 [INIT] Importing GFPGAN handler module...")

# ======================================================
#  GFPGAN MODEL URL
# ======================================================
MODEL_URL = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
MODEL_NAME = "GFPGANv1.4.pth"


# ======================================================
#  ENDPOINT HANDLER
# ======================================================
class EndpointHandler:
    def __init__(self, path="."):
        logger.debug("🚀 [INIT] Starting GFPGAN EndpointHandler initialization...")
        logger.debug(f"📁 Working directory: {os.getcwd()}")
        logger.debug(f"📂 Handler path argument: {path}")

        model_path = os.path.join(path, MODEL_NAME)
        logger.debug(f"🔗 [MODEL] Expected model path: {model_path}")

        # ------------------------------------------------------
        # Download model if missing
        # ------------------------------------------------------
        if not os.path.exists(model_path):
            try:
                logger.debug(f"📥 [DOWNLOAD] Model not found locally — fetching from {MODEL_URL}")
                r = requests.get(MODEL_URL, stream=True)
                r.raise_for_status()
                with open(model_path, "wb") as f:
                    for chunk in r.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
                logger.debug("✅ [MODEL] Downloaded GFPGAN weights successfully.")
            except Exception as e:
                logger.error(f"💥 [ERROR] Failed to download GFPGAN weights: {e}")
                raise

        # ------------------------------------------------------
        # Initialize GFPGANer (same as official Gradio demo)
        # ------------------------------------------------------
        try:
            logger.debug("🧠 [MODEL] Initializing GFPGANer (upscale=2, arch='clean')...")
            self.restorer = GFPGANer(
                model_path=model_path,
                upscale=2,                   # Rescaling factor = 2
                arch="clean",
                channel_multiplier=2,
                bg_upsampler=None
            )
            logger.debug("✅ [MODEL] GFPGAN model initialized successfully.")
        except Exception as e:
            logger.error(f"💥 [ERROR] Model initialization failed: {e}")
            raise

    # ======================================================
    #  INFERENCE CALL
    # ======================================================
    def __call__(self, data):
        logger.debug("⚙️ [INFER] Starting inference...")
        logger.debug(f"📥 Incoming data type: {type(data)}")

        # ------------------------------------------------------
        # Handle both JSON base64 and raw bytes
        # ------------------------------------------------------
        try:
            if isinstance(data, dict) and "inputs" in data:
                logger.debug("📦 Detected JSON base64 input")
                image_bytes = base64.b64decode(data["inputs"])
            elif isinstance(data, (bytes, bytearray)):
                logger.debug("📦 Detected raw bytes input")
                image_bytes = data
            else:
                raise ValueError("Unsupported input format — expected bytes or base64 data")

            logger.debug(f"🧾 [BYTES] Received {len(image_bytes)} bytes")
        except Exception as e:
            logger.error(f"💥 [ERROR] Input parsing failed: {e}")
            return {"error": f"Invalid input: {e}"}

        # ------------------------------------------------------
        # Decode image
        # ------------------------------------------------------
        try:
            img_np = np.array(Image.open(io.BytesIO(image_bytes)).convert("RGB"))
            logger.debug(f"🖼️ [IMAGE] Loaded image of shape: {img_np.shape}")
        except Exception as e:
            logger.error(f"💥 [ERROR] Failed to load image: {e}")
            return {"error": f"Image loading failed: {e}"}

        # ------------------------------------------------------
        # Run GFPGAN restoration
        # ------------------------------------------------------
        try:
            cropped_faces, restored_faces, restored_img = self.restorer.enhance(
                img_np,
                has_aligned=False,
                only_center_face=False,
                paste_back=True  # Matches GFPGAN web demo
            )
            logger.debug("✅ [RESTORE] Face restoration completed successfully.")
        except Exception as e:
            logger.error(f"💥 [ERROR] GFPGAN enhancement failed: {e}")
            return {"error": f"Enhancement failed: {e}"}

        # ------------------------------------------------------
        # Encode result as base64 PNG
        # ------------------------------------------------------
        try:
            _, buffer = cv2.imencode(".png", restored_img[:, :, ::-1])  # BGR→RGB
            img_base64 = base64.b64encode(buffer).decode("utf-8")
            logger.debug("📤 [ENCODE] Encoded restored image successfully.")
            return {"image": img_base64}
        except Exception as e:
            logger.error(f"💥 [ERROR] Failed to encode image: {e}")
            return {"error": f"Encoding failed: {e}"}