CodeCow/stable-diffusion-webui-forge
1
0
Fork 0
mirror of https://github.com/lllyasviel/stable-diffusion-webui-forge.git synced 2026-07-21 21:01:24 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

taesd

Browse Source
This commit is contained in:
Haoming 2025-10-30 12:15:29 +08:00
parent 4f572b8e2f
commit 22d55f8496
5 changed files with 27 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
backend/diffusion_engine/base.py
View File

@ -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

2
backend/diffusion_engine/chroma.py
View File

@ -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)

2
backend/diffusion_engine/flux.py
View File

@ -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

2
backend/diffusion_engine/lumina.py
View File

@ -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)

50
modules/sd_vae_taesd.py
View File

@ -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