diff --git a/modules/onnx_impl/__init__.py b/modules/onnx_impl/__init__.py index a38645391..7c64aa675 100644 --- a/modules/onnx_impl/__init__.py +++ b/modules/onnx_impl/__init__.py @@ -4,6 +4,7 @@ import numpy as np import torch import diffusers import onnxruntime as ort +import optimum.onnxruntime initialized = False @@ -51,7 +52,7 @@ class TorchCompatibleModule: dtype = torch.float32 def to(self, *_, **__): - return self + raise NotImplementedError def type(self, *_, **__): return self @@ -122,6 +123,10 @@ class VAE(TorchCompatibleModule): def config(self): return VAEConfig(self.pipeline.vae_decoder.config) + @property + def device(self): + return self.pipeline.vae_decoder.device + def encode(self, latent_sample: torch.Tensor, return_dict: bool): # pylint: disable=unused-argument latents_np = latent_sample.cpu().numpy() return [ @@ -183,6 +188,11 @@ def preprocess_pipeline(p, refiner_enabled: bool): return shared.sd_model +def ORTDiffusionModelPart_to(self, *args, **kwargs): + self.parent_model = self.parent_model.to(*args, **kwargs) + return self + + def initialize(): global initialized # pylint: disable=global-statement @@ -232,6 +242,8 @@ def initialize(): diffusers.ORTStableDiffusionXLPipeline = diffusers.OnnxStableDiffusionXLPipeline diffusers.ORTStableDiffusionXLImg2ImgPipeline = diffusers.OnnxStableDiffusionXLImg2ImgPipeline + optimum.onnxruntime.modeling_diffusion._ORTDiffusionModelPart.to = ORTDiffusionModelPart_to # pylint: disable=protected-access + initialized = True diff --git a/modules/onnx_impl/pipelines/__init__.py b/modules/onnx_impl/pipelines/__init__.py index b5e68d2ef..919284b2d 100644 --- a/modules/onnx_impl/pipelines/__init__.py +++ b/modules/onnx_impl/pipelines/__init__.py @@ -36,8 +36,8 @@ class PipelineBase(TorchCompatibleModule, diffusers.DiffusionPipeline, metaclass self.model_type = self.__class__.__name__ def to(self, *args, **kwargs): - if self.__class__ == OnnxRawPipeline: - return super().to(*args, **kwargs) + if self.__class__ == OnnxRawPipeline: # cannot move pipeline which is not preprocessed. + return self expected_modules, _ = self._get_signature_keys(self) for name in expected_modules: diff --git a/modules/onnx_impl/pipelines/onnx_stable_diffusion_xl_img2img_pipeline.py b/modules/onnx_impl/pipelines/onnx_stable_diffusion_xl_img2img_pipeline.py index bb37c0466..5606da06e 100644 --- a/modules/onnx_impl/pipelines/onnx_stable_diffusion_xl_img2img_pipeline.py +++ b/modules/onnx_impl/pipelines/onnx_stable_diffusion_xl_img2img_pipeline.py @@ -2,7 +2,7 @@ from typing import Optional, Dict, Any import numpy as np import torch import onnxruntime as ort -import optimum +import optimum.onnxruntime from modules.onnx_impl.pipelines import CallablePipelineBase from modules.onnx_impl.pipelines.utils import randn_tensor @@ -27,7 +27,9 @@ class OnnxStableDiffusionXLImg2ImgPipeline(CallablePipelineBase, optimum.onnxrun model_save_dir = None, add_watermarker: Optional[bool] = None ): - super(optimum.onnxruntime.ORTStableDiffusionXLImg2ImgPipeline, self).__init__(vae_decoder, text_encoder, unet, config, tokenizer, scheduler, feature_extractor, vae_encoder, text_encoder_2, tokenizer_2, use_io_binding, model_save_dir, add_watermarker) + optimum.onnxruntime.ORTStableDiffusionXLImg2ImgPipeline.__init__(self, vae_decoder, text_encoder, unet, config, tokenizer, scheduler, feature_extractor, vae_encoder, text_encoder_2, tokenizer_2, use_io_binding, model_save_dir, add_watermarker) + super().__init__() + del self.image_processor # This image processor requires np array. In order to share same workflow with non-XL pipelines, delete it. def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, generator=None): batch_size = batch_size * num_images_per_prompt diff --git a/modules/onnx_impl/pipelines/onnx_stable_diffusion_xl_pipeline.py b/modules/onnx_impl/pipelines/onnx_stable_diffusion_xl_pipeline.py index 72656704f..452e4f892 100644 --- a/modules/onnx_impl/pipelines/onnx_stable_diffusion_xl_pipeline.py +++ b/modules/onnx_impl/pipelines/onnx_stable_diffusion_xl_pipeline.py @@ -1,6 +1,6 @@ from typing import Optional, Dict, Any import onnxruntime as ort -import optimum +import optimum.onnxruntime from modules.onnx_impl.pipelines import CallablePipelineBase from modules.onnx_impl.pipelines.utils import prepare_latents @@ -25,7 +25,9 @@ class OnnxStableDiffusionXLPipeline(CallablePipelineBase, optimum.onnxruntime.OR model_save_dir = None, add_watermarker: Optional[bool] = None ): - super(optimum.onnxruntime.ORTStableDiffusionXLPipeline, self).__init__(vae_decoder, text_encoder, unet, config, tokenizer, scheduler, feature_extractor, vae_encoder, text_encoder_2, tokenizer_2, use_io_binding, model_save_dir, add_watermarker) + optimum.onnxruntime.ORTStableDiffusionXLPipeline.__init__(self, vae_decoder, text_encoder, unet, config, tokenizer, scheduler, feature_extractor, vae_encoder, text_encoder_2, tokenizer_2, use_io_binding, model_save_dir, add_watermarker) + super().__init__() + del self.image_processor # This image processor requires np array. In order to share same workflow with non-XL pipelines, delete it. def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, generator, latents=None): return prepare_latents(self.scheduler.init_noise_sigma, batch_size, height, width, dtype, generator, latents, num_channels_latents, self.vae_scale_factor) diff --git a/modules/sd_models.py b/modules/sd_models.py index 9e4937690..aa7705a63 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -798,7 +798,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No shared.log.error(f'Diffusers Failed loading {op}: {checkpoint_info.path} {e}') return elif 'ONNX' in model_type: # forced pipeline - sd_model = pipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) + sd_model = pipeline.from_pretrained(checkpoint_info.path) else: err1, err2, err3 = None, None, None try: # 1 - autopipeline, best choice but not all pipelines are available