EricRollei-Eric_Image_Processing_Nodes/scripts/memory_utils.py

"""
Memory management utilities for advanced image processing
Helps prevent memory issues when working with large models and images
"""

import torch
import gc
import psutil
import numpy as np
from typing import Optional, Any, Dict
import logging

logger = logging.getLogger(__name__)

class MemoryManager:
    """Utility class for managing memory during image processing"""
    
    @staticmethod
    def get_memory_info() -> Dict[str, Any]:
        """Get current memory usage information"""
        info = {
            'system_ram_total': psutil.virtual_memory().total / (1024**3),  # GB
            'system_ram_available': psutil.virtual_memory().available / (1024**3),  # GB
            'system_ram_percent': psutil.virtual_memory().percent,
        }
        
        if torch.cuda.is_available():
            try:
                gpu_memory = torch.cuda.get_device_properties(0).total_memory / (1024**3)  # GB
                gpu_allocated = torch.cuda.memory_allocated(0) / (1024**3)  # GB
                gpu_cached = torch.cuda.memory_reserved(0) / (1024**3)  # GB
                
                info.update({
                    'gpu_memory_total': gpu_memory,
                    'gpu_memory_allocated': gpu_allocated,
                    'gpu_memory_cached': gpu_cached,
                    'gpu_memory_free': gpu_memory - gpu_cached,
                })
            except Exception as e:
                logger.warning(f"Could not get GPU memory info: {e}")
                
        return info
    
    @staticmethod
    def cleanup_memory():
        """Clean up memory (both CPU and GPU)"""
        # Python garbage collection
        gc.collect()
        
        # PyTorch CUDA cache cleanup
        if torch.cuda.is_available():
            try:
                torch.cuda.empty_cache()
                torch.cuda.ipc_collect()
            except Exception as e:
                logger.warning(f"GPU cleanup failed: {e}")
    
    @staticmethod
    def safe_device_selection(preferred_device: str = 'auto', 
                            min_gpu_memory_gb: float = 4.0) -> str:
        """
        Safely select device based on available memory
        
        Args:
            preferred_device: 'auto', 'cuda', or 'cpu'
            min_gpu_memory_gb: Minimum GPU memory required for CUDA
        """
        if preferred_device == 'cpu':
            return 'cpu'
            
        if not torch.cuda.is_available():
            return 'cpu'
            
        try:
            gpu_props = torch.cuda.get_device_properties(0)
            total_memory_gb = gpu_props.total_memory / (1024**3)
            
            # Check if we have enough free memory
            torch.cuda.empty_cache()
            allocated_gb = torch.cuda.memory_allocated(0) / (1024**3)
            free_memory_gb = total_memory_gb - allocated_gb
            
            if free_memory_gb < min_gpu_memory_gb:
                logger.info(f"GPU has {free_memory_gb:.1f}GB free, need {min_gpu_memory_gb}GB. Using CPU.")
                return 'cpu'
            else:
                logger.info(f"GPU has {free_memory_gb:.1f}GB free. Using CUDA.")
                return 'cuda'
                
        except Exception as e:
            logger.warning(f"Error checking GPU memory: {e}. Using CPU.")
            return 'cpu'
    
    @staticmethod
    def estimate_memory_requirements(image_shape: tuple, 
                                   model_params: int,
                                   processing_factor: float = 4.0) -> Dict[str, float]:
        """
        Estimate memory requirements for processing
        
        Args:
            image_shape: (H, W, C) shape of input image
            model_params: Number of model parameters
            processing_factor: Memory multiplier for intermediate tensors
        """
        h, w, c = image_shape
        
        # Image memory (input + output + intermediate)
        image_memory_mb = (h * w * c * 4 * processing_factor) / (1024**2)  # float32
        
        # Model memory
        model_memory_mb = (model_params * 4) / (1024**2)  # float32 parameters
        
        # Total estimate
        total_mb = image_memory_mb + model_memory_mb
        
        return {
            'image_memory_mb': image_memory_mb,
            'model_memory_mb': model_memory_mb,
            'total_memory_mb': total_mb,
            'total_memory_gb': total_mb / 1024
        }
    
    @staticmethod
    def optimize_tile_size(image_shape: tuple,
                          max_memory_gb: float = 2.0,
                          min_tile_size: int = 64,
                          max_tile_size: int = 1024) -> int:
        """
        Calculate optimal tile size based on available memory
        
        Args:
            image_shape: (H, W, C) shape of input image
            max_memory_gb: Maximum memory to use for processing
            min_tile_size: Minimum allowed tile size
            max_tile_size: Maximum allowed tile size
        """
        h, w, c = image_shape
        max_memory_mb = max_memory_gb * 1024
        
        # Start with max tile size and reduce if needed
        for tile_size in range(max_tile_size, min_tile_size - 1, -64):
            # Estimate memory for this tile size
            tile_memory_mb = (tile_size * tile_size * c * 4 * 8) / (1024**2)  # Processing factor
            
            if tile_memory_mb <= max_memory_mb:
                return tile_size
                
        return min_tile_size

class SafeModelLoader:
    """Context manager for safe model loading and cleanup"""
    
    def __init__(self, device: str = 'auto', cleanup_on_exit: bool = True):
        self.device = MemoryManager.safe_device_selection(device)
        self.cleanup_on_exit = cleanup_on_exit
        self.initial_memory = None
        
    def __enter__(self):
        self.initial_memory = MemoryManager.get_memory_info()
        logger.info(f"Starting model loading on {self.device}")
        logger.info(f"Initial memory: {self.initial_memory.get('gpu_memory_free', 'N/A')} GB GPU free")
        return self.device
        
    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.cleanup_on_exit:
            MemoryManager.cleanup_memory()
            
        final_memory = MemoryManager.get_memory_info()
        logger.info(f"Final memory: {final_memory.get('gpu_memory_free', 'N/A')} GB GPU free")
        
        if exc_type is not None:
            logger.error(f"Model loading failed: {exc_val}")

def safe_torch_operation(func):
    """Decorator for safe PyTorch operations with memory cleanup"""
    def wrapper(*args, **kwargs):
        try:
            # Clean up before operation
            MemoryManager.cleanup_memory()
            
            # Run operation
            result = func(*args, **kwargs)
            
            # Clean up after operation
            MemoryManager.cleanup_memory()
            
            return result
            
        except Exception as e:
            # Clean up on error
            MemoryManager.cleanup_memory()
            logger.error(f"Operation {func.__name__} failed: {e}")
            raise
            
    return wrapper

# Convenience functions
def get_safe_batch_size(image_size: int, base_batch_size: int = 4) -> int:
    """Get safe batch size based on image dimensions"""
    memory_info = MemoryManager.get_memory_info()
    
    if 'gpu_memory_free' in memory_info:
        free_gb = memory_info['gpu_memory_free']
        if free_gb < 2:
            return 1
        elif free_gb < 4:
            return max(1, base_batch_size // 2)
    
    # Adjust based on image size
    if image_size > 1024:
        return 1
    elif image_size > 512:
        return max(1, base_batch_size // 2)
    else:
        return base_batch_size

def check_memory_requirements(required_gb: float) -> bool:
    """Check if we have enough memory for an operation"""
    memory_info = MemoryManager.get_memory_info()
    
    # Check system RAM
    available_ram = memory_info.get('system_ram_available', 0)
    if available_ram < required_gb:
        return False
    
    # Check GPU memory if using CUDA
    if torch.cuda.is_available():
        gpu_free = memory_info.get('gpu_memory_free', 0)
        if gpu_free < required_gb:
            return False
            
    return True

# Example usage
if __name__ == "__main__":
    # Test memory utilities
    print("Memory Info:")
    info = MemoryManager.get_memory_info()
    for key, value in info.items():
        if 'memory' in key and isinstance(value, float):
            print(f"  {key}: {value:.2f} GB")
        else:
            print(f"  {key}: {value}")
    
    print(f"\nSafe device: {MemoryManager.safe_device_selection()}")
    
    # Test tile size optimization
    tile_size = MemoryManager.optimize_tile_size((2048, 2048, 3))
    print(f"Optimal tile size for 2048x2048 image: {tile_size}")
    
    # Test memory requirements estimation
    reqs = MemoryManager.estimate_memory_requirements((1024, 1024, 3), 1000000)
    print(f"Memory requirements: {reqs['total_memory_gb']:.2f} GB")