From dbd0bf4bacad82d2eaa25dc48d000a9fa1ca8620 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Tue, 24 Oct 2023 20:12:40 +0300 Subject: [PATCH] Add custom pipeline support for Diffusers backend --- modules/modelloader.py | 4 +++- modules/sd_models.py | 3 +++ modules/shared.py | 1 + modules/shared_items.py | 1 + modules/ui_models.py | 7 ++++--- 5 files changed, 12 insertions(+), 4 deletions(-) diff --git a/modules/modelloader.py b/modules/modelloader.py index f5a4e1f35..dee78797d 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -183,7 +183,7 @@ def download_civit_model(model_url: str, model_name: str, model_path: str, model return f'CivitAI download: name={model_name} url={model_url} path={model_path}' -def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config: Dict[str, str] = None, token = None, variant = None, revision = None, mirror = None): +def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config: Dict[str, str] = None, token = None, variant = None, revision = None, mirror = None, custom_pipeline = None): if hub_id is None or len(hub_id) == 0: return None from diffusers import DiffusionPipeline @@ -204,6 +204,8 @@ def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config download_config["revision"] = revision if mirror is not None and len(mirror) > 0: download_config["mirror"] = mirror + if custom_pipeline is not None and len(custom_pipeline) > 0: + download_config["custom_pipeline"] = custom_pipeline shared.log.debug(f"Diffusers downloading: {hub_id} {download_config}") if token is not None and len(token) > 2: shared.log.debug(f"Diffusers authentication: {token}") diff --git a/modules/sd_models.py b/modules/sd_models.py index 9c84935ea..a8828f176 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -806,6 +806,9 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No else: diffusers_load_config['variant'] = shared.opts.diffusers_model_load_variant + if shared.opts.diffusers_pipeline == 'Custom Diffusers Pipeline': + diffusers_load_config['custom_pipeline'] = shared.opts.custom_diffusers_pipeline + if shared.opts.data.get('sd_model_checkpoint', '') == 'model.ckpt' or shared.opts.data.get('sd_model_checkpoint', '') == '': shared.opts.data['sd_model_checkpoint'] = "runwayml/stable-diffusion-v1-5" diff --git a/modules/shared.py b/modules/shared.py index 7b7c10e61..66dbb820f 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -337,6 +337,7 @@ options_templates.update(options_section(('diffusers', "Diffusers Settings"), { "diffusers_attention_slicing": OptionInfo(False, "Enable attention slicing"), "diffusers_model_load_variant": OptionInfo("default", "Diffusers model loading variant", gr.Radio, {"choices": ['default', 'fp32', 'fp16']}), "diffusers_vae_load_variant": OptionInfo("default", "Diffusers VAE loading variant", gr.Radio, {"choices": ['default', 'fp32', 'fp16']}), + "custom_diffusers_pipeline": OptionInfo('hf-internal-testing/diffusers-dummy-pipeline', 'Custom Diffusers pipeline to use'), "diffusers_lora_loader": OptionInfo("diffusers" if cmd_opts.use_openvino else "sequential apply", "Diffusers LoRA loading variant", gr.Radio, {"choices": ['diffusers', 'sequential apply', 'merge and apply']}), "diffusers_force_zeros": OptionInfo(True, "Force zeros for prompts when empty"), "diffusers_aesthetics_score": OptionInfo(False, "Require aesthetics score"), diff --git a/modules/shared_items.py b/modules/shared_items.py index 54287099f..85fba87a3 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -37,6 +37,7 @@ def get_pipelines(): 'Stable Diffusion XL Img2Img': getattr(diffusers, 'StableDiffusionXLImg2ImgPipeline', None), 'Stable Diffusion XL Inpaint': getattr(diffusers, 'StableDiffusionXLInpaintPipeline', None), 'Stable Diffusion XL Instruct': getattr(diffusers, 'StableDiffusionXLInstructPix2PixPipeline', None), + 'Custom Diffusers Pipeline': getattr(diffusers, 'DiffusionPipeline', None), # 'Test': getattr(diffusers, 'TestPipeline', None), # 'Kandinsky V1', 'Kandinsky V2', 'DeepFloyd IF', 'Shap-E', 'Kandinsky V1 Img2Img', 'Kandinsky V2 Img2Img', 'DeepFloyd IF Img2Img', 'Shap-E Img2Img', } diff --git a/modules/ui_models.py b/modules/ui_models.py index 717fe15d8..e84300963 100644 --- a/modules/ui_models.py +++ b/modules/ui_models.py @@ -203,9 +203,9 @@ 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): + 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.diffusers_dir, token=token, variant=variant, revision=revision, mirror=mirror) + 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}"') @@ -227,6 +227,7 @@ def create_ui(): 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') + hf_custom_pipeline = gr.Textbox('', label = 'Custom pipeline', placeholder='optional pipeline for downloads') with gr.Column(scale=1): gr.HTML('
') hf_download_model_btn = gr.Button(value="Download model", variant='primary') @@ -239,7 +240,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], 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 = []