Merge dd9a330860 into 1dae34b0af

Change single blocks to use num_blocks value
Add flexibility to block swapping for Flex model
2026-04-10 23:01:22 +00:00 · 2026-03-31 02:09:24 +00:00 · 2025-03-23 21:30:45 -04:00 · 2025-03-23 21:28:08 -04:00
1 changed files with 2 additions and 2 deletions
--- a/library/flux_models.py
+++ b/library/flux_models.py
@@ -971,8 +971,8 @@ class Flux(nn.Module):

    def enable_block_swap(self, num_blocks: int, device: torch.device):
        self.blocks_to_swap = num_blocks
-        double_blocks_to_swap = num_blocks // 2
-        single_blocks_to_swap = (num_blocks - double_blocks_to_swap) * 2
+        double_blocks_to_swap = min(self.num_double_blocks - 2, num_blocks // 2)
+        single_blocks_to_swap = (num_blocks - (num_blocks // 2)) * 2

        assert double_blocks_to_swap <= self.num_double_blocks - 2 and single_blocks_to_swap <= self.num_single_blocks - 2, (
            f"Cannot swap more than {self.num_double_blocks - 2} double blocks and {self.num_single_blocks - 2} single blocks. "
Author	SHA1	Message	Date
Dave Lage	21187eba7e	Merge `dd9a330860` into `1dae34b0af`	2026-03-31 02:09:24 +00:00
rockerBOO	dd9a330860	Change single blocks to use num_blocks value	2025-03-23 21:30:45 -04:00
rockerBOO	af5adb2b61	Add flexibility to block swapping for Flex model	2025-03-23 21:28:08 -04:00