PNDMScheduler Does Not Compatible with DDPMPipeline

[KeepNoob]

Describe the bug

I was trying to test different schedulers under DDPMPipeline. And an error occurred if I use PNDMScheduler beforehand I have found that PNDMScheduler should be compatible with DDPMPipeline following the official tutorial.

pipeline = DDPMPipeline(unet=model, scheduler=noise_scheduler)
pipeline.scheduler.compatibles

And the output is this:

[diffusers.schedulers.scheduling_ddim.DDIMScheduler,
 diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler,
 diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler,
 diffusers.schedulers.scheduling_lms_discrete.LMSDiscreteScheduler,
 diffusers.schedulers.scheduling_dpmsolver_singlestep.DPMSolverSinglestepScheduler,
 diffusers.schedulers.scheduling_deis_multistep.DEISMultistepScheduler,
 diffusers.utils.dummy_torch_and_torchsde_objects.DPMSolverSDEScheduler,
 diffusers.schedulers.scheduling_k_dpm_2_discrete.KDPM2DiscreteScheduler,
 diffusers.schedulers.scheduling_k_dpm_2_ancestral_discrete.KDPM2AncestralDiscreteScheduler,
 diffusers.schedulers.scheduling_unipc_multistep.UniPCMultistepScheduler,
 diffusers.schedulers.scheduling_heun_discrete.HeunDiscreteScheduler,
 diffusers.schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteScheduler,
 diffusers.schedulers.scheduling_pndm.PNDMScheduler,
 diffusers.schedulers.scheduling_ddpm.DDPMScheduler]

And the main error massage is that

File [c:\Users\User\miniconda3\envs\Pytorch\lib\site-packages\diffusers\pipelines\ddpm\pipeline_ddpm.py:117](file:///C:/Users/User/miniconda3/envs/Pytorch/lib/site-packages/diffusers/pipelines/ddpm/pipeline_ddpm.py:117), in DDPMPipeline.__call__(self, batch_size, generator, num_inference_steps, output_type, return_dict)
    [114](file:///C:/Users/User/miniconda3/envs/Pytorch/lib/site-packages/diffusers/pipelines/ddpm/pipeline_ddpm.py:114)     model_output = self.unet(image, t).sample
    [116](file:///C:/Users/User/miniconda3/envs/Pytorch/lib/site-packages/diffusers/pipelines/ddpm/pipeline_ddpm.py:116)     # 2. compute previous image: x_t -> x_t-1
--> [117](file:///C:/Users/User/miniconda3/envs/Pytorch/lib/site-packages/diffusers/pipelines/ddpm/pipeline_ddpm.py:117)     image = self.scheduler.step(model_output, t, image, generator=generator).prev_sample
    [119](file:///C:/Users/User/miniconda3/envs/Pytorch/lib/site-packages/diffusers/pipelines/ddpm/pipeline_ddpm.py:119) image = (image / 2 + 0.5).clamp(0, 1)
    [120](file:///C:/Users/User/miniconda3/envs/Pytorch/lib/site-packages/diffusers/pipelines/ddpm/pipeline_ddpm.py:120) image = image.cpu().permute(0, 2, 3, 1).numpy()
TypeError: PNDMScheduler.step() got an unexpected keyword argument 'generator'

In DDPMPipeline.scheduler.step( ) function, it takes generator as the argument. But in class diffusers.PNDMScheduler step( ) function shown in official doc, the function only takes model_output (torch.FloatTensor), timestep (int), sample (torch.FloatTensor), return_dict (bool).
Moreover, I also find out that HeunDiscreteScheduler has the same problem

Reproduction

from diffusers import DDIMScheduler
from diffusers import UNet2DModel
from diffusers import PNDMScheduler
model = UNet2DModel(
    sample_size=config.image_size,
    in_channels=1,
    out_channels=1,
    layers_per_block=2,
    block_out_channels=(128,128,256,512),
    down_block_types=(
        "DownBlock2D",
        "DownBlock2D",
        "AttnDownBlock2D",
        "DownBlock2D",
    ),
    up_block_types=(
        "UpBlock2D",
        "AttnUpBlock2D",
        "UpBlock2D",
        "UpBlock2D",
    ),
)
noise_scheduler = PNDMScheduler()
pipeline = DDPMPipeline(unet=model, scheduler=noise_scheduler)
print(pipeline.scheduler.compatibles)
images = pipeline(
        batch_size=1,
        num_inference_steps = 50).images 

Logs

TypeError                                 Traceback (most recent call last)
Input In [58], in <cell line: 26>()
     24 pipeline = DDPMPipeline(unet=model, scheduler=noise_scheduler)
     25 pipeline.scheduler.compatibles
---> 26 images = pipeline(
     27         batch_size=1,
     28         num_inference_steps = 50).images

File c:\Users\User\miniconda3\envs\Pytorch\lib\site-packages\torch\utils\_contextlib.py:115, in context_decorator.<locals>.decorate_context(*args, **kwargs)
    112 @functools.wraps(func)
    113 def decorate_context(*args, **kwargs):
    114     with ctx_factory():
--> 115         return func(*args, **kwargs)

File c:\Users\User\miniconda3\envs\Pytorch\lib\site-packages\diffusers\pipelines\ddpm\pipeline_ddpm.py:117, in DDPMPipeline.__call__(self, batch_size, generator, num_inference_steps, output_type, return_dict)
    114     model_output = self.unet(image, t).sample
    116     # 2. compute previous image: x_t -> x_t-1
--> 117     image = self.scheduler.step(model_output, t, image, generator=generator).prev_sample
    119 image = (image / 2 + 0.5).clamp(0, 1)
    120 image = image.cpu().permute(0, 2, 3, 1).numpy()

TypeError: PNDMScheduler.step() got an unexpected keyword argument 'generator'

System Info

I cannot use this command but I can show some info.
diffusers 0.26.3 pyhd8ed1ab_0 conda-forge
pytorch 2.0.1 py3.10_cuda11.8_cudnn8_0 pytorch
OS: Windows 11

Who can help?

@yiyixuxu @dn

[sayakpaul]

@yiyixuxu could you give this a look?

[yiyixuxu]

hey thanks for the issue!
normally we inspect the signature of the scheduler and only pass generator if it would be accepted e.g. stable diffusion pipeline here.

diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

Line 580 in 80ff4ba

accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())

do you want to open a PR to add this method to DDPMPipeline?

[KeepNoob]

@yiyixuxu I have opened a PR, could you give this a look?

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[sayakpaul]

I think you opened the PR to the wrong repository. Could you please check? @KeepNoob

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.