Metadata-Version: 2.5
Name: mflux-lora
Version: 0.1.0
Summary: One-command LoRA training, sampling, and evaluation for mflux on Apple Silicon
Author: mflux-lora contributors
License: MIT
License-File: LICENSE
Keywords: apple-silicon,fine-tuning,flux,image-generation,lora,mflux,mlx
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: MacOS
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Multimedia :: Graphics
Requires-Python: >=3.10
Requires-Dist: mflux<0.24.0,>=0.23.0
Requires-Dist: pillow<13.0,>=12.3.0
Requires-Dist: rich>=13.0.0
Requires-Dist: typer>=0.12.0
Provides-Extra: dev
Requires-Dist: pytest-cov>=5.0; extra == 'dev'
Requires-Dist: pytest>=8.0; extra == 'dev'
Description-Content-Type: text/markdown

# mflux-lora

One-command LoRA training, sampling, and evaluation for [mflux](https://github.com/mflux-community/mflux) on Apple Silicon.

`mflux-lora` is a thin, opinionated wrapper around mflux. It prepares a dataset, writes a safe training config, calls `mflux-train`, extracts the adapter, and gives you a repeatable evaluation grid. It does not reimplement the diffusion models: all inference and training math comes from mflux/MLX.

> Unofficial community tool. Not affiliated with the mflux maintainers.

## Why

mflux exposes powerful primitives, but a first LoRA run still means hand-writing a training JSON, wiring captions, choosing LoRA targets, and finding the checkpoint afterwards. `mflux-lora` turns that into:

```bash
uvx mflux-lora train ./photos --trigger ohwx --preset balanced
uvx mflux-lora sample --lora lora-run/adapter.safetensors --prompt "ohwx, a portrait photo"
uvx mflux-lora eval --lora lora-run/adapter.safetensors --trigger ohwx
```

## Install

No install needed with `uvx`:

```bash
uvx mflux-lora --help
```

Or install as a tool:

```bash
uv tool install mflux-lora
mflux-lora --help
```

Requirements:

- macOS on Apple Silicon
- Python 3.10+
- 32 GB unified memory minimum (48 GB recommended)
- Free disk: ~10 GB for q8 FLUX.2 4B weights, plus checkpoints

The first run downloads model weights from Hugging Face. Set `HF_ENDPOINT=https://hf-mirror.com` if direct access is slow, or pass `--model-path` to reuse a local model directory.

## Quick start

```bash
# 1. Check the machine
uvx mflux-lora doctor

# 2. Train (images + optional .txt captions in ./photos)
uvx mflux-lora train ./photos \
  --trigger ohwx \
  --preset balanced \
  --output ./lora-run

# 3. Sample
uvx mflux-lora sample \
  --lora ./lora-run/adapter.safetensors \
  --prompt "ohwx, a close-up portrait photo, natural light" \
  --size 512x768 \
  --seed 42 \
  --output ./sample.png

# 4. Evaluate angles / seeds / scales
uvx mflux-lora eval \
  --lora ./lora-run/adapter.safetensors \
  --trigger ohwx \
  --output-dir ./eval
```

`eval` writes one image per prompt × seed × scale plus `eval/contact.jpg`.

Use `--dry-run` on `train` or `demo` to inspect the plan and generated config without loading a model.

## Dataset

Point `train` at a folder:

```text
photos/
  01.jpg
  01.txt
  02.png
  02.txt
```

- Images: `.jpg`, `.jpeg`, `.png`, `.webp`
- Caption: a sibling `.txt` with the same stem
- Missing captions are generated from `--caption-template` and prefixed with `--trigger`
- The trigger is prepended to existing captions if it is not already the first word
- By default images are copied to `<output>/dataset`; use `--in-place` to write captions next to the source images

Recommended for identity LoRA:

| Item | Recommendation |
|---|---|
| Count | 20–60 images (80–150 for very high fidelity) |
| Angles | 25% front, 40% three-quarter, 20% profile, 15% other |
| Resolution | 768–1024 px longest side |
| Consistency | Same person, similar age/makeup; avoid heavy retouching |
| Captions | Describe angle, expression, outfit, light; keep the trigger token first |

For angle control, write captions such as:

```text
ohwx, front view portrait, neutral expression, soft studio light, white shirt
ohwx, three-quarter view portrait, looking left, outdoor daylight, black coat
ohwx, side profile view portrait, window light, no accessories
```

This lets the model bind identity to `ohwx` instead of memorizing one pose.

Inspect a dataset without training:

```bash
uvx mflux-lora dataset validate ./photos
uvx mflux-lora dataset sheet ./photos --output ./sheet.jpg
uvx mflux-lora dataset caption ./photos --trigger ohwx
```

## Supported models

| Key | Train on | Sample on | Weights |
|---|---|---|---|
| `flux2-klein-base-4b` | FLUX.2 klein base 4B | same | bf16, q8, q5, q4 |
| `flux2-klein-base-9b` | FLUX.2 klein base 9B | same | bf16, q8 |
| `z-image-turbo` | Z-Image Turbo | same | bf16, q8 |
| `z-image` | Z-Image base | same | bf16 |
| `ernie-image-turbo` | ERNIE-Image Turbo | same | bf16, q8 |
| `ernie-image` | ERNIE-Image base | same | bf16 |
| `krea-2-raw` | Krea 2 Raw | Krea 2 Turbo | bf16, q8 |

Default: `flux2-klein-base-4b` with `q8` (≈8.6 GB download). Use `--weights bf16` to download the full checkpoint and quantize at load.

```bash
uvx mflux-lora models
```

## Presets

| Preset | Resolution | Rank | Targets | Epochs |
|---|---|---|---|---|
| `fast` | 512 | 8 | attention | 10 |
| `balanced` | 512 | 8 | attention | 20 |
| `quality` | 768 | 16 | attention + MLP/FF | 20 |
| `max` | 1024 | 32 | attention + MLP/FF | 30 |

Override any of them:

```bash
uvx mflux-lora train ./photos --preset balanced --resolution 768 --rank 16 --epochs 30
```

Models that do not define FF/MLP targets (Z-Image) automatically fall back to attention targets.

## Memory guide

Rough peak memory with q8 weights and gradient checkpointing:

| Mac memory | Suggested |
|---|---|
| 32 GB | `flux2-klein-base-4b`, `fast`/`balanced`, 512 px |
| 48 GB | 4B at 768 px, 9B at 512 px |
| 64 GB+ | 9B at 768 px, Krea 2 Raw at 512 px |

Flags for tight machines:

```bash
uvx mflux-lora train ./photos \
  --low-ram \
  --cache-limit-gb 12 \
  --resolution 512
```

`--low-ram` trades speed for memory. `--cache-limit-gb` bounds MLX's buffer cache and is usually enough without the full low-RAM overhead.

## Commands

### train

```bash
uvx mflux-lora train IMAGES [options]
```

Key options: `--model`, `--preset`, `--trigger`, `--weights`, `--output`, `--resolution`, `--rank`, `--epochs`, `--targets`, `--lr`, `--batch-size`, `--grad-accum`, `--seed`, `--model-path`, `--low-ram`, `--cache-limit-gb`, `--save-frequency`, `--no-preview`, `--dry-run`, `--in-place`.

The generated config is `<output>/train.json`. Checkpoints are `<output>/train*/checkpoints/*.zip`. The best/latest adapter is copied to `<output>/adapter.safetensors`.

### sample

```bash
uvx mflux-lora sample \
  --lora ./lora-run/adapter.safetensors \
  --prompt "ohwx, a portrait photo" \
  --scale 1.0 \
  --size 512x768 \
  --seed 42 \
  --output ./out.png
```

- `--lora` accepts either an adapter `.safetensors` or a mflux checkpoint `.zip`
- `--no-bake-lora` is the default; q8 baking can drop part of the update
- `--compare-base` writes a base/LoRA side-by-side image
- Use `--model-path` to reuse a local model directory

### eval

```bash
uvx mflux-lora eval \
  --lora ./lora-run/adapter.safetensors \
  --trigger ohwx \
  --seeds 42,123 \
  --scales 0.8,1.0 \
  --size 512x768 \
  --output-dir ./eval
```

Without `--prompts-file`, it generates front / three-quarter / profile / full-body prompts. With `--include-base`, the contact sheet includes a base-model row for comparison.

### Other commands

```bash
uvx mflux-lora dataset validate ./photos
uvx mflux-lora dataset sheet ./photos
uvx mflux-lora doctor
uvx mflux-lora models
uvx mflux-lora demo --dry-run
```

`demo` uses two bundled images and a 2-epoch run to verify the whole pipeline. It still downloads the model weights on the first real run.

## How the config is built

`mflux-lora` writes a normal mflux training config. You can edit it or pass it directly:

```bash
mflux-train --config ./lora-run/train.json
```

The wrapper chooses:

- model-specific scheduler steps / guidance / timestep range
- LoRA targets for the selected model
- `cache_limit_gb`, gradient checkpointing, preview cadence
- a checkpoint every `--save-frequency` steps

If you already have a config, use mflux directly; this tool is for the common path.

## Evaluation and choosing a checkpoint

Do not pick a checkpoint by training loss. Generate the same prompts and seeds for several checkpoints/scales and compare:

```bash
uvx mflux-lora eval --lora ./lora-run/train/checkpoints/0000300_checkpoint.zip --trigger ohwx --scales 0.7,1.0 --output-dir ./eval-300
uvx mflux-lora eval --lora ./lora-run/train/checkpoints/0000500_checkpoint.zip --trigger ohwx --scales 0.7,1.0 --output-dir ./eval-500
```

For quantitative identity checks, run an ArcFace/InsightFace embedding comparison against held-out reference photos. The package intentionally does not bundle a face-recognition model.

If a late checkpoint starts ignoring the prompt (wrong pose, collages), lower the LoRA scale or pick an earlier checkpoint. This is a known overfitting/strength failure mode, not a bug in the wrapper.

## Troubleshooting

**Pillow import error on Python 3.13**

Pillow 13 requires Python 3.14. `mflux-lora` pins `pillow<13`; if you install mflux separately, also pin it:

```bash
uv tool install mflux --with "pillow<13"
```

**Hugging Face downloads are slow**

```bash
export HF_ENDPOINT=https://hf-mirror.com
```

If the mirror is unavailable, direct Hugging Face plus multiple connections is usually faster.

**OOM or the machine starts swapping**

- Lower `--resolution` (512 → 448)
- Add `--low-ram`
- Lower `--cache-limit-gb` (18 → 12 → 8)
- Close other MLX/GPU jobs
- Keep `--gradient-checkpointing` on (default)

**Training is much slower than expected**

Check `vm.swapusage`. If swap is nearly full, another process or session is competing for unified memory. Stop other jobs before restarting training.

**LoRA has little effect**

- Include the trigger token in the prompt
- Use `--no-bake-lora` (default)
- Raise `--scale` from 0.7 to 0.9/1.0
- Make sure the trigger appears first in captions

## Development

```bash
git clone <this repo>
cd mflux-lora
uv sync
uv run pytest
uv run mflux-lora --help
uv build
uv publish
```

## License

MIT. mflux is MIT-licensed; the bundled demo images come from mflux's training example assets. Model weights and generated images are subject to their own licenses and, for real people, to likeness and publicity rights.
