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

vae

Browse Source
This commit is contained in:
Haoming 2026-04-27 11:55:39 +08:00
parent 1a29b8151c
commit 32adaae7b7
10 changed files with 35 additions and 46 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
backend/diffusion_engine/anima.py
View File

@ -19,7 +19,6 @@ class Anima(ForgeDiffusionEngine):
clip = CLIP(model_dict={"qwen3_06b": huggingface_components["text_encoder"]}, tokenizer_dict={"qwen3_06b": huggingface_components["tokenizer"], "t5xxl": huggingface_components["tokenizer_2"]})
vae = VAE(model=huggingface_components["vae"], is_wan=True)
vae.first_stage_model.latent_format = self.model_config.latent_format
k_predictor = PredictionDiscreteFlow(estimated_config)

2
backend/diffusion_engine/base.py
View File

@ -25,6 +25,8 @@ class ForgeDiffusionEngine:
matched_guesses = []
def __init__(self, estimated_config, huggingface_components):
huggingface_components["vae"].latent_format = estimated_config.latent_format
self.model_config = estimated_config
self.is_inpaint = estimated_config.inpaint_model()

2
backend/diffusion_engine/chroma.py
View File

@ -19,7 +19,9 @@ class Chroma(ForgeDiffusionEngine):
clip = CLIP(model_dict={"t5xxl": huggingface_components["text_encoder"]}, tokenizer_dict={"t5xxl": huggingface_components["tokenizer"]})
vae = VAE(model=huggingface_components["vae"])
k_predictor = PredictionFlux(mu=1.0)
unet = UnetPatcher.from_model(model=huggingface_components["transformer"], diffusers_scheduler=None, k_predictor=k_predictor, config=estimated_config)
self.text_processing_engine_t5 = T5TextProcessingEngine(

1
backend/diffusion_engine/qwen.py
View File

@ -27,7 +27,6 @@ class QwenImage(ForgeDiffusionEngine):
clip = CLIP(model_dict={"qwen25_7b": huggingface_components["text_encoder"]}, tokenizer_dict={"qwen25_7b": huggingface_components["tokenizer"]})
vae = VAE(model=huggingface_components["vae"], is_wan=True)
vae.first_stage_model.latent_format = self.model_config.latent_format
k_predictor = PredictionDiscreteFlow(estimated_config)

1
backend/diffusion_engine/wan.py
View File

@ -24,7 +24,6 @@ class Wan(ForgeDiffusionEngine):
clip = CLIP(model_dict={"umt5xxl": huggingface_components["text_encoder"]}, tokenizer_dict={"umt5xxl": huggingface_components["tokenizer"]})
vae = VAE(model=huggingface_components["vae"], is_wan=True)
vae.first_stage_model.latent_format = self.model_config.latent_format
k_predictor = PredictionDiscreteFlow(estimated_config)

3
backend/huggingface/CabalResearch/Mugen/vae/config.json
View File

@ -36,5 +36,6 @@
"UpDecoderBlock2D"
],
"use_post_quant_conv": true,
"use_quant_conv": true
"use_quant_conv": true,
"mugen": true
}

13
backend/nn/_vae.py Normal file
View File

@ -0,0 +1,13 @@
import torch
from modules_forge.packages.huggingface_guess.latent import LatentFormat
class ProcessLatent:
latent_format: LatentFormat = None
def process_in(self, latent: torch.Tensor) -> torch.Tensor:
return self.latent_format.process_in(latent)
def process_out(self, latent: torch.Tensor) -> torch.Tensor:
return self.latent_format.process_out(latent)

35
backend/nn/vae.py
View File

@ -8,6 +8,7 @@ from einops import rearrange
from backend import memory_management
from backend.attention import attention_function_vae
from backend.nn._vae import ProcessLatent
def nonlinearity(x):
@ -283,12 +284,14 @@ class Decoder(nn.Module):
return h
class IntegratedAutoencoderKL(nn.Module, ConfigMixin):
class IntegratedAutoencoderKL(nn.Module, ProcessLatent, ConfigMixin):
config_name = "config.json"
@register_to_config
def __init__(self, in_channels=3, out_channels=3, down_block_types=("DownEncoderBlock2D",), up_block_types=("UpDecoderBlock2D",), block_out_channels=(64,), layers_per_block=1, act_fn="silu", latent_channels=4, norm_num_groups=32, sample_size=32, scaling_factor=0.18215, shift_factor=0.0, latents_mean=None, latents_std=None, force_upcast=True, use_quant_conv=True, use_post_quant_conv=True):
def __init__(self, in_channels=3, out_channels=3, block_out_channels=(64,), layers_per_block=1, latent_channels=4, use_quant_conv=True, use_post_quant_conv=True, **kwargs):
del kwargs
super().__init__()
ch = block_out_channels[0]
ch_mult = [x // ch for x in block_out_channels]
self.encoder = Encoder(double_z=True, z_channels=latent_channels, resolution=256, in_channels=in_channels, out_ch=out_channels, ch=ch, ch_mult=ch_mult, num_res_blocks=layers_per_block, attn_resolutions=[], dropout=0.0)
@ -296,11 +299,6 @@ class IntegratedAutoencoderKL(nn.Module, ConfigMixin):
self.quant_conv = nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1) if use_quant_conv else None
self.post_quant_conv = nn.Conv2d(latent_channels, latent_channels, 1) if use_post_quant_conv else None
self.embed_dim = latent_channels
self.scaling_factor = scaling_factor
self.shift_factor = shift_factor
if not isinstance(self.shift_factor, float):
self.shift_factor = 0.0
def encode(self, x):
z = self.encoder(x)
@ -318,19 +316,15 @@ class IntegratedAutoencoderKL(nn.Module, ConfigMixin):
x = self.decoder(z)
return x
def process_in(self, latent):
return (latent - self.shift_factor) * self.scaling_factor
def process_out(self, latent):
return (latent / self.scaling_factor) + self.shift_factor
class AutoencoderKLFlux2(IntegratedAutoencoderKL):
config_name = "config.json"
@register_to_config
def __init__(self, in_channels=3, out_channels=3, down_block_types=("DownEncoderBlock2D",), up_block_types=("UpDecoderBlock2D",), block_out_channels=(64,), layers_per_block=1, act_fn="silu", latent_channels=4, norm_num_groups=32, sample_size=32, scaling_factor=0.18215, shift_factor=0.0, latents_mean=None, latents_std=None, force_upcast=True, use_quant_conv=True, use_post_quant_conv=True, *, ech: int = None, dch: int = None):
def __init__(self, in_channels=3, out_channels=3, block_out_channels=(64,), layers_per_block=1, latent_channels=4, use_quant_conv=True, use_post_quant_conv=True, *, mugen: bool = False, ech: int = None, dch: int = None, **kwargs):
del kwargs
super().__init__()
ch = block_out_channels[0]
ch_mult = [x // ch for x in block_out_channels]
self.encoder = Encoder(double_z=True, z_channels=latent_channels, resolution=256, in_channels=in_channels, out_ch=out_channels, ch=ech or ch, ch_mult=ch_mult, num_res_blocks=layers_per_block, attn_resolutions=[], dropout=0.0)
@ -338,11 +332,6 @@ class AutoencoderKLFlux2(IntegratedAutoencoderKL):
self.quant_conv = nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1) if use_quant_conv else None
self.post_quant_conv = nn.Conv2d(latent_channels, latent_channels, 1) if use_post_quant_conv else None
self.embed_dim = latent_channels
self.scaling_factor = scaling_factor
self.shift_factor = shift_factor
if not isinstance(self.shift_factor, float):
self.shift_factor = 0.0
self.bn_eps = 1e-4
self.bn_momentum = 0.1
@ -356,7 +345,7 @@ class AutoencoderKLFlux2(IntegratedAutoencoderKL):
)
self.bn.eval()
self.mugen = False # 32 <-> 128
self.mugen = mugen # 32 <-> 128
def encode(self, x):
z = super().encode(x)
@ -393,12 +382,6 @@ class AutoencoderKLFlux2(IntegratedAutoencoderKL):
return super().decode(z)
def process_in(self, latent):
return latent
def process_out(self, latent):
return latent
def preprocess_decode(self, latent: torch.Tensor):
packed_channels: int = latent.size(1)
latent_channels: int = 128

11
backend/nn/wan_vae.py
View File

@ -9,6 +9,7 @@ from diffusers.configuration_utils import ConfigMixin, register_to_config
from einops import rearrange
from backend.attention import attention_function_vae
from backend.nn._vae import ProcessLatent
from backend.operations import ForgeOperations as ops
CACHE_T = 2
@ -363,7 +364,7 @@ def count_cache_layers(model):
return count
class WanVAE(nn.Module, ConfigMixin):
class WanVAE(nn.Module, ProcessLatent, ConfigMixin):
config_name = "config.json"
@register_to_config
@ -382,8 +383,6 @@ class WanVAE(nn.Module, ConfigMixin):
self.conv2 = CausalConv3d(z_dim, z_dim, 1)
self.decoder = Decoder3d(base_dim, z_dim, conv_out_channels, dim_mult, num_res_blocks, attn_scales, self.temporal_upsample, dropout)
self.latent_format = None
def encode(self, x):
conv_idx = [0]
t = x.shape[2]
@ -419,9 +418,3 @@ class WanVAE(nn.Module, ConfigMixin):
out_ = self.decoder(x[:, :, 1 + 2 * (i - 1) : 1 + 2 * i, :, :], feat_cache=feat_map, feat_idx=conv_idx)
out += out_
return torch.cat(out, 2)
def process_in(self, latent):
return self.latent_format.process_in(latent)
def process_out(self, latent):
return self.latent_format.process_out(latent)

12
backend/patcher/vae.py
View File

@ -131,14 +131,14 @@ class VAE:
self.downscale_ratio = 8
self.downscale_index_formula = None
self.latent_dim = 2
self.latent_channels = int(model.config.latent_channels) # 4 | 16
self.latent_channels = 32 if is_mugen else int(model.config.latent_channels) # 4 | 16
self.memory_used_encode = lambda shape, dtype: (1767 * shape[2] * shape[3]) * memory_management.dtype_size(dtype)
self.memory_used_decode = lambda shape, dtype: (2178 * shape[2] * shape[3] * 64) * memory_management.dtype_size(dtype)
if is_flux2 or is_mugen:
self.upscale_ratio = 8 if is_mugen else 16
self.downscale_ratio = 8 if is_mugen else 16
self.latent_channels = 32 if is_mugen else 128
if is_flux2:
self.upscale_ratio = 16
self.downscale_ratio = 16
self.latent_channels = 128
self.memory_used_decode = lambda shape, dtype: (2178 * shape[2] * shape[3] * 64) * memory_management.dtype_size(dtype) * 4.0
else:
@ -153,8 +153,6 @@ class VAE:
self.output_channels = 3
self.first_stage_model = model.eval()
if is_mugen:
self.first_stage_model.mugen = True
self.device = device or memory_management.vae_device()
offload_device = memory_management.vae_offload_device()