memory_optimization/__init__.py

"""
Memory Optimization Module
Unified memory management system with shared Qwen model integration.
"""
from .config import MemoryOptimizationConfig
from .manager import UnifiedMemoryManager
from .tensor_pool import TensorPool
from .model_cache import ModelCache
from .cleanup import MemoryCleanup

# Global unified memory manager instance
_unified_memory_manager = None

def get_unified_memory_manager() -> UnifiedMemoryManager:
    """Get the global unified memory manager instance"""
    global _unified_memory_manager
    if _unified_memory_manager is None:
        _unified_memory_manager = UnifiedMemoryManager()
    return _unified_memory_manager

# Convenience functions for backward compatibility
def get_shared_model(model_name: str, model_type: str = "transformer", **kwargs):
    """Get shared model instance"""
    return get_unified_memory_manager().get_shared_model(model_name, model_type, **kwargs)

def get_tensor(shape, dtype=None, requires_grad: bool = False, module_name: str = "default"):
    """Get optimized tensor from unified pool"""
    return get_unified_memory_manager().get_tensor(shape, dtype, requires_grad, module_name)

def return_tensor(tensor, module_name: str = "default") -> None:
    """Return tensor to unified pool"""
    get_unified_memory_manager().return_tensor(tensor, module_name)

def clear_memory() -> None:
    """Clear all memory"""
    get_unified_memory_manager().clear_all_memory()

def get_memory_stats():
    """Get memory statistics"""
    return get_unified_memory_manager().get_memory_stats()

__all__ = [
    'MemoryOptimizationConfig',
    'UnifiedMemoryManager',
    'TensorPool',
    'ModelCache',
    'MemoryCleanup',
    'get_unified_memory_manager',
    'get_shared_model',
    'get_tensor',
    'return_tensor',
    'clear_memory',
    'get_memory_stats',
]