diff --git a/backend/diffusion_engine/anima.py b/backend/diffusion_engine/anima.py index 9fe38b31..6ecefa50 100644 --- a/backend/diffusion_engine/anima.py +++ b/backend/diffusion_engine/anima.py @@ -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) diff --git a/backend/diffusion_engine/base.py b/backend/diffusion_engine/base.py index 3a85a9e2..00cbbe78 100644 --- a/backend/diffusion_engine/base.py +++ b/backend/diffusion_engine/base.py @@ -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() diff --git a/backend/diffusion_engine/chroma.py b/backend/diffusion_engine/chroma.py index b72f0702..99067eb6 100644 --- a/backend/diffusion_engine/chroma.py +++ b/backend/diffusion_engine/chroma.py @@ -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( diff --git a/backend/diffusion_engine/qwen.py b/backend/diffusion_engine/qwen.py index 3a0c330d..79389d63 100644 --- a/backend/diffusion_engine/qwen.py +++ b/backend/diffusion_engine/qwen.py @@ -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) diff --git a/backend/diffusion_engine/wan.py b/backend/diffusion_engine/wan.py index 94a3826d..7cecf6d1 100644 --- a/backend/diffusion_engine/wan.py +++ b/backend/diffusion_engine/wan.py @@ -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) diff --git a/backend/huggingface/CabalResearch/Mugen/vae/config.json b/backend/huggingface/CabalResearch/Mugen/vae/config.json index c3f38eb4..62551b9f 100644 --- a/backend/huggingface/CabalResearch/Mugen/vae/config.json +++ b/backend/huggingface/CabalResearch/Mugen/vae/config.json @@ -36,5 +36,6 @@ "UpDecoderBlock2D" ], "use_post_quant_conv": true, - "use_quant_conv": true + "use_quant_conv": true, + "mugen": true } diff --git a/backend/nn/_vae.py b/backend/nn/_vae.py new file mode 100644 index 00000000..90a37e1b --- /dev/null +++ b/backend/nn/_vae.py @@ -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) diff --git a/backend/nn/vae.py b/backend/nn/vae.py index aae21630..d8a7d6d8 100644 --- a/backend/nn/vae.py +++ b/backend/nn/vae.py @@ -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 diff --git a/backend/nn/wan_vae.py b/backend/nn/wan_vae.py index 3fd92e04..45f1b863 100644 --- a/backend/nn/wan_vae.py +++ b/backend/nn/wan_vae.py @@ -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) diff --git a/backend/patcher/vae.py b/backend/patcher/vae.py index a8732d6e..adbb83a7 100644 --- a/backend/patcher/vae.py +++ b/backend/patcher/vae.py @@ -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()