SnapOtter/packages/ai/python/upscale.py

"""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()