VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-18 01:30:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add flexibility to block swapping for Flex model

Browse Source
This commit is contained in:
rockerBOO
2025-03-23 21:28:08 -04:00
parent 6364379f17
commit af5adb2b61
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
library/flux_models.py
View File

@@ -968,7 +968,7 @@ 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
double_blocks_to_swap = min(self.num_double_blocks - 2, num_blocks // 2)
single_blocks_to_swap = (num_blocks - double_blocks_to_swap) * 2
assert double_blocks_to_swap <= self.num_double_blocks - 2 and single_blocks_to_swap <= self.num_single_blocks - 2, (