From 5f8685a91a53ae4f2b27b4f2132a7b7e4f5f6788 Mon Sep 17 00:00:00 2001
From: BunnyAI <119530153+WingeD123@users.noreply.github.com>
Date: Wed, 8 Apr 2026 01:14:58 +0800
Subject: [PATCH] Fix MoCha with context window

---
 nodes_sampler.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nodes_sampler.py b/nodes_sampler.py
index dfe17ade..59c6e490 100644
--- a/nodes_sampler.py
+++ b/nodes_sampler.py
@@ -1258,7 +1258,7 @@ def predict_with_cfg(z, cfg_scale, positive_embeds, negative_embeds, timestep, i
                         latent_end = latent_frames
                         mask_end = latent_end + 1
                         partial_latents = mocha_embeds[:, context_window]  # windowed latents
-                        mask_frame = mocha_embeds[:, latent_end:mask_end]  # single mask frame
+                        mask_frame = mocha_embeds[:, -mocha_num_refs-1:-mocha_num_refs]  # single mask frame
                         ref_frames = mocha_embeds[:, -mocha_num_refs:]     # reference frames
 
                         partial_mocha_embeds = torch.cat([partial_latents, mask_frame, ref_frames], dim=1)