From 68985b96cca8a4b340f2eb4d365abd3239ad3e68 Mon Sep 17 00:00:00 2001 From: Kevin Gilper Date: Wed, 28 Oct 2026 16:26:22 -0400 Subject: [PATCH] Run native enhancer before the edit wrapper --- docs/krea-v12-native-studies.md | 10 +++++ kg_krea_v12/diagnostics.py | 3 ++ kg_krea_v12/enhancer.py | 11 ++++-- tests/test_krea_v12_diagnostics.py | 4 +- tests/test_krea_v12_enhancer.py | 62 ++++++++++++++++++++++++++++-- 5 files changed, 83 insertions(+), 7 deletions(-) diff --git a/docs/krea-v12-native-studies.md b/docs/krea-v12-native-studies.md index ba6a2f9..e64ae88 100644 --- a/docs/krea-v12-native-studies.md +++ b/docs/krea-v12-native-studies.md @@ -379,3 +379,13 @@ foreign state is preserved. There is no execution endpoint or persistent hook. This diagnostic measures numerical transfer. Synthetic source tests do not establish native neural behavior or useful lighting, and a nonzero measured prediction difference does not by itself establish a calibrated semantic axis. + +The enhancer registers in the cloned MODEL's transformer-options wrapper map, +before the native edit wrapper. Comfy also has a separate ModelPatcher wrapper +registry that is merged later during sampling preparation. Registering the +enhancer only there puts it after native edit, whose forward bypasses its inner +executor. Options-level registration preserves the intended enhancer-before-edit +execution order. Duplicate enhancer/edit attachments are rejected in either +registry. The bounded diagnostic additionally rejects unmerged ModelPatcher +wrappers/callbacks, so a later native merge cannot silently add unreviewed calls. +EnhancerOFF/strength0 retains original MODEL identity and performs no registration. diff --git a/kg_krea_v12/diagnostics.py b/kg_krea_v12/diagnostics.py index 7bdce59..28ba930 100644 --- a/kg_krea_v12/diagnostics.py +++ b/kg_krea_v12/diagnostics.py @@ -209,6 +209,9 @@ class BoundedDiagnostic: if len(shape) != 5 or shape[:3] != (1, 16, 1) or not all(1 <= size <= 128 for size in shape[-2:]): raise RuntimeError('V12 diagnostic target exceeds its single-image spatial bound') options = guider.model_patcher.model_options + if (getattr(guider.model_patcher, 'wrappers', {}) + or getattr(guider.model_patcher, 'callbacks', {})): + raise RuntimeError('V12 diagnostic refuses unmerged MODEL wrappers/callbacks') if set(options) - {'transformer_options'}: raise RuntimeError('V12 diagnostic refuses extra MODEL options/callbacks') wrappers = options.get('transformer_options', {}).get('wrappers', {}) diff --git a/kg_krea_v12/enhancer.py b/kg_krea_v12/enhancer.py index 86b4379..8927e3a 100644 --- a/kg_krea_v12/enhancer.py +++ b/kg_krea_v12/enhancer.py @@ -133,9 +133,11 @@ def apply_enhancer(model, enabled=False, strength=1.0): _check_architecture(dm, compatible) original_options_signature = _signature(model.model_options) original_wrappers = _wrappers(model.model_options) - if any(key in original_wrappers for key in (WRAPPER_KEY, *_OTHER_KEYS)): + pending_wrappers = getattr(model, 'wrappers', {}).get('diffusion_model', {}) + if any(key in registry for registry in (original_wrappers, pending_wrappers) + for key in (WRAPPER_KEY, *_OTHER_KEYS)): raise RuntimeError('V12 enhancer refuses an already enhanced MODEL') - if 'krea2_edit' in original_wrappers: + if any('krea2_edit' in registry for registry in (original_wrappers, pending_wrappers)): raise RuntimeError('V12 enhancer must precede the native edit patch, which bypasses inner wrappers') config = {'enabled': True, 'strength': float(strength), 'source_blob': enhancer_blob} @@ -195,7 +197,10 @@ def apply_enhancer(model, enabled=False, strength=1.0): raise RuntimeError('V12 enhancer requires an independent native MODEL clone') to = patched.model_options.setdefault('transformer_options', {}) to[WRAPPER_KEY] = config.copy() - patched.add_wrapper_with_key('diffusion_model', WRAPPER_KEY, wrapper) + # Native edit registers in this options map and bypasses its inner executor. + # ModelPatcher.wrappers is merged later, after options-level registrations; + # putting the enhancer there would let native edit bypass it entirely. + to.setdefault('wrappers', {}).setdefault('diffusion_model', {}).setdefault(WRAPPER_KEY, []).append(wrapper) if _signature(model.model_options) != original_options_signature: raise RuntimeError('V12 enhancer clone changed original MODEL wrapper ownership') return patched, ('Krea native enhancer V12: pinned text-fusion helper attached; ' diff --git a/tests/test_krea_v12_diagnostics.py b/tests/test_krea_v12_diagnostics.py index 879035e..93403c2 100644 --- a/tests/test_krea_v12_diagnostics.py +++ b/tests/test_krea_v12_diagnostics.py @@ -285,7 +285,7 @@ class DiagnosticTests(unittest.TestCase): self.assertEqual(s.state,'failed') def test_wrong_step_count_sampler_cfg_callbacks_and_extra_wrappers_reject(self): - for change in ('steps','sampler','cfg','callback','wrapper','transformer','target'): + for change in ('steps','sampler','cfg','callback','wrapper','transformer','target','pendingWrapper','pendingCallback'): bundle=self.prepared();s,study,g,host,sampler=bundle schedule=[8,7,6,5,4,3,2,1,0] if change=='steps':schedule=[2,1,0] @@ -294,6 +294,8 @@ class DiagnosticTests(unittest.TestCase): elif change=='callback':g.model_patcher.model_options['sampler_post_cfg_function']=[lambda:None] elif change=='wrapper':g.model_patcher.model_options['transformer_options']['wrappers']['predict_noise']={} elif change=='transformer':g.model_patcher.model_options['transformer_options']['patches']={} + elif change=='pendingWrapper':g.model_patcher.wrappers={'diffusion_model':{'foreign':[lambda *a:None]}} + elif change=='pendingCallback':g.model_patcher.callbacks={'on_load_after':{'foreign':[lambda *a:None]}} else:study.context.native['target_latent']['samples']=Array(np.zeros((1,16,1,129,1))) with self.subTest(change=change),self.assertRaises(RuntimeError): s.validate(study,g,sampler,schedule,host[1]) diff --git a/tests/test_krea_v12_enhancer.py b/tests/test_krea_v12_enhancer.py index 9f73f3a..2bdfb2f 100644 --- a/tests/test_krea_v12_enhancer.py +++ b/tests/test_krea_v12_enhancer.py @@ -28,15 +28,20 @@ class Fusion: class Model: - def __init__(self, dm=None, options=None): + def __init__(self, dm=None, options=None, wrappers=None): self.model = NS(diffusion_model=dm or NS(txtfusion=Fusion(), txtlayers=12, txtdim=2560)) self.model_options = {} if options is None else options + self.wrappers = {} if wrappers is None else wrappers + self.callbacks = {} def clone(self): - return Model(self.model.diffusion_model, copy.deepcopy(self.model_options)) + cloned = Model(self.model.diffusion_model, copy.deepcopy(self.model_options), copy.deepcopy(self.wrappers)) + cloned.callbacks = copy.deepcopy(self.callbacks) + return cloned def add_wrapper_with_key(self, kind, key, wrapper): - self.model_options.setdefault('transformer_options', {}).setdefault('wrappers', {}).setdefault(kind, {}).setdefault(key, []).append(wrapper) + # Native ModelPatcher API stores these separately until sampling prep. + self.wrappers.setdefault(kind, {}).setdefault(key, []).append(wrapper) class Executor: @@ -79,6 +84,57 @@ class EnhancerV12Tests(unittest.TestCase): self.assertIs(result, self.model) self.assertEqual(self.model.model_options, {}) + def test_options_registration_precedes_native_edit_after_pending_merge(self): + foreign = lambda executor, *args, **kwargs: executor(*args, **kwargs) + self.model.add_wrapper_with_key('diffusion_model', 'foreign', foreign) + self.model.callbacks={'on_load_after':{'foreign':[foreign]}} + enhanced, options, wrapper = self.attached() + self.assertEqual(list(enhanced.wrappers['diffusion_model']), ['foreign']) + self.assertEqual(list(options['wrappers']['diffusion_model']), [enhancer.WRAPPER_KEY]) + self.assertEqual(self.model.model_options, {}) + self.assertEqual(enhanced.callbacks,self.model.callbacks) + self.assertIsNot(enhanced.callbacks,self.model.callbacks) + observed=[] + dm=self.model.model.diffusion_model + def native_edit(executor, **kwargs): + observed.append('native-edit-bypasses-inner') + return dm.txtfusion.forward(NS(shape=(1,2,12,2560)), + transformer_options=kwargs['transformer_options']) + options['wrappers']['diffusion_model']['krea2_edit']=[native_edit] + # Ordinary native preparation preserves options keys, then appends the + # separately registered ModelPatcher wrappers. No source/body emulation. + for kind, entries in enhanced.wrappers.items(): + for key, functions in entries.items(): + options['wrappers'].setdefault(kind,{}).setdefault(key,[]).extend(functions) + chain=[fn for group in options['wrappers']['diffusion_model'].values() for fn in group] + class Chain: + class_obj=dm + def __init__(self,index=0):self.index=index + def __call__(self,*args,**kwargs): + return chain[self.index](Chain(self.index+1),*args,**kwargs) + result=Chain()(transformer_options=options) + self.assertEqual(result.value,'delegated-helper') + self.assertEqual(observed,['native-edit-bypasses-inner']) + self.assertEqual(len(self.calls),1) + self.assertEqual(list(options['wrappers']['diffusion_model']), + [enhancer.WRAPPER_KEY,'krea2_edit','foreign']) + self.assertEqual(list(self.model.wrappers['diffusion_model']),['foreign']) + + def test_pending_duplicate_enhancer_and_edit_rejected_before_clone(self): + for key in (enhancer.WRAPPER_KEY,*enhancer._OTHER_KEYS,'krea2_edit'): + with self.subTest(key=key): + model=Model() + model.add_wrapper_with_key('diffusion_model',key,lambda *a:None) + with patch.object(model,'clone',side_effect=AssertionError('must reject first')): + with self.assertRaises(RuntimeError):enhancer.apply_enhancer(model,True,1) + + def test_options_duplicate_enhancer_and_edit_rejected_before_clone(self): + for key in (enhancer.WRAPPER_KEY,*enhancer._OTHER_KEYS,'krea2_edit'): + with self.subTest(key=key): + model=Model(options={'transformer_options':{'wrappers':{'diffusion_model':{key:[lambda *a:None]}}}}) + with patch.object(model,'clone',side_effect=AssertionError('must reject first')): + with self.assertRaises(RuntimeError):enhancer.apply_enhancer(model,True,1) + def test_active_invokes_pinned_helper_and_preserves_model_options_and_forward(self): foreign = lambda *args: None self.model.add_wrapper_with_key('diffusion_model', 'foreign', foreign)