From 192b9db5b35107c310a9ea4402fbf58c98d0e8cd Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Tue, 30 Jan 2024 15:01:35 +0900 Subject: [PATCH] cleanup --- modules/olive.py | 26 +++++++++++++++----------- modules/onnx.py | 23 ++++++++++------------- modules/onnx_ep.py | 4 ++-- modules/onnx_pipelines.py | 29 ++++++++++++++--------------- modules/onnx_utils.py | 10 +++++----- 5 files changed, 46 insertions(+), 46 deletions(-) diff --git a/modules/olive.py b/modules/olive.py index e6726c56a..c805fc92e 100644 --- a/modules/olive.py +++ b/modules/olive.py @@ -1,7 +1,7 @@ import os +from typing import Type, Callable, TypeVar, Dict, Any import torch import diffusers -from typing import Type, Callable, TypeVar, Dict, Any from transformers.models.clip.modeling_clip import CLIPTextModel, CLIPTextModelWithProjection @@ -62,10 +62,14 @@ config = OliveOptimizerConfig() def get_variant(): from modules.shared import opts + if opts.diffusers_model_load_variant == 'default': from modules import devices + if devices.dtype == torch.float16: return 'fp16' + + return None elif opts.diffusers_model_load_variant == 'fp32': return None else: @@ -114,7 +118,7 @@ class RandomDataLoader: # ----------------------------------------------------------------------------- -def text_encoder_inputs(_, torch_dtype): +def text_encoder_inputs(batchsize, torch_dtype): input_ids = torch.zeros((config.batch_size, 77), dtype=torch_dtype) return { "input_ids": input_ids, @@ -131,7 +135,7 @@ def text_encoder_conversion_inputs(model): return text_encoder_inputs(1, torch.int32) -def text_encoder_data_loader(data_dir, _, *args, **kwargs): +def text_encoder_data_loader(data_dir, batchsize, *_, **__): return RandomDataLoader(text_encoder_inputs, config.batch_size, torch.int32) @@ -140,7 +144,7 @@ def text_encoder_data_loader(data_dir, _, *args, **kwargs): # ----------------------------------------------------------------------------- -def text_encoder_2_inputs(_, torch_dtype): +def text_encoder_2_inputs(batchsize, torch_dtype): return { "input_ids": torch.zeros((config.batch_size, 77), dtype=torch_dtype), "output_hidden_states": True, @@ -156,7 +160,7 @@ def text_encoder_2_conversion_inputs(model): return text_encoder_2_inputs(1, torch.int64) -def text_encoder_2_data_loader(data_dir, _, *args, **kwargs): +def text_encoder_2_data_loader(data_dir, batchsize, *_, **__): return RandomDataLoader(text_encoder_2_inputs, config.batch_size, torch.int64) @@ -165,7 +169,7 @@ def text_encoder_2_data_loader(data_dir, _, *args, **kwargs): # ----------------------------------------------------------------------------- -def unet_inputs(_, torch_dtype, is_conversion_inputs=False): +def unet_inputs(batchsize, torch_dtype, is_conversion_inputs=False): if config.is_sdxl: inputs = { "sample": torch.rand((2 * config.batch_size, 4, config.height // 8, config.width // 8), dtype=torch_dtype), @@ -219,7 +223,7 @@ def unet_conversion_inputs(model): return tuple(unet_inputs(1, torch.float32, True).values()) -def unet_data_loader(data_dir, _, *args, **kwargs): +def unet_data_loader(data_dir, batchsize, *_, **__): return RandomDataLoader(unet_inputs, config.batch_size, torch.float16) @@ -228,7 +232,7 @@ def unet_data_loader(data_dir, _, *args, **kwargs): # ----------------------------------------------------------------------------- -def vae_encoder_inputs(_, torch_dtype): +def vae_encoder_inputs(batchsize, torch_dtype): return { "sample": torch.rand((config.batch_size, 3, config.height, config.width), dtype=torch_dtype), "return_dict": False, @@ -256,7 +260,7 @@ def vae_encoder_conversion_inputs(model): return tuple(vae_encoder_inputs(1, torch.float32).values()) -def vae_encoder_data_loader(data_dir, _, *args, **kwargs): +def vae_encoder_data_loader(data_dir, batchsize, *_, **__): return RandomDataLoader(vae_encoder_inputs, config.batch_size, torch.float16) @@ -265,7 +269,7 @@ def vae_encoder_data_loader(data_dir, _, *args, **kwargs): # ----------------------------------------------------------------------------- -def vae_decoder_inputs(_, torch_dtype): +def vae_decoder_inputs(batchsize, torch_dtype): return { "latent_sample": torch.rand((config.batch_size, 4, config.height // 8, config.width // 8), dtype=torch_dtype), "return_dict": False, @@ -293,5 +297,5 @@ def vae_decoder_conversion_inputs(model): return tuple(vae_decoder_inputs(1, torch.float32).values()) -def vae_decoder_data_loader(data_dir, _, *args, **kwargs): +def vae_decoder_data_loader(data_dir, batchsize, *_, **__): return RandomDataLoader(vae_decoder_inputs, config.batch_size, torch.float16) diff --git a/modules/onnx.py b/modules/onnx.py index 00af55ddc..6001c9f91 100644 --- a/modules/onnx.py +++ b/modules/onnx.py @@ -1,13 +1,10 @@ +from typing import Any, Dict, Optional import torch import diffusers import onnxruntime as ort -from typing import Any, Dict, Optional initialized = False -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 DynamicSessionOptions(ort.SessionOptions): @@ -46,27 +43,27 @@ class DynamicSessionOptions(ort.SessionOptions): return sess_options -class OnnxFakeModule: +class TorchCompatibleModule: device = torch.device("cpu") dtype = torch.float32 - def to(self, *args, **kwargs): + def to(self, *_, **__): return self - def type(self, *args, **kwargs): + def type(self, *_, **__): return self -class OnnxTemporalModel(OnnxFakeModule): +class TemporalModule(TorchCompatibleModule): """ Replace the models which are not able to be moved to CPU. """ - previous_provider: Any + provider: Any path: str sess_options: ort.SessionOptions - def __init__(self, previous_provider: Any, path: str, sess_options: ort.SessionOptions): - self.previous_provider = previous_provider + def __init__(self, provider: Any, path: str, sess_options: ort.SessionOptions): + self.provider = provider self.path = path self.sess_options = sess_options @@ -77,12 +74,12 @@ class OnnxTemporalModel(OnnxFakeModule): if device is not None and device.type != "cpu": from modules.onnx_ep import TORCH_DEVICE_TO_EP - provider = TORCH_DEVICE_TO_EP[device.type] if device.type in TORCH_DEVICE_TO_EP else self.previous_provider + provider = TORCH_DEVICE_TO_EP[device.type] if device.type in TORCH_DEVICE_TO_EP else self.provider return OnnxRuntimeModel.load_model(self.path, provider, DynamicSessionOptions.from_sess_options(self.sess_options)) return self -class OnnxRuntimeModel(OnnxFakeModule, diffusers.OnnxRuntimeModel): +class OnnxRuntimeModel(TorchCompatibleModule, diffusers.OnnxRuntimeModel): config = {} # dummy def named_modules(self): # dummy diff --git a/modules/onnx_ep.py b/modules/onnx_ep.py index b50e05ee6..0ae94b67b 100644 --- a/modules/onnx_ep.py +++ b/modules/onnx_ep.py @@ -100,10 +100,10 @@ def install_execution_provider(ep: ExecutionProvider): packages.append("onnxruntime-gpu") elif ep == ExecutionProvider.ROCm: if "linux" not in sys.platform: - log.warn("ROCMExecutionProvider is not supported on Windows.") + log.warning("ROCMExecutionProvider is not supported on Windows.") return - packages.append(get_onnxruntime_source_for_rocm()) + packages.append(get_onnxruntime_source_for_rocm(None)) elif ep == ExecutionProvider.OpenVINO: if installed("openvino"): uninstall("openvino") diff --git a/modules/onnx_pipelines.py b/modules/onnx_pipelines.py index 93b53d0f5..e44ba8ba9 100644 --- a/modules/onnx_pipelines.py +++ b/modules/onnx_pipelines.py @@ -1,15 +1,15 @@ import os import json -import torch import shutil import inspect +from abc import ABCMeta +from typing import Union, Optional, Callable, Type, Tuple, List, Any, Dict from packaging import version +import torch import numpy as np import diffusers import onnxruntime as ort import optimum.onnxruntime -from abc import ABCMeta -from typing import Union, Optional, Callable, Type, Tuple, List, Any, Dict from diffusers.pipelines.onnx_utils import ORT_TO_NP_TYPE from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput from diffusers.image_processor import VaeImageProcessor, PipelineImageInput @@ -19,7 +19,7 @@ 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 -from modules.onnx import DynamicSessionOptions, OnnxFakeModule, submodels_sd, submodels_sdxl, submodels_sdxl_refiner +from modules.onnx import DynamicSessionOptions, TorchCompatibleModule from modules.onnx_utils import extract_device, move_inference_session, check_diffusers_cache, check_pipeline_sdxl, check_cache_onnx, load_init_dict, load_submodel, load_submodels, patch_kwargs, load_pipeline, get_base_constructor from modules.onnx_ep import ExecutionProvider, EP_TO_NAME, get_provider @@ -45,7 +45,7 @@ CONVERSION_PASS_UNET = { } -class OnnxPipelineBase(OnnxFakeModule, diffusers.DiffusionPipeline, metaclass=ABCMeta): +class OnnxPipelineBase(TorchCompatibleModule, diffusers.DiffusionPipeline, metaclass=ABCMeta): model_type: str sd_model_hash: str sd_checkpoint_info: CheckpointInfo @@ -145,7 +145,6 @@ class OnnxRawPipeline(OnnxPipelineBase): self.is_refiner = self._is_sdxl and "Img2Img" not in constructor.__name__ and "Img2Img" in diffusers.DiffusionPipeline.load_config(path)["_class_name"] self.constructor = OnnxStableDiffusionXLImg2ImgPipeline if self.is_refiner else constructor self.model_type = self.constructor.__name__ - self.submodels = (submodels_sdxl_refiner if self.is_refiner else submodels_sdxl) if self._is_sdxl else submodels_sd def derive_properties(self, pipeline: diffusers.DiffusionPipeline): pipeline.sd_model_hash = self.sd_model_hash @@ -183,7 +182,7 @@ class OnnxRawPipeline(OnnxPipelineBase): for submodel in submodels: log.info(f"\nConverting {submodel}") - with open(os.path.join(sd_configs_path, "olive", 'sdxl' if self._is_sdxl else 'sd', f"{submodel}.json"), "r") as config_file: + with open(os.path.join(sd_configs_path, "olive", 'sdxl' if self._is_sdxl else 'sd', f"{submodel}.json"), "r", encoding="utf-8") as config_file: conversion_config = json.load(config_file) conversion_config["input_model"]["config"]["model_path"] = os.path.abspath(in_dir) conversion_config["passes"] = { @@ -194,7 +193,7 @@ class OnnxRawPipeline(OnnxPipelineBase): 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: + with open(os.path.join("footprints", f"{submodel}_{EP_TO_NAME[shared.opts.onnx_execution_provider]}_footprints.json"), "r", encoding="utf-8") as footprint_file: footprints = json.load(footprint_file) conversion_footprint = None for _, footprint in footprints.items(): @@ -254,7 +253,7 @@ class OnnxRawPipeline(OnnxPipelineBase): if k not in model_index: model_index[k] = v - with open(os.path.join(out_dir, "model_index.json"), 'w') as file: + with open(os.path.join(out_dir, "model_index.json"), 'w', encoding="utf-8") as file: json.dump(model_index, file) return out_dir @@ -289,7 +288,7 @@ class OnnxRawPipeline(OnnxPipelineBase): 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: + with open(os.path.join(sd_configs_path, "olive", 'sdxl' if self._is_sdxl else 'sd', f"{submodel}.json"), "r", encoding="utf-8") as config_file: olive_config: Dict[str, Dict[str, Dict]] = json.load(config_file) for flow in olive_config["pass_flows"]: @@ -310,7 +309,7 @@ class OnnxRawPipeline(OnnxPipelineBase): 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: + with open(os.path.join("footprints", f"{submodel}_{EP_TO_NAME[shared.opts.onnx_execution_provider]}_footprints.json"), "r", encoding="utf-8") as footprint_file: footprints = json.load(footprint_file) processor_final_pass_footprint = None for _, footprint in footprints.items(): @@ -370,7 +369,7 @@ class OnnxRawPipeline(OnnxPipelineBase): if k not in model_index: model_index[k] = v - with open(os.path.join(out_dir, "model_index.json"), 'w') as file: + with open(os.path.join(out_dir, "model_index.json"), 'w', encoding="utf-8") as file: json.dump(model_index, file) return out_dir @@ -638,7 +637,7 @@ class OnnxStableDiffusionPipeline(diffusers.OnnxStableDiffusionPipeline, OnnxPip has_nsfw_concept = None - if not output_type == "latent": + if output_type != "latent": # image = self.vae_decoder(latent_sample=latents)[0] # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1 image = np.concatenate( @@ -826,7 +825,7 @@ class OnnxStableDiffusionImg2ImgPipeline(diffusers.OnnxStableDiffusionImg2ImgPip has_nsfw_concept = None - if not output_type == "latent": + if output_type != "latent": # image = self.vae_decoder(latent_sample=latents)[0] # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1 image = np.concatenate( @@ -1040,7 +1039,7 @@ class OnnxStableDiffusionInpaintPipeline(diffusers.OnnxStableDiffusionInpaintPip has_nsfw_concept = None - if not output_type == "latent": + if output_type != "latent": # image = self.vae_decoder(latent_sample=latents)[0] # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1 image = np.concatenate( diff --git a/modules/onnx_utils.py b/modules/onnx_utils.py index 73919d121..534e08241 100644 --- a/modules/onnx_utils.py +++ b/modules/onnx_utils.py @@ -18,7 +18,7 @@ def extract_device(args: List, kwargs: Dict): def move_inference_session(session: ort.InferenceSession, device: torch.device): - from modules.onnx import DynamicSessionOptions, OnnxTemporalModel + from modules.onnx import DynamicSessionOptions, TemporalModule from modules.onnx_ep import TORCH_DEVICE_TO_EP previous_provider = session._providers @@ -29,15 +29,15 @@ def move_inference_session(session: ort.InferenceSession, device: torch.device): try: return diffusers.OnnxRuntimeModel.load_model(path, provider, DynamicSessionOptions.from_sess_options(session._sess_options)) except Exception: - return OnnxTemporalModel(previous_provider, path, session._sess_options) + return TemporalModule(previous_provider, path, session._sess_options) 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) + for item in extracted: + merged.update(item) merged = merged.items() R: Dict[str, Tuple[str]] = {} @@ -71,7 +71,7 @@ def check_cache_onnx(path: os.PathLike) -> bool: init_dict = None - with open(init_dict_path, "r") as file: + with open(init_dict_path, "r", encoding="utf-8") as file: init_dict = file.read() if "OnnxRuntimeModel" not in init_dict: