diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 51e295e02..05b92aeef 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -6,6 +6,7 @@ import modules.shared as shared import modules.sd_samplers as sd_samplers import modules.sd_models as sd_models import modules.sd_vae as sd_vae +import modules.taesd.sd_vae_taesd as sd_vae_taesd import modules.images as images from modules.lora_diffusers import lora_state, unload_diffusers_lora from modules.processing import StableDiffusionProcessing @@ -43,6 +44,14 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro else: return latents + def taesd_vae_decode(latents, model, output_type='np'): + shared.log.debug('Diffusers VAE decode: name=TAESD') + decoded = torch.zeros((len(latents), 3, p.height, p.width), dtype=devices.dtype_vae, device=devices.device) + for i in range(len(output.images)): + decoded[i] = (sd_vae_taesd.decode(latents[i]) * 2.0) - 1.0 + images = model.image_processor.postprocess(decoded, output_type=output_type) + return images + def set_pipeline_args(model, prompt: str, negative_prompt: str, prompt_2: typing.Optional[str] =None, negative_prompt_2: typing.Optional[str] = None, is_refiner: bool = False, **kwargs): args = {} @@ -168,7 +177,10 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro return results if shared.sd_refiner is None or not p.enable_hr: - output.images = vae_decode(output.images, shared.sd_model) + if shared.opts.diffusers_taesd_vae_output: + output.images = taesd_vae_decode(output.images, shared.sd_model) + else: + output.images = vae_decode(output.images, shared.sd_model) if lora_state['active']: unload_diffusers_lora() diff --git a/modules/shared.py b/modules/shared.py index 25993ea99..ee601e0df 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -396,6 +396,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), { options_templates.update(options_section(('diffusers', "Diffusers Settings"), { "diffusers_allow_safetensors": OptionInfo(True, 'Diffusers allow loading from safetensors files'), + "diffusers_taesd_vae_output": OptionInfo(False, 'Use TAESD VAE for quick image outputs'), "diffusers_pipeline": OptionInfo(pipelines[0], 'Diffusers pipeline', gr.Dropdown, lambda: {"choices": pipelines}), "diffusers_move_base": OptionInfo(False, "Move base model to CPU when using refiner"), "diffusers_move_refiner": OptionInfo(True, "Move refiner model to CPU when not in use"),