feat: Context windows - add causal_window_fix to improve blending of context windows (CORE-100)#13563
feat: Context windows - add causal_window_fix to improve blending of context windows (CORE-100)#13563drozbay wants to merge 3 commits intoComfy-Org:masterfrom
Conversation
drozbay
requested review from
Kosinkadink,
comfyanonymous and
guill
as code owners
📝 WalkthroughWalkthroughAdds a 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
comfy/context_windows.py (2)
66-75:⚠️ Potential issue | 🟠 Major
retain_index_listclobbers the anchor slot instead of the intended window position.When
causal_anchor_indexis set,indicesbecomes[anchor_idx] + index_list, so window-local position 0 now corresponds to the anchor frame. Theretain_index_listblock right below still indexes bothwindowandfullwith the same numbers:if retain_index_list: idx = tuple([slice(None)] * dim + [retain_index_list]) window[idx] = full[idx]So with
cond_retain_index_list = "0"(the documented "use the initial start image for each window" feature) andcausal_window_fix=True, this overwrites the anchor slot withfull[0]— and that slot is then discarded by thenarrow(...)strip inevaluate_context_windows. The first real frame of the window keeps its original sliced value, silently breakingcond_retain_index_listfor any non-initial window.Both features are enabled together by default on
ContextWindowsManualNode(causal_window_fix=True,cond_retain_index_list=""), so the typical workflow doesn't trigger this, but anyone enablingcond_retain_index_listwill hit it.🛠️ Suggested fix — shift the retain destinations to skip the anchor
indices = self.index_list anchor_idx = getattr(self, 'causal_anchor_index', None) - if anchor_idx is not None and anchor_idx >= 0: + anchor_offset = 0 + if anchor_idx is not None and anchor_idx >= 0: indices = [anchor_idx] + list(indices) + anchor_offset = 1 idx = tuple([slice(None)] * dim + [indices]) window = full[idx] if retain_index_list: - idx = tuple([slice(None)] * dim + [retain_index_list]) - window[idx] = full[idx] + src_idx = tuple([slice(None)] * dim + [retain_index_list]) + dst_positions = [i + anchor_offset for i in retain_index_list] + dst_idx = tuple([slice(None)] * dim + [dst_positions]) + window[dst_idx] = full[src_idx] return window.to(device)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@comfy/context_windows.py` around lines 66 - 75, The retain block currently writes full[retain_index_list] back into window using the same indices even when you prepended the causal_anchor_index into indices, so the anchor at local position 0 gets clobbered; fix by keeping the source indices for full as retain_index_list but offsetting the destination indices into window when causal_anchor_index is present (e.g. compute dest_retain = retain_index_list if anchor_idx is None else [i+1 for i in retain_index_list]) and use that for the window assignment (window[dest_retain] = full[retain_index_list]); update the code around indices, anchor_idx, retain_index_list, window and full accordingly so evaluate_context_windows no longer discards the intended retained frames.
100-141:⚠️ Potential issue | 🟠 MajorHandle anchor in
slice_condoffset/scale branches to match fast-path behavior.The fast path at lines 114-116 calls
window.get_tensor()which prependscausal_anchor_indexwhen present. However, thetemporal_offset > 0andtemporal_scale > 1branches (lines 119-141) rebuild indices fromwindow.index_listdirectly, bypassing this anchor logic.WAN22_Animate exercises these non-default branches with
temporal_offset=1(face_pixel_values, pose_latents) andtemporal_scale=4(face_pixel_values), making the inconsistency active code in models using the new Animate architecture.Apply the same anchor-aware index handling to offset/scale paths for consistency:
Suggested fix
# skip leading latent positions that have no corresponding conditioning (e.g. reference frames) if temporal_offset > 0: indices = [i - temporal_offset for i in window.index_list[temporal_offset:]] indices = [i for i in indices if 0 <= i] else: indices = list(window.index_list) + anchor_idx = getattr(window, 'causal_anchor_index', None) + if anchor_idx is not None and anchor_idx >= 0: + prepend = anchor_idx - temporal_offset + if prepend >= 0: + indices = [prepend] + indices + if not indices: return None🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@comfy/context_windows.py` around lines 100 - 141, slice_cond's temporal_offset>0 and temporal_scale>1 branches rebuild indices directly and miss the window.causal_anchor_index handling that window.get_tensor() applies; update those branches to detect anchor = getattr(window, "causal_anchor_index", None) and prepend the anchor-aware indices the same way the fast path does: for offset branch compute anchor_index = anchor - temporal_offset (and only include if 0 <= anchor_index) and insert it at the start of indices before any scaling; for temporal_scale>1, when expanding scaled indices also produce and prepend the corresponding scaled anchor positions (anchor_index * temporal_scale + k) if they are within cond_size so the offset/scale paths mirror window.get_tensor()'s anchor behavior.
🧹 Nitpick comments (1)
comfy/context_windows.py (1)
327-353: Anchor inject/strip flow looks correct.
anchor_idx = window.index_list[0] - 1with the0 <= anchor_idx < x_in.size(self.dim)guard correctly skips the first window (anchor would be -1) and avoids out-of-range on small tensors. Thenarrow(self.dim, 1, shape-1)matches the prepend-at-position-0 contract fromget_tensor, sosub_conds_outis back tocontext_lengthalongself.dimbeforecombine_context_window_resultsindexes it viawindow.index_list. Nice and tight.One small nit —
window.causal_anchor_indexis a dynamically-attached attribute used as an out-of-band channel toget_tensor. Threading it explicitly (e.g., a parameter or local override) would make the dependency obvious to future readers, but the windows are freshly constructed eachexecute()so there's no actual lifetime hazard today. Fine to leave as-is.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@comfy/context_windows.py` around lines 327 - 353, The anchor injection/strip flow is correct; no functional change is required — the causal anchor logic using anchor_idx, window.causal_anchor_index, get_tensor, and the subsequent narrow(self.dim, 1, ...) on sub_conds_out correctly preserves context_length for combine_context_window_results which indexes by window.index_list; you can leave the implementation as-is, or optionally make the dependency explicit by threading causal_anchor_index into get_tensor (e.g., add an explicit parameter to get_tensor and propagate it from the caller) if you prefer to avoid the dynamic attribute for future clarity.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@comfy/context_windows.py`:
- Around line 66-75: The retain block currently writes full[retain_index_list]
back into window using the same indices even when you prepended the
causal_anchor_index into indices, so the anchor at local position 0 gets
clobbered; fix by keeping the source indices for full as retain_index_list but
offsetting the destination indices into window when causal_anchor_index is
present (e.g. compute dest_retain = retain_index_list if anchor_idx is None else
[i+1 for i in retain_index_list]) and use that for the window assignment
(window[dest_retain] = full[retain_index_list]); update the code around indices,
anchor_idx, retain_index_list, window and full accordingly so
evaluate_context_windows no longer discards the intended retained frames.
- Around line 100-141: slice_cond's temporal_offset>0 and temporal_scale>1
branches rebuild indices directly and miss the window.causal_anchor_index
handling that window.get_tensor() applies; update those branches to detect
anchor = getattr(window, "causal_anchor_index", None) and prepend the
anchor-aware indices the same way the fast path does: for offset branch compute
anchor_index = anchor - temporal_offset (and only include if 0 <= anchor_index)
and insert it at the start of indices before any scaling; for temporal_scale>1,
when expanding scaled indices also produce and prepend the corresponding scaled
anchor positions (anchor_index * temporal_scale + k) if they are within
cond_size so the offset/scale paths mirror window.get_tensor()'s anchor
behavior.
---
Nitpick comments:
In `@comfy/context_windows.py`:
- Around line 327-353: The anchor injection/strip flow is correct; no functional
change is required — the causal anchor logic using anchor_idx,
window.causal_anchor_index, get_tensor, and the subsequent narrow(self.dim, 1,
...) on sub_conds_out correctly preserves context_length for
combine_context_window_results which indexes by window.index_list; you can leave
the implementation as-is, or optionally make the dependency explicit by
threading causal_anchor_index into get_tensor (e.g., add an explicit parameter
to get_tensor and propagate it from the caller) if you prefer to avoid the
dynamic attribute for future clarity.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 9f99b27f-ed22-48b2-bf5f-8ba671642ed4
📒 Files selected for processing (2)
comfy/context_windows.pycomfy_extras/nodes_context_windows.py
|
Responding to Coderabbit comments:
|
alexisrolland
changed the title
There was a problem hiding this comment.
🧹 Nitpick comments (2)
comfy/context_windows.py (2)
334-340: Looped/wrap-around windows skip the fix whenindex_list[0] == 0.For
closed_loop/UNIFORM_LOOPEDschedules a window can legitimately start at sub-frame0after wrap-shifting (e.g.,[0, 1, 2, ...]representing a continuation of an earlier wrap), in which caseanchor_idx = -1andanchor_appliedstaysFalse— so this particular window misses the causal fix even though it's not the "true" first window of the video. The non-looped case (STATIC_STANDARD, the one shown in the PR images) is unaffected, so this is just an edge case worth being aware of.If you want to cover it, you could wrap to
num_frames - 1for looped schedules:♻️ Optional: wrap anchor for looped windows
if self.causal_window_fix: - anchor_idx = window.index_list[0] - 1 - if 0 <= anchor_idx < x_in.size(self.dim): + first = window.index_list[0] + if first > 0: + anchor_idx = first - 1 + elif self.closed_loop: + anchor_idx = x_in.size(self.dim) - 1 + else: + anchor_idx = -1 + if 0 <= anchor_idx < x_in.size(self.dim): window.causal_anchor_index = anchor_idx anchor_applied = True🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@comfy/context_windows.py` around lines 334 - 340, The causal-anchor logic skips looped windows that start at index 0 because anchor_idx becomes -1; update the block in the method using self.causal_window_fix to detect looped schedules (e.g., check self.schedule == UNIFORM_LOOPED or closed_loop mode) and when anchor_idx < 0 set anchor_idx = x_in.size(self.dim) - 1 (or num_frames - 1) before assigning window.causal_anchor_index and setting anchor_applied = True so wrap-around windows receive the causal anchor fix; keep existing bounds check for non-looped schedules.
163-177: Handler default is safe for node-level integration, but document if this is a new default.The node module already explicitly passes
causal_window_fixto the handler (line 55), so the handler-level default doesn't affect the node's current behavior. However, if thecausal_window_fix=Truedefault is new to this PR, external custom nodes that instantiateIndexListContextHandlerdirectly without naming this parameter would silently inherit the new behavior. Consider adding a note in release documentation so custom node authors are aware of the default, or document the parameter in the handler's docstring.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@comfy/context_windows.py` around lines 163 - 177, The handler's __init__ now defaults causal_window_fix=True which may change behavior for external users who instantiate IndexListContextHandler directly; update the IndexListContextHandler class docstring (or the __init__ docstring) to explicitly document the causal_window_fix parameter, its default value (True), and its behavioral effect, and mention that node-level code still passes causal_window_fix explicitly (so current node behavior is unchanged) so external custom nodes are informed of the new default.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@comfy/context_windows.py`:
- Around line 334-340: The causal-anchor logic skips looped windows that start
at index 0 because anchor_idx becomes -1; update the block in the method using
self.causal_window_fix to detect looped schedules (e.g., check self.schedule ==
UNIFORM_LOOPED or closed_loop mode) and when anchor_idx < 0 set anchor_idx =
x_in.size(self.dim) - 1 (or num_frames - 1) before assigning
window.causal_anchor_index and setting anchor_applied = True so wrap-around
windows receive the causal anchor fix; keep existing bounds check for non-looped
schedules.
- Around line 163-177: The handler's __init__ now defaults
causal_window_fix=True which may change behavior for external users who
instantiate IndexListContextHandler directly; update the IndexListContextHandler
class docstring (or the __init__ docstring) to explicitly document the
causal_window_fix parameter, its default value (True), and its behavioral
effect, and mention that node-level code still passes causal_window_fix
explicitly (so current node behavior is unchanged) so external custom nodes are
informed of the new default.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: dfe46824-29e4-4f4b-af93-3d59bdf0de6a
📒 Files selected for processing (1)
comfy/context_windows.py
3 participants
Adds a causal_window_fix toggle to context windows that prepends a throwaway anchor frame at sub-position 0 of every non-initial window, absorbing the model's learned "scene start" content bias on a slot that gets stripped before output.
Without this fix, windows starting mid-video have their first real frame contaminated by the bias the model applies to position 0 (since the model was trained on clips that genuinely start at frame 0), causing visible seams or quality degradation at window boundaries. The anchor is the latent immediately preceding the window's first frame, prepended via an anchor-aware get_tensor and stripped from sub_conds_out after the model forward pass.
Having this option enabled does add one extra latent frame to the sampling for each window beyond the starting window.
The option is exposed as a Boolean input on ContextWindowsManualNode (default True). It will be enabled by default with model-specific nodes like WanContextWindowsManualNode.
Example workflow:
This is a very basic example workflow using Wan2.1 I2V with standard_static window mode and splitting the I2V conditionings over several windows, low overlap to emphasize the effect.
droz_contextwindows_causalwindowfix_test.json
Output:
ComfyUI_00217_.mp4
Single frame comparison:
