From df021b3982c5af36cdac82fa6608b91ad50e26db Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Wed, 25 Oct 2023 19:18:14 +0900 Subject: [PATCH] support more backends --- installer.py | 3 +-- launch.py | 10 ++++++++-- modules/olive.py | 32 +++++++++++++++++++++++++------- modules/sd_models.py | 4 ++-- modules/shared_items.py | 5 +++-- requirements.txt | 4 ---- 6 files changed, 39 insertions(+), 19 deletions(-) diff --git a/installer.py b/installer.py index d5517b915..412f156f1 100644 --- a/installer.py +++ b/installer.py @@ -583,8 +583,7 @@ def install_packages(): install(clip_package, 'clip') invisiblewatermark_package = os.environ.get('INVISIBLEWATERMARK_PACKAGE', "git+https://github.com/patrickvonplaten/invisible-watermark.git@remove_onnxruntime_depedency") install(invisiblewatermark_package, 'invisible-watermark') - install('olive-ai[directml]', 'olive-ai', ignore=True) - install('onnxruntime-directml==1.16.1', 'onnxruntime-directml', ignore=True) + install('olive-ai', 'olive-ai', ignore=True) install('pi-heif', 'pi_heif', ignore=True) tensorflow_package = os.environ.get('TENSORFLOW_PACKAGE', 'tensorflow==2.13.0') install(tensorflow_package, 'tensorflow-rocm' if 'rocm' in tensorflow_package else 'tensorflow', ignore=True) diff --git a/launch.py b/launch.py index 4617e5f60..31a8e9f70 100755 --- a/launch.py +++ b/launch.py @@ -23,8 +23,14 @@ python = sys.executable # used by some extensions to run python skip_install = False # parsed by some extensions -def init_args(): - global parser, args # pylint: disable=global-statement +try: + import torch._dynamo +except ModuleNotFoundError: + sys.modules["torch._dynamo"] = {} # HACK torch 1.13.1 does not have _dynamo. will be removed. + + +def init_modules(): + global parser, args, script_path, extensions_dir # pylint: disable=global-statement import modules.cmd_args parser = modules.cmd_args.parser installer.add_args(parser) diff --git a/modules/olive.py b/modules/olive.py index 71558e331..325c56a10 100644 --- a/modules/olive.py +++ b/modules/olive.py @@ -4,6 +4,7 @@ import torch import shutil import diffusers import numpy as np +import onnxruntime as ort from typing import Union, Optional, Callable, List from transformers.models.clip.modeling_clip import CLIPTextModel, CLIPTextModelWithProjection from installer import log, args @@ -13,14 +14,23 @@ from modules.sd_models import CheckpointInfo submodels = ("text_encoder", "unet", "vae_encoder", "vae_decoder",) -execution_provider = "CUDAExecutionProvider" +available_execution_providers = ort.get_available_providers() +execution_provider = "CUDAExecutionProvider" if "CUDAExecutionProvider" in available_execution_providers else "CPUExecutionProvider" +execution_provider_options = { + "device_id": int(cmd_opts.device_id or 0), +} if args.use_directml: execution_provider = "DmlExecutionProvider" elif args.use_rocm: - execution_provider = "ROCmExecutionProvider" -provider = (execution_provider, { - "device_id": int(cmd_opts.device_id or 0), -}) + if "ROCMExecutionProvider" in available_execution_providers: + execution_provider = "ROCMExecutionProvider" + execution_provider_options["tunable_op_enable"] = 1 + execution_provider_options["tunable_op_tuning_enable"] = 1 + else: + log.warning("Currently, there's no pypi release for onnxruntime-rocm. Please download and install .whl file from https://download.onnxruntime.ai/ The inference will be fall back to CPU.") +elif args.use_ipex or args.use_openvino: + execution_provider = "OpenVINOExecutionProvider" +provider = (execution_provider, execution_provider_options,) class OnnxRuntimeModel(diffusers.OnnxRuntimeModel): config = {} @@ -38,6 +48,12 @@ class OnnxStableDiffusionPipeline(diffusers.OnnxStableDiffusionPipeline): sd_checkpoint_info: CheckpointInfo sd_model_checkpoint: str + @staticmethod + def from_pretrained(*args, **kwargs): + if "provider" not in kwargs: + kwargs["provider"] = provider + return diffusers.OnnxStableDiffusionPipeline.from_pretrained(*args, **kwargs) + def apply(self, dummy_pipeline): self.sd_model_hash = dummy_pipeline.sd_model_hash self.sd_checkpoint_info = dummy_pipeline.sd_checkpoint_info @@ -229,7 +245,9 @@ class OlivePipeline(diffusers.DiffusionPipeline): out_dir = os.path.join(opts.olive_cached_models_path, f"{self.original_filename}-{width}w-{height}h") if os.path.isdir(out_dir): del self.unoptimized - return OnnxStableDiffusionPipeline.from_pretrained(out_dir, provider=provider).apply(self) + return OnnxStableDiffusionPipeline.from_pretrained( + out_dir, + ).apply(self) try: if opts.olive_cache_optimized: @@ -279,7 +297,7 @@ class OlivePipeline(diffusers.DiffusionPipeline): del self.unoptimized for submodel in submodels: kwargs[submodel] = diffusers.OnnxRuntimeModel.from_pretrained( - os.path.dirname(optimized_model_paths[submodel]), provider=provider, + os.path.dirname(optimized_model_paths[submodel]), ) pipeline = OnnxStableDiffusionPipeline( diff --git a/modules/sd_models.py b/modules/sd_models.py index 179d28c17..71952d552 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -793,8 +793,8 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No if os.path.isdir(checkpoint_info.path): if shared.opts.olive_sideloaded_models_path in checkpoint_info.path: try: - from modules.olive import OnnxStableDiffusionPipeline, provider - sd_model = OnnxStableDiffusionPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.olive_sideloaded_models_path, provider=provider) + from modules.olive import OnnxStableDiffusionPipeline + sd_model = OnnxStableDiffusionPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.olive_sideloaded_models_path) sd_model.model_type = sd_model.__class__.__name__ except Exception as e: shared.log.error(f'Failed loading {op}: {checkpoint_info.path} olive={e}') diff --git a/modules/shared_items.py b/modules/shared_items.py index 1b004dad8..5e7a75427 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -26,7 +26,7 @@ def list_crossattention(): def get_pipelines(): import diffusers - from modules.olive import OlivePipeline + from modules.olive import OnnxStableDiffusionPipeline, OlivePipeline from installer import log pipelines = { # note: not all pipelines can be used manually as they require prior pipeline next to decoder pipeline 'Autodetect': None, @@ -39,7 +39,8 @@ def get_pipelines(): 'Stable Diffusion XL Img2Img': getattr(diffusers, 'StableDiffusionXLImg2ImgPipeline', None), 'Stable Diffusion XL Inpaint': getattr(diffusers, 'StableDiffusionXLInpaintPipeline', None), 'Stable Diffusion XL Instruct': getattr(diffusers, 'StableDiffusionXLInstructPix2PixPipeline', None), - 'ONNX Stable Diffusion (Olive)': OlivePipeline, + 'ONNX Stable Diffusion': OnnxStableDiffusionPipeline, + 'ONNX Stable Diffusion with Olive': OlivePipeline, 'Latent Consistency Model': getattr(diffusers, 'LatentConsistencyModelPipeline', None), 'PixArt Alpha': getattr(diffusers, 'PixArtAlphaPipeline', None), 'UniDiffuser': getattr(diffusers, 'UniDiffuserPipeline', None), diff --git a/requirements.txt b/requirements.txt index 6af77209b..d78c266b5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -65,7 +65,3 @@ Pillow==9.5.0 timm==0.9.7 pydantic==1.10.13 typing-extensions==4.8.0 - -torch==1.13.1 -torchvision==0.14.1 -torch-directml==0.1.13.1.dev230413