From 296b7c7b6d7111c5ee07a09576e7e1e6255aeeff Mon Sep 17 00:00:00 2001
From: Mihail Karaev <karaevmihail@gmail.com>
Date: Wed, 17 Dec 2025 11:40:14 +0000
Subject: [PATCH] Small fixes

---
 comfy/ldm/kandinsky5/model.py       | 7 +++++--
 comfy/ldm/kandinsky5/utils_nabla.py | 2 +-
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/comfy/ldm/kandinsky5/model.py b/comfy/ldm/kandinsky5/model.py
index f4e02af70..24a06da0a 100644
--- a/comfy/ldm/kandinsky5/model.py
+++ b/comfy/ldm/kandinsky5/model.py
@@ -391,19 +391,22 @@ class Kandinsky5(nn.Module):
         if T > NABLA_THR:
             assert self.patch_size[0] == 1
 
+            # pro video model uses lower P at higher resolutions
+            P = 0.7 if self.model_dim == 4096 and H * W >= 14080 else 0.9
+
             freqs = freqs.view(freqs.shape[0], *visual_shape[1:], *freqs.shape[2:])
             visual_embed, freqs = fractal_flatten(visual_embed, freqs, visual_shape[1:])
             pt, ph, pw = self.patch_size
             T, H, W = T // pt, H // ph, W // pw
 
-            wT, wW, wH = 11, 11, 3
+            wT, wW, wH = 11, 3, 3
             sta_mask = fast_sta_nabla(T, H // 8, W // 8, wT, wH, wW, device=x.device)
 
             sparse_params = dict(
                 sta_mask=sta_mask.unsqueeze_(0).unsqueeze_(0),
                 attention_type="nabla",
                 to_fractal=True,
-                P=0.8,
+                P=P,
                 wT=wT, wW=wW, wH=wH,
                 add_sta=True,
                 visual_shape=(T, H, W),
diff --git a/comfy/ldm/kandinsky5/utils_nabla.py b/comfy/ldm/kandinsky5/utils_nabla.py
index 5e2bc4076..a346736b2 100644
--- a/comfy/ldm/kandinsky5/utils_nabla.py
+++ b/comfy/ldm/kandinsky5/utils_nabla.py
@@ -143,4 +143,4 @@ def nabla(query, key, value, sparse_params=None):
         .contiguous()
     )
     out = out.flatten(-2, -1)
-    return out
\ No newline at end of file
+    return out