src/utils.py

"""Utility functions for the application."""

import logging
import random
from typing import Optional, Dict, Any
from pathlib import Path

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def generate_seed() -> int:
    """Generate a random seed for image generation.
    
    Returns:
        int: Random seed value.
    """
    return random.randint(0, 2**32 - 1)


def validate_prompt(prompt: str, max_length: int = 500) -> tuple[bool, Optional[str]]:
    """Validate a prompt string.
    
    Args:
        prompt: The prompt to validate.
        max_length: Maximum allowed length.
        
    Returns:
        Tuple of (is_valid, error_message).
    """
    if not prompt or not prompt.strip():
        return False, "Prompt cannot be empty"
    
    if len(prompt) > max_length:
        return False, f"Prompt exceeds maximum length of {max_length} characters"
    
    return True, None


def validate_dimensions(width: int, height: int) -> tuple[bool, Optional[str]]:
    """Validate image dimensions.
    
    Args:
        width: Image width.
        height: Image height.
        
    Returns:
        Tuple of (is_valid, error_message).
    """
    min_dim = 256
    max_dim = 2048
    
    if width < min_dim or width > max_dim:
        return False, f"Width must be between {min_dim} and {max_dim}"
    
    if height < min_dim or height > max_dim:
        return False, f"Height must be between {min_dim} and {max_dim}"
    
    # Check if dimensions are multiples of 8 (common requirement)
    if width % 8 != 0 or height % 8 != 0:
        return False, "Dimensions must be multiples of 8"
    
    return True, None


def validate_params(params: Dict[str, Any]) -> tuple[bool, Optional[str]]:
    """Validate generation parameters.
    
    Args:
        params: Dictionary of parameters to validate.
        
    Returns:
        Tuple of (is_valid, error_message).
    """
    # Validate guidance scale
    guidance = params.get("guidance_scale", 7.5)
    if not 1.0 <= guidance <= 20.0:
        return False, "Guidance scale must be between 1.0 and 20.0"
    
    # Validate steps
    steps = params.get("num_inference_steps", 50)
    if not 10 <= steps <= 100:
        return False, "Number of steps must be between 10 and 100"
    
    # Validate dimensions
    width = params.get("width", 1024)
    height = params.get("height", 1024)
    is_valid, error = validate_dimensions(width, height)
    if not is_valid:
        return False, error
    
    return True, None


def format_timestamp(timestamp: float) -> str:
    """Format a timestamp to a readable string.
    
    Args:
        timestamp: Unix timestamp.
        
    Returns:
        Formatted timestamp string.
    """
    from datetime import datetime
    return datetime.fromtimestamp(timestamp).strftime("%Y-%m-%d %H:%M:%S")


def ensure_directory(path: Path) -> Path:
    """Ensure a directory exists, creating it if necessary.
    
    Args:
        path: Path to the directory.
        
    Returns:
        Path object of the created/existing directory.
    """
    path.mkdir(parents=True, exist_ok=True)
    return path