backend/config.py

"""
Configuration module for Architecture AI Enhancer
Centralized configuration management for the entire application
"""
import os
from pathlib import Path
from typing import Optional
from pydantic_settings import BaseSettings


class Settings(BaseSettings):
    """
    Application settings with environment variable support
    
    All settings can be overridden via environment variables
    """
    
    # Application Metadata
    APP_NAME: str = "Architecture AI Enhancer"
    APP_VERSION: str = "1.0.0"
    DEBUG: bool = False
    
    # Server Configuration
    HOST: str = "0.0.0.0"
    PORT: int = 8000
    WORKERS: int = 1
    
    # CORS Settings
    CORS_ORIGINS: list = [
        "http://localhost:3000",
        "http://localhost:5173",
        "http://127.0.0.1:3000",
        "http://127.0.0.1:5173"
    ]
    
    # Path Configuration
    BASE_DIR: Path = Path(__file__).parent
    MODELS_DIR: Path = BASE_DIR / "models"
    LORA_DIR: Path = MODELS_DIR / "lora"
    BASE_MODEL_DIR: Path = MODELS_DIR / "base"
    DATASETS_DIR: Path = BASE_DIR / "datasets"
    INPUT_DIR: Path = DATASETS_DIR / "input"
    TARGET_DIR: Path = DATASETS_DIR / "target"
    PROCESSED_DIR: Path = DATASETS_DIR / "processed"
    OUTPUT_DIR: Path = BASE_DIR / "output"
    ENHANCED_DIR: Path = OUTPUT_DIR / "enhanced"
    LOGS_DIR: Path = OUTPUT_DIR / "logs"
    
    # Model Configuration
    BASE_MODEL: str = "runwayml/stable-diffusion-v1-5"  # Lighter model (~4GB RAM vs ~12GB for SDXL)
    LORA_MODEL_NAME: str = "office_style.safetensors"
    VAE_MODEL: Optional[str] = None
    
    # Training Hyperparameters
    LORA_RANK: int = 8
    LEARNING_RATE: float = 1e-4
    TRAIN_STEPS: int = 1000
    BATCH_SIZE: int = 1
    GRADIENT_ACCUMULATION_STEPS: int = 4
    MAX_GRAD_NORM: float = 1.0
    WARMUP_STEPS: int = 100
    SAVE_STEPS: int = 250
    
    # Inference Configuration
    IMG2IMG_STRENGTH: float = 0.3  # Range: 0.2-0.35
    GUIDANCE_SCALE: float = 5.5    # Range: 4-7
    NUM_INFERENCE_STEPS: int = 30
    
    # Image Processing
    MAX_IMAGE_SIZE: int = 2048
    UPSCALE_FACTOR: int = 2
    TARGET_RESOLUTION: int = 512  # SD 1.5 works best at 512x512
    
    # Prompts
    DEFAULT_PROMPT: str = (
        "ultra realistic architectural visualization, "
        "professional photography, high detail, sharp focus, "
        "natural lighting, modern office interior, "
        "clean lines, photorealistic rendering"
    )
    
    NEGATIVE_PROMPT: str = (
        "distorted walls, cartoon, illustration, painting, drawing, "
        "unrealistic proportions, blurry, low quality, artifacts, "
        "oversaturated, noise, grain, ugly, deformed"
    )
    
    # Device Configuration
    DEVICE: str = "cuda"  # "cuda" or "cpu" - auto-detected at runtime
    MIXED_PRECISION: str = "fp16"  # "fp16", "bf16", or "no"
    ENABLE_ATTENTION_SLICING: bool = True  # Reduce VRAM usage
    ENABLE_VAE_SLICING: bool = True  # Reduce VRAM usage
    
    # Upload Limits
    MAX_UPLOAD_SIZE: int = 25 * 1024 * 1024  # 25 MB
    ALLOWED_EXTENSIONS: set = {".png", ".jpg", ".jpeg", ".webp"}
    
    class Config:
        env_file = ".env"
        case_sensitive = True


# Global settings instance
settings = Settings()


def ensure_directories():
    """
    Create all necessary directories if they don't exist
    
    This function should be called on application startup
    """
    directories = [
        settings.MODELS_DIR,
        settings.LORA_DIR,
        settings.BASE_MODEL_DIR,
        settings.DATASETS_DIR,
        settings.INPUT_DIR,
        settings.TARGET_DIR,
        settings.PROCESSED_DIR,
        settings.OUTPUT_DIR,
        settings.ENHANCED_DIR,
        settings.LOGS_DIR,
    ]
    
    for directory in directories:
        directory.mkdir(parents=True, exist_ok=True)
        print(f"✓ Ensured directory exists: {directory}")


def get_lora_path() -> Optional[Path]:
    """
    Get the path to the trained LoRA model if it exists
    
    Returns:
        Path to LoRA model or None if not found
    """
    lora_path = settings.LORA_DIR / settings.LORA_MODEL_NAME
    return lora_path if lora_path.exists() else None


def validate_image_file(filename: str) -> bool:
    """
    Validate if a file has an allowed image extension
    
    Args:
        filename: Name of the file to validate
        
    Returns:
        True if valid, False otherwise
    """
    return Path(filename).suffix.lower() in settings.ALLOWED_EXTENSIONS


if __name__ == "__main__":
    # Test configuration
    ensure_directories()
    print(f"\n{settings.APP_NAME} v{settings.APP_VERSION}")
    print(f"Base Model: {settings.BASE_MODEL}")
    print(f"LoRA Path: {get_lora_path()}")