90 lines
2.2 KiB
Python
90 lines
2.2 KiB
Python
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock
|
|
|
|
import torch
|
|
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
|
|
|
mock_model_management = MagicMock()
|
|
mock_model_management.intermediate_device.return_value = "cpu"
|
|
|
|
mock_comfy = MagicMock()
|
|
mock_comfy.model_management = mock_model_management
|
|
|
|
mock_hooks = MagicMock()
|
|
|
|
|
|
class DummyHook:
|
|
def __init__(self, *args, **kwargs):
|
|
pass
|
|
|
|
def clone(self):
|
|
return self
|
|
|
|
|
|
class DummyHookGroup:
|
|
def add(self, hook):
|
|
self.hook = hook
|
|
|
|
|
|
mock_hooks.TransformerOptionsHook = DummyHook
|
|
mock_hooks.HookGroup = DummyHookGroup
|
|
mock_hooks.EnumHookScope = MagicMock()
|
|
mock_hooks.set_hooks_for_conditioning = (
|
|
lambda conditioning, hooks=None, append_hooks=True: conditioning
|
|
)
|
|
mock_comfy.hooks = mock_hooks
|
|
|
|
server = MagicMock()
|
|
server.PromptServer.instance.routes.post = lambda path: (lambda fn: fn)
|
|
server.PromptServer.instance.routes.get = lambda path: (lambda fn: fn)
|
|
|
|
sys.modules["comfy"] = mock_comfy
|
|
sys.modules["comfy.model_management"] = mock_model_management
|
|
sys.modules["comfy.hooks"] = mock_hooks
|
|
sys.modules["server"] = server
|
|
|
|
from A1111_Prompt_Node.nodes.prompt_node import A1111PromptNode
|
|
|
|
|
|
class FakeCondStageModel:
|
|
pass
|
|
|
|
|
|
class FakeClip:
|
|
def __init__(self):
|
|
self.cond_stage_model = FakeCondStageModel()
|
|
|
|
|
|
class CacheProbePromptNode(A1111PromptNode):
|
|
def _get_downstream_steps(self, prompt, start_node_id, debug=False):
|
|
return 2
|
|
|
|
def _encode_with_break_isolation(self, clip, prompt_text, normalization, is_sdxl, debug):
|
|
seq_len = 2 if prompt_text == "short" else 5
|
|
cond = torch.ones((1, seq_len, 3))
|
|
pooled = torch.ones((1, 3))
|
|
return cond, pooled
|
|
|
|
|
|
class TestPromptNodeCache(unittest.TestCase):
|
|
def test_scheduled_padding_does_not_mutate_encoded_cache(self):
|
|
node = CacheProbePromptNode()
|
|
clip = FakeClip()
|
|
|
|
node.encode(clip, "[short:long:0.5]", prompt={}, unique_id="1")
|
|
|
|
cached_shapes = {
|
|
key[2]: cond.shape[1]
|
|
for key, (cond, _pooled) in node._encoded_cache.items()
|
|
}
|
|
self.assertEqual(cached_shapes["short"], 2)
|
|
self.assertEqual(cached_shapes["long"], 5)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|