Add custom pipeline support for Diffusers backend

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}")

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"
 

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"),

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',
     }

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('<br>')
                     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 = []