diff --git a/modules/sd_models.py b/modules/sd_models.py
index bed80bdb5..35507dcec 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -947,8 +947,12 @@ def add_noise_pred_to_diffusers_callback(pipe):
         pipe.prior_pipe._callback_tensor_inputs.append("predicted_image_embedding") # pylint: disable=protected-access
     elif hasattr(pipe, "scheduler") and "flow" in pipe.scheduler.__class__.__name__.lower():
         pipe._callback_tensor_inputs.append("noise_pred") # pylint: disable=protected-access
+    elif hasattr(pipe, "scheduler") and hasattr(pipe.scheduler, "config") and getattr(pipe.scheduler.config, "prediction_type", "none") == "flow_prediction":
+        pipe._callback_tensor_inputs.append("noise_pred") # pylint: disable=protected-access
     elif hasattr(pipe, "default_scheduler") and "flow" in pipe.default_scheduler.__class__.__name__.lower():
         pipe._callback_tensor_inputs.append("noise_pred") # pylint: disable=protected-access
+    elif hasattr(pipe, "default_scheduler") and hasattr(pipe.default_scheduler, "config") and getattr(pipe.default_scheduler.config, "prediction_type", "none") == "flow_prediction":
+        pipe._callback_tensor_inputs.append("noise_pred") # pylint: disable=protected-access
     return pipe