Move simga calcs to do_set_current_image

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:

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

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: