Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 35 additions & 1 deletion main.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,41 @@

faulthandler.enable(file=sys.stderr, all_threads=False)

import comfy_aimdo.control
try:
import comfy_aimdo.control
except ImportError:
logging.info("comfy-aimdo not installed. DynamicVRAM will not be available.")
Comment on lines +34 to +37
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 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.py

Repository: Comfy-Org/ComfyUI

Length of output: 2611


🏁 Script executed:

# Search for all uses of comfy_aimdo in the codebase
rg "comfy_aimdo" --type py -n

Repository: 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 2

Repository: 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)}")
EOF

Repository: 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.py

Repository: 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.py

Repository: Comfy-Org/ComfyUI

Length of output: 1439


🏁 Script executed:

# Check execution.py to understand how aimdo_enabled guards work
sed -n '530,550p' execution.py

Repository: 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 -50

Repository: 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.py

Repository: 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 2

Repository: 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.py

Repository: 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.

import types

class _AimdoStub:
"""No-op stub for comfy_aimdo submodules when the package is not installed.

All attribute accesses return self and all calls return False (falsy),
which prevents any aimdo-guarded code paths from activating while still
allowing the import to succeed on non-NVIDIA hardware (e.g. AMD/ROCm).
"""
def __getattr__(self, name):
return _AimdoStub()

def __call__(self, *args, **kwargs):
return False

def __bool__(self):
return False

def __int__(self):
return 0

def __float__(self):
return 0.0

_aimdo_pkg = types.ModuleType('comfy_aimdo')
for _submod_name in ('control', 'model_vbar', 'torch', 'host_buffer', 'model_mmap'):
_stub = _AimdoStub()
setattr(_aimdo_pkg, _submod_name, _stub)
sys.modules[f'comfy_aimdo.{_submod_name}'] = _stub
sys.modules['comfy_aimdo'] = _aimdo_pkg
import comfy_aimdo # Now resolves to the stub from sys.modules

if enables_dynamic_vram():
comfy_aimdo.control.init()
Expand Down