""" Lemontropia Suite - AI Image Upscaling with Real-ESRGAN Optional AI-powered upscaling for game icons and textures. Real-ESRGAN is specifically designed for low-resolution game graphics and produces excellent results for rendered icons (not pixel art). """ import logging from pathlib import Path from typing import Optional, Tuple import numpy as np try: from PIL import Image PIL_AVAILABLE = True except ImportError: PIL_AVAILABLE = False logger = logging.getLogger(__name__) # Try to import Real-ESRGAN try: import torch from basicsr.archs.rrdbnet_arch import RRDBNet from realesrgan import RealESRGANer REALESRGAN_AVAILABLE = True except ImportError: REALESRGAN_AVAILABLE = False logger.info("Real-ESRGAN not available. Install with: pip install realesrgan") class AIIconUpscaler: """ AI-powered upscaler for game icons using Real-ESRGAN. Real-ESRGAN is trained specifically on game textures and produces excellent results for low-resolution rendered graphics (not pixel art). Usage: upscaler = AIIconUpscaler() if upscaler.is_available(): result = upscaler.upscale(image, scale=4) """ # Model download URLs MODEL_URLS = { 'RealESRGAN_x4plus': 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth', 'RealESRGAN_x4plus_anime_6B': 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus_anime_6B.pth', 'RealESRGAN_x2plus': 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x2plus.pth', } def __init__(self, model_name: str = 'RealESRGAN_x4plus', device: str = 'cpu'): """ Initialize AI upscaler. Args: model_name: Which model to use device: 'cpu' or 'cuda' (GPU) """ self.model_name = model_name self.device = device self.upsampler = None if REALESRGAN_AVAILABLE: self._init_model() def _init_model(self): """Initialize the Real-ESRGAN model.""" try: # Model parameters if 'anime' in self.model_name: # Anime model (6B parameters) model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4) elif 'x2' in self.model_name: # 2x upscale model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2) else: # 4x upscale (default) model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4) # Get model path model_path = self._get_model_path() if not model_path.exists(): logger.warning(f"Model not found: {model_path}") logger.info(f"Download from: {self.MODEL_URLS.get(self.model_name)}") return # Initialize upsampler self.upsampler = RealESRGANer( scale=4 if 'x4' in self.model_name else 2, model_path=str(model_path), model=model, tile=0, # No tiling (process whole image) tile_pad=10, pre_pad=0, half=False if self.device == 'cpu' else True, device=torch.device(self.device) ) logger.info(f"Real-ESRGAN initialized: {self.model_name} on {self.device}") except Exception as e: logger.error(f"Failed to initialize Real-ESRGAN: {e}") self.upsampler = None def _get_model_path(self) -> Path: """Get path to model file.""" # Store models in user's home directory model_dir = Path.home() / ".lemontropia" / "models" model_dir.mkdir(parents=True, exist_ok=True) return model_dir / f"{self.model_name}.pth" def is_available(self) -> bool: """Check if AI upscaler is available and ready.""" return REALESRGAN_AVAILABLE and self.upsampler is not None def upscale(self, image: Image.Image, scale: int = 4) -> Optional[Image.Image]: """ Upscale an image using Real-ESRGAN. Args: image: PIL Image to upscale scale: Upscale factor (2 or 4) Returns: Upscaled PIL Image or None if failed """ if not self.is_available(): logger.error("AI upscaler not available") return None try: # Convert PIL to numpy img_np = np.array(image) # Remove alpha channel if present (Real-ESRGAN expects RGB) has_alpha = img_np.shape[-1] == 4 if has_alpha: alpha = img_np[:, :, 3] img_rgb = img_np[:, :, :3] else: img_rgb = img_np # Upscale output, _ = self.upsampler.enhance(img_rgb, outscale=scale) # Restore alpha channel if needed if has_alpha: # Upscale alpha with simple resize from PIL import Image as PILImage alpha_pil = PILImage.fromarray(alpha) alpha_upscaled = alpha_pil.resize( (output.shape[1], output.shape[0]), PILImage.Resampling.LANCZOS ) alpha_np = np.array(alpha_upscaled) output = np.dstack([output, alpha_np]) # Convert back to PIL result = Image.fromarray(output) return result except Exception as e: logger.error(f"Upscaling failed: {e}") return None def get_info(self) -> dict: """Get information about the upscaler status.""" return { 'available': self.is_available(), 'model': self.model_name, 'device': self.device, 'model_path': str(self._get_model_path()), 'dependencies': { 'torch': torch.__version__ if REALESRGAN_AVAILABLE else 'not installed', 'realesrgan': REALESRGAN_AVAILABLE } } # Convenience function def upscale_with_ai(image: Image.Image, scale: int = 4) -> Optional[Image.Image]: """Quick AI upscale using Real-ESRGAN.""" upscaler = AIIconUpscaler() return upscaler.upscale(image, scale) def check_ai_upscaler(): """Check if AI upscaler is available and print status.""" upscaler = AIIconUpscaler() info = upscaler.get_info() print("=" * 50) print("AI Upscaler Status (Real-ESRGAN)") print("=" * 50) print(f"Available: {info['available']}") print(f"Model: {info['model']}") print(f"Device: {info['device']}") print(f"Model Path: {info['model_path']}") print(f"PyTorch: {info['dependencies']['torch']}") print("=" * 50) if not info['available']: print("\nTo install AI upscaler:") print("1. pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu") print("2. pip install realesrgan") print("3. Download model from:") for name, url in AIIconUpscaler.MODEL_URLS.items(): print(f" {name}: {url}") return info['available'] if __name__ == "__main__": check_ai_upscaler()