diff --git a/modules/olive.py b/modules/olive.py index d5134657a..71558e331 100644 --- a/modules/olive.py +++ b/modules/olive.py @@ -8,12 +8,9 @@ from typing import Union, Optional, Callable, List from transformers.models.clip.modeling_clip import CLIPTextModel, CLIPTextModelWithProjection from installer import log, args from modules.shared import opts, cmd_opts -from modules.paths import models_path, sd_configs_path +from modules.paths import sd_configs_path from modules.sd_models import CheckpointInfo -temp_dir = os.path.join(models_path, "OliveTemp") -cache_dir = os.path.join(models_path, "OliveCache") - submodels = ("text_encoder", "unet", "vae_encoder", "vae_decoder",) execution_provider = "CUDAExecutionProvider" @@ -25,7 +22,18 @@ provider = (execution_provider, { "device_id": int(cmd_opts.device_id or 0), }) +class OnnxRuntimeModel(diffusers.OnnxRuntimeModel): + config = {} + + def named_modules(self): + return () + + +diffusers.OnnxRuntimeModel = OnnxRuntimeModel + + class OnnxStableDiffusionPipeline(diffusers.OnnxStableDiffusionPipeline): + model_type: str sd_model_hash: str sd_checkpoint_info: CheckpointInfo sd_model_checkpoint: str @@ -192,9 +200,9 @@ class OlivePipeline(diffusers.DiffusionPipeline): self.original_filename = os.path.basename(path) self.unoptimized = pipeline del pipeline - if not os.path.exists(temp_dir): - os.mkdir(temp_dir) - self.unoptimized.save_pretrained(temp_dir) + if not os.path.exists(opts.olive_temp_dir): + os.mkdir(opts.olive_temp_dir) + self.unoptimized.save_pretrained(opts.olive_temp_dir) @staticmethod def from_pretrained(pretrained_model_name_or_path, **kwargs): @@ -218,7 +226,7 @@ class OlivePipeline(diffusers.DiffusionPipeline): if width != height: log.warning("Olive received different width and height. The quality of the result is not guaranteed.") - out_dir = os.path.join(cache_dir, f"{self.original_filename}-{width}w-{height}h") + out_dir = os.path.join(opts.olive_cached_models_path, f"{self.original_filename}-{width}w-{height}h") if os.path.isdir(out_dir): del self.unoptimized return OnnxStableDiffusionPipeline.from_pretrained(out_dir, provider=provider).apply(self) @@ -226,7 +234,7 @@ class OlivePipeline(diffusers.DiffusionPipeline): try: if opts.olive_cache_optimized: shutil.copytree( - temp_dir, out_dir, ignore=shutil.ignore_patterns("weights.pb", "*.onnx", "*.safetensors", "*.ckpt") + opts.olive_temp_dir, out_dir, ignore=shutil.ignore_patterns("weights.pb", "*.onnx", "*.safetensors", "*.ckpt") ) optimize_config["width"] = width @@ -260,7 +268,7 @@ class OlivePipeline(diffusers.DiffusionPipeline): ).model_path log.info(f"Optimized {submodel}") - shutil.rmtree(temp_dir) + shutil.rmtree(opts.olive_temp_dir) kwargs = { "tokenizer": self.unoptimized.tokenizer, @@ -295,9 +303,9 @@ class OlivePipeline(diffusers.DiffusionPipeline): if os.path.isfile(weights_src_path): weights_dst_path = os.path.join(dst_parent, (os.path.basename(dst_path) + ".data")) shutil.copyfile(weights_src_path, weights_dst_path) - except Exception as e: + except Exception: log.error(f"Failed to optimize model '{self.original_filename}'.") - shutil.rmtree(temp_dir, ignore_errors=True) + shutil.rmtree(opts.olive_temp_dir, ignore_errors=True) shutil.rmtree(out_dir, ignore_errors=True) pipeline = None shutil.rmtree("cache", ignore_errors=True) @@ -312,7 +320,7 @@ class OlivePipeline(diffusers.DiffusionPipeline): optimize_config = { "is_sdxl": False, - "source": os.path.abspath(temp_dir), + "source": os.path.abspath(opts.olive_temp_dir), "width": 512, "height": 512, diff --git a/modules/sd_models.py b/modules/sd_models.py index fe0add649..179d28c17 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -147,6 +147,7 @@ def list_models(): model_list = list(modelloader.load_models(model_path=model_path, model_url=None, command_path=shared.opts.ckpt_dir, ext_filter=ext_filter, download_name=None, ext_blacklist=[".vae.ckpt", ".vae.safetensors"])) if shared.backend == shared.Backend.DIFFUSERS: model_list += modelloader.load_diffusers_models(model_path=os.path.join(models_path, 'Diffusers'), command_path=shared.opts.diffusers_dir, clear=True) + model_list += modelloader.load_diffusers_models(model_path=shared.opts.olive_sideloaded_models_path, command_path=shared.opts.olive_sideloaded_models_path) for filename in sorted(model_list, key=str.lower): checkpoint_info = CheckpointInfo(filename) if checkpoint_info.name is not None: @@ -790,6 +791,38 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No shared.log.debug(f'Diffusers loading: path="{checkpoint_info.path}"') pipeline, model_type = detect_pipeline(checkpoint_info.path, op) if os.path.isdir(checkpoint_info.path): + if shared.opts.olive_sideloaded_models_path in checkpoint_info.path: + try: + from modules.olive import OnnxStableDiffusionPipeline, provider + sd_model = OnnxStableDiffusionPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.olive_sideloaded_models_path, provider=provider) + sd_model.model_type = sd_model.__class__.__name__ + except Exception as e: + shared.log.error(f'Failed loading {op}: {checkpoint_info.path} olive={e}') + return + else: + err1 = None + err2 = None + err3 = None + try: # try autopipeline first, best choice but not all pipelines are available + sd_model = diffusers.AutoPipelineForText2Image.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) + sd_model.model_type = sd_model.__class__.__name__ + except Exception as e: + err1 = e + try: # try diffusion pipeline next second-best choice, works for most non-linked pipelines + if err1 is not None: + sd_model = diffusers.DiffusionPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) + sd_model.model_type = sd_model.__class__.__name__ + except Exception as e: + err2 = e + try: # try basic pipeline next just in case + if err2 is not None: + sd_model = diffusers.StableDiffusionPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) + sd_model.model_type = sd_model.__class__.__name__ + except Exception as e: + err3 = e # ignore last error + if err3 is not None: + shared.log.error(f'Failed loading {op}: {checkpoint_info.path} auto={err1} diffusion={err2}') + return if model_type in ['InstaFlow']: # forced pipeline sd_model = pipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) else: diff --git a/modules/shared.py b/modules/shared.py index 0ba4a6f4b..bd2805327 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -455,6 +455,7 @@ options_templates.update(options_section(('system-paths', "System Paths"), { "embeddings_dir": OptionInfo(os.path.join(paths.models_path, 'embeddings'), "Folder with textual inversion embeddings", folder=True), "hypernetwork_dir": OptionInfo(os.path.join(paths.models_path, 'hypernetworks'), "Folder with Hypernetwork models", folder=True), "control_dir": OptionInfo(os.path.join(paths.models_path, 'control'), "Folder with Control models", folder=True), + "olive_temp_dir": OptionInfo(os.path.join(paths.models_path, 'Olive', 'temp'), "Directory for olive optimization process", folder=True), "codeformer_models_path": OptionInfo(os.path.join(paths.models_path, 'Codeformer'), "Folder with codeformer models", folder=True), "gfpgan_models_path": OptionInfo(os.path.join(paths.models_path, 'GFPGAN'), "Folder with GFPGAN models", folder=True), "esrgan_models_path": OptionInfo(os.path.join(paths.models_path, 'ESRGAN'), "Folder with ESRGAN models", folder=True), @@ -464,6 +465,8 @@ options_templates.update(options_section(('system-paths', "System Paths"), { "swinir_models_path": OptionInfo(os.path.join(paths.models_path, 'SwinIR'), "Folder with SwinIR models", folder=True), "ldsr_models_path": OptionInfo(os.path.join(paths.models_path, 'LDSR'), "Folder with LDSR models", folder=True), "clip_models_path": OptionInfo(os.path.join(paths.models_path, 'CLIP'), "Folder with CLIP models", folder=True), + "olive_cached_models_path": OptionInfo(os.path.join(paths.models_path, 'Olive', 'cache'), "Folder with olive optimized cached models", folder=True), + "olive_sideloaded_models_path": OptionInfo(os.path.join(paths.models_path, 'Olive', 'sideloaded'), "Folder with olive optimized sideloaded models", folder=True), "other_paths_sep_options": OptionInfo("

Other paths

", "", gr.HTML), "openvino_cache_path": OptionInfo('cache', "Directory for OpenVINO cache", folder=True), diff --git a/modules/ui_models.py b/modules/ui_models.py index 4f09066ac..cc0f743cf 100644 --- a/modules/ui_models.py +++ b/modules/ui_models.py @@ -373,10 +373,10 @@ def create_ui(): def hf_select(evt: gr.SelectData, data): return data[evt.index[0]][0] - def hf_download_model(hub_id: str, token, variant, revision, mirror, custom_pipeline): + def hf_download_model(hub_id: str, token, variant, revision, mirror, olive_optimized): from modules.modelloader import download_diffusers_model - download_diffusers_model(hub_id, cache_dir=opts.diffusers_dir, token=token, variant=variant, revision=revision, mirror=mirror, custom_pipeline=custom_pipeline) - from modules.sd_models import list_models # pylint: disable=W0621 + download_diffusers_model(hub_id, cache_dir=opts.olive_sideloaded_models_path if olive_optimized else opts.diffusers_dir, token=token, variant=variant, revision=revision, mirror=mirror) + from modules.sd_models import list_models # pylint: disable=W0621 list_models() log.info(f'Diffuser model downloaded: model="{hub_id}"') return f'Diffuser model downloaded: model="{hub_id}"' @@ -392,8 +392,9 @@ def create_ui(): hf_selected = gr.Textbox('', label='Select model', placeholder='select model from search results or enter model name manually') with gr.Column(scale=1): with gr.Row(): - hf_variant = gr.Textbox(opts.cuda_dtype.lower(), label='Specify model variant', placeholder='') - hf_revision = gr.Textbox('', label='Specify model revision', placeholder='') + hf_variant = gr.Textbox(opts.cuda_dtype.lower(), label = 'Specify model variant', placeholder='') + hf_revision = gr.Textbox('', label = 'Specify model revision', placeholder='') + hf_olive = gr.Checkbox(False, label = 'Olive optimized') with gr.Row(): hf_token = gr.Textbox('', label='Huggingface token', placeholder='optional access token for private or gated models') hf_mirror = gr.Textbox('', label='Huggingface mirror', placeholder='optional mirror site for downloads') @@ -410,7 +411,7 @@ def create_ui(): hf_search_text.submit(fn=hf_search, inputs=[hf_search_text], outputs=[hf_results]) hf_search_btn.click(fn=hf_search, inputs=[hf_search_text], outputs=[hf_results]) hf_results.select(fn=hf_select, inputs=[hf_results], outputs=[hf_selected]) - hf_download_model_btn.click(fn=hf_download_model, inputs=[hf_selected, hf_token, hf_variant, hf_revision, hf_mirror, hf_custom_pipeline], outputs=[models_outcome]) + hf_download_model_btn.click(fn=hf_download_model, inputs=[hf_selected, hf_token, hf_variant, hf_revision, hf_mirror, hf_olive], outputs=[models_outcome]) with gr.Tab(label="CivitAI"): data = []