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

Remove latent.

Browse Source
This commit is contained in:
ZHOU He
2025-11-19 13:31:22 +08:00
parent 3c8e8d3753
commit 129c370356
1 changed files with 27 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
node.py
+27 -29
View File
@@ -549,33 +549,30 @@ class ExtractAttn:
return {
"required": {
"model": ("MODEL",),
"latent": ("LATENT",),
"kv_start_step": ("INT",{"default":1,"min":0}),
"kv_end_step": ("INT",{"default":9,"min":0}),
"x_start_step": ("INT",{"default":1,"min":0}),
"x_end_step": ("INT",{"default":9,"min":0}),
"mask_start_step": ("INT",{"default":8,"min":0}),
"mask_end_step": ("INT",{"default":9,"min":0}),
"kv_start_step": ("INT",{"default":1,"min":1}),
"kv_end_step": ("INT",{"default":9,"min":1}),
"x_start_step": ("INT",{"default":1,"min":1}),
"x_end_step": ("INT",{"default":9,"min":1}),
"mask_start_step": ("INT",{"default":8,"min":1}),
"mask_end_step": ("INT",{"default":9,"min":1}),
},
"optional":{
"fg_condition_mask": ("MASK",),
}
}
@classmethod
def go(cls, model, latent, kv_start_step, kv_end_step, x_start_step, x_end_step, mask_start_step, mask_end_step, fg_condition_mask = None):
def go(cls, model, kv_start_step, kv_end_step, x_start_step, x_end_step, mask_start_step, mask_end_step, fg_condition_mask = None):
attn_data = [{}]
index = len(attn_data) - 1
def prepare_callback(model, max_steps, x0_output_dict=None):
nonlocal attn_data, fg_condition_mask, latent
nonlocal attn_data, fg_condition_mask
attn_data[index] = {}
sampled_step = -1
sampling_step = -1
layer = 0
w = latent["samples"].shape[-1] // model.model.diffusion_model.patch_size
h = latent["samples"].shape[-2] // model.model.diffusion_model.patch_size
w = 0
h = 0
#text_len = fg_condition_mask.shape[-1]
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
attn_dtype = torch.bfloat16 if model_management.should_use_bf16(device=device) else torch.float16
@@ -606,7 +603,7 @@ class ExtractAttn:
attn_data[index][sampling_step] = {"layers": {}}
# 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:
attn_data[index][sampling_step]["layers"][layer] = {"k": k.cpu(), "v": v.cpu()}
attn_data[index][sampling_step]["layers"][layer] = {"k": k[:, :, text_len:,:].cpu(), "v": v[:, :, text_len:,:].cpu()}
# Apply rotary positional encoding for Flux models
if pe is not None:
q = orig_apply_rope1(q, pe)
@@ -635,12 +632,12 @@ class ExtractAttn:
origin_callback = orig_prepare_callback(model, max_steps, x0_output_dict)
def callback_change_sampled_step(step, x0, x, total_steps):
nonlocal max_steps
nonlocal sampling_step, sampled_step, layer, fg_bg_mask, model
nonlocal sampling_step, sampled_step, layer, fg_bg_mask, model, w, h
B,L,H,W = x0.shape
w = W // model.model.diffusion_model.patch_size
h = H // model.model.diffusion_model.patch_size
sampled_step = step
if fg_condition_mask is not None and sampling_step >= mask_start_step and sampling_step < mask_end_step:
B,L,H,W = x0.shape
w = W // model.model.diffusion_model.patch_size
h = H // model.model.diffusion_model.patch_size
bitmask = fg_bg_mask[:,0] >= fg_bg_mask[:,1]
bitmask = rearrange(bitmask, "b (h w) -> b h w", h=h, w=w)
fg_bg_mask = None
@@ -823,25 +820,23 @@ class GenConsistent:
"model": ("MODEL",),
"subject": ("ATTN",),
"cross_sim": ("CROSSSIM",),
"kv_start_step": ("INT",{"default":1,"min":0}),
"kv_end_step": ("INT",{"default":5,"min":0}),
"x_start_step": ("INT",{"default":1,"min":0}),
"x_end_step": ("INT",{"default":5,"min":0}),
"kv_start_step": ("INT",{"default":1,"min":1}),
"kv_end_step": ("INT",{"default":5,"min":1}),
"x_start_step": ("INT",{"default":1,"min":1}),
"x_end_step": ("INT",{"default":5,"min":1}),
"ada_tome_alpha": ("FLOAT",{"default":0.0,"min":0,"max":1,"step":0.1}),
"latent":("LATENT",),
},
}
@classmethod
def go(cls, model, subject, cross_sim, kv_start_step, kv_end_step, x_start_step, x_end_step, ada_tome_alpha, latent):
def go(cls, model, subject, cross_sim, kv_start_step, kv_end_step, x_start_step, x_end_step, ada_tome_alpha):
def prepare_callback(model, max_steps, x0_output_dict=None):
nonlocal subject, latent
nonlocal subject
sampled_step = -1
sampling_step = -1
layer = 0
w = latent["samples"].shape[-1] // model.model.diffusion_model.patch_size
h = latent["samples"].shape[-2] // model.model.diffusion_model.patch_size
w = 0
h = 0
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
attn_dtype = torch.bfloat16 if model_management.should_use_bf16(device=device) else torch.float16
@@ -989,7 +984,10 @@ class GenConsistent:
origin_callback = orig_prepare_callback(model, max_steps, x0_output_dict)
def callback_change_sampled_step(step, x0, x, total_steps):
nonlocal max_steps
nonlocal sampling_step, sampled_step, layer, fg_bg_mask, model
nonlocal sampling_step, sampled_step, layer, fg_bg_mask, model, w, h
B,L,H,W = x0.shape
w = W // model.model.diffusion_model.patch_size
h = H // model.model.diffusion_model.patch_size
sampled_step = step
origin_callback(step, x0, x, total_steps)
return callback_change_sampled_step