Metadata-Version: 2.4
Name: mlx-diffuser
Version: 0.1.0
Summary: Diffusion training and sampling on Apple silicon with MLX.
Project-URL: Homepage, https://github.com/AmirHossein-razlighi/mlx_diffusion
Project-URL: Documentation, https://AmirHossein-razlighi.github.io/mlx_diffusion/
Project-URL: Repository, https://github.com/AmirHossein-razlighi/mlx_diffusion
Project-URL: Issues, https://github.com/AmirHossein-razlighi/mlx_diffusion/issues
Project-URL: Changelog, https://github.com/AmirHossein-razlighi/mlx_diffusion/blob/main/CHANGELOG.md
Author-email: Amirhossein <arazlighi@gmail.com>
License: Apache-2.0
License-File: LICENSE
Keywords: apple-silicon,diffusion,dit,flow-matching,generative-ai,image-generation,metal,mlx,stable-diffusion
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Operating System :: MacOS
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Requires-Python: >=3.11
Requires-Dist: mlx>=0.31.0
Requires-Dist: numpy>=1.26.0
Requires-Dist: pillow>=11.0.0
Provides-Extra: dev
Requires-Dist: mypy>=1.13.0; extra == 'dev'
Requires-Dist: pytest>=8.3.0; extra == 'dev'
Requires-Dist: ruff>=0.8.0; extra == 'dev'
Provides-Extra: docs
Requires-Dist: mkdocs-material>=9.5.0; extra == 'docs'
Requires-Dist: mkdocstrings[python]>=0.27.0; extra == 'docs'
Provides-Extra: hub
Requires-Dist: huggingface-hub>=0.25.0; extra == 'hub'
Description-Content-Type: text/markdown

# mlx-diffuser

**Diffusion & flow models on Apple silicon, powered by [MLX](https://github.com/ml-explore/mlx).**
Train from scratch, fine-tune, or run inference — for image, video, and discrete
modalities — from one small, readable codebase.

[![CI](https://github.com/AmirHossein-razlighi/mlx_diffusion/actions/workflows/ci.yml/badge.svg)](https://github.com/AmirHossein-razlighi/mlx_diffusion/actions/workflows/ci.yml)
![Python](https://img.shields.io/badge/python-3.11%2B-blue)
![License](https://img.shields.io/badge/license-Apache--2.0-green)

---

If you know PyTorch and 🤗 `diffusers`, you already know this library:
`Model.from_pretrained(...)`, `pipe(...)`, `nn.Module` everywhere. The difference
is underneath — unified memory, `mx.compile`, fused Metal kernels, and built-in
weight quantization, so large models fit and run fast on a Mac.

## Install

```bash
pip install mlx-diffuser          # core
pip install "mlx-diffuser[hub]"   # + Hugging Face Hub loading
```

Requires Apple silicon (M-series) and Python 3.11+.

## Generate

```python
from mlx_diffusion import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("path/or/hub-id", dtype="bf16", quantize=4)
images = pipe([1, 2, 3], num_inference_steps=50, guidance_scale=4.0, seed=0)
```

…or from the terminal:

```bash
mlx-diffuser generate path/or/hub-id --labels 1,2,3 --steps 50 --out samples/
```

## Train from scratch

```python
from mlx_diffusion import DiT, DiTConfig, DiffusionTrainer
from mlx_diffusion.schedulers import FlowMatchEulerScheduler
from mlx_diffusion.training import batch_iterator

model = DiT(DiTConfig(in_channels=3, hidden_size=384, depth=12, num_heads=6))
trainer = DiffusionTrainer(model, FlowMatchEulerScheduler(), lr=1e-4, ema_decay=0.999)
trainer.fit(batch_iterator(data, batch_size=32), steps=10_000)
model.save_pretrained("my-model")
```

## Fine-tune with LoRA (on your laptop)

```python
from mlx_diffusion import DiT, inject_lora, save_lora
from mlx_diffusion.schedulers import FlowMatchEulerScheduler
from mlx_diffusion import DiffusionTrainer

model = DiT.from_pretrained("my-model")
inject_lora(model, rank=8)                 # base frozen; only adapters train
trainer = DiffusionTrainer(model, FlowMatchEulerScheduler(), lr=5e-3)
trainer.fit(batch_iterator(data, batch_size=8), steps=1000)
save_lora(model, "my-lora", rank=8, alpha=16)
```

```bash
mlx-diffuser train --data ./photos --base my-model --lora --out my-lora --steps 1000
```

## What's inside

- **Models** — `DiT` (image/video/text via config), `UNet2D` (Stable-Diffusion
  style, text-conditionable), `AutoencoderKL` (VAE for latent diffusion).
- **Schedulers** — `DDPM`, `DDIM`, `EulerDiscrete`, `FlowMatchEuler`; one object
  covers both training (`add_noise`/`get_target`) and sampling (`step`).
- **Training** — `DiffusionTrainer` (compiled step, EMA, grad clip, min-SNR
  weighting), full fine-tune, and LoRA.
- **Apple-silicon speed** — fused attention, `mx.compile`, lazy evaluation,
  unified-memory controls, and 4/8-bit weight quantization.
- **CLI** — `generate`, `train`, `convert`.

## Design

The library rests on three orthogonal pieces — a **process** (the noise/flow math),
a **network** (a plain `nn.Module`), and a **pipeline** (inference glue). Keeping
them independent is what lets the same network train under DDPM today and
flow-matching tomorrow. See [DESIGN.md](DESIGN.md) and the
[documentation](https://AmirHossein-razlighi.github.io/mlx_diffusion/).

## License

Apache-2.0.
