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

ui/src/app/jobs/new/SimpleJob.tsx

@@ -330,6 +330,14 @@ export default function SimpleJob({
                 />
               </FormGroup>
             )}
+            {modelArch?.additionalSections?.includes('model.model_kwargs.kv_cache') && (
+              <Checkbox
+                label="KV Cache"
+                docKey="model.model_kwargs.kv_cache"
+                checked={jobConfig.config.process[0].model.model_kwargs.kv_cache || false}
+                onChange={value => setJobConfig(value, 'config.process[0].model.model_kwargs.kv_cache')}
+              />
+            )}
             {modelArch?.additionalSections?.includes('model.qie.match_target_res') && (
               <Checkbox
                 label="Match Target Res"

ui/src/app/jobs/new/options.ts

@@ -33,6 +33,7 @@ type AdditionalSections =
   | 'model.qie.match_target_res'
   | 'model.assistant_lora_path'
   | 'model.unconditional_lora_path'
+  | 'model.model_kwargs.kv_cache'
   | 'ideogram_4_prompt';
 
 type ModelGroup = 'image' | 'instruction' | 'video' | 'experimental' | 'audio';
@@ -1104,7 +1105,8 @@ export const modelArchs: ModelArch[] = [
       'config.process[0].model.model_kwargs': [
         {
           edit: true,
-          match_target_res: false,
+          match_target_res: true,
+          kv_cache: true,
         },
         {},
       ],
@@ -1142,7 +1144,8 @@ export const modelArchs: ModelArch[] = [
       'config.process[0].model.model_kwargs': [
         {
           edit: true,
-          match_target_res: false,
+          match_target_res: true,
+          kv_cache: true,
         },
         {},
       ],
@@ -1157,6 +1160,7 @@ export const modelArchs: ModelArch[] = [
       'model.layer_offloading',
       'model.assistant_lora_path',
       'model.qie.match_target_res',
+      'model.model_kwargs.kv_cache',
     ],
   },
   {
@@ -1209,6 +1213,7 @@ export const modelArchs: ModelArch[] = [
       'model.low_vram',
       'model.layer_offloading',
       'model.qie.match_target_res',
+      'model.model_kwargs.kv_cache',
     ],
   },
 ].sort((a, b) => {

ui/src/docs.tsx

@@ -340,6 +340,17 @@ const docs: { [key: string]: ConfigDoc } = {
       </>
     ),
   },
+  'model.model_kwargs.kv_cache': {
+    title: 'KV Cache',
+    description: (
+      <>
+        This will enable KV Cache for control images in a model that supports it. LoRAs trained with this on
+        need to also be inferenced with it, and vice versa. This does not speed up or slow down training, but on inference,
+        the control images only need to be processed once for the entire generation, vs being processed for every step.
+        Which leads to a significant speedup on inference.
+      </>
+    ),
+  },
 };
 
 export const getDoc = (key: string | null | undefined): ConfigDoc | null => {