Memory Management

class VRAMState(Enum):
    DISABLED = 0    # No VRAM present: no need to move models to VRAM
    NO_VRAM = 1     # Very low VRAM: enable all options to save VRAM
    LOW_VRAM = 2    # Limited VRAM: selective model loading
    NORMAL_VRAM = 3 # Standard operation
    HIGH_VRAM = 4   # Keep models in VRAM
    SHARED = 5      # Shared CPU/GPU memory (e.g., integrated graphics)

from comfy import model_management

# Get total memory
total_vram = model_management.get_total_memory()
print(f"Total VRAM: {total_vram / (1024**3):.2f} GB")

# Get free memory
free_vram = model_management.get_free_memory()
print(f"Free VRAM: {free_vram / (1024**3):.2f} GB")

# Get both hardware and PyTorch memory
mem_total, mem_torch = model_management.get_free_memory(
    device, 
    torch_free_too=True
)

class LoadedModel:
    def __init__(self, model):
        self.device = model.load_device
        self.real_model = None
        self.currently_used = True
        
    def model_memory(self):
        """Total model size"""
        return self.model.model_size()
    
    def model_loaded_memory(self):
        """Currently loaded portion"""
        return self.model.loaded_size()
    
    def model_offloaded_memory(self):
        """Portion offloaded to CPU/disk"""
        return self.model.model_size() - self.model.loaded_size()

def load_models_gpu(
    models, 
    memory_required=0, 
    force_patch_weights=False,
    minimum_memory_required=None, 
    force_full_load=False
):
    # 1. Calculate total memory required
    # 2. Free memory if needed
    # 3. Load models with partial loading if necessary
    # 4. Track loaded models
    pass

if vram_set_state == VRAMState.LOW_VRAM:
    loaded_memory = loaded_model.model_loaded_memory()
    current_free_mem = get_free_memory(torch_dev) + loaded_memory
    
    # Calculate how much of the model to load
    lowvram_model_memory = max(
        0,
        (current_free_mem - minimum_memory_required),
        min(
            current_free_mem * MIN_WEIGHT_MEMORY_RATIO,
            current_free_mem - minimum_inference_memory()
        )
    )
    
    loaded_model.model_load(lowvram_model_memory)

def free_memory(
    memory_required, 
    device, 
    keep_loaded=[], 
    for_dynamic=False, 
    ram_required=0
):
    # 1. Garbage collect
    cleanup_models_gc()
    
    # 2. Find candidate models to unload
    can_unload = []
    for shift_model in current_loaded_models:
        if shift_model not in keep_loaded:
            can_unload.append(shift_model)
    
    # 3. Sort by offloaded memory (unload least active first)
    for model in sorted(can_unload):
        if get_free_memory(device) < memory_required:
            model.model_unload()
    
    # 4. Clear cache
    soft_empty_cache()

EXTRA_RESERVED_VRAM = 400 * 1024 * 1024  # 400 MB base

if WINDOWS:
    EXTRA_RESERVED_VRAM = 600 * 1024 * 1024  # 600 MB on Windows
    if total_vram > (15 * 1024):  # 16GB+ cards
        EXTRA_RESERVED_VRAM += 100 * 1024 * 1024

def minimum_inference_memory():
    return (1024 * 1024 * 1024) * 0.8 + extra_reserved_memory()

NUM_STREAMS = 2  # Default for NVIDIA and AMD

if args.async_offload is not None:
    NUM_STREAMS = args.async_offload

def get_offload_stream(device):
    if NUM_STREAMS == 0:
        return None
    
    if device in STREAMS:
        ss = STREAMS[device]
        ss[stream_counter].wait_stream(current_stream(device))
        stream_counter = (stream_counter + 1) % len(ss)
        return ss[stream_counter]

STREAM_CAST_BUFFERS = {}

def get_cast_buffer(offload_stream, device, size, ref):
    cast_buffer = STREAM_CAST_BUFFERS.get(offload_stream, None)
    if cast_buffer is None or cast_buffer.numel() < size:
        with offload_stream.as_context(offload_stream):
            cast_buffer = torch.empty((size), dtype=torch.int8, device=device)
            STREAM_CAST_BUFFERS[offload_stream] = cast_buffer
    return cast_buffer

MAX_PINNED_MEMORY = get_total_memory(torch.device("cpu")) * 0.95

def pin_memory(tensor):
    if not is_device_cpu(tensor.device):
        return False
    
    size = tensor.nbytes
    if (TOTAL_PINNED_MEMORY + size) > MAX_PINNED_MEMORY:
        return False
    
    ptr = tensor.data_ptr()
    if torch.cuda.cudart().cudaHostRegister(ptr, size, 1) == 0:
        PINNED_MEMORY[ptr] = size
        TOTAL_PINNED_MEMORY += size
        return True

def unet_offload_device():
    if vram_state == VRAMState.HIGH_VRAM:
        return get_torch_device()
    else:
        return torch.device("cpu")

def unet_inital_load_device(parameters, dtype):
    torch_dev = get_torch_device()
    if vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.SHARED:
        return torch_dev
    
    model_size = dtype_size(dtype) * parameters
    if model_size < get_free_memory(torch_dev):
        return torch_dev
    return torch.device("cpu")

def text_encoder_device():
    if args.gpu_only:
        return get_torch_device()
    elif vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.NORMAL_VRAM:
        if should_use_fp16(prioritize_performance=False):
            return get_torch_device()
    return torch.device("cpu")

def vae_device():
    if args.cpu_vae:
        return torch.device("cpu")
    return get_torch_device()

def vae_offload_device():
    if args.gpu_only:
        return get_torch_device()
    return torch.device("cpu")

def cleanup_models_gc():
    do_gc = False
    reset_cast_buffers()
    
    for cur in current_loaded_models:
        if cur.is_dead():
            logging.info(
                f"Potential memory leak detected with model {cur.real_model().__class__.__name__}"
            )
            do_gc = True
            break
    
    if do_gc:
        gc.collect()
        soft_empty_cache()