From e65c4d0fb69251e692390574c49873297dc4bae5 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Wed, 16 Sep 2026 13:46:39 -0600
Subject: [PATCH] Fix issue with Krea2 patch size when training with shift

---
 extensions_built_in/diffusion_models/krea2/src/mmdit.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/extensions_built_in/diffusion_models/krea2/src/mmdit.py b/extensions_built_in/diffusion_models/krea2/src/mmdit.py
index f4036be..2d66b48 100644
--- a/extensions_built_in/diffusion_models/krea2/src/mmdit.py
+++ b/extensions_built_in/diffusion_models/krea2/src/mmdit.py
@@ -121,6 +121,11 @@ class SingleMMDiTConfig:
     txtheads: int = 20
     txtkvheads: int = 20
 
+    @property
+    def patch_size(self) -> int:
+        # trainer probes config.patch_size for the dynamic-shift token count
+        return self.patch
+
 
 class SimpleModulation(torch.nn.Module):
     def __init__(self, dim: int):