From ae113cae1e1bcbfc7815f62554e29f00719cd6d5 Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Mon, 5 Feb 2024 22:11:24 +0900 Subject: [PATCH] update onnx pipelines. --- .../pipelines/onnx_stable_diffusion_img2img_pipeline.py | 3 ++- .../onnx_impl/pipelines/onnx_stable_diffusion_pipeline.py | 3 ++- .../pipelines/onnx_stable_diffusion_upscale_pipeline.py | 8 ++++++-- modules/onnx_impl/pipelines/utils.py | 3 ++- 4 files changed, 12 insertions(+), 5 deletions(-) diff --git a/modules/onnx_impl/pipelines/onnx_stable_diffusion_img2img_pipeline.py b/modules/onnx_impl/pipelines/onnx_stable_diffusion_img2img_pipeline.py index a8edb4f45..6d8ea5946 100644 --- a/modules/onnx_impl/pipelines/onnx_stable_diffusion_img2img_pipeline.py +++ b/modules/onnx_impl/pipelines/onnx_stable_diffusion_img2img_pipeline.py @@ -157,7 +157,8 @@ class OnnxStableDiffusionImg2ImgPipeline(diffusers.OnnxStableDiffusionImg2ImgPip # call the callback, if provided if callback is not None and i % callback_steps == 0: - callback(i, t, latents) + step_idx = i // getattr(self.scheduler, "order", 1) + callback(step_idx, t, latents) has_nsfw_concept = None diff --git a/modules/onnx_impl/pipelines/onnx_stable_diffusion_pipeline.py b/modules/onnx_impl/pipelines/onnx_stable_diffusion_pipeline.py index 0985142f0..254abf72c 100644 --- a/modules/onnx_impl/pipelines/onnx_stable_diffusion_pipeline.py +++ b/modules/onnx_impl/pipelines/onnx_stable_diffusion_pipeline.py @@ -128,7 +128,8 @@ class OnnxStableDiffusionPipeline(diffusers.OnnxStableDiffusionPipeline, Callabl # call the callback, if provided if callback is not None and i % callback_steps == 0: - callback(i, t, latents) + step_idx = i // getattr(self.scheduler, "order", 1) + callback(step_idx, t, latents) has_nsfw_concept = None diff --git a/modules/onnx_impl/pipelines/onnx_stable_diffusion_upscale_pipeline.py b/modules/onnx_impl/pipelines/onnx_stable_diffusion_upscale_pipeline.py index 8b7aad6f1..ae7ef4fd4 100644 --- a/modules/onnx_impl/pipelines/onnx_stable_diffusion_upscale_pipeline.py +++ b/modules/onnx_impl/pipelines/onnx_stable_diffusion_upscale_pipeline.py @@ -8,7 +8,7 @@ from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion_upscale import preprocess from diffusers.image_processor import PipelineImageInput from modules.onnx_impl.pipelines import CallablePipelineBase -from modules.onnx_impl.pipelines.utils import prepare_latents +from modules.onnx_impl.pipelines.utils import prepare_latents, randn_tensor class OnnxStableDiffusionUpscalePipeline(diffusers.OnnxStableDiffusionUpscalePipeline, CallablePipelineBase): @@ -102,7 +102,11 @@ class OnnxStableDiffusionUpscalePipeline(diffusers.OnnxStableDiffusionUpscalePip # 5. Add noise to image noise_level = np.array([noise_level]).astype(np.int64) - noise = generator.randn(*image.shape).astype(latents_dtype) + noise = randn_tensor( + image.shape, + latents_dtype, + generator, + ) image = self.low_res_scheduler.add_noise( torch.from_numpy(image), torch.from_numpy(noise), torch.from_numpy(noise_level) diff --git a/modules/onnx_impl/pipelines/utils.py b/modules/onnx_impl/pipelines/utils.py index 67b82aee0..273634689 100644 --- a/modules/onnx_impl/pipelines/utils.py +++ b/modules/onnx_impl/pipelines/utils.py @@ -2,6 +2,7 @@ from typing import Union, List import numpy as np import torch + def extract_generator_seed(generator: Union[torch.Generator, List[torch.Generator]]) -> List[int]: if isinstance(generator, list): generator = [g.seed() for g in generator] @@ -10,7 +11,7 @@ def extract_generator_seed(generator: Union[torch.Generator, List[torch.Generato return generator -def randn_tensor(shape, dtype, generator: Union[torch.Generator, List[torch.Generator], int, List[int]]): +def randn_tensor(shape, dtype: np.dtype, generator: Union[torch.Generator, List[torch.Generator], int, List[int]]): if hasattr(generator, "seed") or (isinstance(generator, list) and hasattr(generator[0], "seed")): generator = extract_generator_seed(generator) if len(generator) == 1: