EMA-VQGAN (16384-f8-d256)
This repository contains the official pre-trained weights for EMA-VQGAN-16384-f8-d256 (Stage 1 Image Quantization Model), serving as the visual discrete token backbone for the TextArtT5 project.
🔗 Main Repository
For full source code, training pipelines, dataset details, and the Stage 2 Autoregressive Transformer, please visit our main GitHub repository:
Model Summary
- Model Type: EMA VQGAN (Vector Quantized GAN with Exponential Moving Average)
- Codebook Size (N): 16,384
- Downsampling Factor (f): 8 (Compresses 256 × 256 × 3 images into 32 × 32 token grids)
- Latent Dimension (d): 256
- Parameters: ~94.5M
- Validation rFID: 2.21
- Training Corpus: Combined dataset including ImageNet, Conceptual Captions (CC3M), MS-COCO, Flickr, and CelebA / CelebA-Dialog.
Quick Inference Example
Make sure you have cloned the TextArtT5 GitHub repository to access the model architecture classes.
import torch
from PIL import Image
from torchvision.transforms import v2
from src.model.vq_model import EMAVQModel
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float32
# 1. Load Pre-trained Model directly from Hugging Face
model = EMAVQModel.from_pretrained("TrungNt14/EMAVQ-16384-f8-d256").to(device=device, dtype=dtype)
model.eval()
# 2. Image Preprocessing Pipeline
image_size = (256, 256)
transform = v2.Compose([
v2.Resize(max(image_size), interpolation=v2.InterpolationMode.BICUBIC),
v2.CenterCrop(image_size),
v2.ToImage(),
v2.ToDtype(dtype, scale=True),
v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])
# 3. Process Input Image
img = Image.open("path/to/your/image.png").convert("RGB")
img_tensor = transform(img).unsqueeze(0).to(device)
# 4. Quantize & Reconstruct
with torch.no_grad():
# Encode to discrete codebook latents
quant, _, info = model.encode(img_tensor)
# Decode back to image space
reconstruction = model.decode(quant)