caption-flow 0.5.6__tar.gz → 0.6.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (78) hide show
  1. {caption_flow-0.5.6 → caption_flow-0.6.0}/PKG-INFO +107 -68
  2. {caption_flow-0.5.6 → caption_flow-0.6.0}/README.md +50 -12
  3. {caption_flow-0.5.6 → caption_flow-0.6.0}/pyproject.toml +1 -1
  4. {caption_flow-0.5.6 → caption_flow-0.6.0}/setup.py +49 -46
  5. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/__init__.py +1 -1
  6. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/models.py +8 -5
  7. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/orchestrator.py +21 -13
  8. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/processors/webdataset.py +187 -15
  9. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/storage/exporter.py +20 -6
  10. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/storage/manager.py +14 -2
  11. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/utils/chunk_tracker.py +45 -11
  12. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/utils/image_processor.py +8 -3
  13. caption_flow-0.6.0/src/caption_flow/utils/output_policy.py +207 -0
  14. caption_flow-0.6.0/src/caption_flow/utils/retry_queue.py +50 -0
  15. caption_flow-0.6.0/src/caption_flow/utils/unicode_utils.py +30 -0
  16. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/utils/vllm_config.py +39 -11
  17. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/workers/caption.py +177 -206
  18. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/workers/openai_compatible.py +71 -361
  19. caption_flow-0.6.0/src/caption_flow/workers/pipeline.py +127 -0
  20. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow.egg-info/PKG-INFO +107 -68
  21. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow.egg-info/SOURCES.txt +15 -0
  22. caption_flow-0.6.0/src/caption_flow.egg-info/requires.txt +92 -0
  23. caption_flow-0.6.0/tests/test_caption_pipeline.py +413 -0
  24. caption_flow-0.6.0/tests/test_chunk_completion_index.py +193 -0
  25. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_config_reload.py +20 -0
  26. caption_flow-0.6.0/tests/test_dependency_config.py +71 -0
  27. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_exporter.py +76 -3
  28. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_openai_compatible_worker.py +2 -1
  29. caption_flow-0.6.0/tests/test_orchestrator_storage_order.py +32 -0
  30. caption_flow-0.6.0/tests/test_result_acceptance.py +229 -0
  31. caption_flow-0.6.0/tests/test_retry_queue.py +132 -0
  32. caption_flow-0.6.0/tests/test_unicode_utils.py +69 -0
  33. caption_flow-0.6.0/tests/test_webdataset_cache_wait.py +118 -0
  34. caption_flow-0.6.0/tests/test_webdataset_prefetch.py +140 -0
  35. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_webdataset_ranges.py +31 -2
  36. caption_flow-0.6.0/tests/test_webshart_dependency.py +27 -0
  37. caption_flow-0.6.0/tests/test_webshart_export_contract.py +183 -0
  38. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_worker_caption.py +19 -28
  39. caption_flow-0.5.6/src/caption_flow.egg-info/requires.txt +0 -89
  40. {caption_flow-0.5.6 → caption_flow-0.6.0}/LICENSE +0 -0
  41. {caption_flow-0.5.6 → caption_flow-0.6.0}/setup.cfg +0 -0
  42. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/cli.py +0 -0
  43. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/monitor.py +0 -0
  44. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/processors/__init__.py +0 -0
  45. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/processors/base.py +0 -0
  46. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/processors/huggingface.py +0 -0
  47. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/processors/local_filesystem.py +0 -0
  48. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/storage/__init__.py +0 -0
  49. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/utils/__init__.py +0 -0
  50. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/utils/auth.py +0 -0
  51. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/utils/caption_utils.py +0 -0
  52. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/utils/certificates.py +0 -0
  53. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/utils/checkpoint_tracker.py +0 -0
  54. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/utils/json_utils.py +0 -0
  55. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/utils/prompt_template.py +0 -0
  56. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/viewer.py +0 -0
  57. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/workers/base.py +0 -0
  58. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow/workers/data.py +0 -0
  59. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow.egg-info/dependency_links.txt +0 -0
  60. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow.egg-info/entry_points.txt +0 -0
  61. {caption_flow-0.5.6 → caption_flow-0.6.0}/src/caption_flow.egg-info/top_level.txt +0 -0
  62. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_caption_utils.py +0 -0
  63. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_certificates.py +0 -0
  64. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_cli.py +0 -0
  65. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_duplicate_job_assignments.py +0 -0
  66. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_fix_verification.py +0 -0
  67. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_huggingface_ranges.py +0 -0
  68. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_image_processor.py +0 -0
  69. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_json_utils.py +0 -0
  70. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_main.py +0 -0
  71. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_monitor.py +0 -0
  72. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_processors.py +0 -0
  73. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_range_level_distribution.py +0 -0
  74. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_storage_components.py +0 -0
  75. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_viewer.py +0 -0
  76. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_vllm_config.py +0 -0
  77. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_worker_reconnection_complete.py +0 -0
  78. {caption_flow-0.5.6 → caption_flow-0.6.0}/tests/test_worker_reconnection_sequence.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: caption-flow
3
- Version: 0.5.6
3
+ Version: 0.6.0
4
4
  Summary: Self-contained distributed community captioning system
5
5
  Author: bghira
6
6
  Author-email: bghira <bghira@users.github.com>
@@ -16,82 +16,83 @@ Classifier: Programming Language :: Python :: 3.13
16
16
  Requires-Python: >=3.12,<3.14
17
17
  Description-Content-Type: text/markdown
18
18
  License-File: LICENSE
19
- Requires-Dist: websockets<17.0,>=16.0
20
- Requires-Dist: pyarrow<26.0.0,>=21.0.0
21
- Requires-Dist: click<9.0.0,>=8.2.0
22
- Requires-Dist: pydantic<3.0.0,>=2.12.0
23
- Requires-Dist: aiofiles<26.0.0,>=24.1.0
24
- Requires-Dist: rich<16.0.0,>=14.0.0
25
- Requires-Dist: cryptography<50.0.0,>=45.0.0
26
- Requires-Dist: PyYAML<7.0.0,>=6.0.2
27
- Requires-Dist: certbot<6.0.0,>=5.0.0
28
- Requires-Dist: numpy<3.0.0,>=2.2.0
29
- Requires-Dist: Pillow<13.0.0,>=11.3.0
19
+ Requires-Dist: websockets<18.0,>=17.4
20
+ Requires-Dist: pyarrow<27.0.0,>=26.0.1
21
+ Requires-Dist: click<9.0.0,>=8.5.1
22
+ Requires-Dist: pydantic<3.0.0,>=2.14.3
23
+ Requires-Dist: aiofiles<26.0.0,>=25.1.0
24
+ Requires-Dist: rich<16.0.0,>=15.0.0
25
+ Requires-Dist: cryptography<51.0.0,>=50.0.2
26
+ Requires-Dist: PyYAML<7.0.0,>=6.0.3
27
+ Requires-Dist: certbot<6.0.0,>=5.10.0
28
+ Requires-Dist: numpy<2.4.0,>=2.3.5; python_version < "3.13"
29
+ Requires-Dist: numpy<2.5.0,>=2.4.6; python_version >= "3.13"
30
+ Requires-Dist: Pillow<14.0.0,>=13.0.0
30
31
  Requires-Dist: trainingsample<0.4.0,>=0.3.2
31
- Requires-Dist: pandas<4.0.0,>=3.0.0
32
- Requires-Dist: datasets<6.0.0,>=5.0.0
33
- Requires-Dist: boto3<2.0.0,>=1.43.0
34
- Requires-Dist: webshart<0.6.0,>=0.5.2
35
- Requires-Dist: pylance<9.0.0,>=8.0.0
36
- Requires-Dist: duckdb<2.0.0,>=1.5.0
37
- Requires-Dist: aiohttp<4.0.0,>=3.13.3
38
- Requires-Dist: fastapi<0.137.0,>=0.133.0
39
- Requires-Dist: uvicorn<1.0.0,>=0.35.0
40
- Requires-Dist: huggingface-hub<2.0.0,>=1.5.0
41
- Requires-Dist: psutil<8.0.0,>=7.0.0
42
- Requires-Dist: requests<3.0.0,>=2.32.0
43
- Requires-Dist: tqdm<5.0.0,>=4.67.0
44
- Requires-Dist: urwid<5.0.0,>=3.0.2
32
+ Requires-Dist: pandas<4.0.0,>=3.1.1
33
+ Requires-Dist: datasets<6.0.0,>=5.2.0
34
+ Requires-Dist: boto3<2.0.0,>=1.44.3
35
+ Requires-Dist: webshart<0.7.0,>=0.6.0
36
+ Requires-Dist: pylance<9.0.0,>=8.0.1
37
+ Requires-Dist: duckdb<2.0.0,>=1.5.6
38
+ Requires-Dist: aiohttp<4.0.0,>=3.14.5
39
+ Requires-Dist: fastapi<0.137.0,>=0.136.3
40
+ Requires-Dist: uvicorn<1.0.0,>=0.54.0
41
+ Requires-Dist: huggingface-hub<2.0.0,>=1.33.0
42
+ Requires-Dist: psutil<8.0.0,>=7.2.2
43
+ Requires-Dist: requests<3.0.0,>=2.34.2
44
+ Requires-Dist: tqdm<5.0.0,>=4.70.1
45
+ Requires-Dist: urwid<6.0.0,>=5.1.3
45
46
  Provides-Extra: vllm
46
- Requires-Dist: vllm<0.26.0,>=0.25.1; extra == "vllm"
47
- Requires-Dist: torch==2.11.0; extra == "vllm"
48
- Requires-Dist: torchvision==0.26.0; extra == "vllm"
47
+ Requires-Dist: vllm<0.34.0,>=0.33.0; extra == "vllm"
48
+ Requires-Dist: torch==2.14.0; extra == "vllm"
49
+ Requires-Dist: torchvision==0.29.0; extra == "vllm"
49
50
  Requires-Dist: torchaudio==2.11.0; extra == "vllm"
50
- Requires-Dist: transformers<6.0.0,>=5.5.3; extra == "vllm"
51
+ Requires-Dist: transformers<5.21.0,>=5.16.1; extra == "vllm"
51
52
  Requires-Dist: qwen-vl-utils<0.1.0,>=0.0.14; extra == "vllm"
52
53
  Provides-Extra: captioning
53
- Requires-Dist: vllm<0.26.0,>=0.25.1; extra == "captioning"
54
- Requires-Dist: torch==2.11.0; extra == "captioning"
55
- Requires-Dist: torchvision==0.26.0; extra == "captioning"
54
+ Requires-Dist: vllm<0.34.0,>=0.33.0; extra == "captioning"
55
+ Requires-Dist: torch==2.14.0; extra == "captioning"
56
+ Requires-Dist: torchvision==0.29.0; extra == "captioning"
56
57
  Requires-Dist: torchaudio==2.11.0; extra == "captioning"
57
- Requires-Dist: transformers<6.0.0,>=5.5.3; extra == "captioning"
58
+ Requires-Dist: transformers<5.21.0,>=5.16.1; extra == "captioning"
58
59
  Requires-Dist: qwen-vl-utils<0.1.0,>=0.0.14; extra == "captioning"
59
60
  Provides-Extra: openai
60
61
  Provides-Extra: cpu
61
- Requires-Dist: torch>=2.11.0; extra == "cpu"
62
- Requires-Dist: torchvision>=0.26.0; extra == "cpu"
62
+ Requires-Dist: torch>=2.14.0; extra == "cpu"
63
+ Requires-Dist: torchvision>=0.29.0; extra == "cpu"
63
64
  Requires-Dist: torchaudio>=2.11.0; extra == "cpu"
64
65
  Provides-Extra: cuda
65
- Requires-Dist: torch>=2.11.0; extra == "cuda"
66
- Requires-Dist: torchvision>=0.26.0; extra == "cuda"
66
+ Requires-Dist: torch>=2.14.0; extra == "cuda"
67
+ Requires-Dist: torchvision>=0.29.0; extra == "cuda"
67
68
  Requires-Dist: torchaudio>=2.11.0; extra == "cuda"
68
69
  Provides-Extra: cuda13
69
- Requires-Dist: torch>=2.11.0; extra == "cuda13"
70
- Requires-Dist: torchvision>=0.26.0; extra == "cuda13"
70
+ Requires-Dist: torch>=2.14.0; extra == "cuda13"
71
+ Requires-Dist: torchvision>=0.29.0; extra == "cuda13"
71
72
  Requires-Dist: torchaudio>=2.11.0; extra == "cuda13"
72
73
  Provides-Extra: rocm
73
- Requires-Dist: torch>=2.11.0; extra == "rocm"
74
- Requires-Dist: torchvision>=0.26.0; extra == "rocm"
74
+ Requires-Dist: torch>=2.14.0; extra == "rocm"
75
+ Requires-Dist: torchvision>=0.29.0; extra == "rocm"
75
76
  Requires-Dist: torchaudio>=2.11.0; extra == "rocm"
76
77
  Provides-Extra: apple
77
- Requires-Dist: torch>=2.13.0; extra == "apple"
78
- Requires-Dist: torchvision>=0.28.0; extra == "apple"
78
+ Requires-Dist: torch>=2.14.0; extra == "apple"
79
+ Requires-Dist: torchvision>=0.29.0; extra == "apple"
79
80
  Requires-Dist: torchaudio>=2.11.0; extra == "apple"
80
- Requires-Dist: vllm-metal==0.1.0; (platform_system == "Darwin" and platform_machine == "arm64" and python_version >= "3.12") and extra == "apple"
81
+ Requires-Dist: qwen-vl-utils<0.1.0,>=0.0.14; extra == "apple"
81
82
  Provides-Extra: dev
82
- Requires-Dist: pytest<9.0.0,>=8.0.0; extra == "dev"
83
- Requires-Dist: pytest-asyncio<2.0.0,>=1.1.0; extra == "dev"
84
- Requires-Dist: pytest-cov<7.0.0,>=6.0.0; extra == "dev"
85
- Requires-Dist: transformers<6.0.0,>=5.5.3; extra == "dev"
86
- Requires-Dist: black<26.0.0,>=25.0.0; extra == "dev"
87
- Requires-Dist: ruff<1.0.0,>=0.12.0; extra == "dev"
88
- Requires-Dist: mypy<2.0.0,>=1.17.0; extra == "dev"
83
+ Requires-Dist: pytest<10.0.0,>=9.2.0; extra == "dev"
84
+ Requires-Dist: pytest-asyncio<2.0.0,>=1.5.0; extra == "dev"
85
+ Requires-Dist: pytest-cov<8.0.0,>=7.1.0; extra == "dev"
86
+ Requires-Dist: transformers<5.21.0,>=5.16.1; extra == "dev"
87
+ Requires-Dist: black<27.0.0,>=26.10.1; extra == "dev"
88
+ Requires-Dist: ruff<1.0.0,>=0.17.5; extra == "dev"
89
+ Requires-Dist: mypy<3.0.0,>=2.4.1; extra == "dev"
89
90
  Provides-Extra: all
90
- Requires-Dist: vllm<0.26.0,>=0.25.1; extra == "all"
91
- Requires-Dist: torch==2.11.0; extra == "all"
92
- Requires-Dist: torchvision==0.26.0; extra == "all"
91
+ Requires-Dist: vllm<0.34.0,>=0.33.0; extra == "all"
92
+ Requires-Dist: torch==2.14.0; extra == "all"
93
+ Requires-Dist: torchvision==0.29.0; extra == "all"
93
94
  Requires-Dist: torchaudio==2.11.0; extra == "all"
94
- Requires-Dist: transformers<6.0.0,>=5.5.3; extra == "all"
95
+ Requires-Dist: transformers<5.21.0,>=5.16.1; extra == "all"
95
96
  Requires-Dist: qwen-vl-utils<0.1.0,>=0.0.14; extra == "all"
96
97
  Dynamic: author
97
98
  Dynamic: license-file
@@ -128,6 +129,18 @@ pip install --upgrade pip
128
129
  pip install "caption-flow[vllm]"
129
130
  ```
130
131
 
132
+ The Linux GPU worker extra targets vLLM 0.33.x with PyTorch 2.14.0,
133
+ torchvision 0.29.0, and torchaudio 2.11.0 (as required by vLLM), plus
134
+ Transformers 5.16.1–5.20.x and Pillow 13. To upgrade an existing worker from a checkout,
135
+ run `pip install --upgrade -e ".[vllm]"` in its environment and restart it.
136
+
137
+ The tested worker stack uses NumPy 2.3.x on Python 3.12 and 2.4.x
138
+ on Python 3.13, and FastAPI below 0.137. These bounds are shared with the
139
+ base install so adding the worker extra remains compatible. CI resolves
140
+ the full Linux worker and development dependencies on both Python versions.
141
+ Lance 8.x, DuckDB 1.x, and Hugging Face Hub 1.x remain bounded to their
142
+ existing major versions pending separate migration testing.
143
+
131
144
  An OpenAI-compatible API worker does not need the GPU dependencies:
132
145
 
133
146
  ```bash
@@ -139,26 +152,36 @@ For an orchestrator or monitor-only install, use `pip install -e .`.
139
152
  SimpleTuner. Terminal image previews remain optional because the current
140
153
  `term-image` release requires an older Pillow major than CaptionFlow uses.
141
154
 
142
- On Apple Silicon, install the MPS-compatible PyTorch chain and pinned Metal
143
- plugin with:
155
+ On Apple Silicon with macOS 14 or newer, install the MPS-compatible PyTorch
156
+ chain and image utilities with:
144
157
 
145
158
  ```bash
146
159
  pip install -e ".[apple]"
147
160
  ```
148
161
 
149
- CaptionFlow can use vLLM on Apple Silicon too, but the normal Linux `vllm`
150
- wheel is not the Apple install path. The pinned `vllm-metal` dependency is
151
- selected automatically on native arm64 Python 3.12+. For a ready-to-run
152
- Metal worker, use the upstream installer, which also builds/installs the
153
- Apple-specific vLLM core:
162
+ For native vLLM Metal captioning, use **macOS 15+, Apple Silicon, and Python
163
+ 3.12**. The current Metal wheels are published on
164
+ [GitHub releases](https://github.com/vllm-project/vllm-metal/releases), while
165
+ PyPI still carries the obsolete 0.1.0 plugin. From this checkout, install the
166
+ matching stable vLLM 0.33.0 core and Metal 0.33.0 wheels into your environment:
154
167
 
155
168
  ```bash
156
- curl -fsSL https://raw.githubusercontent.com/vllm-project/vllm-metal/main/install.sh | bash
157
- source ~/.venv-vllm-metal/bin/activate
158
- pip install -e .
169
+ python -m pip install --upgrade -r requirements.metal.txt
170
+ python -m pip install -e ".[dev]"
171
+ python -m pip check
172
+ python examples/metal_smoke.py
159
173
  ```
160
174
 
161
- Do not combine `.[apple]` with the Linux `.[vllm]` extra.
175
+ The smoke test downloads `mlx-community/Qwen3-VL-2B-Instruct-4bit` on first
176
+ use and runs two images through CaptionFlow's native caption pipeline. It
177
+ checks that the Metal plugin and MLX GPU are active, and that the captions
178
+ identify the images' shapes and different colors. Pass `--model /path/to/model`
179
+ to use a local checkpoint. The test does not connect to an orchestrator.
180
+
181
+ The stable recipe uses the core wheel's PyTorch 2.14.0 pin. Upstream's installer
182
+ can select newer development builds and override Torch with nightly packages;
183
+ those are separate from this recipe. `.[apple]` alone does not install a vLLM
184
+ runtime, and should not be combined with the Linux `.[vllm]` extra.
162
185
 
163
186
  For native macOS CPU vLLM instead, follow the
164
187
  [official source-build instructions](https://docs.vllm.ai/en/stable/getting_started/installation/cpu/?device=apple).
@@ -220,7 +243,7 @@ Options:
220
243
  * **csv**: exports CSV-compatible data columns to the `--output` path containing incomplete metadata
221
244
  * **json**: creates a `.json` file for each sample inside the `--output` subdirectory containing **complete** metadata; useful for webdatasets
222
245
  * **txt**: creates `.txt` file for each sample inside the `--output` subdirectory containing ONLY captions
223
- * **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.
246
+ * **webshart**: updates an **existing per-shard metadata `.json` file** by writing captions under the plural `captions` key. Strings, native JSON objects, and mixed lists are supported. JSON objects stored as text in Lance output columns are decoded during export, preserving nested values instead of writing Python dictionary representations. Native objects supplied through `Caption.outputs` are encoded as JSON only for Lance persistence. Requires webshart 0.6.0 or newer for native JSON caption support (included in the base dependencies). 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.
224
247
  * **huggingface_hub**: creates a dataset on Hugging Face Hub, possibly `--private` and `--nsfw` where necessary
225
248
  * **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.
226
249
 
@@ -279,6 +302,13 @@ prompts, sampling, and output fields. Existing `vllm:` configurations are also
279
302
  accepted by API workers for backward compatibility; endpoint-local `model`
280
303
  values take precedence over the broadcast model name.
281
304
 
305
+ Direct vLLM and API workers share JSON validation, output transformations,
306
+ refusal detection, and semantic retries. Configure those under
307
+ `inference.output_processing`; use `response_format` and `retry_response_format`
308
+ for constrained decoding on either backend. See the
309
+ [structured JSON recipe](recipes/STRUCTURED_JSON.md) for examples and native
310
+ vLLM tuning details.
311
+
282
312
  ---
283
313
 
284
314
  ## dataset formats
@@ -288,6 +318,15 @@ values take precedence over the broadcast model name.
288
318
  * webdatasets shards containing full image data; also can be hosted on the hub
289
319
  * local folder filled with images; orchestrator will serve the data to workers
290
320
 
321
+ For indexed Webshart reads, set `orchestrator.webshart_prefetch_shards: 1` to
322
+ download the next shard in the background while captioning the current one.
323
+ The default is `0` (disabled). This works with both direct vLLM and API workers,
324
+ does not move the sample cursor, and is ignored for explicit remote-range reads.
325
+ Size `shard_cache_gb` to hold the current shard plus the lookahead. Prefetch
326
+ retries are spaced by `webshart_cache_retry_interval` (30 seconds by default);
327
+ a prefetch failure never marks samples processed or prevents reading the
328
+ current cached shard.
329
+
291
330
  ## configuration path
292
331
 
293
332
  ### config discovery order
@@ -27,6 +27,18 @@ pip install --upgrade pip
27
27
  pip install "caption-flow[vllm]"
28
28
  ```
29
29
 
30
+ The Linux GPU worker extra targets vLLM 0.33.x with PyTorch 2.14.0,
31
+ torchvision 0.29.0, and torchaudio 2.11.0 (as required by vLLM), plus
32
+ Transformers 5.16.1–5.20.x and Pillow 13. To upgrade an existing worker from a checkout,
33
+ run `pip install --upgrade -e ".[vllm]"` in its environment and restart it.
34
+
35
+ The tested worker stack uses NumPy 2.3.x on Python 3.12 and 2.4.x
36
+ on Python 3.13, and FastAPI below 0.137. These bounds are shared with the
37
+ base install so adding the worker extra remains compatible. CI resolves
38
+ the full Linux worker and development dependencies on both Python versions.
39
+ Lance 8.x, DuckDB 1.x, and Hugging Face Hub 1.x remain bounded to their
40
+ existing major versions pending separate migration testing.
41
+
30
42
  An OpenAI-compatible API worker does not need the GPU dependencies:
31
43
 
32
44
  ```bash
@@ -38,26 +50,36 @@ For an orchestrator or monitor-only install, use `pip install -e .`.
38
50
  SimpleTuner. Terminal image previews remain optional because the current
39
51
  `term-image` release requires an older Pillow major than CaptionFlow uses.
40
52
 
41
- On Apple Silicon, install the MPS-compatible PyTorch chain and pinned Metal
42
- plugin with:
53
+ On Apple Silicon with macOS 14 or newer, install the MPS-compatible PyTorch
54
+ chain and image utilities with:
43
55
 
44
56
  ```bash
45
57
  pip install -e ".[apple]"
46
58
  ```
47
59
 
48
- CaptionFlow can use vLLM on Apple Silicon too, but the normal Linux `vllm`
49
- wheel is not the Apple install path. The pinned `vllm-metal` dependency is
50
- selected automatically on native arm64 Python 3.12+. For a ready-to-run
51
- Metal worker, use the upstream installer, which also builds/installs the
52
- Apple-specific vLLM core:
60
+ For native vLLM Metal captioning, use **macOS 15+, Apple Silicon, and Python
61
+ 3.12**. The current Metal wheels are published on
62
+ [GitHub releases](https://github.com/vllm-project/vllm-metal/releases), while
63
+ PyPI still carries the obsolete 0.1.0 plugin. From this checkout, install the
64
+ matching stable vLLM 0.33.0 core and Metal 0.33.0 wheels into your environment:
53
65
 
54
66
  ```bash
55
- curl -fsSL https://raw.githubusercontent.com/vllm-project/vllm-metal/main/install.sh | bash
56
- source ~/.venv-vllm-metal/bin/activate
57
- pip install -e .
67
+ python -m pip install --upgrade -r requirements.metal.txt
68
+ python -m pip install -e ".[dev]"
69
+ python -m pip check
70
+ python examples/metal_smoke.py
58
71
  ```
59
72
 
60
- Do not combine `.[apple]` with the Linux `.[vllm]` extra.
73
+ The smoke test downloads `mlx-community/Qwen3-VL-2B-Instruct-4bit` on first
74
+ use and runs two images through CaptionFlow's native caption pipeline. It
75
+ checks that the Metal plugin and MLX GPU are active, and that the captions
76
+ identify the images' shapes and different colors. Pass `--model /path/to/model`
77
+ to use a local checkpoint. The test does not connect to an orchestrator.
78
+
79
+ The stable recipe uses the core wheel's PyTorch 2.14.0 pin. Upstream's installer
80
+ can select newer development builds and override Torch with nightly packages;
81
+ those are separate from this recipe. `.[apple]` alone does not install a vLLM
82
+ runtime, and should not be combined with the Linux `.[vllm]` extra.
61
83
 
62
84
  For native macOS CPU vLLM instead, follow the
63
85
  [official source-build instructions](https://docs.vllm.ai/en/stable/getting_started/installation/cpu/?device=apple).
@@ -119,7 +141,7 @@ Options:
119
141
  * **csv**: exports CSV-compatible data columns to the `--output` path containing incomplete metadata
120
142
  * **json**: creates a `.json` file for each sample inside the `--output` subdirectory containing **complete** metadata; useful for webdatasets
121
143
  * **txt**: creates `.txt` file for each sample inside the `--output` subdirectory containing ONLY captions
122
- * **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.
144
+ * **webshart**: updates an **existing per-shard metadata `.json` file** by writing captions under the plural `captions` key. Strings, native JSON objects, and mixed lists are supported. JSON objects stored as text in Lance output columns are decoded during export, preserving nested values instead of writing Python dictionary representations. Native objects supplied through `Caption.outputs` are encoded as JSON only for Lance persistence. Requires webshart 0.6.0 or newer for native JSON caption support (included in the base dependencies). 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.
123
145
  * **huggingface_hub**: creates a dataset on Hugging Face Hub, possibly `--private` and `--nsfw` where necessary
124
146
  * **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.
125
147
 
@@ -178,6 +200,13 @@ prompts, sampling, and output fields. Existing `vllm:` configurations are also
178
200
  accepted by API workers for backward compatibility; endpoint-local `model`
179
201
  values take precedence over the broadcast model name.
180
202
 
203
+ Direct vLLM and API workers share JSON validation, output transformations,
204
+ refusal detection, and semantic retries. Configure those under
205
+ `inference.output_processing`; use `response_format` and `retry_response_format`
206
+ for constrained decoding on either backend. See the
207
+ [structured JSON recipe](recipes/STRUCTURED_JSON.md) for examples and native
208
+ vLLM tuning details.
209
+
181
210
  ---
182
211
 
183
212
  ## dataset formats
@@ -187,6 +216,15 @@ values take precedence over the broadcast model name.
187
216
  * webdatasets shards containing full image data; also can be hosted on the hub
188
217
  * local folder filled with images; orchestrator will serve the data to workers
189
218
 
219
+ For indexed Webshart reads, set `orchestrator.webshart_prefetch_shards: 1` to
220
+ download the next shard in the background while captioning the current one.
221
+ The default is `0` (disabled). This works with both direct vLLM and API workers,
222
+ does not move the sample cursor, and is ignored for explicit remote-range reads.
223
+ Size `shard_cache_gb` to hold the current shard plus the lookahead. Prefetch
224
+ retries are spaced by `webshart_cache_retry_interval` (30 seconds by default);
225
+ a prefetch failure never marks samples processed or prevents reading the
226
+ current cached shard.
227
+
190
228
  ## configuration path
191
229
 
192
230
  ### config discovery order
@@ -1,5 +1,5 @@
1
1
  [build-system]
2
- requires = ["setuptools>=68.0.0", "wheel"]
2
+ requires = ["setuptools>=77.0.3", "wheel"]
3
3
  build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
@@ -23,58 +23,61 @@ def get_version() -> str:
23
23
 
24
24
 
25
25
  base_deps = [
26
- "websockets>=16.0,<17.0",
27
- "pyarrow>=21.0.0,<26.0.0",
28
- "click>=8.2.0,<9.0.0",
29
- "pydantic>=2.12.0,<3.0.0",
30
- "aiofiles>=24.1.0,<26.0.0",
31
- "rich>=14.0.0,<16.0.0",
32
- "cryptography>=45.0.0,<50.0.0",
33
- "PyYAML>=6.0.2,<7.0.0",
34
- "certbot>=5.0.0,<6.0.0",
35
- "numpy>=2.2.0,<3.0.0",
36
- "Pillow>=11.3.0,<13.0.0",
26
+ "websockets>=17.4,<18.0",
27
+ "pyarrow>=26.0.1,<27.0.0",
28
+ "click>=8.5.1,<9.0.0",
29
+ "pydantic>=2.14.3,<3.0.0",
30
+ "aiofiles>=25.1.0,<26.0.0",
31
+ "rich>=15.0.0,<16.0.0",
32
+ "cryptography>=50.0.2,<51.0.0",
33
+ "PyYAML>=6.0.3,<7.0.0",
34
+ "certbot>=5.10.0,<6.0.0",
35
+ # vLLM's mistral-common requires NumPy <2.4 on Python 3.12;
36
+ # its pinned numba 0.65 requires NumPy <2.5 on Python 3.13.
37
+ "numpy>=2.3.5,<2.4.0; python_version < '3.13'",
38
+ "numpy>=2.4.6,<2.5.0; python_version >= '3.13'",
39
+ "Pillow>=13.0.0,<14.0.0",
37
40
  "trainingsample>=0.3.2,<0.4.0",
38
- "pandas>=3.0.0,<4.0.0",
39
- "datasets>=5.0.0,<6.0.0",
40
- "boto3>=1.43.0,<2.0.0",
41
- "webshart>=0.5.2,<0.6.0",
42
- "pylance>=8.0.0,<9.0.0",
43
- "duckdb>=1.5.0,<2.0.0",
44
- "aiohttp>=3.13.3,<4.0.0",
45
- "fastapi>=0.133.0,<0.137.0",
46
- "uvicorn>=0.35.0,<1.0.0",
47
- "huggingface-hub>=1.5.0,<2.0.0",
48
- "psutil>=7.0.0,<8.0.0",
49
- "requests>=2.32.0,<3.0.0",
50
- "tqdm>=4.67.0,<5.0.0",
51
- "urwid>=3.0.2,<5.0.0",
41
+ "pandas>=3.1.1,<4.0.0",
42
+ "datasets>=5.2.0,<6.0.0",
43
+ "boto3>=1.44.3,<2.0.0",
44
+ "webshart>=0.6.0,<0.7.0",
45
+ "pylance>=8.0.1,<9.0.0",
46
+ "duckdb>=1.5.6,<2.0.0",
47
+ "aiohttp>=3.14.5,<4.0.0",
48
+ # vLLM 0.33 requires FastAPI below 0.137.
49
+ "fastapi>=0.136.3,<0.137.0",
50
+ "uvicorn>=0.54.0,<1.0.0",
51
+ "huggingface-hub>=1.33.0,<2.0.0",
52
+ "psutil>=7.2.2,<8.0.0",
53
+ "requests>=2.34.2,<3.0.0",
54
+ "tqdm>=4.70.1,<5.0.0",
55
+ "urwid>=5.1.3,<6.0.0",
52
56
  ]
53
57
 
54
- # vLLM 0.25.1 currently selects the PyTorch 2.11.0 family. Keep these
55
- # explicit so a worker install cannot silently combine incompatible torch
56
- # packages when another dependency broadens its requirements.
58
+ # Match vLLM 0.33's published requirements, including torchaudio 2.11.0.
59
+ # Keep these explicit so a worker install cannot silently combine incompatible
60
+ # torch packages when another dependency broadens its requirements.
57
61
  vllm_deps = [
58
- "vllm>=0.25.1,<0.26.0",
59
- "torch==2.11.0",
60
- "torchvision==0.26.0",
62
+ "vllm>=0.33.0,<0.34.0",
63
+ "torch==2.14.0",
64
+ "torchvision==0.29.0",
61
65
  "torchaudio==2.11.0",
62
- "transformers>=5.5.3,<6.0.0",
66
+ "transformers>=5.16.1,<5.21.0",
63
67
  "qwen-vl-utils>=0.0.14,<0.1.0",
64
68
  ]
65
69
 
66
70
  PYTORCH_DEPENDENCIES = [
67
- "torch>=2.11.0",
68
- "torchvision>=0.26.0",
71
+ "torch>=2.14.0",
72
+ "torchvision>=0.29.0",
69
73
  "torchaudio>=2.11.0",
70
74
  ]
71
75
 
72
76
  APPLE_PYTORCH_DEPENDENCIES = [
73
- "torch>=2.13.0",
74
- "torchvision>=0.28.0",
75
- "torchaudio>=2.11.0",
76
- # vLLM Metal publishes Apple Silicon wheels for Python 3.12+.
77
- "vllm-metal==0.1.0; platform_system == 'Darwin' and platform_machine == 'arm64' and python_version >= '3.12'",
77
+ *PYTORCH_DEPENDENCIES,
78
+ "qwen-vl-utils>=0.0.14,<0.1.0",
79
+ # Install the matching vLLM core and Metal wheels from upstream GitHub
80
+ # releases separately; PyPI's vllm-metal 0.1.0 is obsolete.
78
81
  ]
79
82
 
80
83
  extras_require = {
@@ -89,13 +92,13 @@ extras_require = {
89
92
  "rocm": PYTORCH_DEPENDENCIES,
90
93
  "apple": APPLE_PYTORCH_DEPENDENCIES,
91
94
  "dev": [
92
- "pytest>=8.0.0,<9.0.0",
93
- "pytest-asyncio>=1.1.0,<2.0.0",
94
- "pytest-cov>=6.0.0,<7.0.0",
95
- "transformers>=5.5.3,<6.0.0",
96
- "black>=25.0.0,<26.0.0",
97
- "ruff>=0.12.0,<1.0.0",
98
- "mypy>=1.17.0,<2.0.0",
95
+ "pytest>=9.2.0,<10.0.0",
96
+ "pytest-asyncio>=1.5.0,<2.0.0",
97
+ "pytest-cov>=7.1.0,<8.0.0",
98
+ "transformers>=5.16.1,<5.21.0",
99
+ "black>=26.10.1,<27.0.0",
100
+ "ruff>=0.17.5,<1.0.0",
101
+ "mypy>=2.4.1,<3.0.0",
99
102
  ],
100
103
  "all": [*vllm_deps],
101
104
  }
@@ -1,6 +1,6 @@
1
1
  """CaptionFlow - Distributed community captioning system."""
2
2
 
3
- __version__ = "0.5.6"
3
+ __version__ = "0.6.0"
4
4
 
5
5
  from .monitor import Monitor
6
6
  from .orchestrator import Orchestrator
@@ -6,7 +6,7 @@ import os
6
6
  from dataclasses import dataclass, field
7
7
  from datetime import datetime
8
8
  from enum import Enum
9
- from typing import Any, Dict, List, Optional, Tuple
9
+ from typing import Any, Dict, List, Optional, Tuple, Union
10
10
 
11
11
  from PIL import Image
12
12
 
@@ -128,7 +128,9 @@ class Caption:
128
128
  caption_count: int = 1 # Number of captions generated for this item
129
129
  caption: Optional[str] = None
130
130
  captions: Optional[List[str]] = None
131
- outputs: Dict[str, List[str]] = field(default_factory=dict)
131
+ outputs: Dict[str, Union[str, Dict[str, Any], List[Union[str, Dict[str, Any]]]]] = field(
132
+ default_factory=dict
133
+ )
132
134
  quality_score: Optional[float] = None
133
135
  quality_scores: Optional[List[float]] = None
134
136
 
@@ -179,12 +181,13 @@ class ProcessingStage:
179
181
  dtype: Optional[str] = None
180
182
  gpu_memory_utilization: Optional[float] = None
181
183
 
182
- # Optional semantic fallback for empty/refused caption responses. API
183
- # workers can omit the image on the fallback request when an upstream
184
- # provider rejects the image before inference.
184
+ # Caption policy and recovery apply to every inference backend.
185
185
  retry_prompt: Optional[str] = None
186
186
  retry_without_image: bool = False
187
187
  retry_sampling: Optional[Dict[str, Any]] = None
188
+ output_processing: Optional[Dict[str, Any]] = None
189
+ response_format: Optional[Dict[str, Any]] = None
190
+ retry_response_format: Optional[Dict[str, Any]] = None
188
191
 
189
192
 
190
193
  @dataclass
@@ -51,7 +51,7 @@ class Orchestrator:
51
51
  assert (
52
52
  processor_type is not None
53
53
  ), "You must supply processor_type in your orchestrator dataset configuration."
54
- processor_config = ProcessorConfig(processor_type=processor_type, config=config)
54
+ self.processor_config = ProcessorConfig(processor_type=processor_type, config=config)
55
55
 
56
56
  # Initialize processor
57
57
  if processor_type == "webdataset":
@@ -70,7 +70,6 @@ class Orchestrator:
70
70
  caption_buffer_size=storage_config.get("caption_buffer_size", 1000),
71
71
  )
72
72
  self.auth = AuthManager(config.get("auth", {}))
73
- self.processor.initialize(processor_config, self.storage)
74
73
 
75
74
  # Processing configuration
76
75
  self.chunks_per_request = config.get("chunks_per_request", 2)
@@ -132,6 +131,8 @@ class Orchestrator:
132
131
 
133
132
  # Initialize storage
134
133
  await self.storage.initialize()
134
+ # Checkpoints can include results that had not reached a durable flush.
135
+ self.processor.initialize(self.processor_config, self.storage)
135
136
 
136
137
  # Start background tasks
137
138
  asyncio.create_task(self._heartbeat_loop())
@@ -410,7 +411,9 @@ class Orchestrator:
410
411
 
411
412
  elif msg_type == "work_complete":
412
413
  unit_id = data["unit_id"]
413
- self.processor.mark_completed(unit_id, worker_id)
414
+ if self.processor.mark_completed(unit_id, worker_id) is False:
415
+ logger.warning(f"Work unit {unit_id} reported complete before all results arrived")
416
+ return
414
417
  logger.debug(f"Work unit {unit_id} completed by worker {worker_id}")
415
418
  await self._send_activity(f"Chunk completed: {unit_id}")
416
419
 
@@ -428,9 +431,8 @@ class Orchestrator:
428
431
  logger.debug(f"Heartbeat from {worker_id}: {data}")
429
432
 
430
433
  async def _handle_results_submission(self, worker_id: str, data: Dict):
431
- """Process results submission from worker - fires off async task and returns immediately."""
432
- # Fire and forget - process in background
433
- asyncio.create_task(self._process_result_async(worker_id, data))
434
+ """Accept results in connection order before handling a completion message."""
435
+ await self._process_result_async(worker_id, data)
434
436
 
435
437
  async def _process_result_async(self, worker_id: str, data: Dict):
436
438
  """Actually process the result in background."""
@@ -444,6 +446,9 @@ class Orchestrator:
444
446
  shard_name = job_id.shard_id
445
447
  chunk_name = job_id.chunk_id
446
448
 
449
+ if data.get("error") or not data.get("outputs"):
450
+ raise ValueError("A failed result cannot be recorded as completed work")
451
+
447
452
  result = WorkResult(
448
453
  unit_id=data["unit_id"],
449
454
  source_id=shard_name,
@@ -455,11 +460,6 @@ class Orchestrator:
455
460
  processing_time_ms=data.get("processing_time_ms", 0),
456
461
  )
457
462
 
458
- # Let processor handle any custom processing - this updates chunk tracker
459
- # IMPORTANT: Call this BEFORE saving to storage so chunk tracker is updated
460
- # regardless of whether the item is a duplicate
461
- processed = self.processor.handle_result(result)
462
-
463
463
  # Create caption record for storage
464
464
  total_outputs = sum(len(v) for v in result.outputs.values())
465
465
 
@@ -480,7 +480,7 @@ class Orchestrator:
480
480
  caption = Caption(
481
481
  job_id=job_id,
482
482
  dataset=result.dataset,
483
- shard=processed["source_id"],
483
+ shard=result.source_id,
484
484
  chunk_id=chunk_name,
485
485
  item_key=item_key,
486
486
  captions=result.outputs.get("captions", []),
@@ -501,6 +501,11 @@ class Orchestrator:
501
501
  # Save to storage (might skip if duplicate)
502
502
  saved = await self.storage.save_caption(caption)
503
503
 
504
+ # Progress may only follow storage acceptance. A duplicate is also
505
+ # accepted: its caption is already present. Startup reconstructs
506
+ # progress from durable storage if buffered results are lost.
507
+ self.processor.handle_result(result)
508
+
504
509
  # Update contributor stats only if actually saved
505
510
  if saved:
506
511
  contributor = await self.storage.get_contributor(worker_user)
@@ -509,8 +514,11 @@ class Orchestrator:
509
514
  await self.storage.save_contributor(contributor)
510
515
 
511
516
  except Exception as e:
517
+ if data.get("unit_id"):
518
+ self.processor.mark_failed(data["unit_id"], worker_id, "Result acceptance failed")
512
519
  logger.error(
513
- f"Error processing result from {worker_id} for unit {data.get('unit_id', 'unknown')}: {e}",
520
+ f"Error processing result from {worker_id} "
521
+ f"for unit {data.get('unit_id', 'unknown')}: {e}",
514
522
  exc_info=True,
515
523
  )
516
524