Default MSRoPE to the shared scaling factor

src/models/dit_blocks.py

@@ -267,13 +267,13 @@ def set_real_rope(model: nn.Module, enabled: bool) -> None:
 
 
 def set_uniform_rope_scaling(model: nn.Module, enabled: bool) -> None:
-    """Scale both spatial axes by one shared factor instead of per-axis ones.
+    """Scale both spatial axes by one shared factor (default) or per-axis ones.
 
-    Per-axis factors (the DyPE default) normalize a rectangular canvas to the
-    square training window in phase space, so content dilates along the longer
-    axis by the aspect ratio. One shared max factor preserves the canvas
-    aspect in phase space at the cost of a shorter phase span on the short
-    axis. Instance-local launch policy; not checkpoint state.
+    Per-axis factors (the DyPE paper's behavior) normalize a rectangular
+    canvas to the square training window in phase space, so content dilates
+    along the longer axis by the aspect ratio. One shared max factor preserves
+    the canvas aspect in phase space at the cost of a shorter phase span on
+    the short axis. Instance-local launch policy; not checkpoint state.
     """
     if type(enabled) is not bool:
         raise ValueError("uniform RoPE scaling policy must be boolean")
@@ -288,6 +288,9 @@ class MSRoPE(nn.Module):
     The 896p training buckets reach 1184 pixels on either axis, or 74 patches
     at 16 pixels per patch. Training always uses the original frequencies.
     DyPE follows https://arxiv.org/html/2510.20766v2, equations 9 and 13–17.
+    Rectangular canvases default to one shared scaling factor so the phase
+    aspect matches the canvas aspect; set_uniform_rope_scaling(model, False)
+    restores the paper's per-axis factors.
     """
 
     def __init__(
@@ -311,7 +314,7 @@ class MSRoPE(nn.Module):
         # 0-based from the corner. Text starts at max(height, width) on both axes.
         self.centered = centered
         # Launch policy set by set_uniform_rope_scaling; not checkpoint state.
-        self.uniform_scaling = False
+        self.uniform_scaling = True
 
         # Precompute frequency tables for both spatial axes
         pos_index = torch.arange(2560)  # maximum position index