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