Spaces:
Sleeping
Sleeping
File size: 3,279 Bytes
343eed9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | """
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()
|