121 lines
4.2 KiB
Python
121 lines
4.2 KiB
Python
"""Load node implementations without importing ComfyUI or the pack entry point.
|
|
|
|
The stubs model only the host boundary; tensor operations under test remain real
|
|
CPU PyTorch. Every temporary sys.modules entry is restored after the import.
|
|
"""
|
|
|
|
import copy
|
|
import importlib.util
|
|
from pathlib import Path
|
|
import sys
|
|
import types
|
|
from unittest.mock import patch
|
|
|
|
import numpy # Load dependencies before isolating sys.modules.
|
|
import torch
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def load_implementation(relative_path, name, stubs=None):
|
|
spec = importlib.util.spec_from_file_location(name, ROOT / relative_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
with patch.dict(sys.modules, {**(stubs or {}), name: module}):
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def lanpaint_module():
|
|
comfy = types.ModuleType("comfy")
|
|
comfy.model_patcher = types.ModuleType("comfy.model_patcher")
|
|
comfy.model_patcher.create_model_options_clone = copy.deepcopy
|
|
comfy.patcher_extension = types.ModuleType("comfy.patcher_extension")
|
|
comfy.patcher_extension.WrappersMP = types.SimpleNamespace(SAMPLER_SAMPLE="sampler_sample")
|
|
return load_implementation(
|
|
"gen2/lanpaint_nodes/lanpaint_soft_denoise.py", "_gen2_test_lanpaint",
|
|
{"comfy": comfy, "comfy.model_patcher": comfy.model_patcher,
|
|
"comfy.patcher_extension": comfy.patcher_extension},
|
|
)
|
|
|
|
|
|
def klein_module():
|
|
comfy = types.ModuleType("comfy")
|
|
comfy.ldm = types.ModuleType("comfy.ldm")
|
|
comfy.ldm.flux = types.ModuleType("comfy.ldm.flux")
|
|
comfy.ldm.flux.model = types.ModuleType("comfy.ldm.flux.model")
|
|
comfy.ldm.flux.layers = types.ModuleType("comfy.ldm.flux.layers")
|
|
|
|
class Flux:
|
|
def forward_orig(self):
|
|
return "native forward_orig"
|
|
|
|
def _forward(self):
|
|
return "native _forward"
|
|
|
|
comfy.ldm.flux.model.Flux = Flux
|
|
# Full transformer execution is outside these CPU wrapper lifecycle tests.
|
|
comfy.ldm.flux.layers.timestep_embedding = lambda value, dim: value
|
|
names = {
|
|
"comfy": comfy, "comfy.ldm": comfy.ldm,
|
|
"comfy.ldm.flux": comfy.ldm.flux,
|
|
"comfy.ldm.flux.model": comfy.ldm.flux.model,
|
|
"comfy.ldm.flux.layers": comfy.ldm.flux.layers,
|
|
}
|
|
module = load_implementation(
|
|
"experimental/f2k/flux2_klein_fix.py", "_gen2_test_f2k", names,
|
|
)
|
|
return module, Flux
|
|
|
|
|
|
class ModelPatcherDouble:
|
|
"""Minimal independent options/attachment clone for node-boundary tests."""
|
|
|
|
def __init__(self):
|
|
self.wrappers = {}
|
|
self.attachments = {}
|
|
self.model_options = {}
|
|
|
|
def clone(self):
|
|
result = type(self)()
|
|
result.wrappers = {kind: dict(items) for kind, items in self.wrappers.items()}
|
|
result.attachments = dict(self.attachments)
|
|
result.model_options = dict(self.model_options)
|
|
return result
|
|
|
|
def remove_wrappers_with_key(self, kind, key):
|
|
self.wrappers.get(kind, {}).pop(key, None)
|
|
|
|
def add_wrapper_with_key(self, kind, key, wrapper):
|
|
self.wrappers.setdefault(kind, {})[key] = wrapper
|
|
|
|
def remove_attachments(self, key):
|
|
self.attachments.pop(key, None)
|
|
|
|
def set_attachments(self, key, value):
|
|
self.attachments[key] = value
|
|
|
|
def set_model_unet_function_wrapper(self, wrapper):
|
|
self.model_options["model_function_wrapper"] = wrapper
|
|
|
|
|
|
class LanPaintRuntimeDouble:
|
|
"""Deterministic runtime with explicit source noising and captured calls."""
|
|
|
|
def __init__(self, source=None, generated=10.0):
|
|
self.latent_image = torch.full((1, 1, 1, 3), 2.0) if source is None else source
|
|
self.noise = torch.full_like(self.latent_image, 3.0)
|
|
self.sigmas = torch.tensor([1.0, 0.5, 0.0])
|
|
self.PaintMethod = "test"
|
|
self.LanPaint_early_stop = False
|
|
self.generated = generated
|
|
self.calls = []
|
|
sampling = types.SimpleNamespace(
|
|
noise_scaling=lambda sigma, noise, latent: latent + sigma * noise,
|
|
)
|
|
self.inner_model = types.SimpleNamespace(inner_model=types.SimpleNamespace(model_sampling=sampling))
|
|
|
|
def __call__(self, x, sigma, **kwargs):
|
|
self.calls.append((x.detach().clone(), sigma.detach().clone(), kwargs))
|
|
return torch.full_like(x, self.generated)
|