refactor taesd and add multiple variants in settings

Signed-off-by: Vladimir Mandic <mandic00@live.com>

modules/loader.py

@@ -74,7 +74,7 @@ timer.startup.record("diffusers")
 
 try:
     import pillow_jxl # pylint: disable=W0611,C0411
-except:
+except Exception:
     pass
 from PIL import Image # pylint: disable=W0611,C0411
 timer.startup.record("pillow")

modules/model_flux.py

@@ -241,7 +241,8 @@ def load_transformer(file_path): # triggered by opts.sd_unet change
             if transformer is not None:
                 return transformer
         shared.log.info(f'Load module: type=UNet/Transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} quant=none dtype={devices.dtype}')
-        # shared.log.warning('Load module: type=UNet/Transformer does not support load-time quantization') # TODO flux transformer from-single-file with quant
+        # TODO flux transformer from-single-file with quant
+        # shared.log.warning('Load module: type=UNet/Transformer does not support load-time quantization')
         transformer = diffusers.FluxTransformer2DModel.from_single_file(file_path, **diffusers_load_config)
     if transformer is None:
         shared.log.error('Failed to load UNet model')

modules/processing_args.py

@@ -1,7 +1,6 @@
 import typing
 import os
 import re
-import copy
 import math
 import time
 import inspect

modules/processing_vae.py

@@ -239,6 +239,8 @@ def vae_decode(latents, model, output_type='np', full_quality=True, width=None,
         decoded = full_vqgan_decode(latents=latents, model=model)
     else:
         decoded = taesd_vae_decode(latents=latents)
+        if torch.is_tensor(decoded):
+            decoded = 2.0 * decoded - 1.0 # typical normalized range
 
     if torch.is_tensor(decoded):
         if hasattr(model, 'video_processor'):

modules/sd_samplers_common.py

@@ -59,7 +59,7 @@ def single_sample_to_image(sample, approximation=None):
                 except Exception:
                     pass
             x_sample = sd_vae_taesd.decode(sample)
-            x_sample = (1.0 + x_sample) / 2.0 # preview requires smaller range
+            # x_sample = (1.0 + x_sample) / 2.0 # preview requires smaller range
         elif shared.sd_model_type == 'sc' and approximation != 3:
             x_sample = sd_vae_stablecascade.decode(sample)
         elif approximation == 0: # Simple
@@ -67,19 +67,24 @@ def single_sample_to_image(sample, approximation=None):
         elif approximation == 1: # Approximate
             x_sample = sd_vae_approx.nn_approximation(sample) * 0.5 + 0.5
             if shared.sd_model_type == "sdxl":
-                x_sample = x_sample[[2,1,0], :, :] # BGR to RGB
+                x_sample = x_sample[[2, 1, 0], :, :] # BGR to RGB
         elif approximation == 3: # Full VAE
             x_sample = processing.decode_first_stage(shared.sd_model, sample.unsqueeze(0))[0]
         else:
             warn_once(f"Unknown latent decode type: {approximation}")
             return Image.new(mode="RGB", size=(512, 512))
         try:
-            if x_sample.shape[0] > 4:
-                return Image.new(mode="RGB", size=(512, 512))
-            if x_sample.dtype == torch.bfloat16:
-                x_sample.to(torch.float16)
-            transform = T.ToPILImage()
-            image = transform(x_sample)
+            if isinstance(x_sample, Image.Image):
+                image = x_sample
+            else:
+                if x_sample.shape[0] > 4 or x_sample.shape[0] == 4:
+                    return Image.new(mode="RGB", size=(512, 512))
+                if x_sample.dtype == torch.bfloat16:
+                    x_sample = x_sample.to(torch.float16)
+                if len(x_sample.shape) == 4:
+                    x_sample = x_sample[0]
+                transform = T.ToPILImage()
+                image = transform(x_sample)
         except Exception as e:
             warn_once(f'Preview: {e}')
             image = Image.new(mode="RGB", size=(512, 512))

modules/sd_samplers_diffusers.py

@@ -7,8 +7,8 @@ from modules import shared, errors
 from modules.sd_samplers_common import SamplerData, flow_models
 
 
-debug = shared.log.trace if os.environ.get('SD_SAMPLER_DEBUG', None) is not None else lambda *args, **kwargs: None
-debug('Trace: SAMPLER')
+debug = os.environ.get('SD_SAMPLER_DEBUG', None) is not None
+debug_log = shared.log.trace if debug else lambda *args, **kwargs: None
 
 try:
     from diffusers import (
@@ -178,17 +178,17 @@ class DiffusionSampler:
             model.default_scheduler = copy.deepcopy(model.scheduler)
         for key, value in config.get('All', {}).items(): # apply global defaults
             self.config[key] = value
-        debug(f'Sampler: all="{self.config}"')
+        debug_log(f'Sampler: all="{self.config}"')
         if hasattr(model.default_scheduler, 'scheduler_config'): # find model defaults
             orig_config = model.default_scheduler.scheduler_config
         else:
             orig_config = model.default_scheduler.config
-        debug(f'Sampler: diffusers="{self.config}"')
-        debug(f'Sampler: original="{orig_config}"')
+        debug_log(f'Sampler: diffusers="{self.config}"')
+        debug_log(f'Sampler: original="{orig_config}"')
         for key, value in orig_config.items(): # apply model defaults
             if key in self.config:
                 self.config[key] = value
-        debug(f'Sampler: default="{self.config}"')
+        debug_log(f'Sampler: default="{self.config}"')
         for key, value in config.get(name, {}).items(): # apply diffusers per-scheduler defaults
             self.config[key] = value
         for key, value in kwargs.items(): # apply user args, if any
@@ -267,15 +267,22 @@ class DiffusionSampler:
             if key not in possible:
                 # shared.log.warning(f'Sampler: sampler="{name}" config={self.config} invalid={key}')
                 del self.config[key]
-        debug(f'Sampler: name="{name}"')
-        debug(f'Sampler: config={self.config}')
-        debug(f'Sampler: signature={possible}')
-        # shared.log.debug(f'Sampler: sampler="{name}" config={self.config}')
-        sampler = constructor(**self.config)
+        debug_log(f'Sampler: name="{name}"')
+        debug_log(f'Sampler: config={self.config}')
+        debug_log(f'Sampler: signature={possible}')
+        # shared.log.debug_log(f'Sampler: sampler="{name}" config={self.config}')
+        try:
+            sampler = constructor(**self.config)
+        except Exception as e:
+            shared.log.error(f'Sampler: sampler="{name}" {e}')
+            if debug:
+                errors.display(e, 'Samplers')
+            self.sampler = None
+            return
         accept_sigmas = "sigmas" in set(inspect.signature(sampler.set_timesteps).parameters.keys())
         accepts_timesteps = "timesteps" in set(inspect.signature(sampler.set_timesteps).parameters.keys())
         accept_scale_noise = hasattr(sampler, "scale_noise")
-        debug(f'Sampler: sampler="{name}" sigmas={accept_sigmas} timesteps={accepts_timesteps}')
+        debug_log(f'Sampler: sampler="{name}" sigmas={accept_sigmas} timesteps={accepts_timesteps}')
         if ('Flux' in model.__class__.__name__) and (not accept_sigmas):
             shared.log.warning(f'Sampler: sampler="{name}" does not accept sigmas')
             self.sampler = None
@@ -289,5 +296,5 @@ class DiffusionSampler:
             if not hasattr(self.sampler, 'dc_ratios'):
                 pass
                 # self.sampler.dc_ratios = self.sampler.cascade_polynomial_regression(test_CFG=6.0, test_NFE=10, cpr_path='tmp/sd2.1.npy')
-        # shared.log.debug(f'Sampler: class="{self.sampler.__class__.__name__}" config={self.sampler.config}')
+        # shared.log.debug_log(f'Sampler: class="{self.sampler.__class__.__name__}" config={self.sampler.config}')
         self.sampler.name = name

modules/sd_vae_taesd.py

@@ -5,217 +5,132 @@ Tiny AutoEncoder for Stable Diffusion
 https://github.com/madebyollin/taesd
 """
 import os
+import threading
 from PIL import Image
 import torch
-import torch.nn as nn
 from modules import devices, paths
 
 
-taesd_models = {
-    'sd-decoder': None,
-    'sd-encoder': None,
-    'sdxl-decoder': None,
-    'sdxl-encoder': None,
-    'sd3-decoder': None,
-    'sd3-encoder': None,
-    'f1-decoder': None,
-    'f1-encoder': None,
+TAESD_MODELS = {
+    'TAESD 1.3 Mocha Croissant': { 'fn': 'taesd_13_', 'uri': 'https://github.com/madebyollin/taesd/raw/7f572ca629c9b0d3c9f71140e5f501e09f9ea280', 'model': None },
+    'TAESD 1.2 Chocolate-Dipped Shortbread': { 'fn': 'taesd_12_', 'uri': 'https://github.com/madebyollin/taesd/raw/8909b44e3befaa0efa79c5791e4fe1c4d4f7884e', 'model': None },
+    'TAESD 1.1 Fruit Loops': { 'fn': 'taesd_11_', 'uri': 'https://github.com/madebyollin/taesd/raw/3e8a8a2ab4ad4079db60c1c7dc1379b4cc0c6b31', 'model': None },
+    'TAESD 1.0': { 'fn': 'taesd_10_', 'uri': 'https://github.com/madebyollin/taesd/raw/88012e67cf0454e6d90f98911fe9d4aef62add86', 'model': None },
 }
-previous_warnings = False
+CQYAN_MODELS = {
+    'Hybrid-Tiny SD': {
+        'sd': { 'repo': 'cqyan/hybrid-sd-tinyvae', 'model': None },
+        'sdxl': { 'repo': 'cqyan/hybrid-sd-tinyvae-xl', 'model': None },
+    },
+    'Hybrid-Small SD': {
+        'sd': { 'repo': 'cqyan/hybrid-sd-small-vae', 'model': None },
+        'sdxl': { 'repo': 'cqyan/hybrid-sd-small-vae-xl', 'model': None },
+    },
+}
+
+prev_warnings = False
+prev_cls = ''
+prev_type = ''
+prev_model = ''
+lock = threading.Lock()
 
 
-def conv(n_in, n_out, **kwargs):
-    return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs)
-
-class Clamp(nn.Module):
-    def forward(self, x):
-        return torch.tanh(x / 3) * 3
-
-class Block(nn.Module):
-    def __init__(self, n_in, n_out):
-        super().__init__()
-        self.conv = nn.Sequential(conv(n_in, n_out), nn.ReLU(), conv(n_out, n_out), nn.ReLU(), conv(n_out, n_out))
-        self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity()
-        self.fuse = nn.ReLU()
-    def forward(self, x):
-        return self.fuse(self.conv(x) + self.skip(x))
-
-def Encoder(latent_channels=4):
-    return nn.Sequential(
-        conv(3, 64), Block(64, 64),
-        conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64),
-        conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64),
-        conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64),
-        conv(64, latent_channels),
-    )
-
-def Decoder(latent_channels=4):
+def warn_once(msg):
     from modules import shared
-    if shared.opts.live_preview_taesd_layers == 1:
-        return nn.Sequential(
-            Clamp(), conv(latent_channels, 64), nn.ReLU(),
-            Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
-            Block(64, 64), Block(64, 64), Block(64, 64), nn.Identity(), conv(64, 64, bias=False),
-            Block(64, 64), Block(64, 64), Block(64, 64), nn.Identity(), conv(64, 64, bias=False),
-            Block(64, 64), conv(64, 3),
-        )
-    elif shared.opts.live_preview_taesd_layers == 2:
-        return nn.Sequential(
-            Clamp(), conv(latent_channels, 64), nn.ReLU(),
-            Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
-            Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
-            Block(64, 64), Block(64, 64), Block(64, 64), nn.Identity(), conv(64, 64, bias=False),
-            Block(64, 64), conv(64, 3),
-        )
-    else:
-        return nn.Sequential(
-            Clamp(), conv(latent_channels, 64), nn.ReLU(),
-            Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
-            Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
-            Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
-            Block(64, 64), conv(64, 3),
-        )
+    global prev_warnings # pylint: disable=global-statement
+    if not prev_warnings:
+        prev_warnings = True
+        shared.log.error(f'Decode: type="taesd" variant="{shared.opts.taesd_variant}": {msg}')
+    return Image.new('RGB', (8, 8), color = (0, 0, 0))
 
 
-class TAESD(nn.Module): # pylint: disable=abstract-method
-    latent_magnitude = 3
-    latent_shift = 0.5
-
-    def __init__(self, encoder_path="taesd_encoder.pth", decoder_path="taesd_decoder.pth", latent_channels=None):
-        """Initialize pretrained TAESD on the given device from the given checkpoints."""
-        super().__init__()
-        if latent_channels is None:
-            latent_channels = self.guess_latent_channels(str(decoder_path), str(encoder_path))
-        self.encoder = Encoder(latent_channels)
-        self.decoder = Decoder(latent_channels)
-        if encoder_path is not None:
-            self.encoder.load_state_dict(torch.load(encoder_path, map_location="cpu"), strict=False)
-        if decoder_path is not None:
-            self.decoder.load_state_dict(torch.load(decoder_path, map_location="cpu"), strict=False)
-
-    def guess_latent_channels(self, decoder_path, encoder_path):
-        """guess latent channel count based on encoder filename"""
-        if "taef1" in encoder_path or "taef1" in decoder_path:
-            return 16
-        if "taesd3" in encoder_path or "taesd3" in decoder_path:
-            return 16
-        return 4
-
-    @staticmethod
-    def scale_latents(x):
-        """raw latents -> [0, 1]"""
-        return x.div(2 * TAESD.latent_magnitude).add(TAESD.latent_shift).clamp(0, 1)
-
-    @staticmethod
-    def unscale_latents(x):
-        """[0, 1] -> raw latents"""
-        return x.sub(TAESD.latent_shift).mul(2 * TAESD.latent_magnitude)
-
-
-def download_model(model_path):
-    model_name = os.path.basename(model_path)
-    model_url = f'https://github.com/madebyollin/taesd/raw/main/{model_name}'
-    if not os.path.exists(model_path):
-        from modules.shared import log
-        os.makedirs(os.path.dirname(model_path), exist_ok=True)
-        log.info(f'Downloading TAESD decoder: {model_path}')
-        torch.hub.download_url_to_file(model_url, model_path)
-
-
-def model(model_class = 'sd', model_type = 'decoder'):
-    vae = taesd_models[f'{model_class}-{model_type}']
-    if vae is None:
-        model_path = os.path.join(paths.models_path, "TAESD", f"tae{model_class}_{model_type}.pth")
-        download_model(model_path)
-        if os.path.exists(model_path):
-            from modules.shared import log
-            taesd_models[f'{model_class}-{model_type}'] = TAESD(decoder_path=model_path, encoder_path=None) if model_type == 'decoder' else TAESD(encoder_path=model_path, decoder_path=None)
-            vae = taesd_models[f'{model_class}-{model_type}']
-            vae.eval()
-            vae.to(devices.device, devices.dtype_vae)
-            log.info(f"Load VAE-TAESD: model={model_path}")
-        else:
-            raise FileNotFoundError(f'TAESD model not found: {model_path}')
-    if vae is None:
+def get_model(model_type = 'decoder'):
+    global prev_cls, prev_type, prev_model # pylint: disable=global-statement
+    from modules import shared
+    cls = shared.sd_model_type
+    if cls == 'ldm':
+        cls = 'sd'
+    folder = os.path.join(paths.models_path, "TAESD")
+    os.makedirs(folder, exist_ok=True)
+    if 'sd' not in cls and 'f1' not in cls:
+        warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported')
         return None
+    if shared.opts.taesd_variant.startswith('TAESD'):
+        cfg = TAESD_MODELS[shared.opts.taesd_variant]
+        if (cls == prev_cls) and (model_type == prev_type) and (shared.opts.taesd_variant == prev_model) and (cfg['model'] is not None):
+            return cfg['model']
+        fn = os.path.join(folder, cfg['fn'] + cls + '_' + model_type + '.pth')
+        if not os.path.exists(fn):
+            uri = cfg['uri'] + '/tae' + cls + '_' + model_type + '.pth'
+            try:
+                shared.log.info(f'Decode: type="taesd" variant="{shared.opts.taesd_variant}": uri="{uri}" fn="{fn}" download')
+                torch.hub.download_url_to_file(uri, fn)
+            except Exception as e:
+                warn_once(f'download uri={uri} {e}')
+        if os.path.exists(fn):
+            prev_cls = cls
+            prev_type = model_type
+            prev_model = shared.opts.taesd_variant
+            shared.log.debug(f'Decode: type="taesd" variant="{shared.opts.taesd_variant}" fn="{fn}" load')
+            from modules.taesd.taesd import TAESD
+            TAESD_MODELS[shared.opts.taesd_variant]['model'] = TAESD(decoder_path=fn if model_type=='decoder' else None, encoder_path=fn if model_type=='encoder' else None)
+            return TAESD_MODELS[shared.opts.taesd_variant]['model']
+    elif shared.opts.taesd_variant.startswith('Hybrid'):
+        cfg = CQYAN_MODELS[shared.opts.taesd_variant].get(cls, None)
+        if (cls == prev_cls) and (model_type == prev_type) and (shared.opts.taesd_variant == prev_model) and (cfg['model'] is not None):
+            return cfg['model']
+        if cfg is None:
+            warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported')
+            return None
+        repo = cfg['repo']
+        prev_cls = cls
+        prev_type = model_type
+        prev_model = shared.opts.taesd_variant
+        shared.log.debug(f'Decode: type="taesd" variant="{shared.opts.taesd_variant}" id="{repo}" load')
+        dtype = devices.dtype_vae if devices.dtype_vae != torch.bfloat16 else torch.float16 # taesd does not support bf16
+        if 'tiny' in repo:
+            from diffusers.models import AutoencoderTiny
+            vae = AutoencoderTiny.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir, torch_dtype=dtype)
+        else:
+            from modules.taesd.hybrid_small import AutoencoderSmall
+            vae = AutoencoderSmall.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir, torch_dtype=dtype)
+        vae = vae.to(devices.device, dtype=dtype)
+        CQYAN_MODELS[shared.opts.taesd_variant][cls]['model'] = vae
+        return vae
     else:
-        return vae.decoder if model_type == 'decoder' else vae.encoder
+        warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported')
+    return None
 
 
 def decode(latents):
-    global previous_warnings # pylint: disable=global-statement
-    from modules import shared
-    model_class = shared.sd_model_type
-    if model_class == 'ldm':
-        model_class = 'sd'
-    dtype = devices.dtype_vae if devices.dtype_vae != torch.bfloat16 else torch.float16 # taesd does not support bf16
-    if 'sd' not in model_class and 'f1' not in model_class:
-        if not previous_warnings:
-            previous_warnings = True
-            shared.log.warning(f'TAESD unsupported model type: {model_class}')
-        # return Image.new('RGB', (8, 8), color = (0, 0, 0))
-        return latents
-    vae = taesd_models.get(f'{model_class}-decoder', None)
-    if vae is None:
-        model_path = os.path.join(paths.models_path, "TAESD", f"tae{model_class}_decoder.pth")
-        download_model(model_path)
-        if os.path.exists(model_path):
-            taesd_models[f'{model_class}-decoder'] = TAESD(decoder_path=model_path, encoder_path=None)
-            shared.log.debug(f'VAE load: type=taesd model="{model_path}"')
-            vae = taesd_models[f'{model_class}-decoder']
-            vae.decoder.to(devices.device, dtype)
-        else:
-            shared.log.error(f'VAE load: type=taesd model="{model_path}" not found')
+    with lock:
+        from modules import shared
+        vae = get_model(model_type='decoder')
+        if vae is None or max(latents.shape) > 256: # safetey check of large tensors
             return latents
-    if vae is None:
-        return latents
-    try:
-        size = max(latents.shape[-1], latents.shape[-2])
-        if size > 256:
-            return latents
-        with devices.inference_context():
-            latents = latents.detach().clone().to(devices.device, dtype)
-            if len(latents.shape) == 3:
-                latents = latents.unsqueeze(0)
-                image = vae.decoder(latents).clamp(0, 1).detach()
-                image = 2.0 * image - 1.0 # typical normalized range except for preview which runs denormalization
-                return image[0]
-            elif len(latents.shape) == 4:
-                image = vae.decoder(latents).clamp(0, 1).detach()
-                image = 2.0 * image - 1.0 # typical normalized range except for preview which runs denormalization
-                return image
-            else:
-                if not previous_warnings:
-                    shared.log.error(f'TAESD decode unsupported latent type: {latents.shape}')
-                    previous_warnings = True
-                return latents
-    except Exception as e:
-        if not previous_warnings:
-            shared.log.error(f'VAE decode taesd: {e}')
-            previous_warnings = True
-        return latents
+        try:
+            with devices.inference_context():
+                tensor = latents.unsqueeze(0) if len(latents.shape) == 3 else latents
+                tensor = tensor.half().detach().clone().to(devices.device, dtype=vae.dtype)
+                if shared.opts.taesd_variant.startswith('TAESD'):
+                    image = vae.decoder(tensor).clamp(0, 1).detach()
+                    return image[0]
+                else:
+                    image = vae.decode(tensor, return_dict=False)[0]
+                    image = (image / 2.0 + 0.5).clamp(0, 1).detach()
+                    return image
+        except Exception as e:
+            return warn_once(f'decode {e}')
 
 
 def encode(image):
-    global previous_warnings # pylint: disable=global-statement
-    from modules import shared
-    model_class = shared.sd_model_type
-    if model_class == 'ldm':
-        model_class = 'sd'
-    if 'sd' not in model_class and 'f1' not in model_class:
-        if not previous_warnings:
-            previous_warnings = True
-            shared.log.warning(f'TAESD unsupported model type: {model_class}')
-        return Image.new('RGB', (8, 8), color = (0, 0, 0))
-    vae = taesd_models[f'{model_class}-encoder']
-    if vae is None:
-        model_path = os.path.join(paths.models_path, "TAESD", f"tae{model_class}_encoder.pth")
-        download_model(model_path)
-        if os.path.exists(model_path):
-            shared.log.debug(f'VAE load: type=taesd model="{model_path}"')
-            taesd_models[f'{model_class}-encoder'] = TAESD(encoder_path=model_path, decoder_path=None)
-            vae = taesd_models[f'{model_class}-encoder']
-            vae.encoder.to(devices.device, devices.dtype_vae)
-    # image = vae.scale_latents(image)
-    latents = vae.encoder(image)
-    return latents.detach()
+    with lock:
+        vae = get_model(model_type='encoder')
+        if vae is None:
+            return image
+        try:
+            with devices.inference_context():
+                latents = vae.encoder(image)
+            return latents.detach()
+        except Exception as e:
+            return warn_once(f'encode {e}')

modules/shared.py

@@ -794,8 +794,10 @@ options_templates.update(options_section(('live-preview', "Live Previews"), {
     "show_progress_every_n_steps": OptionInfo(1, "Live preview display period", gr.Slider, {"minimum": 0, "maximum": 32, "step": 1}),
     "show_progress_type": OptionInfo("Approximate", "Live preview method", gr.Radio, {"choices": ["Simple", "Approximate", "TAESD", "Full VAE"]}),
     "live_preview_refresh_period": OptionInfo(500, "Progress update period", gr.Slider, {"minimum": 0, "maximum": 5000, "step": 25}),
-    "live_preview_taesd_layers": OptionInfo(3, "TAESD decode layers", gr.Slider, {"minimum": 1, "maximum": 3, "step": 1}),
+    "taesd_variant": OptionInfo(shared_items.sd_taesd_items()[0], "TAESD variant", gr.Dropdown, {"choices": shared_items.sd_taesd_items()}),
+    "taesd_layers": OptionInfo(3, "TAESD decode layers", gr.Slider, {"minimum": 1, "maximum": 3, "step": 1}),
     "live_preview_downscale": OptionInfo(True, "Downscale high resolution live previews"),
+
     "logmonitor_show": OptionInfo(True, "Show log view"),
     "logmonitor_refresh_period": OptionInfo(5000, "Log view update period", gr.Slider, {"minimum": 0, "maximum": 30000, "step": 25}),
     "notification_audio_enable": OptionInfo(False, "Play a notification upon completion"),

modules/shared_items.py

@@ -8,6 +8,10 @@ def sd_vae_items():
     return ["Automatic", "None"] + list(modules.sd_vae.vae_dict)
 
 
+def sd_taesd_items():
+    import modules.sd_vae_taesd
+    return list(modules.sd_vae_taesd.TAESD_MODELS.keys()) + list(modules.sd_vae_taesd.CQYAN_MODELS.keys())
+
 def refresh_vae_list():
     import modules.sd_vae
     modules.sd_vae.refresh_vae_list()

modules/taesd/hybrid_small.py

@@ -0,0 +1,506 @@
+# pylint: disable=no-member,unused-argument,attribute-defined-outside-init
+
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Dict, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders.single_file_model import FromOriginalModelMixin
+from diffusers.utils.accelerate_utils import apply_forward_hook
+from diffusers.models.attention_processor import (
+    ADDED_KV_ATTENTION_PROCESSORS,
+    CROSS_ATTENTION_PROCESSORS,
+    Attention,
+    AttentionProcessor,
+    AttnAddedKVProcessor,
+    AttnProcessor,
+)
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.autoencoders.vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder
+
+
+class AutoencoderSmall(ModelMixin, ConfigMixin, FromOriginalModelMixin):
+    r"""
+    A VAE model with KL loss for encoding images into latents and decoding latent representations into images.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+
+    Parameters:
+        in_channels (int, *optional*, defaults to 3): Number of channels in the input image.
+        out_channels (int,  *optional*, defaults to 3): Number of channels in the output.
+        down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
+            Tuple of downsample block types.
+        up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
+            Tuple of upsample block types.
+        block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`):
+            Tuple of block output channels.
+        act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
+        latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space.
+        sample_size (`int`, *optional*, defaults to `32`): Sample input size.
+        scaling_factor (`float`, *optional*, defaults to 0.18215):
+            The component-wise standard deviation of the trained latent space computed using the first batch of the
+            training set. This is used to scale the latent space to have unit variance when training the diffusion
+            model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the
+            diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 1
+            / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image
+            Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) paper.
+        force_upcast (`bool`, *optional*, default to `True`):
+            If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE
+            can be fine-tuned / trained to a lower range without loosing too much precision in which case
+            `force_upcast` can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        down_block_types: Tuple[str] = ("DownEncoderBlock2D",),
+        up_block_types: Tuple[str] = ("UpDecoderBlock2D",),
+        block_out_channels: Tuple[int] = (64,),
+        encoder_block_out_channels: Tuple[int] = None,
+        decoder_block_out_channels: Tuple[int] = None,
+        layers_per_block: int = 1,
+        act_fn: str = "silu",
+        latent_channels: int = 4,
+        norm_num_groups: int = 32,
+        sample_size: int = 32,
+        scaling_factor: float = 0.18215,
+        latents_mean: Optional[Tuple[float]] = None,
+        latents_std: Optional[Tuple[float]] = None,
+        force_upcast: float = True,
+    ):
+        super().__init__()
+
+        if encoder_block_out_channels is not None or decoder_block_out_channels is not None:
+            if encoder_block_out_channels is None:
+                raise NotImplementedError
+            if decoder_block_out_channels is None:
+                raise NotImplementedError
+
+        else:
+            encoder_block_out_channels = block_out_channels
+            decoder_block_out_channels = block_out_channels
+            self.config.encoder_block_out_channels = self.config.decoder_block_out_channels = block_out_channels
+
+
+        # pass init params to Encoder
+        self.encoder = Encoder(
+            in_channels=in_channels,
+            out_channels=latent_channels,
+            down_block_types=down_block_types,
+            block_out_channels=encoder_block_out_channels,
+            layers_per_block=layers_per_block,
+            act_fn=act_fn,
+            norm_num_groups=norm_num_groups,
+            double_z=True,
+        )
+
+        # pass init params to Decoder
+        self.decoder = Decoder(
+            in_channels=latent_channels,
+            out_channels=out_channels,
+            up_block_types=up_block_types,
+            block_out_channels=decoder_block_out_channels,
+            layers_per_block=layers_per_block,
+            norm_num_groups=norm_num_groups,
+            act_fn=act_fn,
+        )
+
+        self.quant_conv = nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1)
+        self.post_quant_conv = nn.Conv2d(latent_channels, latent_channels, 1)
+
+        self.use_slicing = False
+        self.use_tiling = False
+
+        # only relevant if vae tiling is enabled
+        self.tile_sample_min_size = self.config.sample_size
+        sample_size = (
+            self.config.sample_size[0]
+            if isinstance(self.config.sample_size, (list, tuple))
+            else self.config.sample_size
+        )
+        self.tile_latent_min_size = int(sample_size / (2 ** (len(self.config.encoder_block_out_channels) - 1)))
+        self.tile_overlap_factor = 0.25
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, (Encoder, Decoder)):
+            module.gradient_checkpointing = value
+
+    def enable_tiling(self, use_tiling: bool = True):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+        """
+        self.use_tiling = use_tiling
+
+    def disable_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing
+        decoding in one step.
+        """
+        self.enable_tiling(False)
+
+    def enable_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.use_slicing = True
+
+    def disable_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing
+        decoding in one step.
+        """
+        self.use_slicing = False
+
+    @property
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True)
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
+    def set_default_attn_processor(self):
+        """
+        Disables custom attention processors and sets the default attention implementation.
+        """
+        if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+            processor = AttnAddedKVProcessor()
+        elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+            processor = AttnProcessor()
+        else:
+            raise ValueError(
+                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
+            )
+
+        self.set_attn_processor(processor)
+
+    @apply_forward_hook
+    def encode(
+        self, x: torch.FloatTensor, return_dict: bool = True
+    ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+        """
+        Encode a batch of images into latents.
+
+        Args:
+            x (`torch.FloatTensor`): Input batch of images.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+                The latent representations of the encoded images. If `return_dict` is True, a
+                [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
+        """
+        if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
+            return self.tiled_encode(x, return_dict=return_dict)
+
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self.encoder(x_slice) for x_slice in x.split(1)]
+            h = torch.cat(encoded_slices)
+        else:
+            h = self.encoder(x)
+
+        moments = self.quant_conv(h)
+        posterior = DiagonalGaussianDistribution(moments)
+
+        if not return_dict:
+            return (posterior,)
+
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def _decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]:
+        if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
+            return self.tiled_decode(z, return_dict=return_dict)
+
+        z = self.post_quant_conv(z)
+        dec = self.decoder(z)
+
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+    @apply_forward_hook
+    def decode(
+        self, z: torch.FloatTensor, return_dict: bool = True, generator=None
+    ) -> Union[DecoderOutput, torch.FloatTensor]:
+        """
+        Decode a batch of images.
+
+        Args:
+            z (`torch.FloatTensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+
+        """
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z).sample
+
+        if not return_dict:
+            return (decoded,)
+
+        return DecoderOutput(sample=decoded)
+
+    def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[2], b.shape[2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
+        return b
+
+    def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[3], b.shape[3], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
+        return b
+
+    def tiled_encode(self, x: torch.FloatTensor, return_dict: bool = True) -> AutoencoderKLOutput:
+        r"""Encode a batch of images using a tiled encoder.
+
+        When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
+        steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is
+        different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
+        tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
+        output, but they should be much less noticeable.
+
+        Args:
+            x (`torch.FloatTensor`): Input batch of images.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`:
+                If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain
+                `tuple` is returned.
+        """
+        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_latent_min_size - blend_extent
+
+        # Split the image into 512x512 tiles and encode them separately.
+        rows = []
+        for i in range(0, x.shape[2], overlap_size):
+            row = []
+            for j in range(0, x.shape[3], overlap_size):
+                tile = x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
+                tile = self.encoder(tile)
+                tile = self.quant_conv(tile)
+                row.append(tile)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=3))
+
+        moments = torch.cat(result_rows, dim=2)
+        posterior = DiagonalGaussianDistribution(moments)
+
+        if not return_dict:
+            return (posterior,)
+
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def tiled_decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]:
+        r"""
+        Decode a batch of images using a tiled decoder.
+
+        Args:
+            z (`torch.FloatTensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+        """
+        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_sample_min_size - blend_extent
+
+        # Split z into overlapping 64x64 tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, z.shape[2], overlap_size):
+            row = []
+            for j in range(0, z.shape[3], overlap_size):
+                tile = z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
+                tile = self.post_quant_conv(tile)
+                decoded = self.decoder(tile)
+                row.append(decoded)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=3))
+
+        dec = torch.cat(result_rows, dim=2)
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        sample_posterior: bool = False,
+        return_dict: bool = True,
+        generator: Optional[torch.Generator] = None,
+    ) -> Union[DecoderOutput, torch.FloatTensor]:
+        r"""
+        Args:
+            sample (`torch.FloatTensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+        """
+        x = sample
+        posterior = self.encode(x).latent_dist
+        if sample_posterior:
+            z = posterior.sample(generator=generator)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z).sample
+
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections
+    def fuse_qkv_projections(self):
+        """
+        Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query,
+        key, value) are fused. For cross-attention modules, key and value projection matrices are fused.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+        """
+        self.original_attn_processors = None
+
+        for _, attn_processor in self.attn_processors.items():
+            if "Added" in str(attn_processor.__class__.__name__):
+                raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
+
+        self.original_attn_processors = self.attn_processors
+
+        for module in self.modules():
+            if isinstance(module, Attention):
+                module.fuse_projections(fuse=True)
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections
+    def unfuse_qkv_projections(self):
+        """Disables the fused QKV projection if enabled.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+
+        """
+        if self.original_attn_processors is not None:
+            self.set_attn_processor(self.original_attn_processors)

modules/taesd/taesd.py

@@ -0,0 +1,88 @@
+import torch
+import torch.nn as nn
+from modules import devices
+
+
+def conv(n_in, n_out, **kwargs):
+    return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs)
+
+class Clamp(nn.Module):
+    def forward(self, x):
+        return torch.tanh(x / 3) * 3
+
+class Block(nn.Module):
+    def __init__(self, n_in, n_out):
+        super().__init__()
+        self.conv = nn.Sequential(conv(n_in, n_out), nn.ReLU(), conv(n_out, n_out), nn.ReLU(), conv(n_out, n_out))
+        self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity()
+        self.fuse = nn.ReLU()
+    def forward(self, x):
+        return self.fuse(self.conv(x) + self.skip(x))
+
+def Encoder(latent_channels=4):
+    return nn.Sequential(
+        conv(3, 64), Block(64, 64),
+        conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64),
+        conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64),
+        conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64),
+        conv(64, latent_channels),
+    )
+
+def Decoder(latent_channels=4):
+    from modules import shared
+    if shared.opts.taesd_layers == 1:
+        return nn.Sequential(
+            Clamp(), conv(latent_channels, 64), nn.ReLU(),
+            Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
+            Block(64, 64), Block(64, 64), Block(64, 64), nn.Identity(), conv(64, 64, bias=False),
+            Block(64, 64), Block(64, 64), Block(64, 64), nn.Identity(), conv(64, 64, bias=False),
+            Block(64, 64), conv(64, 3),
+        )
+    elif shared.opts.taesd_layers == 2:
+        return nn.Sequential(
+            Clamp(), conv(latent_channels, 64), nn.ReLU(),
+            Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
+            Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
+            Block(64, 64), Block(64, 64), Block(64, 64), nn.Identity(), conv(64, 64, bias=False),
+            Block(64, 64), conv(64, 3),
+        )
+    else:
+        return nn.Sequential(
+            Clamp(), conv(latent_channels, 64), nn.ReLU(),
+            Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
+            Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
+            Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
+            Block(64, 64), conv(64, 3),
+        )
+
+
+class TAESD(nn.Module): # pylint: disable=abstract-method
+    latent_magnitude = 3
+    latent_shift = 0.5
+
+    def __init__(self, encoder_path=None, decoder_path=None, latent_channels=None):
+        super().__init__()
+        self.dtype = devices.dtype_vae if devices.dtype_vae != torch.bfloat16 else torch.float16 # taesd does not support bf16
+        if latent_channels is None:
+            latent_channels = self.guess_latent_channels(str(decoder_path), str(encoder_path))
+        self.encoder = Encoder(latent_channels)
+        self.decoder = Decoder(latent_channels)
+        if encoder_path is not None:
+            self.encoder.load_state_dict(torch.load(encoder_path, map_location="cpu"), strict=False)
+            self.encoder.eval()
+            self.encoder = self.encoder.to(devices.device, dtype=self.dtype)
+        if decoder_path is not None:
+            self.decoder.load_state_dict(torch.load(decoder_path, map_location="cpu"), strict=False)
+            self.decoder.eval()
+            self.decoder = self.decoder.to(devices.device, dtype=self.dtype)
+
+    def guess_latent_channels(self, decoder_path, encoder_path):
+        return 16 if ("f1" in encoder_path or "f1" in decoder_path) or ("sd3" in encoder_path or "sd3" in decoder_path) else 4
+
+    @staticmethod
+    def scale_latents(x):
+        return x.div(2 * TAESD.latent_magnitude).add(TAESD.latent_shift).clamp(0, 1) # raw latents -> [0, 1]
+
+    @staticmethod
+    def unscale_latents(x):
+        return x.sub(TAESD.latent_shift).mul(2 * TAESD.latent_magnitude) # [0, 1] -> raw latents