onnx sdxl vae

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
 
 

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:

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

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)

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