comfy-aimdo

What it is and what it does

comfy-aimdo is a custom PyTorch VRAM allocator designed to handle GPU memory pressure by dynamically offloading model weights to system memory. Instead of the standard PyTorch allocator, it uses CUDA's virtual address reservation APIs to create Virtual Base Address Registers (VBARs) for models—reserving address space without consuming VRAM upfront. Tensors are allocated within these VBARs and only faulted into actual GPU memory when needed via an explicit `fault()` call. If VRAM is insufficient, the allocator falls back to temporary GPU tensors that are garbage-collected after use, effectively spilling to system memory.

The allocator implements a priority system where more recently created VBARs take precedence, and within a VBAR, lower addresses have higher priority. When a weight is evicted due to memory pressure, a watermark is set to prevent repeatedly faulting in already-offloaded weights. Applications can also call `prioritize()` to promote an existing model to top priority. The design assumes regular weight access patterns and recommends flushing the PyTorch caching allocator between model runs to avoid fragmentation. This is specialized infrastructure for scenarios like ComfyUI where multiple large models must coexist with limited VRAM.

Use it for:

• Load and run multiple large language or diffusion models sequentially on a single GPU without OOM errors
• Implement dynamic model swapping in inference pipelines where model priority changes based on workflow order
• Optimize VRAM utilization in multi-model workflows by offloading lower-priority weights to system memory
• Reduce memory fragmentation in long-running applications that repeatedly load and unload different models

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

A PyTorch VRAM allocator that dynamically offloads model weights to system memory when GPU memory is under pressure, using virtual address reservation to minimize fragmentation.

Yes, if you are building or using ComfyUI-like workflows with multiple large models on limited VRAM and are willing to accept the constraints: Nvidia GPU only, specific PyTorch/CUDA versions, and active management of model priority and allocator flushing. The license status is unclear, so verify terms first. Not suitable for general PyTorch projects or non-Nvidia hardware.

Install

comfy-aimdo on PyPI

pip install comfy-aimdo

uv add comfy-aimdo

poetry add comfy-aimdo

Installing comfy-aimdo

Before you install

Low friction installation with no runtime dependencies. Active maintenance with recent releases; however, the package is young (first release January 2026) and relatively niche, so production stability remains unproven.

License in practice

License status is unclear—no SPDX identifier or raw license text is available. Verify the actual license terms in the repository before using in commercial or proprietary projects.

Quickstart

pip install comfy-aimdo

import comfy_aimdo
# Create a VBAR, allocate tensors, call fault() to load weights on demand
# See examples/example.py in the repository for full usage pattern

Requires Nvidia GPU, PyTorch 2.8+, CUDA 12.8+, and Python 3.9+. Windows 11+ or Linux only.

Verify before relying

• Performance overhead of the offloading mechanism compared to standard PyTorch allocation
• Compatibility with specific PyTorch versions beyond the stated 2.8+ minimum
• Real-world stability and fragmentation behavior under sustained production workloads

Package facts

Evidence: comfy_aimdo-0.4.13-py3-none-any.whl

Tags

More Artificial Intelligence packages

Further reading