Schedulers not compatible with OnnxStableDiffusionPipeline: TypeError: unsupported operand

Describe the bug

Hi, I tried to use different schedulers with OnnxStableDiffusionPipeline, but it throw errors. Schedulers are not compatible with numpy used in onnx pipeline.

Onnx checkpoints converted with: convert_stable_diffusion_checkpoint_to_onnx.py

I have found a solution, but it is probably not optimal, because torch usage inside pipeline call.

Solution works with:

• CompVis/stable-diffusion-v1-4
• hakurei/waifu-diffusion

First error:

File "C:\...\diffusers\pipelines\stable_diffusion\pipeline_onnx_stable_diffusion.py", line 152, in __call__
  latents = latents * self.scheduler.init_noise_sigma
TypeError: unsupported operand type(s) for *: 'numpy.ndarray' and 'Tensor'

If I cast latents to torch.tensor, before init_noise_sigma:

latents = torch.tensor(latents)
latents = latents * self.scheduler.init_noise_sigma

File "C:\...\diffusers\pipelines\stable_diffusion\pipeline_onnx_stable_diffusion.py", line 171, in __call__
  noise_pred = self.unet(
File "C:\...\diffusers\onnx_utils.py", line 46, in __call__
  return self.model.run(None, inputs)
File "C:\...\onnxruntime\capi\onnxruntime_inference_collection.py", line 200, in run
  return self._sess.run(output_names, input_feed, run_options)
onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Unexpected input data type. Actual: (tensor(double)) , expected: (tensor(int64))

If I add dtype=np.int64 to timestep in unet args:

# predict the noise residual
noise_pred = self.unet(
  sample=latent_model_input,
  timestep=np.array([t], dtype=np.int64),
  encoder_hidden_states=text_embeddings,
)

File "C:\...\diffusers\pipelines\stable_diffusion\pipeline_onnx_stable_diffusion.py", line 184, in __call__
  latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
File "C:\...\diffusers\schedulers\scheduling_lms_discrete.py", line 224, in step
  pred_original_sample = sample - sigma * model_output
TypeError: unsupported operand type(s) for -: 'numpy.ndarray' and 'Tensor'

And if I cast latents to torch.tensor, before passing them into scheduler.step, it is working again:

# compute the previous noisy sample x_t -> x_t-1
latents = torch.tensor(latents)
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
latents = np.array(latents)

Reproduction

from diffusers import OnnxStableDiffusionPipeline, LMSDiscreteScheduler

lms = LMSDiscreteScheduler()

pipe = OnnxStableDiffusionPipeline.from_pretrained(
    model_path,
    provider="DmlExecutionProvider",
    scheduler=lms,
    local_files_only=True,
)

image = pipe("prompt")[0]

Logs

No response

System Info

• diffusers version: 0.6.0

• Platform: Windows-10-10.0.19044-SP0

• Python version: 3.10.7

• PyTorch version (GPU?): 1.12.1+cpu (False)

• Huggingface_hub version: 0.10.1

• Transformers version: 4.23.1

• Using GPU in script?: DmlExecutionProvider

• Using distributed or parallel set-up in script?: <fill in>

• GPU: AMD RX 6900 XT 16GB

[anton-l]

@averad https://github.com/huggingface/diffusers/pull/1173 should fix the scheduler issues once merged

[averad]

@anton-l you and your team members are the best, thank you.