Metadata-Version: 2.4
Name: visage-image
Version: 1.0.0
Summary: A 256px text-to-image model for face generation: single-stream rectified-flow DiT over an f16x96 VAE latent.
License-Expression: Apache-2.0
Project-URL: Repository, https://github.com/Nikita-Sherstnev/Visage-Image
Project-URL: Weights, https://huggingface.co/Sherstnev/Visage-Image
Project-URL: Prompting, https://github.com/Nikita-Sherstnev/Visage-Image/blob/main/PROMPTING.md
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.9
Requires-Dist: torchvision
Requires-Dist: transformers>=4.56
Requires-Dist: huggingface-hub
Requires-Dist: safetensors
Requires-Dist: pillow
Provides-Extra: train
Requires-Dist: pyyaml; extra == "train"
Requires-Dist: timm; extra == "train"
Requires-Dist: torchmetrics; extra == "train"
Provides-Extra: data
Requires-Dist: opencv-python; extra == "data"
Requires-Dist: ultralytics; extra == "data"
Provides-Extra: dev
Requires-Dist: ruff; extra == "dev"
Dynamic: license-file

# Visage-Image

**[Model card and weights](https://huggingface.co/Sherstnev/Visage-Image)** · **[Architecture](https://github.com/Nikita-Sherstnev/Visage-Image/blob/main/ARCHITECTURE.md)** · **[Prompting](https://github.com/Nikita-Sherstnev/Visage-Image/blob/main/PROMPTING.md)** · **[Data](https://github.com/Nikita-Sherstnev/Visage-Image/blob/main/DATA.md)**

![Five 256px samples: smiling, angry, surprised, laughing, tongue out](https://raw.githubusercontent.com/Nikita-Sherstnev/Visage-Image/main/assets/showcase.png)

## Introduction

A 256px text-to-image model for face generation: a 237M-parameter single-stream diffusion transformer
over the latents of a 128M-parameter f16 autoencoder (16×16×96 latents), conditioned on
layer 8 of Qwen3-0.6B.

This model started as an experiment in training a DiT on data from a limited domain. With so many models available for face detection and for tasks like age and gaze estimation, preparing a custom dataset was fairly simple, time-consuming, but simple. I first tried the FLUX.2 VAE, but it turned out that the plastic look is an inherent vice of that model, it turns anything it touches into plastic. It is optimized for learnability, its latent channels have a lot of redundancy, and it fakes details with GAN-style generation, which makes them look plausible rather than faithful. So I trained my own autoencoder using the latest ideas and recipes, and it turned out quite well. Qwen3-0.6B makes the model robust to variations in prompt wording, and unlike most small general-purpose DiTs, it generates quite good faces. In fact, to train a very specific DiT you don't need a big dataset: a set of similar images that differ only in details is enough, because a DiT interpolates between them very well. The first version of this model was like that. I found a "good region" densely populated with similar images. Samples from it were very realistic, and outside it they were a mess.

The [model card](https://huggingface.co/Sherstnev/Visage-Image) has the highlights, samples and evaluation. [ARCHITECTURE.md](https://github.com/Nikita-Sherstnev/Visage-Image/blob/main/ARCHITECTURE.md) covers each design choice and what was tried instead.

## Repository layout

| path | what |
|---|---|
| `visage/dit.py` | DiT: 18 single-stream blocks, width 768, adaLN timestep conditioning, Qwen3-0.6B captions with learned bucket pads and 16 registers, spiral 2D RoPE |
| `visage/vae.py` | VAE: Haar-wavelet conv encoder with self-attention at f8 and f16 → 96ch f16 latent; two-stage ViT decoder (f16 → f8) + 8×8 patch head + two residual conv refiners |
| `visage/sampling.py` | Qwen3 text encoder (first 8 layers), rectified-flow sampler (Euler / SDE churn, CFG) |
| `visage/loading.py` | `load_dit` / `load_vae` / `load_text_encoder` from a model directory or hub repo id |
| `visage/training.py` | manifest dataset, config loading, Muon + AdamW, LR schedule |
| `visage/adversarial.py` | latent critic for adversarial post-training: trainable heads on a frozen copy of the DiT |
| `scripts/generate.py` | text-to-image CLI |
| `scripts/make_showcase.py` | rebuilds `assets/showcase.png`; holds each tile's prompt and seed |
| `PROMPTING.md` | long-caption format, vocabulary and example prompts |
| `DATA.md` | how the images were curated and captioned, and which model labels each slot |
| `scripts/crop_faces.py` | square head crops from a photo folder, plus a `train.jsonl` manifest |
| `scripts/train_dit.py`, `scripts/train_vae.py` | training loops; configs in `configs/` (`dit_adversarial.yaml` is the release's last stage) |

## Quick Start

```bash
pip install -e .            # add [train] for training: pip install -e ".[train]"
python scripts/generate.py "a young woman with long red hair and freckles, soft window light" --n 4
```

The weights download from the [Hugging Face repo](https://huggingface.co/Sherstnev/Visage-Image) on first use. `--model` also takes a local copy:
`huggingface-cli download Sherstnev/Visage-Image --local-dir weights`, then `--model weights`. The loaders read one folder per
component: `transformer/` and `vae/` hold `config.json` + `model.safetensors`, and `text_encoder/config.json` names
the frozen Qwen3 encoder. The root `config.json` records the sampler settings tuned for the weights.

To use the model from Python without the repo scripts, install the package from PyPI:

```bash
pip install visage-image
```

```python
from visage import load_dit, load_sampling, load_text_encoder, load_vae, sample

text_encoder, dit, vae = load_text_encoder(), load_dit(), load_vae()
text, mask = text_encoder(["a young woman with long red hair and freckles, soft window light"])
latents = sample(dit, text, mask, latent_size=dit.config.image_size, seed=0, **load_sampling())
images = vae.decode(latents)  # (N, 3, 256, 256) in [-1, 1]
```

Recommended sampler settings, which the scripts read from that file: `--solver sde --churn 0.3 --cfg 3 --steps 50 --shift 3.3`.
- **Euler**: crisper but less realistic.
- **churn**: applied every step, so raising `--steps` without lowering churn washes the image out.
- **shift**: depends on the token count. Retune it if you change resolution.
- **cfg**: depends on the caption-dropout rate.

Detailed structured captions follow best, but
short and casual prompts were trained from step 0 and work. Bare prompts default to sharp; no sharpness word is needed.
[PROMPTING.md](https://github.com/Nikita-Sherstnev/Visage-Image/blob/main/PROMPTING.md) lists the long-caption slots in order, with common values and examples.
256px only: the VAE loses 3.2 dB on eye crops at 512.

The strip at the top is `python scripts/make_showcase.py`: default sampler settings, with the
prompts and seeds listed in the script, one image per seed. Any tile is also
`python scripts/generate.py "<prompt>" --n 1 --seed <seed>`.

## Train

The data is a JSONL manifest with one `{"image": "path.jpg", "caption": "..."}` per line.
Relative paths are resolved against the manifest's directory. The VAE ignores captions.
`scripts/crop_faces.py` builds one from a folder of photos with the same detector and head
margins as the release data (YOLO face weights, AGPL-3.0; the download line is in the script's
`--help`). Captions come out empty and have to be filled in before DiT training.

```bash
pip install -e ".[data]"
python scripts/crop_faces.py photos/ -o data/          # -> data/images/..., data/train.jsonl
python scripts/train_vae.py configs/vae.yaml   # -> runs/vae/vae/ (+ latent stats)
python scripts/train_dit.py configs/dit.yaml   # uses the VAE in `vae:` -> runs/dit/transformer/ (EMA)
```

**VAE**
- Loss: Charbonnier + 3·LPIPS(VGG) + 0.4·DINOv2 alignment. The alignment is mcos + mdms on a middle-layer feature, with margins that loosen over 200k steps to 0.2 of their full slack, then hold.
- From step 80k, a phase-template term removes the decoder's fixed 8px sub-pixel pattern.
- No KL and no GAN. Horizontal flips. 400k steps at batch 8.
- Optimizer: Muon on the decoder's hidden matrices, AdamW for the rest. Cosine LR schedule.

**DiT**
- Rectified flow with velocity target and logit-normal timesteps, shifted by √(tokens/64).
- 20% caption dropout to a zero embedding. The release drew dropout from sharp images only and
  omitted the sharpness word on sharp captions, so an unspecified prompt defaults to sharp; this
  loop drops uniformly.
- Muon + AdamW, EMA 0.999 every 10 steps, batch 64×2.
- The release then ran `configs/dit_adversarial.yaml`: 1000 updates of latent adversarial post-training on the
  anneal checkpoint (see [ARCHITECTURE.md](https://github.com/Nikita-Sherstnev/Visage-Image/blob/main/ARCHITECTURE.md#training)).

Checkpoint names (200k, 220k, 222k) count batches of 64, as the training repo did. `scripts/train_dit.py`
counts optimizer updates of 64×2, so 222k there is step 111,000.

The DiT loop encodes images and captions on the fly, which keeps the code minimal. For a
long run, cache the latents and text states once instead.

**Compute.** The release was trained on one NVIDIA RTX 5090 (32 GB):

| model | steps | throughput | wall-clock |
|---|---|---|---|
| VAE | 400k at batch 8 | ~10.7 it/s | ~10.5 h |
| DiT, base run | 200k batches of 64 | ~5.1 batches/s (~330 images/s) | ~11 h |
| DiT, sharp anneal | 20k batches of 64 | ~5.1 batches/s | ~1.1 h |
| DiT, adversarial | 2k batches of 64 | ~2 batches/s | ~0.3 h |
| **total** | | | **~23 GPU-hours** |

The DiT figures are with latents and text states cached in advance.

Lint and format with `ruff check . && ruff format .` (config in `pyproject.toml`).

## License

The code and the weights are [Apache-2.0](https://github.com/Nikita-Sherstnev/Visage-Image/blob/main/LICENSE).

## References

**DiT**
- Rectified flow: Liu et al., *Flow Straight and Fast* — [arXiv:2209.03003](https://arxiv.org/abs/2209.03003); Lipman et al., *Flow Matching* — [arXiv:2210.02747](https://arxiv.org/abs/2210.02747)
- Logit-normal timesteps, resolution-dependent shift: Esser et al., *Scaling Rectified Flow Transformers* (SD3) — [arXiv:2403.03206](https://arxiv.org/abs/2403.03206)
- DiT, adaLN-Zero: Peebles & Xie — [arXiv:2212.09748](https://arxiv.org/abs/2212.09748)
- Text and image tokens in one sequence: Bao et al., *All are Worth Words* (U-ViT) — [arXiv:2209.12152](https://arxiv.org/abs/2209.12152); joint text–image attention: SD3 (above); single-stream blocks at scale: FLUX.1 — [github.com/black-forest-labs/flux](https://github.com/black-forest-labs/flux)
- Spiral RoPE: Liu et al., *Rotate Your Rotary Positional Embeddings in the 2D Plane* — [arXiv:2602.03227](https://arxiv.org/abs/2602.03227)
- RoPE: Su et al., *RoFormer* — [arXiv:2104.09864](https://arxiv.org/abs/2104.09864)
- Register tokens: Darcet et al., *Vision Transformers Need Registers* — [arXiv:2309.16588](https://arxiv.org/abs/2309.16588)
- QK-norm: Dehghani et al., *Scaling Vision Transformers to 22B* — [arXiv:2302.05442](https://arxiv.org/abs/2302.05442)
- SwiGLU: Shazeer, *GLU Variants Improve Transformer* — [arXiv:2002.05202](https://arxiv.org/abs/2002.05202)
- Text encoder: Qwen3 Technical Report — [arXiv:2505.09388](https://arxiv.org/abs/2505.09388)
- Classifier-free guidance: Ho & Salimans — [arXiv:2207.12598](https://arxiv.org/abs/2207.12598)
- SDE churn sampler: Karras et al., *EDM* — [arXiv:2206.00364](https://arxiv.org/abs/2206.00364)

**VAE**
- Losses (no KL/GAN, DINOv2 mcos + mdms alignment), global skip: *Qwen-Image-VAE-2.0* — [arXiv:2605.13565](https://arxiv.org/abs/2605.13565)
- Wavelet encoder: Sadat et al., *LiteVAE* — [arXiv:2405.14477](https://arxiv.org/abs/2405.14477)
- Space-to-channel residual downsampling: Chen et al., *DC-AE* — [arXiv:2410.10733](https://arxiv.org/abs/2410.10733)
- Alignment target: Oquab et al., *DINOv2* — [arXiv:2304.07193](https://arxiv.org/abs/2304.07193)
- Perceptual loss: Zhang et al., *LPIPS* — [arXiv:1801.03924](https://arxiv.org/abs/1801.03924)
- LayerScale: Touvron et al., *CaiT* — [arXiv:2103.17239](https://arxiv.org/abs/2103.17239)
- iFID: [arXiv:2603.05630](https://arxiv.org/abs/2603.05630)

**Optimizer**
- Muon: Jordan et al. — [kellerjordan.github.io/posts/muon](https://kellerjordan.github.io/posts/muon/); RMS-matched LR: Liu et al., *Muon is Scalable* — [arXiv:2502.16982](https://arxiv.org/abs/2502.16982)
