| """Inference wrapper for the Semantic VAE. |
| |
| Images are BCHW tensors in [-1, 1]. ``encode`` and ``decode`` use normalized |
| latents; their ``_raw`` variants use the underlying VAE representation. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import copy |
| from pathlib import Path |
| from typing import Mapping |
|
|
| import timm |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from safetensors import safe_open |
| from safetensors.torch import load_file |
|
|
|
|
| DINO_MODEL_NAME = "vit_base_patch14_dinov2.lvd142m" |
| DINO_MEAN = (0.485, 0.456, 0.406) |
| DINO_STD = (0.229, 0.224, 0.225) |
| DINO_PATCH_SIZE = 14 |
| LATENT_DOWNSAMPLE_FACTOR = 16 |
| SAFETENSORS_FORMAT = "semantic_vae_full_v1" |
|
|
|
|
| def _dino_spatial_size(image_size: tuple[int, int]) -> tuple[int, int]: |
| height, width = image_size |
| half_stride = LATENT_DOWNSAMPLE_FACTOR // 2 |
| patches = ( |
| max(1, (size + half_stride) // LATENT_DOWNSAMPLE_FACTOR) |
| for size in (height, width) |
| ) |
| return tuple(size * DINO_PATCH_SIZE for size in patches) |
|
|
|
|
| def _group_norm(channels: int) -> nn.GroupNorm: |
| return nn.GroupNorm(32, channels, eps=1e-6, affine=True) |
|
|
|
|
| class ResnetBlock(nn.Module): |
| def __init__(self, in_ch: int, out_ch: int) -> None: |
| super().__init__() |
| self.norm1 = _group_norm(in_ch) |
| self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1) |
| self.norm2 = _group_norm(out_ch) |
| self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1) |
| self.shortcut = ( |
| nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity() |
| ) |
|
|
| 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 self.shortcut(x) + h |
|
|
|
|
| class AttnBlock(nn.Module): |
| def __init__(self, channels: int) -> None: |
| super().__init__() |
| self.norm = _group_norm(channels) |
| self.qkv = nn.Conv2d(channels, channels * 3, 1) |
| self.proj_out = nn.Conv2d(channels, channels, 1) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| batch, channels, height, width = x.shape |
| q, k, v = ( |
| self.qkv(self.norm(x)) |
| .reshape(batch, 3, 1, channels, height * width) |
| .transpose(-2, -1) |
| .unbind(1) |
| ) |
| h = F.scaled_dot_product_attention(q, k, v) |
| h = h.transpose(-2, -1).reshape(batch, channels, height, width) |
| return x + self.proj_out(h) |
|
|
|
|
| class Upsample(nn.Module): |
| def __init__(self, channels: int) -> None: |
| super().__init__() |
| self.conv = nn.Conv2d(channels, channels, 3, padding=1) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.conv(F.interpolate(x, scale_factor=2.0, mode="nearest")) |
|
|
|
|
| class Decoder(nn.Module): |
| def __init__( |
| self, |
| z_channels: int = 64, |
| ch: int = 128, |
| num_res_blocks: int = 2, |
| ) -> None: |
| super().__init__() |
| ch_mult = (1, 1, 2, 2, 4) |
| num_resolutions = len(ch_mult) |
| block_in = ch * ch_mult[-1] |
| current_resolution = 16 |
|
|
| self.conv_in = nn.Conv2d(z_channels, block_in, 3, padding=1) |
| self.mid = nn.ModuleList( |
| [ |
| ResnetBlock(block_in, block_in), |
| AttnBlock(block_in), |
| ResnetBlock(block_in, block_in), |
| ] |
| ) |
|
|
| self.up = nn.ModuleList() |
| for level in reversed(range(num_resolutions)): |
| block_out = ch * ch_mult[level] |
| blocks = nn.ModuleList() |
| for _ in range(num_res_blocks + 1): |
| blocks.append(ResnetBlock(block_in, block_out)) |
| block_in = block_out |
| if current_resolution == 16: |
| blocks.append(AttnBlock(block_in)) |
| if level != 0: |
| blocks.append(Upsample(block_in)) |
| current_resolution *= 2 |
| self.up.append(blocks) |
|
|
| self.norm_out = _group_norm(block_in) |
| self.conv_out = nn.Conv2d(block_in, 3, 3, padding=1) |
|
|
| def forward(self, latent: torch.Tensor) -> torch.Tensor: |
| hidden = self.conv_in(latent) |
| for block in self.mid: |
| hidden = block(hidden) |
| for level in self.up: |
| for block in level: |
| hidden = block(hidden) |
| return self.conv_out(F.silu(self.norm_out(hidden))) |
|
|
|
|
| class SemanticVAE(nn.Module): |
| """Inference portion of the trained DINOv2-B semantic autoencoder.""" |
|
|
| def __init__( |
| self, |
| *, |
| pretrained: bool = True, |
| latent_dim: int = 64, |
| encoder_layers: int = 6, |
| decoder_ch: int = 128, |
| decoder_num_res_blocks: int = 2, |
| ) -> None: |
| super().__init__() |
| self.encoder = timm.create_model( |
| DINO_MODEL_NAME, |
| pretrained=pretrained, |
| num_classes=0, |
| dynamic_img_size=True, |
| dynamic_img_pad=True, |
| ) |
| self.semantic_encoder = copy.deepcopy(self.encoder) |
| self.semantic_encoder.requires_grad_(False) |
| self.semantic_encoder.eval() |
|
|
| self.encoder_layers = encoder_layers |
| self.latent_dim = latent_dim |
| self.decoder_ch = decoder_ch |
| self.decoder_num_res_blocks = decoder_num_res_blocks |
| embed_dim = self.encoder.embed_dim |
| self.feature_norms = nn.ModuleList( |
| nn.LayerNorm(embed_dim) for _ in range(encoder_layers) |
| ) |
| branch_dim = latent_dim // 2 |
| self.encoder_projection = nn.Conv2d( |
| embed_dim * encoder_layers, branch_dim, 1 |
| ) |
| self.semantic_projection = nn.Conv2d(embed_dim, branch_dim, 1) |
| self.decoder = Decoder( |
| z_channels=latent_dim, |
| ch=decoder_ch, |
| num_res_blocks=decoder_num_res_blocks, |
| ) |
| self.register_buffer( |
| "dino_mean", torch.tensor(DINO_MEAN).view(1, 3, 1, 1), persistent=False |
| ) |
| self.register_buffer( |
| "dino_std", torch.tensor(DINO_STD).view(1, 3, 1, 1), persistent=False |
| ) |
| self.register_buffer( |
| "latent_mean", torch.zeros(1, latent_dim, 1, 1), persistent=True |
| ) |
| self.register_buffer( |
| "latent_std", torch.ones(1, latent_dim, 1, 1), persistent=True |
| ) |
| self.register_buffer( |
| "latent_stats_samples", torch.tensor(0, dtype=torch.int64), persistent=True |
| ) |
|
|
| def train(self, mode: bool = True) -> SemanticVAE: |
| super().train(mode) |
| self.semantic_encoder.eval() |
| return self |
|
|
| def _dino_input(self, pixels: torch.Tensor) -> torch.Tensor: |
| pixels = F.interpolate( |
| pixels, |
| size=_dino_spatial_size(pixels.shape[-2:]), |
| mode="bicubic", |
| align_corners=False, |
| antialias=True, |
| ) |
| return (pixels.add(1.0).mul(0.5) - self.dino_mean) / self.dino_std |
|
|
| @property |
| def has_latent_stats(self) -> bool: |
| return self.latent_stats_samples.item() > 0 |
|
|
| def set_latent_stats( |
| self, mean: torch.Tensor, std: torch.Tensor, *, samples: int |
| ) -> None: |
| self.latent_mean.copy_(mean.detach().view_as(self.latent_mean)) |
| self.latent_std.copy_(std.detach().view_as(self.latent_std)) |
| self.latent_stats_samples.fill_(samples) |
|
|
| def _require_latent_stats(self) -> None: |
| if not self.has_latent_stats: |
| raise RuntimeError( |
| "This model has no latent statistics. Use encode_raw/decode_raw, " |
| "or load a full .safetensors export containing latent statistics." |
| ) |
|
|
| def normalize_latents(self, latent: torch.Tensor) -> torch.Tensor: |
| self._require_latent_stats() |
| return (latent - self.latent_mean.to(latent.dtype)) / self.latent_std.to( |
| latent.dtype |
| ) |
|
|
| def denormalize_latents(self, latent: torch.Tensor) -> torch.Tensor: |
| self._require_latent_stats() |
| return latent * self.latent_std.to(latent.dtype) + self.latent_mean.to( |
| latent.dtype |
| ) |
|
|
| def encode_raw(self, pixels: torch.Tensor) -> torch.Tensor: |
| """Encode BCHW pixels to an unnormalized 16x-downsampled latent.""" |
| dino_input = self._dino_input(pixels) |
| features = self.encoder.forward_intermediates( |
| dino_input, |
| indices=self.encoder_layers, |
| norm=False, |
| output_fmt="NCHW", |
| intermediates_only=True, |
| ) |
| normalized = [ |
| norm(feature.permute(0, 2, 3, 1)) |
| .permute(0, 3, 1, 2) |
| .contiguous() |
| for feature, norm in zip(features, self.feature_norms, strict=True) |
| ] |
| encoder_latent = self.encoder_projection(torch.cat(normalized, dim=1)) |
| with torch.no_grad(): |
| semantic_feature = self.semantic_encoder.forward_intermediates( |
| dino_input, |
| indices=1, |
| norm=True, |
| output_fmt="NCHW", |
| intermediates_only=True, |
| )[0] |
| semantic_latent = self.semantic_projection(semantic_feature) |
| return torch.cat((encoder_latent, semantic_latent), dim=1) |
|
|
| def encode(self, pixels: torch.Tensor) -> torch.Tensor: |
| """Encode BCHW pixels to a normalized latent.""" |
| return self.normalize_latents(self.encode_raw(pixels)) |
|
|
| def decode_raw( |
| self, |
| latent: torch.Tensor, |
| output_size: tuple[int, int] | None = None, |
| ) -> torch.Tensor: |
| """Decode an unnormalized latent to pixels in [-1, 1].""" |
| reconstruction = torch.tanh(self.decoder(latent)) |
| if output_size is not None and reconstruction.shape[-2:] != output_size: |
| reconstruction = F.interpolate( |
| reconstruction, |
| size=output_size, |
| mode="bicubic", |
| align_corners=False, |
| antialias=True, |
| ) |
| return reconstruction |
|
|
| def decode( |
| self, |
| latent: torch.Tensor, |
| output_size: tuple[int, int] | None = None, |
| ) -> torch.Tensor: |
| """Decode a normalized BCHW latent to pixels in [-1, 1].""" |
| return self.decode_raw(self.denormalize_latents(latent), output_size) |
|
|
| def forward(self, pixels: torch.Tensor) -> torch.Tensor: |
| return self.decode(self.encode(pixels), output_size=pixels.shape[-2:]) |
|
|
|
|
| def load_vae( |
| checkpoint_path: str | Path, |
| *, |
| device: str | torch.device = "cpu", |
| dtype: torch.dtype = torch.float32, |
| pretrained: bool | None = None, |
| latent_stats_path: str | Path | None = None, |
| latent_dim: int = 64, |
| encoder_layers: int = 6, |
| decoder_ch: int = 128, |
| decoder_num_res_blocks: int = 2, |
| ) -> SemanticVAE: |
| """Load an eval-mode Semantic VAE from a full export or trainer checkpoint.""" |
| checkpoint_path = Path(checkpoint_path) |
| is_safetensors = checkpoint_path.suffix == ".safetensors" |
| if is_safetensors: |
| with safe_open(checkpoint_path, framework="pt", device="cpu") as handle: |
| metadata = handle.metadata() or {} |
| if metadata.get("format") != SAFETENSORS_FORMAT: |
| raise RuntimeError( |
| f"Unsupported Semantic VAE safetensors format: " |
| f"{metadata.get('format')!r}" |
| ) |
| latent_dim = int(metadata.get("latent_dim", latent_dim)) |
| encoder_layers = int(metadata.get("encoder_layers", encoder_layers)) |
| decoder_ch = int(metadata.get("decoder_ch", decoder_ch)) |
| decoder_num_res_blocks = int( |
| metadata.get("decoder_num_res_blocks", decoder_num_res_blocks) |
| ) |
|
|
| if pretrained is None: |
| pretrained = not is_safetensors |
| model = SemanticVAE( |
| pretrained=pretrained, |
| latent_dim=latent_dim, |
| encoder_layers=encoder_layers, |
| decoder_ch=decoder_ch, |
| decoder_num_res_blocks=decoder_num_res_blocks, |
| ) |
| state: Mapping[str, torch.Tensor] |
| if is_safetensors: |
| state = load_file(checkpoint_path, device="cpu") |
| model.load_state_dict(state, strict=True) |
| else: |
| checkpoint = torch.load( |
| checkpoint_path, map_location="cpu", weights_only=True, mmap=True |
| ) |
| if isinstance(checkpoint, Mapping) and "model" in checkpoint: |
| state = checkpoint["model"] |
| else: |
| state = checkpoint |
|
|
| incompatible = model.load_state_dict(state, strict=False) |
| if incompatible.unexpected_keys: |
| raise RuntimeError( |
| "Unexpected checkpoint keys: " + ", ".join(incompatible.unexpected_keys) |
| ) |
|
|
| required_prefixes = ( |
| "encoder.patch_embed.", |
| "feature_norms.", |
| "encoder_projection.", |
| "semantic_projection.", |
| "decoder.", |
| ) |
| missing_learned = [ |
| name |
| for name in incompatible.missing_keys |
| if name.startswith(required_prefixes) |
| ] |
| if missing_learned: |
| raise RuntimeError( |
| "Checkpoint is missing learned parameters: " |
| + ", ".join(missing_learned) |
| ) |
|
|
| if latent_stats_path is None: |
| candidate = ( |
| checkpoint_path.parent |
| / "latent_stats" |
| / f"semantic-{checkpoint_path.stem}.pt" |
| ) |
| if candidate.is_file(): |
| latent_stats_path = candidate |
| if latent_stats_path is not None: |
| stats = torch.load( |
| latent_stats_path, map_location="cpu", weights_only=True |
| ) |
| model.set_latent_stats( |
| stats["mean"], stats["std"], samples=int(stats["samples"]) |
| ) |
|
|
| model.requires_grad_(False) |
| return model.to(device=device, dtype=dtype).eval() |
|
|
|
|
| def _default_device() -> str: |
| if torch.cuda.is_available(): |
| return "cuda" |
| if torch.backends.mps.is_available(): |
| return "mps" |
| return "cpu" |
|
|
|
|
| def main() -> None: |
| """Reconstruct an image from its normalized latent.""" |
| from PIL import Image, ImageOps |
| from torchvision.transforms.functional import pil_to_tensor, to_pil_image |
|
|
| parser = argparse.ArgumentParser(description=main.__doc__) |
| parser.add_argument( |
| "--checkpoint", |
| type=Path, |
| default=Path("semantic_vae_step_00050000.safetensors"), |
| ) |
| parser.add_argument("input", type=Path) |
| parser.add_argument("--output", type=Path, default=Path("reconstructed.png")) |
| parser.add_argument("--size", type=int, default=1024) |
| parser.add_argument("--device", default=_default_device()) |
| args = parser.parse_args() |
|
|
| with Image.open(args.input) as source: |
| source = source.convert("RGBA") |
| background = Image.new("RGBA", source.size, "white") |
| image = Image.alpha_composite(background, source).convert("RGB") |
| image = ImageOps.fit( |
| image, |
| (args.size, args.size), |
| method=Image.Resampling.LANCZOS, |
| ) |
| pixels = pil_to_tensor(image).float().div(127.5).sub(1.0).unsqueeze(0) |
|
|
| model = load_vae(args.checkpoint, device=args.device) |
| pixels = pixels.to(args.device) |
| with torch.inference_mode(): |
| latent = model.encode(pixels) |
| reconstruction = model.decode(latent, output_size=(args.size, args.size)) |
| mse = F.mse_loss(reconstruction.float(), pixels.float()) |
| psnr = 10.0 * torch.log10(mse.new_tensor(4.0) / mse) |
|
|
| result = reconstruction[0].float().cpu().add(1.0).mul(0.5).clamp(0.0, 1.0) |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| to_pil_image(result).save(args.output) |
| print( |
| f"Saved {args.output} from latent {tuple(latent.shape)} " |
| f"using {args.device}; PSNR: {psnr.item():.2f} dB" |
| ) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
