File size: 10,248 Bytes
bee1996 8a432f7 bee1996 5dcf26e bee1996 5dcf26e bee1996 5dcf26e bee1996 5dcf26e bee1996 5dcf26e bee1996 8a432f7 bee1996 8a432f7 bee1996 8a432f7 bee1996 8a432f7 bee1996 8a432f7 bee1996 8a432f7 bee1996 8a432f7 bee1996 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 | """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_<variant>.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()
|