VisualInference's picture
Card: research notes moved to moss-image-mlx/research/turbo
6871462 verified
|
Raw History Blame Contribute Delete
6.39 kB
---
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.