fix: make comfy_aimdo import optional for AMD/ROCm GPU support#13266
fix: make comfy_aimdo import optional for AMD/ROCm GPU support#13266octo-patch wants to merge 1 commit intoComfy-Org:masterfrom
Conversation
When comfy-aimdo is not installed (e.g. on AMD/ROCm systems), the bare 'import comfy_aimdo.control' at module level causes an immediate ModuleNotFoundError that prevents ComfyUI from starting at all. This change wraps the import in a try/except block. On ImportError, a minimal no-op stub is installed in sys.modules so that subsequent 'import comfy_aimdo.*' statements in other modules (execution.py, comfy/ops.py, comfy/model_patcher.py, comfy/pinned_memory.py, etc.) also succeed without modification. The stub returns False from all calls, which: - Prevents comfy_aimdo.control.init_device() from signalling success - Keeps comfy.memory_management.aimdo_enabled=False (its default) - Allows all aimdo-guarded code paths to be safely skipped Fixes Comfy-Org#13182
octo-patch
requested review from
Kosinkadink,
comfyanonymous and
guill
as code owners
📝 WalkthroughWalkthroughThe 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@main.py`:
- Around line 34-37: The import block for comfy_aimdo.control currently swallows
all ImportError instances; change it so only a missing-module error falls back:
catch ModuleNotFoundError (or catch ImportError as e and re-raise unless e.name
indicates the comfy_aimdo module is missing) for the import of
comfy_aimdo.control and log the "not installed" message only in that case, while
allowing other ImportError failures in comfy_aimdo.control to propagate so
DynamicVRAM does not silently disappear.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
| try: | ||
| import comfy_aimdo.control | ||
| except ImportError: | ||
| logging.info("comfy-aimdo not installed. DynamicVRAM will not be available.") |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# Check the git history and diff around main.py to confirm if lines 34-37 are newly introduced
git log --oneline -n 20 main.py 2>/dev/null || echo "Git history unavailable"
git diff HEAD~1 main.py 2>/dev/null | head -100 || echo "Git diff unavailable - checking file directly"Repository: Comfy-Org/ComfyUI
Length of output: 131
🏁 Script executed:
# Look at the actual main.py file to see the import block and surrounding context
wc -l main.py
head -80 main.pyRepository: Comfy-Org/ComfyUI
Length of output: 2611
🏁 Script executed:
# Search for all uses of comfy_aimdo in the codebase
rg "comfy_aimdo" --type py -nRepository: Comfy-Org/ComfyUI
Length of output: 2364
🏁 Script executed:
# Search for model_mmap usage specifically
rg "model_mmap\|ModelMMAP" --type py -n -A 2 -B 2Repository: Comfy-Org/ComfyUI
Length of output: 43
🏁 Script executed:
# Verify Python ImportError.name behavior and availability
python3 << 'EOF'
import sys
print(f"Python version: {sys.version}")
# Test 1: ImportError.name attribute
try:
import nonexistent_module_xyz
except ImportError as e:
print(f"ImportError.name: {getattr(e, 'name', 'ATTR_NOT_PRESENT')}")
print(f"ImportError attributes: {dir(e)}")
print(f"ImportError str: {str(e)}")
# Test 2: ImportError with submodule
try:
import os.nonexistent_submod
except ImportError as e:
print(f"Submodule ImportError.name: {getattr(e, 'name', 'ATTR_NOT_PRESENT')}")
print(f"Submodule ImportError str: {str(e)}")
EOFRepository: Comfy-Org/ComfyUI
Length of output: 867
🏁 Script executed:
# Check comfy/utils.py to understand model_mmap usage pattern
sed -n '80,100p' comfy/utils.pyRepository: Comfy-Org/ComfyUI
Length of output: 676
🏁 Script executed:
# Check if there are any guards or conditionals around the model_mmap usage
rg "model_mmap" -B 5 -A 10 comfy/utils.pyRepository: Comfy-Org/ComfyUI
Length of output: 1439
🏁 Script executed:
# Check execution.py to understand how aimdo_enabled guards work
sed -n '530,550p' execution.pyRepository: Comfy-Org/ComfyUI
Length of output: 1329
🏁 Script executed:
# Check if load_safetensors is always called or if there's a guard
rg "load_safetensors" -B 3 -A 3 --type py | head -50Repository: Comfy-Org/ComfyUI
Length of output: 696
🏁 Script executed:
# Check the full context of comfy/utils.py load_safetensors to see if there's any guard we missed
sed -n '85,120p' comfy/utils.pyRepository: Comfy-Org/ComfyUI
Length of output: 1583
🏁 Script executed:
# Verify if comfy_aimdo.model_mmap is imported conditionally anywhere
rg "aimdo_enabled|comfy_aimdo.model_mmap" --type py -B 2 -A 2Repository: Comfy-Org/ComfyUI
Length of output: 5693
🏁 Script executed:
# Final check: confirm that when load_safetensors is called, aimdo_enabled must be True
sed -n '100,140p' comfy/utils.py | grep -A 10 -B 5 "load_safetensors"Repository: Comfy-Org/ComfyUI
Length of output: 839
🏁 Script executed:
# Verify what happens in main.py when aimdo_enabled is set
sed -n '240,260p' main.pyRepository: Comfy-Org/ComfyUI
Length of output: 1214
Only fall back when comfy_aimdo is actually missing.
Line 36 catches every ImportError, so a broken comfy_aimdo.control install on a supported system gets silently reclassified as "not installed" and DynamicVRAM just disappears. Gate the fallback on the missing-module case instead of swallowing all import failures.
Proposed fix
try:
import comfy_aimdo.control
-except ImportError:
+except ImportError as e:
+ if getattr(e, "name", None) not in {"comfy_aimdo", "comfy_aimdo.control"}:
+ raise
logging.info("comfy-aimdo not installed. DynamicVRAM will not be available.")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@main.py` around lines 34 - 37, The import block for comfy_aimdo.control
currently swallows all ImportError instances; change it so only a missing-module
error falls back: catch ModuleNotFoundError (or catch ImportError as e and
re-raise unless e.name indicates the comfy_aimdo module is missing) for the
import of comfy_aimdo.control and log the "not installed" message only in that
case, while allowing other ImportError failures in comfy_aimdo.control to
propagate so DynamicVRAM does not silently disappear.
1 participant
Fixes #13182
Problem
Since commit
f8acd9c40,main.pycontains a bare top-level import:comfy-aimdois an NVIDIA-specific package for the DynamicVRAM feature. On AMD/ROCm systems (and any other non-NVIDIA setup where it is not installed), this import raisesModuleNotFoundErrorimmediately at startup, preventing ComfyUI from running at all.The subsequent imports in
comfy/ops.py,comfy/model_patcher.py,comfy/pinned_memory.py, andexecution.pyalso referencecomfy_aimdo.*at module level, and would fail in turn once the first import is attempted.Solution
Wrap the
import comfy_aimdo.controlin atry/except ImportErrorblock. On failure, a minimal no-op stub is installed intosys.modulesso that all subsequentimport comfy_aimdo.*statements in other modules succeed without requiring changes to those files.The stub class:
_AimdoStub()from all attribute accesses (__getattr__)False(falsy) from all calls (__call__)__bool__,__int__,__float__returningFalse/0/0.0This means:
comfy_aimdo.control.init_device(...)returnsFalse→aimdo_enabledstaysFalsecomfy.memory_management.aimdo_enabledremainsFalse(its default value)if comfy.memory_management.aimdo_enabled:guards correctly skip Dynamic VRAM code pathscomfy_aimdo.control.get_total_vram_usage()returns0(correct forcomfy/windows.pyRAM calculation)No other files need to be changed because all actual uses of
comfy_aimdo.*functions are already guarded by runtime checks that prevent them from executing when aimdo is not active.Testing
trybranch succeeds and behavior is unchanged.elsebranch at line 229).