Skip to content

Commit ae78a65

Browse files
committed
use return_dict=False and return [1] element for 1-step pixart model, which works for both lcm and dmd
1 parent 2dae26c commit ae78a65

File tree

1 file changed

+1
-6
lines changed

1 file changed

+1
-6
lines changed

src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py

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

942942
# compute previous image: x_t -> x_t-1
943943
if num_inference_steps == 1:
944-
try:
945-
# For LCM one step sampling
946-
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).denoised
947-
except:
948-
# For DMD one step sampling: https://arxiv.org/abs/2311.18828
949-
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]
950945
else:
951946
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
952947

0 commit comments

Comments
 (0)