diff --git a/modules/processing_callbacks.py b/modules/processing_callbacks.py index b896893f4..39a6c66d8 100644 --- a/modules/processing_callbacks.py +++ b/modules/processing_callbacks.py @@ -106,23 +106,12 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict = {} else: shared.state.current_latent = kwargs['latents'] + shared.state.current_noise_pred = kwargs.get("noise_pred", None) + if shared.state.current_noise_pred is None: + shared.state.current_noise_pred = kwargs.get("predicted_image_embedding", None) if hasattr(pipe, "scheduler") and hasattr(pipe.scheduler, "sigmas"): - noise_pred = None - if kwargs.get("noise_pred", None) is not None: - noise_pred = kwargs.get("noise_pred") - elif kwargs.get("predicted_image_embedding", None) is not None: - noise_pred = kwargs.get("predicted_image_embedding") - if noise_pred is not None: - sigma = pipe.scheduler.sigmas[step] - sigma_next = pipe.scheduler.sigmas[step + 1] - original_sample = shared.state.current_latent - (noise_pred * (sigma_next-sigma)) - 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 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)) + shared.state.current_sigma = pipe.scheduler.sigmas[step] + shared.state.current_sigma_next = pipe.scheduler.sigmas[step + 1] except Exception as e: shared.log.error(f'Callback: {e}') if shared.cmd_opts.profile and shared.profiler is not None: diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py index dc58a2419..4d1666ff0 100644 --- a/modules/sd_samplers.py +++ b/modules/sd_samplers.py @@ -62,6 +62,10 @@ def create_sampler(name, model): model.prior_pipe.scheduler.config.clip_sample = False config = {k: v for k, v in model.scheduler.config.items() if not k.startswith('_')} shared.log.debug(f'Sampler: sampler=default class={current}: {config}') + if "flow" in model.scheduler.__class__.__name__.lower(): + shared.state.prediction_type = "flow_prediction" + elif hasattr(model.scheduler, "config") and hasattr(model.scheduler.config, "prediction_type"): + shared.state.prediction_type = model.scheduler.config.prediction_type return model.scheduler config = find_sampler_config(name) if config is None or config.constructor is None: @@ -94,6 +98,10 @@ def create_sampler(name, model): if hasattr(model, "prior_pipe") and hasattr(model.prior_pipe, "scheduler"): model.prior_pipe.scheduler = sampler.sampler model.prior_pipe.scheduler.config.clip_sample = False + if "flow" in model.scheduler.__class__.__name__.lower(): + shared.state.prediction_type = "flow_prediction" + elif hasattr(model.scheduler, "config") and hasattr(model.scheduler.config, "prediction_type"): + shared.state.prediction_type = model.scheduler.config.prediction_type clean_config = {k: v for k, v in sampler.config.items() if v is not None and v is not False} shared.log.debug(f'Sampler: sampler="{sampler.name}" class="{model.scheduler.__class__.__name__} config={clean_config}') return sampler.sampler diff --git a/modules/shared_state.py b/modules/shared_state.py index a3312ec33..2bcf8e23a 100644 --- a/modules/shared_state.py +++ b/modules/shared_state.py @@ -17,10 +17,14 @@ class State: sampling_step = 0 sampling_steps = 0 current_latent = None + current_noise_pred = None + current_sigma = None + current_sigma_next = None current_image = None current_image_sampling_step = 0 id_live_preview = 0 textinfo = None + prediction_type = "epsilon" api = False time_start = None need_restart = False @@ -102,6 +106,9 @@ class State: self.current_image = None self.current_image_sampling_step = 0 self.current_latent = None + self.current_noise_pred = None + self.current_sigma = None + self.current_sigma_next = None self.id_live_preview = 0 self.interrupted = False self.job = title @@ -113,6 +120,7 @@ class State: self.sampling_step = 0 self.skipped = False self.textinfo = None + self.prediction_type = "epsilon" self.api = api if api is not None else self.api self.time_start = time.time() if self.debug_output: @@ -152,7 +160,14 @@ class State: from modules.shared import opts import modules.sd_samplers # pylint: disable=W0621 try: - image = modules.sd_samplers.samples_to_image_grid(self.current_latent) if opts.show_progress_grid else modules.sd_samplers.sample_to_image(self.current_latent) + sample = self.current_latent + if self.current_noise_pred is not None and self.current_sigma is not None and self.current_sigma_next is not None: + original_sample = sample - (self.current_noise_pred * (self.current_sigma_next-self.current_sigma)) + if self.prediction_type in {"epsilon", "flow_prediction"}: + sample = original_sample - (self.current_noise_pred * self.current_sigma) + elif self.prediction_type == "v_prediction": + sample = self.current_noise_pred * (-self.current_sigma / (self.current_sigma**2 + 1) ** 0.5) + (original_sample / (self.current_sigma**2 + 1)) + image = modules.sd_samplers.samples_to_image_grid(sample) if opts.show_progress_grid else modules.sd_samplers.sample_to_image(sample) self.assign_current_image(image) self.current_image_sampling_step = self.sampling_step except Exception: