Skip to content

Commit dc62e69

Browse files
authored
[fix bug] PixArt inference_steps=1 (#11079)
* fix bug when pixart-dmd inference with `num_inference_steps=1` * use return_dict=False and return [1] element for 1-step pixart model, which works for both lcm and dmd
1 parent 56f7400 commit dc62e69

File tree

1 file changed

+1
-2
lines changed

1 file changed

+1
-2
lines changed

src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -941,8 +941,7 @@ def __call__(
941941

942942
# compute previous image: x_t -> x_t-1
943943
if num_inference_steps == 1:
944-
# For DMD one step sampling: https://arxiv.org/abs/2311.18828
945-
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).pred_original_sample
944+
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[1]
946945
else:
947946
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
948947

0 commit comments

Comments
 (0)