diff --git a/modules/olive.py b/modules/olive.py index 8aef3c502..e6726c56a 100644 --- a/modules/olive.py +++ b/modules/olive.py @@ -1,7 +1,7 @@ import os import torch import diffusers -from typing import Type, Callable, Dict, Any +from typing import Type, Callable, TypeVar, Dict, Any from transformers.models.clip.modeling_clip import CLIPTextModel, CLIPTextModelWithProjection @@ -35,7 +35,10 @@ class ENVStore: def __delattr__(self, name: str) -> None: if name not in self.__class__.__annotations__: return - os.environ.pop(f"SDNEXT_OLIVE_{name}") + key = f"SDNEXT_OLIVE_{name}" + if key not in os.environ: + return + os.environ.pop(key) class OliveOptimizerConfig(ENVStore): @@ -43,8 +46,8 @@ class OliveOptimizerConfig(ENVStore): is_sdxl: bool - vae_id: str - vae_subfolder: str + vae: str + vae_sdxl_fp16_fix: bool width: int height: int @@ -80,6 +83,15 @@ def get_loader_arguments(): return {} +T = TypeVar("T") +def from_pretrained(cls: Type[T], pretrained_model_name_or_path: os.PathLike, *args, **kwargs) -> T: + pretrained_model_name_or_path = str(pretrained_model_name_or_path) + if pretrained_model_name_or_path.endswith(".onnx"): + cls = diffusers.OnnxRuntimeModel + pretrained_model_name_or_path = os.path.dirname(pretrained_model_name_or_path) + return cls.from_pretrained(pretrained_model_name_or_path, *args, **kwargs, **get_loader_arguments()) + + # ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. @@ -111,7 +123,7 @@ def text_encoder_inputs(_, torch_dtype): def text_encoder_load(model_name): - model = CLIPTextModel.from_pretrained(model_name, subfolder="text_encoder", **get_loader_arguments()) + model = from_pretrained(CLIPTextModel, model_name, subfolder="text_encoder") return model @@ -136,7 +148,7 @@ def text_encoder_2_inputs(_, torch_dtype): def text_encoder_2_load(model_name): - model = CLIPTextModelWithProjection.from_pretrained(model_name, subfolder="text_encoder_2", **get_loader_arguments()) + model = from_pretrained(CLIPTextModelWithProjection, model_name, subfolder="text_encoder_2") return model @@ -199,7 +211,7 @@ def unet_inputs(_, torch_dtype, is_conversion_inputs=False): def unet_load(model_name): - model = diffusers.UNet2DConditionModel.from_pretrained(model_name, subfolder="unet", **get_loader_arguments()) + model = from_pretrained(diffusers.UNet2DConditionModel, model_name, subfolder="unet") return model @@ -225,11 +237,18 @@ def vae_encoder_inputs(_, torch_dtype): def vae_encoder_load(model_name): subfolder = "vae_encoder" if os.path.isdir(os.path.join(model_name, "vae_encoder")) else "vae" - if config.vae_id is not None: - model_name = config.vae_id - subfolder = config.vae_subfolder - model = diffusers.AutoencoderKL.from_pretrained(model_name, subfolder=subfolder, **get_loader_arguments()) + + if config.vae_sdxl_fp16_fix: + model_name = "madebyollin/sdxl-vae-fp16-fix" + subfolder = "" + + if config.vae is None: + model = from_pretrained(diffusers.AutoencoderKL, model_name, subfolder=subfolder) + else: + model = diffusers.AutoencoderKL.from_single_file(config.vae) + model.forward = lambda sample, return_dict: model.encode(sample, return_dict)[0].sample() + return model @@ -255,11 +274,18 @@ def vae_decoder_inputs(_, torch_dtype): def vae_decoder_load(model_name): subfolder = "vae_decoder" if os.path.isdir(os.path.join(model_name, "vae_decoder")) else "vae" - if config.vae_id is not None: - model_name = config.vae_id - subfolder = config.vae_subfolder - model = diffusers.AutoencoderKL.from_pretrained(model_name, subfolder=subfolder, **get_loader_arguments()) + + if config.vae_sdxl_fp16_fix: + model_name = "madebyollin/sdxl-vae-fp16-fix" + subfolder = "" + + if config.vae is None: + model = from_pretrained(diffusers.AutoencoderKL, model_name, subfolder=subfolder) + else: + model = diffusers.AutoencoderKL.from_single_file(config.vae) + model.forward = model.decode + return model diff --git a/modules/onnx_pipelines.py b/modules/onnx_pipelines.py index b2d40e46f..bfb54afa2 100644 --- a/modules/onnx_pipelines.py +++ b/modules/onnx_pipelines.py @@ -15,7 +15,7 @@ from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput from diffusers.image_processor import VaeImageProcessor, PipelineImageInput from installer import log from modules import shared -from modules.paths import sd_configs_path +from modules.paths import sd_configs_path, models_path from modules.sd_models import CheckpointInfo from modules.processing import StableDiffusionProcessing from modules.olive import config @@ -24,6 +24,11 @@ from modules.onnx_utils import extract_device, move_inference_session, check_dif from modules.onnx_ep import ExecutionProvider, EP_TO_NAME, get_provider +SUBMODELS_SD = ("text_encoder", "unet", "vae_encoder", "vae_decoder",) +SUBMODELS_SDXL = ("text_encoder", "text_encoder_2", "unet", "vae_encoder", "vae_decoder",) +SUBMODELS_SDXL_REFINER = ("text_encoder_2", "unet", "vae_encoder", "vae_decoder",) + + class OnnxPipelineBase(OnnxFakeModule, diffusers.DiffusionPipeline, metaclass=ABCMeta): model_type: str sd_model_hash: str @@ -40,7 +45,7 @@ class OnnxPipelineBase(OnnxFakeModule, diffusers.DiffusionPipeline, metaclass=AB expected_modules, _ = self._get_signature_keys(self) for name in expected_modules: if not hasattr(self, name): - log.warn(f"Pipeline does not have module '{name}'.") + log.warning(f"Pipeline does not have module '{name}'.") continue module = getattr(self, name) @@ -89,7 +94,6 @@ class OnnxRawPipeline(OnnxPipelineBase): original_filename: str constructor: Type[OnnxPipelineBase] - submodels: List[str] init_dict: Dict[str, Tuple[str]] = {} scheduler: Any = None # for Img2Img @@ -116,8 +120,8 @@ class OnnxRawPipeline(OnnxPipelineBase): del pipeline self.init_dict = load_init_dict(constructor, path) except Exception: - log.error(f'Failed to load ONNX pipeline: is_sdxl={self._is_sdxl}') - log.warn('Model load failed. Please check Diffusers pipeline in Compute Settings.') + log.error(f'ONNX: Failed to load ONNX pipeline: is_sdxl={self._is_sdxl}') + log.warning('ONNX: You cannot load this model using the pipeline you selected. Please check Diffusers pipeline in Compute Settings.') return if "vae" in self.init_dict: del self.init_dict["vae"] @@ -134,7 +138,7 @@ class OnnxRawPipeline(OnnxPipelineBase): pipeline.scheduler = self.scheduler return pipeline - def convert(self, in_dir: os.PathLike): + def convert(self, submodels: List[str], in_dir: os.PathLike): if not shared.cmd_opts.debug: ort.set_default_logger_severity(3) @@ -144,105 +148,98 @@ class OnnxRawPipeline(OnnxPipelineBase): if os.path.isdir(out_dir): # if model is ONNX format or had already converted. return out_dir + from olive.workflows import run try: - from olive.workflows import run - try: - from olive.model import ONNXModel - except ImportError: - from olive.model import ONNXModelHandler as ONNXModel + from olive.model import ONNXModel + except ImportError: + from olive.model import ONNXModelHandler as ONNXModel - shutil.rmtree("cache", ignore_errors=True) - shutil.rmtree("footprints", ignore_errors=True) + shutil.rmtree("cache", ignore_errors=True) + shutil.rmtree("footprints", ignore_errors=True) - if shared.opts.onnx_cache_converted: - shutil.copytree( - in_dir, out_dir, ignore=shutil.ignore_patterns("weights.pb", "*.onnx", "*.safetensors", "*.ckpt") - ) + if shared.opts.onnx_cache_converted: + shutil.copytree( + in_dir, out_dir, ignore=shutil.ignore_patterns("weights.pb", "*.onnx", "*.safetensors", "*.ckpt") + ) - converted_model_paths = {} + converted_model_paths = {} - for submodel in self.submodels: - log.info(f"\nConverting {submodel}") + for submodel in submodels: + log.info(f"\nConverting {submodel}") - with open(os.path.join(sd_configs_path, "onnx", f"{'sdxl' if self._is_sdxl else 'sd'}_{submodel}.json"), "r") as config_file: - conversion_config = json.load(config_file) - conversion_config["input_model"]["config"]["model_path"] = os.path.abspath(in_dir) - conversion_config["engine"]["execution_providers"] = [shared.opts.onnx_execution_provider] + with open(os.path.join(sd_configs_path, "onnx", f"{'sdxl' if self._is_sdxl else 'sd'}_{submodel}.json"), "r") as config_file: + conversion_config = json.load(config_file) + conversion_config["input_model"]["config"]["model_path"] = os.path.abspath(in_dir) + conversion_config["engine"]["execution_providers"] = [shared.opts.onnx_execution_provider] - run(conversion_config) + run(conversion_config) - with open(os.path.join("footprints", f"{submodel}_{EP_TO_NAME[shared.opts.onnx_execution_provider]}_footprints.json"), "r") as footprint_file: - footprints = json.load(footprint_file) - conversion_footprint = None - for _, footprint in footprints.items(): - if footprint["from_pass"] == "OnnxConversion": - conversion_footprint = footprint + with open(os.path.join("footprints", f"{submodel}_{EP_TO_NAME[shared.opts.onnx_execution_provider]}_footprints.json"), "r") as footprint_file: + footprints = json.load(footprint_file) + conversion_footprint = None + for _, footprint in footprints.items(): + if footprint["from_pass"] == "OnnxConversion": + conversion_footprint = footprint - assert conversion_footprint, "Failed to convert model" + assert conversion_footprint, "Failed to convert model" - converted_model_paths[submodel] = ONNXModel( - **conversion_footprint["model_config"]["config"] - ).model_path + converted_model_paths[submodel] = ONNXModel( + **conversion_footprint["model_config"]["config"] + ).model_path - log.info(f"Converted {submodel}") + log.info(f"Converted {submodel}") - for submodel in self.submodels: - src_path = converted_model_paths[submodel] - src_parent = os.path.dirname(src_path) - dst_parent = os.path.join(out_dir, submodel) - dst_path = os.path.join(dst_parent, "model.onnx") - if not os.path.isdir(dst_parent): - os.mkdir(dst_parent) - shutil.copyfile(src_path, dst_path) + for submodel in submodels: + src_path = converted_model_paths[submodel] + src_parent = os.path.dirname(src_path) + dst_parent = os.path.join(out_dir, submodel) + dst_path = os.path.join(dst_parent, "model.onnx") + if not os.path.isdir(dst_parent): + os.mkdir(dst_parent) + shutil.copyfile(src_path, dst_path) - data_src_path = os.path.join(src_parent, (os.path.basename(src_path) + ".data")) - if os.path.isfile(data_src_path): - data_dst_path = os.path.join(dst_parent, (os.path.basename(dst_path) + ".data")) - shutil.copyfile(data_src_path, data_dst_path) + data_src_path = os.path.join(src_parent, (os.path.basename(src_path) + ".data")) + if os.path.isfile(data_src_path): + data_dst_path = os.path.join(dst_parent, (os.path.basename(dst_path) + ".data")) + shutil.copyfile(data_src_path, data_dst_path) - weights_src_path = os.path.join(src_parent, "weights.pb") - if os.path.isfile(weights_src_path): - weights_dst_path = os.path.join(dst_parent, "weights.pb") - shutil.copyfile(weights_src_path, weights_dst_path) - del converted_model_paths + weights_src_path = os.path.join(src_parent, "weights.pb") + if os.path.isfile(weights_src_path): + weights_dst_path = os.path.join(dst_parent, "weights.pb") + shutil.copyfile(weights_src_path, weights_dst_path) + del converted_model_paths - kwargs = {} + kwargs = {} - init_dict = self.init_dict.copy() - for submodel in self.submodels: - kwargs[submodel] = diffusers.OnnxRuntimeModel.load_model( - os.path.join(out_dir, submodel, "model.onnx"), - provider=get_provider(), - ) if self._is_sdxl else diffusers.OnnxRuntimeModel.from_pretrained( - os.path.join(out_dir, submodel), - provider=get_provider(), - ) - if submodel in init_dict: - del init_dict[submodel] # already loaded as OnnxRuntimeModel. - kwargs.update(load_submodels(in_dir, self._is_sdxl, init_dict)) # load others. - constructor = get_base_constructor(self.constructor, self.is_refiner) - kwargs = patch_kwargs(constructor, kwargs) + init_dict = self.init_dict.copy() + for submodel in submodels: + kwargs[submodel] = diffusers.OnnxRuntimeModel.load_model( + os.path.join(out_dir, submodel, "model.onnx"), + provider=get_provider(), + ) if self._is_sdxl else diffusers.OnnxRuntimeModel.from_pretrained( + os.path.join(out_dir, submodel), + provider=get_provider(), + ) + if submodel in init_dict: + del init_dict[submodel] # already loaded as OnnxRuntimeModel. + kwargs.update(load_submodels(in_dir, self._is_sdxl, init_dict)) # load others. + constructor = get_base_constructor(self.constructor, self.is_refiner) + kwargs = patch_kwargs(constructor, kwargs) - pipeline = constructor(**kwargs) - model_index = json.loads(pipeline.to_json_string()) - del pipeline + pipeline = constructor(**kwargs) + model_index = json.loads(pipeline.to_json_string()) + del pipeline - for k, v in init_dict.items(): # copy missing submodels. (ORTStableDiffusionXLPipeline) - if k not in model_index: - model_index[k] = v + for k, v in init_dict.items(): # copy missing submodels. (ORTStableDiffusionXLPipeline) + if k not in model_index: + model_index[k] = v - with open(os.path.join(out_dir, "model_index.json"), 'w') as file: - json.dump(model_index, file) + with open(os.path.join(out_dir, "model_index.json"), 'w') as file: + json.dump(model_index, file) - return out_dir - except Exception as e: - log.error(f"Failed to convert model '{self.original_filename}'.") - log.error(e) # for test. - shutil.rmtree(shared.opts.onnx_temp_dir, ignore_errors=True) - shutil.rmtree(out_dir, ignore_errors=True) - return None + return out_dir - def optimize(self, in_dir: os.PathLike): + def run_olive(self, submodels: List[str], in_dir: os.PathLike): if not shared.cmd_opts.debug: ort.set_default_logger_severity(4) @@ -253,130 +250,122 @@ class OnnxRawPipeline(OnnxPipelineBase): if not shared.opts.olive_cache_optimized: out_dir = shared.opts.onnx_temp_dir + from olive.workflows import run try: - from olive.workflows import run - try: - from olive.model import ONNXModel - except ImportError: - from olive.model import ONNXModelHandler as ONNXModel + from olive.model import ONNXModel + except ImportError: + from olive.model import ONNXModelHandler as ONNXModel - shutil.rmtree("cache", ignore_errors=True) - shutil.rmtree("footprints", ignore_errors=True) + shutil.rmtree("cache", ignore_errors=True) + shutil.rmtree("footprints", ignore_errors=True) - if shared.opts.olive_cache_optimized: - shutil.copytree( - in_dir, out_dir, ignore=shutil.ignore_patterns("weights.pb", "*.onnx", "*.safetensors", "*.ckpt") - ) + if shared.opts.olive_cache_optimized: + shutil.copytree( + in_dir, out_dir, ignore=shutil.ignore_patterns("weights.pb", "*.onnx", "*.safetensors", "*.ckpt") + ) - optimized_model_paths = {} + optimized_model_paths = {} - for submodel in self.submodels: - log.info(f"\nProcessing {submodel}") + for submodel in submodels: + log.info(f"\nProcessing {submodel}") - with open(os.path.join(sd_configs_path, "olive", 'sdxl' if self._is_sdxl else 'sd', f"{submodel}.json"), "r") as config_file: - olive_config: Dict[str, Dict[str, Dict]] = json.load(config_file) + with open(os.path.join(sd_configs_path, "olive", 'sdxl' if self._is_sdxl else 'sd', f"{submodel}.json"), "r") as config_file: + olive_config: Dict[str, Dict[str, Dict]] = json.load(config_file) - for flow in olive_config["pass_flows"]: - for i in range(len(flow)): - flow[i] = flow[i].replace("AutoExecutionProvider", shared.opts.onnx_execution_provider) - olive_config["input_model"]["config"]["model_path"] = os.path.abspath(os.path.join(in_dir, submodel, "model.onnx")) - olive_config["engine"]["execution_providers"] = [shared.opts.onnx_execution_provider] + for flow in olive_config["pass_flows"]: + for i in range(len(flow)): + flow[i] = flow[i].replace("AutoExecutionProvider", shared.opts.onnx_execution_provider) + olive_config["input_model"]["config"]["model_path"] = os.path.abspath(os.path.join(in_dir, submodel, "model.onnx")) + olive_config["engine"]["execution_providers"] = [shared.opts.onnx_execution_provider] - for pass_key in olive_config["passes"]: - if olive_config["passes"][pass_key]["type"] == "OrtTransformersOptimization": - float16 = shared.opts.olive_float16 and not (submodel == "vae_encoder" and shared.opts.olive_vae_encoder_float32) - olive_config["passes"][pass_key]["config"]["float16"] = float16 - if shared.opts.onnx_execution_provider == ExecutionProvider.CUDA or shared.opts.onnx_execution_provider == ExecutionProvider.ROCm: - if version.parse(ort.__version__) < version.parse("1.17.0"): - olive_config["passes"][pass_key]["config"]["optimization_options"] = {"enable_skip_group_norm": False} - if float16: - olive_config["passes"][pass_key]["config"]["keep_io_types"] = False + for pass_key in olive_config["passes"]: + if olive_config["passes"][pass_key]["type"] == "OrtTransformersOptimization": + float16 = shared.opts.olive_float16 and not (submodel == "vae_encoder" and shared.opts.olive_vae_encoder_float32) + olive_config["passes"][pass_key]["config"]["float16"] = float16 + if shared.opts.onnx_execution_provider == ExecutionProvider.CUDA or shared.opts.onnx_execution_provider == ExecutionProvider.ROCm: + if version.parse(ort.__version__) < version.parse("1.17.0"): + olive_config["passes"][pass_key]["config"]["optimization_options"] = {"enable_skip_group_norm": False} + if float16: + olive_config["passes"][pass_key]["config"]["keep_io_types"] = False - run(olive_config) + run(olive_config) - with open(os.path.join("footprints", f"{submodel}_{EP_TO_NAME[shared.opts.onnx_execution_provider]}_footprints.json"), "r") as footprint_file: - footprints = json.load(footprint_file) - processor_final_pass_footprint = None - for _, footprint in footprints.items(): - if footprint["from_pass"] == olive_config["passes"][olive_config["pass_flows"][-1][-1]]["type"]: - processor_final_pass_footprint = footprint + with open(os.path.join("footprints", f"{submodel}_{EP_TO_NAME[shared.opts.onnx_execution_provider]}_footprints.json"), "r") as footprint_file: + footprints = json.load(footprint_file) + processor_final_pass_footprint = None + for _, footprint in footprints.items(): + if footprint["from_pass"] == olive_config["passes"][olive_config["pass_flows"][-1][-1]]["type"]: + processor_final_pass_footprint = footprint - assert processor_final_pass_footprint, "Failed to optimize model" + assert processor_final_pass_footprint, "Failed to optimize model" - optimized_model_paths[submodel] = ONNXModel( - **processor_final_pass_footprint["model_config"]["config"] - ).model_path + optimized_model_paths[submodel] = ONNXModel( + **processor_final_pass_footprint["model_config"]["config"] + ).model_path - log.info(f"Processed {submodel}") + log.info(f"Processed {submodel}") - for submodel in self.submodels: - src_path = optimized_model_paths[submodel] - src_parent = os.path.dirname(src_path) - dst_parent = os.path.join(out_dir, submodel) - dst_path = os.path.join(dst_parent, "model.onnx") - if not os.path.isdir(dst_parent): - os.mkdir(dst_parent) - shutil.copyfile(src_path, dst_path) + for submodel in submodels: + src_path = optimized_model_paths[submodel] + src_parent = os.path.dirname(src_path) + dst_parent = os.path.join(out_dir, submodel) + dst_path = os.path.join(dst_parent, "model.onnx") + if not os.path.isdir(dst_parent): + os.mkdir(dst_parent) + shutil.copyfile(src_path, dst_path) - data_src_path = os.path.join(src_parent, (os.path.basename(src_path) + ".data")) - if os.path.isfile(data_src_path): - data_dst_path = os.path.join(dst_parent, (os.path.basename(dst_path) + ".data")) - shutil.copyfile(data_src_path, data_dst_path) + data_src_path = os.path.join(src_parent, (os.path.basename(src_path) + ".data")) + if os.path.isfile(data_src_path): + data_dst_path = os.path.join(dst_parent, (os.path.basename(dst_path) + ".data")) + shutil.copyfile(data_src_path, data_dst_path) - weights_src_path = os.path.join(src_parent, "weights.pb") - if os.path.isfile(weights_src_path): - weights_dst_path = os.path.join(dst_parent, "weights.pb") - shutil.copyfile(weights_src_path, weights_dst_path) - del optimized_model_paths + weights_src_path = os.path.join(src_parent, "weights.pb") + if os.path.isfile(weights_src_path): + weights_dst_path = os.path.join(dst_parent, "weights.pb") + shutil.copyfile(weights_src_path, weights_dst_path) + del optimized_model_paths - kwargs = {} + kwargs = {} - init_dict = self.init_dict.copy() - for submodel in self.submodels: - kwargs[submodel] = diffusers.OnnxRuntimeModel.load_model( - os.path.join(out_dir, submodel, "model.onnx"), - provider=get_provider(), - ) if self._is_sdxl else diffusers.OnnxRuntimeModel.from_pretrained( - os.path.join(out_dir, submodel), - provider=get_provider(), - ) - if submodel in init_dict: - del init_dict[submodel] # already loaded as OnnxRuntimeModel. - kwargs.update(load_submodels(in_dir, self._is_sdxl, init_dict)) # load others. - constructor = get_base_constructor(self.constructor, self.is_refiner) - kwargs = patch_kwargs(constructor, kwargs) + init_dict = self.init_dict.copy() + for submodel in submodels: + kwargs[submodel] = diffusers.OnnxRuntimeModel.load_model( + os.path.join(out_dir, submodel, "model.onnx"), + provider=get_provider(), + ) if self._is_sdxl else diffusers.OnnxRuntimeModel.from_pretrained( + os.path.join(out_dir, submodel), + provider=get_provider(), + ) + if submodel in init_dict: + del init_dict[submodel] # already loaded as OnnxRuntimeModel. + kwargs.update(load_submodels(in_dir, self._is_sdxl, init_dict)) # load others. + constructor = get_base_constructor(self.constructor, self.is_refiner) + kwargs = patch_kwargs(constructor, kwargs) - pipeline = constructor(**kwargs) - model_index = json.loads(pipeline.to_json_string()) - del pipeline + pipeline = constructor(**kwargs) + model_index = json.loads(pipeline.to_json_string()) + del pipeline - for k, v in init_dict.items(): # copy missing submodels. (ORTStableDiffusionXLPipeline) - if k not in model_index: - model_index[k] = v + for k, v in init_dict.items(): # copy missing submodels. (ORTStableDiffusionXLPipeline) + if k not in model_index: + model_index[k] = v - with open(os.path.join(out_dir, "model_index.json"), 'w') as file: - json.dump(model_index, file) + with open(os.path.join(out_dir, "model_index.json"), 'w') as file: + json.dump(model_index, file) - return out_dir - except Exception as e: - log.error(f"Failed to optimize model '{self.original_filename}'.") - log.error(e) # for test. - shutil.rmtree(shared.opts.onnx_temp_dir, ignore_errors=True) - shutil.rmtree(out_dir, ignore_errors=True) - return None + return out_dir def preprocess(self, p: StableDiffusionProcessing): - in_dir = self.path if os.path.isdir(self.path) else shared.opts.onnx_temp_dir disable_classifier_free_guidance = p.cfg_scale < 0.01 config.from_diffusers_cache = self.from_diffusers_cache - if self._is_sdxl and not shared.opts.diffusers_vae_upcast: - log.info("ONNX: VAE override set: id=madebyollin/sdxl-vae-fp16-fix, subfolder=") - config.vae_id = "madebyollin/sdxl-vae-fp16-fix" - config.vae_subfolder = "" - config.is_sdxl = self._is_sdxl + config.vae = os.path.join(models_path, "VAE", shared.opts.sd_vae) + if not os.path.isfile(config.vae): + del config.vae + config.vae_sdxl_fp16_fix = self._is_sdxl and not shared.opts.diffusers_vae_upcast + config.width = p.width config.height = p.height config.batch_size = p.batch_size @@ -389,39 +378,66 @@ class OnnxRawPipeline(OnnxPipelineBase): config.time_ids_size = 5 if not disable_classifier_free_guidance and "turbo" in str(self.path).lower(): - log.warning("It looks like you are trying to run a Turbo model with CFG Scale, which will lead to 'size mismatch' or 'unexpected parameter' error.") + log.warning("ONNX: It looks like you are trying to run a Turbo model with CFG Scale, which will lead to 'size mismatch' or 'unexpected parameter' error.") + + try: + converted_dir = self.convert( + (SUBMODELS_SDXL_REFINER if self.is_refiner else SUBMODELS_SDXL) if self._is_sdxl else SUBMODELS_SD, + self.path if os.path.isdir(self.path) else shared.opts.onnx_temp_dir + ) + except Exception: + log.error(f'ONNX: Failed to convert model: model={self.original_filename}') + shutil.rmtree(shared.opts.onnx_temp_dir, ignore_errors=True) + shutil.rmtree(os.path.join(shared.opts.onnx_cached_models_path, self.original_filename), ignore_errors=True) + return kwargs = { "provider": get_provider(), } - - converted_dir = self.convert(in_dir) - if converted_dir is None: - log.error('Failed to convert model.') - return out_dir = converted_dir - if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "olive-ai": + submodels_for_olive = [] + + if shared.opts.cuda_compile_backend == "olive-ai": + if "Text Encoder" in shared.opts.cuda_compile: + if not self.is_refiner: + submodels_for_olive.append("text_encoder") + if self._is_sdxl: + submodels_for_olive.append("text_encoder_2") + if "Model" in shared.opts.cuda_compile: + submodels_for_olive.append("unet") + if "VAE" in shared.opts.cuda_compile: + submodels_for_olive.append("vae_encoder") + submodels_for_olive.append("vae_decoder") + + if len(submodels_for_olive) == 0: + log.warning("Olive: Skipping olive run.") + else: log.warning("Olive implementation is experimental. It contains potentially an issue and is subject to change at any time.") + + in_dir = converted_dir + if p.width != p.height: - log.warning("Olive: different width and height are detected. The quality of the result is not guaranteed.") + log.warning("Olive: Different width and height are detected. The quality of the result is not guaranteed.") + if shared.opts.olive_static_dims: sess_options = DynamicSessionOptions() - sess_options_config = { + sess_options.enable_static_dims({ "is_sdxl": self._is_sdxl, "is_refiner": self.is_refiner, "hidden_batch_size": p.batch_size if disable_classifier_free_guidance else p.batch_size * 2, "height": p.height, "width": p.width, - } - sess_options.enable_static_dims(sess_options_config) + }) kwargs["sess_options"] = sess_options - optimized_dir = self.optimize(converted_dir) - if optimized_dir is None: - log.error('Olive: failed to optimize pipeline. The generation will fall back to unoptimized one.') - return self.derive_properties(load_pipeline(self.constructor, converted_dir, **kwargs)) - out_dir = optimized_dir + + try: + out_dir = self.run_olive(submodels_for_olive, in_dir) + except Exception: + log.error(f"Olive: Failed to run olive passes: model='{self.original_filename}'.") + shutil.rmtree(shared.opts.onnx_temp_dir, ignore_errors=True) + shutil.rmtree(os.path.join(shared.opts.onnx_cached_models_path, self.original_filename), ignore_errors=True) pipeline = self.derive_properties(load_pipeline(self.constructor, out_dir, **kwargs)) diff --git a/modules/onnx_utils.py b/modules/onnx_utils.py index 501cd6339..73919d121 100644 --- a/modules/onnx_utils.py +++ b/modules/onnx_utils.py @@ -8,10 +8,12 @@ import onnxruntime as ort def extract_device(args: List, kwargs: Dict): device = kwargs.get("device", None) + if device is None: for arg in args: if isinstance(arg, torch.device): device = arg + return device @@ -22,6 +24,7 @@ def move_inference_session(session: ort.InferenceSession, device: torch.device): previous_provider = session._providers provider = TORCH_DEVICE_TO_EP[device.type] if device.type in TORCH_DEVICE_TO_EP else previous_provider path = session._model_path + if provider is not None: try: return diffusers.OnnxRuntimeModel.load_model(path, provider, DynamicSessionOptions.from_sess_options(session._sess_options)) @@ -32,15 +35,19 @@ def move_inference_session(session: ort.InferenceSession, device: torch.device): def load_init_dict(cls: Type[diffusers.DiffusionPipeline], path: os.PathLike): merged: Dict[str, Any] = {} extracted = cls.extract_init_dict(diffusers.DiffusionPipeline.load_config(path)) + for dict in extracted: merged.update(dict) + merged = merged.items() R: Dict[str, Tuple[str]] = {} + for k, v in merged: if isinstance(v, list): if k not in cls.__init__.__annotations__: continue R[k] = v + return R @@ -56,24 +63,33 @@ def check_pipeline_sdxl(cls: Type[diffusers.DiffusionPipeline]) -> bool: def check_cache_onnx(path: os.PathLike) -> bool: if not os.path.isdir(path): return False + init_dict_path = os.path.join(path, "model_index.json") + if not os.path.isfile(init_dict_path): return False + init_dict = None + with open(init_dict_path, "r") as file: init_dict = file.read() + if "OnnxRuntimeModel" not in init_dict: return False + return True def load_submodel(path: os.PathLike, is_sdxl: bool, submodel_name: str, item: List[Union[str, None]], **kwargs_ort): lib, atr = item + if lib is None or atr is None: return None + library = importlib.import_module(lib) attribute = getattr(library, atr) path = os.path.join(path, submodel_name) + if issubclass(attribute, diffusers.OnnxRuntimeModel): return diffusers.OnnxRuntimeModel.load_model( os.path.join(path, "model.onnx"), @@ -82,11 +98,13 @@ def load_submodel(path: os.PathLike, is_sdxl: bool, submodel_name: str, item: Li path, **kwargs_ort, ) + return attribute.from_pretrained(path) def load_submodels(path: os.PathLike, is_sdxl: bool, init_dict: Dict[str, Type], **kwargs_ort): loaded = {} + for k, v in init_dict.items(): if not isinstance(v, list): loaded[k] = v @@ -95,6 +113,7 @@ def load_submodels(path: os.PathLike, is_sdxl: bool, init_dict: Dict[str, Type], loaded[k] = load_submodel(path, is_sdxl, k, v, **kwargs_ort) except Exception: pass + return loaded @@ -102,6 +121,7 @@ def patch_kwargs(cls: Type[diffusers.DiffusionPipeline], kwargs: Dict) -> Dict: if cls == diffusers.OnnxStableDiffusionPipeline or cls == diffusers.OnnxStableDiffusionImg2ImgPipeline or cls == diffusers.OnnxStableDiffusionInpaintPipeline: kwargs["safety_checker"] = None kwargs["requires_safety_checker"] = False + if cls == diffusers.OnnxStableDiffusionXLPipeline or cls == diffusers.OnnxStableDiffusionXLImg2ImgPipeline: kwargs["config"] = {} @@ -118,6 +138,8 @@ def load_pipeline(cls: Type[diffusers.DiffusionPipeline], path: os.PathLike, **k def get_base_constructor(cls: Type[diffusers.DiffusionPipeline], is_refiner: bool): if cls == diffusers.OnnxStableDiffusionImg2ImgPipeline or cls == diffusers.OnnxStableDiffusionInpaintPipeline: return diffusers.OnnxStableDiffusionPipeline + if cls == diffusers.OnnxStableDiffusionXLImg2ImgPipeline and not is_refiner: return diffusers.OnnxStableDiffusionXLPipeline + return cls diff --git a/modules/shared.py b/modules/shared.py index 1fe026b42..547d33944 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -444,6 +444,7 @@ options_templates.update(options_section(('diffusers', "Diffusers Settings"), { "onnx_sep": OptionInfo("

ONNX Runtime

", "", gr.HTML), "onnx_execution_provider": OptionInfo(get_default_execution_provider().value, 'Execution Provider', gr.Dropdown, lambda: {"choices": available_execution_providers }), + "onnx_show_menu": OptionInfo(False, 'ONNX show onnx-specific menu (restart required)'), "onnx_cache_converted": OptionInfo(True, 'ONNX cache converted models'), "onnx_unload_base": OptionInfo(False, 'ONNX unload base model when processing refiner'), })) diff --git a/modules/ui.py b/modules/ui.py index 6b94fafcb..545dce933 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -368,6 +368,13 @@ def create_ui(startup_timer = None): interfaces += [(interrogate_interface, "Interrogate", "interrogate")] interfaces += [(train_interface, "Train", "train")] interfaces += [(models_interface, "Models", "models")] + if shared.opts.onnx_show_menu: + with gr.Blocks(analytics_enabled=False) as onnx_interface: + if shared.backend == shared.Backend.DIFFUSERS: + from modules import ui_onnx + ui_onnx.create_ui() + timer.startup.record("ui-onnx") + interfaces += [(onnx_interface, "ONNX", "onnx")] interfaces += script_callbacks.ui_tabs_callback() interfaces += [(settings_interface, "System", "system")] diff --git a/modules/ui_onnx.py b/modules/ui_onnx.py index 00dea2203..54112657e 100644 --- a/modules/ui_onnx.py +++ b/modules/ui_onnx.py @@ -19,7 +19,6 @@ def create_ui(): from modules.paths import sd_configs_path from modules.onnx_ep import ExecutionProvider, install_execution_provider from modules.onnx_utils import check_diffusers_cache - from modules.olive import config as olive_config with gr.Blocks(analytics_enabled=False) as ui: with gr.Row(): @@ -48,25 +47,6 @@ def create_ui(): ep_install.click(fn=install_execution_provider, inputs=ep_checkbox) - with gr.TabItem("Override VAE", id="force_vae"): - gr.Markdown("Ignore baked-in vae and replace it with what you want.") - - onnx_vae_id = gr.Textbox(label="Huggingface VAE ID", info="Leave empty for default (baked-in vae).", value="") - onnx_vae_subfolder = gr.Textbox(label="VAE subfolder", info="Leave empty for root. Default: vae", value="vae") - onnx_vae_apply_button = gr.Button(value="Apply") - - def onnx_vae_apply(id: str, subfolder: str): - olive_config.vae_id = id - olive_config.vae_subfolder = subfolder - if id == "": - log.info("ONNX: VAE override unset.") - del olive_config.vae_id - olive_config.vae_subfolder = "vae" - else: - log.info(f"ONNX: VAE override set: id={id}, subfolder={subfolder}") - - onnx_vae_apply_button.click(fn=onnx_vae_apply, inputs=[onnx_vae_id, onnx_vae_subfolder,]) - if opts.cuda_compile_backend == "olive-ai": import olive.passes as olive_passes from olive.hardware.accelerator import AcceleratorSpec, Device