Create utils.py

utils.py

@@ -0,0 +1,897 @@
+#!/usr/bin/env python3
+"""
+Trouter-Imagine-1 Utilities and Helper Functions
+Apache 2.0 License
+
+Comprehensive utility module providing:
+- Prompt enhancement and optimization
+- Image post-processing
+- Metadata management
+- Performance monitoring
+- Configuration management
+- Quality assessment
+- Batch processing helpers
+- File management
+- API wrappers
+- Advanced preprocessing
+"""
+
+import torch
+from PIL import Image, ImageEnhance, ImageFilter, ImageDraw, ImageFont
+import numpy as np
+from typing import List, Dict, Tuple, Optional, Union
+import json
+import os
+import hashlib
+from pathlib import Path
+from datetime import datetime
+import re
+import logging
+from dataclasses import dataclass, asdict
+import time
+from collections import defaultdict
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+# ============================================================================
+# DATA CLASSES FOR CONFIGURATION
+# ============================================================================
+
+@dataclass
+class GenerationConfig:
+    """Configuration for image generation"""
+    prompt: str
+    negative_prompt: str = ""
+    width: int = 512
+    height: int = 512
+    num_inference_steps: int = 30
+    guidance_scale: float = 7.5
+    seed: Optional[int] = None
+    num_images: int = 1
+    
+    def to_dict(self) -> Dict:
+        return asdict(self)
+    
+    @classmethod
+    def from_dict(cls, data: Dict) -> 'GenerationConfig':
+        return cls(**data)
+    
+    def validate(self) -> bool:
+        """Validate configuration parameters"""
+        if self.width % 8 != 0 or self.height % 8 != 0:
+            raise ValueError("Width and height must be multiples of 8")
+        if self.num_inference_steps < 1:
+            raise ValueError("num_inference_steps must be at least 1")
+        if self.guidance_scale < 0:
+            raise ValueError("guidance_scale must be positive")
+        return True
+
+
+@dataclass
+class GenerationMetadata:
+    """Metadata for generated images"""
+    prompt: str
+    negative_prompt: str
+    model_id: str
+    width: int
+    height: int
+    num_inference_steps: int
+    guidance_scale: float
+    seed: int
+    timestamp: str
+    generation_time: float
+    scheduler: str = "unknown"
+    
+    def to_json(self) -> str:
+        return json.dumps(asdict(self), indent=2)
+    
+    @classmethod
+    def from_json(cls, json_str: str) -> 'GenerationMetadata':
+        return cls(**json.loads(json_str))
+
+
+# ============================================================================
+# PROMPT ENHANCEMENT
+# ============================================================================
+
+class PromptEnhancer:
+    """Enhance and optimize prompts for better generation"""
+    
+    QUALITY_BOOSTERS = [
+        "highly detailed",
+        "professional",
+        "4k",
+        "ultra detailed",
+        "sharp focus",
+        "intricate details"
+    ]
+    
+    STYLE_KEYWORDS = {
+        "photo": ["photography", "realistic", "photorealistic", "sharp focus"],
+        "art": ["digital art", "concept art", "artistic", "detailed"],
+        "paint": ["oil painting", "painterly", "brushstrokes", "canvas"],
+        "anime": ["anime style", "manga", "cel shaded", "vibrant"],
+        "3d": ["3d render", "octane render", "unreal engine", "cgi"]
+    }
+    
+    NEGATIVE_DEFAULTS = [
+        "blurry", "low quality", "distorted", "deformed", 
+        "ugly", "bad anatomy", "watermark", "signature"
+    ]
+    
+    @staticmethod
+    def enhance_prompt(
+        prompt: str,
+        style: Optional[str] = None,
+        add_quality: bool = True,
+        add_details: bool = True
+    ) -> str:
+        """
+        Enhance a prompt with quality boosters and style keywords
+        
+        Args:
+            prompt: Base prompt
+            style: Style to apply (photo, art, paint, anime, 3d)
+            add_quality: Add quality boosters
+            add_details: Add detail-related keywords
+        
+        Returns:
+            Enhanced prompt
+        """
+        enhanced = prompt.strip()
+        
+        # Add style keywords
+        if style and style.lower() in PromptEnhancer.STYLE_KEYWORDS:
+            style_words = PromptEnhancer.STYLE_KEYWORDS[style.lower()]
+            enhanced += ", " + ", ".join(style_words[:2])
+        
+        # Add quality boosters
+        if add_quality:
+            quality_words = PromptEnhancer.QUALITY_BOOSTERS[:3]
+            enhanced += ", " + ", ".join(quality_words)
+        
+        return enhanced
+    
+    @staticmethod
+    def build_negative_prompt(
+        base_negative: str = "",
+        include_defaults: bool = True,
+        subject_type: Optional[str] = None
+    ) -> str:
+        """
+        Build a comprehensive negative prompt
+        
+        Args:
+            base_negative: User-provided negative prompt
+            include_defaults: Include default negative terms
+            subject_type: Type of subject (person, animal, landscape, etc.)
+        
+        Returns:
+            Enhanced negative prompt
+        """
+        negatives = []
+        
+        if base_negative:
+            negatives.append(base_negative)
+        
+        if include_defaults:
+            negatives.extend(PromptEnhancer.NEGATIVE_DEFAULTS)
+        
+        # Subject-specific negatives
+        subject_negatives = {
+            "person": ["extra limbs", "extra fingers", "fused fingers", "bad hands"],
+            "animal": ["extra legs", "incorrect anatomy", "fused limbs"],
+            "face": ["asymmetric eyes", "crossed eyes", "bad teeth"],
+            "landscape": ["oversaturated", "underexposed", "poor composition"]
+        }
+        
+        if subject_type and subject_type.lower() in subject_negatives:
+            negatives.extend(subject_negatives[subject_type.lower()])
+        
+        return ", ".join(negatives)
+    
+    @staticmethod
+    def extract_keywords(prompt: str) -> List[str]:
+        """Extract important keywords from a prompt"""
+        # Remove common words
+        stop_words = {'a', 'an', 'the', 'in', 'on', 'at', 'with', 'by', 'for'}
+        words = prompt.lower().split()
+        keywords = [w.strip('.,!?;:') for w in words if w not in stop_words]
+        return keywords
+    
+    @staticmethod
+    def validate_prompt(prompt: str) -> Tuple[bool, List[str]]:
+        """
+        Validate a prompt and return warnings
+        
+        Returns:
+            (is_valid, list_of_warnings)
+        """
+        warnings = []
+        
+        if len(prompt.strip()) < 3:
+            warnings.append("Prompt is very short, consider adding more detail")
+        
+        if len(prompt) > 500:
+            warnings.append("Prompt is very long, may be truncated")
+        
+        # Check for common issues
+        if "high quality" in prompt.lower() and "low quality" in prompt.lower():
+            warnings.append("Contradictory quality terms detected")
+        
+        # Check for excessive punctuation
+        if prompt.count(',') > 20:
+            warnings.append("Too many commas, consider simplifying")
+        
+        return len(warnings) == 0, warnings
+
+
+# ============================================================================
+# IMAGE POST-PROCESSING
+# ============================================================================
+
+class ImageProcessor:
+    """Post-processing utilities for generated images"""
+    
+    @staticmethod
+    def enhance_image(
+        image: Image.Image,
+        brightness: float = 1.0,
+        contrast: float = 1.0,
+        saturation: float = 1.0,
+        sharpness: float = 1.0
+    ) -> Image.Image:
+        """
+        Enhance image with various adjustments
+        
+        Args:
+            image: Input PIL Image
+            brightness: Brightness factor (1.0 = no change)
+            contrast: Contrast factor
+            saturation: Color saturation factor
+            sharpness: Sharpness factor
+        
+        Returns:
+            Enhanced image
+        """
+        enhanced = image
+        
+        if brightness != 1.0:
+            enhancer = ImageEnhance.Brightness(enhanced)
+            enhanced = enhancer.enhance(brightness)
+        
+        if contrast != 1.0:
+            enhancer = ImageEnhance.Contrast(enhanced)
+            enhanced = enhancer.enhance(contrast)
+        
+        if saturation != 1.0:
+            enhancer = ImageEnhance.Color(enhanced)
+            enhanced = enhancer.enhance(saturation)
+        
+        if sharpness != 1.0:
+            enhancer = ImageEnhance.Sharpness(enhanced)
+            enhanced = enhancer.enhance(sharpness)
+        
+        return enhanced
+    
+    @staticmethod
+    def apply_filter(
+        image: Image.Image,
+        filter_type: str = "none"
+    ) -> Image.Image:
+        """
+        Apply various filters to image
+        
+        Args:
+            image: Input image
+            filter_type: Type of filter (blur, sharpen, edge_enhance, smooth, detail)
+        
+        Returns:
+            Filtered image
+        """
+        filters = {
+            "blur": ImageFilter.BLUR,
+            "sharpen": ImageFilter.SHARPEN,
+            "edge_enhance": ImageFilter.EDGE_ENHANCE,
+            "edge_enhance_more": ImageFilter.EDGE_ENHANCE_MORE,
+            "smooth": ImageFilter.SMOOTH,
+            "smooth_more": ImageFilter.SMOOTH_MORE,
+            "detail": ImageFilter.DETAIL
+        }
+        
+        if filter_type.lower() in filters:
+            return image.filter(filters[filter_type.lower()])
+        
+        return image
+    
+    @staticmethod
+    def upscale_simple(
+        image: Image.Image,
+        scale: int = 2,
+        method: str = "lanczos"
+    ) -> Image.Image:
+        """Simple upscaling using PIL"""
+        methods = {
+            "lanczos": Image.LANCZOS,
+            "bicubic": Image.BICUBIC,
+            "bilinear": Image.BILINEAR,
+            "nearest": Image.NEAREST
+        }
+        
+        resample = methods.get(method.lower(), Image.LANCZOS)
+        new_size = (image.width * scale, image.height * scale)
+        return image.resize(new_size, resample=resample)
+    
+    @staticmethod
+    def add_watermark(
+        image: Image.Image,
+        text: str,
+        position: str = "bottom-right",
+        opacity: int = 128
+    ) -> Image.Image:
+        """Add text watermark to image"""
+        watermark = image.copy()
+        draw = ImageDraw.Draw(watermark, 'RGBA')
+        
+        # Try to load a font
+        try:
+            font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 20)
+        except:
+            font = ImageFont.load_default()
+        
+        # Calculate position
+        bbox = draw.textbbox((0, 0), text, font=font)
+        text_width = bbox[2] - bbox[0]
+        text_height = bbox[3] - bbox[1]
+        
+        positions = {
+            "top-left": (10, 10),
+            "top-right": (image.width - text_width - 10, 10),
+            "bottom-left": (10, image.height - text_height - 10),
+            "bottom-right": (image.width - text_width - 10, image.height - text_height - 10),
+            "center": ((image.width - text_width) // 2, (image.height - text_height) // 2)
+        }
+        
+        pos = positions.get(position, positions["bottom-right"])
+        
+        # Draw with opacity
+        draw.text(pos, text, fill=(255, 255, 255, opacity), font=font)
+        
+        return watermark
+    
+    @staticmethod
+    def create_comparison(
+        images: List[Image.Image],
+        labels: Optional[List[str]] = None,
+        padding: int = 10
+    ) -> Image.Image:
+        """Create side-by-side comparison of images"""
+        if not images:
+            raise ValueError("No images provided")
+        
+        # Ensure all images have same height
+        max_height = max(img.height for img in images)
+        resized_images = []
+        
+        for img in images:
+            if img.height != max_height:
+                ratio = max_height / img.height
+                new_width = int(img.width * ratio)
+                img = img.resize((new_width, max_height), Image.LANCZOS)
+            resized_images.append(img)
+        
+        # Calculate total width
+        total_width = sum(img.width for img in resized_images) + padding * (len(resized_images) - 1)
+        
+        # Create comparison image
+        comparison = Image.new('RGB', (total_width, max_height), color='white')
+        
+        x_offset = 0
+        for i, img in enumerate(resized_images):
+            comparison.paste(img, (x_offset, 0))
+            
+            # Add label if provided
+            if labels and i < len(labels):
+                draw = ImageDraw.Draw(comparison)
+                draw.text((x_offset + 10, 10), labels[i], fill='white')
+            
+            x_offset += img.width + padding
+        
+        return comparison
+    
+    @staticmethod
+    def get_image_stats(image: Image.Image) -> Dict:
+        """Get statistical information about an image"""
+        img_array = np.array(image)
+        
+        stats = {
+            "size": image.size,
+            "mode": image.mode,
+            "mean_brightness": np.mean(img_array),
+            "std_brightness": np.std(img_array),
+            "min_value": np.min(img_array),
+            "max_value": np.max(img_array)
+        }
+        
+        if len(img_array.shape) == 3:
+            stats["mean_per_channel"] = np.mean(img_array, axis=(0, 1)).tolist()
+        
+        return stats
+
+
+# ============================================================================
+# METADATA MANAGEMENT
+# ============================================================================
+
+class MetadataManager:
+    """Manage image metadata"""
+    
+    @staticmethod
+    def embed_metadata(
+        image: Image.Image,
+        metadata: Union[Dict, GenerationMetadata]
+    ) -> Image.Image:
+        """Embed metadata into image"""
+        from PIL import PngImagePlugin
+        
+        png_info = PngImagePlugin.PngInfo()
+        
+        if isinstance(metadata, GenerationMetadata):
+            metadata = asdict(metadata)
+        
+        for key, value in metadata.items():
+            png_info.add_text(key, str(value))
+        
+        return image, png_info
+    
+    @staticmethod
+    def extract_metadata(image_path: str) -> Dict:
+        """Extract metadata from saved image"""
+        image = Image.open(image_path)
+        metadata = {}
+        
+        if hasattr(image, 'text'):
+            metadata = dict(image.text)
+        
+        return metadata
+    
+    @staticmethod
+    def save_metadata_json(
+        metadata: Union[Dict, GenerationMetadata],
+        filepath: str
+    ):
+        """Save metadata to separate JSON file"""
+        if isinstance(metadata, GenerationMetadata):
+            metadata = asdict(metadata)
+        
+        with open(filepath, 'w') as f:
+            json.dump(metadata, f, indent=2)
+    
+    @staticmethod
+    def load_metadata_json(filepath: str) -> Dict:
+        """Load metadata from JSON file"""
+        with open(filepath, 'r') as f:
+            return json.load(f)
+
+
+# ============================================================================
+# PERFORMANCE MONITORING
+# ============================================================================
+
+class PerformanceMonitor:
+    """Monitor and log generation performance"""
+    
+    def __init__(self):
+        self.generation_times = []
+        self.memory_usage = []
+        self.start_time = None
+    
+    def start(self):
+        """Start timing"""
+        self.start_time = time.time()
+    
+    def stop(self) -> float:
+        """Stop timing and return elapsed time"""
+        if self.start_time is None:
+            return 0.0
+        elapsed = time.time() - self.start_time
+        self.generation_times.append(elapsed)
+        self.start_time = None
+        return elapsed
+    
+    def get_gpu_memory(self) -> Dict:
+        """Get current GPU memory usage"""
+        if not torch.cuda.is_available():
+            return {"available": False}
+        
+        return {
+            "allocated": torch.cuda.memory_allocated() / 1024**3,  # GB
+            "reserved": torch.cuda.memory_reserved() / 1024**3,
+            "max_allocated": torch.cuda.max_memory_allocated() / 1024**3
+        }
+    
+    def get_statistics(self) -> Dict:
+        """Get performance statistics"""
+        if not self.generation_times:
+            return {"no_data": True}
+        
+        return {
+            "total_generations": len(self.generation_times),
+            "total_time": sum(self.generation_times),
+            "average_time": np.mean(self.generation_times),
+            "min_time": min(self.generation_times),
+            "max_time": max(self.generation_times),
+            "std_time": np.std(self.generation_times)
+        }
+    
+    def reset(self):
+        """Reset all statistics"""
+        self.generation_times = []
+        self.memory_usage = []
+        self.start_time = None
+
+
+# ============================================================================
+# CONFIGURATION MANAGEMENT
+# ============================================================================
+
+class ConfigManager:
+    """Manage configuration files"""
+    
+    @staticmethod
+    def load_config(filepath: str) -> Dict:
+        """Load configuration from JSON file"""
+        with open(filepath, 'r') as f:
+            return json.load(f)
+    
+    @staticmethod
+    def save_config(config: Dict, filepath: str):
+        """Save configuration to JSON file"""
+        with open(filepath, 'w') as f:
+            json.dump(config, f, indent=2)
+    
+    @staticmethod
+    def create_default_config() -> Dict:
+        """Create default configuration"""
+        return {
+            "model_id": "OpenTrouter/Trouter-Imagine-1",
+            "device": "cuda",
+            "dtype": "float16",
+            "defaults": {
+                "width": 512,
+                "height": 512,
+                "num_inference_steps": 30,
+                "guidance_scale": 7.5
+            },
+            "optimization": {
+                "attention_slicing": True,
+                "vae_slicing": True,
+                "xformers": True
+            },
+            "output": {
+                "format": "png",
+                "quality": 95,
+                "save_metadata": True
+            }
+        }
+    
+    @staticmethod
+    def validate_config(config: Dict) -> Tuple[bool, List[str]]:
+        """Validate configuration"""
+        errors = []
+        
+        required_keys = ["model_id", "device", "defaults"]
+        for key in required_keys:
+            if key not in config:
+                errors.append(f"Missing required key: {key}")
+        
+        if "device" in config:
+            valid_devices = ["cuda", "cpu", "mps"]
+            if config["device"] not in valid_devices:
+                errors.append(f"Invalid device: {config['device']}")
+        
+        return len(errors) == 0, errors
+
+
+# ============================================================================
+# BATCH PROCESSING HELPERS
+# ============================================================================
+
+class BatchProcessor:
+    """Helper for batch processing operations"""
+    
+    @staticmethod
+    def load_prompts_from_file(filepath: str) -> List[str]:
+        """Load prompts from text file (one per line)"""
+        with open(filepath, 'r', encoding='utf-8') as f:
+            prompts = [line.strip() for line in f if line.strip() and not line.startswith('#')]
+        return prompts
+    
+    @staticmethod
+    def load_prompts_from_json(filepath: str) -> List[Dict]:
+        """Load prompts and configs from JSON file"""
+        with open(filepath, 'r') as f:
+            data = json.load(f)
+        
+        if isinstance(data, list):
+            return data
+        elif isinstance(data, dict) and "prompts" in data:
+            return data["prompts"]
+        else:
+            raise ValueError("Invalid JSON format")
+    
+    @staticmethod
+    def save_batch_results(
+        results: List[Tuple[Image.Image, Dict]],
+        output_dir: str,
+        prefix: str = "batch"
+    ):
+        """Save batch generation results"""
+        output_path = Path(output_dir)
+        output_path.mkdir(parents=True, exist_ok=True)
+        
+        for i, (image, metadata) in enumerate(results):
+            # Save image
+            image_file = output_path / f"{prefix}_{i:04d}.png"
+            image.save(image_file)
+            
+            # Save metadata
+            metadata_file = output_path / f"{prefix}_{i:04d}_metadata.json"
+            with open(metadata_file, 'w') as f:
+                json.dump(metadata, f, indent=2)
+    
+    @staticmethod
+    def create_batch_report(
+        results: List[Tuple[Image.Image, Dict]],
+        output_file: str
+    ):
+        """Create a summary report of batch processing"""
+        report = {
+            "total_images": len(results),
+            "timestamp": datetime.now().isoformat(),
+            "images": []
+        }
+        
+        for i, (_, metadata) in enumerate(results):
+            report["images"].append({
+                "index": i,
+                "prompt": metadata.get("prompt", ""),
+                "generation_time": metadata.get("generation_time", 0),
+                "parameters": {
+                    "width": metadata.get("width", 0),
+                    "height": metadata.get("height", 0),
+                    "steps": metadata.get("num_inference_steps", 0),
+                    "guidance": metadata.get("guidance_scale", 0)
+                }
+            })
+        
+        # Calculate statistics
+        times = [m.get("generation_time", 0) for _, m in results]
+        if times:
+            report["statistics"] = {
+                "total_time": sum(times),
+                "average_time": np.mean(times),
+                "min_time": min(times),
+                "max_time": max(times)
+            }
+        
+        with open(output_file, 'w') as f:
+            json.dump(report, f, indent=2)
+
+
+# ============================================================================
+# FILE MANAGEMENT
+# ============================================================================
+
+class FileManager:
+    """Utilities for file management"""
+    
+    @staticmethod
+    def create_directory_structure(base_dir: str) -> Dict[str, Path]:
+        """Create organized directory structure"""
+        base = Path(base_dir)
+        
+        dirs = {
+            "outputs": base / "outputs",
+            "metadata": base / "metadata",
+            "configs": base / "configs",
+            "logs": base / "logs",
+            "temp": base / "temp"
+        }
+        
+        for dir_path in dirs.values():
+            dir_path.mkdir(parents=True, exist_ok=True)
+        
+        return dirs
+    
+    @staticmethod
+    def generate_filename(
+        prompt: str,
+        timestamp: bool = True,
+        max_length: int = 50
+    ) -> str:
+        """Generate filename from prompt"""
+        # Clean prompt
+        clean = re.sub(r'[^\w\s-]', '', prompt.lower())
+        clean = re.sub(r'[-\s]+', '_', clean)
+        clean = clean[:max_length]
+        
+        if timestamp:
+            ts = datetime.now().strftime("%Y%m%d_%H%M%S")
+            return f"{ts}_{clean}.png"
+        
+        return f"{clean}.png"
+    
+    @staticmethod
+    def get_file_hash(filepath: str) -> str:
+        """Calculate MD5 hash of file"""
+        hash_md5 = hashlib.md5()
+        with open(filepath, "rb") as f:
+            for chunk in iter(lambda: f.read(4096), b""):
+                hash_md5.update(chunk)
+        return hash_md5.hexdigest()
+    
+    @staticmethod
+    def cleanup_temp_files(temp_dir: str, older_than_hours: int = 24):
+        """Clean up temporary files older than specified hours"""
+        temp_path = Path(temp_dir)
+        if not temp_path.exists():
+            return
+        
+        cutoff_time = time.time() - (older_than_hours * 3600)
+        
+        for file in temp_path.glob("*"):
+            if file.is_file() and file.stat().st_mtime < cutoff_time:
+                file.unlink()
+                logger.info(f"Deleted old temp file: {file}")
+
+
+# ============================================================================
+# QUALITY ASSESSMENT
+# ============================================================================
+
+class QualityAssessor:
+    """Assess image quality"""
+    
+    @staticmethod
+    def calculate_sharpness(image: Image.Image) -> float:
+        """Calculate image sharpness using Laplacian variance"""
+        img_array = np.array(image.convert('L'))
+        laplacian = np.array([[0, 1, 0], [1, -4, 1], [0, 1, 0]])
+        
+        # Convolve
+        from scipy import signal
+        filtered = signal.convolve2d(img_array, laplacian, mode='valid')
+        variance = np.var(filtered)
+        
+        return float(variance)
+    
+    @staticmethod
+    def calculate_brightness(image: Image.Image) -> float:
+        """Calculate average brightness"""
+        img_array = np.array(image.convert('L'))
+        return float(np.mean(img_array))
+    
+    @staticmethod
+    def calculate_contrast(image: Image.Image) -> float:
+        """Calculate image contrast"""
+        img_array = np.array(image.convert('L'))
+        return float(np.std(img_array))
+    
+    @staticmethod
+    def assess_quality(image: Image.Image) -> Dict:
+        """Comprehensive quality assessment"""
+        return {
+            "sharpness": QualityAssessor.calculate_sharpness(image),
+            "brightness": QualityAssessor.calculate_brightness(image),
+            "contrast": QualityAssessor.calculate_contrast(image),
+            "resolution": f"{image.width}x{image.height}",
+            "aspect_ratio": image.width / image.height
+        }
+
+
+# ============================================================================
+# UTILITY FUNCTIONS
+# ============================================================================
+
+def seed_everything(seed: int):
+    """Set all random seeds for reproducibility"""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def get_optimal_resolution(
+    target_pixels: int,
+    aspect_ratio: str = "1:1"
+) -> Tuple[int, int]:
+    """
+    Calculate optimal resolution for target pixel count
+    
+    Args:
+        target_pixels: Target total pixels (e.g., 512*512 = 262144)
+        aspect_ratio: Desired aspect ratio (e.g., "16:9", "4:3", "1:1")
+    
+    Returns:
+        (width, height) tuple
+    """
+    ratios = {
+        "1:1": (1, 1),
+        "4:3": (4, 3),
+        "3:4": (3, 4),
+        "16:9": (16, 9),
+        "9:16": (9, 16),
+        "3:2": (3, 2),
+        "2:3": (2, 3)
+    }
+    
+    ratio_w, ratio_h = ratios.get(aspect_ratio, (1, 1))
+    
+    # Calculate dimensions
+    height = int(np.sqrt(target_pixels * ratio_h / ratio_w))
+    width = int(height * ratio_w / ratio_h)
+    
+    # Round to nearest multiple of 8
+    width = (width // 8) * 8
+    height = (height // 8) * 8
+    
+    return width, height
+
+
+def estimate_generation_time(
+    width: int,
+    height: int,
+    steps: int,
+    device: str = "cuda",
+    gpu_model: str = "RTX 3080"
+) -> float:
+    """
+    Estimate generation time based on parameters
+    
+    Returns:
+        Estimated time in seconds
+    """
+    # Base time per step (seconds) for different GPUs at 512x512
+    base_times = {
+        "RTX 4090": 0.04,
+        "RTX 3090": 0.07,
+        "RTX 3080": 0.10,
+        "RTX 2080": 0.15,
+        "M1 Max": 0.25
+    }
+    
+    base_time = base_times.get(gpu_model, 0.10)
+    
+    # Scale by resolution
+    pixel_factor = (width * height) / (512 * 512)
+    
+    # Estimate
+    estimated = base_time * steps * pixel_factor
+    
+    return estimated
+
+
+# Export main classes and functions
+__all__ = [
+    'GenerationConfig',
+    'GenerationMetadata',
+    'PromptEnhancer',
+    'ImageProcessor',
+    'MetadataManager',
+    'PerformanceMonitor',
+    'ConfigManager',
+    'BatchProcessor',
+    'FileManager',
+    'QualityAssessor',
+    'seed_everything',
+    'get_optimal_resolution',
+    'estimate_generation_time'
+]