PNDMScheduler.set_timesteps crashes with NumPy broadcasting error for fewer than 4 inference steps #14953

Describe the bug

PNDMScheduler.set_timesteps() raises an internal NumPy broadcasting error when num_inference_steps is 1, 2, or 3 with the default skip_prk_steps=False.

The error occurs during construction of the PRK timesteps. The scheduler uses self.pndm_order = 4, but for fewer than four inference steps, the following expression produces arrays with incompatible shapes:

prk_timesteps = np.array(self._timesteps[-self.pndm_order :]).repeat(2) + np.tile(
    np.array([0, self.config.num_train_timesteps // num_inference_steps // 2]),
    self.pndm_order,
)

For num_inference_steps=3, the left-hand array has shape (6,), whereas the right-hand array has shape (8,), resulting in a ValueError.

Observed behavior:

• num_inference_steps=1: broadcasting error, shapes (2,) and (8,)
• num_inference_steps=2: broadcasting error, shapes (4,) and (8,)
• num_inference_steps=3: broadcasting error, shapes (6,) and (8,)
• num_inference_steps=4: succeeds

This behavior is reproducible with all three timestep-spacing modes (leading, linspace, and trailing). The same low-step configurations succeed when skip_prk_steps=True.

Expected behavior: If PRK requires at least four inference steps, set_timesteps() should raise a descriptive validation error rather than exposing an internal NumPy broadcasting error. If fewer steps are intended to be supported, the PRK timestep construction should handle those cases correctly.

A possible validation is:

if not self.config.skip_prk_steps and num_inference_steps < self.pndm_order:
    raise ValueError(
        f"`num_inference_steps` must be >= {self.pndm_order} when "
        f"`skip_prk_steps=False`, got {num_inference_steps}."
    )

I searched existing Diffusers issues and PRs for this specific failure and did not find an exact duplicate.

Reproduction

Minimal reproduction:

from diffusers import PNDMScheduler

scheduler = PNDMScheduler()
scheduler.set_timesteps(3)

Reproduction across inference-step counts:

from diffusers import PNDMScheduler

for n in [1, 2, 3, 4]:
    try:
        scheduler = PNDMScheduler()
        scheduler.set_timesteps(n)
        print(f"n={n}: OK")
    except Exception as e:
        print(f"n={n}: {type(e).__name__}: {e}")

Output:

n=1: ValueError: operands could not be broadcast together with shapes (2,) (8,)
n=2: ValueError: operands could not be broadcast together with shapes (4,) (8,)
n=3: ValueError: operands could not be broadcast together with shapes (6,) (8,)
n=4: OK

Control case (PRK disabled):

from diffusers import PNDMScheduler

for n in [1, 2, 3]:
    scheduler = PNDMScheduler(skip_prk_steps=True)
    scheduler.set_timesteps(n)
    print(n, scheduler.timesteps.tolist())

The control case succeeds for all three values.

Existing test coverage:

python -m pytest tests/schedulers/test_scheduler_pndm.py -q

Result:

32 passed, 1 skipped

The existing PNDM scheduler tests pass, but do not appear to cover the low-step PRK failure.

Logs

Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "E:\diffusers\src\diffusers\schedulers\scheduling_pndm.py", line 217, in set_timesteps
    prk_timesteps = np.array(self._timesteps[-self.pndm_order :]).repeat(2) + np.tile(
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: operands could not be broadcast together with shapes (6,) (8,)

System Info

• 🤗 Diffusers version: 0.41.0.dev0
• Platform: Windows-10-10.0.22631-SP0
• Running on Google Colab?: No
• Python version: 3.11.16
• PyTorch version (GPU?): 2.14.1+cpu (False)
• Huggingface_hub version: 1.33.0
• Transformers version: 5.18.0
• Accelerate version: 1.15.0
• PEFT version: 0.21.2
• Safetensors version: 0.8.0
• xFormers version: not installed
• Accelerator: NA
• Using GPU in script?: No
• Using distributed or parallel set-up in script?: No

Who can help?

@yiyixuxu @sayakpaul @DN6

I am open to make a PR if allowed and solve this issue