Integrate Olive into compile backend.

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
 

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"]

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)

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(

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

modules/shared.py

@@ -440,11 +440,6 @@ options_templates.update(options_section(('diffusers', "Diffusers Settings"), {
     "onnx_sep": OptionInfo("<h2>ONNX Runtime</h2>", "", 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("<h3>Olive</h3>", "", 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"), {