Visage-Image

GitHub · Architecture · Prompting · Data · ComfyUI

Introduction

Visage-Image is a 256px text-to-image model for face generation. It has three parts:

• DiT (236.8M parameters): a single-stream rectified-flow transformer.
• VAE (128.4M parameters): an f16 autoencoder that turns each 256px image into a 16×16×96 latent.
• Text encoder (frozen, not included): the first 8 layers of Qwen3-0.6B, at the hub revision recorded in text_encoder/config.json.

The code and the weights are released under Apache-2.0.

Highlights

A transformer VAE decoder. The VAEs that latent diffusion usually runs on (Stable Diffusion, SDXL, FLUX.1/FLUX.2, Qwen-Image, Sana's DC-AE) all decode with convolutions, and their last layers run at full image resolution. The Visage-Image decoder is a ViT. It works on the 16×16 latent grid, splits each token into its 2×2 children for a second ViT stage at 32×32, and emits pixels only at the end, through a linear head that paints each 8×8 patch. Two small residual conv refiners then smooth across patch seams. In a head-to-head test with the encoder and losses held fixed, the ViT decoder beat the conv decoder on every validation metric and trained about 5× faster.

A wavelet encoder. Following LiteVAE, the 26M-parameter encoder is given a three-level Haar wavelet pyramid instead of learning frequency decomposition from raw pixels. Two parameter-free shortcuts fold the pixels straight into the deepest stage and into the latent.

No KL. The latent is deterministic. DINOv2 feature alignment keeps it organised in place of a KL penalty.

The result. At 25% lower rate (f16×96 against FLUX.2's f8×32), the VAE ties the FLUX.2 VAE on reconstruction, eye detail included, and beats it on iFID. It loses on rFID (see the table below).

A small single-stream DiT. Text and image tokens share one sequence and one set of weights. Captions come from a mid-layer of a small LLM (layer 8 of Qwen3-0.6B) rather than T5 or CLIP. They are padded with learned per-position embeddings to fixed length buckets, so a sample never depends on what it is batched with. Sixteen register tokens and a 2D spiral RoPE complete the design.

ARCHITECTURE.md covers each choice, the measurements behind it and what was tried instead.

Quick Start

The code is on GitHub and on PyPI as visage-image. The weights in this repository download on first use.

pip install visage-image

import torch
from torchvision.utils import save_image

from visage import load_dit, load_text_encoder, load_vae, sample

prompt = "a young woman with long red hair and freckles, soft window light"
text, mask = load_text_encoder()([prompt] * 4)  # Qwen3-0.6B, first 8 layers
dit = load_dit()
latents = sample(dit, text, mask, seed=0)  # SDE, churn 0.3, CFG 3, 50 steps, shift 3.3
with torch.no_grad():
    images = load_vae().decode(latents).float()
save_image((images + 1) / 2, "visage.png", nrow=4, padding=0)

Or from the command line, in a clone of the GitHub repo: python scripts/generate.py "<prompt>" --n 4. PROMPTING.md lists the caption format and example prompts.

ComfyUI. ComfyUI-Visage adds loader, text-encode, sampler and decode nodes, with an example workflow. Install it from ComfyUI Manager (search "Visage") or with comfy node install comfyui-visage.

folder	contents
transformer/	DiT EMA weights, step 222k (220k + 2k adversarial post-training steps), and its config
vae/	VAE step 400k, with per-channel latent statistics, and its config
text_encoder/	the frozen encoder's record: model, layer, caption length and hub revision; its weights come from Qwen3-0.6B

The text states were stored in float8 (e4m3) during training, and the text encoder rounds through float8 the same way, so inference sees exactly what training saw.

Showcase

Most sheets vary one caption slot from row to row; the short sheets use short prompts. The full prompt is printed on the left of each row. Four seeds per row, shared across rows, with the default sampler (SDE, churn 0.3, CFG 3, 50 steps).

expression

age

turn gaze

short common

More sample sheets

hair

eyes

gaze

turn

pitch

skin

photo

shot

background

extras

short

Evaluation

VAE vs FLUX.2 VAE

The Visage-Image VAE is f16×96, or 0.375 latent floats per pixel. FLUX.2 is f8×32, or 0.5 floats per pixel, so the Visage-Image VAE works at 25% lower rate. All numbers are at 256px.

metric	Visage-Image VAE	FLUX.2 VAE
whole-image PSNR (512 held-out images)	38.69	38.67	tie (+0.01 ± 0.05 dB)
whole-image MAE	0.00802	0.00837	Visage-Image; better on 74% of images
eye-crop PSNR, paired (990 crops)	−0.02 dB	0 (reference)	tie
iFID ↓ (10,016 images)	17.96	21.99	Visage-Image
rFID ↓ (10,016 images)	1.77	0.93	FLUX.2
parameters	128.4M	84.0M

The eye-crop metric is the paired PSNR difference on tight eye crops from images neither VAE was trained on. That's where fine face detail fails first and where whole-image metrics are least sensitive.

iFID is FID on images decoded from the midpoint between a latent and its nearest neighbour (arXiv:2603.05630). It measures latent-space smoothness. It does not predict how fast a DiT learns: in short DiT screens trained on each latent, FLUX.2's latent still reaches a lower generation FID (101.3 vs 190.7 for this VAE), despite losing on iFID.

Not at 512px. The decoder runs at 512, but the quality doesn't carry over: −3.2 dB on eye crops and iFID 28.6 vs 25.5 against FLUX.2. At 128px, rFID rises to 5.46.

Other properties

• Deterministic latent: there's no KL term and no sampling. encode returns the posterior mean, so there's no noise tax on reconstruction.
• DiT-friendly grid: f16 gives 256 tokens at 256px with patch size 1, so the DiT needs no patchify.
• Pre-normalized latents: per-channel latent statistics are stored in the checkpoint. encode returns zero-mean, unit-std latents and decode accepts them.
• Stable across training: iFID was 17.73 at 50k, 17.71 at 170k, 17.89 at 230k and 17.96 at 400k steps, while eye-crop PSNR went from −2.27 to −0.02 dB against FLUX.2. Earlier recipes lost latent-space quality as reconstruction improved.

DiT

• Sharp anneal: ties a casual-prompt fine-tune of the 200k model by eye overall and gives better eyes.
• Bare prompts default to sharp. Sharp training captions carry no sharpness word, and caption dropout draws only from sharp images, so an unspecified prompt defaults to sharp. An earlier internal model needed a "Very sharp," prefix; this one doesn't.
• Adversarial post-training is a clear win by eye. 222k beat 220k 16–0 (4 ties) in a blind A/B on 20 random training-style captions, same seed per pair. The difference is sharpness; conditioning sheets for every prompt slot came out the same. Running it longer (224k, 226k) started to change layout and identity, and the held-out flow loss rose only at high noise, so 222k is the release.
• Held-out FID (2,048 held-out images, default sampler): 22.04 at 222k vs 23.46 at 220k; KID ×10³ 5.15 vs 6.04 (std 0.3). The adversarial step improved both, by more than the checkpoint-to-checkpoint noise (0.25 FID).
• Attribute benchmark (28 targets, SigLIP-scored): 77.0% mean pass at 222k (75.8% at 220k). Known weak spots: sunglasses (35%, often drawn pushed up on the head), angry (38%, rendered as yelling rather than a furrowed brow), wet hair (46%).

Training

• VAE: Charbonnier + 3·LPIPS + 0.4·DINOv2 alignment (Qwen-Image-VAE-2.0 losses), no KL, no GAN. The alignment margins loosen over 200k steps to 0.2 of their slack, then hold. From step 80k, a phase-template penalty removes a faint fixed 2×4 px hatch from the decoder's patch head. Muon on the decoder, horizontal flips, batch 8, 400k steps.
• DiT:
  • rectified flow, logit-normal timesteps with a resolution shift, 20% caption dropout
  • Muon + AdamW, EMA 0.999 every 10 steps, batch 128 (64 × 2 accumulation); step counts, checkpoint names included, count batches of 64
  • 200k steps at 256px from scratch, with long, short and casual prompts mixed 50/50 from step 0
  • then a 20k-step cosine anneal on the sharp and very-sharp images only
  • then 2k steps of latent adversarial post-training (LADD-style): a critic made of small trainable heads on blocks 5, 11 and 17 of a frozen copy of the 220k DiT, run on the noised, mean-centred one-step prediction, at weight 0.006 next to the flow loss, flat LR, same sharp subset (configs/dit_adversarial.yaml).
• Compute: one NVIDIA RTX 5090 (32 GB), about 23 GPU-hours in total.
  • VAE: about 10.5 hours (400k steps at about 10.7 it/s, batch 8)
  • DiT: about 12.5 hours (11 hours for the 200k base run at about 5.1 batches of 64 per second, or about 330 images/s, plus 1.1 hours for the anneal and about 0.3 hours of adversarial post-training), on cached latents and text states

Long captions are structured descriptions of the subject: age, pose, gaze, expression, hair, lighting, background. They are rendered from per-image attributes, several of them measured by dedicated models (head pose, gaze, eye closure, age) rather than read off by a VLM. The remaining slots come from Qwen3.5-4B fine-tuned on hand-corrected labels to write the caption format. DATA.md has the full pipeline.

Limitations and intended use

• Scope: trained on face crops at 256px only. It is not a general text-to-image model, and the VAE does not hold up at 512px.
• Real people: captions contain no names, and each identity was capped at 1000 images per training epoch. Prompting with the name of a public figure does not reproduce their face: each seed gives a different person, as it does with no name in the prompt.
• Out-of-scope uses: impersonation, non-consensual intimate imagery, harassment, deceptive media, and any sexual content depicting minors.
• Prompt adherence: pose, gaze and expression follow the prompt less reliably than appearance attributes do. See the attribute benchmark above for known weak spots. For tongues, write "tongue out" or "sticking tongue out"; "tongue visible" mostly fails. "Biting lip" does not work: no training caption describes it.

License

The code and the weights are released under Apache-2.0.