From 52efa909c57dfa9915fa7d10bd584d3afee589a1 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 14 Feb 2024 14:51:20 +0300 Subject: [PATCH] Experimental Stable Cascade support --- modules/sd_models.py | 18 ++++++++++++++++-- modules/shared_items.py | 1 + 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/modules/sd_models.py b/modules/sd_models.py index 28992fe5e..c522152b6 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -594,6 +594,10 @@ def detect_pipeline(f: str, op: str = 'model', warning=True): if shared.backend == shared.Backend.ORIGINAL: warn(f'Model detected as PixArt Alpha model, but attempting to load using backend=original: {op}={f} size={size} MB') guess = 'PixArt Alpha' + if 'stable-cascade' in f.lower() or 'stablecascade' in f.lower(): + if shared.backend == shared.Backend.ORIGINAL: + warn(f'Model detected as Stable Cascade model, but attempting to load using backend=original: {op}={f} size={size} MB') + guess = 'Stable Cascade' # switch for specific variant if guess == 'Stable Diffusion' and 'inpaint' in f.lower(): guess = 'Stable Diffusion Inpaint' @@ -813,14 +817,24 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No files = shared.walk_files(checkpoint_info.path, ['.safetensors', '.bin', '.ckpt']) if 'variant' not in diffusers_load_config and any('diffusion_pytorch_model.fp16' in f for f in files): # deal with diffusers lack of variant fallback when loading diffusers_load_config['variant'] = 'fp16' - if model_type in ['InstaFlow']: # forced pipeline + if model_type in ['Stable Cascade']: # forced pipeline + try: + from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline, StableCascadeCombinedPipeline + # set prior manually for now + prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) + decoder = StableCascadeDecoderPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) + sd_model = StableCascadeCombinedPipeline(tokenizer=decoder.tokenizer, text_encoder=decoder.text_encoder, decoder=decoder.decoder, scheduler=decoder.scheduler, vqgan=decoder.vqgan, prior_prior=prior.prior, prior_scheduler=prior.scheduler, feature_extractor=prior.feature_extractor, image_encoder=prior.image_encoder) + except Exception as e: + shared.log.error(f'Diffusers Failed loading {op}: {checkpoint_info.path} {e}') + return + elif model_type in ['InstaFlow']: # forced pipeline try: pipeline = diffusers.utils.get_class_from_dynamic_module('instaflow_one_step', module_file='pipeline.py') sd_model = pipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) except Exception as e: shared.log.error(f'Diffusers Failed loading {op}: {checkpoint_info.path} {e}') return - if model_type in ['SegMoE']: # forced pipeline + elif model_type in ['SegMoE']: # forced pipeline try: from modules.segmoe.segmoe_model import SegMoEPipeline sd_model = SegMoEPipeline(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) diff --git a/modules/shared_items.py b/modules/shared_items.py index 24bf66caa..162f1ec52 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -54,6 +54,7 @@ def get_pipelines(): 'PixArt Alpha': getattr(diffusers, 'PixArtAlphaPipeline', None), 'UniDiffuser': getattr(diffusers, 'UniDiffuserPipeline', None), 'Wuerstchen': getattr(diffusers, 'WuerstchenCombinedPipeline', None), + 'Stable Cascade': getattr(diffusers, 'StableCascadeCombinedPipeline', None), 'Kandinsky 2.1': getattr(diffusers, 'KandinskyPipeline', None), 'Kandinsky 2.2': getattr(diffusers, 'KandinskyV22Pipeline', None), 'Kandinsky 3': getattr(diffusers, 'Kandinsky3Pipeline', None),