Model management

How ComfyUI decides where models live in memory, when to move them between CPU and GPU, and which dtypes/attention backends to use.

Purpose

Run any supported model on any supported hardware, from 1 GB GPUs to multi-GPU rigs. Most users never set a memory flag and ComfyUI gets it right; the same scheduler also drives async offload, fp8/bf16/fp16 selection, and per-model split loading.

Layout

comfy/
├── model_management.py       # 1,829 lines — VRAMState, dtype/attention probes, scheduling
├── memory_management.py      # Lower-level allocator helpers
├── pinned_memory.py          # Pinned host memory pool for fast H2D copies
├── model_patcher.py          # 1,744 lines — ModelPatcher, the per-model handle
├── ops.py                    # disable_weight_init / manual_cast / MixedPrecisionOps
└── float.py                  # FP8/BF16 cast helpers

The DynamicVRAM scheduler (newer) lives in the external comfy_aimdo package; ComfyUI activates it when enables_dynamic_vram() returns true (see comfy/cli_args.py).

Key abstractions

Type / function	File	What it is
VRAMState	comfy/model_management.py	Enum: DISABLED / NO_VRAM / LOW_VRAM / NORMAL_VRAM / HIGH_VRAM / SHARED
vram_state, set_vram_to, cpu_state	comfy/model_management.py	Globals set at startup based on flags + detected hardware
get_torch_device	comfy/model_management.py	The active inference device
intermediate_device, intermediate_dtype	comfy/model_management.py	Where node outputs should live so the next node can use them cheaply
get_free_memory, total_vram	comfy/model_management.py	Memory probes
load_models_gpu, unload_all_models	comfy/model_management.py	Move a list of ModelPatchers to GPU (and evict others if needed)
ModelPatcher	comfy/model_patcher.py	Per-model handle. Owns weights + patches. The single most central type in comfy/
ModelPatcherDynamic	comfy/model_patcher.py	DynamicVRAM-enabled subclass; replaces ModelPatcher when conditions allow
disable_weight_init	comfy/ops.py	The default ops mode: nn.Linear/nn.Conv2d that don't initialize weights at build
manual_cast_class	comfy/ops.py	An ops mode that casts weights on the fly to a target dtype
MixedPrecisionOps	comfy/ops.py	Per-layer quantization (see Quantization)
pick_operations	comfy/ops.py	The function that decides which ops mode to use given a model config

VRAM scheduling

stateDiagram-v2
    [*] --> Probe
    Probe --> NORMAL_VRAM: enough VRAM
    Probe --> HIGH_VRAM: --highvram or --gpu-only
    Probe --> LOW_VRAM: --lowvram or auto-detected low memory
    Probe --> NO_VRAM: --novram
    Probe --> SHARED: integrated GPU (Apple, MPS, etc.)
    Probe --> DISABLED: --cpu

    NORMAL_VRAM --> Run
    HIGH_VRAM --> Run
    LOW_VRAM --> Split: weights split-loaded
    NO_VRAM --> Split
    SHARED --> Run
    DISABLED --> CPU_only

    Run --> Run: load_models_gpu(...)
    Split --> Split: per-block on/off device
    CPU_only --> CPU_only

In NORMAL_VRAM and above, models sit in GPU memory; load_models_gpu evicts other models if there isn't room. In LOW_VRAM and below, models are split-loaded per block — the diffusion network is moved one chunk at a time.

The actual heuristic for picking a state is in model_management.py's vram_state_management. It's responsive: when --reserve-vram carves off OS memory, when an estimate from comfy_aimdo overrides it, etc.

ModelPatcher

ModelPatcher (in comfy/model_patcher.py) wraps each loaded model. It owns:

• The underlying nn.Module (self.model).
• A list of "patches" (LoRA, hooks) keyed by parameter name.
• Move-to-device logic that knows how to apply patches lazily: weights are kept on CPU "clean" and only the active patched values get moved to GPU.
• Clone semantics: .clone() produces a new patcher that shares the underlying weights but has its own patch list. This is how nodes can apply a LoRA without modifying the loaded model.

ModelPatcherDynamic is the DynamicVRAM variant that integrates with comfy_aimdo for finer-grained per-block streaming. It replaces the legacy patcher when:

• Pytorch ≥ 2.8
• An NVIDIA GPU is present (and not WSL)
• comfy_aimdo.control.init_device(...) succeeds

This is set up at import time in main.py.

Async offload

--async-offload [N] enables overlapping weight uploads with compute via N CUDA streams (default 2). On NVIDIA it's on by default; --disable-async-offload opts out. The mechanism uses pinned host memory (comfy/pinned_memory.py) and per-stream prefetch in ModelPatcher.

Dtype selection

ComfyUI picks dtypes based on:

1. Explicit overrides: --fp32-unet, --fp16-unet, --bf16-unet, --fp8_e4m3fn-unet, --fp8_e5m2-unet, --fp8_e8m0fnu-unet, plus VAE/text-enc variants.
2. Hardware probes in model_management.py: should_use_fp16, should_use_bf16, supports_fp8_compute, get_supported_float8_types.
3. The model's own preferences (BaseModel.manual_cast_dtype).

The decision flows into pick_operations in comfy/ops.py, which returns the nn.Module mode used by every layer. Three modes:

• disable_weight_init — plain modules without parameter init; loaded later from state dict.
• manual_cast_class — cast weights on the fly when called; used when stored dtype ≠ compute dtype.
• MixedPrecisionOps — per-layer quantization. See Quantization.

Attention backends

The --use-*-attention flags (and the xformers autodetect) pick an attention implementation:

Flag	Backend
--use-pytorch-cross-attention	torch.nn.functional.scaled_dot_product_attention
--use-flash-attention	FlashAttention
--use-sage-attention	Sage attention
--use-split-cross-attention	Split (chunked) attention
--use-quad-cross-attention	Sub-quadratic attention
(default)	xformers if available, else PyTorch SDPA

Wired up in comfy/ldm/modules/attention.py (the optimized_attention dispatcher) by reading comfy.cli_args.args.

CPU / fallback paths

• --cpu runs everything on CPU — slow but works without a GPU.
• --cpu-vae keeps the VAE on CPU even when the rest is on GPU; useful for memory-tight high-resolution decodes.
• --directml uses torch-directml on Windows for AMD/Intel.
• Apple Silicon uses MPS automatically (it shows up as cpu_state == CPUState.MPS).
• Ascend NPU and Cambricon MLU paths exist via env vars set in main.py (ASCEND_RT_VISIBLE_DEVICES, etc.).

Smart memory and unloading

prompt_worker in main.py calls comfy.model_management.unload_all_models() and soft_empty_cache() between prompts based on flags from the queue (free_memory, unload_models). The garbage-collection cadence is throttled to once every 10 seconds.

--disable-smart-memory forces aggressive CPU offload between models — useful when you have many models and limited GPU.

Integration points

• Imported by every loader in comfy/sd.py.
• Driven by flags from comfy/cli_args.py.
• Underlies the Sampling pipeline: load_models_gpu is called before each sample.
• DynamicVRAM hooks are configured in main.py.
• Test surface in tests-unit/comfy_test/.

Where to start a change

• Adding hardware support: probes go in comfy/model_management.py (e.g., is_intel_xpu, is_ascend_npu). New env vars belong in main.py.
• New dtype path: add to pick_operations in comfy/ops.py; register the flag in comfy/cli_args.py.
• New attention backend: extend comfy/ldm/modules/attention.py's dispatcher and add a flag.