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

Moved x0 logic to the model class instead of recalculating it multiple times in the trainer

Browse Source
This commit is contained in:
Jaret Burkett committed 2026-10-21 11:38:09 -06:00
1 parent c22f21fa95
commit 15303bf876
4 files changed
+32 -9
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

toolkit/models/diffusion_feature_extraction.py
+15 -5
View File
@@ -552,7 +552,9 @@ class DiffusionFeatureExtractor4(nn.Module):
# min 0.001
tv = torch.clamp(tv, min=0.001)
if self.sd_ref is not None and getattr(self.sd_ref(), "x0_pred", False):
if self.sd_ref is not None and getattr(self.sd_ref(), "get_x0_prediction", None) is not None:
x0 = self.sd_ref().get_x0_prediction(noise_pred, noisy_latents, tv)
elif self.sd_ref is not None and getattr(self.sd_ref(), "x0_pred", False):
x0 = noise_pred
else:
# step latent
@@ -781,7 +783,9 @@ class DiffusionFeatureExtractor6(nn.Module):
# min 0.001
tv = torch.clamp(tv, min=0.001)
if self.sd_ref is not None and getattr(self.sd_ref(), "x0_pred", False):
if self.sd_ref is not None and getattr(self.sd_ref(), "get_x0_prediction", None) is not None:
x0 = self.sd_ref().get_x0_prediction(noise_pred, noisy_latents, tv)
elif self.sd_ref is not None and getattr(self.sd_ref(), "x0_pred", False):
x0 = noise_pred
else:
# step latent
@@ -960,7 +964,9 @@ class DiffusionFeatureExtractor7(nn.Module):
target_0_1 = (tensors + 1) / 2 # 0 to 1
if not self.do_partial_step:
if getattr(self.sd_ref(), "x0_pred", False):
if getattr(self.sd_ref(), "get_x0_prediction", None) is not None:
x0 = self.sd_ref().get_x0_prediction(noise_pred, noisy_latents, tv)
elif getattr(self.sd_ref(), "x0_pred", False):
x0 = noise_pred
else:
# step latent
@@ -1164,7 +1170,9 @@ class DiffusionFeatureExtractor9(nn.Module):
target_0_1 = (tensors + 1) / 2 # 0 to 1
if not self.do_partial_step:
if getattr(self.sd_ref(), "x0_pred", False):
if getattr(self.sd_ref(), "get_x0_prediction", None) is not None:
x0 = self.sd_ref().get_x0_prediction(noise_pred, noisy_latents, tv)
elif getattr(self.sd_ref(), "x0_pred", False):
x0 = noise_pred
else:
# step latent
@@ -1327,7 +1335,9 @@ class DiffusionFeatureExtractor10(nn.Module):
target_0_1 = (tensors + 1) / 2 # 0 to 1
if not self.do_partial_step:
if getattr(self.sd_ref(), "x0_pred", False):
if self.sd_ref is not None and getattr(self.sd_ref(), "get_x0_prediction", None) is not None:
x0 = self.sd_ref().get_x0_prediction(noise_pred, noisy_latents, tv)
elif self.sd_ref is not None and getattr(self.sd_ref(), "x0_pred", False):
x0 = noise_pred
else:
# step latent