UniPC and DPM-Solver with use_flow_sigmas condition the model on integer timesteps while integrating with float sigmas #15167

Describe the bug

With use_flow_sigmas=True, UniPCMultistepScheduler and DPMSolverMultistepScheduler store timesteps as torch.int64, while the update uses the exact float sigmas. Pipelines pass timesteps to the transformer, so the network is conditioned on floor(1000 * sigma) / 1000, not on the sigma the solver integrates with. In effect the sampler solves a slightly different ODE.

FlowMatchEulerDiscreteScheduler keeps timesteps as floats and does not have this mismatch. Flow-matching transformers such as FLUX.1-dev, SD3.5 and Wan2.1 take a continuous time as input.

The difference does not go away with more steps, so it does not show up in a refinement check. Mean PSNR of the final image against Euler with 1,600 steps (FlowMatchEulerDiscreteScheduler, fp32 solver state, seed 0, the same noise for all solvers, first prompts of DrawBench):

On FLUX.1-dev, stock UniPC with 200 and 400 steps agree with each other to 36.09 dB, so they look converged while being about 26 dB from the Euler solution.

Relevant lines on main: self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64) at the end of set_timesteps in scheduling_unipc_multistep.py and scheduling_dpmsolver_multistep.py.

A possible fix: keep the timesteps as floats (sigmas[:-1] * num_train_timesteps) when use_flow_sigmas=True, as FlowMatchEulerDiscreteScheduler does. Models that really are trained on discrete timesteps would not be affected, since this applies only to the flow-sigma path. I measured the effect only on the two models above, and with UniPC only; I have not measured DPM-Solver. If the integer behavior is intended (for example, to match an original implementation), it would help to document it.

Reproduction

No model is needed to see the mismatch:

from diffusers import UniPCMultistepScheduler, DPMSolverMultistepScheduler, FlowMatchEulerDiscreteScheduler

kw = dict(prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=3.0)
for name, s in (("UniPC", UniPCMultistepScheduler(**kw)), ("DPM++", DPMSolverMultistepScheduler(**kw))):
    s.set_timesteps(8)
    print(name, s.timesteps.dtype, s.timesteps.tolist())
    print("  1000*sigma", [round(x, 3) for x in (s.sigmas[:-1] * 1000).tolist()])

e = FlowMatchEulerDiscreteScheduler(shift=3.0)
e.set_timesteps(8)
print("Euler", e.timesteps.dtype, [round(x, 3) for x in e.timesteps.tolist()])
print("  1000*sigma", [round(x, 3) for x in (e.sigmas[:-1] * 1000).tolist()])

Logs

UniPC torch.int64 [999, 954, 900, 833, 750, 643, 500, 301]
  1000*sigma [999.999, 954.595, 900.12, 833.555, 750.375, 643.469, 500.999, 301.678]
DPM++ torch.int64 [999, 954, 899, 832, 749, 642, 499, 299]
  1000*sigma [999.666, 954.198, 899.64, 832.963, 749.625, 642.49, 499.667, 299.76]
Euler torch.float32 [1000.0, 947.543, 882.788, 800.837, 693.793, 548.046, 337.972, 8.929]
  1000*sigma [1000.0, 947.543, 882.788, 800.837, 693.793, 548.046, 337.972, 8.929]

System Info

diffusers 0.40.0, torch 2.11.0+cu128, Python 3.11, Linux. The two lines above are also on main as of 2026-11-02.