backend/services/upscaler.py

"""
Upscaler Service

Provides image upscaling functionality using Real-ESRGAN or similar models.
This enhances the resolution of generated images for higher quality output.
"""
import logging
from typing import Optional
import numpy as np
from PIL import Image
import cv2


logger = logging.getLogger(__name__)


# Global upscaler instance
_upscaler = None


def get_upscaler():
    """
    Get or create the upscaler model instance
    
    This uses Real-ESRGAN for high-quality upscaling.
    Falls back to simple interpolation if Real-ESRGAN is not available.
    
    Returns:
        Upscaler instance or None
    """
    global _upscaler
    
    if _upscaler is not None:
        return _upscaler
    
    try:
        from realesrgan import RealESRGANer
        from basicsr.archs.rrdbnet_arch import RRDBNet
        
        logger.info("Loading Real-ESRGAN model...")
        
        # Initialize the model
        # Using RealESRGAN_x4plus for 4x upscaling
        model = RRDBNet(
            num_in_ch=3,
            num_out_ch=3,
            num_feat=64,
            num_block=23,
            num_grow_ch=32,
            scale=4
        )
        
        # You can download the model weights from:
        # https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth
        model_path = "models/base/RealESRGAN_x4plus.pth"
        
        _upscaler = RealESRGANer(
            scale=4,
            model_path=model_path,
            model=model,
            tile=400,  # Tile size for processing large images
            tile_pad=10,
            pre_pad=0,
            half=True  # Use FP16 for faster inference
        )
        
        logger.info("Real-ESRGAN loaded successfully")
        return _upscaler
        
    except ImportError:
        logger.warning("Real-ESRGAN not available. Will use fallback upscaling.")
        return None
    except Exception as e:
        logger.warning(f"Failed to load Real-ESRGAN: {e}. Using fallback.")
        return None


def upscale_with_realesrgan(image: Image.Image, scale: int = 2) -> Image.Image:
    """
    Upscale image using Real-ESRGAN
    
    Args:
        image: Input PIL Image
        scale: Upscaling factor
        
    Returns:
        Upscaled PIL Image
    """
    upscaler = get_upscaler()
    
    if upscaler is None:
        raise RuntimeError("Real-ESRGAN not available")
    
    # Convert PIL to numpy array (BGR for OpenCV)
    img_np = np.array(image)
    img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
    
    # Upscale
    output, _ = upscaler.enhance(img_bgr, outscale=scale)
    
    # Convert back to PIL (RGB)
    output_rgb = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
    result = Image.fromarray(output_rgb)
    
    return result


def upscale_with_lanczos(image: Image.Image, scale: int = 2) -> Image.Image:
    """
    Fallback upscaling using Lanczos interpolation
    
    This is a simple but effective upscaling method when Real-ESRGAN
    is not available.
    
    Args:
        image: Input PIL Image
        scale: Upscaling factor
        
    Returns:
        Upscaled PIL Image
    """
    width, height = image.size
    new_size = (width * scale, height * scale)
    
    logger.info(f"Upscaling with Lanczos: {image.size} -> {new_size}")
    
    return image.resize(new_size, Image.LANCZOS)


def upscale_image(image: Image.Image, scale: int = 2) -> Image.Image:
    """
    Upscale an image using the best available method
    
    Tries Real-ESRGAN first, falls back to Lanczos if unavailable.
    
    Args:
        image: Input PIL Image
        scale: Upscaling factor (2 or 4 recommended)
        
    Returns:
        Upscaled PIL Image
    """
    try:
        # Try Real-ESRGAN
        logger.info(f"Upscaling image by {scale}x")
        result = upscale_with_realesrgan(image, scale=scale)
        logger.info("Upscaling completed with Real-ESRGAN")
        return result
        
    except (RuntimeError, Exception) as e:
        # Fallback to Lanczos
        logger.info(f"Using fallback upscaling method: {e}")
        return upscale_with_lanczos(image, scale=scale)


def adaptive_upscale(
    image: Image.Image,
    target_size: Optional[int] = None,
    max_scale: int = 4
) -> Image.Image:
    """
    Upscale image adaptively to reach a target size
    
    This function calculates the appropriate scale factor to reach
    the target size without exceeding max_scale.
    
    Args:
        image: Input PIL Image
        target_size: Target resolution for the longest edge
        max_scale: Maximum upscaling factor
        
    Returns:
        Upscaled PIL Image
    """
    if target_size is None:
        target_size = 2048  # Default target
    
    current_size = max(image.size)
    
    if current_size >= target_size:
        logger.info("Image already at target size or larger")
        return image
    
    # Calculate required scale
    required_scale = target_size / current_size
    
    # Clamp to max_scale and round to nearest power of 2
    if required_scale <= 1:
        scale = 1
    elif required_scale <= 2:
        scale = 2
    else:
        scale = min(4, max_scale)
    
    logger.info(f"Adaptive upscaling: {scale}x (current: {current_size}px, target: {target_size}px)")
    
    return upscale_image(image, scale=scale)