Upload 3 files

pyproject.toml

@@ -0,0 +1,11 @@
+[project]
+name = "semantic-vae"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "safetensors>=0.8.0",
+    "timm>=1.0.29",
+    "torch>=2.13.0",
+]

semantic_vae.py

@@ -0,0 +1,452 @@
+"""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()

semantic_vae_step_00050000.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e24402fc26c5806b6ccac0baf575d77cd6dee1ac413b2d84bfb5c6c43333825c
+size 859679660