Keep the execution helpers importable without torch

Moving comfy_torch_device into _execution_helpers put a top-level torch
import there, which broke the standard-library test run in CI. torch is
now imported inside that one function, and its tests skip where torch is
absent.
This commit is contained in:
ausboss
2026-09-24 00:22:16 -05:00
parent 84ee1fd366
commit 01e77fafd3
2 changed files with 14 additions and 5 deletions
+8 -4
View File
@@ -11,8 +11,6 @@ compute device live here too, so no helper module keeps its own copy.
from __future__ import annotations
import torch
def raise_if_interrupted() -> None:
"""Abort the running node as soon as the queue is cancelled.
@@ -98,8 +96,14 @@ def warn_once(
print(f"[AusBoss] {message}")
def comfy_torch_device() -> torch.device:
"""The device ComfyUI computes on; CUDA when available outside ComfyUI."""
def comfy_torch_device():
"""The device ComfyUI computes on; CUDA when available outside ComfyUI.
torch is imported here, not at the top, so the rest of this module stays
importable where torch is absent (the standard-library tests in CI).
"""
import torch
try:
from comfy.model_management import get_torch_device
+6 -1
View File
@@ -8,7 +8,10 @@ from pathlib import Path
import sys
from unittest.mock import patch
import torch
try:
import torch
except ImportError: # CI runs this file with the standard library alone.
torch = None
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))
@@ -162,6 +165,7 @@ class WarnOnceTests(unittest.TestCase):
self.assertEqual(len(kept), 5)
@unittest.skipIf(torch is None, "needs torch")
class ComfyTorchDeviceTests(unittest.TestCase):
def test_comfy_chooses_the_device(self):
entries = fake_comfy(model_management={"get_torch_device": lambda: "cpu"})
@@ -174,6 +178,7 @@ class ComfyTorchDeviceTests(unittest.TestCase):
self.assertEqual(comfy_torch_device().type, expected)
@unittest.skipIf(torch is None, "the helper modules need torch")
class SharedSeamTests(unittest.TestCase):
def test_the_helpers_share_one_copy_of_each_seam(self):
from nodes import _inpaint_crop_helpers, _interpolate_helpers, _lama_helpers, _mask_helpers