From 907d1b0b0f649bd569bbf23430e107ab37832d73 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sat, 2 Aug 2025 00:34:10 +0300
Subject: [PATCH] 5B + context windows

---
 nodes.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/nodes.py b/nodes.py
index 0724300..30f6b74 100644
--- a/nodes.py
+++ b/nodes.py
@@ -2654,6 +2654,8 @@ class WanVideoSampler:
                             partial_audio_proj = audio_proj[:, c]
 
                         partial_latent_model_input = latent_model_input[:, c]
+                        if latents_to_insert is not None and c[0] != 0:
+                            partial_latent_model_input[:, 0:num_latents_to_insert] = latents_to_insert
 
                         partial_unianim_data = None
                         if unianim_data is not None:
@@ -2673,8 +2675,10 @@ class WanVideoSampler:
 
                         if len(timestep.shape) != 1:
                             partial_timestep = timestep[:, c]
+                            partial_timestep[:, :num_latents_to_insert] = 0
                         else:
                             partial_timestep = timestep
+                        #print("Partial timestep:", partial_timestep)
 
                         noise_pred_context, new_teacache = predict_with_cfg(
                             partial_latent_model_input,