engine/optimizations/memory_manager.py

"""
Memory Manager pour DARKMEDIA-X
Gère l'allocation et libération de mémoire GPU/CPU
"""

import gc
import os
from typing import Optional

import psutil
import torch


class MemoryManager:
    """Gère la mémoire GPU et CPU pour éviter les surcharges."""

    MIN_VRAM_MB = 1024  # Minimum 1GB de VRAM libre

    @staticmethod
    def init_cuda():
        """Initialise CUDA avec les paramètres d'optimisation."""
        try:
            if torch.cuda.is_available():
                # Configuration pour éviter les erreurs de pagination
                os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb=512"
                torch.cuda.empty_cache()
                torch.cuda.synchronize()
                print(f"✅ CUDA initialisé: {torch.cuda.get_device_name(0)}")
                return True
        except Exception as e:
            print(f"⚠️ CUDA init échoué: {e}")
        return False

    @staticmethod
    def get_available_vram() -> int:
        """Retourne la VRAM disponible en MB."""
        try:
            if torch.cuda.is_available():
                return torch.cuda.mem_get_info()[0] // (1024 * 1024)
        except:
            pass
        return 0

    @staticmethod
    def get_available_ram() -> int:
        """Retourne la RAM disponible en MB."""
        return int(psutil.virtual_memory().available / (1024 * 1024))

    @staticmethod
    def check_memory_available(required_mb: int = 2048) -> bool:
        """Vérifie si assez de mémoire est disponible."""
        vram = MemoryManager.get_available_vram()
        ram = MemoryManager.get_available_ram()

        # Sur GPU: vérifier VRAM
        if vram > 0 and vram < required_mb:
            print(f"⚠️ VRAM insuffisante: {vram}MB < {required_mb}MB requis")
            return False

        # Sur CPU: vérifier RAM
        if vram == 0 and ram < required_mb:
            print(f"⚠️ RAM insuffisante: {ram}MB < {required_mb}MB requis")
            return False

        return True

    @staticmethod
    def clear_memory(force: bool = False):
        """Nettoie la mémoire GPU et CPU."""
        gc.collect()

        try:
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                if force:
                    torch.cuda.synchronize()
        except Exception as e:
            print(f"⚠️ Erreur nettoyage CUDA: {e}")

    @staticmethod
    def get_device() -> str:
        """Retourne le device optimal (cuda ou cpu)."""
        try:
            if torch.cuda.is_available():
                vram = MemoryManager.get_available_vram()
                if vram > MemoryManager.MIN_VRAM_MB:
                    return "cuda"
        except:
            pass
        return "cpu"

    @staticmethod
    def print_memory_status():
        """Affiche l'état de la mémoire."""
        vram = MemoryManager.get_available_vram()
        ram = MemoryManager.get_available_ram()

        status = "💾 Mémoire: "
        if vram > 0:
            status += f"VRAM={vram}MB "
        status += f"RAM={ram}MB"

        print(status)


def init_memory():
    """Initialise la gestion de mémoire au démarrage."""
    MemoryManager.init_cuda()
    MemoryManager.clear_memory()
    MemoryManager.print_memory_status()