Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-18 16:54:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

cleanup

Browse Source
This commit is contained in:
Seunghoon Lee
2024-01-30 15:01:35 +09:00
parent 250f145bca
commit 192b9db5b3
5 changed files with 46 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/olive.py
+15 -11
View File
@@ -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)
modules/onnx.py
+10 -13
View File
@@ -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
modules/onnx_ep.py
+2 -2
View File
@@ -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")
modules/onnx_pipelines.py
+14 -15
View File
@@ -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(
modules/onnx_utils.py
+5 -5
View File
@@ -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: