upscale-ui/src/config.py

"""Configuration settings for Real-ESRGAN Web UI."""

import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional


# Base paths
BASE_DIR = Path(__file__).parent.parent
DATA_DIR = BASE_DIR / "data"
MODELS_DIR = DATA_DIR / "models"
TEMP_DIR = DATA_DIR / "temp"
CHECKPOINTS_DIR = DATA_DIR / "checkpoints"
OUTPUT_DIR = DATA_DIR / "output"
DATABASE_PATH = DATA_DIR / "upscale.db"

# Ensure directories exist
for dir_path in [MODELS_DIR, TEMP_DIR, CHECKPOINTS_DIR, OUTPUT_DIR]:
    dir_path.mkdir(parents=True, exist_ok=True)


@dataclass
class ModelInfo:
    """Information about a Real-ESRGAN or GFPGAN model."""

    name: str
    scale: int
    filename: str
    url: str
    size_mb: float
    description: str
    model_type: str = "realesrgan"  # "realesrgan" or "gfpgan"
    netscale: int = 4  # Network scale factor
    num_block: int = 23  # Number of RRDB blocks (6 for anime models)
    num_grow_ch: int = 32  # Growth channels


# Model registry with download URLs
MODEL_REGISTRY: dict[str, ModelInfo] = {
    "RealESRGAN_x4plus": ModelInfo(
        name="RealESRGAN_x4plus",
        scale=4,
        filename="RealESRGAN_x4plus.pth",
        url="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
        size_mb=63.9,
        description="Best quality for general photos (4x)",
        num_block=23,
    ),
    "RealESRGAN_x2plus": ModelInfo(
        name="RealESRGAN_x2plus",
        scale=2,
        filename="RealESRGAN_x2plus.pth",
        url="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
        size_mb=63.9,
        description="General photos (2x upscaling)",
        netscale=2,
        num_block=23,
    ),
    "RealESRGAN_x4plus_anime_6B": ModelInfo(
        name="RealESRGAN_x4plus_anime_6B",
        scale=4,
        filename="RealESRGAN_x4plus_anime_6B.pth",
        url="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth",
        size_mb=17.0,
        description="Optimized for anime/illustrations (4x)",
        num_block=6,
    ),
    "realesr-animevideov3": ModelInfo(
        name="realesr-animevideov3",
        scale=4,
        filename="realesr-animevideov3.pth",
        url="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth",
        size_mb=17.0,
        description="Anime video with temporal consistency (4x)",
        num_block=6,
    ),
    "realesr-general-x4v3": ModelInfo(
        name="realesr-general-x4v3",
        scale=4,
        filename="realesr-general-x4v3.pth",
        url="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",
        size_mb=63.9,
        description="General purpose with denoise control (4x)",
        num_block=23,
    ),
    "GFPGANv1.4": ModelInfo(
        name="GFPGANv1.4",
        scale=1,
        filename="GFPGANv1.4.pth",
        url="https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.4.pth",
        size_mb=332.0,
        description="Face enhancement and restoration",
        model_type="gfpgan",
    ),
}


@dataclass
class RTX4090Config:
    """RTX 4090 optimized settings (24GB VRAM)."""

    tile_size: int = 0  # 0 = no tiling, use full image
    tile_pad: int = 10  # Padding for tile boundaries
    pre_pad: int = 0  # Pre-padding for input
    half: bool = True  # FP16 precision
    device: str = "cuda"
    gpu_id: int = 0

    # Thresholds for when to enable tiling
    max_pixels_no_tile: int = 8294400  # ~4K (3840x2160)

    def should_tile(self, width: int, height: int) -> bool:
        """Determine if tiling is needed based on image size."""
        return width * height > self.max_pixels_no_tile

    def get_tile_size(self, width: int, height: int) -> int:
        """Get appropriate tile size for the image."""
        if not self.should_tile(width, height):
            return 0  # No tiling
        # For very large images, use 512px tiles
        return 512


@dataclass
class VideoCodecConfig:
    """Video encoding configuration."""

    encoder: str
    crf: int
    preset: str
    nvenc_encoder: Optional[str] = None
    description: str = ""


# Video codec presets (maximum quality settings)
VIDEO_CODECS: dict[str, VideoCodecConfig] = {
    "H.264": VideoCodecConfig(
        encoder="libx264",
        crf=18,
        preset="slow",
        nvenc_encoder="h264_nvenc",
        description="Universal compatibility, excellent quality",
    ),
    "H.265": VideoCodecConfig(
        encoder="libx265",
        crf=20,
        preset="slow",
        nvenc_encoder="hevc_nvenc",
        description="50% smaller files, great quality",
    ),
    "AV1": VideoCodecConfig(
        encoder="libsvtav1",
        crf=23,
        preset="4",
        nvenc_encoder="av1_nvenc",
        description="Best compression, future-proof",
    ),
}


@dataclass
class AppConfig:
    """Main application configuration."""

    # Server settings (configurable via environment variables)
    server_name: str = os.getenv("UPSCALE_HOST", "0.0.0.0")
    server_port: int = int(os.getenv("UPSCALE_PORT", "7860"))
    share: bool = os.getenv("UPSCALE_SHARE", "").lower() in ("true", "1", "yes")

    # Queue settings
    max_queue_size: int = 20
    concurrency_limit: int = 1  # Single job at a time for GPU efficiency

    # Processing defaults
    default_model: str = "RealESRGAN_x4plus"
    default_scale: int = 4
    default_face_enhance: bool = False
    default_output_format: str = "png"

    # Video defaults
    default_video_codec: str = "H.265"
    default_video_crf: int = 20
    default_video_preset: str = "slow"

    # History settings
    max_history_items: int = 1000
    thumbnail_size: tuple[int, int] = (256, 256)

    # Checkpoint settings
    checkpoint_interval: int = 100  # Save every N frames

    # GPU config
    gpu: RTX4090Config = field(default_factory=RTX4090Config)


# Global config instance
config = AppConfig()


def get_model_path(model_name: str) -> Path:
    """Get the path to a model file."""
    if model_name not in MODEL_REGISTRY:
        raise ValueError(f"Unknown model: {model_name}")
    return MODELS_DIR / MODEL_REGISTRY[model_name].filename


def get_available_models() -> list[str]:
    """Get list of available model names for UI dropdown."""
    return [
        name
        for name, info in MODEL_REGISTRY.items()
        if info.model_type == "realesrgan"
    ]


def get_model_choices() -> list[tuple[str, str]]:
    """Get model choices for Gradio dropdown (label, value)."""
    return [
        (f"{info.description}", name)
        for name, info in MODEL_REGISTRY.items()
        if info.model_type == "realesrgan"
    ]