Skip to content

Commit f76282a

Browse files
authored
Fix handling handling of 0-step denoising process (#6544)
## Summary #6522 introduced a change in behavior in cases where start/end were set such that there are 0 timesteps. This PR reverts that change. cc @StAlKeR7779 ## QA Instructions Run with euler, 5 steps, start: 0.0, end: 0.05. I ran this test before #6522, after #6522, and on this branch. This branch restores the behavior to pre-#6522 i.e. noise is injected even if no denoising steps are applied. ## Checklist - [x] _The PR has a short but descriptive title, suitable for a changelog_ - [x] _Tests added / updated (if applicable)_ - [x] _Documentation added / updated (if applicable)_
2 parents dc23beb + 9a3b8c6 commit f76282a

File tree

4 files changed

+15
-11
lines changed

4 files changed

+15
-11
lines changed

invokeai/app/invocations/denoise_latents.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -625,6 +625,7 @@ def init_scheduler(
625625
t_start_idx *= scheduler.order
626626
t_end_idx *= scheduler.order
627627

628+
init_timestep = timesteps[t_start_idx : t_start_idx + 1]
628629
timesteps = timesteps[t_start_idx : t_start_idx + t_end_idx]
629630

630631
scheduler_step_kwargs: Dict[str, Any] = {}
@@ -647,7 +648,7 @@ def init_scheduler(
647648
if isinstance(scheduler, TCDScheduler):
648649
scheduler_step_kwargs.update({"eta": 1.0})
649650

650-
return timesteps, scheduler_step_kwargs
651+
return timesteps, init_timestep, scheduler_step_kwargs
651652

652653
def prep_inpaint_mask(
653654
self, context: InvocationContext, latents: torch.Tensor
@@ -813,7 +814,7 @@ def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
813814
dtype=unet.dtype,
814815
)
815816

816-
timesteps, scheduler_step_kwargs = self.init_scheduler(
817+
timesteps, init_timestep, scheduler_step_kwargs = self.init_scheduler(
817818
scheduler,
818819
device=unet.device,
819820
steps=self.steps,
@@ -825,6 +826,7 @@ def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
825826
result_latents = pipeline.latents_from_embeddings(
826827
latents=latents,
827828
timesteps=timesteps,
829+
init_timestep=init_timestep,
828830
noise=noise,
829831
seed=seed,
830832
mask=mask,

invokeai/app/invocations/tiled_multi_diffusion_denoise_latents.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -252,7 +252,7 @@ def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
252252
)
253253
)
254254

255-
timesteps, scheduler_step_kwargs = DenoiseLatentsInvocation.init_scheduler(
255+
timesteps, init_timestep, scheduler_step_kwargs = DenoiseLatentsInvocation.init_scheduler(
256256
scheduler,
257257
device=unet.device,
258258
steps=self.steps,
@@ -269,6 +269,7 @@ def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
269269
scheduler_step_kwargs=scheduler_step_kwargs,
270270
noise=noise,
271271
timesteps=timesteps,
272+
init_timestep=init_timestep,
272273
callback=step_callback,
273274
)
274275

invokeai/backend/stable_diffusion/diffusers_pipeline.py

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,7 @@ def latents_from_embeddings(
273273
noise: Optional[torch.Tensor],
274274
seed: int,
275275
timesteps: torch.Tensor,
276+
init_timestep: torch.Tensor,
276277
callback: Callable[[PipelineIntermediateState], None],
277278
control_data: list[ControlNetData] | None = None,
278279
ip_adapter_data: Optional[list[IPAdapterData]] = None,
@@ -298,6 +299,8 @@ def latents_from_embeddings(
298299
HACK(ryand): seed is only used in a particular case when `noise` is None, but we need to re-generate the
299300
same noise used earlier in the pipeline. This should really be handled in a clearer way.
300301
timesteps: The timestep schedule for the denoising process.
302+
init_timestep: The first timestep in the schedule. This is used to determine the initial noise level, so
303+
should be populated if you want noise applied *even* if timesteps is empty.
301304
callback: A callback function that is called to report progress during the denoising process.
302305
control_data: ControlNet data.
303306
ip_adapter_data: IP-Adapter data.
@@ -312,17 +315,16 @@ def latents_from_embeddings(
312315
SD UNet model.
313316
is_gradient_mask: A flag indicating whether `mask` is a gradient mask or not.
314317
"""
315-
if timesteps.shape[0] == 0:
318+
if init_timestep.shape[0] == 0:
316319
return latents
317320

318321
orig_latents = latents.clone()
322+
319323
batch_size = latents.shape[0]
324+
batched_init_timestep = init_timestep.expand(batch_size)
320325

321326
# noise can be None if the latents have already been noised (e.g. when running the SDXL refiner).
322327
if noise is not None:
323-
# batched_init_timestep should have shape (batch_size, 1).
324-
batched_init_timestep = timesteps[0:1].expand(batch_size)
325-
326328
# TODO(ryand): I'm pretty sure we should be applying init_noise_sigma in cases where we are starting with
327329
# full noise. Investigate the history of why this got commented out.
328330
# latents = noise * self.scheduler.init_noise_sigma # it's like in t2l according to diffusers

invokeai/backend/stable_diffusion/multi_diffusion_pipeline.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -44,20 +44,19 @@ def multi_diffusion_denoise(
4444
scheduler_step_kwargs: dict[str, Any],
4545
noise: Optional[torch.Tensor],
4646
timesteps: torch.Tensor,
47+
init_timestep: torch.Tensor,
4748
callback: Callable[[PipelineIntermediateState], None],
4849
) -> torch.Tensor:
4950
self._check_regional_prompting(multi_diffusion_conditioning)
5051

51-
if timesteps.shape[0] == 0:
52+
if init_timestep.shape[0] == 0:
5253
return latents
5354

5455
batch_size, _, latent_height, latent_width = latents.shape
56+
batched_init_timestep = init_timestep.expand(batch_size)
5557

5658
# noise can be None if the latents have already been noised (e.g. when running the SDXL refiner).
5759
if noise is not None:
58-
# batched_init_timestep should have shape (batch_size, 1).
59-
batched_init_timestep = timesteps[0:1].expand(batch_size)
60-
6160
# TODO(ryand): I'm pretty sure we should be applying init_noise_sigma in cases where we are starting with
6261
# full noise. Investigate the history of why this got commented out.
6362
# latents = noise * self.scheduler.init_noise_sigma # it's like in t2l according to diffusers

0 commit comments

Comments
 (0)