mirror of
https://github.com/vladmandic/automatic
synced 2026-09-20 01:31:13 +02:00
Add ui for olive
This commit is contained in:
Seunghoon Lee
@@ -1,3 +1,4 @@
|
||||
import sys
|
||||
from enum import Enum
|
||||
from typing import Tuple, List
|
||||
import onnxruntime as ort
|
||||
@@ -69,3 +70,45 @@ def get_provider() -> Tuple:
|
||||
from modules.shared import opts
|
||||
|
||||
return (opts.onnx_execution_provider, get_execution_provider_options(),)
|
||||
|
||||
|
||||
def install_execution_provider(ep: ExecutionProvider):
|
||||
from installer import pip, uninstall, installed
|
||||
from modules.shared import log
|
||||
|
||||
if installed("onnxruntime"):
|
||||
uninstall("onnxruntime")
|
||||
if installed("onnxruntime-directml"):
|
||||
uninstall("onnxruntime-directml")
|
||||
if installed("onnxruntime-gpu"):
|
||||
uninstall("onnxruntime-gpu")
|
||||
if installed("onnxruntime-training"):
|
||||
uninstall("onnxruntime-training")
|
||||
if installed("onnxruntime-openvino"):
|
||||
uninstall("onnxruntime-openvino")
|
||||
|
||||
packages = ["onnxruntime"] # Failed to load olive: cannot import name '__version__' from 'onnxruntime'
|
||||
|
||||
if ep == ExecutionProvider.DirectML:
|
||||
packages.append("onnxruntime-directml")
|
||||
elif ep == ExecutionProvider.CUDA:
|
||||
packages.append("onnxruntime-gpu")
|
||||
elif ep == ExecutionProvider.ROCm:
|
||||
if "linux" not in sys.platform:
|
||||
log.warn("ROCMExecutionProvider is not supported on Windows.")
|
||||
return
|
||||
|
||||
try:
|
||||
major, minor = sys.version_info
|
||||
cp_str = f"{major}{minor}"
|
||||
packages.append(f"https://download.onnxruntime.ai/onnxruntime_training-1.16.3%2Brocm56-cp{cp_str}-cp{cp_str}-manylinux_2_17_x86_64.manylinux2014_x86_64.whl")
|
||||
except Exception:
|
||||
log.warn("Failed to install onnxruntime for ROCm.")
|
||||
elif ep == ExecutionProvider.OpenVINO:
|
||||
if installed("openvino"):
|
||||
uninstall("openvino")
|
||||
packages.append("openvino")
|
||||
packages.append("onnxruntime-openvino")
|
||||
|
||||
pip(f"install --upgrade {' '.join(packages)}")
|
||||
log.info("Please restart SD.Next.")
|
||||
|
||||
@@ -72,23 +72,17 @@ class OnnxRawPipeline(OnnxPipelineBase):
|
||||
def __init__(self, constructor: Type[OnnxPipelineBase], path: os.PathLike):
|
||||
self.model_type = constructor.__name__
|
||||
self._is_sdxl = check_pipeline_sdxl(constructor)
|
||||
self.is_refiner = self._is_sdxl and "Img2Img" in diffusers.DiffusionPipeline.load_config(path)["_class_name"]
|
||||
self.from_huggingface_cache = shared.opts.diffusers_dir in os.path.abspath(path)
|
||||
self.path = path
|
||||
self.original_filename = os.path.basename(path)
|
||||
|
||||
self.constructor = construct_refiner_pipeline if self.is_refiner else constructor
|
||||
self.submodels = (submodels_sdxl_refiner if self.is_refiner else submodels_sdxl) if self._is_sdxl else submodels_sd
|
||||
if os.path.isdir(path):
|
||||
self.is_refiner = self._is_sdxl and "Img2Img" in diffusers.DiffusionPipeline.load_config(path)["_class_name"]
|
||||
self.init_dict = load_init_dict(constructor, path)
|
||||
self.scheduler = load_submodel(self.path, None, "scheduler", self.init_dict["scheduler"])
|
||||
else:
|
||||
try:
|
||||
cls = None
|
||||
if self._is_sdxl:
|
||||
cls = diffusers.StableDiffusionXLPipeline
|
||||
else:
|
||||
cls = diffusers.StableDiffusionPipeline
|
||||
cls = diffusers.StableDiffusionXLPipeline if self._is_sdxl else diffusers.StableDiffusionPipeline
|
||||
pipeline = cls.from_single_file(path)
|
||||
self.scheduler = pipeline.scheduler
|
||||
if os.path.isdir(shared.opts.onnx_temp_dir):
|
||||
@@ -96,12 +90,16 @@ class OnnxRawPipeline(OnnxPipelineBase):
|
||||
os.mkdir(shared.opts.onnx_temp_dir)
|
||||
pipeline.save_pretrained(shared.opts.onnx_temp_dir)
|
||||
del pipeline
|
||||
self.is_refiner = self._is_sdxl and "Img2Img" in diffusers.DiffusionPipeline.load_config(shared.opts.onnx_temp_dir)["_class_name"]
|
||||
self.init_dict = load_init_dict(constructor, shared.opts.onnx_temp_dir)
|
||||
except Exception:
|
||||
log.error('Failed to load pipeline to optimize.')
|
||||
if "vae" in self.init_dict:
|
||||
del self.init_dict["vae"]
|
||||
|
||||
self.constructor = construct_refiner_pipeline if self.is_refiner else constructor
|
||||
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
|
||||
pipeline.sd_checkpoint_info = self.sd_checkpoint_info
|
||||
@@ -237,9 +235,9 @@ class OnnxRawPipeline(OnnxPipelineBase):
|
||||
optimized_model_paths = {}
|
||||
|
||||
for submodel in self.submodels:
|
||||
log.info(f"\nOptimizing {submodel}")
|
||||
log.info(f"\nProcessing {submodel}")
|
||||
|
||||
with open(os.path.join(sd_configs_path, "olive", f"{'sdxl' if self._is_sdxl else 'sd'}_{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") as config_file:
|
||||
olive_config = json.load(config_file)
|
||||
pass_key = f"optimize_{shared.opts.onnx_execution_provider}"
|
||||
olive_config["pass_flows"] = [[pass_key]]
|
||||
|
||||
@@ -376,6 +376,12 @@ def create_ui(startup_timer = None):
|
||||
interfaces += [(extensions_interface, "Extensions", "extensions")]
|
||||
timer.startup.record("ui-extensions")
|
||||
|
||||
if shared.opts.diffusers_pipeline.startswith("ONNX"):
|
||||
from modules import ui_onnx
|
||||
|
||||
onnx_interface = ui_onnx.create_ui()
|
||||
interfaces += [(onnx_interface, "ONNX", "onnx")]
|
||||
|
||||
shared.tab_names = []
|
||||
for _interface, label, _ifid in interfaces:
|
||||
shared.tab_names.append(label)
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
import os
|
||||
import json
|
||||
from typing import Dict, List, Union
|
||||
import gradio as gr
|
||||
from olive.passes import REGISTRY
|
||||
|
||||
|
||||
def get_recursively(d: Union[Dict, List], *args):
|
||||
if len(args) == 0:
|
||||
return d
|
||||
return get_recursively(d.get(args[0]), *args[1:])
|
||||
|
||||
|
||||
def create_ui():
|
||||
from modules.ui_components import DropdownMulti
|
||||
from modules.shared import log, opts, cmd_opts
|
||||
from modules.paths import sd_configs_path
|
||||
from modules.onnx_ep import ExecutionProvider, install_execution_provider
|
||||
|
||||
with gr.Blocks(analytics_enabled=False) as ui:
|
||||
with gr.Row():
|
||||
with gr.Tabs(elem_id="tabs_onnx"):
|
||||
with gr.TabItem("Manage execution providers", id="onnxep"):
|
||||
choices = []
|
||||
|
||||
for ep in ExecutionProvider:
|
||||
choices.append(ep)
|
||||
|
||||
ep_default = None
|
||||
if cmd_opts.use_directml:
|
||||
ep_default = ExecutionProvider.DirectML
|
||||
elif cmd_opts.use_cuda:
|
||||
ep_default = ExecutionProvider.CUDA
|
||||
elif cmd_opts.use_rocm:
|
||||
ep_default = ExecutionProvider.ROCm
|
||||
elif cmd_opts.use_openvino:
|
||||
ep_default = ExecutionProvider.OpenVINO
|
||||
|
||||
ep_checkbox = gr.Radio(label="Execution provider", value=ep_default, choices=choices)
|
||||
ep_install = gr.Button(value="Install")
|
||||
gr.Text("Warning! If you are trying to reinstall, it may not work due to permission issue.")
|
||||
|
||||
ep_install.click(fn=install_execution_provider, inputs=ep_checkbox)
|
||||
|
||||
if opts.cuda_compile_backend == "olive-ai":
|
||||
with gr.Tabs(elem_id="tabs_olive"):
|
||||
with gr.TabItem("Customize pass flow", id="pass_flow"):
|
||||
with gr.Tabs(elem_id="tabs_model_type"):
|
||||
with gr.TabItem("Stable Diffusion", id="sd"):
|
||||
sd_config_path = os.path.join(sd_configs_path, "olive", "sd")
|
||||
sd_submodels = os.listdir(sd_config_path)
|
||||
sd_configs: Dict[str, Dict] = {}
|
||||
|
||||
with gr.Tabs(elem_id="tabs_sd_submodel"):
|
||||
def sd_create_change_listener(*args):
|
||||
def listener(v: Dict):
|
||||
get_recursively(sd_configs, *args[:-1])[args[-1]] = v
|
||||
return listener
|
||||
|
||||
for submodel in sd_submodels:
|
||||
config: Dict = None
|
||||
|
||||
with open(os.path.join(sd_config_path, submodel), "r") as file:
|
||||
config = json.load(file)
|
||||
sd_configs[submodel] = config
|
||||
|
||||
submodel_name = submodel[:-5]
|
||||
with gr.TabItem(submodel_name, id=f"sd_{submodel_name}"):
|
||||
pass_flows = DropdownMulti(label="Pass flow", value=sd_configs[submodel]["pass_flows"][0], choices=sd_configs[submodel]["passes"].keys())
|
||||
pass_flows.change(fn=sd_create_change_listener(submodel, "pass_flows", 0), inputs=pass_flows)
|
||||
|
||||
with gr.Tabs(elem_id=f"tabs_sd_{submodel_name}_pass"):
|
||||
for k in sd_configs[submodel]["passes"]:
|
||||
with gr.TabItem(k, id=f"sd_{submodel_name}_pass_{k}"):
|
||||
pass_type = gr.Dropdown(label="Type", value=sd_configs[submodel]["passes"][k]["type"], choices=(x.__name__ for x in tuple(REGISTRY.values())))
|
||||
|
||||
pass_type.change(fn=sd_create_change_listener(submodel, "passes", k, "type"), inputs=pass_type)
|
||||
|
||||
def sd_save():
|
||||
for k, v in sd_configs.items():
|
||||
with open(os.path.join(sd_config_path, k), "w") as file:
|
||||
json.dump(v, file)
|
||||
log.info("Olive: config for SD was saved.")
|
||||
|
||||
sd_save_button = gr.Button(value="Save")
|
||||
sd_save_button.click(fn=sd_save)
|
||||
|
||||
with gr.TabItem("Stable Diffusion XL", id="sdxl"):
|
||||
sdxl_config_path = os.path.join(sd_configs_path, "olive", "sdxl")
|
||||
sdxl_submodels = os.listdir(sdxl_config_path)
|
||||
sdxl_configs: Dict[str, Dict] = {}
|
||||
|
||||
with gr.Tabs(elem_id="tabs_sdxl_submodel"):
|
||||
def sdxl_create_change_listener(*args):
|
||||
def listener(v: Dict):
|
||||
get_recursively(sdxl_configs, *args[:-1])[args[-1]] = v
|
||||
return listener
|
||||
|
||||
for submodel in sdxl_submodels:
|
||||
config: Dict = None
|
||||
|
||||
with open(os.path.join(sdxl_config_path, submodel), "r") as file:
|
||||
config = json.load(file)
|
||||
sdxl_configs[submodel] = config
|
||||
|
||||
submodel_name = submodel[:-5]
|
||||
with gr.TabItem(submodel_name, id=f"sdxl_{submodel_name}"):
|
||||
pass_flows = DropdownMulti(label="Pass flow", value=sdxl_configs[submodel]["pass_flows"][0], choices=sdxl_configs[submodel]["passes"].keys())
|
||||
pass_flows.change(fn=sdxl_create_change_listener(submodel, "pass_flows", 0), inputs=pass_flows)
|
||||
|
||||
with gr.Tabs(elem_id=f"tabs_sdxl_{submodel_name}_pass"):
|
||||
for k in sdxl_configs[submodel]["passes"]:
|
||||
with gr.TabItem(k, id=f"sdxl_{submodel_name}_pass_{k}"):
|
||||
pass_type = gr.Dropdown(label="Type", value=sdxl_configs[submodel]["passes"][k]["type"], choices=(x.__name__ for x in tuple(REGISTRY.values())))
|
||||
|
||||
pass_type.change(fn=sdxl_create_change_listener(submodel, "passes", k, "type"), inputs=pass_type)
|
||||
|
||||
def sdxl_save():
|
||||
for k, v in sdxl_configs.items():
|
||||
with open(os.path.join(sdxl_config_path, k), "w") as file:
|
||||
json.dump(v, file)
|
||||
log.info("Olive: config for SDXL was saved.")
|
||||
|
||||
sdxl_save_button = gr.Button(value="Save")
|
||||
sdxl_save_button.click(fn=sdxl_save)
|
||||
return ui
|
||||
Reference in New Issue
Block a user