primal_vae / modeling_primal_vae.py

Add gradient checkpointing to the vae

5dcf26e verified about 22 hours ago

10.2 kB

"""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()

	"""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()