qwen_image: fixes for TREAD

simpletuner/helpers/models/qwen_image/model.py

@@ -305,6 +305,32 @@ def setup_training_noise_schedule(self):
 
         return self.config, self.noise_schedule
 
+    def tread_init(self):
+        """
+        Initialize the TREAD model training method.
+        """
+        from simpletuner.helpers.training.tread import TREADRouter
+
+        if (
+            getattr(self.config, "tread_config", None) is None
+            or getattr(self.config, "tread_config", None) is {}
+            or getattr(self.config, "tread_config", {}).get("routes", None) is None
+        ):
+            logger.error("TREAD training requires you to configure the routes in the TREAD config")
+            import sys
+
+            sys.exit(1)
+
+        self.unwrap_model(model=self.model).set_router(
+            TREADRouter(
+                seed=getattr(self.config, "seed", None) or 42,
+                device=self.accelerator.device,
+            ),
+            self.config.tread_config["routes"],
+        )
+
+        logger.info("TREAD training is enabled")
+
     def _get_model_flavour(self) -> Optional[str]:
         return getattr(self.config, "model_flavour", None)
 

simpletuner/helpers/models/qwen_image/transformer.py

@@ -1087,6 +1087,28 @@ def forward(
         router = self._tread_router
         use_routing = self.training and len(routes) > 0 and torch.is_grad_enabled()
 
+        # Pre-normalize route indices to handle negative values
+        if use_routing:
+            num_blocks = len(self.transformer_blocks)
+
+            def _to_pos(idx: int) -> int:
+                return idx if idx >= 0 else num_blocks + idx
+
+            routes = [
+                {
+                    **r,
+                    "start_layer_idx": _to_pos(r["start_layer_idx"]),
+                    "end_layer_idx": _to_pos(r["end_layer_idx"]),
+                }
+                for r in routes
+            ]
+
+        # TREAD state tracking
+        route_ptr = 0
+        routing_now = False
+        tread_mask_info = None
+        saved_tokens = None
+
         grad_enabled = torch.is_grad_enabled()
         musubi_manager = self._musubi_block_swap
         musubi_offload_active = False
@@ -1105,26 +1127,21 @@ def forward(
 
         capture_idx = 0
         for index_block, block in enumerate(self.transformer_blocks):
-            # TREAD routing for this layer
-            if use_routing:
-                # Check if this layer should use routing
-                for route in routes:
-                    start_idx = route["start_layer_idx"]
-                    end_idx = route["end_layer_idx"]
-                    # Handle negative indices
-                    if start_idx < 0:
-                        start_idx = len(self.transformer_blocks) + start_idx
-                    if end_idx < 0:
-                        end_idx = len(self.transformer_blocks) + end_idx
-
-                    if start_idx <= index_block <= end_idx:
-                        mask_info = router.get_mask(
-                            hidden_states.shape[1], route["selection_ratio"], force_keep_mask=force_keep_mask
-                        )
-                        hidden_states = router.start_route(hidden_states, mask_info)
-                        break
             if musubi_offload_active and musubi_manager.is_managed_block(index_block):
                 musubi_manager.stream_in(block, hidden_states.device)
+
+            # TREAD: START a route?
+            if use_routing and route_ptr < len(routes) and index_block == routes[route_ptr]["start_layer_idx"]:
+                mask_ratio = routes[route_ptr]["selection_ratio"]
+                tread_mask_info = router.get_mask(
+                    hidden_states,
+                    mask_ratio=mask_ratio,
+                    force_keep=force_keep_mask,
+                )
+                saved_tokens = hidden_states.clone()
+                hidden_states = router.start_route(hidden_states, tread_mask_info)
+                routing_now = True
+
             if torch.is_grad_enabled() and self.gradient_checkpointing:
 
                 def create_custom_forward(module, mod_idx):
@@ -1181,24 +1198,16 @@ def custom_forward(
 
             _store_hidden_state(hidden_states_buffer, f"layer_{capture_idx}", hidden_states)
             capture_idx += 1
-            # TREAD end routing for this layer
-            if use_routing:
-                # Check if this layer should end routing
-                for route in routes:
-                    start_idx = route["start_layer_idx"]
-                    end_idx = route["end_layer_idx"]
-                    # Handle negative indices
-                    if start_idx < 0:
-                        start_idx = len(self.transformer_blocks) + start_idx
-                    if end_idx < 0:
-                        end_idx = len(self.transformer_blocks) + end_idx
-
-                    if start_idx <= index_block <= end_idx:
-                        mask_info = router.get_mask(
-                            hidden_states.shape[1], route["selection_ratio"], force_keep_mask=force_keep_mask
-                        )
-                        hidden_states = router.end_route(hidden_states, mask_info)
-                        break
+
+            # TREAD: END the current route?
+            if routing_now and index_block == routes[route_ptr]["end_layer_idx"]:
+                hidden_states = router.end_route(
+                    hidden_states,
+                    tread_mask_info,
+                    original_x=saved_tokens,
+                )
+                routing_now = False
+                route_ptr += 1
 
             # controlnet residual
             if controlnet_block_samples is not None:

simpletuner/simpletuner_sdk/server/services/field_registry/sections/model.py

@@ -699,13 +699,14 @@ def _quant_label(value: str) -> str:
                         "flux",
                         "flux2",
                         "hidream",
-                        "wan",
-                        "wan_s2v",
+                        "kandinsky5-image",
+                        "kandinsky5-video",
                         "pixart",
+                        "qwen_image",
                         "sana",
                         "sd3",
-                        "kandinsky5-image",
-                        "kandinsky5-video",
+                        "wan",
+                        "wan_s2v",
                     ],
                 )
             ],
@@ -718,13 +719,14 @@ def _quant_label(value: str) -> str:
                 "flux",
                 "flux2",
                 "hidream",
-                "wan",
-                "wan_s2v",
+                "kandinsky5-image",
+                "kandinsky5-video",
                 "pixart",
+                "qwen_image",
                 "sana",
                 "sd3",
-                "kandinsky5-image",
-                "kandinsky5-video",
+                "wan",
+                "wan_s2v",
             ],
             order=23,
         )

tests/test_tread_transformers.py

@@ -370,6 +370,31 @@ def test_sana_tread_integration(self):
             self.assertIsNotNone(model._tread_router)
             self.assertIsNotNone(model._tread_routes)
 
+    def test_qwen_image_tread_integration(self):
+        """Test TREAD integration in Qwen Image transformer."""
+        with patch("simpletuner.helpers.models.qwen_image.transformer.TREADRouter"):
+            from simpletuner.helpers.models.qwen_image.transformer import QwenImageTransformer2DModel
+
+            # Create model with minimal config
+            model = QwenImageTransformer2DModel(
+                patch_size=2,
+                in_channels=16,
+                out_channels=16,
+                num_layers=2,
+                num_attention_heads=12,
+                attention_head_dim=64,
+                joint_attention_dim=2048,
+            )
+
+            # Test set_router method
+            model.set_router(self.mock_router, self.test_routes)
+            self.assertEqual(model._tread_router, self.mock_router)
+            self.assertEqual(model._tread_routes, self.test_routes)
+
+            # Test TREAD attributes exist
+            self.assertIsNotNone(model._tread_router)
+            self.assertIsNotNone(model._tread_routes)
+
 
 class TestTREADModelInitialization(unittest.TestCase):
     """Test TREAD initialization methods in model classes."""
@@ -570,6 +595,29 @@ def test_sana_tread_init(self):
                 # Verify TREADRouter was called with correct parameters
                 mock_tread_router.assert_called_once_with(seed=42, device="mps")
 
+    def test_qwen_image_tread_init(self):
+        """Test TREAD initialization in Qwen Image model."""
+        with patch("simpletuner.helpers.training.tread.TREADRouter") as mock_tread_router:
+            with patch("simpletuner.helpers.models.qwen_image.model.logger"):
+                from simpletuner.helpers.models.qwen_image.model import QwenImage
+
+                # Create mock model instance
+                model_instance = Mock()
+                model_instance.config = self.mock_config
+                model_instance.accelerator = Mock()
+                model_instance.accelerator.device = "cuda"
+                model_instance.unwrap_model.return_value = Mock()
+
+                # Test tread_init method exists
+                self.assertTrue(hasattr(QwenImage, "tread_init"))
+
+                # Test successful initialization
+                model_instance.__class__ = QwenImage
+                QwenImage.tread_init(model_instance)
+
+                # Verify TREADRouter was called with correct parameters
+                mock_tread_router.assert_called_once_with(seed=42, device="cuda")
+
 
 if __name__ == "__main__":
     unittest.main()