topazy/vqvae_transformer_onearthenergy

Files

File	Model	Size	Params
vqvae/vqgan_epoch_25.pth	VQVAE, epoch 25	710 MB	177.5 M
transformer/transformer_epoch_50.pth	Transformer, epoch 50	754 MB	185.3 M

pip install huggingface_hub

from huggingface_hub import hf_hub_download

vqvae = hf_hub_download("topazy/vqvae_transformer_onearthenergy",
                        "vqvae/vqgan_epoch_25.pth")
trans = hf_hub_download("topazy/vqvae_transformer_onearthenergy",
                        "transformer/transformer_epoch_50.pth")

Loading into the models:

import torch
from omegaconf import OmegaConf
from lpu3dnet.frame import vqgan, transformer

cfg_v = OmegaConf.load("lpu3dnet/config/ex12/vqgan.yaml")
cfg_t = OmegaConf.load("lpu3dnet/config/ex12/transformer.yaml")

m_vqvae = vqgan.VQGAN(cfg_v)
m_vqvae.load_state_dict(torch.load(vqvae, map_location="cpu"))   # strict=True

m_trans = transformer.Transformer(cfg_t)
m_trans.load_state_dict(torch.load(trans, map_location="cpu"))   # strict=True

Input convention. Binary volumes, 1 = pore, 0 = grain, shape (B, 1, 64, 64, 64), float32. From the raw 8-bit tifs: img.astype('float32') / 255 > 0.5. Decoder output is segmented with two-class Otsu, not a fixed 0.5 threshold.

Links

• Code: https://github.com/ZihanRen/vqvae_transformer_newenergy (placeholder)
• Data: NETL EDX, DOI 10.2172/1963265

License

MIT. Built on nanoGPT; evaluation uses OpenPNM and PoreSpy.