vae/__init__.py

from __future__ import annotations

import os
from pathlib import Path
from typing import Optional, Union, Sequence, List

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image


IMAGE_SIZE = 512
LATENT_CHANNELS = 8

C512 = 32
C256 = 64
C128 = 96
C64 = 128

NUM_RES_BLOCKS = 2

LATENT_DTYPE = torch.float16
MODEL_DTYPE = torch.float32

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def valid_num_groups(channels: int, preferred: int = 16) -> int:
    upper = min(channels, preferred)
    for groups in range(upper, 0, -1):
        if channels % groups == 0:
            return groups
    return 1


class ResBlock(nn.Module):
    def __init__(self, channels: int):
        super().__init__()
        groups = valid_num_groups(channels)

        self.norm1 = nn.GroupNorm(groups, channels, eps=1e-5)
        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1)
        self.norm2 = nn.GroupNorm(groups, channels, eps=1e-5)
        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = self.conv1(F.silu(self.norm1(x)))
        x = self.conv2(F.silu(self.norm2(x)))
        return residual + x


class Encoder(nn.Module):
    def __init__(self):
        super().__init__()

        self.input = nn.Conv2d(3, C512, 3, padding=1)

        self.block0 = nn.Sequential(
            ResBlock(C512),
            ResBlock(C512),
        )

        self.down1 = nn.Conv2d(C512, C256, 4, stride=2, padding=1)
        self.block1 = nn.Sequential(
            ResBlock(C256),
            ResBlock(C256),
        )

        self.down2 = nn.Conv2d(C256, C128, 4, stride=2, padding=1)
        self.block2 = nn.Sequential(
            ResBlock(C128),
            ResBlock(C128),
        )

        self.down3 = nn.Conv2d(C128, C64, 4, stride=2, padding=1)
        self.block3 = nn.Sequential(
            ResBlock(C64),
            ResBlock(C64),
        )

        self.final_norm = nn.GroupNorm(valid_num_groups(C64), C64)
        self.mu = nn.Conv2d(C64, LATENT_CHANNELS, 3, padding=1)
        self.logvar = nn.Conv2d(C64, LATENT_CHANNELS, 3, padding=1)

    def forward(self, x: torch.Tensor):
        x = self.block0(self.input(x))
        x = self.block1(F.silu(self.down1(x)))
        x = self.block2(F.silu(self.down2(x)))
        x = self.block3(F.silu(self.down3(x)))
        x = F.silu(self.final_norm(x))

        mu = self.mu(x)
        logvar = torch.clamp(self.logvar(x), -10.0, 10.0)

        return mu, logvar


class Decoder(nn.Module):
    def __init__(self):
        super().__init__()

        self.input = nn.Conv2d(LATENT_CHANNELS, C64, 3, padding=1)

        self.block3 = nn.Sequential(
            ResBlock(C64),
            ResBlock(C64),
        )

        self.up1 = nn.Conv2d(C64, C128, 3, padding=1)
        self.block2 = nn.Sequential(
            ResBlock(C128),
            ResBlock(C128),
        )

        self.up2 = nn.Conv2d(C128, C256, 3, padding=1)
        self.block1 = nn.Sequential(
            ResBlock(C256),
            ResBlock(C256),
        )

        self.up3 = nn.Conv2d(C256, C512, 3, padding=1)
        self.block0 = nn.Sequential(
            ResBlock(C512),
            ResBlock(C512),
        )

        self.final_norm = nn.GroupNorm(valid_num_groups(C512), C512)
        self.output = nn.Conv2d(C512, 3, 3, padding=1)

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        x = self.block3(self.input(z))

        x = F.interpolate(x, scale_factor=2, mode="nearest")
        x = self.block2(F.silu(self.up1(x)))

        x = F.interpolate(x, scale_factor=2, mode="nearest")
        x = self.block1(F.silu(self.up2(x)))

        x = F.interpolate(x, scale_factor=2, mode="nearest")
        x = self.block0(F.silu(self.up3(x)))

        x = F.silu(self.final_norm(x))
        return torch.tanh(self.output(x))


class VAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()

    def encode(self, x: torch.Tensor):
        return self.encoder(x)

    def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor):
        std = torch.exp(0.5 * logvar)
        return mu + torch.randn_like(std) * std

    def decode(self, z: torch.Tensor):
        return self.decoder(z)

    def forward(self, x: torch.Tensor, sample: bool = True):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar) if sample else mu
        return self.decode(z), mu, logvar, z


_MODEL: Optional[VAE] = None
_CHECKPOINT_PATH: Optional[Path] = None


def _candidate_checkpoints() -> List[Path]:
    candidates = []

    env_path = os.environ.get("QARVEXIUM_VAE_CHECKPOINT")
    if env_path:
        candidates.append(Path(env_path))

    package_dir = Path(__file__).resolve().parent

    candidates.extend([
        package_dir / "qvae.pt",
        package_dir / "checkpoints" / "qvae.pt",
    ])

    result = []
    seen = set()

    for path in candidates:
        path = path.expanduser().resolve()
        if path not in seen:
            seen.add(path)
            result.append(path)

    return result


def _find_checkpoint() -> Path:
    candidates = _candidate_checkpoints()

    for path in candidates:
        if path.is_file():
            return path

    searched = "\n".join(f"  - {p}" for p in candidates)

    raise FileNotFoundError(
        "Could not find the Qarvexium VAE checkpoint.\n\n"
        f"Searched:\n{searched}\n\n"
        "Set QARVEXIUM_VAE_CHECKPOINT to the absolute path "
        "of your checkpoint."
    )


def _extract_state_dict(checkpoint):
    if not isinstance(checkpoint, dict):
        raise RuntimeError("Unsupported checkpoint format.")

    if "model" in checkpoint:
        state_dict = checkpoint["model"]
    elif "state_dict" in checkpoint:
        state_dict = checkpoint["state_dict"]
    elif "model_state_dict" in checkpoint:
        state_dict = checkpoint["model_state_dict"]
    else:
        state_dict = checkpoint

    if not isinstance(state_dict, dict):
        raise RuntimeError("Checkpoint model state is not a state_dict.")

    return {
        key[len("module."):] if key.startswith("module.") else key: value
        for key, value in state_dict.items()
    }


def load_model(
    checkpoint_path: Optional[Union[str, os.PathLike]] = None,
    force_reload: bool = False,
) -> VAE:
    global _MODEL, _CHECKPOINT_PATH

    requested_path = (
        Path(checkpoint_path).expanduser().resolve()
        if checkpoint_path is not None
        else None
    )

    if (
        not force_reload
        and _MODEL is not None
        and (requested_path is None or requested_path == _CHECKPOINT_PATH)
    ):
        return _MODEL

    if requested_path is None:
        requested_path = _find_checkpoint()

    if not requested_path.is_file():
        raise FileNotFoundError(
            f"Checkpoint does not exist:\n{requested_path}"
        )

    print(
        f"[Qarvexium VAE] Loading checkpoint: {requested_path}",
        flush=True,
    )

    checkpoint = torch.load(
        requested_path,
        map_location="cpu",
        weights_only=False,
    )

    if isinstance(checkpoint, dict):
        saved_config = checkpoint.get("config")

        if saved_config is not None:
            expected = {
                "image_size": IMAGE_SIZE,
                "latent_channels": LATENT_CHANNELS,
                "c512": C512,
                "c256": C256,
                "c128": C128,
                "c64": C64,
            }

            for key, expected_value in expected.items():
                saved_value = saved_config.get(key)
                if saved_value is not None and saved_value != expected_value:
                    raise RuntimeError(
                        "Checkpoint architecture mismatch:\n"
                        f"  {key}: checkpoint={saved_value}, "
                        f"package={expected_value}"
                    )

    model = VAE()
    state_dict = _extract_state_dict(checkpoint)

    try:
        model.load_state_dict(state_dict, strict=True)
    except RuntimeError as error:
        raise RuntimeError(
            "The checkpoint still does not match the "
            "Qarvexium VAE architecture.\n\n"
            f"{error}"
        ) from error

    parameter_count = sum(p.numel() for p in model.parameters())

    if isinstance(checkpoint, dict):
        saved_parameters = checkpoint.get("parameters")
        if (
            saved_parameters is not None
            and int(saved_parameters) != parameter_count
        ):
            raise RuntimeError(
                "Checkpoint parameter count mismatch:\n"
                f"  checkpoint = {saved_parameters:,}\n"
                f"  package    = {parameter_count:,}"
            )

    model.float().eval().to(DEVICE)

    _MODEL = model
    _CHECKPOINT_PATH = requested_path

    print(
        f"[Qarvexium VAE] Loaded successfully "
        f"({parameter_count:,} parameters)",
        flush=True,
    )
    print(f"[Qarvexium VAE] Model dtype: {MODEL_DTYPE}", flush=True)
    print(f"[Qarvexium VAE] Latent dtype: {LATENT_DTYPE}", flush=True)
    print(f"[Qarvexium VAE] Device: {DEVICE}", flush=True)

    return model


def _prepare_image(image: Image.Image) -> torch.Tensor:
    if not isinstance(image, Image.Image):
        raise TypeError("Expected PIL.Image.Image")

    image = image.convert("RGB").resize(
        (IMAGE_SIZE, IMAGE_SIZE),
        resample=Image.Resampling.BICUBIC,
    )

    array = np.asarray(image, dtype=np.uint8)

    if array.shape != (IMAGE_SIZE, IMAGE_SIZE, 3):
        raise ValueError(
            f"Expected image shape ({IMAGE_SIZE}, {IMAGE_SIZE}, 3), "
            f"got {array.shape}"
        )

    array = np.transpose(array, (2, 0, 1))
    array = np.ascontiguousarray(array)

    tensor = torch.from_numpy(array)
    tensor = tensor.float() / 127.5 - 1.0

    return tensor.unsqueeze(0)


def _tensor_to_pil(tensor: torch.Tensor) -> Image.Image:
    if tensor.ndim == 4:
        if tensor.shape[0] != 1:
            raise ValueError("Expected a single image tensor with batch size 1.")
        tensor = tensor[0]

    if tensor.ndim != 3:
        raise ValueError(
            f"Expected CHW tensor, got shape {tuple(tensor.shape)}"
        )

    tensor = ((tensor.float().clamp(-1.0, 1.0) + 1.0) / 2.0)
    tensor = (
        tensor.permute(1, 2, 0)
        .cpu()
        .numpy()
    )

    array = (tensor * 255.0).round().astype(np.uint8)
    return Image.fromarray(array, mode="RGB")


def _validate_latent(
    latent: torch.Tensor,
    allow_batch: bool = True,
) -> torch.Tensor:
    if not isinstance(latent, torch.Tensor):
        raise TypeError("latent must be a torch.Tensor")

    if latent.ndim == 3:
        latent = latent.unsqueeze(0)

    if latent.ndim != 4:
        raise ValueError(
            "Expected latent shape (8,64,64) or (B,8,64,64), "
            f"got {tuple(latent.shape)}"
        )

    if not allow_batch and latent.shape[0] != 1:
        raise ValueError("Expected batch size 1.")

    expected = (LATENT_CHANNELS, 64, 64)

    if tuple(latent.shape[1:]) != expected:
        raise ValueError(
            "Invalid latent shape.\n"
            f"Expected: (B, {LATENT_CHANNELS}, 64, 64)\n"
            f"Got:      {tuple(latent.shape)}"
        )

    return latent


@torch.inference_mode()
def encode(
    image: Image.Image,
    checkpoint_path: Optional[Union[str, os.PathLike]] = None,
) -> torch.Tensor:
    model = load_model(checkpoint_path)
    x = _prepare_image(image).to(DEVICE, dtype=MODEL_DTYPE)
    mu, _ = model.encode(x)
    return mu.detach().to(LATENT_DTYPE).cpu()


@torch.inference_mode()
def decode(
    latent: torch.Tensor,
    checkpoint_path: Optional[Union[str, os.PathLike]] = None,
) -> Image.Image:
    model = load_model(checkpoint_path)
    latent = _validate_latent(latent, allow_batch=False)
    z = latent.to(DEVICE, dtype=MODEL_DTYPE)
    return _tensor_to_pil(model.decode(z))


def encode_path(
    path: Union[str, os.PathLike],
    checkpoint_path: Optional[Union[str, os.PathLike]] = None,
) -> torch.Tensor:
    path = Path(path).expanduser().resolve()

    if not path.is_file():
        raise FileNotFoundError(f"Image not found:\n{path}")

    with Image.open(path) as image:
        image = image.convert("RGB").copy()

    return encode(image, checkpoint_path=checkpoint_path)


def reconstruct_path(
    path: Union[str, os.PathLike],
    checkpoint_path: Optional[Union[str, os.PathLike]] = None,
) -> Image.Image:
    return decode(
        encode_path(path, checkpoint_path=checkpoint_path),
        checkpoint_path=checkpoint_path,
    )


def reconstruct(
    image: Image.Image,
    checkpoint_path: Optional[Union[str, os.PathLike]] = None,
) -> Image.Image:
    return decode(
        encode(image, checkpoint_path=checkpoint_path),
        checkpoint_path=checkpoint_path,
    )


@torch.inference_mode()
def encode_batch(
    images: Sequence[Image.Image],
    checkpoint_path: Optional[Union[str, os.PathLike]] = None,
) -> torch.Tensor:
    if len(images) == 0:
        raise ValueError("images cannot be empty.")

    model = load_model(checkpoint_path)

    batch = torch.stack(
        [_prepare_image(image)[0] for image in images],
        dim=0,
    ).to(DEVICE, dtype=MODEL_DTYPE)

    mu, _ = model.encode(batch)
    return mu.detach().to(LATENT_DTYPE).cpu()


@torch.inference_mode()
def decode_batch(
    latents: torch.Tensor,
    checkpoint_path: Optional[Union[str, os.PathLike]] = None,
) -> List[Image.Image]:
    model = load_model(checkpoint_path)

    if not isinstance(latents, torch.Tensor):
        raise TypeError("latents must be a torch.Tensor")

    if latents.ndim != 4:
        raise ValueError(
            "Expected shape (B, 8, 64, 64), "
            f"got {tuple(latents.shape)}"
        )

    expected = (LATENT_CHANNELS, 64, 64)

    if tuple(latents.shape[1:]) != expected:
        raise ValueError(
            "Invalid latent shape.\n"
            f"Expected: (B, {LATENT_CHANNELS}, 64, 64)\n"
            f"Got:      {tuple(latents.shape)}"
        )

    z = latents.to(DEVICE, dtype=MODEL_DTYPE)
    reconstruction = model.decode(z)

    reconstruction = ((reconstruction.float().clamp(-1.0, 1.0) + 1.0) / 2.0)
    reconstruction = (
        reconstruction.permute(0, 2, 3, 1)
        .cpu()
        .numpy()
    )

    arrays = (reconstruction * 255.0).round().astype(np.uint8)

    return [
        Image.fromarray(array, mode="RGB")
        for array in arrays
    ]


def latent_info(latent: torch.Tensor) -> dict:
    latent = _validate_latent(latent)

    values_per_image = latent[0].numel()
    fp32_bytes = values_per_image * 4
    fp16_bytes = values_per_image * 2

    return {
        "shape": tuple(latent.shape),
        "dtype": str(latent.dtype),
        "device": str(latent.device),
        "values_per_image": values_per_image,
        "fp32_bytes_per_image": fp32_bytes,
        "fp16_bytes_per_image": fp16_bytes,
        "fp32_kib_per_image": fp32_bytes / 1024.0,
        "fp16_kib_per_image": fp16_bytes / 1024.0,
        "storage_reduction": 2.0,
    }


def model_info(
    checkpoint_path: Optional[Union[str, os.PathLike]] = None,
) -> dict:
    model = load_model(checkpoint_path)

    parameter_count = sum(
        parameter.numel()
        for parameter in model.parameters()
    )

    latent_values = LATENT_CHANNELS * 64 * 64

    return {
        "parameters": parameter_count,
        "parameters_millions": parameter_count / 1_000_000.0,
        "device": str(DEVICE),
        "model_dtype": str(MODEL_DTYPE),
        "latent_dtype": str(LATENT_DTYPE),
        "image_size": IMAGE_SIZE,
        "latent_channels": LATENT_CHANNELS,
        "latent_shape": (LATENT_CHANNELS, 64, 64),
        "latent_values": latent_values,
        "latent_fp32_kib": latent_values * 4 / 1024.0,
        "latent_fp16_kib": latent_values * 2 / 1024.0,
        "architecture": "32 / 64 / 96 / 128",
        "checkpoint": str(_CHECKPOINT_PATH) if _CHECKPOINT_PATH else None,
    }


__all__ = [
    "VAE",
    "Encoder",
    "Decoder",
    "encode",
    "decode",
    "encode_path",
    "reconstruct",
    "reconstruct_path",
    "encode_batch",
    "decode_batch",
    "load_model",
    "model_info",
    "latent_info",
    "IMAGE_SIZE",
    "LATENT_CHANNELS",
    "LATENT_DTYPE",
    "MODEL_DTYPE",
    "DEVICE",
]