TrungNt14/EMAVQ-16384-f8-d256

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:

👉 GitHub: TextArtT5 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)