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

Migrate to a new DTO for latents to carry more information that a normal tensor such as audio.

Browse Source
This commit is contained in:
Jaret Burkett committed 2026-08-30 10:30:21 -06:00
1 parent 2a69c1e7de
commit 764b5064fb
8 files changed
+369 -215
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

jobs/process/BaseSDTrainProcess.py
+3 -35
View File
@@ -1133,34 +1133,11 @@ class BaseSDTrainProcess(BaseTrainProcess):
latents = self.sd.encode_images(imgs)
batch.latents = latents
if self.train_config.standardize_latents:
if self.sd.is_xl or self.sd.is_vega or self.sd.is_ssd:
target_mean_list = [-0.1075, 0.0231, -0.0135, 0.2164]
target_std_list = [0.8979, 0.7505, 0.9150, 0.7451]
else:
target_mean_list = [0.2949, -0.3188, 0.0807, 0.1929]
target_std_list = [0.8560, 0.9629, 0.7778, 0.6719]
latents_channel_mean = latents.mean(dim=(2, 3), keepdim=True)
latents_channel_std = latents.std(dim=(2, 3), keepdim=True)
latents = (latents - latents_channel_mean) / latents_channel_std
target_mean = torch.tensor(target_mean_list, device=self.device_torch, dtype=dtype)
target_std = torch.tensor(target_std_list, device=self.device_torch, dtype=dtype)
# expand them to match dim
target_mean = target_mean.unsqueeze(0).unsqueeze(2).unsqueeze(3)
target_std = target_std.unsqueeze(0).unsqueeze(2).unsqueeze(3)
latents = latents * target_std + target_mean
batch.latents = latents
# show_latents(latents, self.sd.vae, 'latents')
if batch.unconditional_tensor is not None and batch.unconditional_latents is None:
unconditional_imgs = batch.unconditional_tensor
unconditional_imgs = unconditional_imgs.to(self.device_torch, dtype=dtype)
unconditional_latents = self.sd.encode_images(unconditional_imgs)
batch.unconditional_latents = unconditional_latents * self.train_config.latent_multiplier
batch.unconditional_latents = unconditional_latents
unaugmented_latents = None
if self.train_config.loss_target == 'differential_noise':
@@ -1391,33 +1368,24 @@ class BaseSDTrainProcess(BaseTrainProcess):
noise = noise * noise_multiplier
with self.timer('make_noisy_latents'):
latent_multiplier = self.train_config.latent_multiplier
# handle adaptive scaling mased on std
if self.train_config.adaptive_scaling_factor:
std = latents.std(dim=(2, 3), keepdim=True)
normalizer = 1 / (std + 1e-6)
latent_multiplier = normalizer
latents = latents * (1 / (std + 1e-6))
latents = latents * latent_multiplier
if self.train_config.do_blank_stabilization:
# zero out latents with blank prompts
blank_latent = torch.zeros_like(latents)
for i, prompt in enumerate(conditioned_prompts):
if prompt.strip() == '':
latents[i] = blank_latent[i]
batch.latents = latents
# normalize latents to a mean of 0 and an std of 1
# mean_zero_latents = latents - latents.mean()
# latents = mean_zero_latents / mean_zero_latents.std()
if batch.unconditional_latents is not None:
batch.unconditional_latents = batch.unconditional_latents * self.train_config.latent_multiplier
noisy_latents = self.sd.add_noise(latents, noise, timesteps)
# determine scaled noise