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

Scale DFE 9 to a velocity equiv weight to match flow matching gradient strength. Probably need to rework all DFEs to do this as the math checks out.

Browse Source
This commit is contained in:
Jaret Burkett
2026-04-28 09:10:02 -06:00
parent 1fc4ad3979
commit acc6a36214
1 changed files with 11 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
toolkit/models/diffusion_feature_extraction.py
View File

@@ -1118,8 +1118,8 @@ class DiffusionFeatureExtractor9(nn.Module):
# min 0.001 # min 0.001
tv = torch.clamp(tv, min=0.001) tv = torch.clamp(tv, min=0.001)
else: else:
# step is random 0.05 to 0.02 # step is random 0.1 to 0.25
step = torch.rand_like(tv) * 0.03 + 0.02 step = torch.rand_like(tv) * 0.15 + 0.1
next_step = tv - step next_step = tv - step
next_step = torch.clamp(next_step, min=0.0) next_step = torch.clamp(next_step, min=0.0)
stepped_latents = noisy_latents + (next_step - tv) * noise_pred stepped_latents = noisy_latents + (next_step - tv) * noise_pred
@@ -1147,17 +1147,19 @@ class DiffusionFeatureExtractor9(nn.Module):
pred_images = pred_images.to(device, dtype=dtype) pred_images = pred_images.to(device, dtype=dtype)
pred = self.get_pred(pred_images) pred = self.get_pred(pred_images)
loss = torch.nn.functional.mse_loss( perceptual_loss = torch.nn.functional.mse_loss(
pred.float(), target.float() pred.float(), target.float(), reduction="none"
) * 100.0 )
velocity_equiv_weight = (1.0 / torch.clamp(tv, min=0.001) ** 2)
loss_perceptual = (perceptual_loss * velocity_equiv_weight).mean()
if self.do_partial_step: if self.do_partial_step:
loss = loss * 10.0 loss_perceptual = loss_perceptual * 10.0
if 'loss' not in self.losses: if 'loss' not in self.losses:
self.losses['loss'] = loss.item() self.losses['loss'] = loss_perceptual.item()
else: else:
self.losses['loss'] += loss.item() self.losses['loss'] += loss_perceptual.item()
with torch.no_grad(): with torch.no_grad():
if self.step % self.log_every == 0 and self.step > 0: if self.step % self.log_every == 0 and self.step > 0:
print(f"DFE losses:") print(f"DFE losses:")
@@ -1170,7 +1172,7 @@ class DiffusionFeatureExtractor9(nn.Module):
# total_loss += mse_loss # total_loss += mse_loss
self.step += 1 self.step += 1
return loss return loss_perceptual
def load_dfe(model_path, vae=None, sd: 'BaseModel' = None) -> DiffusionFeatureExtractor: def load_dfe(model_path, vae=None, sd: 'BaseModel' = None) -> DiffusionFeatureExtractor:
if model_path == "v3": if model_path == "v3":