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

Add sapiens2 as a diffusion feature extractor

Browse Source
This commit is contained in:
Jaret Burkett
2026-04-27 15:59:03 -06:00
parent 67d67f8c1d
commit 1fc4ad3979
3 changed files with 1110 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -365,6 +365,15 @@ class SDTrainer(BaseSDTrainProcess):
# must be set to train for gradient checkpointing to work
self.dfe.vision_encoder.train()
self.dfe.vision_encoder.gradient_checkpointing = True
elif hasattr(self.dfe, 'model') and self.train_config.gradient_checkpointing:
if hasattr(self.dfe.model, 'enable_gradient_checkpointing'):
self.dfe.model.train()
self.dfe.model.enable_gradient_checkpointing()
elif hasattr(self.dfe.model, 'gradient_checkpointing'):
self.dfe.model.train()
self.dfe.model.gradient_checkpointing = True
else:
print_acc("Warning: Could not enable gradient checkpointing on diffusion feature extractor model.")
else:
self.dfe.eval()
@@ -664,7 +673,7 @@ class SDTrainer(BaseSDTrainProcess):
dfe_loss += torch.nn.functional.mse_loss(pred_feature_list[i], target_feature_list[i], reduction="mean")
additional_loss += dfe_loss * self.train_config.diffusion_feature_extractor_weight * 100.0
elif self.dfe.version in [3, 4, 5, 6, 7, 8]:
elif self.dfe.version in [3, 4, 5, 6, 7, 8, 9, 10]:
dfe_loss = self.dfe(
noise=noise,
noise_pred=noise_pred,