ThereforeGames
df17db844b
Store lazy-compile and model-adapter cached outputs as weakref.ref and add resolve/prune helpers so caches don't strongly retain base models. Eviction and lookup now resolve weakrefs, prune dead entries, and only dispose non-collected patchers. Added tests to assert caches don't retain base models and to ensure pre-quantized 'as_needed' inputs still trigger prior-cache cleanup. Bumped package version to 2.2.1 and updated CHANGELOG.md to document the fix.
310 lines
12 KiB
Python
310 lines
12 KiB
Python
import importlib
|
|
import gc
|
|
import sys
|
|
import unittest
|
|
import weakref
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest import mock
|
|
|
|
import torch
|
|
|
|
|
|
COMFY_ROOT = Path(__file__).resolve().parents[3]
|
|
CUSTOM_NODES_ROOT = COMFY_ROOT / "custom_nodes"
|
|
PACKAGE_NAME = Path(__file__).resolve().parents[1].name
|
|
sys.path.insert(0, str(COMFY_ROOT))
|
|
sys.path.insert(0, str(CUSTOM_NODES_ROOT))
|
|
|
|
from comfy.model_patcher import ModelPatcher
|
|
|
|
model_adapter = importlib.import_module(f"{PACKAGE_NAME}.int8_model_adapter")
|
|
|
|
|
|
class ModelAdapterCacheTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.last_model_ref = model_adapter._INT8_MODEL_ADAPTER_LAST_MODEL_REF
|
|
self.last_model_id = model_adapter._INT8_MODEL_ADAPTER_LAST_MODEL_ID
|
|
self.cache_limit = model_adapter._INT8_MODEL_ADAPTER_OUTPUT_CACHE_LIMIT
|
|
model_adapter._INT8_MODEL_ADAPTER_LAST_MODEL_REF = None
|
|
model_adapter._INT8_MODEL_ADAPTER_LAST_MODEL_ID = None
|
|
model_adapter._INT8_MODEL_ADAPTER_OUTPUT_CACHE_LIMIT = 1
|
|
|
|
def tearDown(self):
|
|
model_adapter._INT8_MODEL_ADAPTER_LAST_MODEL_REF = self.last_model_ref
|
|
model_adapter._INT8_MODEL_ADAPTER_LAST_MODEL_ID = self.last_model_id
|
|
model_adapter._INT8_MODEL_ADAPTER_OUTPUT_CACHE_LIMIT = self.cache_limit
|
|
|
|
def test_output_caches_are_local_to_each_base_model(self):
|
|
first_model = torch.nn.Module()
|
|
second_model = torch.nn.Module()
|
|
|
|
first_cache = model_adapter._get_output_cache(first_model)
|
|
second_cache = model_adapter._get_output_cache(second_model)
|
|
|
|
self.assertIsNot(first_cache, second_cache)
|
|
|
|
def test_output_cache_does_not_retain_base_model(self):
|
|
class CachedOutput:
|
|
pass
|
|
|
|
shared_model = torch.nn.Module()
|
|
cached_output = CachedOutput()
|
|
cached_output.model = shared_model
|
|
model_adapter._remember_cached_output(shared_model, ("output",), cached_output)
|
|
shared_model_ref = weakref.ref(shared_model)
|
|
cached_output_ref = weakref.ref(cached_output)
|
|
|
|
del cached_output
|
|
del shared_model
|
|
gc.collect()
|
|
|
|
self.assertIsNone(cached_output_ref())
|
|
self.assertIsNone(shared_model_ref())
|
|
|
|
def test_architecture_change_clears_prior_adapter_output(self):
|
|
first_model = torch.nn.Module()
|
|
second_model = torch.nn.Module()
|
|
first_cache = model_adapter._get_output_cache(first_model)
|
|
cached_model_patcher = mock.Mock()
|
|
first_cache[("first",)] = cached_model_patcher
|
|
model_adapter._INT8_MODEL_ADAPTER_LAST_MODEL_REF = weakref.ref(first_model)
|
|
model_adapter._INT8_MODEL_ADAPTER_LAST_MODEL_ID = id(first_model)
|
|
|
|
with mock.patch.object(model_adapter, "_cleanup_adapter_cache_memory") as cleanup_memory:
|
|
model_adapter._prepare_model_cache(second_model)
|
|
|
|
self.assertEqual(first_cache, {})
|
|
cached_model_patcher.unpatch_model.assert_called_once_with(unpatch_weights=True)
|
|
cleanup_memory.assert_called_once_with()
|
|
|
|
def test_as_needed_prequantized_model_still_cleans_prior_model_cache(self):
|
|
first_model = torch.nn.Module()
|
|
first_cache = model_adapter._get_output_cache(first_model)
|
|
cached_model_patcher = mock.Mock()
|
|
first_cache[("first",)] = cached_model_patcher
|
|
model_adapter._INT8_MODEL_ADAPTER_LAST_MODEL_REF = weakref.ref(first_model)
|
|
model_adapter._INT8_MODEL_ADAPTER_LAST_MODEL_ID = id(first_model)
|
|
|
|
second_model = torch.nn.Module()
|
|
second_model.diffusion_model = torch.nn.Module()
|
|
quantized_module = torch.nn.Linear(4, 4, bias=False)
|
|
quantized_module.quant_format = "int8_tensorwise"
|
|
second_model.diffusion_model.layer = quantized_module
|
|
second_model_patcher = ModelPatcher(second_model, torch.device("cpu"), torch.device("cpu"))
|
|
|
|
with mock.patch.object(model_adapter, "_cleanup_adapter_cache_memory") as cleanup_memory:
|
|
output = model_adapter.INT8ModelAdapter().apply_quantization(
|
|
second_model_patcher,
|
|
enable_quantization=model_adapter.QUANTIZATION_CONTROL_AS_NEEDED,
|
|
model_type="auto",
|
|
quantization_mode="int8_convrot",
|
|
log_progress=False,
|
|
)[0]
|
|
|
|
self.assertIs(output, second_model_patcher)
|
|
self.assertEqual(first_cache, {})
|
|
cached_model_patcher.unpatch_model.assert_called_once_with(unpatch_weights=True)
|
|
cleanup_memory.assert_called_once_with()
|
|
|
|
def test_cache_room_is_created_before_requantization(self):
|
|
cached_model_patcher = mock.Mock()
|
|
cache = {("first",): cached_model_patcher}
|
|
|
|
with mock.patch.object(model_adapter, "_cleanup_adapter_cache_memory") as cleanup_memory:
|
|
model_adapter._make_output_cache_room(cache)
|
|
|
|
self.assertEqual(cache, {})
|
|
cached_model_patcher.unpatch_model.assert_called_once_with(unpatch_weights=True)
|
|
cleanup_memory.assert_called_once_with()
|
|
|
|
def test_cache_eviction_does_not_copy_quantized_weight_backups_into_source_modules(self):
|
|
shared_model = torch.nn.Module()
|
|
shared_model.diffusion_model = torch.nn.Module()
|
|
shared_model.diffusion_model.layer = torch.nn.Linear(4, 4, bias=False, dtype=torch.float16)
|
|
shared_model.model_lowvram = False
|
|
shared_model.lowvram_patch_counter = 0
|
|
shared_model.model_loaded_weight_memory = 1
|
|
shared_model.model_offload_buffer_memory = 0
|
|
shared_model.current_weight_patches_uuid = "quantized"
|
|
shared_model.device = torch.device("cpu")
|
|
|
|
original_module = shared_model.diffusion_model.layer
|
|
base_model_patcher = ModelPatcher(shared_model, torch.device("cpu"), torch.device("cpu"))
|
|
quantized_model_patcher = base_model_patcher.clone()
|
|
quantized_module = torch.nn.Linear(4, 4, bias=False)
|
|
quantized_module.weight = torch.nn.Parameter(
|
|
torch.ones((4, 4), dtype=torch.int8),
|
|
requires_grad=False,
|
|
)
|
|
quantized_model_patcher.add_object_patch("diffusion_model.layer", quantized_module)
|
|
quantized_model_patcher.patch_model(load_weights=False)
|
|
quantized_model_patcher.backup["diffusion_model.layer.weight"] = SimpleNamespace(
|
|
weight=quantized_module.weight.detach().clone(),
|
|
inplace_update=False,
|
|
)
|
|
|
|
model_adapter._dispose_cached_adapter_output(quantized_model_patcher)
|
|
replacement_model_patcher = base_model_patcher.clone()
|
|
replacement_model_patcher.unpatch_model(unpatch_weights=True)
|
|
|
|
self.assertIs(shared_model.diffusion_model.layer, original_module)
|
|
self.assertEqual(shared_model.diffusion_model.layer.weight.dtype, torch.float16)
|
|
self.assertEqual(replacement_model_patcher.backup, {})
|
|
|
|
def test_as_needed_ignores_quantization_temporarily_installed_by_sibling_patcher(self):
|
|
shared_model = torch.nn.Module()
|
|
shared_model.diffusion_model = torch.nn.Module()
|
|
shared_model.diffusion_model.layer = torch.nn.Linear(4, 4, bias=False)
|
|
base_model_patcher = ModelPatcher(shared_model, torch.device("cpu"), torch.device("cpu"))
|
|
quantized_model_patcher = base_model_patcher.clone()
|
|
quantized_module = torch.nn.Linear(4, 4, bias=False)
|
|
quantized_module._is_quantized = True
|
|
quantized_module._quant_format = "asym_w4a8_int8"
|
|
quantized_model_patcher.add_object_patch("diffusion_model.layer", quantized_module)
|
|
quantized_model_patcher.patch_model(load_weights=False)
|
|
|
|
self.assertIs(shared_model.diffusion_model.layer, quantized_module)
|
|
self.assertTrue(model_adapter._model_has_quantized_modules(quantized_model_patcher))
|
|
self.assertFalse(model_adapter._model_has_quantized_modules(base_model_patcher))
|
|
|
|
def test_as_needed_recognizes_quantized_underlying_module_while_object_patch_is_active(self):
|
|
shared_model = torch.nn.Module()
|
|
shared_model.diffusion_model = torch.nn.Module()
|
|
quantized_module = torch.nn.Linear(4, 4, bias=False)
|
|
quantized_module.quant_format = "int8_tensorwise"
|
|
shared_model.diffusion_model.layer = quantized_module
|
|
base_model_patcher = ModelPatcher(shared_model, torch.device("cpu"), torch.device("cpu"))
|
|
patched_model_patcher = base_model_patcher.clone()
|
|
patched_model_patcher.add_object_patch("diffusion_model.layer", torch.nn.Identity())
|
|
patched_model_patcher.patch_model(load_weights=False)
|
|
|
|
self.assertTrue(model_adapter._model_has_quantized_modules(base_model_patcher))
|
|
|
|
def test_stale_toolkit_class_identity_does_not_block_object_patch_cleanup(self):
|
|
stale_quantized_module = torch.nn.Linear(4, 4, bias=False)
|
|
stale_quantized_module._is_quantized = True
|
|
stale_quantized_module._quant_format = "asym_w4a8_int8"
|
|
model_patcher = SimpleNamespace(
|
|
object_patches={"diffusion_model.layer": stale_quantized_module},
|
|
)
|
|
|
|
model_adapter._clear_prior_int8_object_patches(model_patcher)
|
|
|
|
self.assertEqual(model_patcher.object_patches, {})
|
|
|
|
def test_prior_adapter_reset_fully_unpatches_weights(self):
|
|
quantized_module = model_adapter.Int8TensorwiseOps.Linear(4, 4, bias=False)
|
|
original_module = torch.nn.Linear(4, 4, bias=False)
|
|
shared_model = torch.nn.Module()
|
|
shared_model.diffusion_model = torch.nn.Module()
|
|
shared_model.diffusion_model.layer = quantized_module
|
|
setattr(
|
|
shared_model,
|
|
model_adapter._INT8_MODEL_ADAPTER_ORIGINAL_MODULES_KEY,
|
|
{"diffusion_model.layer": original_module},
|
|
)
|
|
model_patcher = SimpleNamespace(
|
|
model=shared_model,
|
|
object_patches={"diffusion_model.layer": quantized_module},
|
|
object_patches_backup={},
|
|
get_attachment=lambda _key: None,
|
|
set_attachments=lambda _key, _value: None,
|
|
unpatch_model=mock.Mock(),
|
|
)
|
|
|
|
model_adapter._reset_prior_int8_object_patches(model_patcher)
|
|
|
|
model_patcher.unpatch_model.assert_called_once_with(unpatch_weights=True)
|
|
|
|
def test_prior_adapter_modules_are_restored_even_when_the_live_class_changed(self):
|
|
original_module = torch.nn.Linear(4, 4, bias=False)
|
|
stale_module = torch.nn.Linear(4, 4, bias=False)
|
|
diffusion_model = torch.nn.Module()
|
|
diffusion_model.layer = stale_module
|
|
shared_model = torch.nn.Module()
|
|
shared_model.diffusion_model = diffusion_model
|
|
setattr(
|
|
shared_model,
|
|
model_adapter._INT8_MODEL_ADAPTER_ORIGINAL_MODULES_KEY,
|
|
{"diffusion_model.layer": original_module},
|
|
)
|
|
|
|
model_patcher = SimpleNamespace(
|
|
model=shared_model,
|
|
object_patches={},
|
|
object_patches_backup={},
|
|
get_attachment=lambda _key: None,
|
|
set_attachments=lambda _key, _value: None,
|
|
unpatch_model=mock.Mock(),
|
|
)
|
|
|
|
restored_count = model_adapter._reset_prior_int8_object_patches(model_patcher)
|
|
|
|
self.assertEqual(restored_count, 1)
|
|
self.assertIs(shared_model.diffusion_model.layer, original_module)
|
|
|
|
def test_checkpoint_full_precision_mm_source_hint_does_not_block_requantization(self):
|
|
diffusion_model = torch.nn.Module()
|
|
diffusion_model.regular = torch.nn.Linear(256, 16, bias=False)
|
|
diffusion_model.checkpoint_full_precision = torch.nn.Linear(256, 16, bias=False)
|
|
diffusion_model.checkpoint_full_precision._full_precision_mm_config = True
|
|
|
|
candidate_names = {
|
|
module_name
|
|
for module_name, _module in model_adapter._collect_int8_candidates(diffusion_model, ())
|
|
}
|
|
|
|
self.assertIn("regular", candidate_names)
|
|
self.assertIn("checkpoint_full_precision", candidate_names)
|
|
|
|
def test_anima_dtype_guard_does_not_replace_preprocessor(self):
|
|
class DiffusionModel(torch.nn.Module):
|
|
def preprocess_text_embeds(self, text_embeds, text_ids):
|
|
return text_embeds + text_ids.to(text_embeds.dtype)
|
|
|
|
diffusion_model = DiffusionModel()
|
|
adapter_state = {"model_type": "anima", "log_progress": False}
|
|
transformer_options = {"int8_model_adapter": adapter_state}
|
|
text_embeds = torch.ones(1, dtype=torch.float32)
|
|
text_ids = torch.ones(1, dtype=torch.long)
|
|
|
|
class Executor:
|
|
class_obj = SimpleNamespace(diffusion_model=diffusion_model)
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
self.assert_preprocessor_is_unmodified()
|
|
return diffusion_model.preprocess_text_embeds(text_embeds, kwargs["t5xxl_ids"])
|
|
|
|
def assert_preprocessor_is_unmodified(self):
|
|
if "preprocess_text_embeds" in diffusion_model.__dict__:
|
|
raise AssertionError("Anima preprocessor was replaced on the module instance")
|
|
|
|
result = model_adapter._int8_model_adapter_notice_wrapper(
|
|
Executor(),
|
|
transformer_options=transformer_options,
|
|
t5xxl_ids=text_ids,
|
|
)
|
|
|
|
self.assertTrue(torch.equal(result, torch.full((1,), 2.0)))
|
|
self.assertNotIn("preprocess_text_embeds", diffusion_model.__dict__)
|
|
|
|
def test_anima_dtype_guard_rejects_float_token_ids_before_forward(self):
|
|
adapter_state = {"model_type": "anima", "log_progress": False}
|
|
transformer_options = {"int8_model_adapter": adapter_state}
|
|
executor = mock.Mock()
|
|
executor.class_obj = SimpleNamespace(diffusion_model=torch.nn.Module())
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "Anima requires integer token ids"):
|
|
model_adapter._int8_model_adapter_notice_wrapper(
|
|
executor,
|
|
transformer_options=transformer_options,
|
|
t5xxl_ids=torch.ones(1, dtype=torch.float32),
|
|
)
|
|
|
|
executor.assert_not_called()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|