From 5a35e9dfd8ab34137306e705ab470852c9eb5d9b Mon Sep 17 00:00:00 2001 From: bghira <59658056+bghira@users.noreply.github.com> Date: Tue, 20 Jan 2026 11:41:25 -0600 Subject: [PATCH 1/3] qwen_image: fixes for TREAD --- .../helpers/models/qwen_image/model.py | 26 ++++++ .../helpers/models/qwen_image/transformer.py | 81 ++++++++++--------- 2 files changed, 71 insertions(+), 36 deletions(-) diff --git a/simpletuner/helpers/models/qwen_image/model.py b/simpletuner/helpers/models/qwen_image/model.py index 4ece95ed2..5451afb62 100644 --- a/simpletuner/helpers/models/qwen_image/model.py +++ b/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) diff --git a/simpletuner/helpers/models/qwen_image/transformer.py b/simpletuner/helpers/models/qwen_image/transformer.py index de9e35933..b223d61eb 100644 --- a/simpletuner/helpers/models/qwen_image/transformer.py +++ b/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: From ebc6b05abf4f6d9f79871eb9204e5b177446cf74 Mon Sep 17 00:00:00 2001 From: bghira <59658056+bghira@users.noreply.github.com> Date: Tue, 20 Jan 2026 11:41:37 -0600 Subject: [PATCH 2/3] test --- tests/test_tread_transformers.py | 48 ++++++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/tests/test_tread_transformers.py b/tests/test_tread_transformers.py index 4bf1bb730..985c72ef1 100644 --- a/tests/test_tread_transformers.py +++ b/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() From 01039958f8261807ae6197854b21ea6448f8aafe Mon Sep 17 00:00:00 2001 From: bghira <59658056+bghira@users.noreply.github.com> Date: Tue, 20 Jan 2026 11:44:20 -0600 Subject: [PATCH 3/3] enable qwen tread access via UI --- .../services/field_registry/sections/model.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/model.py b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/model.py index 8f8ccb023..d7f9b392b 100644 --- a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/model.py +++ b/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, )