"""Image upscaling with Real-ESRGAN.""" import sys import json import os import types # basicsr imports torchvision.transforms.functional_tensor which was removed # in torchvision >= 0.17. This shim must exist before basicsr is imported. try: import torchvision.transforms.functional_tensor # noqa: F401 except (ImportError, ModuleNotFoundError): try: import torchvision.transforms.functional as _F import torchvision.transforms _shim = types.ModuleType("torchvision.transforms.functional_tensor") for _attr in dir(_F): if not _attr.startswith("_"): setattr(_shim, _attr, getattr(_F, _attr)) sys.modules["torchvision.transforms.functional_tensor"] = _shim torchvision.transforms.functional_tensor = _shim except ImportError: pass def emit_progress(percent, stage): """Emit structured progress to stderr for bridge.ts to capture.""" print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True) _MODELS_BASE = os.environ.get("MODELS_PATH", "/opt/models") REALESRGAN_MODEL_PATH = os.environ.get( "REALESRGAN_MODEL_PATH", os.path.join(_MODELS_BASE, "realesrgan", "RealESRGAN_x4plus.pth"), ) GFPGAN_MODEL_PATH = os.environ.get( "GFPGAN_MODEL_PATH", os.path.join(_MODELS_BASE, "gfpgan", "GFPGANv1.3.pth"), ) def apply_denoise(img, strength): """Apply denoising to a PIL image. Uses OpenCV when available, falls back to PIL.""" if strength <= 0: return img try: import numpy as np import cv2 from PIL import Image arr = np.array(img) # Map 0-1 strength to filter parameter (3-15 range) h = int(3 + strength * 12) if len(arr.shape) == 3 and arr.shape[2] >= 3: denoised = cv2.fastNlMeansDenoisingColored(arr, None, h, h, 7, 21) else: denoised = cv2.fastNlMeansDenoising(arr, None, h, 7, 21) return Image.fromarray(denoised) except ImportError: from PIL import ImageFilter radius = max(0.5, strength * 1.5) return img.filter(ImageFilter.GaussianBlur(radius=radius)) def main(): input_path = sys.argv[1] output_path = sys.argv[2] settings = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {} scale = settings.get("scale", 2) model_choice = settings.get("model", "auto") face_enhance = settings.get("faceEnhance", False) denoise_strength = float(settings.get("denoise", 0)) output_format = settings.get("format", "png") quality = int(settings.get("quality", 95)) try: emit_progress(5, "Opening image") from PIL import Image img = Image.open(input_path) new_size = (img.width * scale, img.height * scale) method = "lanczos" result = None # Try Real-ESRGAN if requested if model_choice in ("auto", "realesrgan"): try: emit_progress(10, "Loading AI model") # Redirect stdout to stderr for the ENTIRE AI pipeline. # Libraries like basicsr, realesrgan, gfpgan, and torch print # download progress and init messages to stdout which would # corrupt our JSON result. stdout_fd = None try: stdout_fd = os.dup(1) os.dup2(2, 1) except OSError: # os.dup may fail on Windows with piped stdio stdout_fd = None sys.stdout = sys.stderr try: from basicsr.archs.rrdbnet_arch import RRDBNet from realesrgan import RealESRGANer from gpu import gpu_available import numpy as np import torch if not os.path.exists(REALESRGAN_MODEL_PATH): raise FileNotFoundError( f"RealESRGAN model not found: {REALESRGAN_MODEL_PATH}" ) use_gpu = gpu_available() device = torch.device("cuda" if use_gpu else "cpu") # RealESRGAN_x4plus is a 4x model internally ai_model = RRDBNet( num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4, ) tile_size = 512 upsampler = RealESRGANer( scale=4, model_path=REALESRGAN_MODEL_PATH, model=ai_model, tile=tile_size, tile_pad=10, half=use_gpu, device=device, ) emit_progress(20, "AI model loaded") img_array = np.array(img.convert("RGB")) emit_progress(30, "Enhancing image with AI") try: output_array, _ = upsampler.enhance(img_array, outscale=scale) except RuntimeError as oom_err: if "out of memory" not in str(oom_err).lower(): raise torch.cuda.empty_cache() print( f"[upscale] OOM with tile={tile_size}, retrying with tile=256", file=sys.stderr, flush=True, ) emit_progress(35, "Retrying with smaller tiles") upsampler.tile = 256 output_array, _ = upsampler.enhance(img_array, outscale=scale) emit_progress(80, "AI enhancement complete") result = Image.fromarray(output_array) method = "realesrgan" if face_enhance: emit_progress(82, "Enhancing faces") from gfpgan import GFPGANer if not os.path.exists(GFPGAN_MODEL_PATH): raise FileNotFoundError( f"GFPGAN model not found at {GFPGAN_MODEL_PATH}. " "Install the upscale-enhance feature or disable faceEnhance." ) face_enhancer = GFPGANer( model_path=GFPGAN_MODEL_PATH, upscale=scale, arch="clean", channel_multiplier=2, bg_upsampler=upsampler, ) try: _, _, face_output = face_enhancer.enhance( img_array, has_aligned=False, only_center_face=False, paste_back=True, ) except RuntimeError as oom_err: if "out of memory" not in str(oom_err).lower(): raise torch.cuda.empty_cache() print( "[upscale] OOM during face enhance, retrying with tile=256", file=sys.stderr, flush=True, ) emit_progress(84, "Retrying face enhancement") upsampler.tile = 256 face_enhancer.bg_upsampler = upsampler _, _, face_output = face_enhancer.enhance( img_array, has_aligned=False, only_center_face=False, paste_back=True, ) result = Image.fromarray(face_output) emit_progress(88, "Face enhancement complete") finally: # Restore stdout after ALL AI processing if stdout_fd is not None: os.dup2(stdout_fd, 1) os.close(stdout_fd) sys.stdout = sys.__stdout__ except (ImportError, FileNotFoundError, RuntimeError, OSError) as e: import traceback print(f"[upscale] Real-ESRGAN failed: {e}", file=sys.stderr, flush=True) traceback.print_exc(file=sys.stderr) print(json.dumps({ "success": False, "error": ( f"Real-ESRGAN is not available: {e}. " "Install the upscale-enhance feature or use model=lanczos for basic upscaling." ), })) sys.exit(1) if result is None and model_choice == "lanczos": emit_progress(50, "Upscaling with Lanczos") result = img.resize(new_size, Image.LANCZOS) method = "lanczos" if result is None: raise RuntimeError(f"Requested model '{model_choice}' is not available") # Denoise if denoise_strength > 0: emit_progress(90, "Reducing noise") result = apply_denoise(result, denoise_strength) # Determine final output path based on format base_path = output_path.rsplit(".", 1)[0] EXT_MAP = { "jpeg": ".jpg", "jpg": ".jpg", "png": ".png", "webp": ".webp", "tiff": ".tiff", "gif": ".gif", } final_path = base_path + EXT_MAP.get(output_format, ".png") # Save with format-specific options emit_progress(95, "Saving result") save_kwargs = {} if output_format in ("jpeg", "jpg"): result = result.convert("RGB") # Strip alpha for JPEG save_kwargs["quality"] = quality save_kwargs["optimize"] = True elif output_format == "webp": save_kwargs["quality"] = quality elif output_format == "tiff": save_kwargs["compression"] = "tiff_lzw" elif output_format == "gif": result = result.convert("P", palette=Image.ADAPTIVE, colors=256) result.save(final_path, **save_kwargs) # Get actual dimensions of the saved result actual_w, actual_h = result.size print( json.dumps( { "success": True, "scale": scale, "width": actual_w, "height": actual_h, "method": method, "output_path": final_path, "format": output_format, } ) ) except ImportError: print( json.dumps( { "success": False, "error": "Pillow is not installed. Install with: pip install Pillow", } ) ) sys.exit(1) except Exception as e: print(json.dumps({"success": False, "error": str(e)})) sys.exit(1) if __name__ == "__main__": main()