diff --git a/src/diffusers/schedulers/scheduling_ddim.py b/src/diffusers/schedulers/scheduling_ddim.py index 14356eafda..13c9b3b4a5 100644 --- a/src/diffusers/schedulers/scheduling_ddim.py +++ b/src/diffusers/schedulers/scheduling_ddim.py @@ -463,7 +463,10 @@ class DDIMScheduler(SchedulerMixin, ConfigMixin): prev_sample = prev_sample + variance if not return_dict: - return (prev_sample,) + return ( + prev_sample, + pred_original_sample, + ) return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) diff --git a/src/diffusers/schedulers/scheduling_ddim_cogvideox.py b/src/diffusers/schedulers/scheduling_ddim_cogvideox.py index ec5c5f3e1c..5c13175293 100644 --- a/src/diffusers/schedulers/scheduling_ddim_cogvideox.py +++ b/src/diffusers/schedulers/scheduling_ddim_cogvideox.py @@ -394,7 +394,10 @@ class CogVideoXDDIMScheduler(SchedulerMixin, ConfigMixin): prev_sample = a_t * sample + b_t * pred_original_sample if not return_dict: - return (prev_sample,) + return ( + prev_sample, + pred_original_sample, + ) return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) diff --git a/src/diffusers/schedulers/scheduling_ddim_parallel.py b/src/diffusers/schedulers/scheduling_ddim_parallel.py index 0cf84b694d..64412709ae 100644 --- a/src/diffusers/schedulers/scheduling_ddim_parallel.py +++ b/src/diffusers/schedulers/scheduling_ddim_parallel.py @@ -480,7 +480,10 @@ class DDIMParallelScheduler(SchedulerMixin, ConfigMixin): prev_sample = prev_sample + variance if not return_dict: - return (prev_sample,) + return ( + prev_sample, + pred_original_sample, + ) return DDIMParallelSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) diff --git a/src/diffusers/schedulers/scheduling_ddpm.py b/src/diffusers/schedulers/scheduling_ddpm.py index 81a770edf6..468fdf61a9 100644 --- a/src/diffusers/schedulers/scheduling_ddpm.py +++ b/src/diffusers/schedulers/scheduling_ddpm.py @@ -492,7 +492,10 @@ class DDPMScheduler(SchedulerMixin, ConfigMixin): pred_prev_sample = pred_prev_sample + variance if not return_dict: - return (pred_prev_sample,) + return ( + pred_prev_sample, + pred_original_sample, + ) return DDPMSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample) diff --git a/src/diffusers/schedulers/scheduling_ddpm_parallel.py b/src/diffusers/schedulers/scheduling_ddpm_parallel.py index 5dfcf3c17a..f377ee6e8c 100644 --- a/src/diffusers/schedulers/scheduling_ddpm_parallel.py +++ b/src/diffusers/schedulers/scheduling_ddpm_parallel.py @@ -500,7 +500,10 @@ class DDPMParallelScheduler(SchedulerMixin, ConfigMixin): pred_prev_sample = pred_prev_sample + variance if not return_dict: - return (pred_prev_sample,) + return ( + pred_prev_sample, + pred_original_sample, + ) return DDPMParallelSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample) diff --git a/src/diffusers/schedulers/scheduling_edm_euler.py b/src/diffusers/schedulers/scheduling_edm_euler.py index 4b823c0d28..04a09b114d 100644 --- a/src/diffusers/schedulers/scheduling_edm_euler.py +++ b/src/diffusers/schedulers/scheduling_edm_euler.py @@ -360,7 +360,10 @@ class EDMEulerScheduler(SchedulerMixin, ConfigMixin): self._step_index += 1 if not return_dict: - return (prev_sample,) + return ( + prev_sample, + pred_original_sample, + ) return EDMEulerSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) diff --git a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py index 485e919e9c..4df43a160c 100644 --- a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py @@ -435,7 +435,10 @@ class EulerAncestralDiscreteScheduler(SchedulerMixin, ConfigMixin): self._step_index += 1 if not return_dict: - return (prev_sample,) + return ( + prev_sample, + pred_original_sample, + ) return EulerAncestralDiscreteSchedulerOutput( prev_sample=prev_sample, pred_original_sample=pred_original_sample diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index 5c39583356..7083571a60 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -677,7 +677,10 @@ class EulerDiscreteScheduler(SchedulerMixin, ConfigMixin): self._step_index += 1 if not return_dict: - return (prev_sample,) + return ( + prev_sample, + pred_original_sample, + ) return EulerDiscreteSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) diff --git a/src/diffusers/schedulers/scheduling_lms_discrete.py b/src/diffusers/schedulers/scheduling_lms_discrete.py index aed8c5828c..0a09004554 100644 --- a/src/diffusers/schedulers/scheduling_lms_discrete.py +++ b/src/diffusers/schedulers/scheduling_lms_discrete.py @@ -507,7 +507,10 @@ class LMSDiscreteScheduler(SchedulerMixin, ConfigMixin): self._step_index += 1 if not return_dict: - return (prev_sample,) + return ( + prev_sample, + pred_original_sample, + ) return LMSDiscreteSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) diff --git a/src/diffusers/schedulers/scheduling_unclip.py b/src/diffusers/schedulers/scheduling_unclip.py index 6e1580290f..22a53b0e73 100644 --- a/src/diffusers/schedulers/scheduling_unclip.py +++ b/src/diffusers/schedulers/scheduling_unclip.py @@ -320,7 +320,10 @@ class UnCLIPScheduler(SchedulerMixin, ConfigMixin): pred_prev_sample = pred_prev_sample + variance if not return_dict: - return (pred_prev_sample,) + return ( + pred_prev_sample, + pred_original_sample, + ) return UnCLIPSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample)