Views
No views yet
| Component | Value |
|---|---|
| Encoder | ViT-B (vit_b_enc) |
| Decoder | Patched UNet (unet_patched) |
| Quantizer | FSQ (codebook: 8-8-8-6-5, vocab: 15,360) |
| Image size | 448×448 px |
| Patch size | 16×16 px |
| Token grid | 28×28 = 784 tokens per image |
| Input channels | 2 (Digital Surface Model, Digital Terrain Model) |
| Latent dim | 5 |
config.json for exact normalization parameters.1import torch
2from huggingface_hub import hf_hub_download
3from terratorch.models.backbones.terramind.tokenizer.vqvae import DiVAE
4
5# Download weights and config
6weights_path = hf_hub_download(repo_id="YOUR_REPO_ID", filename="tokenizer.pt")
7
8# Instantiate model
9tokenizer = DiVAE(
10 image_size=448,
11 patch_size=16,
12 n_channels=2,
13 enc_type="vit_b_enc",
14 dec_type="unet_patched",
15 quant_type="fsq",
16 codebook_size="8-8-8-6-5",
17 latent_dim=5,
18 post_mlp=True,
19 norm_codes=True,
20)
21
22# Load weights
23state_dict = torch.load(weights_path, map_location="cpu")
24tokenizer.load_state_dict(state_dict)
25tokenizer.eval()
26
27# Encode: image → tokens
28x = torch.randn(1, 2, 448, 448)
29quant, code_loss, tokens = tokenizer.encode(x)
30print(tokens.shape) # (1, 28, 28)
31
32# Decode: tokens → reconstruction (diffusion sampling)
33recon = tokenizer(x, timesteps=50)ahn-best-epoch-0002.ckpt