From 917d9e6c86816f4e5095101c17926cb39dd7c533 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 25 Dec 2024 18:43:10 +0300 Subject: [PATCH] Live previews add flow_prediction --- modules/processing_callbacks.py | 2 +- modules/sd_models.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/modules/processing_callbacks.py b/modules/processing_callbacks.py index c05529ee9..b896893f4 100644 --- a/modules/processing_callbacks.py +++ b/modules/processing_callbacks.py @@ -119,7 +119,7 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict = {} if "flow" in pipe.scheduler.__class__.__name__.lower(): shared.state.current_latent = original_sample - (noise_pred * sigma) elif hasattr(pipe.scheduler, "config") and hasattr(pipe.scheduler.config, "prediction_type"): - if pipe.scheduler.config.prediction_type == "epsilon": + if pipe.scheduler.config.prediction_type in {"epsilon", "flow_prediction"}: shared.state.current_latent = original_sample - (noise_pred * sigma) elif pipe.scheduler.config.prediction_type == "v_prediction": shared.state.current_latent = noise_pred * (-sigma / (sigma**2 + 1) ** 0.5) + (original_sample / (sigma**2 + 1)) diff --git a/modules/sd_models.py b/modules/sd_models.py index fa6398d3e..a52e15e7e 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1345,6 +1345,8 @@ def add_noise_pred_to_diffusers_callback(pipe): pipe.prior_pipe._callback_tensor_inputs.append("predicted_image_embedding") elif hasattr(pipe, "scheduler") and "flow" in pipe.scheduler.__class__.__name__.lower(): pipe._callback_tensor_inputs.append("noise_pred") + elif hasattr(pipe, "default_scheduler") and "flow" in pipe.default_scheduler.__class__.__name__.lower(): + pipe._callback_tensor_inputs.append("noise_pred") return pipe