diff --git a/modules/onnx.py b/modules/onnx.py
index abc2b00ab..5c17c59ae 100644
--- a/modules/onnx.py
+++ b/modules/onnx.py
@@ -26,6 +26,36 @@ class OnnxRuntimeModel(OnnxFakeModule, diffusers.OnnxRuntimeModel):
return ()
+def optimize_pipeline(p, refiner_enabled: bool):
+ from modules import shared, sd_models
+
+ if "ONNX" not in shared.opts.diffusers_pipeline:
+ shared.log.warning(f"Unsupported pipeline for 'olive-ai' compile backend: {shared.opts.diffusers_pipeline}. You should select one of the ONNX pipelines.")
+ return
+
+ if shared.opts.cuda_compile_backend == "olive-ai":
+ compile_height = p.height
+ compile_width = p.width
+ if (shared.compiled_model_state is None or
+ shared.compiled_model_state.height != compile_height
+ or shared.compiled_model_state.width != compile_width
+ or shared.compiled_model_state.batch_size != p.batch_size):
+ shared.log.info("Olive: Parameter change detected")
+ shared.log.info("Olive: Recompiling base model")
+ sd_models.unload_model_weights(op='model')
+ sd_models.reload_model_weights(op='model')
+ if refiner_enabled:
+ shared.log.info("Olive: Recompiling refiner")
+ sd_models.unload_model_weights(op='refiner')
+ sd_models.reload_model_weights(op='refiner')
+ shared.compiled_model_state.height = compile_height
+ shared.compiled_model_state.width = compile_width
+ shared.compiled_model_state.batch_size = p.batch_size
+
+ if hasattr(shared.sd_model, "preprocess"):
+ shared.sd_model = shared.sd_model.preprocess(p)
+
+
def initialize():
global initialized
diff --git a/modules/onnx_ep.py b/modules/onnx_ep.py
index 25e28332b..4ba6c93c5 100644
--- a/modules/onnx_ep.py
+++ b/modules/onnx_ep.py
@@ -57,7 +57,7 @@ def get_execution_provider_options():
elif opts.onnx_execution_provider == ExecutionProvider.OpenVINO:
from modules.intel.openvino import get_device as get_raw_openvino_device
raw_openvino_device = get_raw_openvino_device()
- if opts.onnx_olive_float16 and not opts.openvino_hetero_gpu:
+ if opts.olive_float16 and not opts.openvino_hetero_gpu:
raw_openvino_device = f"{raw_openvino_device}_FP16"
execution_provider_options["device_type"] = raw_openvino_device
del execution_provider_options["device_id"]
diff --git a/modules/onnx_pipelines.py b/modules/onnx_pipelines.py
index 0feda0e6d..6bcbc85d8 100644
--- a/modules/onnx_pipelines.py
+++ b/modules/onnx_pipelines.py
@@ -19,6 +19,7 @@ from installer import log
from modules import shared
from modules.paths import sd_configs_path
from modules.sd_models import CheckpointInfo
+from modules.sd_models_compile import CompiledModelState
from modules.processing import StableDiffusionProcessing
from modules.olive import config
from modules.onnx import OnnxFakeModule, submodels_sd, submodels_sdxl, submodels_sdxl_refiner
@@ -35,11 +36,6 @@ class OnnxPipelineBase(OnnxFakeModule, diffusers.DiffusionPipeline, metaclass=AB
def __init__(self):
self.model_type = self.__class__.__name__
- def override_processing(self, p: StableDiffusionProcessing):
- disable_classifier_free_guidance = p.cfg_scale < 0.01 or "turbo" in self.sd_checkpoint_info.model_name.lower()
- if disable_classifier_free_guidance:
- p.cfg_scale = 0.0
-
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, **_):
return OnnxRawPipeline(
@@ -76,7 +72,7 @@ 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 = "refiner" in str(path).lower()
+ 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)
@@ -223,7 +219,7 @@ class OnnxRawPipeline(OnnxPipelineBase):
if os.path.isdir(out_dir): # already optimized (cached)
return out_dir
- if not shared.opts.onnx_cache_optimized:
+ if not shared.opts.olive_cache_optimized:
out_dir = shared.opts.onnx_temp_dir
try:
@@ -233,7 +229,7 @@ class OnnxRawPipeline(OnnxPipelineBase):
shutil.rmtree("cache", ignore_errors=True)
shutil.rmtree("footprints", ignore_errors=True)
- if shared.opts.onnx_cache_optimized:
+ if shared.opts.olive_cache_optimized:
shutil.copytree(
in_dir, out_dir, ignore=shutil.ignore_patterns("weights.pb", "*.onnx", "*.safetensors", "*.ckpt")
)
@@ -248,12 +244,12 @@ class OnnxRawPipeline(OnnxPipelineBase):
pass_key = f"optimize_{shared.opts.onnx_execution_provider}"
olive_config["pass_flows"] = [[pass_key]]
olive_config["input_model"]["config"]["model_path"] = os.path.abspath(os.path.join(in_dir, submodel, "model.onnx"))
- olive_config["passes"][pass_key]["config"]["float16"] = shared.opts.onnx_olive_float16
+ olive_config["passes"][pass_key]["config"]["float16"] = shared.opts.olive_float16
olive_config["engine"]["execution_providers"] = [shared.opts.onnx_execution_provider]
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 shared.opts.onnx_olive_float16:
+ if shared.opts.olive_float16:
olive_config["passes"][pass_key]["config"]["keep_io_types"] = False
run(olive_config)
@@ -330,6 +326,7 @@ class OnnxRawPipeline(OnnxPipelineBase):
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_huggingface_cache = self.from_huggingface_cache
config.use_fp16_fixed_vae = self._is_sdxl and not shared.opts.diffusers_vae_upcast
@@ -343,9 +340,12 @@ class OnnxRawPipeline(OnnxPipelineBase):
config.cross_attention_dim = 2048 if self._is_sdxl else (256 + p.height)
config.time_ids_size = 6 if self._is_sdxl and not self.is_refiner else 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.")
+
kwargs = {
"provider": get_provider(),
- "sess_options": get_sess_options(p.batch_size if p.cfg_scale < 0.01 or "turbo" in str(self.path).lower() else p.batch_size * 2, p.height, p.width, self._is_sdxl),
+ "sess_options": get_sess_options(p.batch_size if disable_classifier_free_guidance else p.batch_size * 2, p.height, p.width, self._is_sdxl),
}
converted_dir = self.convert(in_dir)
diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py
index 287ea6e15..39322195b 100644
--- a/modules/processing_diffusers.py
+++ b/modules/processing_diffusers.py
@@ -7,7 +7,7 @@ import torch
import torchvision.transforms.functional as TF
import diffusers
from modules import shared, devices, processing, sd_samplers, sd_models, images, errors, masking, prompt_parser_diffusers, sd_hijack_hypertile, processing_correction, processing_vae
-from modules.onnx_pipelines import OnnxStableDiffusionPipeline
+from modules.onnx import optimize_pipeline as onnx_optimize_pipeline
debug = shared.log.trace if os.environ.get('SD_DIFFUSERS_DEBUG', None) is not None else lambda *args, **kwargs: None
@@ -21,37 +21,24 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
orig_pipeline = shared.sd_model
results = []
- if hasattr(shared.sd_model, 'preprocess'):
- shared.sd_model = shared.sd_model.preprocess(p)
-
- if hasattr(shared.sd_model, 'override_processing'):
- shared.sd_model.override_processing(p)
-
def is_txt2img():
return sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.TEXT_2_IMAGE
def is_refiner_enabled():
return p.enable_hr and p.refiner_steps > 0 and p.refiner_start > 0 and p.refiner_start < 1 and shared.sd_refiner is not None
- def resize_images():
- if getattr(p, 'image', None) is not None and getattr(p, 'init_images', None) is None:
- p.init_images = [p.image]
- if getattr(p, 'init_images', None) is not None and len(p.init_images) > 0:
- tgt_width, tgt_height = 8 * math.ceil(p.init_images[0].width / 8), 8 * math.ceil(p.init_images[0].height / 8)
- if p.init_images[0].size != (tgt_width, tgt_height):
- shared.log.debug(f'Resizing init images: original={p.init_images[0].width}x{p.init_images[0].height} target={tgt_width}x{tgt_height}')
- p.init_images = [images.resize_image(1, image, tgt_width, tgt_height, upscaler_name=None) for image in p.init_images]
- p.height = tgt_height
- p.width = tgt_width
- sd_hijack_hypertile.hypertile_set(p)
- if getattr(p, 'mask', None) is not None and p.mask.size != (tgt_width, tgt_height):
- p.mask = images.resize_image(1, p.mask, tgt_width, tgt_height, upscaler_name=None)
- if getattr(p, 'init_mask', None) is not None and p.init_mask.size != (tgt_width, tgt_height):
- p.init_mask = images.resize_image(1, p.init_mask, tgt_width, tgt_height, upscaler_name=None)
- if getattr(p, 'mask_for_overlay', None) is not None and p.mask_for_overlay.size != (tgt_width, tgt_height):
- p.mask_for_overlay = images.resize_image(1, p.mask_for_overlay, tgt_width, tgt_height, upscaler_name=None)
- return tgt_width, tgt_height
- return p.width, p.height
+ if getattr(p, 'init_images', None) is not None and len(p.init_images) > 0:
+ tgt_width, tgt_height = 8 * math.ceil(p.init_images[0].width / 8), 8 * math.ceil(p.init_images[0].height / 8)
+ if p.init_images[0].width != tgt_width or p.init_images[0].height != tgt_height:
+ shared.log.debug(f'Resizing init images: original={p.init_images[0].width}x{p.init_images[0].height} target={tgt_width}x{tgt_height}')
+ p.init_images = [images.resize_image(1, image, tgt_width, tgt_height, upscaler_name=None) for image in p.init_images]
+ p.height = tgt_height
+ p.width = tgt_width
+ hypertile_set(p)
+ if getattr(p, 'mask', None) is not None and p.mask.size != (tgt_width, tgt_height):
+ p.mask = images.resize_image(1, p.mask, tgt_width, tgt_height, upscaler_name=None)
+ if getattr(p, 'mask_for_overlay', None) is not None and p.mask_for_overlay.size != (tgt_width, tgt_height):
+ p.mask_for_overlay = images.resize_image(1, p.mask_for_overlay, tgt_width, tgt_height, upscaler_name=None)
def hires_resize(latents): # input=latents output=pil
if not torch.is_tensor(latents):
@@ -226,7 +213,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
generator = [torch.Generator(generator_device).manual_seed(s) for s in p.seeds]
prompts, negative_prompts, prompts_2, negative_prompts_2 = fix_prompts(prompts, negative_prompts, prompts_2, negative_prompts_2)
parser = 'Fixed attention'
- if shared.opts.prompt_attention != 'Fixed attention' and 'StableDiffusion' in model.__class__.__name__ and not isinstance(model, OnnxStableDiffusionPipeline):
+ if shared.opts.prompt_attention != 'Fixed attention' and 'StableDiffusion' in model.__class__.__name__ and 'Onnx' not in model.__class__.__name__:
try:
prompt_parser_diffusers.encode_prompts(model, p, prompts, negative_prompts, kwargs.get("num_inference_steps", 1), kwargs.pop("clip_skip", None))
parser = shared.opts.prompt_attention
@@ -477,6 +464,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
return max(1, int(steps))
shared.sd_model = update_pipeline(shared.sd_model, p)
+ onnx_optimize_pipeline(p, is_refiner_enabled())
base_args = set_pipeline_args(
model=shared.sd_model,
prompts=p.prompts,
@@ -548,6 +536,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
if (latent_scale_mode is not None or p.hr_force) and p.denoising_strength > 0:
p.ops.append('hires')
shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE)
+ onnx_optimize_pipeline(p, is_refiner_enabled())
recompile_model(hires=True)
update_sampler(shared.sd_model, second_pass=True)
hires_args = set_pipeline_args(
diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py
index 26dc3c7a1..615cc24f5 100644
--- a/modules/sd_models_compile.py
+++ b/modules/sd_models_compile.py
@@ -152,12 +152,22 @@ def compile_stablefast(sd_model):
def compile_torch(sd_model):
+ if shared.opts.cuda_compile_backend == "olive-ai":
+ if shared.compiled_model_state is None:
+ shared.compiled_model_state = CompiledModelState()
+ return sd_model
try:
import torch._dynamo # pylint: disable=unused-import,redefined-outer-name
torch._dynamo.reset() # pylint: disable=protected-access
shared.log.debug(f"Model compile available backends: {torch._dynamo.list_backends()}") # pylint: disable=protected-access
if shared.opts.cuda_compile_backend == "openvino_fx":
optimize_openvino()
+ """
+ elif shared.opts.cuda_compile_backend == "olive-ai":
+ if shared.compiled_model_state is None:
+ shared.compiled_model_state = CompiledModelState()
+ return sd_model
+ """
log_level = logging.WARNING if shared.opts.cuda_compile_verbose else logging.CRITICAL # pylint: disable=protected-access
if hasattr(torch, '_logging'):
torch._logging.set_logs(dynamo=log_level, aot=log_level, inductor=log_level) # pylint: disable=protected-access
diff --git a/modules/shared.py b/modules/shared.py
index a2b027256..81a78729e 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -440,11 +440,6 @@ 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_cache_converted": OptionInfo(True, 'Cache converted models'),
-
- "onnx_olive_sep": OptionInfo("Olive
", "", gr.HTML),
- "onnx_enable_olive": OptionInfo(False, 'Optimize ONNX pipeline using Olive'),
- "onnx_olive_float16": OptionInfo(True, 'Olive use FP16 on optimization (will use FP32 if unchecked)'),
- "onnx_cache_optimized": OptionInfo(True, 'Olive cache optimized models'),
}))
options_templates.update(options_section(('system-paths', "System Paths"), {