medlat 0.1.0

A comprehensive collection of quantizers, continuous autoencoders, and tokenizers for medical imaging and general visual data processing

MedLat

MedLat (medlat) is a PyTorch library for first-stage models (tokenizers, continuous and discrete autoencoders, quantizers) and second-stage generators (autoregressive and diffusion-based), aimed at medical imaging and general visual data. Everything is exposed through a single model registry: instantiate by string ID, list IDs programmatically, and optionally read paper/code links via get_model_info.

2D and 3D: Registered models are built to cover standard 2D computer vision (natural images, 2D slices, projections) and 3D medical volumetric data (e.g. CT/MRI-style volumes). The same registry IDs apply; switch spatial dimensionality with constructor arguments such as dims=2 vs dims=3 (and matching img_size / spatial layout) on the builders you instantiate—so one codebase serves both pipelines.

Project page (SEO, FAQ, llms.txt): open docs/index.html locally, or publish the docs/ folder to GitHub Pages (see docs/sitemap.xml / docs/robots.txt).

---

What this repo contains

Layer	Role
First stage	Maps images to a compact representation: continuous latents, discrete codes, or token sequences.
Generator	Models the distribution in that space (tokens or latents): MaskGIT, MAR, DiT, LDM-style UNet, or ADM.
Utilities	Gaussian diffusion construction (medlat.diffusion), training helpers, ViT cores, alignments to foundation models, and GenWrapper for encoder ↔ generator wiring.

Typical flow: train or load a first-stage model → attach a generator whose input size matches the tokenizer/VAE output (n_embed, embed_dim, seq_len, stride, etc.) → sample or train end-to-end using the example notebooks.

---

Installation

pip install -e .

Core dependencies (see pyproject.toml): PyTorch, NumPy, Einops, timm, OmegaConf, MONAI. Optional [dev] extras include pytest, black, flake8, mypy.

Some code paths import extra packages (e.g. wavelets in quantizers); install those if you hit import errors for a specific model.

---

Quick start

from medlat import available_models, get_model, GenWrapper

# Discover registered IDs (hundreds of concrete configs)
print(list(available_models()))

# Metadata only (no instantiation): description, paper_url, code_url when set
from medlat import get_model_info
print(get_model_info("continuous.aekl.f8_d16"))

# First stage: tokenizer or VAE (dims=2 for 2D CV, dims=3 for volumetric 3D medical data)
tokenizer = get_model("token.titok.s_128", img_size=128, dims=2)
# Other examples: continuous.aekl.f8_d16, discrete.vq.f8_d8_e16384, token.vmae.s_p8_d16, ...

# Generator: match tokenizer latent interface (example for discrete + autoregressive)
generator = get_model(
    "maskgit.b",
    img_size=128,
    seq_len=tokenizer.num_latent_tokens,
    num_tokens=tokenizer.n_embed,
    in_channels=tokenizer.embed_dim,
)

# Glue: encoding/decoding and optional latent scale handling
wrapper = GenWrapper(generator, tokenizer)

GenWrapper selects encode/decode routes based on discrete vs continuous first stage and autoregressive vs non-autoregressive generator; see medlat/modules/wrapper.py for details.

---

Package layout (medlat/)

Path	Contents
medlat/registry.py	ModelRegistry, register_model, get_model, available_models, get_model_info.
medlat/first_stage/continuous/	KL autoencoders, MAISI, MedVAE, VAVAE, DCAE, soft-VQ / WQ-VAE style builders, plus one ViT AE entry used in HCVQ workflows.
medlat/first_stage/discrete/	Full VQ / RQ / FSQ / LFQ / BSQ / QINCo-family models and MaskGIT-style VQ presets.
medlat/first_stage/discrete/quantizer/	Standalone quantizer modules also registered under discrete.quantizer.*.
medlat/first_stage/token/	TiTok, MAETok, VMAE, DeTok, SoftVQ, ViTA.
medlat/generators/autoregressive/	MaskGIT, MAR.
medlat/generators/non_autoregressive/	DiT, LDM UNet, ADM (UNet + classifiers).
medlat/diffusion/	create_gaussian_diffusion, schedules, resampling (used with diffusion generators).
medlat/modules/	GenWrapper, ViT cores, positional embeddings, alignments, I/O helpers.

Subpackages often include a short README.md (e.g. generators/.../dit/README.md) with extra context for that family.

---

Naming conventions (registry IDs)

These patterns appear throughout IDs; exact letters vary by family.

• f{N} — Spatial downsampling factor of the bottleneck (e.g. f8 ≈ 1/8 resolution per spatial axis).
• d{N} — Latent channel width or embedding dimension (family-specific).
• e{N} — Codebook size for vector quantization.
• l{N} / b{N} — Levels or bit width for scalar / binary quantizers (FSQ, LFQ, BSQ).
• Scale tags — e.g. s / b / l / xl for model depth (TiTok, DeTok, DiT, MaskGIT, MAR).

Use available_models("token."), available_models("continuous."), etc. to list subsets.

---

Model families (registry)

Below, family means a group of related architectures; each family has many registered configs (different factors, widths, and presets). The table lists representative IDs and where they live in code.

A. Tokenizers & patch sequence models (token.*)

Family	Idea	Representative IDs
TiTok	Compact tokenization for generative modeling	token.titok.s_128, token.titok.b_128_p8_e2e, token.titok.l_256_e2e, …
MAETok	MAE-style tokenizer variants	token.maetok.s_256, token.maetok.b_512_p8, token.maetok.l_128, …
VMAE	ViT/VMAE-style tokenizers, symmetric and asymmetric	token.vmae.s_p8_d16, token.vmae.t_p8_d16, token.vmae.b_p16_d32, …
DeTok	Depth/width scale grid (ss, sb, …, xlxl)	token.detok.ss, token.detok.bb, token.detok.xlxl, …
SoftVQ	Soft vector quantization tokenizers	token.softvq.s_t32_d32, token.softvq.bl_t64_d32, …
ViTA	ViT-based tokenizer preset	token.vita.reconmae

Registration: medlat/first_stage/token/register.py, vmae/register.py, detok/detok.py, vita/vita.py.

---

B. Continuous autoencoders (continuous.*)

Family	Idea	Representative IDs
AEKL	LDM-style KL autoencoder (conv encoder/decoder)	continuous.aekl.f4_d3 … continuous.aekl.f32_d64
MAISI	MONAI MAISI-style 3D-friendly KL AE	continuous.maisi.f4_d3
MedVAE	KL AE + BiomedCLIP alignment	continuous.medvae.f8_d16, continuous.medvae.f8_d32
VAVAE	KL AE + alignment to DINOv2, MAE, or BiomedCLIP	continuous.vavae.f8_d32_dinov2, continuous.vavae.f16_d32_mae, …
DCAE	EfficientViT DC-AE (strong compression, optional HF weights)	continuous.dcae.f32c32, continuous.dcae.f64c128, continuous.dcae.f128c512
Soft VQ / WQ-VAE	Soft or warped quantization in a continuous wrapper	continuous.soft_vq.f8_d16_e16384, continuous.wqvae.f8_d4_e16384, variants with _dinov2 / _biomedclip

Registration: medlat/first_stage/continuous/register.py.

---

C. Discrete VAEs & full quantizer stacks (discrete.* except discrete.quantizer.*)

Family	Idea	Representative IDs
VQ-VAE / VQGAN-style	Conv VQ with LDM-style encoder/decoder	discrete.vq.f4_d3_e8192, discrete.vq.f8_d8_e16384, discrete.vq.f16_d32_e16384, …
MaskGIT VQ	Preset aligned with MaskGIT-style setups	discrete.maskgit.vq.f16_d256_e1024
MS-RQ	Multi-scale residual quantization	discrete.msrq.f16_d32_e4096
LFQ	Lookup-free / implicit codebook quantization	discrete.lfq.f16_d10_b10, many f/d/b sweeps
Simple QINCo / QINCo / RQ variants	Residual and improved QINCo-style quantizers	discrete.simple_qinco.*, discrete.qinco.*, discrete.rsimple_qinco.*, discrete.rqvae.*
SimVQ	SimVQ-style discrete bottleneck	discrete.simvq.f4_d3_e8192, …
FSQ	Finite scalar quantization	discrete.fsq.f8_d4_l16384, discrete.fsq.f16_d256_e8, …
BSQ	Binary spherical quantization	discrete.bsq.f16_d10_b10, …
HCVQ	Hybrid conv/ViT quantizer presets (residual, soft, grouped, SD-VQ, MS-RQ, VAE)	discrete.hcvq.residual_vq.S_16, discrete.hcvq.sd_vq.S_16, …

A ViT-based AE entry used alongside HCVQ-style pipelines is registered as discrete.hcvq.vae.S_16 (builder in continuous/register.py).

Registration: medlat/first_stage/discrete/register.py.

---

D. Standalone quantizers (discrete.quantizer.*)

Use these when composing custom VQ models or for ablations. Registered names include:

vector_quantizer, vector_quantizer2, gumbel_quantizer, simple_qinco, residual_quantizer, grouped_residual_quantizer, msrq_vector_quantizer2, msrq_vector_quantizer3d, lookup_free_quantizer, binary_spherical_quantizer, finite_scalar_quantizer, soft_vector_quantizer, …

Full names are discrete.quantizer.<name>. Implementation: medlat/first_stage/discrete/quantizer/quantize.py.

---

E. Autoregressive generators

Family	Idea	IDs
MaskGIT	Masked token modeling	maskgit.b, maskgit.l, maskgit.h
MAR	Autoregressive over tokens	mar.b, mar.l, mar.h

Registration: medlat/generators/autoregressive/maskgit/register.py, mar/register.py.

---

F. Non-autoregressive (diffusion) generators

Family	Idea	IDs
DiT	Diffusion transformer (patchified latent)	dit.s_1 … dit.s_8, dit.b_*, dit.l_*, dit.xl_* (depth × patch size)
LDM	Latent diffusion UNet (stride in ID)	ldm.f1 … ldm.f16
ADM	Dhariwal–Nichol–style UNet diffusion + classifiers	adm.diffusion.{64,128,256,512}{C,U}, adm.classifier.{64,128,256,512}C

Diffusion sampling utilities live under medlat/diffusion/ (e.g. create_gaussian_diffusion).

Registration: medlat/generators/non_autoregressive/dit/register.py, ldm/register.py, adm/register.py.

---

Examples (notebooks)

Notebook	Topic
example_tokenizer.ipynb	Train / use first-stage models
example_generator_maskgit.ipynb	MaskGIT pipeline
example_generator_mar.ipynb	MAR pipeline
example_generator_nonautoregressive.ipynb	DiT / LDM / ADM-style generation

---

Discovering models in code

from medlat import available_models

# Everything
all_ids = list(available_models())

# By prefix
token_ids = list(available_models("token."))
vq_ids = list(available_models("discrete.vq."))
dit_ids = list(available_models("dit."))

---

Citation

If you use MedLat in research or a publication, please cite the software. Version and author metadata also appear in pyproject.toml.

BibTeX:

@software{bubeck_medlat_2025,
  author = {Bubeck, Niklas},
  title = {{MedLat}: {PyTorch} library for first-stage models and latent generators},
  url = {https://github.com/niklasbubeck/MedLat},
  version = {0.1.0},
  year = {2025},
}

Plain text (APA-style):

Bubeck, N. (2025). MedLat (Version 0.1.0) [Computer software]. https://github.com/niklasbubeck/MedLat

---

License

MIT — see pyproject.toml for the declared license and author metadata.

---