github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Jaret Burkett
2026-07-09 12:00:55 -06:00
parent 6233efe1bb
commit a4bbe167ce
7 changed files with 177 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
ui/src/app/jobs/new/SimpleJob.tsx
View File

@@ -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"

9
ui/src/app/jobs/new/options.ts
View File

@@ -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) => {