Metadata-Version: 2.4
Name: caption-flow
Version: 0.5.4
Summary: Self-contained distributed community captioning system
Author: bghira
Author-email: bghira <bghira@users.github.com>
License-Expression: AGPL-3.0-or-later
Project-URL: Homepage, https://github.com/bghira/CaptionFlow
Project-URL: Repository, https://github.com/bghira/CaptionFlow
Keywords: captioning,distributed,vllm,dataset,community
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Requires-Python: >=3.11,<3.14
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: websockets<17.0,>=16.0
Requires-Dist: pyarrow<26.0.0,>=21.0.0
Requires-Dist: click<9.0.0,>=8.2.0
Requires-Dist: pydantic<3.0.0,>=2.12.0
Requires-Dist: aiofiles<26.0.0,>=24.1.0
Requires-Dist: rich<16.0.0,>=14.0.0
Requires-Dist: cryptography<50.0.0,>=45.0.0
Requires-Dist: PyYAML<7.0.0,>=6.0.2
Requires-Dist: certbot<6.0.0,>=5.0.0
Requires-Dist: numpy<3.0.0,>=2.2.0
Requires-Dist: Pillow<13.0.0,>=11.3.0
Requires-Dist: pandas<4.0.0,>=3.0.0
Requires-Dist: datasets<6.0.0,>=5.0.0
Requires-Dist: boto3<2.0.0,>=1.43.0
Requires-Dist: webshart<0.6.0,>=0.5.2
Requires-Dist: pylance<9.0.0,>=8.0.0
Requires-Dist: duckdb<2.0.0,>=1.5.0
Requires-Dist: aiohttp<4.0.0,>=3.13.3
Requires-Dist: fastapi<0.137.0,>=0.133.0
Requires-Dist: uvicorn<1.0.0,>=0.35.0
Requires-Dist: huggingface-hub<2.0.0,>=1.5.0
Requires-Dist: opencv-python-headless<6.0.0,>=4.13.0
Requires-Dist: psutil<8.0.0,>=7.0.0
Requires-Dist: requests<3.0.0,>=2.32.0
Requires-Dist: tqdm<5.0.0,>=4.67.0
Requires-Dist: urwid<5.0.0,>=3.0.2
Provides-Extra: vllm
Requires-Dist: vllm<0.26.0,>=0.25.1; extra == "vllm"
Requires-Dist: torch==2.11.0; extra == "vllm"
Requires-Dist: torchvision==0.26.0; extra == "vllm"
Requires-Dist: torchaudio==2.11.0; extra == "vllm"
Requires-Dist: transformers<6.0.0,>=5.5.3; extra == "vllm"
Requires-Dist: qwen-vl-utils<0.1.0,>=0.0.14; extra == "vllm"
Provides-Extra: captioning
Requires-Dist: vllm<0.26.0,>=0.25.1; extra == "captioning"
Requires-Dist: torch==2.11.0; extra == "captioning"
Requires-Dist: torchvision==0.26.0; extra == "captioning"
Requires-Dist: torchaudio==2.11.0; extra == "captioning"
Requires-Dist: transformers<6.0.0,>=5.5.3; extra == "captioning"
Requires-Dist: qwen-vl-utils<0.1.0,>=0.0.14; extra == "captioning"
Provides-Extra: cpu
Requires-Dist: torch>=2.11.0; extra == "cpu"
Requires-Dist: torchvision>=0.26.0; extra == "cpu"
Requires-Dist: torchaudio>=2.11.0; extra == "cpu"
Provides-Extra: cuda
Requires-Dist: torch>=2.11.0; extra == "cuda"
Requires-Dist: torchvision>=0.26.0; extra == "cuda"
Requires-Dist: torchaudio>=2.11.0; extra == "cuda"
Provides-Extra: cuda13
Requires-Dist: torch>=2.11.0; extra == "cuda13"
Requires-Dist: torchvision>=0.26.0; extra == "cuda13"
Requires-Dist: torchaudio>=2.11.0; extra == "cuda13"
Provides-Extra: rocm
Requires-Dist: torch>=2.11.0; extra == "rocm"
Requires-Dist: torchvision>=0.26.0; extra == "rocm"
Requires-Dist: torchaudio>=2.11.0; extra == "rocm"
Provides-Extra: apple
Requires-Dist: torch>=2.13.0; extra == "apple"
Requires-Dist: torchvision>=0.28.0; extra == "apple"
Requires-Dist: torchaudio>=2.11.0; extra == "apple"
Requires-Dist: vllm-metal==0.1.0; (platform_system == "Darwin" and platform_machine == "arm64" and python_version >= "3.12") and extra == "apple"
Provides-Extra: dev
Requires-Dist: pytest<9.0.0,>=8.0.0; extra == "dev"
Requires-Dist: pytest-asyncio<2.0.0,>=1.1.0; extra == "dev"
Requires-Dist: pytest-cov<7.0.0,>=6.0.0; extra == "dev"
Requires-Dist: transformers<6.0.0,>=5.5.3; extra == "dev"
Requires-Dist: black<26.0.0,>=25.0.0; extra == "dev"
Requires-Dist: ruff<1.0.0,>=0.12.0; extra == "dev"
Requires-Dist: mypy<2.0.0,>=1.17.0; extra == "dev"
Provides-Extra: all
Requires-Dist: vllm<0.26.0,>=0.25.1; extra == "all"
Requires-Dist: torch==2.11.0; extra == "all"
Requires-Dist: torchvision==0.26.0; extra == "all"
Requires-Dist: torchaudio==2.11.0; extra == "all"
Requires-Dist: transformers<6.0.0,>=5.5.3; extra == "all"
Requires-Dist: qwen-vl-utils<0.1.0,>=0.0.14; extra == "all"
Dynamic: author
Dynamic: license-file
Dynamic: provides-extra
Dynamic: requires-dist
Dynamic: requires-python

# CaptionFlow

<!-- [![Tests](https://github.com/bghira/CaptionFlow/workflows/tests/badge.svg)](https://github.com/bghira/CaptionFlow/actions/workflows/tests.yml) -->
[![codecov](https://codecov.io/github/bghira/CaptionFlow/graph/badge.svg?token=PRAQPNGYAS)](https://codecov.io/github/bghira/CaptionFlow)
[![PyPI version](https://badge.fury.io/py/caption-flow.svg)](https://badge.fury.io/py/caption-flow)

scalable, fault-tolerant **vLLM-powered image captioning**.

a fast websocket-based orchestrator paired with lightweight gpu workers achieves exceptional performance for batched requests through vLLM.

CaptionFlow is also integrated in [bghira/SimpleTuner](https://github.com/bghira/SimpleTuner), where it powers an end-to-end caption-to-training workflow through the SimpleTuner WebUI. Use CaptionFlow directly when you want a standalone distributed captioning system, or use it through SimpleTuner when you want dataset captioning, caption review/export, and model training managed as one suite.

* **orchestrator**: hands out work in chunked shards, collects captions, checkpoints progress, and keeps simple stats.
* **workers (vLLM)**: connect to the orchestrator, stream in image samples, batch them, and generate 1..N captions per image using prompts supplied by the orchestrator.
* **config-driven**: all components read YAML config; flags can override.

> no conda. just `venv` + `pip`.

---

## install from pypi

```bash
python -m venv .venv
source .venv/bin/activate  # windows: .venv\Scripts\activate
pip install --upgrade pip
pip install "caption-flow[vllm]"
```

For an orchestrator or monitor-only install, use `pip install -e .`.
`.[captioning]` is an alias for `.[vllm]` for integrations such as
SimpleTuner. Terminal image previews remain optional because the current
`term-image` release requires an older Pillow major than CaptionFlow uses.

On Apple Silicon, install the MPS-compatible PyTorch chain and pinned Metal
plugin with:

```bash
pip install -e ".[apple]"
```

CaptionFlow can use vLLM on Apple Silicon too, but the normal Linux `vllm`
wheel is not the Apple install path. The pinned `vllm-metal` dependency is
selected automatically on native arm64 Python 3.12+. For a ready-to-run
Metal worker, use the upstream installer, which also builds/installs the
Apple-specific vLLM core:

```bash
curl -fsSL https://raw.githubusercontent.com/vllm-project/vllm-metal/main/install.sh | bash
source ~/.venv-vllm-metal/bin/activate
pip install -e .
```

Do not combine `.[apple]` with the Linux `.[vllm]` extra.

For native macOS CPU vLLM instead, follow the
[official source-build instructions](https://docs.vllm.ai/en/stable/getting_started/installation/cpu/?device=apple).

## quickstart (single box)

for a full caption-to-training workflow with a web interface, use the SimpleTuner WebUI integration. the standalone flow below is best when you want to run CaptionFlow directly, contribute workers to a cluster, or export captions for your own downstream training pipeline.

1. copy + edit the sample configs

```bash
cp examples/orchestrator/local_image_files.yaml my-orchestrator.yaml
cp examples/worker.yaml my-worker.yaml
cp examples/monitor.yaml my-monitor.yaml   # optional terminal interface
```

set a unique shared token in both `my-orchestrator.yaml` and `my-worker.yaml` (see `auth.worker_tokens` in the orchestrator config and `worker.token` in the worker config).

if you use private hugging face datasets/models, export `HUGGINGFACE_HUB_TOKEN` before starting anything.

2. start the orchestrator

```bash
caption-flow orchestrator --config my-orchestrator.yaml
```

3. start one or more vLLM workers

```bash
# gpu 0 on the same host
caption-flow worker --config my-worker.yaml --gpu-id 0

# your second GPU
caption-flow worker --config my-worker.yaml --gpu-id 1

# on a remote host
caption-flow worker --config my-worker.yaml --server ws://your.hostname.address:8765
```

4. (optional) start the monitor

```bash
caption-flow monitor --config my-monitor.yaml
```

5. export the data

```bash
% caption-flow export --help                                                                                                                                      
Usage: caption-flow export [OPTIONS]

  Export caption data to various formats.

Options:
  --format [jsonl|json|csv|txt|parquet|webshart|lance|huggingface_hub|all] Export format (default: jsonl)
```

* **jsonl**: create JSON line file in the specified `--output` path
* **csv**: exports CSV-compatible data columns to the `--output` path containing incomplete metadata
* **json**: creates a `.json` file for each sample inside the `--output` subdirectory containing **complete** metadata; useful for webdatasets
* **txt**: creates `.txt` file for each sample inside the `--output` subdirectory containing ONLY captions
* **webshart**: updates an **existing per-shard metadata `.json` file** by writing captions under the plural `captions` key. for this format, pass `--output` as the path to the existing shard metadata JSON file when exporting one shard. if you export multiple shards, pass `--output` as a directory containing one existing `{shard_name}.json` file per shard.
* **huggingface_hub**: creates a dataset on Hugging Face Hub, possibly `--private` and `--nsfw` where necessary
* **all**: creates the directory/file-generating export formats in a specified `--output` directory. prefer a directory here; `webshart` is a special case that expects existing per-shard metadata `.json` files rather than creating new metadata files.

> note: `--output` paths ending in `.json` are treated specially for `webshart`. use a directory for normal multi-format exports and an existing shard metadata JSON file only when intentionally updating a `webshart` shard.

---

## how it’s wired

### orchestrator

* **websocket server** (default `0.0.0.0:8765`) with three client roles: workers, data-feeders, and admin.
* **dataset control**: the orchestrator centrally defines the dataset (`huggingface` or `local`) and version/name. it chunk-slices shards and assigns work.
* **data serving to remote workers**: local files can be captioned by remote workers that don't have access to the same files, automatically.
* **vLLM config broadcast**: model, tp size, dtype, max seq len, memory targets, batching, sampling params, and **inference prompts** are all pushed to workers; workers can apply many changes without a model reload.
* **storage + checkpoints**: captions buffer to disk with periodic checkpoints. chunk state is tracked so restarts don’t double-work.
* **auth**: token lists for `worker`, `monitor`, and `admin` roles.

### vLLM worker

* **one process per gpu**. select the device with `--gpu-id` (or `worker.gpu_id` in YAML).
* **gets its marching orders** from the orchestrator: dataset info, model, prompts, batch size, and sampling.
* **resilient**: detects disconnects, abandons the current chunk cleanly, clears queues, reconnects, and resumes.
* **batched generate()**: images are resized down for consistent batching; each image can get multiple captions (one per prompt).

---

## dataset formats

* huggingface hub or local based URL list datasets that are compatible with the datasets library
* huggingface hub datasets that are simple containers of raw image files
* webdatasets shards containing full image data; also can be hosted on the hub
* local folder filled with images; orchestrator will serve the data to workers

## configuration path

### config discovery order

for any component, the CLI looks for config in this order (first match wins):

1. `--config /path/to/file.yaml`
2. `./<component>.yaml` (current directory)
3. `~/.caption-flow/<component>.yaml`
4. `$XDG_CONFIG_HOME/caption-flow/<component>.yaml`
5. `/etc/caption-flow/<component>.yaml`
6. any `$XDG_CONFIG_DIRS` entries under `caption-flow/`
7. `./examples/<component>.yaml` (fallback)

---

## tls / certificates

use the built-in helpers during development:

```bash
# self-signed certs for quick local testing
caption-flow generate_cert --self-signed --domain localhost --output-dir ./certs

# inspect any certificate file
caption-flow inspect_cert ./certs/fullchain.pem
```

then point the orchestrator at the resulting cert/key (or run `--no-ssl` for dev-only ws\://).

---

## tips & notes

* **multi-gpu**: start one worker process per gpu (set `--gpu-id` or `worker.gpu_id`).
* **throughput**: tune `vllm.batch_size` in the orchestrator config (or override with `--batch-size` at worker start). higher isn’t always better; watch VRAM.
* **prompts**: add more strings under `vllm.inference_prompts` to get multiple captions per image; the worker returns only non-empty generations.
* **private HF**: if your dataset/model needs auth, export `HUGGINGFACE_HUB_TOKEN` before `caption-flow worker ...`.
* **self-signed ssl**: pass `--no-verify-ssl` to workers/monitors in dev.
* **recovery**: if you hard-crash mid-run, `caption-flow scan_chunks --fix` can reset abandoned chunks so the orchestrator can reissue them cleanly.

---

## roadmap

* hot config reload via the admin websocket path.
* dedicated data-feeder clients (separate from gpu workers) that push samples into the orchestrator.
* richer monitor TUI.

PRs welcome. keep it simple and fast.

## architecture

```
┌─────────────┐     WebSocket      ┌─────────────┐
│   Worker    │◄──────────────────►│             │
│             │                    │             │     ┌──────────────┐
│             │◄───────────────────│             │────►│Arrow/Parquet │
└─────────────┘   HTTP (img data)  │ Orchestrator│     │   Storage    │
                                   │             │     └──────────────┘
┌─────────────┐                    │             │
│   Worker    │◄──────────────────►│             │
│             │                    │             │
│             │◄───────────────────│             │
└─────────────┘   HTTP (img data)  └─────────────┘
                                           ▲
┌─────────────┐                           │
│   Monitor   │◄──────────────────────────┘
└─────────────┘
```

## Community Clusters

To contribute compute to a cluster:

1. Install caption-flow: `pip install "caption-flow[vllm]"`
2. Get a worker token from the project maintainer
3. Run: `caption-flow worker --server wss://project.domain.com:8765 --token YOUR_TOKEN`

Your contributions will be tracked and attributed in the final dataset!

## License

AGPLv3
