Better live previews

This commit is contained in:
Disty0
2024-12-25 18:19:43 +03:00
parent a5f94fe0dd
commit 5e35d6c3a3
2 changed files with 33 additions and 0 deletions
+18
View File
@@ -105,6 +105,24 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict = {}
shared.state.current_latent = pipe._unpack_latents(kwargs['latents'], height, width, pipe.vae_scale_factor) # pylint: disable=protected-access
else:
shared.state.current_latent = kwargs['latents']
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 == "epsilon":
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))
except Exception as e:
shared.log.error(f'Callback: {e}')
if shared.cmd_opts.profile and shared.profiler is not None: