Added reference token attention isolation (kv_cache) for Krea2 edit training. Same training cost with significant inference speed up. 2x inference speedup.

extensions_built_in/diffusion_models/krea2/krea2.py

@@ -208,6 +208,16 @@ class Krea2Model(BaseModel):
         self.has_multiple_control_images = self.is_edit
         # Reference images keep their own aspect/size (not resized to the target).
         self.use_raw_control_images = self.is_edit
+        # model_kwargs.kv_cache = true: train with an asymmetric attention mask
+        # where the clean reference tokens attend only to each other (never to
+        # text / noisy tokens). Their hidden states then depend only on the
+        # refs + t=0 modulation, so at inference their per-layer K/V can be
+        # computed once and reused across all denoising steps
+        # (OminiControl2-style conditioning feature reuse). Off by default:
+        # the base model was trained fully bidirectional, so a LoRA must be
+        # trained with kv_cache enabled for kv-cached inference (the ComfyUI
+        # node / hub pipeline kv_cache toggles) to work properly.
+        self.kv_cache = bool(self.model_config.model_kwargs.get("kv_cache", False))
 
     @staticmethod
     def get_train_scheduler():
@@ -642,6 +652,7 @@ class Krea2Model(BaseModel):
             context,
             text_mask,
             ref_latents=ref_latents,
+            isolate_refs=self.kv_cache,
         )
         return pred