|
10 | 10 | from ...configuration_utils import FrozenDict |
11 | 11 | from ...schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler |
12 | 12 | from ...utils import deprecate, logging |
13 | | -from ..onnx_utils import OnnxRuntimeModel |
| 13 | +from ..onnx_utils import ORT_TO_NP_TYPE, OnnxRuntimeModel |
14 | 14 | from ..pipeline_utils import DiffusionPipeline |
15 | 15 | from . import StableDiffusionPipelineOutput |
16 | 16 |
|
@@ -391,16 +391,21 @@ def __call__( |
391 | 391 |
|
392 | 392 | t_start = max(num_inference_steps - init_timestep + offset, 0) |
393 | 393 | timesteps = self.scheduler.timesteps[t_start:].numpy() |
| 394 | + timestep_dtype = next( |
| 395 | + (input.type for input in self.unet.model.get_inputs() if input.name == "timestep"), "tensor(float)" |
| 396 | + ) |
| 397 | + timestep_dtype = ORT_TO_NP_TYPE[timestep_dtype] |
394 | 398 |
|
395 | 399 | for i, t in enumerate(self.progress_bar(timesteps)): |
396 | 400 | # expand the latents if we are doing classifier free guidance |
397 | 401 | latent_model_input = np.concatenate([latents] * 2) if do_classifier_free_guidance else latents |
398 | 402 | latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) |
399 | 403 |
|
400 | 404 | # predict the noise residual |
401 | | - noise_pred = self.unet( |
402 | | - sample=latent_model_input, timestep=np.array([t]), encoder_hidden_states=prompt_embeds |
403 | | - )[0] |
| 405 | + timestep = np.array([t], dtype=timestep_dtype) |
| 406 | + noise_pred = self.unet(sample=latent_model_input, timestep=timestep, encoder_hidden_states=prompt_embeds)[ |
| 407 | + 0 |
| 408 | + ] |
404 | 409 |
|
405 | 410 | # perform guidance |
406 | 411 | if do_classifier_free_guidance: |
|
0 commit comments