github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Fixed issue that prevented full fine-tuning of flux2 models when using gradient checkpointing

Browse Source
This commit is contained in:
Jaret Burkett
2026-02-06 16:18:43 -07:00
parent 1422789452
commit e82cf6eec2
1 changed files with 2 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
extensions_built_in/diffusion_models/flux2/src/model.py
View File

@@ -170,8 +170,6 @@ class Flux2(nn.Module):
for block in self.double_blocks: for block in self.double_blocks:
if torch.is_grad_enabled() and self.gradient_checkpointing: if torch.is_grad_enabled() and self.gradient_checkpointing:
img.requires_grad_(True)
txt.requires_grad_(True)
img, txt = ckpt.checkpoint( img, txt = ckpt.checkpoint(
block, block,
img, img,
@@ -180,6 +178,7 @@ class Flux2(nn.Module):
pe_ctx, pe_ctx,
double_block_mod_img, double_block_mod_img,
double_block_mod_txt, double_block_mod_txt,
use_reentrant=False,
) )
else: else:
img, txt = block( img, txt = block(
@@ -196,12 +195,12 @@ class Flux2(nn.Module):
for i, block in enumerate(self.single_blocks): for i, block in enumerate(self.single_blocks):
if torch.is_grad_enabled() and self.gradient_checkpointing: if torch.is_grad_enabled() and self.gradient_checkpointing:
img.requires_grad_(True)
img = ckpt.checkpoint( img = ckpt.checkpoint(
block, block,
img, img,
pe, pe,
single_block_mod, single_block_mod,
use_reentrant=False,
) )
else: else:
img = block( img = block(