From e7e317191a7d1fdb5718301e24bb813b069bce00 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Thu, 4 Jun 2026 10:18:09 +0200 Subject: [PATCH] automated pipeline registrations and tests Signed-off-by: Vladimir Mandic --- data/reference.json | 30 ------- modules/sd_detect.py | 36 ++++----- modules/sd_models.py | 45 +++++------ modules/shared_items.py | 139 +++++++++++++++++---------------- pipelines/generic.py | 9 ++- pipelines/generic_test.py | 38 ++++++++- pipelines/generic_util.py | 13 +++ pipelines/model_anima.py | 1 + pipelines/model_bria.py | 2 + pipelines/model_flex.py | 1 + pipelines/model_flite.py | 3 +- pipelines/model_hidream.py | 1 + pipelines/model_hyimage.py | 1 + pipelines/model_instaflow.py | 22 ++++++ pipelines/model_lens.py | 7 +- pipelines/model_lumina.py | 1 + pipelines/model_omnigen.py | 9 ++- pipelines/model_sdxs.py | 9 ++- pipelines/model_segmoe.py | 22 ++++++ pipelines/model_step1x_edit.py | 1 + pipelines/model_ultraflux.py | 4 +- pipelines/model_vibe.py | 2 +- pipelines/model_xomni.py | 2 + pipelines/model_zetachroma.py | 3 +- 24 files changed, 244 insertions(+), 157 deletions(-) create mode 100644 pipelines/model_instaflow.py create mode 100644 pipelines/model_segmoe.py diff --git a/data/reference.json b/data/reference.json index 89c95feff..b28af9c3d 100644 --- a/data/reference.json +++ b/data/reference.json @@ -821,28 +821,6 @@ "size": 3.64, "date": "2024 October" }, - "aMUSEd 256": { - "path": "huggingface/amused/amused-256", - "skip": true, - "desc": "Amused is a lightweight text to image model based off of the muse architecture. Amused is particularly useful in applications that require a lightweight and fast model such as generating many images quickly at once.", - "preview": "amused--amused-256.jpg", - "extras": "width: 256, height: 256, sampler: Default" - }, - "aMUSEd 512": { - "path": "amused/amused-512", - "desc": "Amused is a lightweight text to image model based off of the muse architecture. Amused is particularly useful in applications that require a lightweight and fast model such as generating many images quickly at once.", - "preview": "amused--amused-512.jpg", - "extras": "width: 512, height: 512, sampler: Default", - "size": 4.29 - }, - "Warp Wuerstchen": { - "path": "warp-ai/wuerstchen", - "desc": "W\u00fcrstchen is a diffusion model whose text-conditional model works in a highly compressed latent space of images. Why is this important? Compressing data can reduce computational costs for both training and inference by magnitudes. Training on 1024x1024 images, is way more expensive than training at 32x32. Usually, other works make use of a relatively small compression, in the range of 4x - 8x spatial compression. W\u00fcrstchen takes this to an extreme. Through its novel design, we achieve a 42x spatial compression. W\u00fcrstchen employs a two-stage compression, what we call Stage A and Stage B. Stage A is a VQGAN, and Stage B is a Diffusion Autoencoder (more details can be found in the paper). A third model, Stage C, is learned in that highly compressed latent space. This training requires fractions of the compute used for current top-performing models, allowing also cheaper and faster inference.", - "preview": "warp-ai--wuerstchen.jpg", - "extras": "sampler: Default, cfg_scale: 4.0, cfg_image: 0.0", - "size": 12.16, - "date": "2023 August" - }, "KOALA 700M": { "path": "huggingface/etri-vilab/koala-700m-llava-cap", "variant": "fp16", @@ -869,14 +847,6 @@ "preview": "KBlueLeaf--HDM-xut-340M-anime.jpg", "extras": "" }, - "Tsinghua UniDiffuser": { - "path": "thu-ml/unidiffuser-v1", - "desc": "UniDiffuser is a unified diffusion framework to fit all distributions relevant to a set of multi-modal data in one transformer. UniDiffuser is able to perform image, text, text-to-image, image-to-text, and image-text pair generation by setting proper timesteps without additional overhead.\nSpecifically, UniDiffuser employs a variation of transformer, called U-ViT, which parameterizes the joint noise prediction network. Other components perform as encoders and decoders of different modalities, including a pretrained image autoencoder from Stable Diffusion, a pretrained image ViT-B/32 CLIP encoder, a pretrained text ViT-L CLIP encoder, and a GPT-2 text decoder finetuned by ourselves.", - "preview": "thu-ml--unidiffuser-v1.jpg", - "extras": "width: 512, height: 512, sampler: Default", - "size": 5.37, - "date": "2023 May" - }, "SalesForce BLIP-Diffusion": { "path": "salesforce/blipdiffusion", "desc": "BLIP-Diffusion, a new subject-driven image generation model that supports multimodal control which consumes inputs of subject images and text prompts. Unlike other subject-driven generation models, BLIP-Diffusion introduces a new multimodal encoder which is pre-trained to provide subject representation.", diff --git a/modules/sd_detect.py b/modules/sd_detect.py index adfecd74b..0717eed70 100644 --- a/modules/sd_detect.py +++ b/modules/sd_detect.py @@ -53,29 +53,29 @@ def guess_by_name(fn, current_guess): elif 'hunyuandit' in fn.lower(): new_guess = 'HunyuanDiT' elif 'pixart-xl' in fn.lower(): - new_guess = 'PixArt Alpha' + new_guess = 'PixArtAlpha' elif 'stable-diffusion-3' in fn.lower(): new_guess = 'Stable Diffusion 3' elif 'stable-cascade' in fn.lower() or 'stablecascade' in fn.lower() or 'wuerstchen3' in fn.lower() or ('sotediffusion' in fn.lower() and "v2" in fn.lower()): new_guess = 'Stable Cascade' elif 'pixart-sigma' in fn.lower(): - new_guess = 'PixArt Sigma' + new_guess = 'PixArtSigma' elif 'sana' in fn.lower(): new_guess = 'Sana' elif 'lumina-next' in fn.lower(): - new_guess = 'Lumina-Next' + new_guess = 'LuminaNext' elif 'lumina-dimoo' in fn.lower(): - new_guess = 'Lumina-DiMOO' + new_guess = 'LuminaDiMOO' elif 'lumina-image-2' in fn.lower(): - new_guess = 'Lumina 2' + new_guess = 'Lumina2' elif 'kolors' in fn.lower(): new_guess = 'Kolors' elif 'auraflow' in fn.lower() or 'pony-v7' in fn.lower(): new_guess = 'AuraFlow' elif 'cogview3' in fn.lower(): - new_guess = 'CogView 3' + new_guess = 'CogView3' elif 'cogview4' in fn.lower(): - new_guess = 'CogView 4' + new_guess = 'CogView4' elif 'meissonic' in fn.lower(): new_guess = 'Meissonic' elif 'omnigen2' in fn.lower(): @@ -93,7 +93,7 @@ def guess_by_name(fn, current_guess): elif 'chroma' in fn.lower() and 'xl' not in fn.lower(): new_guess = 'Chroma' elif 'flux.2' in fn.lower() and 'klein' in fn.lower(): - new_guess = 'FLUX2 Klein' + new_guess = 'FLUX2Klein' elif 'flux.2' in fn.lower(): new_guess = 'FLUX2' elif 'lens' in fn.lower(): @@ -124,19 +124,19 @@ def guess_by_name(fn, current_guess): elif 'nextstep' in fn.lower(): new_guess = 'NextStep' elif 'kandinsky-2-1' in fn.lower(): - new_guess = 'Kandinsky 2.1' + new_guess = 'Kandinsky21' elif 'kandinsky-2-2' in fn.lower(): - new_guess = 'Kandinsky 2.2' + new_guess = 'Kandinsky22' elif 'kandinsky-3' in fn.lower(): - new_guess = 'Kandinsky 3.0' + new_guess = 'Kandinsky30' elif 'kandinsky-5.0' in fn.lower(): - new_guess = 'Kandinsky 5.0' + new_guess = 'Kandinsky50' elif 'hunyuanimage3' in fn.lower() or 'hunyuanimage-3' in fn.lower(): new_guess = 'HunyuanImage3' elif 'hunyuanimage' in fn.lower(): new_guess = 'HunyuanImage' elif 'x-omni' in fn.lower(): - new_guess = 'X-Omni' + new_guess = 'XOmni' elif 'sdxl-turbo' in fn.lower() or 'stable-diffusion-xl' in fn.lower(): new_guess = 'Stable Diffusion XL' elif 'stable-video-diffusion' in fn.lower(): @@ -146,11 +146,11 @@ def guess_by_name(fn, current_guess): elif 'gemini-' in fn.lower() and 'image' in fn.lower(): new_guess = 'NanoBanana' elif 'ernie-image' in fn.lower(): - new_guess = 'ERNIE-Image' + new_guess = 'ERNIEImage' elif 'nucleus-image' in fn.lower() or 'nucleusmoe-image' in fn.lower(): - new_guess = 'Nucleus-Image' + new_guess = 'NucleusImage' elif 'z-image' in fn.lower() or 'z_image' in fn.lower(): - new_guess = 'Z-Image' + new_guess = 'ZImage' elif 'longcat-image' in fn.lower(): new_guess = 'LongCat' elif 'ovis-image' in fn.lower(): @@ -160,11 +160,11 @@ def guess_by_name(fn, current_guess): elif 'sdxs-1b' in fn.lower(): new_guess = 'SDXS' elif 'step1x-edit' in fn.lower(): - new_guess = 'Step1X-Edit' + new_guess = 'Step1XEdit' elif 'vibe-image-edit' in fn.lower(): new_guess = 'VIBE' elif 'joyai-image-edit' in fn.lower() or 'joy-image-edit' in fn.lower(): - new_guess = 'Joy' + new_guess = 'JoyEdit' if debug_load: log.trace(f'Autodetect: method=name file="{fn}" previous="{current_guess}" current="{new_guess}"') return new_guess or current_guess diff --git a/modules/sd_models.py b/modules/sd_models.py index 1a2bc76ca..2bf3d0868 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -11,7 +11,7 @@ import diffusers.loaders.single_file_utils import torch import huggingface_hub as hf from modules.logger import log -from modules import timer, paths, shared, shared_items, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, sd_models_compile, sd_detect, model_quant, sd_hijack_te, sd_hijack_accelerate, sd_hijack_safetensors, sd_hijack_transformers, sd_hijack_hfhub, attention +from modules import timer, paths, shared, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, sd_models_compile, sd_detect, model_quant, sd_hijack_te, sd_hijack_accelerate, sd_hijack_safetensors, sd_hijack_transformers, sd_hijack_hfhub, attention from modules.memstats import memory_stats from modules.shared_helpers import walk_files from modules.modeldata import model_data @@ -344,16 +344,13 @@ def load_diffuser_force(detected_model_type: str, checkpoint_info: CheckpointInf sd_model = load_cascade_combined(checkpoint_info, diffusers_load_config) allow_post_quant = True elif model_type in ['InstaFlow']: - pipeline = diffusers.utils.get_class_from_dynamic_module('instaflow_one_step', module_file='pipeline.py') - shared_items.pipelines['InstaFlow'] = pipeline - sd_model = pipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) + from pipelines.model_instaflow import load_instaflow + sd_model = load_instaflow(checkpoint_info, diffusers_load_config) allow_post_quant = True elif model_type in ['SegMoE']: - from pipelines.segmoe.segmoe_model import SegMoEPipeline - sd_model = SegMoEPipeline(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) - sd_model = sd_model.pipe # segmoe pipe does its stuff in __init__ and __call__ is the original pipeline + from pipelines.model_segmoe import load_segmoe + sd_model = load_segmoe(checkpoint_info, diffusers_load_config) allow_post_quant = True - shared_items.pipelines['SegMoE'] = SegMoEPipeline elif model_type in ['PixArt Sigma']: from pipelines.model_pixart import load_pixart sd_model = load_pixart(checkpoint_info, diffusers_load_config) @@ -362,11 +359,11 @@ def load_diffuser_force(detected_model_type: str, checkpoint_info: CheckpointInf from pipelines.model_sana import load_sana sd_model = load_sana(checkpoint_info, diffusers_load_config) allow_post_quant = False - elif model_type in ['Lumina-Next']: + elif model_type in ['LuminaNext']: from pipelines.model_lumina import load_lumina sd_model = load_lumina(checkpoint_info, diffusers_load_config) allow_post_quant = True - elif model_type in ['Lumina-DiMOO']: + elif model_type in ['LuminaDiMOO']: from pipelines.model_lumina import load_lumina_dimoo sd_model = load_lumina_dimoo(checkpoint_info, diffusers_load_config) allow_post_quant = False @@ -390,7 +387,7 @@ def load_diffuser_force(detected_model_type: str, checkpoint_info: CheckpointInf from pipelines.model_flux2 import load_flux2 sd_model = load_flux2(checkpoint_info, diffusers_load_config) allow_post_quant = False - elif model_type in ['FLUX2 Klein']: + elif model_type in ['FLUX2Klein']: from pipelines.model_flux2_klein import load_flux2_klein sd_model = load_flux2_klein(checkpoint_info, diffusers_load_config) allow_post_quant = False @@ -406,7 +403,7 @@ def load_diffuser_force(detected_model_type: str, checkpoint_info: CheckpointInf from pipelines.model_chroma import load_chroma sd_model = load_chroma(checkpoint_info, diffusers_load_config) allow_post_quant = False - elif model_type in ['Lumina 2']: + elif model_type in ['Lumina2']: from pipelines.model_lumina import load_lumina2 sd_model = load_lumina2(checkpoint_info, diffusers_load_config) allow_post_quant = False @@ -414,11 +411,11 @@ def load_diffuser_force(detected_model_type: str, checkpoint_info: CheckpointInf from pipelines.model_sd3 import load_sd3 sd_model = load_sd3(checkpoint_info, diffusers_load_config) allow_post_quant = False - elif model_type in ['CogView 3']: + elif model_type in ['CogView3']: from pipelines.model_cogview import load_cogview3 sd_model = load_cogview3(checkpoint_info, diffusers_load_config) allow_post_quant = False - elif model_type in ['CogView 4']: + elif model_type in ['CogView4']: from pipelines.model_cogview import load_cogview4 sd_model = load_cogview4(checkpoint_info, diffusers_load_config) allow_post_quant = False @@ -466,7 +463,7 @@ def load_diffuser_force(detected_model_type: str, checkpoint_info: CheckpointInf from pipelines.model_bria import load_bria sd_model = load_bria(checkpoint_info, diffusers_load_config) allow_post_quant = False - elif model_type in ['Step1X-Edit']: + elif model_type in ['Step1XEdit']: from pipelines.model_step1x_edit import load_step1x_edit sd_model = load_step1x_edit(checkpoint_info, diffusers_load_config) allow_post_quant = False @@ -474,7 +471,7 @@ def load_diffuser_force(detected_model_type: str, checkpoint_info: CheckpointInf from pipelines.model_vibe import load_vibe sd_model = load_vibe(checkpoint_info, diffusers_load_config) allow_post_quant = False - elif model_type in ['Joy']: + elif model_type in ['JoyEdit']: from pipelines.model_joy import load_joy sd_model = load_joy(checkpoint_info, diffusers_load_config) allow_post_quant = False @@ -490,19 +487,19 @@ def load_diffuser_force(detected_model_type: str, checkpoint_info: CheckpointInf from pipelines.model_lens import load_lens sd_model = load_lens(checkpoint_info, diffusers_load_config) allow_post_quant = False - elif model_type in ['Kandinsky 2.1']: + elif model_type in ['Kandinsky21']: from pipelines.model_kandinsky import load_kandinsky21 sd_model = load_kandinsky21(checkpoint_info, diffusers_load_config) allow_post_quant = True - elif model_type in ['Kandinsky 2.2']: + elif model_type in ['Kandinsky22']: from pipelines.model_kandinsky import load_kandinsky22 sd_model = load_kandinsky22(checkpoint_info, diffusers_load_config) allow_post_quant = True - elif model_type in ['Kandinsky 3.0']: + elif model_type in ['Kandinsky30']: from pipelines.model_kandinsky import load_kandinsky3 sd_model = load_kandinsky3(checkpoint_info, diffusers_load_config) allow_post_quant = False - elif model_type in ['Kandinsky 5.0']: + elif model_type in ['Kandinsky50']: from pipelines.model_kandinsky import load_kandinsky5 sd_model = load_kandinsky5(checkpoint_info, diffusers_load_config) allow_post_quant = False @@ -518,7 +515,7 @@ def load_diffuser_force(detected_model_type: str, checkpoint_info: CheckpointInf from pipelines.model_hyimage import load_hyimage3 sd_model = load_hyimage3(checkpoint_info, diffusers_load_config) # pylint: disable=assignment-from-none allow_post_quant = False - elif model_type in ['X-Omni']: + elif model_type in ['XOmni']: from pipelines.model_xomni import load_xomni sd_model = load_xomni(checkpoint_info, diffusers_load_config) # pylint: disable=assignment-from-none allow_post_quant = False @@ -530,15 +527,15 @@ def load_diffuser_force(detected_model_type: str, checkpoint_info: CheckpointInf from pipelines.model_prx import load_prx sd_model = load_prx(checkpoint_info, diffusers_load_config) allow_post_quant = False - elif model_type in ['ERNIE-Image']: + elif model_type in ['ERNIEImage']: from pipelines.model_ernie import load_ernie_image sd_model = load_ernie_image(checkpoint_info, diffusers_load_config) allow_post_quant = False - elif model_type in ['Nucleus-Image']: + elif model_type in ['NucleusImage']: from pipelines.model_nucleus import load_nucleus sd_model = load_nucleus(checkpoint_info, diffusers_load_config) allow_post_quant = False - elif model_type in ['Z-Image']: + elif model_type in ['ZImage']: from pipelines.model_z_image import load_z_image sd_model = load_z_image(checkpoint_info, diffusers_load_config) allow_post_quant = False diff --git a/modules/shared_items.py b/modules/shared_items.py index d642d3e16..8ca401be9 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -1,13 +1,16 @@ import diffusers -pipelines = { - # note: not all pipelines can be used manually as they require prior pipeline next to decoder pipeline - 'Autodetect': None, - 'Custom Diffusers Pipeline': getattr(diffusers, 'DiffusionPipeline', None), +class OnlinePipeline(diffusers.DiffusionPipeline): + pass - # standard pipelines - 'Diffusion': getattr(diffusers, 'DiffusionPipeline', None), + +pipelines = { + 'Autodetect': None, + 'AutoPipeline': diffusers.AutoPipelineForText2Image, + 'Diffusion': diffusers.DiffusionPipeline, + + # standard diffusers pipelines 'Stable Diffusion': getattr(diffusers, 'StableDiffusionPipeline', None), 'Stable Diffusion Inpaint': getattr(diffusers, 'StableDiffusionInpaintPipeline', None), 'Stable Diffusion Instruct': getattr(diffusers, 'StableDiffusionInstructPix2PixPipeline', None), @@ -22,58 +25,60 @@ pipelines = { 'Stable Cascade': getattr(diffusers, 'StableCascadeCombinedPipeline', None), 'Stable Diffusion 3': getattr(diffusers, 'StableDiffusion3Pipeline', None), 'Latent Consistency Model': getattr(diffusers, 'LatentConsistencyModelPipeline', None), - 'PixArt Alpha': getattr(diffusers, 'PixArtAlphaPipeline', None), - 'PixArt Sigma': getattr(diffusers, 'PixArtSigmaPipeline', None), - 'HunyuanDiT': getattr(diffusers, 'HunyuanDiTPipeline', None), - 'DeepFloyd IF': getattr(diffusers, 'IFPipeline', None), - 'FLUX': getattr(diffusers, 'FluxPipeline', None), - 'UltraFlux': getattr(diffusers, 'DiffusionPipeline', None), - 'FLEX': getattr(diffusers, 'AutoPipelineForText2Image', None), - 'Chroma': getattr(diffusers, 'ChromaPipeline', None), - 'Sana': getattr(diffusers, 'SanaPipeline', None), - 'Lumina-Next': getattr(diffusers, 'LuminaText2ImgPipeline', None), - 'Lumina 2': getattr(diffusers, 'Lumina2Pipeline', None), 'AuraFlow': getattr(diffusers, 'AuraFlowPipeline', None), - 'Kandinsky 2.1': getattr(diffusers, 'KandinskyCombinedPipeline', None), - 'Kandinsky 2.2': getattr(diffusers, 'KandinskyV22CombinedPipeline', None), - 'Kandinsky 3.0': getattr(diffusers, 'Kandinsky3Pipeline', None), - 'Kandinsky 5.0': getattr(diffusers, 'Kandinsky5T2IPipeline', None), - 'Wuerstchen': getattr(diffusers, 'WuerstchenCombinedPipeline', None), - 'Kolors': getattr(diffusers, 'KolorsPipeline', None), - 'CogView 3': getattr(diffusers, 'CogView3PlusPipeline', None), - 'CogView 4': getattr(diffusers, 'CogView4Pipeline', None), - 'UniDiffuser': getattr(diffusers, 'UniDiffuserPipeline', None), - 'Amused': getattr(diffusers, 'AmusedPipeline', None), - 'HiDream': getattr(diffusers, 'HiDreamImagePipeline', None), - 'OmniGen': getattr(diffusers, 'OmniGenPipeline', None), + 'Chroma': getattr(diffusers, 'ChromaPipeline', None), + 'ChronoEdit': getattr(diffusers, 'ChronoEditPipeline', None), + 'CogView3': getattr(diffusers, 'CogView3PlusPipeline', None), + 'CogView4': getattr(diffusers, 'CogView4Pipeline', None), 'Cosmos': getattr(diffusers, 'Cosmos2TextToImagePipeline', None), - 'WanAI': getattr(diffusers, 'WanPipeline', None), - 'Qwen': getattr(diffusers, 'QwenImagePipeline', None), - 'Joy': getattr(diffusers, 'JoyImageEditPipeline', None), - 'HunyuanImage': getattr(diffusers, 'HunyuanImagePipeline', None), - 'ERNIE-Image': getattr(diffusers, 'ErnieImagePipeline', None), - 'Nucleus-Image': getattr(diffusers, 'NucleusMoEImagePipeline', None), - 'Z-Image': getattr(diffusers, 'ZImagePipeline', None), - 'Lens': getattr(diffusers, 'LensPipeline', None), + 'DeepFloydIF': getattr(diffusers, 'IFPipeline', None), + 'ERNIEImage': getattr(diffusers, 'ErnieImagePipeline', None), + 'FLUX': getattr(diffusers, 'FluxPipeline', None), + 'FLUX2Klein': getattr(diffusers, 'Flux2KleinPipeline', None), 'FLUX2': getattr(diffusers, 'Flux2Pipeline', None), - 'FLUX2 Klein': getattr(diffusers, 'Flux2KleinPipeline', None), - 'LongCat': getattr(diffusers, 'LongCatImagePipeline', None), 'GLM-Image': getattr(diffusers, 'GlmImagePipeline', None), - # dynamically imported and redefined later - 'Meissonic': getattr(diffusers, 'DiffusionPipeline', None), - 'OmniGen2': getattr(diffusers, 'DiffusionPipeline', None), - 'InstaFlow': getattr(diffusers, 'DiffusionPipeline', None), - 'SegMoE': getattr(diffusers, 'DiffusionPipeline', None), - 'FLite': getattr(diffusers, 'DiffusionPipeline', None), - 'Bria': getattr(diffusers, 'DiffusionPipeline', None), - 'X-Omni': getattr(diffusers, 'DiffusionPipeline', None), - 'HunyuanImage3': getattr(diffusers, 'DiffusionPipeline', None), - 'ChronoEdit': getattr(diffusers, 'DiffusionPipeline', None), - 'Anima': getattr(diffusers, 'DiffusionPipeline', None), - 'SDXS': getattr(diffusers, 'DiffusionPipeline', None), - 'Step1X-Edit': getattr(diffusers, 'DiffusionPipeline', None), - 'VIBE': getattr(diffusers, 'DiffusionPipeline', None), - 'Lumina-DiMOO': getattr(diffusers, 'DiffusionPipeline', None), + 'HiDream': getattr(diffusers, 'HiDreamImagePipeline', None), + 'HunyuanDiT': getattr(diffusers, 'HunyuanDiTPipeline', None), + 'HunyuanImage': getattr(diffusers, 'HunyuanImagePipeline', None), + 'JoyEdit': getattr(diffusers, 'JoyImageEditPipeline', None), + 'Kandinsky21': getattr(diffusers, 'KandinskyCombinedPipeline', None), + 'Kandinsky22': getattr(diffusers, 'KandinskyV22CombinedPipeline', None), + 'Kandinsky30': getattr(diffusers, 'Kandinsky3Pipeline', None), + 'Kandinsky50': getattr(diffusers, 'Kandinsky5T2IPipeline', None), + 'Kolors': getattr(diffusers, 'KolorsPipeline', None), + 'LongCat': getattr(diffusers, 'LongCatImagePipeline', None), + 'Lumina2': getattr(diffusers, 'Lumina2Pipeline', None), + 'LuminaNext': getattr(diffusers, 'LuminaText2ImgPipeline', None), + 'NucleusImage': getattr(diffusers, 'NucleusMoEImagePipeline', None), + 'OmniGen': getattr(diffusers, 'OmniGenPipeline', None), + 'PixArtAlpha': getattr(diffusers, 'PixArtAlphaPipeline', None), + 'PixArtSigma': getattr(diffusers, 'PixArtSigmaPipeline', None), + 'Qwen': getattr(diffusers, 'QwenImagePipeline', None), + 'Sana': getattr(diffusers, 'SanaPipeline', None), + 'WanAI': getattr(diffusers, 'WanPipeline', None), + 'ZImage': getattr(diffusers, 'ZImagePipeline', None), + + # pipelines with custom code that is fetched online + 'InstaFlow': OnlinePipeline, + 'Anima': OnlinePipeline, + + # sdnext custom pipelines are dynamically imported and redefined later + 'Bria': None, + 'FLite': None, + 'FLEX': None, + 'HiDreamO1': None, + 'HunyuanImage3': None, + 'Lens': None, + 'LuminaDiMOO': None, + 'Meissonic': None, + 'OmniGen2': None, + 'SDXS': None, + 'SegMoE': None, + 'Step1XEdit': None, + 'UltraFlux': None, + 'VIBE': None, + 'XOmni': None, + 'ZetaChroma': None, } @@ -128,6 +133,7 @@ def list_crossattention(): def get_pipelines(): + from modules.logger import log if hasattr(diffusers, 'OnnxStableDiffusionPipeline') and 'ONNX Stable Diffusion' not in list(pipelines): try: from modules.onnx_impl import initialize_onnx @@ -139,20 +145,21 @@ def get_pipelines(): 'ONNX Stable Diffusion Upscale': getattr(diffusers, 'OnnxStableDiffusionUpscalePipeline', None), } except Exception as e: - from modules.logger import log log.error(f'ONNX initialization error: {e}') onnx_pipelines = {} pipelines.update(onnx_pipelines) - if 'Lens' in pipelines and pipelines['Lens'] is None: - try: - import pipelines.lens as _lens - pipelines['Lens'] = getattr(diffusers, 'LensPipeline', None) - except Exception: - pass - for k, v in pipelines.items(): + stats_builtin = 0 + stats_custom = 0 + for k, v in pipelines.copy().items(): if k != 'Autodetect' and v is None: - from modules.logger import log - log.error(f'Model="{k}" diffusers={diffusers.__version__} path={diffusers.__file__} pipeline not available') + stats_custom += 1 + pipelines[k] = diffusers.DiffusionPipeline + else: + stats_builtin += 1 + if stats_custom > 0: + log.debug(f'Pipelines init: diffusers={stats_builtin} custom={stats_custom}') + else: + log.debug(f'Pipelines init: verified={stats_builtin}') return pipelines @@ -163,10 +170,6 @@ def get_repo(model): return 'stabilityai/stable-diffusion-xl-base-1.0' elif model == 'StableDiffusion3Pipeline' or model == 'Stable Diffusion 3': return 'stabilityai/stable-diffusion-3.5-medium' - elif model == 'FluxPipeline' or model == 'FLUX': - return 'black-forest-labs/FLUX.1-dev' - elif model == 'LensPipeline' or model == 'Lens': - return 'microsoft/Lens' else: return None diff --git a/pipelines/generic.py b/pipelines/generic.py index 7265ddb86..701a38360 100644 --- a/pipelines/generic.py +++ b/pipelines/generic.py @@ -1,6 +1,13 @@ from pipelines.generic_transformer import load_transformer from pipelines.generic_text_encoder import load_text_encoder from pipelines.generic_vae import load_vae_override +from pipelines.generic_util import get_loader, set_pipeline -__all__ = ["load_transformer", "load_text_encoder", "load_vae_override"] +__all__ = [ + "load_transformer", + "load_text_encoder", + "load_vae_override", + "get_loader", + "set_pipeline", +] diff --git a/pipelines/generic_test.py b/pipelines/generic_test.py index 467bdfbc5..9b75090f8 100644 --- a/pipelines/generic_test.py +++ b/pipelines/generic_test.py @@ -1,11 +1,14 @@ import os import importlib + +from scipy import stats from installer import log def test_pipelines(): from modules.sd_checkpoint import CheckpointInfo - log.info('Test pipelines...') + + log.info('Pipelines test...') pipelines = os.listdir("pipelines") pipelines.sort() for filename in pipelines: @@ -21,3 +24,36 @@ def test_pipelines(): load_func(ckpt) except Exception as e: log.error(f"Error: {module_name}.{attr}(): {e}") + + log.info('Pipelines verify...') + from modules.shared_items import get_pipelines + pipelines = get_pipelines() + stats_diffusers = 0 + stats_transformers = 0 + stats_custom = 0 + stats_deprecated = 0 + stats_fallback = 0 + stats_online = 0 + for name, cls in pipelines.items(): + if name == 'Autodetect' or name == 'AutoPipeline' or name == 'Diffusion' or name.startswith('ONNX'): + continue + elif cls is None: + log.warning(f"Pipeline: {name} not available") + elif cls.__name__ == 'DiffusionPipeline': + log.warning(f"Pipeline: {name} using fallback") + else: + if 'deprecated.' in str(cls): + log.warning(f"Pipeline: {name}={cls} deprecated") + stats_deprecated += 1 + if 'diffusers.pipelines.' in str(cls): + stats_diffusers += 1 + elif 'pipelines.' in str(cls): + stats_custom += 1 + elif 'transformers.' in str(cls): + stats_transformers += 1 + elif 'OnlinePipeline' in str(cls): + stats_online += 1 + else: + stats_fallback += 1 + log.warning(f"Pipeline: {name}={cls} not recognized") + log.info(f"Pipelines test: diffusers={stats_diffusers} transformers={stats_transformers} custom={stats_custom} deprecated={stats_deprecated} online={stats_online} fallback={stats_fallback}") diff --git a/pipelines/generic_util.py b/pipelines/generic_util.py index 92d026d1e..f71ff139f 100644 --- a/pipelines/generic_util.py +++ b/pipelines/generic_util.py @@ -9,3 +9,16 @@ def get_loader(component): if component == 'diffusers': return 'runai' if shared.opts.runai_streamer_diffusers else 'default' return 'runai' if shared.opts.runai_streamer_transformers else 'default' + + +def set_pipeline(name, cls): + """Set pipeline class in shared_items.pipelines with logging.""" + import diffusers + from modules.logger import log + from modules import shared_items + if shared_items.pipelines.get(name, None) is not None: + return + shared_items.pipelines[name] = cls + setattr(diffusers, name, cls) + setattr(diffusers.pipelines, name, cls) + log.debug(f'Pipeline: {name}={cls} initialized') diff --git a/pipelines/model_anima.py b/pipelines/model_anima.py index d7e8aa7f1..c4a57cf96 100644 --- a/pipelines/model_anima.py +++ b/pipelines/model_anima.py @@ -91,6 +91,7 @@ def load_anima(checkpoint_info, diffusers_load_config=None): diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["anima"] = AnimaTextToImagePipeline diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["anima"] = AnimaImageToImagePipeline diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["anima"] = AnimaInpaintPipeline + generic.set_pipeline('Anima', AnimaTextToImagePipeline) # UNET dropdown (shared.opts.sd_unet) may redirect the transformer to a # community file that bundles both the transformer and the llm_adapter. diff --git a/pipelines/model_bria.py b/pipelines/model_bria.py index d2ca94075..5f6bb3155 100644 --- a/pipelines/model_bria.py +++ b/pipelines/model_bria.py @@ -44,6 +44,7 @@ def load_bria(checkpoint_info, diffusers_load_config=None): else: cls = diffusers.BriaFiboPipeline diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING['bria-fibo'] = cls + generic.set_pipeline('Bria', cls) if repo_id is None or repo_id.lower() == 'none': return None @@ -74,6 +75,7 @@ def load_bria(checkpoint_info, diffusers_load_config=None): from pipelines.bria import BRIA_SPEC transformer = generic.load_transformer(repo_id, cls_name=BriaTransformer2DModel, load_config=diffusers_load_config, native_spec=BRIA_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config) + generic.set_pipeline('Bria', BriaPipeline) if repo_id is None or repo_id.lower() == 'none': return None diff --git a/pipelines/model_flex.py b/pipelines/model_flex.py index 655502417..dea13de27 100644 --- a/pipelines/model_flex.py +++ b/pipelines/model_flex.py @@ -18,6 +18,7 @@ def load_flex(checkpoint_info, diffusers_load_config=None): text_encoder_2 = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder_2") from pipelines.flex2 import Flex2Pipeline + generic.set_pipeline('FLEX', Flex2Pipeline) if repo_id is None or repo_id.lower() == 'none': return None pipe = Flex2Pipeline.from_pretrained( diff --git a/pipelines/model_flite.py b/pipelines/model_flite.py index cd0fead29..5162aa8cb 100644 --- a/pipelines/model_flite.py +++ b/pipelines/model_flite.py @@ -1,5 +1,4 @@ import sys -import diffusers import transformers from modules import shared, devices, sd_models, model_quant, sd_hijack_te from modules.logger import log @@ -16,7 +15,7 @@ def load_flite(checkpoint_info, diffusers_load_config=None): log.debug(f'Load model: type=FLite repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') from pipelines import f_lite - diffusers.FLitePipeline = f_lite.FLitePipeline + generic.set_pipeline('FLite', f_lite.FLitePipeline) sys.modules['f_lite'] = f_lite dit_model = generic.load_transformer(repo_id, cls_name=f_lite.DiT, load_config=diffusers_load_config, subfolder="dit_model", native_spec=f_lite.FLITE_SPEC) diff --git a/pipelines/model_hidream.py b/pipelines/model_hidream.py index 476c17edf..86cbaa33a 100644 --- a/pipelines/model_hidream.py +++ b/pipelines/model_hidream.py @@ -40,6 +40,7 @@ def load_hidream_o1(checkpoint_info, diffusers_load_config=None): from pipelines.hidream.hidream_o1 import HiDreamO1Pipeline, HiDreamO1ImagePipeline from pipelines.hidream.qwen3_vl_transformers import HiDreamO1Qwen3VLTransformer from pipelines.hidream.scheduler_flashfloweuler import FlashFlowMatchEulerDiscreteScheduler + generic.set_pipeline('HiDreamO1', HiDreamO1Pipeline) load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True, allow_quant=True) log.debug(f'Load model: type=HiDreamO1 repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') diff --git a/pipelines/model_hyimage.py b/pipelines/model_hyimage.py index a7f4057d1..7e6effb38 100644 --- a/pipelines/model_hyimage.py +++ b/pipelines/model_hyimage.py @@ -58,6 +58,7 @@ def load_hyimage3(checkpoint_info, diffusers_load_config=None): # pylint: disabl allow_quant = False load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True, allow_quant=allow_quant) + generic.set_pipeline('HunyuanImage3', transformers.AutoModelForCausalLM) if repo_id is None or repo_id.lower() == 'none': return None pipe = transformers.AutoModelForCausalLM.from_pretrained( diff --git a/pipelines/model_instaflow.py b/pipelines/model_instaflow.py new file mode 100644 index 000000000..0fdb95e59 --- /dev/null +++ b/pipelines/model_instaflow.py @@ -0,0 +1,22 @@ +import diffusers +from modules.logger import log +from modules import shared, devices, sd_models, model_quant +from pipelines import generic + + +def load_instaflow(checkpoint_info, diffusers_load_config=None): + if diffusers_load_config is None: + diffusers_load_config = {} + repo_id = sd_models.path_to_repo(checkpoint_info) + sd_models.hf_auth_check(checkpoint_info) + + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) + log.debug(f'Load model: type=InstaFlow repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + + if repo_id is None or repo_id.lower() == 'none': + return None + + pipeline = diffusers.utils.get_class_from_dynamic_module('instaflow_one_step', module_file='pipeline.py') + generic.set_pipeline('InstaFlow', pipeline) + sd_model = pipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) + return sd_model diff --git a/pipelines/model_lens.py b/pipelines/model_lens.py index e5d6d317f..596fbb95e 100644 --- a/pipelines/model_lens.py +++ b/pipelines/model_lens.py @@ -18,6 +18,10 @@ def load_lens(checkpoint_info, diffusers_load_config=None): transformer = generic.load_transformer(repo_id, cls_name=lens.LensTransformer2DModel, load_config=diffusers_load_config, native_spec=lens.LENS_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=lens.LensGptOssEncoder, load_config=diffusers_load_config, allow_quant=False) + diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["lens"] = lens.LensPipeline + diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["lens"] = lens.LensImg2ImgPipeline + diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["lens"] = lens.LensInpaintPipeline + generic.set_pipeline('Lens', lens.LensPipeline) if repo_id is None or repo_id.lower() == 'none': return None pipe = lens.LensPipeline.from_pretrained( @@ -31,9 +35,6 @@ def load_lens(checkpoint_info, diffusers_load_config=None): "output_type": "np", "enable_reasoner": shared.opts.model_lens_enable_pe, } - diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["lens"] = lens.LensPipeline - diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["lens"] = lens.LensImg2ImgPipeline - diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["lens"] = lens.LensInpaintPipeline sd_hijack_te.init_hijack(pipe) sd_hijack_vae.init_hijack(pipe) diff --git a/pipelines/model_lumina.py b/pipelines/model_lumina.py index 947aa7051..90b3b77e0 100644 --- a/pipelines/model_lumina.py +++ b/pipelines/model_lumina.py @@ -79,6 +79,7 @@ def load_lumina_dimoo(checkpoint_info, diffusers_load_config=None): from pipelines.lumina_dimmo.pipelines import LuminaDiMOOTextPipeline, LuminaDiMOOImagePipeline diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["luminadimoo"] = LuminaDiMOOTextPipeline diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["luminadimoo"] = LuminaDiMOOImagePipeline + generic.set_pipeline('LuminaDiMOO', LuminaDiMOOTextPipeline) if repo_id is None or repo_id.lower() == 'none': return None diff --git a/pipelines/model_omnigen.py b/pipelines/model_omnigen.py index db9e38d18..b3b1bdb0c 100644 --- a/pipelines/model_omnigen.py +++ b/pipelines/model_omnigen.py @@ -1,6 +1,7 @@ import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te from modules.logger import log +from pipelines import generic def load_omnigen(checkpoint_info, diffusers_load_config=None): # pylint: disable=unused-argument @@ -42,10 +43,10 @@ def load_omnigen2(checkpoint_info, diffusers_load_config=None): # pylint: disabl sd_models.hf_auth_check(checkpoint_info) from pipelines.omnigen2 import OmniGen2Pipeline, OmniGen2Transformer2DModel, Qwen2_5_VLForConditionalGeneration - diffusers.OmniGen2Pipeline = OmniGen2Pipeline # monkey-pathch - diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["omnigen2"] = diffusers.OmniGen2Pipeline - diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["omnigen2"] = diffusers.OmniGen2Pipeline - diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["omnigen2"] = diffusers.OmniGen2Pipeline + diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["omnigen2"] = OmniGen2Pipeline + diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["omnigen2"] = OmniGen2Pipeline + diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["omnigen2"] = OmniGen2Pipeline + generic.set_pipeline('OmniGen2', OmniGen2Pipeline) load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='Model') log.debug(f'Load model: type=OmniGen2 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={diffusers_load_config}') diff --git a/pipelines/model_sdxs.py b/pipelines/model_sdxs.py index 8a59e1278..55f775bbf 100644 --- a/pipelines/model_sdxs.py +++ b/pipelines/model_sdxs.py @@ -6,6 +6,11 @@ from modules.logger import log from pipelines import generic +class SDXSPipeline(diffusers.DiffusionPipeline): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def hijack_encode_text(prompt: str | list[str]): jobid = shared.state.begin('TE Encode') t0 = time.time() @@ -38,9 +43,11 @@ def load_sdxs(checkpoint_info, diffusers_load_config=None): text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3_5ForConditionalGeneration, load_config=diffusers_load_config, allow_shared=False) + generic.set_pipeline('SDXS', SDXSPipeline) + if repo_id is None or repo_id.lower() == 'none': return None - pipe = diffusers.DiffusionPipeline.from_pretrained( + pipe = SDXSPipeline.from_pretrained( repo_id, text_encoder=text_encoder, cache_dir=shared.opts.diffusers_dir, diff --git a/pipelines/model_segmoe.py b/pipelines/model_segmoe.py new file mode 100644 index 000000000..3a35a0c2c --- /dev/null +++ b/pipelines/model_segmoe.py @@ -0,0 +1,22 @@ +from modules.logger import log +from modules import shared, devices, sd_models, model_quant +from pipelines import generic + + +def load_segmoe(checkpoint_info, diffusers_load_config=None): + if diffusers_load_config is None: + diffusers_load_config = {} + repo_id = sd_models.path_to_repo(checkpoint_info) + sd_models.hf_auth_check(checkpoint_info) + + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) + log.debug(f'Load model: type=SegMoE repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + + from pipelines.segmoe.segmoe_model import SegMoEPipeline + generic.set_pipeline('SegMoE', SegMoEPipeline) + if repo_id is None or repo_id.lower() == 'none': + return None + + sd_model = SegMoEPipeline(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) + sd_model = sd_model.pipe # segmoe pipe does its stuff in __init__ and __call__ is the original pipeline + return sd_model diff --git a/pipelines/model_step1x_edit.py b/pipelines/model_step1x_edit.py index 113e1a374..59713465b 100644 --- a/pipelines/model_step1x_edit.py +++ b/pipelines/model_step1x_edit.py @@ -19,6 +19,7 @@ def load_step1x_edit(checkpoint_info, diffusers_load_config=None): diffusers.Step1XEditPipeline = Step1XEditPipeline diffusers.Step1XEditTransformer2DModel = Step1XEditTransformer2DModel + generic.set_pipeline('Step1XEdit', Step1XEditPipeline) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen2_5_VLForConditionalGeneration, load_config=diffusers_load_config) from pipelines.step1x import STEP1X_SPEC diff --git a/pipelines/model_ultraflux.py b/pipelines/model_ultraflux.py index 217341b97..c78710445 100644 --- a/pipelines/model_ultraflux.py +++ b/pipelines/model_ultraflux.py @@ -17,6 +17,8 @@ def load_ultraflux(checkpoint_info, diffusers_load_config=None): from pipelines.ultraflux.pipeline_flux import UltraFluxPipeline from pipelines.ultraflux.transformer_flux import FluxTransformer2DModel from pipelines.ultraflux.autoencoder_kl import AutoencoderUltraFluxKL + diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING['ultraflux'] = UltraFluxPipeline + generic.set_pipeline('UltraFlux', UltraFluxPipeline) transformer = generic.load_transformer(repo_id, cls_name=FluxTransformer2DModel, load_config=diffusers_load_config) text_encoder_2 = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder='text_encoder_2') @@ -46,8 +48,6 @@ def load_ultraflux(checkpoint_info, diffusers_load_config=None): if hasattr(pipe.scheduler.config, 'time_shift'): pipe.scheduler.config.time_shift = 4 - diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING['ultraflux'] = UltraFluxPipeline - del text_encoder_2 del transformer del vae diff --git a/pipelines/model_vibe.py b/pipelines/model_vibe.py index 98b893805..5ba897f9c 100644 --- a/pipelines/model_vibe.py +++ b/pipelines/model_vibe.py @@ -18,7 +18,7 @@ def load_vibe(checkpoint_info, diffusers_load_config=None): from pipelines.vibe import VIBESanaEditingModel, VIBESanaEditingPipeline, VIBESanaImagePipeline diffusers.VIBESanaEditingPipeline = VIBESanaEditingPipeline diffusers.VIBESanaEditingModel = VIBESanaEditingModel - + generic.set_pipeline('VIBE', VIBESanaEditingPipeline) sys.modules['vibe.transformer.vibe_sana_editing'] = diffusers # monkey patch since hf model_index.json points to custom class path from pipelines.vibe import VIBE_SPEC diff --git a/pipelines/model_xomni.py b/pipelines/model_xomni.py index f9cd7ced9..70c49072d 100644 --- a/pipelines/model_xomni.py +++ b/pipelines/model_xomni.py @@ -3,6 +3,7 @@ import transformers import diffusers from modules import shared, devices, sd_models, model_quant from modules.logger import log +from pipelines import generic class XOmniPipeline(diffusers.DiffusionPipeline): @@ -112,6 +113,7 @@ def load_xomni(checkpoint_info, diffusers_load_config=None): sd_models.hf_auth_check(checkpoint_info) pipe = XOmniPipeline() + generic.set_pipeline('XOmni', XOmniPipeline) if repo_id is None or repo_id.lower() == 'none': return None diff --git a/pipelines/model_zetachroma.py b/pipelines/model_zetachroma.py index 22e6374a0..3e9a2551a 100644 --- a/pipelines/model_zetachroma.py +++ b/pipelines/model_zetachroma.py @@ -23,8 +23,7 @@ def load_zetachroma(checkpoint_info, diffusers_load_config=None): ) from pipelines import generic, zetachroma - - diffusers.ZetaChromaPipeline = zetachroma.ZetaChromaPipeline + generic.set_pipeline('ZetaChroma', zetachroma.ZetaChromaPipeline) sys.modules["zetachroma"] = zetachroma if repo_id is None or repo_id.lower() == 'none':