185 lines
6.5 KiB
Python
185 lines
6.5 KiB
Python
# SimpleSyrup - workflow-focused ComfyUI extensions for image generation
|
|
# Copyright (C) 2026 Artificial Sweetener and contributors
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
"""Tests for SimpleSyrup's shared torch model device manager."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from types import ModuleType
|
|
from typing import Any, cast
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from simple_syrup.runtime.model_device_manager import (
|
|
TorchModelDeviceManager,
|
|
external_model_inference,
|
|
)
|
|
|
|
|
|
class RecordingModel:
|
|
"""Small PyTorch-style model double."""
|
|
|
|
def __init__(self, device: str | None = None) -> None:
|
|
"""Create a model that records device movement."""
|
|
|
|
self.device = torch.device(device) if device is not None else None
|
|
self.to_calls: list[str] = []
|
|
self.eval_calls = 0
|
|
|
|
def to(self, device: str | torch.device) -> None:
|
|
"""Record model movement."""
|
|
|
|
self.device = torch.device(device)
|
|
self.to_calls.append(str(device))
|
|
|
|
def eval(self) -> None:
|
|
"""Record eval mode selection."""
|
|
|
|
self.eval_calls += 1
|
|
|
|
|
|
def test_manager_cpu_policy_does_not_call_comfy_gpu_loader() -> None:
|
|
"""CPU inference keeps the model on CPU without Comfy GPU loading."""
|
|
|
|
model = RecordingModel()
|
|
managed = TorchModelDeviceManager().manage(model, "model", "source")
|
|
|
|
with TorchModelDeviceManager().inference(managed, "cpu") as loaded:
|
|
assert loaded.device == torch.device("cpu")
|
|
assert loaded.model is model
|
|
|
|
assert model.to_calls == ["cpu"]
|
|
assert model.eval_calls >= 2
|
|
|
|
|
|
def test_manager_auto_policy_uses_comfy_model_patcher(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Auto inference loads owned models through Comfy's model manager."""
|
|
|
|
state: dict[str, object] = {"loaded": []}
|
|
|
|
class FakeModelPatcher:
|
|
"""Fake Comfy ModelPatcher boundary."""
|
|
|
|
def __init__(
|
|
self,
|
|
model: object,
|
|
load_device: torch.device,
|
|
offload_device: torch.device,
|
|
) -> None:
|
|
"""Record patcher construction and mimic Comfy's device attribute."""
|
|
|
|
self.model = model
|
|
self.load_device = load_device
|
|
self.offload_device = offload_device
|
|
cast(Any, model).device = load_device
|
|
|
|
comfy = ModuleType("comfy")
|
|
model_management = ModuleType("comfy.model_management")
|
|
model_patcher = ModuleType("comfy.model_patcher")
|
|
model_management.get_torch_device = lambda: torch.device("cuda:0") # type: ignore[attr-defined]
|
|
model_management.text_encoder_offload_device = lambda: torch.device("cpu") # type: ignore[attr-defined]
|
|
model_management.load_model_gpu = lambda patcher: state["loaded"].append(patcher) # type: ignore[attr-defined]
|
|
model_patcher.ModelPatcher = FakeModelPatcher # type: ignore[attr-defined]
|
|
monkeypatch.setitem(sys.modules, "comfy", comfy)
|
|
monkeypatch.setitem(sys.modules, "comfy.model_management", model_management)
|
|
monkeypatch.setitem(sys.modules, "comfy.model_patcher", model_patcher)
|
|
|
|
model = RecordingModel("cpu")
|
|
manager = TorchModelDeviceManager()
|
|
managed = manager.manage(model, "model", "source")
|
|
|
|
with manager.inference(managed, "auto") as loaded:
|
|
assert loaded.device == torch.device("cuda:0")
|
|
|
|
loaded_patchers = state["loaded"]
|
|
assert isinstance(loaded_patchers, list)
|
|
assert len(loaded_patchers) == 1
|
|
assert managed.patcher is loaded_patchers[0]
|
|
|
|
|
|
def test_manager_auto_policy_bypasses_patcher_for_read_only_device_model(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Models like SAM with read-only `device` properties use bounded `.to(...)`."""
|
|
|
|
class ReadOnlyDeviceModel:
|
|
"""PyTorch-like model with a read-only device property."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Create a model on CPU."""
|
|
|
|
self._device = torch.device("cpu")
|
|
self.to_calls: list[str] = []
|
|
self.eval_calls = 0
|
|
|
|
@property
|
|
def device(self) -> torch.device:
|
|
"""Return the current device without allowing assignment."""
|
|
|
|
return self._device
|
|
|
|
def to(self, device: str | torch.device) -> None:
|
|
"""Record model movement."""
|
|
|
|
self._device = torch.device(device)
|
|
self.to_calls.append(str(device))
|
|
|
|
def eval(self) -> None:
|
|
"""Record eval mode selection."""
|
|
|
|
self.eval_calls += 1
|
|
|
|
state: dict[str, object] = {"loaded": [], "emptied": 0}
|
|
|
|
comfy = ModuleType("comfy")
|
|
model_management = ModuleType("comfy.model_management")
|
|
model_patcher = ModuleType("comfy.model_patcher")
|
|
model_management.get_torch_device = lambda: torch.device("cuda:0") # type: ignore[attr-defined]
|
|
model_management.text_encoder_offload_device = lambda: torch.device("cpu") # type: ignore[attr-defined]
|
|
model_management.load_model_gpu = lambda patcher: state["loaded"].append(patcher) # type: ignore[attr-defined]
|
|
model_management.soft_empty_cache = lambda: state.__setitem__("emptied", 1) # type: ignore[attr-defined]
|
|
model_patcher.ModelPatcher = object # type: ignore[attr-defined]
|
|
monkeypatch.setitem(sys.modules, "comfy", comfy)
|
|
monkeypatch.setitem(sys.modules, "comfy.model_management", model_management)
|
|
monkeypatch.setitem(sys.modules, "comfy.model_patcher", model_patcher)
|
|
|
|
model = ReadOnlyDeviceModel()
|
|
manager = TorchModelDeviceManager()
|
|
managed = manager.manage(model, "sam", "source")
|
|
|
|
with manager.inference(managed, "auto") as loaded:
|
|
assert loaded.device == torch.device("cuda:0")
|
|
assert model.device == torch.device("cuda:0")
|
|
|
|
assert state["loaded"] == []
|
|
assert state["emptied"] == 1
|
|
assert managed.patcher is None
|
|
assert model.device == torch.device("cpu")
|
|
assert model.to_calls == ["cuda:0", "cpu"]
|
|
|
|
|
|
def test_external_model_inference_restores_original_device(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""External raw model movement is bounded to the inference context."""
|
|
|
|
comfy = ModuleType("comfy")
|
|
model_management = ModuleType("comfy.model_management")
|
|
model_management.get_torch_device = lambda: torch.device("cuda:0") # type: ignore[attr-defined]
|
|
monkeypatch.setitem(sys.modules, "comfy", comfy)
|
|
monkeypatch.setitem(sys.modules, "comfy.model_management", model_management)
|
|
|
|
model = RecordingModel("cpu")
|
|
|
|
with external_model_inference(model, "auto") as loaded:
|
|
assert loaded.device == torch.device("cuda:0")
|
|
assert model.device == torch.device("cuda:0")
|
|
|
|
assert model.device == torch.device("cpu")
|
|
assert model.to_calls == ["cuda:0", "cpu"]
|