How to use from the
Use from the
Diffusers library
pip install -U diffusers transformers accelerate
import torch
from diffusers import DiffusionPipeline

# switch to "mps" for apple devices
pipe = DiffusionPipeline.from_pretrained("Gabriel/diffu-moss-r6b", dtype=torch.bfloat16, device_map="cuda")

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(prompt).images[0]

diffu-moss-r6b

A full fine-tune of the DiT of MOSS-Image-Base that draws whole synthetic archive pages: Swedish handwritten records of the 17th to 19th centuries. It is the drawing model of diffu, which makes training pages for page-level handwriting recognition: a page layout (a PAGE/ALTO file, or the lines' boxes as JSON), the text to write in it (or the layout's own) and a sample page go in; a page image and its exact ALTO and PAGE labels come out.

The model reads a MOSS Layout prompt, one text item per line with its exact text and its box, and a sample page given as extra image tokens, whose paper, ink and look it follows. Only the DiT is in this repository. The text encoder, tokenizer, VAE and scheduler are MOSS-Image-Base's, unchanged.

Files

File What
diffusion_pytorch_model.safetensors the DiT's weights, bfloat16 (10.2 GB, 462 tensors, the same names and shapes as Base's DiT)
config.json diffu_moss.transformer.MossTransformer (diffu's own implementation of MOSS-Image's DiT)
export.json takes_reference: true: the model takes a sample page

Use

With diffu (see its docs/INFERENCE.md for running on a Ray or KubeRay cluster):

hf download OpenMOSS-Team/MOSS-Image-Base --revision 85975793d8b3d432f1b6e27f352abaa4c2a4c0d3 --local-dir base
hf download Gabriel/diffu-moss-r6b --local-dir dit
git clone https://github.com/Borg93/diffu2 && cd diffu2 && uv sync --all-packages
uv run diffu-moss-init --base ../base --transformer ../dit --out ../moss-r6b
uv run diffu-run --jobs run.jsonl --draw-settings draw.json --model ../moss-r6b --out out --gpus 1

run.jsonl holds one page job per line, as JSON: the layout (template, a PAGE/ALTO file, or layout, the lines' boxes in pixels), the text (left out, the layout's own), the style, the sample page and an id. diffu's README.md shows two lines; docs/INFERENCE.md section 5 describes every field, with a worked example of each combination.

On Blackwell GPUs the DiT can run in MXFP8, quantized from these weights as the model loads: give the draw settings {"precision": "mxfp8", "compile_blocks": true}. On a GB10 at 1024 x 1024 this draws at 1.38 s a step, against 1.76 s for bfloat16 compiled and 3.46 s uncompiled; its pages look like bfloat16's, slightly less saturated.

For SGLang Diffusion, diffu-moss-export-hf writes this DiT in the official MOSS-Image layout, and diffu's plugin diffu_sglang_moss adds the sample page (docs/INFERENCE.md section 13).

Recommended drawing: about 1 MP, 50 steps, guidance 3, shift 5, block 22 skipped in the unconditional pass (MOSS-Image's own recipe).

Training

  • Data: 3,406 page pairs balanced across eight collections of Swedish archival handwriting ground truth (court, chancery, mining and travel records): 2,031 whole pages and 1,375 windows of pages, about 1 MP each, each paired with a sample page from the same volume.
  • Prompt forms: no sample, the volume named (20 %); the sample with the look in words (20 %); the sample with the volume named (20 %); the sample alone (40 %). In the last 1,000 steps half of the sample-alone pairs were restyled the same way on page and sample (microfilm grey, colourised, tinted), so the look has to come from the sample.
  • Recipe: MOSS-Image's own full fine-tune: Muon on the joint blocks' matrices and AdamW elsewhere, learning rate 5e-5, weight decay 0.01, gradients clipped at 1, one pair a step, float32 weights under bfloat16 autocast; text encoder and VAE frozen.
  • Steps: 1,500 from Base on the sample form, then 1,000 with the restyled pairs (2,500 in all), on one DGX Spark.

Evaluation

On held-out pages of the training volumes, each drawn with its own sample and with samples of other volumes, from the same noise:

Check Result
The page's own sample gives a page of the right kind (colour or grey) 4 of 5
A grey sample makes a colour page grey 3 of 5
Colour and grey changes together 78 %
The sample's handwriting is followed instead of the volume's 1 of 12

Limitations

  • The handwriting does not follow the sample page. Trained on eight volumes, the model recognises a volume from its text and layout and writes that volume's hand. Training on many more volumes is being tested.
  • The look follows the sample only in part (above).
  • Every line has to be checked: diffu reads each drawn line back and marks a line it cannot confirm as unclear in the labels.
  • Only eight volumes' hands and layouts were trained; other kinds of pages are untested.

Licence

Apache-2.0, as MOSS-Image-Base.

Downloads last month
7
Safetensors
Model size
5B params
Tensor type
BF16
ยท
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support

Model tree for Gabriel/diffu-moss-r6b

Finetuned
(46)
this model