ComfyUI-master Update: Errors and Fixes

ComfyUI-master Update: Errors and Fixes (Full Guide)

This document describes the errors that occurred after updating the ComfyUI-master core to the latest version and the fixes applied, with detailed code explanations.

---

Table of contents

1. Overview
2. Summary of changes (created/modified files)
3. Error 1: comfy.checkpoint_pickle removal
4. Error 2: ModuleNotFoundError: comfy_aimdo
5. Summary

---

1. Overview

Item	Description
Scope	Compatibility after updating the ComfyUI-master tree that Forge references (COMFYUI_FOLDER_NAME = "ComfyUI-master")
Issues	2 (removal of checkpoint_pickle; ComfyUI’s unconditional use of comfy_aimdo)
Approach	Align Forge with ComfyUI changes; add stub package where a module is missing

---

2. Summary of changes (created/modified files)

Type	Path	Description
Modified	backend/nn/comfy_utils.py	torch.load changes following removal of comfy.checkpoint_pickle
New	comfy_aimdo/__init__.py	Root of the comfy_aimdo stub package
New	comfy_aimdo/torch.py	Stub for aimdo_to_tensor
New	comfy_aimdo/model_vbar.py	Stubs for VBAR-related APIs and ModelVBAR
New	comfy_aimdo/control.py	Stubs for device/VRAM control APIs

---

3. Error 1: comfy.checkpoint_pickle removal

3.1 Cause

ComfyUI removed comfy/checkpoint_pickle.py and switched to safe loading only: PyTorch 2.4+ torch.load(..., weights_only=True) and torch.serialization.add_safe_globals.

• ComfyUI change: The previous “restricted pickle module” was dropped; loading now always uses weights_only=True.
• Forge: backend/nn/comfy_utils.py still had import comfy.checkpoint_pickle and used pickle_module=comfy.checkpoint_pickle, so it failed when that module no longer existed.

3.2 Modified file: backend/nn/comfy_utils.py

Change A: Remove the import

Before (relevant line only)

import comfy.checkpoint_pickle

After

# (line removed)

Explanation
checkpoint_pickle was removed from the ComfyUI repo, so Forge no longer imports it. Safe checkpoint loading is handled by the torch.load change below.

---

Change B: torch.load usage in load_torch_file

Before

        if safe_load or ALWAYS_SAFE_LOAD:
            pl_sd = torch.load(ckpt, map_location=device, weights_only=True, **torch_args)
        else:
            logging.warning("WARNING: loading {} unsafely, upgrade your pytorch to 2.4 or newer to load this file safely.".format(ckpt))
            pl_sd = torch.load(ckpt, map_location=device, pickle_module=comfy.checkpoint_pickle)

After

        # ComfyUI removed checkpoint_pickle; always use weights_only=True for safety
        pl_sd = torch.load(ckpt, map_location=device, weights_only=True, **torch_args)

Explanation

• Old behavior
  • When safe_load or ALWAYS_SAFE_LOAD was true, loading used weights_only=True (safe).
  • Otherwise it tried to use pickle_module=comfy.checkpoint_pickle for “restricted” loading, which failed at runtime because that module no longer exists.
• New behavior
  • torch.load is always called with weights_only=True only.
  • This matches ComfyUI’s current comfy/utils.py load_torch_file (always safe loading).
• Note
  With weights_only=True, some old .ckpt/.pt files that contain non-tensor Python objects may not load. Converting such files to safetensors (or similar) under PyTorch 2.4+ is the intended approach.

---

3.3 ComfyUI side (reference)

• In ComfyUI-master/comfy/utils.py, torch.load for ckpt/pt is always used with weights_only=True:

  pl_sd = torch.load(ckpt, map_location=device, weights_only=True, **torch_args)

• The same file still has a block that registers legacy pickle-related classes (e.g. ModelCheckpoint) with torch.serialization.add_safe_globals for compatibility.

Forge’s comfy_utils.py is aligned with this “always load with weights_only” policy.

---

4. Error 2: ModuleNotFoundError: No module named 'comfy_aimdo'

4.1 Cause

ComfyUI started unconditionally importing the optional PyPI package comfy_aimdo (AIMDO: AI Model Dynamic Offloader) at top level in:

• comfy/model_management.py → import comfy_aimdo.torch, import comfy_aimdo.model_vbar
• comfy/ops.py → import comfy_aimdo.model_vbar, import comfy_aimdo.torch
• comfy/model_patcher.py → import comfy_aimdo.model_vbar
• comfy/windows.py → import comfy_aimdo.control (used inside get_free_ram())

Forge imports ComfyUI’s comfy package (e.g. model_management, ops), but does not install comfy_aimdo, so startup failed with ModuleNotFoundError when import comfy_aimdo.torch ran.

What comfy_aimdo does (upstream)

• On-demand offloading of weights via “virtual base address registers” (VBAR) for Nvidia GPUs.
• It has its own requirements (e.g. PyTorch 2.6+, CUDA 12.8+) and is separate from Forge’s usual memory/loader.

Forge does not use AIMDO. We added a stub package at the project root that satisfies the imports and API surface so existing memory paths stay unchanged and the import error is resolved.

4.2 Why a stub is enough

• Code paths that use comfy_aimdo are only taken when AIMDO is enabled and weights live in VBAR.
• In our stub, vbar_fault always returns None, so ComfyUI’s if signature is not None: branch (AIMDO path) is never taken and only the normal tensor path is used.
• So a minimal stub that is importable and does not change existing branches is sufficient.

---

4.3 Created files in detail

4.3.1 comfy_aimdo/__init__.py

# Stub package for ComfyUI's optional comfy_aimdo dependency.
# Forge uses its own memory/loader; this allows ComfyUI code to import without the real package.

Explanation
Minimal __init__.py so Python treats comfy_aimdo as a package. Forge uses its own memory/loader, so this file only documents that and does not export anything.

---

4.3.2 comfy_aimdo/torch.py

# Stub: real comfy_aimdo provides AIMDO tensor conversion. Forge does not use AIMDO.
import torch


def aimdo_to_tensor(v_handle, device):
    """Stub: return a tensor on device. Not used when vbar_fault returns None."""
    if isinstance(v_handle, torch.Tensor):
        return v_handle.to(device=device)
    raise RuntimeError("comfy_aimdo stub: aimdo_to_tensor called without real comfy_aimdo")

Explanation

• Upstream: Converts an AIMDO virtual buffer handle v_handle to a tensor on the given device.
• Stub:
  • With our stub, vbar_fault always returns None, so this function is normally never called.
  • If it is called, we treat v_handle as a tensor when it is already a torch.Tensor (transfer to device) and otherwise raise a RuntimeError so misuse without the real comfy_aimdo is visible.

---

4.3.3 comfy_aimdo/model_vbar.py

# Stub: real comfy_aimdo provides VBAR for dynamic weight offloading. Forge does not use it.


def vbar_fault(_v):
    """Return None so ComfyUI uses the normal (non-AIMDO) code path."""
    return None


def vbar_signature_compare(signature, stored):
    """Stub: always False when signature is not None."""
    return False


def vbar_unpin(_v):
    """No-op when AIMDO is not in use."""
    pass


def vbars_reset_watermark_limits():
    """No-op stub for execution.py."""
    pass


class ModelVBAR:
    """Stub: real ModelVBAR manages virtual weight buffers. Forge uses its own loader."""

    def __init__(self, size, device_index):
        self._size = size
        self._device_index = device_index

    def loaded_size(self):
        return 0

Explanation

• vbar_fault(_v)

  • Upstream: Returns the current “signature” of virtual buffer _v for AIMDO transfer decisions.
  • Stub: Always returns None. So ComfyUI’s if signature is not None: is false, the AIMDO path (e.g. aimdo_to_tensor) is never used, and only the normal torch.empty_like / copy_ path runs.

• vbar_signature_compare(signature, stored)

  • Upstream: Compares signatures to see if data is already transferred.
  • Stub: Always False; used in the “resident” branch when signature is not None, which we never take.

• vbar_unpin(_v)

  • Upstream: Unpins a pinned buffer.
  • Stub: No-op (pass); fine because we do not use AIMDO.

• vbars_reset_watermark_limits()

  • Stub no-op for possible calls from ComfyUI’s execution.py.

• ModelVBAR class

  • Upstream: Manages virtual weight buffers by size and device index.
  • Stub: Can be constructed as ModelVBAR(size, device_index) and loaded_size() always returns 0, so model_patcher.py’s _vbar_get / loaded_size() etc. stay type- and interface-consistent.

---

4.3.4 comfy_aimdo/control.py

# Stub: real comfy_aimdo.control provides AIMDO device/VRAM control. Forge does not use it.


def init():
    pass


def init_device(device_index):
    return False


def get_total_vram_usage():
    """Return 0 so comfy.windows.get_free_ram() falls back to normal calculation."""
    return 0


def set_log_debug():
    pass
# ... set_log_critical, set_log_error, set_log_warning, set_log_info, analyze: same (pass)

Explanation

• init()
  Stub no-op for possible AIMDO initialization from ComfyUI’s main.py etc.

• init_device(device_index)
  Upstream returns True when AIMDO is available on that device. Stub always returns False so we are treated as “AIMDO not in use”.

• get_total_vram_usage()

  • comfy/windows.py’s get_free_ram() uses something like total - (committed - comfy_aimdo.control.get_total_vram_usage()).
  • Returning 0 from the stub means “no VRAM usage from AIMDO”, so the effective free RAM calculation is unchanged.

• set_log_* / analyze
  Stub no-ops for log level and analysis entry points; Forge does not use AIMDO logging or analysis.

---

4.4 How the stub gets imported

• In modules/paths.py, script_path (the repo root) is inserted at the front of sys.path.
• So import comfy_aimdo resolves to the comfy_aimdo/ directory under script_path.
• By placing the comfy_aimdo package at the project root, when ComfyUI’s comfy.model_management (etc.) imports comfy_aimdo, the stub is used and the ModuleNotFoundError is avoided, without installing the real package.

---

5. Summary

Issue	Cause	Fix	File(s)
Missing module reference	ComfyUI removed checkpoint_pickle.py	Remove that import and always use torch.load(..., weights_only=True)	backend/nn/comfy_utils.py
No module named 'comfy_aimdo'	ComfyUI now unconditionally imports comfy_aimdo	Add a stub package at project root (no AIMDO behavior)	comfy_aimdo/ (all files)

With these changes, Forge continues to run and run inference with its existing memory/loader even after updating ComfyUI-master to the latest core.