Merge 04b6add331 into 0e4b6e8695

extensions_built_in/diffusion_models/krea2/krea2.py

@@ -69,7 +69,7 @@ KREA2_MMDIT_CONFIG = dict(
     kvheads=12,
     multiplier=4,
     layers=28,
-    patch=2,
+    patch_size=2,
     channels=16,
     txtheads=20,
     txtkvheads=20,
@@ -182,7 +182,7 @@ class Krea2Model(BaseModel):
         self.is_transformer = True
         self.target_lora_modules = ["SingleStreamDiT"]
 
-        self.patch_size = KREA2_MMDIT_CONFIG["patch"]
+        self.patch_size = KREA2_MMDIT_CONFIG["patch_size"]
         self.vae_scale_factor = 8  # Qwen-Image VAE is f8
         # Safety cap on prompt token length (truncation only); embeds are stored
         # per-sample at natural length and padded to the batch max at the model call.

extensions_built_in/diffusion_models/krea2/src/mmdit.py

@@ -14,6 +14,9 @@ Differences from the reference (all training-driven, numerically equivalent):
   - ``enable_gradient_checkpointing`` / ``disable_gradient_checkpointing`` and a
     per-block ``torch.utils.checkpoint`` wrapper are added (gated on
     ``torch.is_grad_enabled()`` so eval/sampling never pays for it).
+  - ``patch`` renamed to ``patch_size`` (diffusers-style name, matching the
+    toolkit's other archs) so the generic trainer's timestep-shift setup can
+    discover the token patch size via ``unet.config.patch_size``.
 """
 
 import math
@@ -110,7 +113,7 @@ class SingleMMDiTConfig:
     heads: int
     multiplier: int
     layers: int
-    patch: int
+    patch_size: int
     channels: int
     bias: bool = False
     theta: float = 1e3
@@ -263,10 +266,10 @@ class Attention(torch.nn.Module):
 
 
 class LastLayer(torch.nn.Module):
-    def __init__(self, features: int, patch: int, channels: int):
+    def __init__(self, features: int, patch_size: int, channels: int):
         super().__init__()
         self.norm = RMSNorm(features)
-        self.linear = torch.nn.Linear(features, patch * patch * channels, bias=True)
+        self.linear = torch.nn.Linear(features, patch_size * patch_size * channels, bias=True)
         self.modulation = SimpleModulation(features)
 
     def forward(self, x: Tensor, tvec: Tensor) -> Tensor:
@@ -427,7 +430,7 @@ class SingleStreamDiT(nn.Module):
             config.features, axes, theta=config.theta, ntk=1.0
         )
         self.first = nn.Linear(
-            config.channels * config.patch**2, config.features, bias=True
+            config.channels * config.patch_size**2, config.features, bias=True
         )
 
         self.blocks = nn.ModuleList(
@@ -461,7 +464,7 @@ class SingleStreamDiT(nn.Module):
             nn.GELU(approximate="tanh"),
             nn.Linear(config.features, config.features),
         )
-        self.last = LastLayer(config.features, config.patch, config.channels)
+        self.last = LastLayer(config.features, config.patch_size, config.channels)
 
         self.tproj = nn.Sequential(
             nn.GELU(approximate="tanh"), nn.Linear(config.features, config.features * 6)

extensions_built_in/diffusion_models/krea2/src/pipeline.py

@@ -172,7 +172,7 @@ def predict_velocity(
     the velocity ``noise - clean`` reshaped back to ``(B, C, h, w)``. No time
     flip / negation: Krea's convention matches toolkit's.
     """
-    patch = model.config.patch
+    patch = model.config.patch_size
     b, c, h, w = latents.shape
 
     if ref_kv_cache is not None and not isolate_refs: