Add dynamic shift model invariants

Author: bghira

simpletuner/helpers/models/flux2/model.py

@@ -87,6 +87,7 @@ class Flux2(ImageModelFoundation):
     ENABLED_IN_WIZARD = True
     PREDICTION_TYPE = PredictionTypes.FLOW_MATCHING
     MODEL_TYPE = ModelTypes.TRANSFORMER
+    USES_DYNAMIC_SHIFT = True
     AUTO_LORA_FORMAT_DETECTION = True
     NATIVE_COMFYUI_LORA_SUPPORT = True  # Flux2 has native ComfyUI LoRA support, no conversion needed
     COMFYUI_LORA_PRESERVE_COMPONENT_PREFIXES = {"transformer"}

tests/test_validation_dynamic_shift.py

@@ -1,10 +1,13 @@
+import inspect
+import math
 import types
 import unittest
 from unittest.mock import MagicMock, patch
 
 import torch
 
 from simpletuner.helpers.models.common import ModelFoundation
+from simpletuner.helpers.models.registry import ModelRegistry
 from simpletuner.helpers.training.validation import Evaluation
 
 
@@ -164,6 +167,70 @@ def set_timesteps(self, num_inference_steps=None, mu=None, **kwargs):
         with self.assertRaises(ValueError):
             eval_helper.get_timestep_schedule(scheduler, latents=torch.zeros(1, 4, 4, 4))
 
+    def _dynamic_shift_model_classes(self):
+        dynamic_shift_model_classes = {}
+        for family, registry_entry in ModelRegistry.model_families().items():
+            if hasattr(registry_entry, "get_real_class"):
+                model_cls = registry_entry.get_real_class()
+            else:
+                model_cls = registry_entry
+            if getattr(model_cls, "USES_DYNAMIC_SHIFT", False):
+                dynamic_shift_model_classes[family] = model_cls
+        return dynamic_shift_model_classes
+
+    def _model_class_declares_patch_size(self, model_cls):
+        component_cls = getattr(model_cls, "MODEL_CLASS", None)
+        self.assertIsNotNone(component_cls, f"{model_cls.__name__} must define MODEL_CLASS")
+        init_method = getattr(component_cls, "__init__", None)
+        try:
+            parameters = inspect.signature(init_method).parameters
+        except (TypeError, ValueError):
+            return False
+        return "patch_size" in parameters
+
+    def test_dynamic_shift_models_expose_patch_geometry_or_sequence_length(self):
+        model_classes = self._dynamic_shift_model_classes()
+        self.assertGreater(len(model_classes), 0)
+
+        for family, model_cls in model_classes.items():
+            with self.subTest(family=family):
+                has_model_sequence_length = "_latent_sequence_length" in model_cls.__dict__
+                has_component_patch_size = self._model_class_declares_patch_size(model_cls)
+                self.assertTrue(
+                    has_model_sequence_length or has_component_patch_size,
+                    f"{model_cls.__name__} uses dynamic shift but does not expose patch geometry",
+                )
+
+                model = model_cls.__new__(model_cls)
+                model.config = types.SimpleNamespace(controlnet=False)
+                model.accelerator = None
+                model.model = types.SimpleNamespace(config=types.SimpleNamespace())
+                if not has_model_sequence_length:
+                    model.model.config.patch_size = 2
+                    model.model.config.patch_size_t = 1
+
+                latent_channels = getattr(model_cls, "LATENT_CHANNEL_COUNT", 4)
+                latents = torch.zeros(1, latent_channels, 8, 8)
+                mu = model.calculate_dynamic_shift_mu(self.scheduler, latents)
+
+                self.assertTrue(math.isfinite(mu), f"{model_cls.__name__} produced non-finite dynamic shift mu")
+
+    def test_dynamic_scheduler_configs_declare_dynamic_shift_flag(self):
+        for family, registry_entry in ModelRegistry.model_families().items():
+            if hasattr(registry_entry, "get_real_class"):
+                model_cls = registry_entry.get_real_class()
+            else:
+                model_cls = registry_entry
+            scheduler_config = getattr(inspect.getmodule(model_cls), "SCHEDULER_CONFIG", None)
+            if not isinstance(scheduler_config, dict) or not scheduler_config.get("use_dynamic_shifting"):
+                continue
+
+            with self.subTest(family=family):
+                self.assertTrue(
+                    getattr(model_cls, "USES_DYNAMIC_SHIFT", False),
+                    f"{model_cls.__name__} scheduler config enables dynamic shift without USES_DYNAMIC_SHIFT",
+                )
+
 
 if __name__ == "__main__":
     unittest.main()