diff --git a/launch.py b/launch.py index 767eaa984..9f4d639f8 100755 --- a/launch.py +++ b/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() diff --git a/modules/olive.py b/modules/olive.py index a9ac3f19f..5dfa87bae 100644 --- a/modules/olive.py +++ b/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")) diff --git a/modules/onnx.py b/modules/onnx.py index c39b964e7..0ab540c3b 100644 --- a/modules/onnx.py +++ b/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 diff --git a/modules/sd_models.py b/modules/sd_models.py index a0ed5b6b8..8ede576c9 100644 --- a/modules/sd_models.py +++ b/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)