diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py
index 957918d..563616a 100644
--- a/extensions_built_in/sd_trainer/SDTrainer.py
+++ b/extensions_built_in/sd_trainer/SDTrainer.py
@@ -919,10 +919,7 @@ class SDTrainer(BaseSDTrainProcess):
                         tv = torch.clamp(tv, min=0.001)
                 
                 # step latent, use here or with do_fft_loss
-                if self.sd.x0_pred:
-                    t0 = noise_pred
-                else:
-                    t0 = noisy_latents - tv * noise_pred
+                t0 = self.sd.get_x0_prediction(noise_pred, noisy_latents, tv)
                 
                 if self.train_config.t0_loss_target:
                     # replace the loss targets and pred
diff --git a/toolkit/models/base_model.py b/toolkit/models/base_model.py
index 29256fd..4e57768 100644
--- a/toolkit/models/base_model.py
+++ b/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
diff --git a/toolkit/models/diffusion_feature_extraction.py b/toolkit/models/diffusion_feature_extraction.py
index 4d0b5f6..910bfbc 100644
--- a/toolkit/models/diffusion_feature_extraction.py
+++ b/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
diff --git a/toolkit/stable_diffusion_model.py b/toolkit/stable_diffusion_model.py
index 24c8043..91f2bfd 100644
--- a/toolkit/stable_diffusion_model.py
+++ b/toolkit/stable_diffusion_model.py
@@ -3276,3 +3276,11 @@ class StableDiffusion:
     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