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

toolkit/models/base_model.py

@@ -1783,3 +1783,11 @@ class BaseModel:
     def scale_loss(self, loss):
         # called to get the loss scaler for the model. Can be overridden in child classes
         return loss
+    
+    def get_x0_prediction(self, model_output: torch.Tensor, model_input: torch.Tensor, sigma: torch.Tensor):
+        # can be overridden in child classes to get the x0 prediction from model output
+        # step latent, use here or with do_fft_loss
+        if self.x0_pred:
+            return model_output
+        else:
+            return model_input - sigma * model_output

toolkit/models/diffusion_feature_extraction.py

@@ -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