From 68d12628368a26f81650412a5f00577bffd2edcd Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 16 Apr 2025 18:40:50 +0300 Subject: [PATCH] Fix live previews with HiDream --- modules/sd_models.py | 4 ++++ 1 file changed, 4 insertions(+) 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