update onnx pipelines.

This commit is contained in:
Seunghoon Lee
2024-02-05 22:11:24 +09:00
parent 8049053462
commit ae113cae1e
4 changed files with 12 additions and 5 deletions
@@ -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
@@ -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
@@ -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)
+2 -1
View File
@@ -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: