diffu-moss-r6b: full fine-tune of MOSS-Image-Base's DiT, step 750 of r6b

README.md

@@ -0,0 +1,98 @@
+---
+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 (from a PAGE/ALTO file), the text to write in it 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 (template, text, style, sample page, id); `docs/INFERENCE.md` section 5
+describes it with worked examples.
+
+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.

config.json

@@ -0,0 +1,12 @@
+{
+  "_class_name": "MossTransformer",
+  "_diffusers_version": "0.42.0.dev0",
+  "dim": 3072,
+  "ffn_dim": 8192,
+  "freq_dim": 256,
+  "latent_channels": 128,
+  "num_heads": 24,
+  "num_layers": 42,
+  "rope_base": 10000.0,
+  "text_dim": 4096
+}

diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2cc4c67a4fbce820dd5743a8fbc345e63ef2d5c574e29dc73c3f6550963c2894
+size 10164664128

export.json

@@ -0,0 +1,6 @@
+{
+ "takes_reference": true,
+ "base": "OpenMOSS-Team/MOSS-Image-Base@85975793d8b3d432f1b6e27f352abaa4c2a4c0d3 (full fine-tunes r5 1500 steps, r6 250, r6b 750)",
+ "checkpoint": "full_r6b/step_000750",
+ "adapter": null
+}