nodes/thatname-Comfyui_CharaConsist
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add expand kv, but have not impl multi-subject and rope yet

Browse Source
This commit is contained in:
ZHOU He
2025-11-16 01:42:27 +08:00
parent a79fe2d5fb
commit a7663bf094
1 changed files with 32 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
node.py
+32 -5
View File
@@ -843,14 +843,41 @@ class GenConsistent:
if sampling_step != sampled_step + 1:
layer = 0
sampling_step = sampled_step + 1
# Store k, v, and x in a dict with layer id as key
if sampling_step >= kv_start_step and sampling_step < kv_end_step:
subject[index][sampling_step]["layers"][layer]
# todo
# Apply rotary positional encoding for Flux models
if pe is not None:
q, k = flux_math.apply_rope(q, k, pe)
# Store k, v, and x in a dict with layer id as key
if sampling_step >= kv_start_step and sampling_step < kv_end_step:
# For all other steps, we may apply the key-value expansion algorithm
# Get cross-sim data for this layer
id_fg_inds = cross_sim.get("id_fg_inds", None)
curr_fg_inds = cross_sim.get("curr_fg_inds", None)
# Check if we have valid cross-sim data and subject data to work with
if (id_fg_inds is not None and curr_fg_inds is not None and
len(id_fg_inds) > 0 and len(curr_fg_inds) > 0 and
sampling_step in subject[index] and
layer in subject[index][sampling_step]["layers"] and
"k" in subject[index][sampling_step]["layers"][layer] and
"v" in subject[index][sampling_step]["layers"][layer]):
ori_saved_key = subject[index][sampling_step]["layers"][layer]["k"].to(k.device)
ori_saved_value = subject[index][sampling_step]["layers"][layer]["v"].to(v.device)
# Ensure we have valid indices within tensor bounds
valid_id_inds = id_fg_inds[id_fg_inds < ori_saved_key.shape[2]]
valid_curr_inds = curr_fg_inds[curr_fg_inds < k.shape[2]]
if len(valid_id_inds) > 0 and len(valid_curr_inds) > 0:
# Apply key-value expansion based on cross-sim data
# Extract the saved key-value pairs using valid indices
saved_key_fg = ori_saved_key[:, :, valid_id_inds]
saved_value_fg = ori_saved_value[:, :, valid_id_inds]
# Concatenate the expanded key-value pairs
k = torch.cat([k, saved_key_fg], dim=2)
v = torch.cat([v, saved_value_fg], dim=2)
x = F.scaled_dot_product_attention(
q, k, v, attn_mask=mask
@@ -891,4 +918,4 @@ class GenConsistent:
model.model_options = copy.copy(model.model_options)
model.model_options["prepare_callback"] = prepare_callback
hook_prepare_callback()
return (model, subject,)
return (model, subject,)