"""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" ]