diff --git a/modules/onnx.py b/modules/onnx.py index 799c0b0f1..d0b0ab0f7 100644 --- a/modules/onnx.py +++ b/modules/onnx.py @@ -33,7 +33,7 @@ def preprocess_pipeline(p, refiner_enabled: bool): 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": + if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "olive-ai": compile_height = p.height compile_width = p.width if (shared.compiled_model_state is None or @@ -62,12 +62,29 @@ def initialize(): if initialized: return - from modules.onnx_pipelines import do_diffusers_hijack + from modules import onnx_pipelines as pipelines # OnnxRuntimeModel Hijack. OnnxRuntimeModel.__module__ = 'diffusers' diffusers.OnnxRuntimeModel = OnnxRuntimeModel - do_diffusers_hijack() + diffusers.OnnxStableDiffusionPipeline = pipelines.OnnxStableDiffusionPipeline + diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion"] = diffusers.OnnxStableDiffusionPipeline + + diffusers.OnnxStableDiffusionImg2ImgPipeline = pipelines.OnnxStableDiffusionImg2ImgPipeline + diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion"] = diffusers.OnnxStableDiffusionImg2ImgPipeline + + diffusers.OnnxStableDiffusionInpaintPipeline = pipelines.OnnxStableDiffusionInpaintPipeline + diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["onnx-stable-diffusion"] = diffusers.OnnxStableDiffusionInpaintPipeline + + diffusers.OnnxStableDiffusionXLPipeline = pipelines.OnnxStableDiffusionXLPipeline + diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion-xl"] = diffusers.OnnxStableDiffusionXLPipeline + + diffusers.OnnxStableDiffusionXLImg2ImgPipeline = pipelines.OnnxStableDiffusionXLImg2ImgPipeline + diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion-xl"] = diffusers.OnnxStableDiffusionXLImg2ImgPipeline + + # Huggingface model compatibility + diffusers.ORTStableDiffusionXLPipeline = diffusers.OnnxStableDiffusionXLPipeline + diffusers.ORTStableDiffusionXLImg2ImgPipeline = diffusers.OnnxStableDiffusionXLImg2ImgPipeline initialized = True diff --git a/modules/onnx_pipelines.py b/modules/onnx_pipelines.py index 954297d53..ddd6f5799 100644 --- a/modules/onnx_pipelines.py +++ b/modules/onnx_pipelines.py @@ -22,7 +22,7 @@ from modules.sd_models import CheckpointInfo from modules.processing import StableDiffusionProcessing from modules.olive import config from modules.onnx import OnnxFakeModule, submodels_sd, submodels_sdxl, submodels_sdxl_refiner -from modules.onnx_utils import check_pipeline_sdxl, load_init_dict, load_submodel, load_submodels, load_pipeline, get_sess_options, patch_kwargs, construct_refiner_pipeline +from modules.onnx_utils import check_pipeline_sdxl, check_cache_onnx, load_init_dict, load_submodel, load_submodels, load_pipeline, get_sess_options, patch_kwargs, construct_refiner_pipeline from modules.onnx_ep import ExecutionProvider, EP_TO_NAME, get_provider @@ -80,8 +80,8 @@ class OnnxRawPipeline(OnnxPipelineBase): self.init_dict = load_init_dict(constructor, path) self.scheduler = load_submodel(self.path, None, "scheduler", self.init_dict["scheduler"]) else: + cls = diffusers.StableDiffusionXLPipeline if self._is_sdxl else diffusers.StableDiffusionPipeline try: - 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): @@ -92,7 +92,9 @@ class OnnxRawPipeline(OnnxPipelineBase): 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.') + log.error(f'Failed to load pipeline to optimize: is_sdxl={self._is_sdxl}') + log.warn('Model load failed. Please check Diffusers pipeline in Compute Settings.') + return if "vae" in self.init_dict: del self.init_dict["vae"] @@ -111,7 +113,9 @@ class OnnxRawPipeline(OnnxPipelineBase): ort.set_default_logger_severity(3) out_dir = os.path.join(shared.opts.onnx_cached_models_path, self.original_filename) - if os.path.isdir(out_dir): # already converted (cached) + if (self.from_huggingface_cache and check_cache_onnx(self.path)): + return self.path + if os.path.isdir(out_dir): # if model is ONNX format or had already converted. return out_dir try: @@ -357,7 +361,7 @@ class OnnxRawPipeline(OnnxPipelineBase): return self.derive_properties(load_pipeline(diffusers.StableDiffusionXLPipeline if self._is_sdxl else diffusers.StableDiffusionPipeline, self.path, **kwargs)) out_dir = converted_dir - if shared.opts.cuda_compile_backend == "olive-ai": + if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "olive-ai": log.warning("Olive implementation is experimental. It contains potentially an issue and is subject to change at any time.") if p.width != p.height: log.warning("Olive detected different width and height. The quality of the result is not guaranteed.") @@ -413,6 +417,9 @@ def prepare_latents( class OnnxStableDiffusionPipeline(diffusers.OnnxStableDiffusionPipeline, OnnxPipelineBase): + __module__ = 'diffusers' + __name__ = 'OnnxStableDiffusionPipeline' + def __init__( self, vae_encoder: diffusers.OnnxRuntimeModel, @@ -570,6 +577,9 @@ class OnnxStableDiffusionPipeline(diffusers.OnnxStableDiffusionPipeline, OnnxPip class OnnxStableDiffusionImg2ImgPipeline(diffusers.OnnxStableDiffusionImg2ImgPipeline, OnnxPipelineBase): + __module__ = 'diffusers' + __name__ = 'OnnxStableDiffusionImg2ImgPipeline' + image_processor: VaeImageProcessor def __init__( @@ -762,6 +772,9 @@ class OnnxStableDiffusionImg2ImgPipeline(diffusers.OnnxStableDiffusionImg2ImgPip class OnnxStableDiffusionInpaintPipeline(diffusers.OnnxStableDiffusionInpaintPipeline, OnnxPipelineBase): + __module__ = 'diffusers' + __name__ = 'OnnxStableDiffusionInpaintPipeline' + def __init__( self, vae_encoder: diffusers.OnnxRuntimeModel, @@ -969,6 +982,9 @@ class OnnxStableDiffusionInpaintPipeline(diffusers.OnnxStableDiffusionInpaintPip class OnnxStableDiffusionXLPipeline(OnnxPipelineBase, optimum.onnxruntime.ORTStableDiffusionXLPipeline): + __module__ = 'optimum.onnxruntime.modeling_diffusion' + __name__ = 'ORTStableDiffusionXLPipeline' + def __init__( self, vae_decoder, @@ -1159,6 +1175,9 @@ class OnnxStableDiffusionXLPipeline(OnnxPipelineBase, optimum.onnxruntime.ORTSta class OnnxStableDiffusionXLImg2ImgPipeline(OnnxPipelineBase, optimum.onnxruntime.ORTStableDiffusionXLImg2ImgPipeline): + __module__ = 'optimum.onnxruntime.modeling_diffusion' + __name__ = 'ORTStableDiffusionXLImg2ImgPipeline' + def __init__( self, vae_decoder, @@ -1344,28 +1363,3 @@ class OnnxStableDiffusionXLImg2ImgPipeline(OnnxPipelineBase, optimum.onnxruntime return (image,) return StableDiffusionXLPipelineOutput(images=image) - - -def do_diffusers_hijack(): - diffusers.OnnxStableDiffusionPipeline = OnnxStableDiffusionPipeline - diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion"] = diffusers.OnnxStableDiffusionPipeline - - OnnxStableDiffusionImg2ImgPipeline.__module__ = 'diffusers' - OnnxStableDiffusionImg2ImgPipeline.__name__ = 'OnnxStableDiffusionImg2ImgPipeline' - diffusers.OnnxStableDiffusionImg2ImgPipeline = OnnxStableDiffusionImg2ImgPipeline - diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion"] = diffusers.OnnxStableDiffusionImg2ImgPipeline - - OnnxStableDiffusionInpaintPipeline.__module__ = 'diffusers' - OnnxStableDiffusionInpaintPipeline.__name__ = 'OnnxStableDiffusionInpaintPipeline' - diffusers.OnnxStableDiffusionInpaintPipeline = OnnxStableDiffusionInpaintPipeline - diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["onnx-stable-diffusion"] = diffusers.OnnxStableDiffusionInpaintPipeline - - OnnxStableDiffusionXLPipeline.__module__ = 'optimum.onnxruntime.modeling_diffusion' - OnnxStableDiffusionXLPipeline.__name__ = 'ORTStableDiffusionXLPipeline' - diffusers.OnnxStableDiffusionXLPipeline = OnnxStableDiffusionXLPipeline - diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion-xl"] = diffusers.OnnxStableDiffusionXLPipeline - - OnnxStableDiffusionXLImg2ImgPipeline.__module__ = 'optimum.onnxruntime.modeling_diffusion' - OnnxStableDiffusionXLImg2ImgPipeline.__name__ = 'ORTStableDiffusionXLImg2ImgPipeline' - diffusers.OnnxStableDiffusionXLImg2ImgPipeline = OnnxStableDiffusionXLImg2ImgPipeline - diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion-xl"] = diffusers.OnnxStableDiffusionXLImg2ImgPipeline diff --git a/modules/onnx_utils.py b/modules/onnx_utils.py index ad2210e67..ee89e373a 100644 --- a/modules/onnx_utils.py +++ b/modules/onnx_utils.py @@ -44,6 +44,20 @@ def check_pipeline_sdxl(cls: Type[diffusers.DiffusionPipeline]) -> bool: return 'XL' in cls.__name__ +def check_cache_onnx(path: os.PathLike) -> bool: + if not os.path.isdir(path): + return False + init_dict_path = os.path.join(path, "model_index.json") + if not os.path.isfile(init_dict_path): + return False + init_dict = None + with open(init_dict_path, "r") as file: + init_dict = file.read() + if "OnnxRuntimeModel" not in init_dict: + return False + return True + + def load_submodel(path: os.PathLike, is_sdxl: bool, submodel_name: str, item: List[Union[str, None]], **kwargs_ort): lib, atr = item if lib is None or atr is None: diff --git a/modules/sd_models.py b/modules/sd_models.py index 8566232e3..7487c839e 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -147,7 +147,6 @@ 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.onnx_sideloaded_models_path, command_path=shared.opts.onnx_sideloaded_models_path, clear=False) for filename in sorted(model_list, key=str.lower): checkpoint_info = CheckpointInfo(filename) if checkpoint_info.name is not None: diff --git a/modules/shared.py b/modules/shared.py index 021d909d6..cda60e5b5 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -471,7 +471,6 @@ options_templates.update(options_section(('system-paths', "System Paths"), { "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), "onnx_cached_models_path": OptionInfo(os.path.join(paths.models_path, 'ONNX', 'cache'), "Folder with ONNX cached models", folder=True), - "onnx_sideloaded_models_path": OptionInfo(os.path.join(paths.models_path, 'ONNX', 'sideloaded'), "Folder with ONNX models from huggingface", 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 0557b9b1f..4f09066ac 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, is_onnx, custom_pipeline): + def hf_download_model(hub_id: str, token, variant, revision, mirror, custom_pipeline): from modules.modelloader import download_diffusers_model - download_diffusers_model(hub_id, cache_dir=opts.onnx_sideloaded_models_path if is_onnx else 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.diffusers_dir, token=token, variant=variant, revision=revision, mirror=mirror, custom_pipeline=custom_pipeline) + 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,9 +392,8 @@ 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_onnx = gr.Checkbox(False, label = 'ONNX model') + hf_variant = gr.Textbox(opts.cuda_dtype.lower(), label='Specify model variant', placeholder='') + hf_revision = gr.Textbox('', label='Specify model revision', placeholder='') 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') @@ -411,7 +410,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_onnx, 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_custom_pipeline], outputs=[models_outcome]) with gr.Tab(label="CivitAI"): data = []