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

Adjust DFE to handle 5 dimension latent spaces

Browse Source
This commit is contained in:
Jaret Burkett
2025-10-27 07:48:44 -06:00
parent 8c12977891
commit 42e5e3cd1c
1 changed files with 16 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -617,33 +617,28 @@ class SDTrainer(BaseSDTrainProcess):
stepped_latents = torch.cat(stepped_chunks, dim=0) stepped_latents = torch.cat(stepped_chunks, dim=0)
stepped_latents = stepped_latents.to(self.sd.vae.device, dtype=self.sd.vae.dtype) stepped_latents = stepped_latents.to(self.sd.vae.device, dtype=self.sd.vae.dtype)
# resize to half the size of the latents sl = stepped_latents
stepped_latents_half = torch.nn.functional.interpolate( if len(sl.shape) == 5:
stepped_latents, # video B,C,T,H,W
size=(stepped_latents.shape[2] // 2, stepped_latents.shape[3] // 2), sl = sl.permute(0, 2, 1, 3, 4) # B,T,C,H,W
mode='bilinear', b, t, c, h, w = sl.shape
align_corners=False sl = sl.reshape(b * t, c, h, w)
) pred_features = self.dfe(sl.float())
pred_features = self.dfe(stepped_latents.float())
pred_features_half = self.dfe(stepped_latents_half.float())
with torch.no_grad(): with torch.no_grad():
target_features = self.dfe(batch.latents.to(self.device_torch, dtype=torch.float32)) bl = batch.latents
batch_latents_half = torch.nn.functional.interpolate( bl = bl.to(self.sd.vae.device)
batch.latents.to(self.device_torch, dtype=torch.float32), if len(bl.shape) == 5:
size=(batch.latents.shape[2] // 2, batch.latents.shape[3] // 2), # video B,C,T,H,W
mode='bilinear', bl = bl.permute(0, 2, 1, 3, 4) # B,T,C,H,W
align_corners=False b, t, c, h, w = bl.shape
) bl = bl.reshape(b * t, c, h, w)
target_features_half = self.dfe(batch_latents_half) target_features = self.dfe(bl.float())
# scale dfe so it is weaker at higher noise levels # scale dfe so it is weaker at higher noise levels
dfe_scaler = 1 - (timesteps.float() / 1000.0).view(-1, 1, 1, 1).to(self.device_torch) dfe_scaler = 1 - (timesteps.float() / 1000.0).view(-1, 1, 1, 1).to(self.device_torch)
dfe_loss = torch.nn.functional.mse_loss(pred_features, target_features, reduction="none") * \ dfe_loss = torch.nn.functional.mse_loss(pred_features, target_features, reduction="none") * \
self.train_config.diffusion_feature_extractor_weight * dfe_scaler self.train_config.diffusion_feature_extractor_weight * dfe_scaler
additional_loss += dfe_loss.mean()
dfe_loss_half = torch.nn.functional.mse_loss(pred_features_half, target_features_half, reduction="none") * \
self.train_config.diffusion_feature_extractor_weight * dfe_scaler
additional_loss += dfe_loss.mean() + dfe_loss_half.mean()
elif self.dfe.version == 2: elif self.dfe.version == 2:
# version 2 # version 2
# do diffusion feature extraction on target # do diffusion feature extraction on target