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

Added weighing to DFE

Browse Source
This commit is contained in:
Jaret Burkett
2025-01-22 08:50:57 -07:00
parent 89dd041b97
commit 04abe57c76
2 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -385,7 +385,8 @@ class SDTrainer(BaseSDTrainProcess):
# do diffusion feature extraction on prediction # do diffusion feature extraction on prediction
pred_features = self.dfe(torch.cat([noise_pred.float(), noise.float()], dim=1)) pred_features = self.dfe(torch.cat([noise_pred.float(), noise.float()], dim=1))
additional_loss += torch.nn.functional.mse_loss(pred_features, target_features, reduction="mean") additional_loss += torch.nn.functional.mse_loss(pred_features, target_features, reduction="mean") * \
self.train_config.diffusion_feature_extractor_weight
if target is None: if target is None:
target = noise target = noise

1
toolkit/config_modules.py
View File

@@ -403,6 +403,7 @@ class TrainConfig:
# diffusion feature extractor # diffusion feature extractor
self.diffusion_feature_extractor_path = kwargs.get('diffusion_feature_extractor_path', None) self.diffusion_feature_extractor_path = kwargs.get('diffusion_feature_extractor_path', None)
self.diffusion_feature_extractor_weight = kwargs.get('diffusion_feature_extractor_weight', 0.1)
# optimal noise pairing # optimal noise pairing
self.optimal_noise_pairing_samples = kwargs.get('optimal_noise_pairing_samples', 1) self.optimal_noise_pairing_samples = kwargs.get('optimal_noise_pairing_samples', 1)