semantic_vae.py

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