diff --git a/CHANGELOG.md b/CHANGELOG.md
index c57a51ad4..198500b9a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,6 @@
# Change Log for SD.Next
-## Update for 2025-01-11
+## Update for 2025-01-12
- [Allegro Video](https://huggingface.co/rhymes-ai/Allegro)
- optimizations: full offload and quantization support
@@ -32,6 +32,10 @@
- move steps, strength, prompt, negative from settings into ui params
- set/restore detailer metadata
- new [detailer wiki](https://github.com/vladmandic/automatic/wiki/Detailer)
+- **Preview**
+ - since different TAESD versions produce different results and latest is not necessarily greatest
+ you can choose TAESD version in settings -> live preview
+ also added is support for another finetuned version of TAESD [Hybrid TinyVAE](https://huggingface.co/cqyan/hybrid-sd-tinyvae-xl)
- **Other**
- **XYZ Grid**: add prompt search&replace options: *primary, refine, detailer, all*
- **SysInfo**: update to collected data and benchmarks
@@ -46,6 +50,7 @@
- sd35 img2img
- samplers test for scale noise before using
- scheduler api
+ - sampler create error handling
- controlnet with hires
- controlnet with batch count
- apply settings skip hidden settings
diff --git a/javascript/logger.js b/javascript/logger.js
index 8fa812b86..08baf1165 100644
--- a/javascript/logger.js
+++ b/javascript/logger.js
@@ -1,4 +1,4 @@
-const timeout = 10000;
+const timeout = 30000;
const log = async (...msg) => {
const dt = new Date();
diff --git a/modules/loader.py b/modules/loader.py
index 63c52d18c..c48afa7a9 100644
--- a/modules/loader.py
+++ b/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")
diff --git a/modules/model_flux.py b/modules/model_flux.py
index f8d112953..8fe147223 100644
--- a/modules/model_flux.py
+++ b/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')
diff --git a/modules/processing_args.py b/modules/processing_args.py
index 9ccd8f9dd..4e51d6d4f 100644
--- a/modules/processing_args.py
+++ b/modules/processing_args.py
@@ -1,7 +1,6 @@
import typing
import os
import re
-import copy
import math
import time
import inspect
diff --git a/modules/processing_vae.py b/modules/processing_vae.py
index 04af9bab1..faaacb21e 100644
--- a/modules/processing_vae.py
+++ b/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'):
diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py
index b01b847d8..62bc4f55c 100644
--- a/modules/sd_samplers_common.py
+++ b/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))
diff --git a/modules/sd_samplers_diffusers.py b/modules/sd_samplers_diffusers.py
index 9f88bf103..6b2de72aa 100644
--- a/modules/sd_samplers_diffusers.py
+++ b/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
diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py
index 4507ee3c8..c8a1b882f 100644
--- a/modules/sd_vae_taesd.py
+++ b/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}')
diff --git a/modules/shared.py b/modules/shared.py
index d27d2326b..f1a9483c2 100644
--- a/modules/shared.py
+++ b/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"),
diff --git a/modules/shared_items.py b/modules/shared_items.py
index 17b7ce1ee..5c1e3aebb 100644
--- a/modules/shared_items.py
+++ b/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()
diff --git a/modules/taesd/hybrid_small.py b/modules/taesd/hybrid_small.py
new file mode 100644
index 000000000..a59b0b4d7
--- /dev/null
+++ b/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.
+
+
+
+ This API is 🧪 experimental.
+
+
+ """
+ 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.
+
+
+
+ This API is 🧪 experimental.
+
+
+
+ """
+ if self.original_attn_processors is not None:
+ self.set_attn_processor(self.original_attn_processors)
diff --git a/modules/taesd/taesd.py b/modules/taesd/taesd.py
new file mode 100644
index 000000000..8e391a8fb
--- /dev/null
+++ b/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