From 22d55f849689c962c83cc27d5d8b229d70ffebc8 Mon Sep 17 00:00:00 2001 From: Haoming Date: Thu, 30 Oct 2025 12:15:29 +0800 Subject: [PATCH] taesd --- backend/diffusion_engine/base.py | 4 +-- backend/diffusion_engine/chroma.py | 2 ++ backend/diffusion_engine/flux.py | 2 ++ backend/diffusion_engine/lumina.py | 2 ++ modules/sd_vae_taesd.py | 50 ++++++++++++------------------ 5 files changed, 27 insertions(+), 33 deletions(-) diff --git a/backend/diffusion_engine/base.py b/backend/diffusion_engine/base.py index af2b02d7..6bff1bc2 100644 --- a/backend/diffusion_engine/base.py +++ b/backend/diffusion_engine/base.py @@ -54,9 +54,9 @@ class ForgeDiffusionEngine: self.cond_stage_model = None self.use_distilled_cfg_scale = False self.is_sd1 = False - self.is_sd2 = False self.is_sdxl = False - self.is_wan = False # also affects the usage of WanVAE + self.is_flux = False # affects the usage of TAESD + self.is_wan = False # affects the usage of WanVAE (B, C, T, H, W) def save_unet(self, filename): import safetensors.torch as sf diff --git a/backend/diffusion_engine/chroma.py b/backend/diffusion_engine/chroma.py index 1c85f9c3..fc9d994c 100644 --- a/backend/diffusion_engine/chroma.py +++ b/backend/diffusion_engine/chroma.py @@ -34,6 +34,8 @@ class Chroma(ForgeDiffusionEngine): self.forge_objects_original = self.forge_objects.shallow_copy() self.forge_objects_after_applying_lora = self.forge_objects.shallow_copy() + self.is_flux = True + @torch.inference_mode() def get_learned_conditioning(self, prompt: list[str]): memory_management.load_model_gpu(self.forge_objects.clip.patcher) diff --git a/backend/diffusion_engine/flux.py b/backend/diffusion_engine/flux.py index 87e2576f..cd681bf4 100644 --- a/backend/diffusion_engine/flux.py +++ b/backend/diffusion_engine/flux.py @@ -59,6 +59,8 @@ class Flux(ForgeDiffusionEngine): self.forge_objects_original = self.forge_objects.shallow_copy() self.forge_objects_after_applying_lora = self.forge_objects.shallow_copy() + self.is_flux = True + def set_clip_skip(self, clip_skip): self.text_processing_engine_l.clip_skip = clip_skip diff --git a/backend/diffusion_engine/lumina.py b/backend/diffusion_engine/lumina.py index f4bcec07..68a587d3 100644 --- a/backend/diffusion_engine/lumina.py +++ b/backend/diffusion_engine/lumina.py @@ -34,6 +34,8 @@ class Lumina2(ForgeDiffusionEngine): self.forge_objects_original = self.forge_objects.shallow_copy() self.forge_objects_after_applying_lora = self.forge_objects.shallow_copy() + self.is_flux = True + @torch.inference_mode() def get_learned_conditioning(self, prompt: list[str]): memory_management.load_model_gpu(self.forge_objects.clip.patcher) diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py index 87605954..9afda1db 100644 --- a/modules/sd_vae_taesd.py +++ b/modules/sd_vae_taesd.py @@ -4,7 +4,9 @@ Tiny AutoEncoder for Stable Diffusion https://github.com/madebyollin/taesd """ + import os + import torch import torch.nn as nn @@ -35,23 +37,11 @@ class Block(nn.Module): def decoder(latent_channels=4): - 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), - ) + 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)) 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), - ) + 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)) class TAESDDecoder(nn.Module): @@ -59,7 +49,6 @@ class TAESDDecoder(nn.Module): latent_shift = 0.5 def __init__(self, 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: @@ -69,8 +58,7 @@ class TAESDDecoder(nn.Module): latent_channels = 4 self.decoder = decoder(latent_channels) - self.decoder.load_state_dict( - torch.load(decoder_path, map_location='cpu' if devices.device.type != 'cuda' else None)) + self.decoder.load_state_dict(torch.load(decoder_path, map_location="cpu" if devices.device.type != "cuda" else None)) class TAESDEncoder(nn.Module): @@ -78,7 +66,6 @@ class TAESDEncoder(nn.Module): latent_shift = 0.5 def __init__(self, encoder_path="taesd_encoder.pth", latent_channels=None): - """Initialize pretrained TAESD on the given device from the given checkpoints.""" super().__init__() if latent_channels is None: @@ -88,33 +75,32 @@ class TAESDEncoder(nn.Module): latent_channels = 4 self.encoder = encoder(latent_channels) - self.encoder.load_state_dict( - torch.load(encoder_path, map_location='cpu' if devices.device.type != 'cuda' else None)) + self.encoder.load_state_dict(torch.load(encoder_path, map_location="cpu" if devices.device.type != "cuda" else None)) def download_model(model_path, model_url): if not os.path.exists(model_path): os.makedirs(os.path.dirname(model_path), exist_ok=True) - print(f'Downloading TAESD model to: {model_path}') + print(f"Downloading TAESD model to: {model_path}") torch.hub.download_url_to_file(model_url, model_path) def decoder_model(): - if not shared.sd_model.is_webui_legacy_model(): - if shared.sd_model.is_wan: - return None + if shared.sd_model.is_flux: model_name = "taef1_decoder.pth" elif shared.sd_model.is_sdxl: model_name = "taesdxl_decoder.pth" - else: + elif shared.sd_model.is_sd1: model_name = "taesd_decoder.pth" + else: + return None loaded_model = sd_vae_taesd_models.get(model_name) if loaded_model is None: model_path = os.path.join(paths_internal.models_path, "VAE-taesd", model_name) - download_model(model_path, 'https://github.com/madebyollin/taesd/raw/main/' + model_name) + download_model(model_path, "https://github.com/madebyollin/taesd/raw/main/" + model_name) if os.path.exists(model_path): loaded_model = TAESDDecoder(model_path) @@ -122,24 +108,26 @@ def decoder_model(): loaded_model.to(devices.device, devices.dtype) sd_vae_taesd_models[model_name] = loaded_model else: - raise FileNotFoundError('TAESD model not found') + raise FileNotFoundError("TAESD model not found...") return loaded_model.decoder def encoder_model(): - if not shared.sd_model.is_webui_legacy_model(): + if shared.sd_model.is_flux: model_name = "taef1_encoder.pth" elif shared.sd_model.is_sdxl: model_name = "taesdxl_encoder.pth" - else: + elif shared.sd_model.is_sd1: model_name = "taesd_encoder.pth" + else: + return None loaded_model = sd_vae_taesd_models.get(model_name) if loaded_model is None: model_path = os.path.join(paths_internal.models_path, "VAE-taesd", model_name) - download_model(model_path, 'https://github.com/madebyollin/taesd/raw/main/' + model_name) + download_model(model_path, "https://github.com/madebyollin/taesd/raw/main/" + model_name) if os.path.exists(model_path): loaded_model = TAESDEncoder(model_path) @@ -147,6 +135,6 @@ def encoder_model(): loaded_model.to(devices.device, devices.dtype) sd_vae_taesd_models[model_name] = loaded_model else: - raise FileNotFoundError('TAESD model not found') + raise FileNotFoundError("TAESD model not found...") return loaded_model.encoder