diff --git a/launch.py b/launch.py index 0806566c8..c03ca6e22 100755 --- a/launch.py +++ b/launch.py @@ -41,14 +41,6 @@ def init_args(): def init_paths(): global script_path, extensions_dir # pylint: disable=global-statement - try: - import olive.workflows # pylint: disable=unused-import - except ModuleNotFoundError: - pass - import modules.cmd_args - parser = modules.cmd_args.parser - installer.add_args(parser) - args, _ = parser.parse_known_args() import modules.paths modules.paths.register_paths() script_path = modules.paths.script_path diff --git a/modules/onnx_impl/__init__.py b/modules/onnx_impl/__init__.py index 40813c7a8..a12220a21 100644 --- a/modules/onnx_impl/__init__.py +++ b/modules/onnx_impl/__init__.py @@ -100,9 +100,9 @@ def preprocess_pipeline(p, refiner_enabled: bool): 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 + return shared.sd_model - if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "olive-ai": + if shared.opts.cuda_compile_backend == "olive-ai" and len(shared.opts.cuda_compile) != 1: compile_height = p.height compile_width = p.width if (shared.compiled_model_state is None or @@ -131,6 +131,8 @@ def preprocess_pipeline(p, refiner_enabled: bool): sd_models.reload_model_weights(op='model') shared.sd_model = shared.sd_model.preprocess(p) + return shared.sd_model + def initialize(): global initialized # pylint: disable=global-statement diff --git a/modules/onnx_impl/pipelines.py b/modules/onnx_impl/pipelines.py index bfebce4ac..9f389be29 100644 --- a/modules/onnx_impl/pipelines.py +++ b/modules/onnx_impl/pipelines.py @@ -370,9 +370,9 @@ class OnnxRawPipeline(PipelineBase): } out_dir = converted_dir - submodels_for_olive = [] - if shared.opts.cuda_compile_backend == "olive-ai": + submodels_for_olive = [] + if "Text Encoder" in shared.opts.cuda_compile: if not self.is_refiner: submodels_for_olive.append("text_encoder") @@ -384,34 +384,34 @@ class OnnxRawPipeline(PipelineBase): submodels_for_olive.append("vae_encoder") submodels_for_olive.append("vae_decoder") - if len(submodels_for_olive) == 0: - log.warning("Olive: Skipping olive run.") - else: - log.warning("Olive implementation is experimental. It contains potentially an issue and is subject to change at any time.") + if len(submodels_for_olive) == 0: + log.warning("Olive: Skipping olive run.") + else: + log.warning("Olive implementation is experimental. It contains potentially an issue and is subject to change at any time.") - in_dir = converted_dir + in_dir = converted_dir - if p.width != p.height: - log.warning("Olive: Different width and height are detected. The quality of the result is not guaranteed.") + if p.width != p.height: + log.warning("Olive: Different width and height are detected. The quality of the result is not guaranteed.") - if shared.opts.olive_static_dims: - sess_options = DynamicSessionOptions() - sess_options.enable_static_dims({ - "is_sdxl": self._is_sdxl, - "is_refiner": self.is_refiner, + if shared.opts.olive_static_dims: + sess_options = DynamicSessionOptions() + sess_options.enable_static_dims({ + "is_sdxl": self._is_sdxl, + "is_refiner": self.is_refiner, - "hidden_batch_size": p.batch_size if disable_classifier_free_guidance else p.batch_size * 2, - "height": p.height, - "width": p.width, - }) - kwargs["sess_options"] = sess_options + "hidden_batch_size": p.batch_size if disable_classifier_free_guidance else p.batch_size * 2, + "height": p.height, + "width": p.width, + }) + kwargs["sess_options"] = sess_options - try: - out_dir = self.run_olive(submodels_for_olive, in_dir) - except Exception: - log.error(f"Olive: Failed to run olive passes: model='{self.original_filename}'.") - shutil.rmtree(shared.opts.onnx_temp_dir, ignore_errors=True) - shutil.rmtree(os.path.join(shared.opts.onnx_cached_models_path, self.original_filename), ignore_errors=True) + try: + out_dir = self.run_olive(submodels_for_olive, in_dir) + except Exception: + log.error(f"Olive: Failed to run olive passes: model='{self.original_filename}'.") + shutil.rmtree(shared.opts.onnx_temp_dir, ignore_errors=True) + shutil.rmtree(os.path.join(shared.opts.onnx_cached_models_path, self.original_filename), ignore_errors=True) pipeline = self.derive_properties(load_pipeline(self.constructor, out_dir, **kwargs)) diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index dcf6d15a7..10c914716 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -27,18 +27,25 @@ def process_diffusers(p: processing.StableDiffusionProcessing): 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 - 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 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 def hires_resize(latents): # input=latents output=pil if not torch.is_tensor(latents): @@ -399,7 +406,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing): p.task_args['sag_scale'] = p.sag_scale else: shared.log.warning(f'SAG incompatible scheduler: current={sd_model.scheduler.__class__.__name__} supported={supported}') - if shared.opts.cuda_compile_backend == "olive-ai": + if sd_model.__class__.__name__ == "OnnxRawPipeline": sd_model = preprocess_onnx_pipeline(p, is_refiner_enabled()) return sd_model @@ -540,7 +547,8 @@ 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) - preprocess_onnx_pipeline(p, is_refiner_enabled()) + if shared.sd_model.__class__.__name__ == "OnnxRawPipeline": + shared.sd_model = preprocess_onnx_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 615cc24f5..8728d3326 100644 --- a/modules/sd_models_compile.py +++ b/modules/sd_models_compile.py @@ -152,22 +152,16 @@ 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 5af58f7e8..59896d356 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -362,7 +362,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), { "cuda_compile_sep": OptionInfo("