diff --git a/invokeai/backend/flux/denoise.py b/invokeai/backend/flux/denoise.py index 72f87e2aef..ebad4cf61c 100644 --- a/invokeai/backend/flux/denoise.py +++ b/invokeai/backend/flux/denoise.py @@ -22,7 +22,18 @@ def denoise( guidance: float, traj_guidance_extension: TrajectoryGuidanceExtension | None, # noqa: F821 ): - step = 0 + # step 0 is the initial state + total_steps = len(timesteps) - 1 + step_callback( + PipelineIntermediateState( + step=0, + order=1, + total_steps=total_steps, + timestep=int(timesteps[0]), + latents=img, + ), + ) + step = 1 # guidance_vec is ignored for schnell. guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype) for t_curr, t_prev in tqdm(list(zip(timesteps[:-1], timesteps[1:], strict=True))): @@ -49,7 +60,7 @@ def denoise( PipelineIntermediateState( step=step, order=1, - total_steps=len(timesteps), + total_steps=total_steps, timestep=int(t_curr), latents=preview_img, ), diff --git a/invokeai/backend/stable_diffusion/diffusers_pipeline.py b/invokeai/backend/stable_diffusion/diffusers_pipeline.py index b3a668518b..646e1a92d3 100644 --- a/invokeai/backend/stable_diffusion/diffusers_pipeline.py +++ b/invokeai/backend/stable_diffusion/diffusers_pipeline.py @@ -366,7 +366,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline): with attn_ctx: callback( PipelineIntermediateState( - step=-1, + step=0, # initial latents order=self.scheduler.order, total_steps=len(timesteps), timestep=self.scheduler.config.num_train_timesteps, @@ -395,7 +395,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline): callback( PipelineIntermediateState( - step=i, + step=i + 1, # final latents order=self.scheduler.order, total_steps=len(timesteps), timestep=int(t), diff --git a/invokeai/backend/stable_diffusion/multi_diffusion_pipeline.py b/invokeai/backend/stable_diffusion/multi_diffusion_pipeline.py index 6c07fc1c2c..63e74de504 100644 --- a/invokeai/backend/stable_diffusion/multi_diffusion_pipeline.py +++ b/invokeai/backend/stable_diffusion/multi_diffusion_pipeline.py @@ -81,7 +81,7 @@ class MultiDiffusionPipeline(StableDiffusionGeneratorPipeline): callback( PipelineIntermediateState( - step=-1, + step=0, order=self.scheduler.order, total_steps=len(timesteps), timestep=self.scheduler.config.num_train_timesteps, @@ -182,7 +182,7 @@ class MultiDiffusionPipeline(StableDiffusionGeneratorPipeline): callback( PipelineIntermediateState( - step=i, + step=i + 1, order=self.scheduler.order, total_steps=len(timesteps), timestep=int(t),