Run native enhancer before the edit wrapper

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.

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', {})

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; '

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])

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)