"""Primal VAE for diffusers (AutoModel remote code). Deterministic autoencoder: RGB in [-1, 1] -> latent_channels at `compression`x spatial compression, latent ~ N(0, 1) by training, no scaling or shift factor. `patchify(z, compression=8)` regroups the latent into the token grid a DiT consumes at that total compression (1ch @ 2x -> 16ch @ 8x, 2ch @ 2x -> 32ch @ 8x). Loads the ai-toolkit training checkpoints directly: `from_pretrained(repo, variant="2x1")` reads `primal_vae_2x1.safetensors` and derives the architecture from the tensor shapes; `weights_name=` picks any other file. Standalone: torch + diffusers + safetensors. """ import math import os from typing import List, Optional import torch import torch.nn as nn import torch.nn.functional as F import torch.utils.checkpoint from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.models.modeling_utils import ModelMixin from huggingface_hub import hf_hub_download from safetensors import safe_open # dilation of block i's first conv; widens the receptive field without downsampling DILATION_CYCLE = (1, 2, 4) class ChannelRMSNorm(nn.Module): """RMSNorm over channels at each position: no spatial statistics, so tiles encode/decode like the whole image.""" def __init__(self, ch: int, eps: float = 1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(ch)) def forward(self, x: torch.Tensor) -> torch.Tensor: return F.rms_norm(x.movedim(1, -1), (x.shape[1],), self.weight.to(x.dtype), self.eps).movedim(-1, 1) class ResBlock(nn.Module): def __init__(self, width: int, dilation: int): super().__init__() self.norm1 = ChannelRMSNorm(width) self.conv1 = nn.Conv2d(width, width, 3, padding=dilation, dilation=dilation) self.norm2 = ChannelRMSNorm(width) self.conv2 = nn.Conv2d(width, width, 3, padding=1) def forward(self, x: torch.Tensor) -> torch.Tensor: h = self.conv1(F.silu(self.norm1(x))) h = self.conv2(F.silu(self.norm2(h))) return x + h class UStack(nn.Module): """conv_in -> [blocks, pixel-unshuffle down] per level -> bottom blocks -> [up, + skip, blocks] per level -> norm/act -> conv_out.""" def __init__(self, in_ch: int, widths: List[int], depths: List[int], up_depths: List[int], out_ch: int): super().__init__() L = len(widths) if len(depths) != L or len(up_depths) != L - 1: raise ValueError("level_depths must match level_widths; up_depths has one entry less") self.conv_in = nn.Conv2d(in_ch, widths[0], 3, padding=1) self.down_blocks = nn.ModuleList([nn.ModuleList([ResBlock(widths[i], DILATION_CYCLE[j % len(DILATION_CYCLE)]) for j in range(depths[i])]) for i in range(L)]) self.downs = nn.ModuleList([nn.Conv2d(4 * widths[i], widths[i + 1], 3, padding=1) for i in range(L - 1)]) self.ups = nn.ModuleList([nn.Conv2d(widths[i + 1], 4 * widths[i], 3, padding=1) for i in range(L - 1)]) self.up_blocks = nn.ModuleList([nn.ModuleList([ResBlock(widths[i], DILATION_CYCLE[j % len(DILATION_CYCLE)]) for j in range(up_depths[i])]) for i in range(L - 1)]) self.norm_out = ChannelRMSNorm(widths[0]) self.conv_out = nn.Conv2d(widths[0], out_ch, 3, padding=1) # set by ModelMixin.enable_gradient_checkpointing(); per ResBlock, only while grad is enabled self.gradient_checkpointing = False self._gradient_checkpointing_func = lambda blk, x: torch.utils.checkpoint.checkpoint(blk, x, use_reentrant=False) def _block(self, blk: ResBlock, x: torch.Tensor) -> torch.Tensor: if self.gradient_checkpointing and torch.is_grad_enabled(): return self._gradient_checkpointing_func(blk, x) return blk(x) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.conv_in(x) skips = [] L = len(self.down_blocks) for i in range(L): for blk in self.down_blocks[i]: x = self._block(blk, x) if i < L - 1: skips.append(x) x = self.downs[i](F.pixel_unshuffle(x, 2)) for i in reversed(range(L - 1)): x = F.pixel_shuffle(self.ups[i](x), 2) + skips[i] for blk in self.up_blocks[i]: x = self._block(blk, x) return self.conv_out(F.silu(self.norm_out(x))) class PrimalVAE(ModelMixin, ConfigMixin): """Primal VAE: `encode` -> [B, latent_channels, H/compression, W/compression] (~N(0,1)), `decode` back to [B, 3, H, W] in [-1, 1]. No posterior, no scaling/shift factor.""" _supports_gradient_checkpointing = True @register_to_config def __init__( self, compression: int = 2, latent_channels: int = 1, level_widths: List[int] = (96, 192, 384, 384), level_depths: List[int] = (1, 1, 2, 4), up_depths: Optional[List[int]] = (1, 1, 2), ): super().__init__() widths, depths = [int(v) for v in level_widths], [int(v) for v in level_depths] ups = [int(v) for v in up_depths] if up_depths is not None else depths[:-1] self.compression = int(compression) self.latent_channels = int(latent_channels) # the latent must tile into the U's 2^(levels-1) and a 4x4 token patch self.latent_multiple = max(4, 2 ** (len(widths) - 1)) pix = 3 * self.compression ** 2 self.encoder = UStack(pix, widths, depths, ups, self.latent_channels) self.decoder = UStack(self.latent_channels, widths, depths, ups, pix) @property def image_multiple(self) -> int: """Image sides must be a multiple of this (16 for the released models).""" return self.compression * self.latent_multiple def token_channels(self, compression: int = 8) -> int: """Channels of the latent patchified to `compression`: 16 for 2x1 at 8x (Flux.1), 32 for 2x2 (Flux.2).""" return self.latent_channels * self._patch(compression) ** 2 def _patch(self, compression: int) -> int: if compression % self.compression: raise ValueError(f"compression {compression} is not a multiple of the model's {self.compression}x") return compression // self.compression def _check_image(self, x: torch.Tensor): H, W = x.shape[-2:] m = self.image_multiple if H % m or W % m: raise ValueError(f"image sides must be multiples of {m}, got {H}x{W}") def encode(self, x: torch.Tensor) -> torch.Tensor: """x: [B, 3, H, W] in [-1, 1] -> z: [B, latent_channels, H/compression, W/compression].""" self._check_image(x) return self.encoder(F.pixel_unshuffle(x.to(self.dtype), self.compression)) def decode(self, z: torch.Tensor) -> torch.Tensor: """z -> x_hat: [B, 3, H, W], nominally in [-1, 1] (not clamped).""" return F.pixel_shuffle(self.decoder(z.to(self.dtype)), self.compression) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.decode(self.encode(x)) def patchify(self, z: torch.Tensor, compression: int = 8) -> torch.Tensor: """Latent -> token grid at `compression`x total: [B, C, h, w] -> [B, C * p^2, h/p, w/p] with p = compression / model compression (space-to-depth). 8 gives the Flux 8x layout from a 2x latent.""" return F.pixel_unshuffle(z, self._patch(compression)) def unpatchify(self, tokens: torch.Tensor) -> torch.Tensor: """Inverse of `patchify`; the patch size follows from the channel count (C = latent_channels * p^2).""" p = math.isqrt(tokens.shape[1] // self.latent_channels) if self.latent_channels * p * p != tokens.shape[1]: raise ValueError(f"{tokens.shape[1]} channels is not latent_channels ({self.latent_channels}) times a square") return F.pixel_shuffle(tokens, p) # ---- loading ----------------------------------------------------------- HUB_KWARGS = ("cache_dir", "force_download", "proxies", "local_files_only", "token", "revision", "subfolder") @staticmethod def config_from_shapes(shapes: dict) -> dict: """Architecture from the tensor shapes alone; the file's header metadata is not consulted.""" widths = [shapes["encoder.conv_in.weight"][0]] while f"encoder.downs.{len(widths) - 1}.weight" in shapes: widths.append(shapes[f"encoder.downs.{len(widths) - 1}.weight"][0]) count = lambda pre, i: len({k.split(".")[3] for k in shapes if k.startswith(f"{pre}.{i}.")}) return { "compression": math.isqrt(shapes["encoder.conv_in.weight"][1] // 3), "latent_channels": shapes["encoder.conv_out.weight"][0], "level_widths": widths, "level_depths": [count("encoder.down_blocks", i) for i in range(len(widths))], "up_depths": [count("encoder.up_blocks", i) for i in range(len(widths) - 1)], } @classmethod def from_pretrained(cls, pretrained_model_name_or_path, variant: str = "2x1", weights_name: Optional[str] = None, torch_dtype=None, **kwargs): """`variant` picks `primal_vae_.safetensors` in the repo ("2x1" or "2x2"); `weights_name` any file.""" name = weights_name or f"primal_vae_{variant}.safetensors" hub = {k: v for k, v in kwargs.items() if k in cls.HUB_KWARGS and v is not None} path = str(pretrained_model_name_or_path) if os.path.isdir(path): file = os.path.join(path, hub.get("subfolder") or "", name) else: file = hf_hub_download(path, name, **hub) with safe_open(file, framework="pt", device="cpu") as f: keys = [k for k in f.keys() if k != "arch_spec"] # arch_spec: ai-toolkit's own derived buffer cfg = cls.config_from_shapes({k: f.get_slice(k).get_shape() for k in keys}) with torch.device("meta"): model = cls(**cfg) state = {k: f.get_tensor(k) for k in keys} if torch_dtype is not None: state = {k: v.to(torch_dtype) for k, v in state.items()} model.load_state_dict(state, assign=True) model.register_to_config(_name_or_path=path) return model.eval()