utils.py

"""
Utility functions for the image generation application.

This module provides helper functions for tasks like image saving,
timestamp generation, and other common operations.
"""

import os
import time
from datetime import datetime
from pathlib import Path
from typing import Optional

import gradio as gr
from PIL import Image

# Create output directory for saved images
OUTPUTS_DIR = Path("outputs")
OUTPUTS_DIR.mkdir(exist_ok=True)


def save_image(image: Image.Image, prompt: str) -> str:
    """
    Save the generated image to disk with a filename based on timestamp and prompt.
    
    Args:
        image: The PIL Image to save
        prompt: The prompt used to generate the image
        
    Returns:
        Path to the saved image
    """
    if image is None:
        return ""
        
    # Create a filename from the timestamp and a shortened version of the prompt
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    # Clean prompt for filename use (first 20 chars, alphanumeric only)
    clean_prompt = "".join(c for c in prompt if c.isalnum() or c.isspace())[:20].strip()
    clean_prompt = clean_prompt.replace(" ", "_")
    
    filename = f"{timestamp}_{clean_prompt}.png"
    filepath = OUTPUTS_DIR / filename
    
    # Save the image
    image.save(filepath)
    return str(filepath)


def format_generation_info(
    prompt: str, 
    negative_prompt: str, 
    seed: int, 
    width: int, 
    height: int, 
    guidance_scale: float, 
    steps: int
) -> str:
    """
    Format generation parameters into a readable string.
    
    Args:
        prompt: Text prompt used
        negative_prompt: Negative prompt used
        seed: Random seed used
        width: Image width
        height: Image height
        guidance_scale: Guidance scale value
        steps: Number of inference steps
        
    Returns:
        Formatted string with generation parameters
    """
    info = f"**Prompt:** {prompt}\n"
    if negative_prompt:
        info += f"**Negative prompt:** {negative_prompt}\n"
    info += f"**Seed:** {seed}\n"
    info += f"**Size:** {width}x{height}\n"
    info += f"**Guidance scale:** {guidance_scale}\n"
    info += f"**Steps:** {steps}\n"
    return info


class GenerationHistory:
    """Manages a history of generated images and their parameters."""
    
    def __init__(self, max_history: int = 10):
        """
        Initialize the generation history.
        
        Args:
            max_history: Maximum number of items to keep in history
        """
        self.history = []
        self.max_history = max_history
        
    def add(
        self, 
        image: Image.Image, 
        prompt: str, 
        negative_prompt: str, 
        seed: int, 
        width: int, 
        height: int, 
        guidance_scale: float, 
        steps: int
    ) -> None:
        """
        Add a new generation to the history.
        
        Args:
            image: Generated image
            prompt: Text prompt used
            negative_prompt: Negative prompt used
            seed: Random seed used
            width: Image width
            height: Image height
            guidance_scale: Guidance scale value
            steps: Number of inference steps
        """
        if image is None:
            return
            
        # Create entry with all relevant information
        entry = {
            "image": image,
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "seed": seed,
            "width": width,
            "height": height,
            "guidance_scale": guidance_scale,
            "steps": steps,
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }
        
        # Add to history and maintain max size
        self.history.append(entry)
        if len(self.history) > self.max_history:
            self.history.pop(0)
            
    def get_latest(self, n: int = 1) -> list:
        """
        Get the latest n entries from history.
        
        Args:
            n: Number of entries to retrieve
            
        Returns:
            List of history entries
        """
        return self.history[-n:] if self.history else []
        
    def clear(self) -> None:
        """Clear the generation history."""
        self.history = []