File size: 6,391 Bytes
4dd94ee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
99d9a77
4dd94ee
6871462
4dd94ee
99d9a77
4dd94ee
99d9a77
adf9175
99d9a77
4dd94ee
 
 
 
adf9175
4dd94ee
 
 
99d9a77
 
 
 
 
4dd94ee
99d9a77
 
 
 
 
adf9175
99d9a77
4dd94ee
99d9a77
4dd94ee
99d9a77
 
 
 
 
 
 
 
 
4dd94ee
6871462
adf9175
99d9a77
adf9175
 
4dd94ee
99d9a77
 
 
 
 
 
 
 
4dd94ee
99d9a77
4dd94ee
99d9a77
4dd94ee
6871462
4dd94ee
 
 
99d9a77
 
 
adf9175
99d9a77
4dd94ee
 
 
 
 
99d9a77
4dd94ee
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-SFT
base_model_relation: adapter
library_name: mlx
pipeline_tag: text-to-image
language:
- en
- zh
tags:
- mlx
- apple-silicon
- moss-image
- text-to-image
- lora
- distillation
- few-step
---

# MOSS-Image-SFT Turbo for Apple Silicon (MLX)

A fast sampler for [OpenMOSS-Team/MOSS-Image-SFT](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT) on Macs. It uses **15 network passes instead of 100**: about **2 minutes** per ~1 MP image on an M4 Pro instead of about 12½ (**6.2×**). It stays close to the image the full 50-step model draws for the same seed, including its textures and lettering.

**Code:** [moss-image-mlx](https://github.com/felix-ab/moss-image-mlx) (`--turbo`) · **Write-up:** [how I designed and measured it](https://github.com/felix-ab/moss-image-mlx/blob/main/research/turbo/docs/technical-writeup.md) · **Base tiers:** [8bit](https://huggingface.co/VisualInference/MOSS-Image-SFT-MLX-8bit), [compact](https://huggingface.co/VisualInference/MOSS-Image-SFT-MLX-compact)

![Wet river pebbles at actual size: the full model at 50 steps above, --turbo below, same prompt and seed](assets/hero.jpg)

*Both crops show the same image region at actual size, from the same prompt and seed. The prompt was not in the training set.*

This repository holds only the trained LoRA: rank 64 on the SFT transformer (252 projections, fp32, 595 MB). The package downloads it for you.

## Quickstart

```bash
pip install -U moss-image-mlx     # 0.2.0 or newer
moss-mlx "A highway LED sign at sunset that reads \"HELLO MLX\"" --turbo -o sign.png
```

- **First run:** the package downloads the 8-bit MOSS-Image tier, this LoRA (pinned to a fixed revision) and the stripe-removal LoRA, then caches them.
- **Sizes:** I trained it at ~1 MP in 9 aspect ratios, from 768×1360 through 1024×1024 to 1360×768. Set the size with `--width` and `--height`.
- **More images:** `--seeds 1,2,3` renders several seeds with one model load; `--batch jobs.jsonl` renders many prompts.
- **Ignored in turbo mode:** `--steps` and `--shift`. `--turbo-weights FILE` uses a local copy of the LoRA.
- **Hardware:** an Apple Silicon Mac. I measured everything on the 8bit tier, the default on Macs with ≥ 24 GB.

## Speed

![Time per image on an M4 Pro: 12 min 23 s for the full model, 1 min 59 s for --turbo](assets/speed.png)

Each network pass takes as long as before; `--turbo` simply needs 15 instead of 100. Peak memory is 9.4 GB (9.3 GB for the full model).

## Results

I compared `--turbo` with the full model's own image (50 steps, CFG 3, with the stripe-removal LoRA) on **32 prompts that were not in the training set**, each at its own aspect ratio, seed 101:

| | `--turbo` | Full model, 50 steps |
|---|---|---|
| Network passes | **15** | 100 |
| Time per ~1 MP image (M4 Pro, 8bit) | **~119 s** | ~743 s |
| Average pixel difference from the full model's image (0–1 scale) | 0.090 | 0 |
| Repeated-texture score, relative to the full model (materials and nature prompts, n = 8) | −0.006 (no extra repetition) | 0 |
| Fine detail (high-frequency energy), relative to the full model | −0.004 | 0 |
| Brightness, relative to the full model | +0.011 | 0 |
| English text read back correctly by OCR (n = 4) | 94.7% | 93.4% |

The package's sampler reproduces the research images bit for bit. [how-the-numbers-work.md](https://github.com/felix-ab/moss-image-mlx/blob/main/research/turbo/docs/how-the-numbers-work.md) defines each measure.

![Four pairs, full model left and --turbo right, same seed: a beekeeper, a novel cover, a Chinese tea wrapper and a croissant](assets/pairs.jpg)

## How it works

![The official sampler's 100 passes vs the 15 of --turbo](assets/schedule.png)

1. **A small LoRA sets the composition** in 5 steps without CFG. I trained it on the full model's own sampling paths for about 2,000 prompts (English and Chinese text, layouts, people, products, places, materials). When the LoRA drifted off a saved path during training, I reran the full model from where the LoRA was and trained on that result. Pulling it straight back to the saved path instead produced double exposures.
2. **The full model takes one step with CFG** among those early steps (the third), where the shapes of large letters form. I tried that step at each of the 6 positions; only the third repaired both a garbled shop sign and an infographic title.
3. **The full model draws the fine detail** in 4 DPM-Solver++(2M) steps, with no distilled steps.
4. **Every pass applies [nynxz's stripe-removal LoRA](https://huggingface.co/nynxz/MOSS-Image)** at 1.5. I applied it to the training data too.

Training cost about $37 on rented A100s.

**Why not distill every step?** I took the full model's saved intermediate states and finished them with a DMD-distilled LoRA, a common way to build few-step models. The more DMD steps it took in the middle of sampling, the more it turned natural textures into repeating patterns:

![Lichen on a stone wall: the full model, then the full model's saved states after 42, 35 and 14 steps finished with 2, 3 and 6 DMD steps](assets/repetition_test.jpg)

So this LoRA learns to follow the full model's own path for each seed, and the full model finishes the image. The [technical write-up](https://github.com/felix-ab/moss-image-mlx/blob/main/research/turbo/docs/technical-writeup.md) has the evidence for each decision and what failed; the [research repo](https://github.com/felix-ab/moss-image-mlx/tree/main/research/turbo) has the full lab notebook.

## Limitations

- **Small, dense text** has a few more errors than with the full model (OCR 95.7% vs 97.0% on 8 prompts).
- **Regular patterns**, like a chain-link fence shadow or rows of machines, can come out irregular.
- **Compositions** usually match the full model for the same seed, but not always.
- **Untested:** sizes outside ~1 MP and the 9 trained aspect ratios; the compact and bf16 tiers with `--turbo`; the Base checkpoint (the CLI refuses `--transformer` in turbo mode).
- **Evidence size:** each comparison uses 1–2 seeds and 7–32 prompts, and I checked the images by eye. I did not compute FID or KID at scale.

## License and credits

Apache-2.0.
- **Base model:** MOSS-Image SFT, © OpenMOSS, Apache-2.0.
- **Stripe-removal LoRA:** by nynxz, Apache-2.0.
- **This work** is independent and not affiliated with OpenMOSS.