qvx-o/qvae

Qarvexium VAE

A Variational Autoencoder (VAE) for efficient image encoding and reconstruction. This model compresses 512×512 RGB images into compact 8-channel 64×64 latent representations, achieving significant dimensionality reduction while maintaining image quality.

Model Details

• Model Type: Variational Autoencoder (VAE)
• Architecture: Convolutional encoder-decoder with residual blocks
• Input Resolution: 512×512 RGB images
• Latent Space: 8 channels × 64×64 spatial dimensions
• Parameters: ~5M parameters
• Compression Ratio: ~6.4× reduction (512²×3 → 64²×8)

Architecture Highlights

• Encoder: Progressive downsampling through 4 stages (32→64→96→128 channels)
• Decoder: Symmetric upsampling with nearest-neighbor interpolation
• Residual Blocks: 2 ResBlocks per stage with GroupNorm and SiLU activation
• Latent Space: Learns both mean (μ) and log-variance for probabilistic encoding

Usage

Installation

pip install torch torchvision pillow numpy

Basic Usage

from PIL import Image
from vae import encode, decode, reconstruct

# Load an image
image = Image.open("input.jpg")

# Encode image to latent representation
latent = encode(image)  # Shape: (1, 8, 64, 64), dtype: float16
print(f"Latent shape: {latent.shape}")

# Decode latent back to image
reconstructed = decode(latent)
reconstructed.save("output.jpg")

# Or do both in one step
reconstructed = reconstruct(image)

Batch Processing

from vae import encode_batch, decode_batch
from PIL import Image

# Encode multiple images
images = [Image.open(f"image_{i}.jpg") for i in range(4)]
latents = encode_batch(images)  # Shape: (4, 8, 64, 64)

# Decode multiple latents
reconstructed_images = decode_batch(latents)
for i, img in enumerate(reconstructed_images):
    img.save(f"reconstructed_{i}.jpg")

File Path Operations

from vae import encode_path, reconstruct_path

# Encode directly from file path
latent = encode_path("input.jpg")

# Reconstruct and save
reconstructed = reconstruct_path("input.jpg")
reconstructed.save("reconstructed.jpg")

Custom Checkpoint Path

from vae import load_model, encode

# Load model with custom checkpoint
model = load_model(checkpoint_path="path/to/custom/checkpoint.pt")

# Or specify checkpoint in encode/decode operations
latent = encode(image, checkpoint_path="path/to/checkpoint.pt")

Model Information

from vae import model_info, latent_info

# Get model details
info = model_info()
print(f"Parameters: {info['parameters']:,}")
print(f"Device: {info['device']}")
print(f"Latent shape: {info['latent_shape']}")

# Get latent tensor details
latent_details = latent_info(latent)
print(f"FP16 storage: {latent_details['fp16_kib_per_image']:.2f} KiB per image")

Data Types

• Model Weights: float32
• Latent Representations: float16 (for storage efficiency)
• Input/Output Images: uint8 (0-255 RGB)

Storage Efficiency

• Original Image: 512×512×3 = 786,432 values (3 MB in FP32, 1.5 MB in FP16)
• Latent Representation: 64×64×8 = 32,768 values (128 KB in FP32, 64 KB in FP16)
• Compression Ratio: ~24× in FP16 format

Checkpoint Loading

The model searches for checkpoints in the following order:

1. Path specified in QARVEXIUM_VAE_CHECKPOINT environment variable
2. vae/best.pt (package directory)
3. vae/checkpoints/best.pt
4. /kaggle/working/coco_vae_10k/checkpoints/best.pt (Kaggle)
5. /content/coco_vae_10k/checkpoints/best.pt (Colab)

Setting Custom Checkpoint Path

export QARVEXIUM_VAE_CHECKPOINT="/path/to/your/checkpoint.pt"

Or in Python:

import os
os.environ['QARVEXIUM_VAE_CHECKPOINT'] = "/path/to/your/checkpoint.pt"

Example Applications

Image Compression Pipeline

from vae import encode, decode
from PIL import Image
import torch

# Compress image to latent
image = Image.open("large_image.jpg")
latent = encode(image)

# Save compressed latent (64 KB vs. 1.5 MB)
torch.save(latent, "compressed.pt")

# Later: load and reconstruct
latent = torch.load("compressed.pt")
reconstructed = decode(latent)
reconstructed.save("decompressed.jpg")

Latent Space Interpolation

from vae import encode, decode
from PIL import Image
import torch

# Encode two images
image1 = Image.open("image1.jpg")
image2 = Image.open("image2.jpg")
latent1 = encode(image1)
latent2 = encode(image2)

# Interpolate in latent space
alpha = 0.5
interpolated = alpha * latent1 + (1 - alpha) * latent2

# Decode interpolated latent
result = decode(interpolated)
result.save("interpolated.jpg")

Dataset Preprocessing

from vae import encode_batch
from PIL import Image
import torch
from pathlib import Path

# Encode entire dataset
image_paths = list(Path("dataset").glob("*.jpg"))
batch_size = 32

all_latents = []
for i in range(0, len(image_paths), batch_size):
    batch_paths = image_paths[i:i+batch_size]
    images = [Image.open(p) for p in batch_paths]
    latents = encode_batch(images)
    all_latents.append(latents)

# Save compressed dataset
torch.save(torch.cat(all_latents), "dataset_latents.pt")

Performance Considerations

Device Selection

from vae import DEVICE
print(f"Using device: {DEVICE}")  # Automatically uses CUDA if available

Memory Usage

• Single image encoding: ~100 MB VRAM
• Batch processing: ~100 MB + (batch_size × 6 MB) VRAM
• Latent storage (FP16): 64 KB per image

Optimization Tips

1. Batch Processing: Use encode_batch() and decode_batch() for multiple images
2. FP16 Storage: Latents are automatically stored in float16 to save space
3. Model Caching: The model is cached after first load (use force_reload=False)

Limitations

• Fixed input size: Images must be resized to 512×512
• Color only: Grayscale images are converted to RGB
• Lossy compression: Reconstruction is approximate, not pixel-perfect
• BICUBIC resampling: Input images are resized using BICUBIC interpolation

License

MIT License