update olive by #2429

launch.py

@@ -29,6 +29,15 @@ except ModuleNotFoundError:
     sys.modules["torch._dynamo"] = {} # HACK torch 1.13.1 does not have _dynamo. will be removed.
 
 
+def init_olive():
+    try:
+        if installer.opts['onnx_enable_olive']:
+            import olive.workflows # pylint: disable=unused-import
+            installer.log.debug('Load olive')
+    except Exception as e:
+        installer.log.error(f'Failed to load olive: {e}')
+
+
 def init_args():
     global parser, args # pylint: disable=global-statement
     import modules.cmd_args
@@ -200,6 +209,7 @@ if __name__ == "__main__":
     except Exception:
         pass
     installer.read_options()
+    init_olive()
     if args.skip_all:
         args.quick = True
     installer.check_python()

modules/olive.py

@@ -9,27 +9,24 @@ from installer import log
 from modules import shared
 from modules.paths import sd_configs_path
 from modules.sd_models import CheckpointInfo
-from modules.onnx import ExecutionProvider, OnnxStableDiffusionPipeline, get_execution_provider_options
+from modules.onnx import ExecutionProvider, get_execution_provider_options
 
 is_available = "olive" in sys.modules # Olive is not available if it is not loaded at startup.
 
 def enable_olive_onchange():
+    from installer import installed, install, uninstall
     if shared.opts.onnx_enable_olive:
-        if "olive" in sys.modules:
-            log.info("You already have Olive installed. No additional installation is required.")
-            return
-        from installer import install
-        install('olive-ai', 'Olive')
-        log.info("Olive is installed. Please restart ui completely to load Olive.")
+        if not installed('olive-ai', reload=True, quiet=True):
+            install('olive-ai', 'olive-ai')
     else:
-        from installer import pip
         global is_available
+        is_available = False
         if "olive" in sys.modules:
             del sys.modules["olive"]
-        is_available = False
         if shared.opts.diffusers_pipeline == 'ONNX Stable Diffusion with Olive':
             shared.opts.diffusers_pipeline = 'ONNX Stable Diffusion'
-        pip('uninstall olive-ai --yes --quiet', ignore=True, quiet=True)
+        if installed('olive-ai', reload=True, quiet=True):
+            uninstall('olive-ai')
 
 submodels = ("text_encoder", "unet", "vae_encoder", "vae_decoder",)
 
@@ -42,6 +39,7 @@ EP_TO_NAME = {
 }
 
 class OlivePipeline(diffusers.DiffusionPipeline):
+    model_type = diffusers.OnnxStableDiffusionPipeline.__name__
     sd_model_hash: str
     sd_checkpoint_info: CheckpointInfo
     sd_model_checkpoint: str
@@ -70,6 +68,12 @@ class OlivePipeline(diffusers.DiffusionPipeline):
     def from_ckpt(*args, **kwargs):
         return OlivePipeline.from_single_file(**args, **kwargs)
 
+    def derive_properties(self, pipeline: diffusers.OnnxStableDiffusionPipeline):
+        pipeline.sd_model_hash = self.sd_model_hash
+        pipeline.sd_checkpoint_info = self.sd_checkpoint_info
+        pipeline.sd_model_checkpoint = self.sd_model_checkpoint
+        return pipeline
+
     def to(self, *args, **kwargs):
         pass
 
@@ -85,11 +89,13 @@ class OlivePipeline(diffusers.DiffusionPipeline):
             log.warning("Olive received different width and height. The quality of the result is not guaranteed.")
 
         out_dir = os.path.join(shared.opts.olive_cached_models_path, f"{self.original_filename}-{width}w-{height}h")
-        if os.path.isdir(out_dir):
+        if os.path.isdir(out_dir): # already optimized (cached)
             del self.unoptimized
-            return OnnxStableDiffusionPipeline.from_pretrained(
-                out_dir,
-            ).apply(self)
+            return self.derive_properties(
+                diffusers.OnnxStableDiffusionPipeline.from_pretrained(
+                    out_dir,
+                )
+            )
 
         try:
             if shared.opts.onnx_cache_optimized:
@@ -146,10 +152,12 @@ class OlivePipeline(diffusers.DiffusionPipeline):
                     provider=(shared.opts.onnx_execution_provider, get_execution_provider_options(),),
                 )
 
-            pipeline = OnnxStableDiffusionPipeline(
-                **kwargs,
-                requires_safety_checker=False,
-            ).apply(self)
+            pipeline = self.derive_properties(
+                diffusers.OnnxStableDiffusionPipeline(
+                    **kwargs,
+                    requires_safety_checker=False,
+                )
+            )
             del kwargs
             if shared.opts.onnx_cache_optimized:
                 pipeline.to_json_file(os.path.join(out_dir, "model_index.json"))

modules/onnx.py

@@ -68,15 +68,15 @@ diffusers.OnnxRuntimeModel = OnnxRuntimeModel
 
 
 class OnnxStableDiffusionPipeline(diffusers.OnnxStableDiffusionPipeline):
-    model_type: str
+    model_type = diffusers.OnnxStableDiffusionPipeline.__name__
     sd_model_hash: str
     sd_checkpoint_info: CheckpointInfo
     sd_model_checkpoint: str
 
     @staticmethod
     def from_pretrained(pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
-        provider = (shared.opts.onnx_execution_provider, get_execution_provider_options(),)
-        init_dict = diffusers.OnnxStableDiffusionPipeline.extract_init_dict(diffusers.DiffusionPipeline.load_config(pretrained_model_name_or_path), **kwargs)[0]
+        kwargs["provider"] = kwargs["provider"] if "provider" in kwargs else (shared.opts.onnx_execution_provider, get_execution_provider_options(),)
+        init_dict = super(OnnxStableDiffusionPipeline, OnnxStableDiffusionPipeline).extract_init_dict(diffusers.DiffusionPipeline.load_config(pretrained_model_name_or_path), **kwargs)[0]
         init_kwargs = {}
         for k, v in init_dict.items():
             if not isinstance(v, list):
@@ -90,16 +90,16 @@ class OnnxStableDiffusionPipeline(diffusers.OnnxStableDiffusionPipeline):
             constructor = getattr(library, constructor_name)
             submodel_kwargs = {}
             if issubclass(constructor, diffusers.OnnxRuntimeModel):
-                submodel_kwargs["provider"] = provider
-            init_kwargs[k] = constructor.from_pretrained(os.path.join(pretrained_model_name_or_path, k), **submodel_kwargs)
+                submodel_kwargs["provider"] = kwargs["provider"]
+            try:
+                init_kwargs[k] = constructor.from_pretrained(
+                    os.path.join(pretrained_model_name_or_path, k),
+                    **submodel_kwargs,
+                )
+            except Exception:
+                pass
         return OnnxStableDiffusionPipeline(**init_kwargs)
 
-    def apply(self, dummy_pipeline):
-        self.sd_model_hash = dummy_pipeline.sd_model_hash
-        self.sd_checkpoint_info = dummy_pipeline.sd_checkpoint_info
-        self.sd_model_checkpoint = dummy_pipeline.sd_model_checkpoint
-        return self
-
     def __call__(
         self,
         prompt: Union[str, List[str]] = None,
@@ -243,3 +243,5 @@ class OnnxStableDiffusionPipeline(diffusers.OnnxStableDiffusionPipeline):
             return (image, has_nsfw_concept)
 
         return diffusers.pipelines.stable_diffusion.StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
+
+diffusers.OnnxStableDiffusionPipeline = OnnxStableDiffusionPipeline

modules/sd_models.py

@@ -147,7 +147,7 @@ def list_models():
     model_list = list(modelloader.load_models(model_path=model_path, model_url=None, command_path=shared.opts.ckpt_dir, ext_filter=ext_filter, download_name=None, ext_blacklist=[".vae.ckpt", ".vae.safetensors"]))
     if shared.backend == shared.Backend.DIFFUSERS:
         model_list += modelloader.load_diffusers_models(model_path=os.path.join(models_path, 'Diffusers'), command_path=shared.opts.diffusers_dir, clear=True)
-        model_list += modelloader.load_diffusers_models(model_path=shared.opts.olive_sideloaded_models_path, command_path=shared.opts.olive_sideloaded_models_path)
+        model_list += modelloader.load_diffusers_models(model_path=shared.opts.olive_sideloaded_models_path, command_path=shared.opts.olive_sideloaded_models_path, clear=False)
     for filename in sorted(model_list, key=str.lower):
         checkpoint_info = CheckpointInfo(filename)
         if checkpoint_info.name is not None:
@@ -790,8 +790,10 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
 
         shared.log.debug(f'Diffusers loading: path="{checkpoint_info.path}"')
         pipeline, model_type = detect_pipeline(checkpoint_info.path, op)
-        if os.path.isdir(checkpoint_info.path):
-            if shared.opts.olive_sideloaded_models_path in checkpoint_info.path:
+        if 'ONNX' in shared.opts.diffusers_pipeline:
+            from modules.onnx import get_execution_provider_options
+            diffusers_load_config['provider'] = (shared.opts.onnx_execution_provider, get_execution_provider_options(),)
+            if shared.opts.diffusers_pipeline == 'ONNX Stable Diffusion with Olive':
                 try:
                     from modules.onnx import OnnxStableDiffusionPipeline
                     sd_model = OnnxStableDiffusionPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.olive_sideloaded_models_path)