"""Cache-determinism regression: the reported Krea2 repro, simulated (v2.16.0, plan S3). User report (2026-09-17): with identical HiFlow parameters, results differ unless model AND node caches are cleared before each run; the HiFlow console shows a different stage entry sigma (0.3872 vs 0.5055) for identical params. Root cause (verified against the installed ComfyUI source, see src/effective_sampling.py docstring): object patches applied by a patch-node clone STAY on the shared BaseModel after that clone is superseded (``load_models_gpu`` detaches stale entries with ``detach(unpatch_all=False)``, model_management.py:962; ``partially_unload`` never restores object patches). Restore-to-original only happens through ``object_patches_backup`` — a dict SHARED between a patcher and its clones (model_patcher.py:428/462) — so a patcher whose own backup is empty and that carries no ``model_sampling`` patch resolves the live (leaked) attribute. This file encodes those semantics in a FakeComfyLifecycle harness and asserts the v2.16.0 contract: HiFlow's resolved model_sampling / stage schedule is a function of the GRAPH's patch chain, not of what ran earlier in the session — even when the live BaseModel attribute has diverged. If a future ComfyUI changes these lifecycle semantics, only THIS harness needs updating. """ from __future__ import annotations import sys import types from pathlib import Path import pytest import torch sys.path.insert(0, str(Path(__file__).parent.parent)) import nodes.hiflow as hfn # noqa: E402 from src.effective_sampling import effective_model_sampling # noqa: E402 from src.hiflow import build_stage_sigmas # noqa: E402 # --------------------------------------------------------------------------- # Fake ComfyUI lifecycle # --------------------------------------------------------------------------- class FakeBaseModel: """The SHARED BaseModel: object patches mutate its live attributes.""" def __init__(self, model_sampling): self.model_sampling = model_sampling class FakeModelPatcher: """ModelPatcher stand-in with the lifecycle semantics that matter. - clone() shares the parent's object_patches_backup dict object (via get_clone_model_override, model_patcher.py:428 + clone 462) and copies object_patches; - patch_model() applies object patches to the SHARED base model, backing up only keys not already backed up (model_patcher.py:1113-1120); - unpatch_model() restores the backup and CLEARS it (1165-1169); - detach(unpatch_all=False) skips the unpatch -> leak (called from model_management.py:962). """ def __init__(self, base_model, object_patches=None): self.model = base_model self.object_patches = dict(object_patches or {}) self.object_patches_backup = {} # replaced by SHARED dict on clone() def clone(self): child = FakeModelPatcher(self.model, self.object_patches) child.object_patches_backup = self.object_patches_backup # SHARED return child def patch_model(self): for key, patch in self.object_patches.items(): old = getattr(self.model, key) if key not in self.object_patches_backup: self.object_patches_backup[key] = old setattr(self.model, key, patch) def unpatch_model(self): for key in list(self.object_patches_backup.keys()): setattr(self.model, key, self.object_patches_backup[key]) self.object_patches_backup.clear() def partially_load(self): # partially_load = unpatch(own backup) + patch(own patches) # (model_patcher.py:1256+). A patcher with no patch for a key and an # empty backup restores nothing for it. self.unpatch_model() self.patch_model() def detach(self, unpatch_all=True): if unpatch_all: self.unpatch_model() def get_model_object(self, name): if name in self.object_patches: return self.object_patches[name] if name in self.object_patches_backup: return self.object_patches_backup[name] return getattr(self.model, name) class FakeLoadedModels: """load_models_gpu semantics (model_management.py:913-1014). Identity-keyed entries; stale same-base entries are detached with ``unpatch_all=False`` (LEAK); the requested patcher always runs model_load -> partially_load (re-applies ITS OWN patches). """ def __init__(self): self.entries = [] def load_models_gpu(self, patcher): stale = [e for e in self.entries if e is not patcher and e.model is patcher.model] for e in stale: e.detach(unpatch_all=False) # LEAK by design self.entries.remove(e) if patcher not in self.entries: self.entries.append(patcher) patcher.partially_load() def _install_fake_calculate_sigmas(monkeypatch): """comfy.samplers.calculate_sigmas stand-in: indexes the sampling's own table so the schedule genuinely depends on WHICH object is resolved.""" def calculate_sigmas(ms, scheduler, steps): assert scheduler == "simple" table = getattr(ms, "sigma_table") idx = torch.linspace(0, len(table) - 1, steps + 1).round().long() return table[idx] fake = types.ModuleType("comfy.samplers") fake.calculate_sigmas = calculate_sigmas monkeypatch.setitem(sys.modules, "comfy.samplers", fake) comfy_mod = sys.modules.get("comfy") if comfy_mod is not None: monkeypatch.setattr( comfy_mod, "samplers", fake, raising=False) # --------------------------------------------------------------------------- # Sampling stand-ins with distinct, measurable schedules # --------------------------------------------------------------------------- def _make_sampling(name, shift): class _Sampling: def __init__(self, shift): self.shift = shift t = torch.linspace(1.0, 0.0, 101) self.sigma_table = shift * t / (1.0 + (shift - 1.0) * t) def timestep(self, sigma): return sigma * 1000.0 _Sampling.__name__ = name return _Sampling(shift) def _checkpoint_patcher(shift=2.0): """A fresh Load-Checkpoint output: patcher over a fresh BaseModel.""" return FakeModelPatcher( FakeBaseModel(_make_sampling("ModelSamplingContinuousFlow", shift))) def _install_dype_patch(patcher): """Stand-in for apply_dype_to_model's schedule patch (patch_utils.py:166-171): a NEW clone + an object patch whose shift differs from the original.""" patcher.object_patches["model_sampling"] = _make_sampling( "DypeModelSamplingFlux", 3.5) def _stage_entry_sigma(model, steps=30, tau=0.6): """What HiFlow's console line reports: the guided stage's entry sigma.""" base_sigmas = hfn._base_sigmas(model, steps) stage = build_stage_sigmas(base_sigmas, tau, 16) return float(stage[0]) # --------------------------------------------------------------------------- # The reported repro # --------------------------------------------------------------------------- @pytest.mark.unit class TestReportedRepro: """Run1 (plain checkpoint) -> Run2 (DyPE clone, different schedule) -> Run3 (restore run-1 graph, caches NOT cleared). Contract delivered by v2.16.0: - a patcher with its OWN patch (DyPE/SEGA in the graph) resolves that patch — always, regardless of live-attr history; - an unpatched patcher resolves the backup original when the lineage holds one (the ComfyUI shared-backup heal path, model_patcher.py:428); - ONLY the un-healable residue (no patch, empty backup, live leak — a live attr inherited from outside the patcher's lineage) falls back to the live attr, exactly what a stock KSampler on the same input sees; the S6 stale-leak warning flags that case instead of silent drift. """ def test_p0_reload_heals_via_shared_backup(self, monkeypatch): """Run3 re-requests the SAME checkpoint patcher object (the normal Load-Checkpoint cache hit): the shared backup restores the original and the schedule matches run 1.""" _install_fake_calculate_sigmas(monkeypatch) loaded = FakeLoadedModels() p1 = _checkpoint_patcher() loaded.load_models_gpu(p1) s1 = _stage_entry_sigma(p1) p2 = p1.clone() _install_dype_patch(p2) loaded.load_models_gpu(p2) s2 = _stage_entry_sigma(p2) # Run 3: the cached checkpoint patcher object again (node cache hit), # DyPE node removed from the graph. loaded.load_models_gpu(p1) s3 = _stage_entry_sigma(p1) assert s2 != s1, "harness sanity: the patched schedule must differ" assert s3 == pytest.approx(s1), ( "re-loading the checkpoint patcher must restore its original " "schedule via the shared backup (unpatch_model on partially_load)" ) def test_restored_dype_params_resolve_own_schedule(self, monkeypatch): """Run3 re-executes the patch node with restored params (a NEW clone): its own patch decides the schedule — not the live attr left by run 2's clone, which had DIFFERENT shift parameters.""" _install_fake_calculate_sigmas(monkeypatch) loaded = FakeLoadedModels() p1 = _checkpoint_patcher() loaded.load_models_gpu(p1) # Run 2: DyPE clone at shift 3.5. p2 = p1.clone() _install_dype_patch(p2) loaded.load_models_gpu(p2) s2 = _stage_entry_sigma(p2) # Run 3: restored params -> fresh clone whose patch is re-derived for # the restored parameters. The live attr is POISONED with a foreign # schedule (as if another run's patch had leaked): resolution must # ignore it. p3 = p1.clone() _install_dype_patch(p3) p3.model.model_sampling = _make_sampling("Leaked", 0.5) loaded.load_models_gpu(p3) s3 = _stage_entry_sigma(p3) assert s3 == pytest.approx(s2), ( "a patcher with its own schedule patch must resolve that patch " "regardless of the live attribute's history" ) def test_unhealable_live_leak_is_ksampler_consistent(self, monkeypatch): """The documented residual: a patcher with no patch and an empty backup over a leaked base resolves the live attr — the SAME schedule a stock KSampler on that input would use (consistent, not silent divergence between our node and ComfyUI's own).""" _install_fake_calculate_sigmas(monkeypatch) loaded = FakeLoadedModels() p1 = _checkpoint_patcher() loaded.load_models_gpu(p1) p2 = p1.clone() _install_dype_patch(p2) loaded.load_models_gpu(p2) # A patcher outside the leak's backup lineage over the same base. p3 = FakeModelPatcher(p1.model) loaded.load_models_gpu(p3) assert effective_model_sampling(p3) is p3.model.model_sampling @pytest.mark.unit class TestResolutionContract: def test_same_patcher_reuse_keeps_schedule_stable(self, monkeypatch): _install_fake_calculate_sigmas(monkeypatch) loaded = FakeLoadedModels() p1 = _checkpoint_patcher() loaded.load_models_gpu(p1) s1 = _stage_entry_sigma(p1) loaded.load_models_gpu(p1) # cache HIT: same patcher object assert _stage_entry_sigma(p1) == pytest.approx(s1) def test_dype_in_chain_is_always_honored(self, monkeypatch): """A DyPE clone in THIS graph always resolves to its own patched schedule — before and after unrelated model loads.""" _install_fake_calculate_sigmas(monkeypatch) loaded = FakeLoadedModels() dype_patcher = _checkpoint_patcher() _install_dype_patch(dype_patcher) loaded.load_models_gpu(dype_patcher) s_direct = _stage_entry_sigma(dype_patcher) # Unrelated load in between (another model's patcher). other = FakeModelPatcher( FakeBaseModel(_make_sampling("Other", 1.5))) loaded.load_models_gpu(other) assert _stage_entry_sigma(dype_patcher) == pytest.approx(s_direct) assert s_direct > 0.0 def test_schedule_matches_resolved_object_table(self, monkeypatch): """The derived base schedule indexes the RESOLVED object's table — the schedule a stock KSampler on the same input would compute.""" _install_fake_calculate_sigmas(monkeypatch) loaded = FakeLoadedModels() p1 = _checkpoint_patcher(shift=2.0) loaded.load_models_gpu(p1) p2 = p1.clone() _install_dype_patch(p2) loaded.load_models_gpu(p2) p3 = FakeModelPatcher(p1.model) sigmas = hfn._base_sigmas(p3, steps=30) resolved = effective_model_sampling(p3) expected = resolved.sigma_table[ torch.linspace(0, 100, 31).round().long()] assert torch.allclose(sigmas, expected)