From 39e5d614ff6449a2cb9f237e09f2f4e37ffbecac Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 13 Nov 2023 16:37:41 -0500 Subject: [PATCH] consistency decoder --- CHANGELOG.md | 3 ++- modules/processing_diffusers.py | 2 +- modules/sd_vae.py | 10 ++++++++-- 3 files changed, 11 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b6a4bbb4b..b8111bfc7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,7 +21,8 @@ - Updated logic for calculating **steps** when using base/hires/refiner workflows - Safe model offloading for non-standard models - Fix **DPM SDE** scheduler - - Better support for SD 1.5 **inpainting** models + - Better support for SD 1.5 **inpainting** models + - Add support for **OpenAI Consistency decoder VAE** - Update to `diffusers==0.23.0` - **Extra networks** - Use multi-threading for 5x load speedup diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 0ea8292ac..9a56fc028 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -98,7 +98,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro model.vae.to(devices.device) latents.to(model.vae.device) - upcast = (model.vae.dtype == torch.float16) and model.vae.config.force_upcast and hasattr(model, 'upcast_vae') + upcast = (model.vae.dtype == torch.float16) and getattr(model.vae.config, 'force_upcast', False) and hasattr(model, 'upcast_vae') if upcast: # this is done by diffusers automatically if output_type != 'latent' model.upcast_vae() latents = latents.to(next(iter(model.vae.post_quant_conv.parameters())).dtype) diff --git a/modules/sd_vae.py b/modules/sd_vae.py index 36aac1df0..9e28e4c73 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -203,10 +203,16 @@ def load_vae_diffusers(model_file, vae_file=None, vae_source="unknown-source"): if os.path.isfile(vae_file): _pipeline, model_type = sd_models.detect_pipeline(model_file, 'vae') diffusers_load_config = { "config_file": paths.sd_default_config if model_type != 'Stable Diffusion XL' else os.path.join(paths.sd_configs_path, 'sd_xl_base.yaml')} - vae = diffusers.AutoencoderKL.from_single_file(vae_file, **diffusers_load_config) + if os.path.getsize(vae_file) > 1310944880: + vae = diffusers.ConsistencyDecoderVAE.from_pretrained('openai/consistency-decoder', **diffusers_load_config) # consistency decoder does not have from single file, so we'll just download it once more + else: + vae = diffusers.AutoencoderKL.from_single_file(vae_file, **diffusers_load_config) vae = vae.to(devices.dtype_vae) else: - vae = diffusers.AutoencoderKL.from_pretrained(vae_file, **diffusers_load_config) + if 'consistency-decoder' in vae_file: + vae = diffusers.ConsistencyDecoderVAE.from_pretrained(vae_file, **diffusers_load_config) + else: + vae = diffusers.AutoencoderKL.from_pretrained(vae_file, **diffusers_load_config) global loaded_vae_file # pylint: disable=global-statement loaded_vae_file = os.path.basename(vae_file) # shared.log.debug(f'Diffusers VAE config: {vae.config}')