File size: 4,955 Bytes
2246c79
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
726b9fe
 
2246c79
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
726b9fe
 
 
2246c79
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
---
license: apache-2.0
base_model: OpenMOSS-Team/MOSS-Image-Base
base_model_relation: finetune
pipeline_tag: text-to-image
language:
  - sv
tags:
  - moss-image
  - handwriting
  - historical-documents
  - archives
  - htr
  - synthetic-data
---

# diffu-moss-r6b

A full fine-tune of the DiT of [MOSS-Image-Base](https://huggingface.co/OpenMOSS-Team/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](https://github.com/Borg93/diffu2), 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):

```bash
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.