43 Commits
Author SHA1 Message Date
Jose Luis Cordova 33c21baff8 Refresh CaptionForge 1.0.2 workflow assets 2026-09-19 19:47:44 -04:00
Jose Luis Cordova b116b08033 Sync CaptionForge 1.0.2 Ollama witness default 2026-09-19 19:16:38 -04:00
Jose Luis Cordova 1bc70ccb2a Polish CaptionForge 1.0.2 release workflows 2026-09-19 17:34:15 -04:00
Jose Luis Cordova 9d5b131b08 Release CaptionForge 1.0.2 2026-09-19 13:08:14 -04:00
Jose Luis Cordova fd5f519aae Harden end-to-end caption cleanup 2026-09-19 13:00:22 -04:00
José Luis Córdova 90753aa7c6 Align replacement helper with boundary-safe default 2026-09-19 09:42:15 -04:00
José Luis Córdova a84724a38a Align replacement helper with boundary-safe default 2026-09-19 09:42:13 -04:00
José Luis Córdova 7fa2871810 Add replace_pairs boundary regression coverage 2026-09-19 09:41:59 -04:00
José Luis Córdova 49080cbae2 Make Ollama replace_pairs boundary-safe 2026-09-19 09:41:41 -04:00
José Luis Córdova 9f76aae7e7 Make Qwen replace_pairs boundary-safe by default 2026-09-19 09:41:35 -04:00
José Luis Córdova 49290e8812 Make Joy replace_pairs boundary-safe by default 2026-09-19 09:41:33 -04:00
José Luis Córdova f7f770f953 Add boundary-safe shared replacement helper 2026-09-19 09:41:13 -04:00
José Luis Córdova a1e613a6ac Add cleanup boundary regression coverage 2026-09-19 09:23:59 -04:00
José Luis Córdova 25425d41b5 Fix Ollama forbidden matching and audit cleanup settings 2026-09-19 09:23:36 -04:00
José Luis Córdova b438fddbc7 Use boundary-safe forbidden matching in Qwen 2026-09-19 09:23:26 -04:00
José Luis Córdova 32a8d4ecc7 Use boundary-safe forbidden matching in Joy 2026-09-19 09:23:24 -04:00
José Luis Córdova af785e3d36 Add boundary-safe shared cleanup matcher 2026-09-19 09:23:06 -04:00
Jose Luis Cordova a373a7e8d9 Restore scoped Joy/Qwen bitsandbytes diagnostics and install dependencies 2026-09-18 14:49:19 -04:00
José Luis Córdova b234e8020c Test explicit reuse of generated dataset inputs 2026-09-18 11:49:20 -04:00
José Luis Córdova 6c111930d5 Allow explicit reuse of generated dataset inputs 2026-09-18 11:49:07 -04:00
José Luis Córdova 31a10619d9 Allow explicit reuse of generated dataset inputs 2026-09-18 11:49:05 -04:00
José Luis Córdova 5ed231eaad Allow explicit reuse of generated dataset inputs 2026-09-18 11:49:02 -04:00
Jose Luis Cordova 435ad73ad6 CaptionForge: add paired dataset export 2026-09-17 20:19:55 -04:00
Jose Luis Cordova 3cf6b3536d CaptionForge: release 1.0.1 2026-09-13 19:03:50 -04:00
Jose Luis Cordova cb5741b604 CaptionForge: harden orchestrator resume and recovery 2026-09-13 17:11:16 -04:00
Jose Luis Cordova 59c4eb31e9 CaptionForge: cap validator image size from planner 2026-09-13 13:09:46 -04:00
Jose Luis Cordova 2a5d62ef8a CaptionForge: improve 1.0.0 project landing page 2026-09-12 14:05:16 -04:00
Jose Luis Cordova 5b1c395a59 CaptionForge: refresh canonical workflows for 1.0.0 2026-09-12 13:24:45 -04:00
Jose Luis Cordova a6a13d389c CaptionForge: relax SHORT length constraint 2026-09-12 13:16:10 -04:00
Jose Luis Cordova e69d9a0dd8 CaptionForge: tighten public Ollama model choices 2026-09-12 10:31:34 -04:00
Jose Luis Cordova 416555fdfb CaptionForge: reconcile remaining release polish 2026-09-12 10:15:26 -04:00
Jose Luis Cordova c26c8a613f CaptionForge: refine Orchestrator output contract 2026-09-12 09:29:17 -04:00
Jose Luis Cordova 093b6908f5 CaptionForge: harden Pass-A source identity and artifacts 2026-09-12 09:00:03 -04:00
Jose Luis Cordova 9b210ad811 CaptionForge: polish documentation and release metadata for 1.0.0 2026-09-11 15:18:52 -04:00
Jose Luis Cordova cd2d0a675f CaptionForge: freeze downstream ownership and production defaults 2026-09-11 00:41:05 -04:00
Jose Luis Cordova 4d2124a8f4 CaptionForge: refine short captions from validated long 2026-09-09 08:10:41 -04:00
Jose Luis Cordova 5568f74dbb CaptionForge: harden seed contract and planner ownership 2026-09-08 22:58:07 -04:00
Jose Luis Cordova dc2d05899b CaptionForge: reconcile planner and capstone contracts 2026-09-07 19:07:26 -04:00
Jose Luis Cordova 3dd61f527c CaptionForge: harden release workflow and output handling 2026-09-07 11:57:05 -04:00
Jose Luis Cordova b09eccbb6e Bump package version to 0.1.1 2026-06-20 22:02:38 -04:00
Jose Luis Cordova 61529735a0 Remove experimental prototypes from public package 2026-06-20 21:57:30 -04:00
Jose Luis Cordova bc45c31450 Update Comfy Registry metadata 2026-06-20 21:05:21 -04:00
Jose Luis Cordova 171c93a7b8 Fix README release notes 2026-06-20 18:08:58 -04:00
62 changed files with 8824 additions and 8907 deletions
+5 -4
View File
@@ -69,7 +69,8 @@ config/.bak/
# CaptionForge local experimental/quarantined code
/nodes/experimental/
/engines/experimental/
# CaptionForge local experimental/quarantined code
/nodes/experimental/
/engines/experimental/
node.zip
# CaptionForge local/private model configuration references
config/captionforge_ollama_models_new.json
config/captionforge_ollama_models_unrestricted.json
config/Ollama_Models_Catalog.md
+34 -21
View File
@@ -1,31 +1,44 @@
Project Folder/File Structure:
# Repository structure
CaptionForge keeps ComfyUI-facing nodes separate from reusable backend code.
Only the explicit production packages and package data listed in
`pyproject.toml` are included in a distribution.
```text
CaptionForge/
├─ __init__.py
├─ __init__.py ComfyUI registration entry point
├─ captionforge_version.py authoritative release version
├─ README.md
├─ pyproject.toml
│
├─ config/
│ └─ captionforge_ollama_models.json user-editable Ollama model tags
├─ nodes/
│ ├─ __init__.py
│ ├─ captionforge_extra_options_CUI_node.py
│ ├─ jlc_captionforge_node.py
│ ├─ jlc_captionforge_pipeline_planner_node.py
| ├─ caption_nodes
│ │ ├─ jlc_joy_caption_node.py
│ │ ├─ jlc_qwen_caption_node.py
│ │ ├─ jlc_ollama_caption_node.py
│
│ ├─ jlc_captionforge_node.py production B/C/D capstone
│ ├─ jlc_captionforge_template_options.py
│ ├─ captionforge_ollama_model_dropdowns.py
│ └─ caption_nodes/
│ ├─ jlc_captionforge_joy_caption_node.py
│ ├─ jlc_captionforge_qwen_caption_node.py
│ └─ jlc_captionforge_ollama_caption_node.py
├─ engines/
│ ├─ __init__.py
│ ├─ captionforge_claim_engine.py
│ ├─ captionforge_model_cache.py
│ ├─ captionforge_pipeline_planner_engine.py
│ ├─ captionforge_prompt_defaults.py
│ ├─ captionforge_caption_prompt_kit.py
│ ├─ captionforge_joy_space_prompt_kit.py
│ ├─ captionforge_model_cache.py
│ ├─ jlc_joy_caption_engine.py
│ ├─ jlc_qwen_caption_engine.py
│
├─ config/
│ ├─ captionforge_ollama_models.json
│
└─ web/
├─ jlc_captionforge_icons.js
└─ ...
│ ├─ captionforge_distiller_engine.py standalone reference/CLI engine
│ └─ captionforge_vlm_validator_engine.py standalone reference/CLI engine
├─ assets/
│ ├─ icons/
│ └─ workflows/ canonical JSON/API JSON/PNG workflows
├─ docs/ release-facing studies
├─ tests/ CPU contract tests
└─ web/ ComfyUI frontend/icon support
```
Local `.backups/`, `.tests/`, `.vscode/`, `.bak/`, `.deprecated/`, and
experimental trees are archival or development-only. They are neither
registered by CaptionForge nor included in the production package.
+400 -430
View File
@@ -1,257 +1,382 @@
# CaptionForge
<p align="center">
<img src="assets/icons/jlc-comfyui-nodes_Logo-0512.png" width="150" alt="CaptionForge logo">
</p>
**Accurate, auditable image captions for LoRA dataset preparation in ComfyUI.**
<h1 align="center">CaptionForge for ComfyUI</h1>
CaptionForge is built around a simple idea: one captioner can be useful, but one captioner is also easy to fool. Instead of asking a single model to describe an image and hoping it gets everything right, CaptionForge can ask multiple independent captioning engines to produce separate “witness accounts” of the same image. Those accounts are then merged by a text-LLM distillation pass that looks for agreement, preserves useful details, and separates likely contradictions or unsupported claims. The resulting draft is checked against the image by a final vision-language model, which acts as the image-aware judge before the final captions are exported.
<p align="center">
<strong>Accurate, auditable multi-model image captions for LoRA dataset preparation.</strong>
</p>
The goal is not magic, and it is not perfection. CaptionForge is meant for automated captioning of large image archives and LoRA training sets where hand-captioning would be too slow, but where the usual hallucinations, omissions, and inconsistencies from a single captioning model are still a problem. The pipeline is intentionally heavier than a normal caption node, so it is best used when caption quality, auditability, and consistency matter enough to justify the extra computation.
The current v0.1.x workflow is tuned primarily for character, fashion, portrait, doll/render, cosplay, pageant, glamour, and style-LoRA datasets, where visible details such as face, hair, eyes, expression, pose, body shape, clothing construction, accessories, colors, materials, lighting, background, framing, and visual style matter.
<p align="center">
<img alt="ComfyUI" src="https://img.shields.io/badge/ComfyUI-Custom%20Nodes-blue">
<img alt="License" src="https://img.shields.io/badge/license-MIT-green">
<img alt="Version" src="https://img.shields.io/badge/version-1.0.2-blue">
</p>
---
<p align="center">
<img src="assets/icons/jlc-comfyui-nodes_Logo-0512.png" width="120">
&nbsp;&nbsp;&nbsp;
<img src="assets/icons/jlc-comfyui-nodes_Logo-Dark-0512.png" width="120">
</p>
## Overview
[![ComfyUI](https://img.shields.io/badge/ComfyUI-Custom%20Nodes-blue)]()
[![License](https://img.shields.io/badge/license-MIT-green)]()
![Status](https://img.shields.io/badge/status-v0.1.x%20preview-orange)
![Version](https://img.shields.io/badge/version-0.1.0-orange)
CaptionForge is a local, model-agnostic captioning framework for building richer, more auditable captions for LoRA dataset preparation inside ComfyUI.
## Starter workflow
The core idea is simple: a single image captioner can be useful, but it should not be treated as authoritative. CaptionForge can collect several independent **Pass A witness captions**, synthesize them with a text LLM, validate the resulting draft against the original image with a vision-language model, and then export three useful caption forms from the same validated semantic result:
A full workflow sample is included as a PNG with embedded ComfyUI workflow metadata:
- `*_long.txt` — authoritative image-validated natural-language caption
- `*_short.txt` — concise AI-compressed natural-language caption intended for modern LoRA training workflows such as FLUX-family training
- `*_taggy.txt` — compact comma-separated caption suited to tag-oriented SD-style training workflows
```text
assets/workflows/CaptionForge_FullWorkflow.png
```
CaptionForge also writes structured JSONL audit records so intermediate evidence, prompts, model settings, and final outputs can be inspected instead of treated as a black box.
> **Current release: CaptionForge 1.0.2.**
> The A/B/C/D semantic pipeline, Planner/Orchestrator authority model, seed contract, and production defaults are frozen for this release.
---
## Sample Workflow
For installation dependencies and Joy Balanced (8-bit) warning handling, see
[Joy 8-bit and clean installation](docs/joy-8bit-installation.md).
<p align="center">
<a href="assets/workflows/CaptionForge_FullWorkflow.png">Download workflow PNG</a>
<a href="assets/workflows/CaptionForge_FullWorkflow_Rel_v1.0.2.png">
<img src="assets/workflows/CaptionForge_FullWorkflow_Rel_v1.0.2.png"
alt="CaptionForge Full Workflow"
width="100%">
</a>
</p>
<p align="center">
<img src="assets/workflows/CaptionForge_FullWorkflow.png" alt="CaptionForge full starter workflow" width="900">
</p>
Canonical workflow files:
A separate JSON export of the same workflow is also included:
- [`CaptionForge_FullWorkflow_Rel_v1.0.2.json`](assets/workflows/CaptionForge_FullWorkflow_Rel_v1.0.2.json) — editable ComfyUI workflow
- [`CaptionForge_FullWorkflow_API_Rel_v1.0.2.json`](assets/workflows/CaptionForge_FullWorkflow_API_Rel_v1.0.2.json) — API-format workflow
- [`CaptionForge_FullWorkflow_Rel_v1.0.2.png`](assets/workflows/CaptionForge_FullWorkflow_Rel_v1.0.2.png) — workflow PNG with embedded metadata
Click the workflow image above to view it at full size.
---
## Why CaptionForge exists
A strong standalone captioner may be enough for many datasets. CaptionForge is for cases where one captioner is not accurate, complete, consistent, or auditable enough.
Different captioning models often notice different details. One may capture face and hair accurately but miss clothing construction. Another may catch materials or accessories but misread pose. A third may notice background or style cues that the others omit.
CaptionForge treats those captions as **witness statements**, not final truth.
The pipeline then:
1. gathers independent witness captions;
2. combines useful evidence into a richer draft;
3. checks that draft against the original image;
4. produces validated LONG, SHORT, and TAGGY forms;
5. records the process in auditable JSONL artifacts.
The goal is not perfect captions. The goal is a more defensible automated captioning process for large datasets where hand-captioning would be impractical.
---
## Production pipeline
```text
assets/workflows/CaptionForge_FullWorkflow.json
```
<p align="center">
<a href="assets/workflows/CaptionForge_FullWorkflow.json">Download workflow JSON</a>
</p>
In ComfyUI, load the workflow by dragging either `CaptionForge_FullWorkflow.png` or `CaptionForge_FullWorkflow.json` onto the canvas.
## Install
Clone CaptionForge into your ComfyUI custom nodes folder:
```bash
git clone https://github.com/Damkohler/CaptionForge.git ComfyUI/custom_nodes/CaptionForge
```
Or copy the repository manually so the folder layout is:
```text
ComfyUI/custom_nodes/CaptionForge/
```
Then restart ComfyUI.
If your ComfyUI environment does not already include the needed Python packages, install CaptionForge dependencies from inside your ComfyUI Python environment. The exact command depends on how your ComfyUI install is managed, but typical options are:
```bash
cd ComfyUI/custom_nodes/CaptionForge
pip install -e .
```
or, if you maintain dependencies manually:
```bash
pip install torch transformers accelerate huggingface-hub pillow numpy safetensors qwen-vl-utils
```
Optional 8-bit loading may require:
```bash
pip install bitsandbytes
```
Ollama-backed stages require a working local Ollama installation and installed Ollama model tags.
Example:
```bash
ollama pull mistral-small:24b
ollama pull gemma4:26b
```
CaptionForge does **not** ship model weights. Joy, Qwen, and Ollama model downloads remain user-controlled.
## What the workflow does
CaptionForge's main pipeline is:
```text
Pass A — raw witness captions
Pass A — witness captions
Joy Caption xN
Qwen Caption xN
optional Ollama VLM Caption xN
Ollama VLM Caption xN
Pass B — text-LLM distillation
combine witness captions
preserve repeated and useful details
separate contradictions and weak claims
build a rich draft caption
Pass B — Distiller
text LLM combines witness evidence
preserves useful repeated and plausible details
builds a rich draft caption
Pass C — image-aware VLM validation
inspect the actual image
keep image-supported details
remove unsupported hallucinations
correct visible errors
produce the authoritative long caption
Pass C — Validator
image-aware VLM inspects the source image
checks the draft against visible evidence
removes unsupported claims where possible
corrects visible errors where possible
produces the authoritative LONG caption
Pass D — deterministic export formatting
write the validated long caption
derive a shorter LoRA-length caption
derive a compact taggy caption
write TXT and JSONL audit records
Pass D — Formatter
text LLM transforms validated LONG into:
SHORT
TAGGY
```
The important distinction is that the expensive semantic work should mostly end at the VLM-validated long caption. The short and taggy outputs are intentionally lighter recipe-style formatting steps derived from that validated caption, not new attempts to reinterpret the image.
The final semantic authority is the image-aware validation pass. Pass D formats that validated result; it is not intended to reinterpret the image from scratch.
## Current status
---
CaptionForge v0.1.0 is a working experimental preview for ComfyUI users and node developers who want to test a multi-pass captioning pipeline.
## Recommended production defaults
It is not presented as a universal replacement for a strong standalone captioner. If JoyCaption, Qwen, Florence, BLIP, WD14, or another captioning tool already gives you exactly what your dataset needs, you may not need CaptionForge. This project is aimed at cases where a single captioner is not accurate, complete, consistent, or auditable enough.
### Pass A — witnesses
Expected v0.1.x realities:
The current recommended witness profile is:
- the workflow is computationally heavy
- large models may be slow
- model choices matter a lot
- output schemas may still evolve
- prompts and defaults may continue to be refined
- not every dataset will benefit equally
- comparison feedback is welcome
- **Joy:** 2 runs per image
- **Qwen:** 1 run per image
- **Ollama:** 1 run per connected Ollama witness node
This a heavy tool. Use it when the extra caption quality and audit trail of large automated jobs are worth the runtime cost.
Multiple Ollama witness nodes may be connected when additional independent VLM caption voices are desired.
## Why use this instead of a standalone captioner?
### Pass B — Distiller
You may want CaptionForge when:
- one captioner notices the face but misses clothing details
- another captioner notices clothing but misreads the pose
- a third captioner catches style or material details the others miss
- you want an LLM to consolidate agreement instead of merely accepting one model's wording
- you want a final VLM to check the draft against the actual image
- you want intermediate JSONL records for debugging and audit
- you want final captions written as sidecars beside the source images
- you need both long natural captions and compact LoRA-style derivatives
The project question is practical:
> Can independent caption witnesses plus text distillation plus image-aware validation produce better dataset captions than a single captioning model alone?
For some datasets, the answer may be yes. For others, a simpler captioner may be enough. CaptionForge is designed to make that comparison visible.
## What CaptionForge tries to optimize
CaptionForge currently favors captions that are:
- rich enough for LoRA training
- visually grounded
- less hallucinated than unvalidated text-only synthesis
- explicit about visible, trainable details
- auditable through JSONL records
- locally runnable
- model-agnostic enough to improve as better captioners, distillers, and validators become available
Useful caption details often include:
- subject type and visible style
- face shape and facial traits
- hair color and hairstyle
- eye color and makeup as separate details
- expression and pose
- hands and body position
- body shape and visible proportions when relevant
- clothing construction, layers, fit, and materials
- accessories, jewelry, nails, props, and distinctive details
- colors, textures, lighting, background, framing, and crop
Visible glamour, swimwear, lingerie, revealing clothing, cleavage, side openings, exposed midriff, or similar styling may be described neutrally when it is actually visible and relevant to the dataset. CaptionForge prompts should not invent hidden anatomy, unseen clothing, explicit acts, or contradicted details.
## Active node families
Node categories are being normalized under:
Default model:
```text
Captioning/CaptionForge
mistral-small:24b
```
with active caption nodes under:
Default production settings:
```text
Captioning/CaptionForge/Caption Nodes
source cap = 1536
num_predict = 3096
temperature = 0.24
top_p = 0.90
top_k = 60
seed = -1
```
### JLC CaptionForge Pipeline Planner
### Pass C — Validator
The central planning node for normal runs.
Default model:
```text
gemma4:26b
```
Default production settings:
```text
num_predict = 2112
temperature = 0.00
top_p = 0.92
top_k = 80
seed = -1
```
### Pass D — Formatter
Default model:
```text
mistral-small:24b
```
Default production settings:
```text
num_predict = 3200
temperature = 0.12
top_p = 0.88
top_k = 50
seed = -1
```
`seed = -1` means CaptionForge does **not** send an explicit seed to Ollama. This is intentionally unseeded behavior and is not guaranteed to be repeatable.
---
## LONG, SHORT, and TAGGY
### LONG
The authoritative image-validated natural-language caption.
Use LONG when maximum descriptive detail is useful for:
- review
- auditing
- experimentation
- highly descriptive datasets
### SHORT
A concise semantic compression of validated LONG.
SHORT is intended as the best starting point for many FLUX-family LoRA workflows.
The Formatter aims for **roughly 100 words, give or take**. CaptionForge does not hard-cut a valid SHORT at a word or character boundary; the caption is allowed to finish naturally.
### TAGGY
A compact comma-separated caption derived from validated LONG.
TAGGY is intended primarily for SD-style or other tag-oriented training workflows.
CaptionForge writes all three forms so a dataset can later be trained or compared using different caption styles without recaptioning the source images.
---
## Main nodes
### CaptionForge Pipeline Planner
The Planner is the control center for the full workflow.
It coordinates:
- input image path or direct image passthrough
- recursive folder traversal
- image or folder input
- recursive traversal
- filename glob filtering
- output directory
- output location
- run name
- overwrite behavior
- Pass A witness run counts
- seed schedules
- sampling schedules
- max image size
- max token budget
- LoRA trigger word
- user caption anchor
- distiller settings
- validator settings
- final export settings
- derived JSONL/TXT/config paths
- witness run counts
- Pass A seed scheduling
- shared Ollama configuration
- Distiller settings
- Validator settings
- Formatter settings
- audit options
### JLC CaptionForge
In the full workflow, **the Planner is authoritative**.
The main capstone/orchestration node.
The Planner now owns `forbidden_phrases` and `replace_pairs` for the full workflow. Joy, Qwen, Ollama Caption, and the Orchestrator retain their corresponding standalone inputs; when a Planner is connected, its values take precedence.
It consumes Pass A raw caption records, runs the distillation and validation stages, and exports final captions. The VLM-validated natural paragraph is the authoritative long caption. Formatting stages should not blindly rewrite that natural caption.
### CaptionForge Joy Caption
### JLC CaptionForge Joy Caption
Python/Hugging Face JoyCaption-family Pass A witness.
Python/Hugging Face JoyCaption/LLaVA-family Pass A witness.
Joy remains a strong first-class caption source and is commonly run twice per image to gain useful diversity.
Joy is treated as a first-class CaptionForge caption source and is often one of the strongest raw caption witnesses.
### JLC CaptionForge Qwen Caption
### CaptionForge Qwen Caption
Python/Hugging Face Qwen-family Pass A witness.
Qwen is useful as a second independent captioning voice, especially when its behavior complements Joy. Optional 8-bit loading may be available where supported.
Qwen provides a complementary caption voice and is commonly run once per image in the production profile.
### JLC CaptionForge Ollama Caption
### CaptionForge Ollama Caption
Ollama-backed VLM Pass A witness.
This node delegates image-caption generation to a local Ollama server rather than loading Hugging Face/PyTorch weights inside ComfyUI. It can use configured Ollama VLM tags such as:
Multiple Ollama witness nodes may be chained to contribute captions from different local VLMs.
### CaptionForge Template Options
Provides shared captioning guidance for witness nodes so they can emphasize LoRA-relevant visual detail without forcing every backend into an identical prompt implementation.
### CaptionForge Orchestrator
The Orchestrator performs the downstream B/C/D work.
It:
- consumes Pass A witness records;
- runs Distiller synthesis;
- runs image-aware validation;
- produces LONG, SHORT, and TAGGY;
- writes final sidecars and audit artifacts.
The Orchestrator exposes exactly five public outputs:
```text
long_captions
short_captions
taggy_captions
final_records
status
```
In standalone mode, Orchestrator-local B/C/D settings are authoritative.
When connected to the Pipeline Planner, Planner settings take precedence.
Cleanup uses boundary-safe whole-word and phrase matching, so a rule such as `old` does not alter `bold`, `holding`, or `gold`. The effective rules are enforced through final LONG, SHORT, and TAGGY generation so downstream models cannot silently reintroduce forbidden or superseded wording.
---
## Output files
For each successfully processed image, CaptionForge writes:
```text
*_long.txt
*_short.txt
*_taggy.txt
```
The run also writes structured JSON/JSONL records for intermediate and final pipeline stages.
`A_RAW_CAPTIONS.jsonl` is the authoritative Pass A witness ledger. Planned CaptionForge execution does **not** create redundant individual Joy/Qwen/Ollama witness `.txt` files.
Optional prompt and raw-response auditing can be enabled when deeper inspection is useful.
---
## Source-image identity
CaptionForge keeps human-readable names separate from collision-safe canonical identity.
For normal folder input:
```text
image = exact source filename stem
image_key = relative dataset path including extension
```
Example:
```text
People/Session 1/_DSC2094-Edit-2.jpg
image = _DSC2094-Edit-2
image_key = People/Session 1/_DSC2094-Edit-2.jpg
```
This preserves leading underscores, spaces, punctuation, capitalization, relative directory hierarchy, and the extension in the canonical key.
It also prevents collisions between files such as:
```text
Set A/photo.jpg
Set A/photo.png
Set B/photo.jpg
```
Canonical JSON/JSONL identity uses portable `/` separators.
---
## Optional IMAGE input
Direct ComfyUI IMAGE input can be used independently or together with Planner folder input.
Filename-less optional tensors receive deterministic names such as:
```text
comfy_image_0000
```
and collision-safe canonical keys such as:
```text
captionforge-optional-image://comfy_image_0000.png
```
The Orchestrator materializes these images into the run's `opt_images` directory so the generated sidecars can be matched back to their corresponding images.
---
## Ollama model choices
The public dropdown configuration lives at:
```text
config/captionforge_ollama_models.json
```
Current public choices include:
### Distiller / Formatter
```text
mistral-small:24b
gpt-oss:20b
VladimirGav/gemma4-26b-16GB-VRAM-Uncensored
```
### Validator / Ollama witness
```text
gemma4:26b
@@ -259,266 +384,111 @@ qwen3.6:35B-A3B
huihui_ai/gemma-4-abliterated:26b
```
Its purpose is to provide access to other raw-caption witness alternatives. It's function is parallel to the Joy Caption and Qwen Caption nodes, and should not be confused with the later VLM validator/capstone role.
### JLC CaptionForge Template Options
Shared prompt-option sidecar for caption nodes.
Template Options let one sidecar node feed consistent LoRA-relevant prompt modifiers into Joy, Qwen, Ollama, and later caption witnesses without duplicating the same option widgets on every caption node.
## Model and memory behavior
CaptionForge uses two model ecosystems:
1. **Python / Hugging Face model folders** for Joy and Qwen witness engines.
2. **Ollama models** for text-LLM distillation, image-aware VLM validation, optional formatting, and Ollama-backed caption witnesses.
Joy and Qwen use Python/Hugging Face engines that integrate with the CaptionForge process-local model cache. Those engines manage Python model residency, reuse, and eviction before loading heavyweight caption models.
Ollama-facing stages are different. Ollama models live in the Ollama daemon, not inside the CaptionForge Python model cache. Before handing work to Ollama, the Ollama Caption node and the CaptionForge capstone clear any resident CaptionForge Python/HF caption models if needed. After that handoff, Ollama owns Ollama model residency.
In short:
The defaults remain:
```text
Joy/Qwen engines:
manage Python-hosted caption models through captionforge_model_cache
Ollama Caption and CaptionForge capstone:
clear Python-hosted models before calling the Ollama daemon
Ollama daemon:
owns Ollama model loading and residency
Distiller -> mistral-small:24b
Validator -> gemma4:26b
Formatter -> mistral-small:24b
Caption -> gemma4:26b
```
## Model locations
Custom Ollama model tags can also be exposed through the node configuration.
Large model weights are intentionally not stored in this repository.
CaptionForge does **not** ship model weights.
Python-based witness models are expected under ComfyUI model folders, for example:
---
## Model handoff and memory management
CaptionForge coordinates Python/Hugging Face caption models and Ollama-backed stages within the same workflow.
Joy and Qwen models are evicted before Ollama handoff when required, reducing stale GPU-memory pressure during long mixed-engine runs.
Large local models may take significant time to load, especially on first use. Later calls are generally faster once the selected backend models are resident.
CaptionForge is intentionally heavier than a one-model caption node. The design target is caption quality and auditability rather than minimum inference count.
---
## Seed behavior
### Pass A
The Pipeline Planner owns Pass A seed scheduling.
The same planned run-seed schedule is reused across images.
### Pass B / C / D
Distiller, Validator, and Formatter each use one stage seed for the run.
For Ollama-backed stages:
```text
ComfyUI/models/LLM/JLC_JoyCaption/
ComfyUI/models/LLM/JLC_QwenCaption/
seed = -1
```
Ollama models must be installed and runnable through Ollama outside this repository.
means no explicit seed is sent.
CaptionForge does not require every supported backend to be installed for every workflow. Users can test smaller subsets first.
Use explicit seeds when reproducibility is required.
## Ollama model dropdown configuration
---
The file:
## Running the full workflow
1. Load [`CaptionForge_FullWorkflow_Rel_v1.0.2.json`](assets/workflows/CaptionForge_FullWorkflow_Rel_v1.0.2.json).
2. Select an image or dataset folder in the **Pipeline Planner**.
3. Select the output folder and run name.
4. Start with the recommended witness counts:
- Joy: 2
- Qwen: 1
- Ollama: 1 per connected Ollama witness
5. Leave Distiller, Validator, and Formatter settings at their defaults unless you have a reason to experiment.
6. Queue the workflow.
CaptionForge will produce the final caption sidecars and preserve the run's structured audit trail.
The optional IMAGE socket is intended for a quick single-image workflow. For multiple images, use **Input - image path** with a folder. CaptionForge 1.0.2 does not claim heterogeneous mixed-aspect IMAGE-list support: generic ComfyUI IMAGE batchers may resize or crop images to a common tensor shape before CaptionForge receives them. Native heterogeneous image-list handling is deferred to future/v2 work.
---
## Validation for 1.0.1
The final 1.0.1 release-preparation pass included:
- Planner ownership/default tests
- Pass A artifact and source-identity tests
- Orchestrator output-contract tests
- SHORT/Formatter behavior tests
- configuration and release-polish tests
- canonical workflow structural verification
- full real-model smoke testing with Joy, Qwen, Ollama, Distiller, Validator, and Formatter
The final audited smoke test processed both Planner folder input and optional IMAGE input through the complete pipeline with:
```text
config/captionforge_ollama_models.json
final_ok = 6
final_failed = 0
```
defines user-editable Ollama model tags for dropdowns used by distiller, validator, formatter, and Ollama caption-witness nodes.
The final release-preparation verification completed with **51 focused tests passing**.
Example:
---
```json
{
"distiller_models": [
"mistral-small:24b",
"VladimirGav/gemma4-26b-16GB-VRAM-Uncensored",
"deepseek-r1:32b",
"tarruda/neuraldaredevil-8b-abliterated:fp16",
"gpt-oss:20b"
],
"validator_models": [
"gemma4:26b",
"qwen3.6:35B-A3B",
"huihui_ai/gemma-4-abliterated:26b"
],
"format_models": [
"mistral-small:24b",
"VladimirGav/gemma4-26b-16GB-VRAM-Uncensored",
"gpt-oss:20b",
"deepseek-r1:32b"
],
"caption_models": [
"gemma4:26b",
"qwen3.6:35B-A3B",
"huihui_ai/gemma-4-abliterated:26b"
],
"defaults": {
"distiller_model": "mistral-small:24b",
"validator_model": "gemma4:26b",
"format_model": "mistral-small:24b",
"caption_model": "gemma4:26b"
},
"include_custom": true
}
```
## Caption quality expectations
Terminology:
CaptionForge reduces many common single-captioner failure modes, but it does not make captioning perfect.
```text
distiller_model text-only LLM for Pass B distillation
validator_model image-aware VLM for Pass C validation
format_model text-only LLM for formatting/taggy conversion when used
caption_model Ollama-backed Pass A image-caption witness model
```
Vision-language models can still miss subtle colors, confuse nearby objects, preserve a plausible but incorrect consensus, omit details, or phrase the same visible fact differently between runs.
Values should be concrete Ollama model tags used exactly as written.
For large LoRA datasets, the goal is to improve the bulk quality and auditability of generated captions enough that remaining small errors are manageable.
## Output layout
---
CaptionForge writes auditable run artifacts and final sidecars during planned runs.
## License
A typical planned run uses this structure:
CaptionForge is released under the **MIT License**.
```text
<output_root>/
<run_name>__working/
<run_name>__A_RAW_CAPTIONS.jsonl
<run_name>__B_DISTILL.jsonl
<run_name>__B_DISTILL_readable.jsonl
<run_name>__B_DISTILL_readable.json
<run_name>__B_DISTILL_prompts.jsonl
<run_name>__C_VLM_VALIDATED.jsonl
<run_name>__C_VLM_VALIDATED_readable/
<run_name>__C_VLM_VALIDATOR_prompts.jsonl
<run_name>__D_FINAL_EXPORT.jsonl
<run_name>__output_paths.json
<run_name>__run_config.json
images/
opt_images/
comfy_image_0000.png
comfy_image_0001.png
```
See [`LICENSE`](LICENSE) for details.
Folder-input images keep their source locations. Optional direct `IMAGE` inputs are saved into the run working area under:
```text
<output_root>/<run_name>__working/images/opt_images/
```
Final caption sidecars are written beside the resolved source image. For folder-input images, that means beside the original image. For optional direct images, that means beside the saved optional image inside `working/images/opt_images/`.
Final sidecars currently include:
```text
<image_stem>_long.txt
<image_stem>_short.txt
<image_stem>_taggy.txt
```
Meaning:
```text
_long.txt the authoritative VLM-validated natural caption
_short.txt a shorter LoRA-length caption derived from the long caption
_taggy.txt a compact comma-separated taggy caption derived from the long caption
```
Long captions are intentional in v0.1.x. The current release-candidate strategy favors preserving visible, trainable detail in the validated long caption, then deriving shorter and taggy outputs from that result.
Exact JSONL schemas may evolve during the preview phase.
## Dependencies
Python dependencies are declared in `pyproject.toml` where applicable.
Typical local use may involve:
```text
torch
transformers
accelerate
huggingface-hub
pillow
numpy
safetensors
qwen-vl-utils
```
Optional quantization support may involve:
```text
bitsandbytes
```
Ollama-backed stages require a working local Ollama installation and installed Ollama model tags.
## Hardware notes
CaptionForge is designed for local workflows, but strong results may require large local models.
Practical performance depends on:
- GPU VRAM
- system RAM
- model size
- quantization mode
- Ollama version
- context length
- image size
- number of Pass A witness runs
- whether models are kept loaded or unloaded between runs
The author's active development environment includes an RTX 4090 Laptop GPU with 16 GB VRAM. Larger models may be slow, may require careful quantization, or may need more capable hardware.
## Experimental branches
Some experimental or unsupported code may exist in the repository for future A/B testing or research.
Experimental branches should be:
- clearly labeled
- kept out of the normal ComfyUI registration path
- not imported by `__init__.py`
- not shown as mainline nodes unless deliberately enabled
- treated as unsupported starting points rather than stable user features
The active public workflow should be the main Planner → Pass A witnesses → Distiller → VLM Validator → Export path.
## Development principles
CaptionForge currently prioritizes:
- local execution
- auditable intermediate records
- JSONL sidecars
- reusable engines separated from ComfyUI node wrappers
- planner-driven workflows
- model cache and VRAM hygiene
- strong defaults for LoRA captioning
- explicit prompt roles
- model-agnostic backends
- visible, trainable detail over generic caption prose
- practical feedback from real datasets
## Feedback wanted
Useful feedback includes:
- comparisons against standalone JoyCaption, Qwen, or other captioners
- examples where CaptionForge improves caption quality
- examples where CaptionForge makes captions worse
- hallucination reports
- missed-detail reports
- model recommendations
- prompt improvements
- broken node reports
- workflow usability feedback
- VRAM/performance observations
- JSONL/audit trail suggestions
Please include enough context to reproduce the issue or evaluate the result: selected nodes, model tags, relevant settings, whether the run used direct IMAGE input or a folder path, and a small sample of generated captions when possible.
## Attribution & License
Concept and implementation by **J. L. Córdova**, with development assistance from **ChatGPT (OpenAI)**.
CaptionForge's Joy/template-option workflow is locally adapted and was inspired in part by the practical template interface pattern used by the public JoyCaption Beta One Hugging Face Space:
```text
https://huggingface.co/spaces/fffiloni/JoyCaption-Beta-One
```
Copyright (c) 2026 J. L. Córdova
Released under the **MIT License**. See [`LICENSE`](./LICENSE) for details.
+9 -21
View File
@@ -8,11 +8,9 @@ CaptionForge — ComfyUI Package Entry Point
- Repository
https://github.com/Damkohler/CaptionForge
- CaptionForge focuses on practical dataset-captioning infrastructure for
LoRA dataset preparation, using multi-engine caption generation, JSONL
audit trails, claim extraction and refinement, text-LLM distillation,
image-aware VLM validation, and consensus-oriented caption improvement
to produce grounded, auditable training captions.
- CaptionForge 1.0 provides a production A/B/C/D caption pipeline for LoRA
dataset preparation: independent witness captions, text-LLM synthesis,
image-aware VLM validation, SHORT/TAGGY formatting, and JSONL audit trails.
- Package Purpose
- This file is the ComfyUI registration entry point for the CaptionForge
@@ -27,12 +25,8 @@ CaptionForge — ComfyUI Package Entry Point
• display the nodes in the Add Node menu
• register the package as a unified CaptionForge node collection
- The package currently registers:
• JLC Qwen Caption
• JLC Joy Caption
• JLC Qwen Caption (Lite)
• JLC Joy Caption (Lite)
• JLC CaptionForge Claim Extractor
- The package registers the Pipeline Planner, Template Options helper,
Joy/Qwen/Ollama Pass-A caption nodes, and the CaptionForge Orchestrator.
- Package Structure
- CaptionForge keeps ComfyUI-facing node wrappers separate from reusable
@@ -46,12 +40,8 @@ CaptionForge — ComfyUI Package Entry Point
export, model cache behavior, and Pass B claim extraction.
- Web / Icon Assets
- JavaScript/icon registration is intentionally disabled in this entry point
for now to avoid frontend branding conflicts while the package structure
stabilizes.
- A future version may re-enable WEB_DIRECTORY after confirming that the
CaptionForge frontend assets do not conflict with other JLC node packages.
- ``WEB_DIRECTORY`` exposes the package's current ComfyUI frontend assets.
Node registration remains entirely explicit through the mappings below.
- Attribution & License
- Concept and implementation by **J. L. Córdova**
@@ -136,12 +126,10 @@ NODE_DISPLAY_NAME_MAPPINGS.update(CAPTIONFORGE_DISPLAY_NAME_MAPPINGS)
WEB_DIRECTORY = "./web"
CAPTIONFORGE_ICON = "⚒"
CAPTIONFORGE_PREFIX = f"{CAPTIONFORGE_ICON} CaptionForge"
print(f"{CAPTIONFORGE_PREFIX} loaded ({len(NODE_CLASS_MAPPINGS)} nodes)")
print(f"CaptionForge loaded ({len(NODE_CLASS_MAPPINGS)} nodes)")
__all__ = [
"NODE_CLASS_MAPPINGS",
"NODE_DISPLAY_NAME_MAPPINGS",
"WEB_DIRECTORY",
]
]
File diff suppressed because it is too large Load Diff
Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.8 MiB

@@ -0,0 +1,403 @@
{
"51": {
"inputs": {
"model": "llama-joycaption-beta-one-hf-llava",
"memory_mode": "Balanced (8-bit)",
"keep_loaded": true,
"caption_template_mode": true,
"caption_type": "JLC LoRA Literal",
"caption_length": "any",
"custom_prompt_mode": false,
"prompt_preset": "default_literal",
"system_prompt": "You are a helpful image-captioning assistant. Describe only what is visible in the image. Do not invent unseen context.",
"custom_prompt": "",
"max_new_tokens": 384,
"temperature": 0.75,
"top_p": 0.9,
"top_k": 50,
"repetition_penalty": 1,
"max_size": 1024,
"forbidden_phrases": "",
"replace_pairs": "",
"download_probe_only": false,
"image": [
"148",
0
],
"pipeline_plan": [
"148",
1
],
"template_options": [
"55",
0
]
},
"class_type": "JLC_CaptionForgeJoy",
"_meta": {
"title": " JLC CaptionForge Joy Caption"
}
},
"52": {
"inputs": {
"model": "Qwen2.5-VL-7B-NSFW-Caption-V3-abliterated",
"qwen_quantization": "Balanced (8-bit)",
"keep_loaded": true,
"caption_template_mode": true,
"caption_type": "LoRA Literal",
"caption_length": "any",
"custom_prompt_mode": false,
"prompt_preset": "default_literal",
"system_prompt": "You are a helpful image-captioning assistant. Describe only what is visible in the image. Do not invent unseen context.",
"custom_prompt": "",
"max_new_tokens": 384,
"temperature": 0.75,
"top_p": 0.9,
"top_k": 50,
"repetition_penalty": 1.08,
"max_size": 1024,
"forbidden_phrases": "",
"replace_pairs": "",
"download_probe_only": false,
"image": [
"51",
0
],
"pipeline_plan": [
"51",
1
],
"template_options": [
"51",
2
]
},
"class_type": "JLC_CaptionForgeQwen",
"_meta": {
"title": " JLC CaptionForge Qwen Caption"
}
},
"55": {
"inputs": {
"name_input__replacement_for_name": "",
"option_01__character_name_required": false,
"option_02__literal_visible_details_only": true,
"option_03__no_roleplay_or_dialogue": true,
"option_04__no_ambiguous_wording": true,
"option_05__avoid_meta_lead_ins": true,
"option_06__describe_facial_features": true,
"option_07__eye_color_vs_makeup": true,
"option_08__describe_makeup_separately": true,
"option_09__describe_expression": true,
"option_10__hair_color_and_style": true,
"option_11__skin_tone_and_texture": true,
"option_12__body_shape_and_proportions": true,
"option_13__visible_anatomy_and_coverage": true,
"option_14__stylized_material_traits": true,
"option_15__clothing_pieces_clearly": true,
"option_16__clothing_materials_textures": true,
"option_17__precise_colors": false,
"option_18__accessories_and_jewelry": false,
"option_19__pose_and_stance": true,
"option_20__hands_and_arms": true,
"option_21__framing_shot_distance": true,
"option_22__viewing_angle": true,
"option_23__background_and_setting": true,
"option_24__lighting": false,
"option_25__image_style_medium": true,
"option_26__watermark_only_if_present": false,
"option_27__visible_text_if_relevant": true,
"option_28__sfw_suggestive_nsfw_rating": false,
"option_29__neutral_sensual_details": true,
"option_30__vulgar_wording_allowed": true,
"option_31__no_polite_euphemisms": true
},
"class_type": "JLC_CaptionForgeExtraOptions",
"_meta": {
"title": " JLC CaptionForge Template Options"
}
},
"56": {
"inputs": {
"model": "custom",
"custom_model_tag": "starnodes/qwen3.8-vl-27b-promptrefiner-abliterated",
"ollama_url": "http://127.0.0.1:11434",
"keep_loaded": true,
"caption_template_mode": true,
"caption_type": "Descriptive",
"caption_length": "any",
"custom_prompt_mode": false,
"prompt_preset": "female_character_lora",
"system_prompt": "You are a multimodal image captioning engine for female character NSFW LoRA dataset preparation.\n\nYour job is to inspect the image and write one strong standalone NSFW caption. Use sexual explicit language without euphemisms. When clearly visible, describe anatomy or actions in explicit terms, using words like pussy, blowjob, masturbating, dildo, fingering, spread, pubic hair, shaved pussy, horny, big breasts, anal, or similar. Do not moralize. No censorship.\n\nRules:\n- Describe only visible image content.\n- Output exactly one paragraph and nothing else.\n- Do not output analysis, reasoning, notes, bullet points, labels, or a thinking trace.\n- Do not roleplay, address the user, ask follow-up questions, or include safety disclaimers.\n- If a detail is uncertain, omit it rather than hedge.\n- Use direct dataset-caption language, not chatty commentary.",
"custom_prompt": "Write one clean caption for a female character LoRA dataset.\n\nCaption priorities:\n- Be specific rather than generic.\n- Prioritize female-character-LoRA-relevant visual traits, preserving NSFW content and explicit language: subject type, pose, facial structure, expression, hair color and hairstyle, eye color, makeup color as distinct from eye color, lip shape and fullness, skin appearance, neck length, shoulder line, bust/chest shape, waist definition, hip width, body proportions, silhouette, legs, clothing construction, accessories, dominant colors, lighting, background, framing, and visual style.\n- Include visible anatomy and figure information when clear and relevant. Do not omit body-shape or facial-structure details merely because they are sensual, revealing, or anatomy-related.\n- If the image presents glamour, sensuality, revealing clothing, lingerie, swimwear, cleavage, underboob, sideboob, bare midriff, exposed hips, thighs, or other sexualized styling, describe that directly and matter-of-factly when visible.\n- Use strong, open dataset-caption language rather than timid euphemisms, but do not turn the caption into roleplay, commentary, or erotic prose.\n- Do not invent nudity, explicit acts, hidden anatomy, or hidden clothing details.\n- Do not force photo, render, doll, anime, realistic, or stylized language unless supported by the image.\n- Avoid meta phrases like \"this image shows\" or \"the picture depicts.\"\n- One paragraph only.\n\nFinal caption:",
"max_new_tokens": 6000,
"temperature": 0.9,
"top_p": 0.6,
"top_k": 80,
"repetition_penalty": 1.03,
"max_size": 1024,
"request_timeout_seconds": 1800,
"forbidden_phrases": "",
"replace_pairs": "",
"download_probe_only": false,
"image": [
"52",
0
],
"pipeline_plan": [
"52",
1
],
"template_options": [
"52",
2
]
},
"class_type": "JLC_CaptionForgeOllamaCaption",
"_meta": {
"title": " JLC CaptionForge Ollama Caption"
}
},
"90": {
"inputs": {
"model": "custom",
"custom_model_tag": "huihui_ai/qwen3.6-abliterated:35b",
"ollama_url": "http://127.0.0.1:11434",
"keep_loaded": true,
"caption_template_mode": true,
"caption_type": "Descriptive",
"caption_length": "any",
"custom_prompt_mode": false,
"prompt_preset": "female_character_lora",
"system_prompt": "You are a multimodal image captioning engine for female character NSFW LoRA dataset preparation.\n\nYour job is to inspect the image and write one strong standalone NSFW caption. Use sexual explicit language without euphemisms. When clearly visible, describe anatomy or actions in explicit terms, using words like pussy, blowjob, masturbating, dildo, fingering, spread, pubic hair, shaved pussy, horny, big breasts, anal, or similar. Do not moralize. No censorship.\n\nRules:\n- Describe only visible image content.\n- Output exactly one paragraph and nothing else.\n- Do not output analysis, reasoning, notes, bullet points, labels, or a thinking trace.\n- Do not roleplay, address the user, ask follow-up questions, or include safety disclaimers.\n- If a detail is uncertain, omit it rather than hedge.\n- Use direct dataset-caption language, not chatty commentary.",
"custom_prompt": "Write one clean caption for a female character LoRA dataset.\n\nCaption priorities:\n- Be specific rather than generic.\n- Prioritize female-character-LoRA-relevant visual traits, preserving NSFW content and explicit language: subject type, pose, facial structure, expression, hair color and hairstyle, eye color, makeup color as distinct from eye color, lip shape and fullness, skin appearance, neck length, shoulder line, bust/chest shape, waist definition, hip width, body proportions, silhouette, legs, clothing construction, accessories, dominant colors, lighting, background, framing, and visual style.\n- Include visible anatomy and figure information when clear and relevant. Do not omit body-shape or facial-structure details merely because they are sensual, revealing, or anatomy-related.\n- If the image presents glamour, sensuality, revealing clothing, lingerie, swimwear, cleavage, underboob, sideboob, bare midriff, exposed hips, thighs, or other sexualized styling, describe that directly and matter-of-factly when visible.\n- Use strong, open dataset-caption language rather than timid euphemisms, but do not turn the caption into roleplay, commentary, or erotic prose.\n- Do not invent nudity, explicit acts, hidden anatomy, or hidden clothing details.\n- Do not force photo, render, doll, anime, realistic, or stylized language unless supported by the image.\n- Avoid meta phrases like \"this image shows\" or \"the picture depicts.\"\n- One paragraph only.\n\nFinal caption:",
"max_new_tokens": 6000,
"temperature": 0.9,
"top_p": 0.6,
"top_k": 80,
"repetition_penalty": 1.03,
"max_size": 1024,
"request_timeout_seconds": 1800,
"forbidden_phrases": "",
"replace_pairs": "",
"download_probe_only": false,
"image": [
"56",
0
],
"pipeline_plan": [
"56",
1
],
"template_options": [
"56",
2
]
},
"class_type": "JLC_CaptionForgeOllamaCaption",
"_meta": {
"title": " JLC CaptionForge Ollama Caption"
}
},
"147": {
"inputs": {
"Input - captions JSONL": "",
"Input - image path": "",
"Input - include caption families": "joy,qwen,ollama",
"Input - max captions per family": 5,
"Input - max total captions": 20,
"Output - folder": "",
"Output - run name": "captionforge_run",
"Output - overwrite outputs": true,
"Ollama - URL": "http://127.0.0.1:11434",
"Ollama - keep loaded": true,
"Ollama - request timeout seconds": 1800,
"LoRA - trigger word": "",
"LoRA - user caption anchor": "",
"Cleanup - forbidden phrases": "",
"Cleanup - replace pairs": "",
"Fat Draft - model": "mistral-small:24b",
"Fat Draft - custom Ollama model": "",
"Fat Draft - prompt": "/no_think\n\nYou are a detail-preserving caption merger for LoRA dataset preparation.\n\nYou receive multiple captions of the same image. You do NOT see the image.\n\nTask:\nMerge all non-contradictory caption details into one deliberately over-complete draft caption.\n\nRules:\n- Do not validate against the image.\n- Do not decide that details are false just because they appear once.\n- Do not summarize aggressively.\n- Preserve concrete details from all captions.\n- Split contradictions by choosing cautious wording or listing the alternative only when needed.\n- Prefer specific visual language over generic language.\n- Keep visible body, clothing, material, accessory, color, pose, lighting, style, and framing details.\n- Preserve doll-like, glossy/plastic-like, material, garment-construction, body-shape, and facial-feature details when present.\n- Use neutral dataset-caption language, including visible sensual styling or revealing clothing when present.\n- Do not add details absent from the captions.\n- Treat subject names or trigger-like identity tokens as optional identity labels. Preserve them only when they appear consistently in the captions; do not let them replace visible description.\n- Output only one paragraph, no notes, no JSON.",
"Fat Draft - max caption chars": 1536,
"Fat Draft - max new tokens": 3096,
"Fat Draft - temperature": 0.24,
"Fat Draft - top p": 0.9,
"Fat Draft - top k": 60,
"Validator - model": "gemma4:26b",
"Validator - custom Ollama model": "",
"Validator - system prompt": "/no_think\nYou are a direct image validation engine. Inspect the image and answer only with the requested caption.",
"Validator - prompt": "/no_think\n\nLook at the image and validate this draft caption.\n\nTask:\nReturn a corrected caption paragraph that keeps only image-supported details.\n\nRules:\n- Output only the corrected caption.\n- One paragraph.\n- No reasoning, no notes, no JSON.\n- Keep all true visible details from the draft.\n- Delete unsupported details.\n- Correct small visible errors.\n- Do not add new details unless needed to correct an error already present.\n- Preserve useful LoRA details: subject, face, hair, eyes, makeup, lips, skin texture, pose, body shape, outfit, accessories, materials, colors, lighting, background, framing, and visual style.\n- Visible sensual styling, revealing clothing, cleavage, thighs, bare skin, swimwear, lingerie, or body-shape details may be described neutrally when present.\n- Do not invent hidden anatomy, unseen clothing, explicit acts, or details contradicted by the image.",
"Validator - max new tokens": 2112,
"Validator - max image size": 1024,
"Validator - temperature": 0,
"Validator - top p": 0.92,
"Validator - top k": 80,
"Formatter - model": "mistral-small:24b",
"Formatter - custom Ollama model": "",
"Formatter - prompt": "/no_think\n\nYou are a LoRA caption format converter. The validated paragraph is your only source of truth.\n\nOutput exactly two labeled lines:\n\nSHORT: <a concise natural-language caption, typically around 100 words, that preserves all LoRA-useful validated details>\n\nTAGGY: <one compact comma-separated caption>\n\nSHORT must preserve the image's distinctive training identity across the whole source:\n1. subject, defining face/hair/body traits, and every major outfit piece/material;\n2. pose/action and key accessories or unusual visible details;\n3. setting, lighting, framing, and visual medium/style.\n\nOmit a category only when absent. Use only source details; never add, infer, euphemize, or correct. Compress wording, not category coverage. Do not copy only the source opening.\n\nAim for roughly 100 words. Keep it concise, but allow modest variation when needed to preserve important information and finish the caption naturally. Do not cut off a sentence merely to satisfy a word-count target.\n\nTAGGY must preserve all concrete LoRA-useful source details as compact comma-separated phrases.\n\nNo markdown, reasoning, notes, or other labels.",
"Formatter - max new tokens": 3200,
"Formatter - temperature": 0.12,
"Formatter - top p": 0.88,
"Formatter - top k": 50,
"Audit - write prompt JSONL": false,
"Audit - preserve raw responses": false,
"Final - TXT export format": "natural",
"Final - write TXT sidecars": true,
"Final - write JSONL": true,
"Dataset - export image and caption": false,
"Dataset - output folder": "",
"Dataset - max image size": 0,
"Dataset - dimension divisor": 16,
"Dataset - image format": "PNG",
"Dataset - JPEG quality": 95,
"Dataset - caption": "short",
"Input - single image": [
"90",
0
],
"pipeline_plan": [
"90",
1
]
},
"class_type": "JLC_CaptionForge",
"_meta": {
"title": " JLC CaptionForge Orchestrator"
}
},
"148": {
"inputs": {
"Planner - enabled": true,
"Input - image path": "",
"Input - recursive": false,
"Input - filename glob": "*",
"Output - folder": "",
"Output - run name": "captionforge_run",
"Output - overwrite outputs": true,
"Ollama - URL": "http://127.0.0.1:11434",
"Ollama - keep loaded": true,
"Ollama - request timeout seconds": 1800,
"LoRA - trigger word": "Fairy",
"LoRA - user caption anchor": "",
"Cleanup - forbidden phrases": "",
"Cleanup - replace pairs": "",
"Caption - Joy runs/image": "2",
"Caption - Qwen runs/image": "1",
"Caption - Ollama runs/image": "1",
"Caption - base seed": 1,
"Caption - seed mode": "increment",
"Caption - temperature schedule": "0.75,0.85,0.95",
"Caption - top p schedule": "0.60",
"Caption - top k schedule": "80",
"Caption - max image size": 1024,
"Caption - max new tokens": 4096,
"Distiller - model": "mistral-small:24b",
"Distiller - custom Ollama model": "",
"Distiller - prompt": "/no_think\n\nYou are a detail-preserving caption merger for LoRA dataset preparation.\n\nYou receive multiple captions of the same image. You do NOT see the image.\n\nTask:\nMerge all non-contradictory caption details into one deliberately over-complete draft caption.\n\nRules:\n- Do not validate against the image.\n- Do not decide that details are false just because they appear once.\n- Do not summarize aggressively.\n- Preserve concrete details from all captions.\n- Split contradictions by choosing cautious wording or listing the alternative only when needed.\n- Prefer specific visual language over generic language.\n- Keep visible body, clothing, material, accessory, color, pose, lighting, style, and framing details.\n- Preserve doll-like, glossy/plastic-like, material, garment-construction, body-shape, and facial-feature details when present.\n- Use neutral dataset-caption language, including visible sensual styling or revealing clothing when present.\n- Do not add details absent from the captions.\n- Treat subject names or trigger-like identity tokens as optional identity labels. Preserve them only when they appear consistently in the captions; do not let them replace visible description.\n- Output only one paragraph, no notes, no JSON.",
"Distiller - seed": -1,
"Distiller - max caption chars for LLM": 1536,
"Distiller - num predict": 3096,
"Distiller - temperature": 0.35,
"Distiller - top p": 0.92,
"Distiller - top k": 80,
"Distiller - write prompt JSONL": false,
"Distiller - preserve raw response": false,
"Validator - model": "gemma4:26b",
"Validator - custom Ollama model": "",
"Validator - system prompt": "/no_think\nYou are a direct image validation engine. Inspect the image and answer only with the requested caption.",
"Validator - prompt": "/no_think\n\nLook at the image and validate this draft caption.\n\nTask:\nReturn a corrected caption paragraph that keeps only image-supported details.\n\nRules:\n- Output only the corrected caption.\n- One paragraph.\n- No reasoning, no notes, no JSON.\n- Keep all true visible details from the draft.\n- Delete unsupported details.\n- Correct small visible errors.\n- Do not add new details unless needed to correct an error already present.\n- Preserve useful LoRA details: subject, face, hair, eyes, makeup, lips, skin texture, pose, body shape, outfit, accessories, materials, colors, lighting, background, framing, and visual style.\n- Visible sensual styling, revealing clothing, cleavage, thighs, bare skin, swimwear, lingerie, or body-shape details may be described neutrally when present.\n- Do not invent hidden anatomy, unseen clothing, explicit acts, or details contradicted by the image.",
"Validator - seed": -1,
"Validator - num predict": 2112,
"Validator - temperature": 0,
"Validator - top p": 0.92,
"Validator - top k": 80,
"Validator - write prompt JSONL": false,
"Validator - preserve raw VLM response": false,
"Formatter - model": "mistral-small:24b",
"Formatter - custom Ollama model": "",
"Formatter - prompt": "/no_think\n\nYou are a LoRA caption format converter. The validated paragraph is your only source of truth.\n\nOutput exactly two labeled lines:\n\nSHORT: <a concise natural-language caption, typically around 100 words, that preserves all LoRA-useful validated details>\n\nTAGGY: <one compact comma-separated caption>\n\nSHORT must preserve the image's distinctive training identity across the whole source:\n1. subject, defining face/hair/body traits, and every major outfit piece/material;\n2. pose/action and key accessories or unusual visible details;\n3. setting, lighting, framing, and visual medium/style.\n\nOmit a category only when absent. Use only source details; never add, infer, euphemize, or correct. Compress wording, not category coverage. Do not copy only the source opening.\n\nAim for roughly 100 words. Keep it concise, but allow modest variation when needed to preserve important information and finish the caption naturally. Do not cut off a sentence merely to satisfy a word-count target.\n\nTAGGY must preserve all concrete LoRA-useful source details as compact comma-separated phrases.\n\nNo markdown, reasoning, notes, or other labels.",
"Formatter - seed": -1,
"Formatter - num predict": 3200,
"Formatter - temperature": 0,
"Formatter - top p": 0.9,
"Formatter - top k": 60,
"Formatter - write prompt JSONL": false,
"Formatter - preserve raw response": false,
"Final - write TXT sidecars": true,
"Final - write JSONL": true,
"Dataset - export image and caption": true,
"Dataset - output folder": "",
"Dataset - max image size": 1536,
"Dataset - dimension divisor": 16,
"Dataset - image format": "PNG",
"Dataset - JPEG quality": 95,
"Dataset - caption": "short",
"Input - single image": [
"151",
0
]
},
"class_type": "JLC_CaptionForge_Pipeline_Planner",
"_meta": {
"title": " JLC CaptionForge Pipeline Planner"
}
},
"151": {
"inputs": {
"image": "jlc_CaptionForge.jpg",
"resize_by": "scale longer dimension",
"multiplier": 1,
"longer_size": 1536,
"shorter_size": 1024,
"width": 1024,
"height": 1024,
"megapixels": 1,
"scale_method": "area",
"divisible_by": 16
},
"class_type": "JLC_LoadAndResizeImage",
"_meta": {
"title": " JLC Load, Resize & Encode Image"
}
},
"156": {
"inputs": {
"mode": "raw value",
"input": [
"147",
0
]
},
"class_type": "DisplayAny",
"_meta": {
"title": "🔧 Display Any"
}
},
"157": {
"inputs": {
"mode": "raw value",
"input": [
"147",
1
]
},
"class_type": "DisplayAny",
"_meta": {
"title": "🔧 Display Any"
}
},
"158": {
"inputs": {
"mode": "raw value",
"input": [
"147",
2
]
},
"class_type": "DisplayAny",
"_meta": {
"title": "🔧 Display Any"
}
}
}
File diff suppressed because one or more lines are too long
Binary file not shown.

After

Width:  |  Height:  |  Size: 2.3 MiB

+3 -2
View File
@@ -1,2 +1,3 @@
# captionforge_version.py
CAPTIONFORGE_VERSION = "0.1.0"
"""Single authoritative package/release version for CaptionForge."""
CAPTIONFORGE_VERSION = "1.0.2"
+8 -10
View File
@@ -1,28 +1,27 @@
{
"_meta": {
"name": "CaptionForge Ollama Model Dropdowns",
"version": "0.1.0",
"version": "1.0.2",
"description": "User-editable Ollama model dropdown configuration for CaptionForge nodes and engines.",
"consumed_by": [
"nodes/captionforge_ollama_model_dropdowns.py",
"CaptionForge Pipeline Planner",
"JLC CaptionForge capstone",
"JLC CaptionForge Orchestrator",
"JLC CaptionForge Ollama Caption"
],
"notes": [
"Values should be concrete Ollama model tags used exactly as written.",
"distiller_models are used for text-only LLM distillation and formatting stages.",
"distiller_models are used for Pass B text-only LLM synthesis.",
"validator_models are used for image-aware VLM validation.",
"format_models are used for Pass D text-only SHORT and TAGGY formatting.",
"caption_models are used by Ollama-backed Pass A caption witness nodes.",
"Set include_custom to true to expose a custom model-tag entry in supported nodes."
]
},
"distiller_models": [
"mistral-small:24b",
"VladimirGav/gemma4-26b-16GB-VRAM-Uncensored",
"deepseek-r1:32b",
"tarruda/neuraldaredevil-8b-abliterated:fp16",
"gpt-oss:20b"
"gpt-oss:20b",
"VladimirGav/gemma4-26b-16GB-VRAM-Uncensored"
],
"validator_models": [
"gemma4:26b",
@@ -31,9 +30,8 @@
],
"format_models": [
"mistral-small:24b",
"VladimirGav/gemma4-26b-16GB-VRAM-Uncensored",
"gpt-oss:20b",
"deepseek-r1:32b"
"VladimirGav/gemma4-26b-16GB-VRAM-Uncensored"
],
"caption_models": [
"gemma4:26b",
@@ -47,4 +45,4 @@
"caption_model": "gemma4:26b"
},
"include_custom": true
}
}
+99
View File
@@ -0,0 +1,99 @@
# Dataset export prototype
This opt-in export mode extends the existing Pipeline Planner and Orchestrator.
## Try it
1. Reload your user-managed ComfyUI instance at port **8189** to load the Python changes.
2. Load `assets/workflows/CaptionForge_FullWorkflow_Rel_v1.0.2.json`, or add fresh
Planner and Orchestrator nodes to an existing workflow.
3. Set the Planner input path and output folder. The prototype follows a
**1536-pixel** caption/Validator maximum, with no enlargement.
4. Enable **Dataset - export image and caption** on the Planner. The supplied
canonical workflow exposes this control; ordinary node defaults leave it off.
5. Choose the divisor, image format, and training caption. Queue a small dataset.
When connected, the Planner owns **every** Dataset setting, including disabled,
zero, and blank values. Orchestrator controls apply only in standalone mode.
Older plans with no Dataset settings keep export disabled.
| Control on both nodes | Default | Behavior |
| --- | --- | --- |
| Dataset - export image and caption | Off | Enables paired image and plain TXT export |
| Dataset - output folder | Blank | Parent folder; blank uses the main Output folder |
| Dataset - max image size | 0 | Follows configured Validator size; positive values override export size only |
| Dataset - dimension divisor | 16 | Integer; rounds both edges down. 1 disables alignment |
| Dataset - image format | PNG | PNG or JPEG, RGB |
| Dataset - JPEG quality | 95 | Applies only to JPEG |
| Dataset - caption | short | short, long, or taggy for the matching plain TXT |
In planned runs, the effective Validator maximum comes from the Planner's
**Caption - max image size**. Standalone runs use **Validator - max image size**.
If both the export override and effective Validator limit are zero, there is no
long-edge cap; divisor alignment still applies. Neither mode enlarges images.
## Output and protection
For an input `portraits/photo.jpg`, PNG export creates:
```text
<selected output parent>/training_dataset/
.captionforge-dataset.json
files/portraits/
photo.jpg.png
photo.jpg.txt
photo.jpg.captionforge.json
photo.jpg_long.txt
photo.jpg_short.txt
photo.jpg_taggy.txt
```
The last three files follow **Final - write TXT sidecars**. The selected plain
training TXT is always written when Dataset export is enabled. With export on,
caption variants are written beside the exported image rather than into the
source archive. With export off, existing v1.0.1 sidecar behavior is preserved.
Original extensions remain in the export stem, so `photo.jpg` and `photo.png`
produce distinct image/caption pairs. Relative folders are retained. Optional
IMAGE inputs use a separate `optional/` namespace. Sources outside the declared
input root use an `external/` namespace derived from their parent directory.
The dataset folder must be empty or already owned by CaptionForge. Untracked
files are not overwritten even when overwrite is enabled. The input path may
not be inside the selected export dataset. Resolved destination paths must stay
within the dataset and cannot refer to the source image, including hard links.
All three witness scanners exclude marked datasets on subsequent runs, even
when export has since been disabled. Do not remove the ownership marker.
## Resize and resume
Both output dimensions are rounded down after proportional size calculation.
This introduces a small aspect-ratio adjustment without cropping. If either
edge would become zero, export reports an error and keeps the captions in the
final record; reduce the divisor to handle such images.
Validator pixels are reused when their dimensions exactly match the requested
export dimensions. Otherwise, export resizes directly from the original to
avoid repeated resampling. Validator retry sizes do not change export size.
With overwrite disabled, completed caption records can receive a new dataset
export without repeating B/C/D model calls. Existing pairs resume only when
source metadata, settings, caption, and output hashes match their receipt.
Changed or incomplete pairs require overwrite to regenerate.
When changing PNG/JPEG format for an existing dataset, choose another output
parent; the prototype refuses to leave duplicate training images sharing one caption.
Export failures
retain caption text and are reported in the final record. Each file is published
by replacement from a temporary file; a receipt written last records completion.
An interrupted pair is not treated as complete.
## Scope
This prototype does not encode training latents or caption embeddings. It does
not change release version numbers, production captioning defaults, or public
node outputs. Export paths and dimensions are included in `final_records` under
`dataset_export`. The existing JSONL audit pipeline remains available.
CPU checks cover resize limits, divisor errors, pairing, collisions, recursive
exclusion, Planner ownership, and caption-preserving resume. Live canvas and
real-model validation remain pending; port 8189 was not responding during development.
+117
View File
@@ -0,0 +1,117 @@
# Joy 8-bit dtype and clean installation
## Why Balanced (8-bit) retains BF16
`Balanced (8-bit)::bf16` describes quantized weights plus Joy's floating-point
configuration, not an all-INT8 computation. Joy's quantized loader uses
`torch_dtype="auto"` (the checkpoint dtype), excludes the vision tower and
multimodal projector from INT8 conversion, and uses BF16 autocast where supported.
The configuration dtype is also part of the cache key; it does not override
`auto` in the quantized loader. This existing behavior is unchanged.
JoyCaption's upstream ComfyUI implementation uses the same `auto` loading and
BF16 autocast. In bitsandbytes 0.46.1, `MatMul8bitLt.forward` explicitly casts
activations to FP16 for INT8 quantization, while retaining the input dtype for
other computation/output handling. This is an expected internal conversion,
not sufficient evidence that the entire Joy model should switch to FP16.
Changing the vision tower, projector and surrounding floating-point computation
just to remove the warning would change numerical behavior without a quality
validation basis.
References:
- [JoyCaption native BF16 usage](https://github.com/fpgaminer/joycaption/blob/main/README.md)
- [Upstream JoyCaption ComfyUI implementation](https://github.com/fpgaminer/joycaption_comfyui/blob/main/nodes.py)
- [bitsandbytes 0.46.1 implementation](https://github.com/bitsandbytes-foundation/bitsandbytes/blob/0.46.1/bitsandbytes/autograd/_functions.py)
The desktop reproduction also showed FP32 input activations. The BF16 cache key
is not a promise that every intermediate tensor is BF16: autocast applies per
operation, and normalization or type promotion can produce FP32 intermediates.
The exact producing layer cannot be identified from the supplied log.
Bitsandbytes accepts these through the same explicit FP16 quantization conversion,
so both observed messages are handled without changing numerical behavior.
See the [current bitsandbytes implementation](https://github.com/bitsandbytes-foundation/bitsandbytes/blob/main/bitsandbytes/autograd/_functions.py).
## Warning scope
The old import-time filters were removed. During Balanced (8-bit) generation
only, CaptionForge ignores these exact `UserWarning` messages or WARNING log records from
`bitsandbytes.autograd._functions`:
```text
MatMul8bitLt: inputs will be cast from torch.bfloat16 to float16 during quantization
MatMul8bitLt: inputs will be cast from torch.float32 to float16 during quantization
```
Other cast dtypes, different messages, other modules, other warning categories,
processor/loading/cleanup diagnostics and Default-mode warnings remain visible.
The caller's filters are restored even if generation fails. Recent bitsandbytes versions emit this through `logger.warning` instead of
`warnings.warn`; both routes are covered. The logging filter matches the fully
formatted message and originating logger, and applies only to the inference
thread. It is removed on normal exit or failure, without changing logger levels,
handlers, propagation, or pre-existing filters. This does not alter tensors,
quantization parameters or model outputs.
The same scope is used around quantized Qwen generation. Qwen's previous
load-time warning filter was process-wide, BF16-only, and did not cover the
bitsandbytes logging route; it has been removed. Non-quantized Qwen generation
does not install the scope.
Python 3.10-3.12 warning filters are process-wide while a `catch_warnings` scope
is active. An identical warning from concurrent bitsandbytes work could therefore
also be suppressed during that interval. This is scoped filtering, not a claim
of thread-local warning isolation.
## Installation declarations
`accelerate` was already required by `pyproject.toml`; declaring it there alone
did not cover Manager's requirements-file installation path. CaptionForge now
ships `requirements.txt` with the same runtime dependencies as
`project.dependencies`, including `accelerate` and `bitsandbytes>=0.46.1`.
The `quantization` extra is retained for compatibility with existing install
commands, but bitsandbytes no longer requires opting into an extra.
[ComfyUI Manager's installation code](https://github.com/Comfy-Org/ComfyUI-Manager/blob/main/glob/manager_core.py)
reads `requirements.txt` and installs its entries. Both declarations are kept
explicit so Manager can process individual requirement lines. A regression test
prevents the lists from drifting. The lower bound does not guarantee compatibility
with every future torch/CUDA/bitsandbytes combination.
For an existing installation, run this from CaptionForge's directory using
**the Python interpreter belonging to ComfyUI Desktop's environment**, then
restart ComfyUI:
```powershell
python -m pip install -r requirements.txt
python -m pip check
```
This affects new Registry installs after a release containing these changes is
published; editing the repository does not update an already-published archive.
## Validation
CPU-only focused checks (Python 3.11/3.12; Python 3.10 also needs `tomli`):
```text
python -m unittest discover -s tests -p test_joy_hardening.py -v
```
They cover six repeated caption bursts through both warnings and logging, near-match diagnostics,
Default mode, filter restoration, exceptions, and matching dependency lists.
Generation tests execute the real `caption_pil` method with test doubles for
processor, model and torch; they do not test numerical inference.
Before releasing, validate on a clean supported ComfyUI Desktop environment:
1. Install the updated package through Manager without manually installing
accelerate/bitsandbytes. Confirm both are present using that environment's
`python -m pip show accelerate bitsandbytes`, then run `python -m pip check`.
2. Run Joy Balanced (8-bit), two runs over three ordinary test images. Confirm
six successful captions, normal progress/error diagnostics, and no BF16 cast
warning bursts. Reuse the loaded model for a subsequent run as well.
3. Check Default mode if sufficient VRAM is available. No warning policy or
numerical settings should change there.
A real clean Desktop install and CUDA caption generation cannot be established
by the CPU-only checks; they remain release validation steps.
+157
View File
@@ -0,0 +1,157 @@
# Fixed-corpus caption quality study
Date: 2026-09-09
Baseline: `5568f74dbb103b99bd6603481a247aee6033f753`
Release disposition: the dual-output Pass-D refinement described below was
adopted for CaptionForge 1.0.0. The corpus name is retained as historical study
provenance; it is not the package version.
## Question
Evaluate the current `_long.txt`, `_short.txt`, and `_taggy.txt` outputs before
the 1.0 release. In particular, determine whether short captions retain useful
LoRA-training detail and whether taggy captions are compact, faithful, and
readable enough to keep without redesign.
## Corpus and method
The fixed corpus is the 19-image local `Release_1.0.1_BigTest01` set. Despite
the historical folder name, this is pre-1.0 release-candidate evidence. It
contains real photographs, digital illustrations, anime-style art, glossy 3D
characters, a ball-jointed doll, portraits, full-body images, stage/runway
scenes, indoor and outdoor settings, simple and detailed backgrounds, varied
poses, and varied clothing/material descriptions.
The stored run used the production model path:
- Pass A: one Joy and one Qwen caption per image
- Pass B: `mistral-small:24b`
- Pass C: `gemma4:26b`
- Pass D: `mistral-small:24b`
The study combined mechanical measurements with a manual image/caption review.
The proposed refinement was then tested directly against all 19 validated long
captions with `mistral-small:24b`, seed 3, temperature 0.12, top-p 0.88, and
top-k 50.
## Existing output measurements
| Metric | Result |
| --- | ---: |
| Images with complete long/short/taggy triplets | 19 |
| Mean long length | 127.3 words |
| Mean short length | 81.6 words |
| Mean short/long ratio | 66.6% |
| Shorts that are exact prefixes of long | 19/19 |
| Shorts that drop content | 16/19 |
| Shorts that drop at least 20 words | 14/19 |
| Shorts identical to long | 3/19 |
| Mean taggy length | 67.8 words / 28.5 items |
| Taggy range | 18–36 items |
| Largest taggy output | 648 characters |
| Exact duplicate taggy items | 0 |
| Taggy outputs reaching a compaction limit | 0 |
## Findings
### Long
`_long` is appropriately treated as authoritative. It consistently contains
the richest set of subject, appearance, clothing, pose, setting, lighting,
framing, and style details. This study did not identify a downstream-formatting
reason to change Pass C.
Manual image review did find isolated Pass-C wording errors (for example, a
doll support stand described as a microphone stand). Those are validator
accuracy issues, not short/taggy divergence, and are outside this formatting
study.
### Short: material problem found
The deterministic implementation is sentence-prefix extraction, not semantic
compression. Every stored short is the beginning of its long caption. Because
the validator commonly places setting, lighting, framing, style, accessories,
and some pose/body details near the end, those categories are systematically
more likely to disappear.
Representative losses include:
- the entire outfit, exposed midriff/thighs, pose, and shiny fabric detail from
the pink bikini/sarong illustration;
- mirror-selfie action, phone/hand placement, background, lighting, and
anime/semi-realistic rendering details;
- body shape, side-view pose, runway, crowd, sign text, reflective floor,
full-body framing, and doll-like style;
- sailor hat, high heels, salute, relaxed arm, full-body framing, and studio
lighting;
- castle, lanterns, trees, moon, stars, river, lighting, and 3D style from the
fairy image.
There is also a concrete limit defect: when a validated long caption is a
single sentence longer than 90 words, the current helper keeps the entire
sentence. Two corpus shorts therefore contain 113 and 119 words despite the
advertised 90-word maximum.
Conclusion: the existing `_short` is materially worse than `_long` for
training whenever useful categories occur late in the validated paragraph. It
is not reliable as the intended LoRA-length semantic summary.
### Taggy: acceptable with minor cleanup
`_taggy` is compact enough for the studied corpus and remains semantically
close to `_long`. The deterministic compactor did not truncate any result, so
it removed no useful content through item or character limits. Exact duplicate
items were absent.
Minor defects remain: occasional dangling fragments (`sides`, `left`,
`visible`), terminal sentence punctuation, and near-duplicates such as
`glossy skin`/`smooth skin` or `bokeh lights`/`bokeh effect`. These do not
justify a new taggy architecture. Terminal punctuation can be safely removed
during deterministic cleanup; semantic near-deduplication should not be made
aggressive because similar phrases can carry distinct material or appearance
information.
## Refinement experiment
The existing Pass-D formatter was asked to return two labeled derivatives in
the same call. The short output was constrained to three concise sentences
covering, when present:
1. subject/appearance and major clothing/material;
2. pose/body/accessories or unusual visible details;
3. setting/lighting/framing and medium/style.
Results:
| Metric | Result |
| --- | ---: |
| Dual responses parsed | 19/19 |
| Shorts within 90 words | 19/19 |
| Short word range | 47–77 |
| Mean short length | 60.9 words |
| Added model calls | 0 |
| Added image-validation calls | 0 |
Manual review found that the refined shorts consistently selected information
from across the whole validated caption rather than stopping after its opening.
They are meaningfully shorter while retaining a balanced training identity.
As expected for compression, not every minor detail survives; the full long and
taggy variants remain available when maximum recall is preferred.
## Decision
Adopt the dual-output Pass-D prompt as the smallest defensible fix:
- `_long` remains the untouched Pass-C authority.
- One existing Pass-D text-only call derives both `_short` and `_taggy`.
- The formatter may only compress validated content; it may not add, infer, or
correct visual claims.
- Deterministic cleanup enforces the short limit and normalizes taggy output.
- Legacy/custom taggy-only formatter responses remain supported, using the
corrected deterministic short helper as a fallback.
- Do not add another VLM pass or another LLM call.
This resolves the demonstrated short-caption failure without redesigning the
successful semantic pipeline or the acceptable taggy path.
+6 -1
View File
@@ -1 +1,6 @@
# empty file
"""Production backend and support engines for CaptionForge.
Active ComfyUI execution is orchestrated by the node wrappers. The standalone
Distiller and Validator modules are retained as reference/CLI implementations,
not as the authoritative production B/C/D path.
"""
+12 -15
View File
@@ -8,19 +8,16 @@ CaptionForge Caption Prompt Kit
- Repository:
https://github.com/Damkohler/CaptionForge
- CaptionForge focuses on practical dataset-captioning infrastructure for
LoRA dataset preparation, using multi-engine caption generation, JSONL
audit trails, claim extraction and refinement, text-LLM distillation,
image-aware VLM validation, and consensus-oriented caption improvement
to produce grounded, auditable training captions.
- CaptionForge 1.0 uses independent Pass-A witnesses, text-LLM synthesis,
image-aware validation, SHORT/TAGGY formatting, and JSONL audit trails to
produce grounded LoRA dataset captions.
- Module Purpose
- The **CaptionForge Caption Prompt Kit** is a shared, dependency-free
prompt builder for non-Joy CaptionForge caption witnesses.
- The **CaptionForge Caption Prompt Kit** is the shared, dependency-free
prompt builder for active Qwen and Ollama caption witnesses.
- It provides a common caption-prompt vocabulary for Qwen, SmolVLM, and
future lightweight or generic caption engines while still allowing
model-dialect-specific wording.
- It retains generic and model-dialect-specific wording helpers so active
witnesses can share controls without sharing identical prompts.
- It defines:
• caption length choices
@@ -65,10 +62,10 @@ CaptionForge Caption Prompt Kit
JSONL audit trails so users can understand what each caption witness was
asked to do.
- Development Status
- CaptionForge v0.1.0 experimental developer-preview infrastructure.
- Caption types, extra options, and dialect wording may evolve as additional
caption engines are tested.
- Production Status
- Active CaptionForge 1.0 Pass-A support code used by Qwen and Ollama
witnesses. ``PROMPT_KIT_VERSION`` below identifies its internal prompt
contract and is intentionally separate from the package release version.
- Attribution & License
- Concept and implementation by **J. L. Córdova**
@@ -94,7 +91,7 @@ MANIFEST = {
"Shared dependency-free prompt builder for non-Joy CaptionForge caption "
"witnesses. Provides caption length choices, caption type choices, "
"extra-option text, dialect normalization, prompt construction helpers, "
"and prompt metadata for Qwen, SmolVLM, Ollama, and future generic "
"and prompt metadata for Qwen, Ollama, and future generic "
"caption engines."
),
}
+141
View File
@@ -0,0 +1,141 @@
"""Shared forbidden-phrase matching helpers for CaptionForge cleanup paths."""
from __future__ import annotations
import re
from collections.abc import Iterable
from typing import Any
def phrase_boundary_pattern(
phrase: str,
*,
case_insensitive: bool = True,
) -> re.Pattern[str] | None:
"""Compile a phrase matcher that respects token boundaries.
Boundary checks are added only when the corresponding phrase edge is a
word character. This preserves literal punctuation in configured phrases
while preventing tokens such as old from matching inside holding or bold.
"""
text = str(phrase or "").strip()
if not text:
return None
pattern = re.escape(text)
if re.match(r"\w", text[0], flags=re.UNICODE):
pattern = r"(?<!\w)" + pattern
if re.match(r"\w", text[-1], flags=re.UNICODE):
pattern = pattern + r"(?!\w)"
flags = re.IGNORECASE if case_insensitive else 0
return re.compile(pattern, flags=flags)
def contains_forbidden_phrase(text: str, forbidden_phrases: Iterable[str]) -> bool:
"""Return True when any configured forbidden phrase matches at boundaries."""
haystack = str(text or "")
for phrase in forbidden_phrases:
pattern = phrase_boundary_pattern(phrase)
if pattern is not None and pattern.search(haystack):
return True
return False
def remove_forbidden_phrases(text: str, forbidden_phrases: Iterable[str]) -> str:
"""Remove configured forbidden phrases without corrupting containing words."""
result = str(text or "")
for phrase in forbidden_phrases:
pattern = phrase_boundary_pattern(phrase)
if pattern is not None:
result = pattern.sub("", result)
return result
def replace_phrases(
text: str,
replacement_rules: Iterable[tuple[str, str]],
*,
case_insensitive: bool = True,
) -> str:
"""Apply replacement rules only at whole-word/phrase boundaries."""
result = str(text or "")
for old, new in replacement_rules:
pattern = phrase_boundary_pattern(old, case_insensitive=case_insensitive)
if pattern is not None:
result = pattern.sub(str(new or ""), result)
return result
def normalize_forbidden_phrases(value: Any) -> list[str]:
"""Normalize UI or plan values into an ordered forbidden-phrase list."""
if value is None:
return []
if isinstance(value, str):
items = value.splitlines()
elif isinstance(value, (list, tuple)):
items = value
else:
items = [value]
return [text for item in items if (text := str(item or "").strip())]
def normalize_replace_pairs(value: Any) -> list[tuple[str, str]]:
"""Normalize ``old=>new`` UI text or serialized plan replacement pairs."""
if value is None:
return []
if isinstance(value, str):
items: Iterable[Any] = value.splitlines()
elif isinstance(value, (list, tuple)):
items = value
else:
items = [value]
pairs: list[tuple[str, str]] = []
for item in items:
if isinstance(item, dict):
old, new = item.get("old", ""), item.get("new", "")
elif isinstance(item, (list, tuple)) and len(item) >= 2:
old, new = item[0], item[1]
else:
line = str(item or "").strip()
if not line or line.startswith("#") or "=>" not in line:
continue
old, new = line.split("=>", 1)
old_text = str(old or "").strip()
if old_text:
pairs.append((old_text, str(new or "").strip()))
return pairs
def resolve_cleanup_settings(
pipeline_plan: Any,
standalone_forbidden_phrases: Any,
standalone_replace_pairs: Any,
) -> tuple[list[str], list[tuple[str, str]]]:
"""Resolve Planner-owned cleanup values, preserving standalone node use."""
plan = pipeline_plan if isinstance(pipeline_plan, dict) else {}
cleanup = plan.get("cleanup") if isinstance(plan.get("cleanup"), dict) else None
if cleanup is not None:
forbidden_value = cleanup.get("forbidden_phrases", [])
replace_value = cleanup.get("replace_pairs", [])
else:
forbidden_value = standalone_forbidden_phrases
replace_value = standalone_replace_pairs
return normalize_forbidden_phrases(forbidden_value), normalize_replace_pairs(replace_value)
def apply_cleanup_contract(
text: str,
forbidden_phrases: Iterable[str],
replace_pairs: Iterable[tuple[str, str]],
) -> str:
"""Apply the shared boundary-safe cleanup contract and repair separators."""
result = replace_phrases(text, replace_pairs)
result = remove_forbidden_phrases(result, forbidden_phrases)
result = re.sub(r"\s+([,.;:!?])", r"\1", result)
result = re.sub(r",\s*,+", ",", result)
result = re.sub(r"([.;:!?])(?:\s*[,.;:!?])+", r"\1", result)
result = re.sub(r"\s+", " ", result)
return result.strip(" ,")
+227
View File
@@ -0,0 +1,227 @@
"""Optional, non-enlarging training image/caption export shared by both nodes."""
from __future__ import annotations
import hashlib
import json
import os
import tempfile
from pathlib import Path
from PIL import Image
from .captionforge_source_identity import optional_image_filename
EXPORT_MARKER = ".captionforge-dataset.json"
MARKER_CONTENT = {"type": "captionforge_training_dataset", "version": 1}
EXPORT_DEFAULTS = {
"enabled": False,
"output_folder": "",
"max_size": 0,
"divisor": 16,
"image_format": "PNG",
"jpeg_quality": 95,
"caption": "short",
}
EXPORT_WIDGETS = {
"enabled": "Dataset - export image and caption",
"output_folder": "Dataset - output folder",
"max_size": "Dataset - max image size",
"divisor": "Dataset - dimension divisor",
"image_format": "Dataset - image format",
"jpeg_quality": "Dataset - JPEG quality",
"caption": "Dataset - caption",
}
def dataset_export_inputs() -> dict:
"""Append optional widgets so older API workflows remain valid."""
specs = {
"enabled": ("BOOLEAN", {"tooltip": "Export a resized image and matching training TXT. Planner owns these controls when connected."}),
"output_folder": ("STRING", {"tooltip": "Parent folder for training_dataset. Blank uses Output - folder. Originals are never replaced."}),
"max_size": ("INT", {"min": 0, "max": 8192, "step": 1, "tooltip": "Maximum long edge; never enlarges. 0 follows the configured Validator size (Planner Caption - max image size)."}),
"divisor": ("INT", {"min": 1, "max": 512, "step": 1, "tooltip": "Round both dimensions DOWN to this multiple after resizing. 1 disables alignment. Images too small for the divisor fail export without enlargement."}),
"image_format": (["PNG", "JPEG"], {"tooltip": "Format of the exported RGB training image."}),
"jpeg_quality": ("INT", {"min": 1, "max": 100, "step": 1, "tooltip": "JPEG quality; ignored for PNG."}),
"caption": (["short", "long", "taggy"], {"tooltip": "Caption written to the matching plain .txt file. Other caption variants remain available."}),
}
return {EXPORT_WIDGETS[key]: (kind, {"default": EXPORT_DEFAULTS[key], **options})
for key, (kind, options) in specs.items()}
def export_settings_from_widgets(widgets: dict) -> dict:
return {key: widgets.get(name, EXPORT_DEFAULTS[key]) for key, name in EXPORT_WIDGETS.items()}
def normalize_export_settings(values: dict | None) -> dict:
result = {**EXPORT_DEFAULTS, **(values or {})}
result["enabled"] = str(result["enabled"]).strip().lower() in {"true", "1", "yes", "on"}
result["output_folder"] = str(result["output_folder"] or "").strip()
for key, low, high in (("max_size", 0, 8192), ("divisor", 1, 512), ("jpeg_quality", 1, 100)):
value = float(result[key])
if not value.is_integer() or not low <= value <= high:
raise ValueError(f"Dataset {key} must be an integer from {low} to {high}.")
result[key] = int(value)
if result["image_format"] not in {"PNG", "JPEG"} or result["caption"] not in {"short", "long", "taggy"}:
raise ValueError("Unsupported dataset image format or caption choice.")
return result
def dataset_root(settings: dict, output_folder: str | Path) -> Path:
return (Path(settings["output_folder"] or output_folder).expanduser() / "training_dataset").resolve()
def is_dataset_export(path: Path) -> bool:
"""Recognize generated datasets, including after export has been disabled."""
resolved = path.resolve()
return any((parent / EXPORT_MARKER).is_file() for parent in (resolved, *resolved.parents))
def prepare_dataset_root(root: Path, input_path: str | Path = "") -> None:
"""Claim only an empty directory or a previously managed export directory."""
root = root.resolve()
if input_path and Path(input_path).resolve().is_relative_to(root):
raise ValueError("Dataset destination contains the input path. Choose another output folder.")
marker = root / EXPORT_MARKER
if marker.exists():
if marker.is_symlink() or json.loads(marker.read_text(encoding="utf-8")) != MARKER_CONTENT:
raise ValueError("Dataset folder has an invalid CaptionForge marker.")
return
if root.exists() and any(root.iterdir()):
raise ValueError(f"Dataset folder is not empty and is not managed by CaptionForge: {root}")
root.mkdir(parents=True, exist_ok=True)
with marker.open("x", encoding="utf-8") as stream:
json.dump(MARKER_CONTENT, stream)
def export_dimensions(source_size: tuple[int, int], max_size: int, divisor: int) -> tuple[int, int]:
"""Use integer arithmetic so exact multiples never lose a pixel to float error."""
width, height = source_size
longest = max(width, height)
scaled = tuple(edge * max_size // longest if 0 < max_size < longest else edge for edge in source_size)
size = tuple((edge // divisor) * divisor for edge in scaled)
if min(size) < 1:
raise ValueError(f"Image {width}x{height} is too small for divisor {divisor} at max size {max_size}; use a smaller divisor.")
return size
def resize_for_export(image: Image.Image, max_size: int, divisor: int) -> Image.Image:
"""Calculate dimensions first and resample once; never clamp an edge upward."""
size = export_dimensions(image.size, max_size, divisor)
return image.resize(size, Image.Resampling.LANCZOS) if size != image.size else image
def export_paths(root: Path, source: Path, input_root: str | Path, image_key: str, image_format: str) -> tuple[Path, Path, Path]:
"""Keep the original extension in the stem to avoid conversion collisions."""
optional_name = optional_image_filename(image_key)
if optional_name:
relative = Path("optional") / optional_name
else:
base = Path(input_root).resolve() if input_root else source.resolve().parent
if base.is_file():
base = base.parent
try:
relative = Path("files") / source.resolve().relative_to(base)
except ValueError:
namespace = hashlib.sha256(str(source.resolve().parent).encode()).hexdigest()[:16]
relative = Path("external") / namespace / source.name
target = root / relative.parent / (relative.name + (".png" if image_format == "PNG" else ".jpg"))
caption = target.with_suffix(".txt")
receipt = target.with_suffix(".captionforge.json")
check_export_destinations(root, source, (target, caption, receipt))
return target, caption, receipt
def check_export_destinations(root: Path, source: Path, paths) -> None:
for path in paths:
if path.is_symlink() or not path.resolve().is_relative_to(root.resolve()):
raise ValueError("Dataset destination escapes its protected folder.")
if path.resolve() == source.resolve() or (path.exists() and path.samefile(source)):
raise ValueError("Dataset export would overwrite a source image.")
def _digest(path: Path) -> str:
return hashlib.sha256(path.read_bytes()).hexdigest()
def export_pair(*, source: Path, image_key: str, input_root: str | Path, root: Path,
settings: dict, validator_max_size: int, captions: dict, overwrite: bool,
prepared_image: Image.Image | None = None, write_variants: bool = False) -> dict:
"""Publish a pair with a receipt written last; incomplete pairs never resume."""
source = source.resolve()
if is_dataset_export(source):
raise ValueError("A generated dataset image cannot be used as its own archive source.")
prepare_dataset_root(root, input_root)
target, caption_path, receipt_path = export_paths(root, source, input_root, image_key, settings["image_format"])
max_size = settings["max_size"] or validator_max_size
caption = str(captions.get(settings["caption"]) or "").strip()
if not caption:
raise ValueError("Selected training caption is empty.")
variants = {style: str(captions.get(style) or "").strip() for style in ("long", "short", "taggy")} if write_variants else {}
variant_paths = {style: target.with_name(f"{target.stem}_{style}.txt") for style in variants}
check_export_destinations(root, source, variant_paths.values())
stat = source.stat()
signature = {"source": str(source), "image_key": image_key, "source_size": stat.st_size,
"source_mtime_ns": stat.st_mtime_ns, "max_size": max_size,
"divisor": settings["divisor"], "format": settings["image_format"],
"jpeg_quality": settings["jpeg_quality"], "caption": caption, "variants": variants}
previous = json.loads(receipt_path.read_text(encoding="utf-8")) if receipt_path.exists() else {}
owned_variants = set(previous.get("owned_variants", previous.get("signature", {}).get("variants", {})))
for style, path in variant_paths.items():
if path.exists() and style not in owned_variants:
raise FileExistsError(f"Dataset caption variant is untracked: {path}")
previous_image = previous.get("image") or previous.get("export", {}).get("image")
if previous_image and Path(previous_image) != target and Path(previous_image).exists():
raise FileExistsError("This source was exported in another image format. Choose a different Dataset output folder to avoid duplicate training images.")
if any(path.exists() for path in (target, caption_path, receipt_path)):
if previous.get("signature", {}).get("source") != str(source) or previous.get("signature", {}).get("image_key") != image_key:
raise FileExistsError(f"Dataset filename is already occupied by another or untracked source: {target}")
if not overwrite:
if (previous.get("signature") == signature and target.is_file() and caption_path.is_file()
and previous.get("image_sha256") == _digest(target)
and previous.get("caption_sha256") == _digest(caption_path)
and all(path.is_file() and previous.get("variant_sha256", {}).get(style) == _digest(path)
for style, path in variant_paths.items())):
return {**previous["export"], "resumed": True}
raise FileExistsError("Dataset pair is incomplete, changed, or uses different settings; enable overwrite to regenerate it.")
if prepared_image is None:
with Image.open(source) as original:
prepared_image = original.convert("RGB")
resized = resize_for_export(prepared_image, max_size, settings["divisor"])
result = {"status": "ok", "image": str(target), "caption": str(caption_path),
"width": resized.width, "height": resized.height, "caption_style": settings["caption"],
"max_size": max_size, "divisor": settings["divisor"], "resumed": False,
"variants": {style: str(path) for style, path in variant_paths.items()}}
target.parent.mkdir(parents=True, exist_ok=True)
temporary: list[Path] = []
try:
for _ in range(3 + len(variants)):
descriptor, name = tempfile.mkstemp(prefix=".captionforge-", dir=target.parent)
os.close(descriptor)
temporary.append(Path(name))
options = {"quality": settings["jpeg_quality"], "subsampling": 0} if settings["image_format"] == "JPEG" else {}
resized.save(temporary[0], format=settings["image_format"], **options)
temporary[1].write_text(caption + "\n", encoding="utf-8")
variant_hashes = {}
for index, (style, text) in enumerate(variants.items(), start=2):
temporary[index].write_text(text + "\n", encoding="utf-8")
variant_hashes[style] = _digest(temporary[index])
receipt = {"signature": signature, "export": result,
"image_sha256": _digest(temporary[0]), "caption_sha256": _digest(temporary[1]),
"variant_sha256": variant_hashes, "owned_variants": sorted(owned_variants | set(variants))}
# Reserve ownership before publishing, without claiming untracked files.
reservation = {"signature": signature, "image": str(target), "owned_variants": receipt["owned_variants"]}
if previous:
temporary[-1].write_text(json.dumps(reservation), encoding="utf-8")
os.replace(temporary[-1], receipt_path)
else:
with receipt_path.open("x", encoding="utf-8") as stream:
json.dump(reservation, stream)
temporary[-1].write_text(json.dumps(receipt, ensure_ascii=False, indent=2), encoding="utf-8")
for staging, destination in zip(temporary, (target, caption_path, *variant_paths.values(), receipt_path)):
os.replace(staging, destination)
finally:
for path in temporary:
path.unlink(missing_ok=True)
return result
+8 -27
View File
@@ -8,11 +8,9 @@ CaptionForge Distiller Engine
- Repository:
https://github.com/Damkohler/CaptionForge
- CaptionForge focuses on practical dataset-captioning infrastructure for
LoRA dataset preparation, using multi-engine caption generation, JSONL
audit trails, claim extraction and refinement, text-LLM distillation,
image-aware VLM validation, and consensus-oriented caption improvement
to produce grounded, auditable training captions.
- CaptionForge 1.0 uses independent Pass-A witnesses, text-LLM synthesis,
image-aware validation, SHORT/TAGGY formatting, and JSONL audit trails to
produce grounded LoRA dataset captions.
- Engine Purpose
- The **CaptionForge Distiller Engine** is the Pass B text-LLM
@@ -72,10 +70,10 @@ CaptionForge Distiller Engine
- Draft captions should be rich enough for LoRA dataset preparation while
remaining auditable through their source claims and metadata.
- Development Status
- CaptionForge v0.1.0 experimental developer-preview infrastructure.
- Prompt contracts, audit fields, and parser behavior may evolve before a
stable CaptionForge release.
- Reference Status
- This CLI-oriented engine is retained for reference, diagnostics, and
controlled experiments. It is not imported or registered by the active
ComfyUI pipeline; production Pass B runs inside ``jlc_captionforge_node``.
- Attribution & License
- Concept and implementation by **J. L. Córdova**
@@ -127,23 +125,6 @@ from pathlib import Path
from typing import Any, Optional
ENGINE_NAME = "captionforge_distiller_engine"
ENGINE_VERSION = "0.2.0"
CAPTIONFORGE_PASS = "B_DISTILL"
DEFAULT_OLLAMA_BASE_URL = "http://127.0.0.1:11434"
MANIFEST = {
"name": "CaptionForge Distiller Engine",
"version": ENGINE_VERSION,
"author": "J. L. Córdova",
"description": (
"CLI-first CaptionForge Pass B pollster/copywriter engine. Consumes Pass A "
"caption JSONL records, groups captions by image, asks a text LLM to "
"vote/organize visual claims, and emits accepted evidence, singleton "
"candidates, rejected conflicts, and rich/taggy draft captions."
),
}
DEFAULT_DISTILLER_INSTRUCTIONS = (
"You are CaptionForge Pass B: a caption ballot pollster and rich-caption copywriter. "
"Your job is not to summarize, prune for brevity, or compress the source captions. "
@@ -1355,7 +1336,7 @@ def process_batch(batch: BatchConfig, config: DistillerConfig) -> int:
def build_parser() -> argparse.ArgumentParser:
p = argparse.ArgumentParser(description="CaptionForge Distiller Engine v0.2.0")
p = argparse.ArgumentParser(description=f"CaptionForge Distiller Engine v{ENGINE_VERSION}")
p.add_argument("--input-jsonl", required=True, help="Pass A captions JSONL input.")
p.add_argument("--output-jsonl", default="", help="Full output JSONL path. Default: <input>_distilled.jsonl")
+7 -9
View File
@@ -8,11 +8,9 @@ CaptionForge Joy Space Prompt Kit
- Repository:
https://github.com/Damkohler/CaptionForge
- CaptionForge focuses on practical dataset-captioning infrastructure for
LoRA dataset preparation, using multi-engine caption generation, JSONL
audit trails, claim extraction and refinement, text-LLM distillation,
image-aware VLM validation, and consensus-oriented caption improvement
to produce grounded, auditable training captions.
- CaptionForge 1.0 uses independent Pass-A witnesses, text-LLM synthesis,
image-aware validation, SHORT/TAGGY formatting, and JSONL audit trails to
produce grounded LoRA dataset captions.
- Module Purpose
- The **CaptionForge Joy Space Prompt Kit** is a small, dependency-free
@@ -69,10 +67,10 @@ CaptionForge Joy Space Prompt Kit
- Prompt metadata should make it clear which caption type, length, options,
name input, and system prompt were used for a Joy caption run.
- Development Status
- CaptionForge v0.1.0 experimental developer-preview infrastructure.
- Joy prompt templates and option lists may evolve as the local CaptionForge
Joy node matures.
- Production Status
- Active CaptionForge 1.0 Pass-A support code. It owns the Joy-specific
prompt contract while model loading and inference remain in the Joy
engine.
- Attribution & License
- Concept and implementation by **J. L. Córdova**
+84
View File
@@ -0,0 +1,84 @@
"""Scoped handling of expected bitsandbytes activation-cast diagnostics."""
from contextlib import contextmanager
from importlib import metadata
import logging
import re
import threading
import warnings
_CAST_MESSAGES = frozenset(
f"MatMul8bitLt: inputs will be cast from torch.{dtype} "
"to float16 during quantization"
for dtype in ("bfloat16", "float32")
)
_BNB_LOGGER = "bitsandbytes.autograd._functions"
_MINIMUM_BNB_VERSION = (0, 46, 1)
def warn_if_suspicious_8bit_stack(node_name: str) -> None:
"""Warn, without blocking inference, when bnb predates the project floor."""
try:
version = metadata.version("bitsandbytes")
except Exception:
return
parts = tuple(int(piece) for piece in re.findall(r"\d+", version)[:3])
normalized = parts + (0,) * (3 - len(parts))
if normalized < _MINIMUM_BNB_VERSION:
warnings.warn(
f"CaptionForge's {node_name} node detected bitsandbytes {version}; "
"this older 8-bit inference stack may cause severe slowdowns or compatibility issues. "
"CaptionForge will continue without modifying packages.",
RuntimeWarning,
stacklevel=2,
)
class _QuantizedCastLogFilter(logging.Filter):
def __init__(self):
super().__init__()
self.thread_id = threading.get_ident()
def filter(self, record):
return not (
record.name == _BNB_LOGGER
and record.levelno == logging.WARNING
and record.thread == self.thread_id
and record.getMessage() in _CAST_MESSAGES
)
@contextmanager
def quantized_inference_warnings(enabled: bool):
"""Preserve diagnostics except known BF16/FP32 casts during 8-bit inference.
No filter is installed at import time or for Default mode. catch_warnings
restores the caller's filters even when generation raises an exception.
On Python 3.10-3.12 warning filters are process-wide during this brief scope;
the exact message, category and module keep the suppression narrowly bounded.
"""
if not enabled:
yield
return
with warnings.catch_warnings():
warnings.filterwarnings(
"ignore",
message=(
r"\AMatMul8bitLt: inputs will be cast from torch\.(?:bfloat16|float32) "
r"to float16 during quantization\Z"
),
category=UserWarning,
module=r"\Abitsandbytes\.autograd\._functions\Z",
)
# Recent bitsandbytes versions use logging rather than warnings.warn.
# Attach to the emitting logger so propagation/ComfyUI handlers remain
# untouched. Limit this route to the current inference thread as well.
logger = logging.getLogger(_BNB_LOGGER)
log_filter = _QuantizedCastLogFilter()
logger.addFilter(log_filter)
try:
yield
finally:
logger.removeFilter(log_filter)
+7 -9
View File
@@ -8,11 +8,9 @@ CaptionForge Global Model Cache Manager
- Repository:
https://github.com/Damkohler/CaptionForge
- CaptionForge focuses on practical dataset-captioning infrastructure for
LoRA dataset preparation, using multi-engine caption generation, JSONL
audit trails, claim extraction and refinement, text-LLM distillation,
image-aware VLM validation, and consensus-oriented caption improvement
to produce grounded, auditable training captions.
- CaptionForge 1.0 uses independent Pass-A witnesses, text-LLM synthesis,
image-aware validation, SHORT/TAGGY formatting, and JSONL audit trails to
produce grounded LoRA dataset captions.
- Engine Purpose
- The **CaptionForge Global Model Cache Manager** provides shared,
@@ -75,10 +73,10 @@ CaptionForge Global Model Cache Manager
- The module favors predictable behavior over aggressive automatic memory
management.
- Development Status
- CaptionForge v0.1.0 experimental developer-preview infrastructure.
- Cache policy and diagnostics may evolve as CaptionForge's supported engine
set matures.
- Production Status
- Active CaptionForge 1.0 runtime support. Its conservative single-model
residency policy coordinates heavyweight local Hugging Face witnesses;
it does not own caption semantics.
- Attribution & License
- Concept and implementation by **J. L. Córdova**
+113 -46
View File
@@ -8,11 +8,9 @@ CaptionForge Pipeline Planner Engine
- Repository:
https://github.com/Damkohler/CaptionForge
- CaptionForge focuses on practical dataset-captioning infrastructure for
LoRA dataset preparation, using multi-engine caption generation, JSONL
audit trails, claim extraction and refinement, text-LLM distillation,
image-aware VLM validation, and consensus-oriented caption improvement
to produce grounded, auditable training captions.
- CaptionForge 1.0 uses independent Pass-A witnesses, text-LLM synthesis,
image-aware validation, SHORT/TAGGY formatting, and JSONL audit trails to
produce grounded LoRA dataset captions.
- Engine Purpose
- The **CaptionForge Pipeline Planner Engine** builds reusable
@@ -71,10 +69,10 @@ CaptionForge Pipeline Planner Engine
- The planner favors explicit dictionaries, predictable filenames, and
auditable JSONL-oriented handoff between passes.
- Development Status
- CaptionForge v0.1.0 experimental developer-preview infrastructure.
- Plan schema, supported witness families, and downstream defaults may evolve
before a stable CaptionForge release.
- Production Status
- Active CaptionForge 1.0 planning support. The Planner is authoritative in
connected workflows, and its shared B/C/D defaults are kept in parity with
the standalone Orchestrator controls.
- Attribution & License
- Concept and implementation by **J. L. Córdova**
@@ -106,17 +104,27 @@ MANIFEST = {
}
import hashlib
import json
import random
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from .captionforge_prompt_defaults import (
DEFAULT_FAT_DRAFT_INSTRUCTIONS,
DEFAULT_TAGGY_FORMATTER_INSTRUCTIONS,
DEFAULT_VALIDATOR_INSTRUCTIONS,
DEFAULT_VALIDATOR_SYSTEM_PROMPT,
)
from .captionforge_cleanup import normalize_forbidden_phrases, normalize_replace_pairs
MAX_SEED_32 = 0xFFFFFFFF
PIPELINE_PLAN_TYPE = "captionforge_pipeline_plan"
PIPELINE_PLAN_VERSION = 5
PIPELINE_PLAN_VERSION = 6
DEFAULT_OLLAMA_URL = "http://127.0.0.1:11434"
MAX_PASS_A_RUNS_PER_MODEL = 5
PASS_A_SEED_NAMESPACE = b"captionforge-pass-a-seed-v1"
@dataclass(frozen=True)
@@ -213,22 +221,26 @@ def _normalize_seed_mode(value: Any) -> str:
def _seed_for_run(base_seed: int, seed_mode: str, index: int) -> int | None:
if base_seed < 0:
if seed_mode == "random":
return random.SystemRandom().randint(0, MAX_SEED_32)
return None
base_seed = max(0, min(int(base_seed), MAX_SEED_32))
index = int(index)
if index < 0:
raise ValueError("CaptionForge Pass-A run index must be nonnegative")
if seed_mode == "fixed":
return base_seed
if seed_mode == "increment":
return min(MAX_SEED_32, base_seed + index)
return (base_seed + index) & MAX_SEED_32
if seed_mode == "decrement":
return max(0, base_seed - index)
return (base_seed - index) & MAX_SEED_32
if seed_mode == "random":
rng = random.Random(base_seed)
out = base_seed
for _ in range(index + 1):
out = rng.randint(0, MAX_SEED_32)
return out
payload = b"\0".join(
(
PASS_A_SEED_NAMESPACE,
str(base_seed).encode("ascii"),
str(index).encode("ascii"),
)
)
return int.from_bytes(hashlib.blake2s(payload, digest_size=4).digest(), "big")
return base_seed
@@ -364,8 +376,8 @@ def build_captionforge_pipeline_plan(
run_name: str = "captionforge_run",
overwrite_outputs: bool = True,
joy_runs_per_image: Any = 2,
qwen_runs_per_image: Any = 2,
ollama_runs_per_image: Any = "Disabled",
qwen_runs_per_image: Any = 1,
ollama_runs_per_image: Any = 1,
ollama_caption_runs_per_image: Any | None = None,
caption_ollama_runs_per_image: Any | None = None,
ollama_vlm_runs_per_image: Any | None = None,
@@ -377,12 +389,21 @@ def build_captionforge_pipeline_plan(
top_p_schedule: str = "",
top_k_schedule: str = "",
max_size: int = 1024,
max_new_tokens: int = 512,
max_new_tokens: int = 4096,
trigger_word: str = "",
user_caption_anchor: str = "",
forbidden_phrases: Any = "",
replace_pairs: Any = "",
ollama_url: str = DEFAULT_OLLAMA_URL,
ollama_keep_loaded: bool = True,
ollama_request_timeout_seconds: int = 1800,
distiller_model_family: str = "Llama",
distiller_base_seed: int | None = None,
distiller_prompt: str = DEFAULT_FAT_DRAFT_INSTRUCTIONS,
# Compatibility-only. Pass B owns one fixed seed; modes are ignored.
distiller_seed_mode: str = "fixed",
# Compatibility-only argument for older callers. Production Pass B is
# always one global fat-draft call per image.
distiller_strategy: str = "single_pass",
distiller_max_caption_chars_for_llm: int = 1536,
distiller_num_predict: int = 3096,
@@ -393,21 +414,36 @@ def build_captionforge_pipeline_plan(
distiller_preserve_raw_response: bool = False,
validator_model_family: str = "Llama Vision",
validator_base_seed: int | None = None,
validator_system_prompt: str = DEFAULT_VALIDATOR_SYSTEM_PROMPT,
validator_prompt: str = DEFAULT_VALIDATOR_INSTRUCTIONS,
# Compatibility-only. Pass C owns one fixed seed; modes are ignored.
validator_seed_mode: str = "fixed",
validator_num_predict: int = 2200,
validator_num_predict: int = 2112,
validator_temperature: float = 0.0,
validator_top_p: float = 0.92,
validator_top_k: int = 80,
validator_write_prompt_jsonl: bool = False,
validator_preserve_raw_vlm_response: bool = False,
formatter_model: str = "mistral-small:24b",
formatter_prompt: str = DEFAULT_TAGGY_FORMATTER_INSTRUCTIONS,
formatter_num_predict: int = 3200,
formatter_temperature: float = 0.12,
formatter_top_p: float = 0.88,
formatter_top_k: int = 50,
formatter_write_prompt_jsonl: bool = False,
formatter_preserve_raw_response: bool = False,
# Compatibility-only argument for older callers. Production always exports
# the invariant long/short/taggy sidecar set.
final_caption_style: str = "narrative",
final_write_txt_sidecars: bool = True,
final_write_jsonl: bool = True,
dataset_export: dict[str, Any] | None = None,
# Legacy compatibility aliases retained for older callers.
distiller_seed: int | None = None,
validator_seed: int | None = None,
distiller_model: str = "llama3.1:8b",
validator_model: str = "llama3.2-vision:11b",
formatter_seed: int | None = None,
distiller_model: str = "mistral-small:24b",
validator_model: str = "gemma4:26b",
captions_per_image: int | None = None,
# Deprecated compatibility-only parameters. Intentionally ignored.
pass_c_deterministic: bool = True,
@@ -415,9 +451,9 @@ def build_captionforge_pipeline_plan(
semantic_profile: str = "",
) -> dict[str, Any]:
joy_runs = _normalize_runs_per_image(joy_runs_per_image, 2)
qwen_runs = _normalize_runs_per_image(qwen_runs_per_image, 2)
qwen_runs = _normalize_runs_per_image(qwen_runs_per_image, 1)
ollama_runs = max(
_normalize_runs_per_image(value, "Disabled")
_normalize_runs_per_image(value, 1)
for value in (
ollama_runs_per_image,
ollama_caption_runs_per_image,
@@ -440,14 +476,14 @@ def build_captionforge_pipeline_plan(
base_seed_n = _coerce_int(base_seed, -1, -1, MAX_SEED_32)
seed_mode_n = _normalize_seed_mode(seed_mode)
if distiller_base_seed is None:
distiller_base_seed = distiller_seed
if validator_base_seed is None:
validator_base_seed = validator_seed
if distiller_seed is None:
distiller_seed = distiller_base_seed
if validator_seed is None:
validator_seed = validator_base_seed
d_seed = base_seed_n if distiller_base_seed is None else _coerce_int(distiller_base_seed, base_seed_n, -1, MAX_SEED_32)
v_seed_default = -1 if base_seed_n < 0 else min(MAX_SEED_32, base_seed_n + 1000003)
v_seed = v_seed_default if validator_base_seed is None else _coerce_int(validator_base_seed, v_seed_default, -1, MAX_SEED_32)
d_seed = _coerce_int(distiller_seed, -1, -1, MAX_SEED_32)
v_seed = _coerce_int(validator_seed, -1, -1, MAX_SEED_32)
f_seed = _coerce_int(formatter_seed, -1, -1, MAX_SEED_32)
run_name_n = _clean_name(run_name)
input_path_n = str(input_path or "").strip()
@@ -458,7 +494,7 @@ def build_captionforge_pipeline_plan(
"top_p_schedule": str(top_p_schedule or "").strip(),
"top_k_schedule": str(top_k_schedule or "").strip(),
"max_size": _coerce_int(max_size, 1024, 0, 4096),
"max_new_tokens": _coerce_int(max_new_tokens, 512, 16, 4096),
"max_new_tokens": _coerce_int(max_new_tokens, 4096, 16, 4096),
"trigger_word": str(trigger_word or "").strip(),
"user_caption_anchor": str(user_caption_anchor or "").strip(),
"output_dir": str(output_dir or "").strip(),
@@ -470,6 +506,14 @@ def build_captionforge_pipeline_plan(
"run_name": run_name_n,
"overwrite_outputs": _coerce_bool(overwrite_outputs, True),
}
forbidden = normalize_forbidden_phrases(forbidden_phrases)
replacements = normalize_replace_pairs(replace_pairs)
cleanup = {
"forbidden_phrases": forbidden,
"replace_pairs": [{"old": old, "new": new} for old, new in replacements],
"matching": "boundary_safe_case_insensitive",
"order": ["replace_pairs", "forbidden_phrases", "normalize_whitespace_punctuation"],
}
shared["captions_per_image"] = (
_coerce_int(captions_per_image, max(joy_runs, qwen_runs, ollama_runs, florence_runs, llama_runs, 1), 1, 100)
if captions_per_image is not None
@@ -483,14 +527,18 @@ def build_captionforge_pipeline_plan(
# Pass A JSONL/audit artifacts, so point it at the run working directory.
shared["output_dir"] = paths["output_dir"]
ollama = {
"url": str(ollama_url or DEFAULT_OLLAMA_URL).strip() or DEFAULT_OLLAMA_URL,
"keep_loaded": _coerce_bool(ollama_keep_loaded, True),
"request_timeout_seconds": _coerce_int(ollama_request_timeout_seconds, 1800, 10, 7200),
}
distiller = {
"backend": "ollama",
"model_family": str(distiller_model_family or "Llama").strip() or "Llama",
"model": str(distiller_model or "llama3.1:8b").strip() or "llama3.1:8b",
"base_seed": d_seed,
"model": str(distiller_model or "mistral-small:24b").strip() or "mistral-small:24b",
"prompt": str(distiller_prompt or DEFAULT_FAT_DRAFT_INSTRUCTIONS),
"seed": d_seed,
"seed_mode": _normalize_seed_mode(distiller_seed_mode),
"strategy": str(distiller_strategy or "single_pass").strip() or "single_pass",
"max_caption_chars_for_llm": _coerce_int(distiller_max_caption_chars_for_llm, 1536, 0, 12000),
"num_predict": _coerce_int(distiller_num_predict, 3096, 64, 12000),
"temperature": _coerce_float(distiller_temperature, 0.24, 0.0, 2.0),
@@ -503,12 +551,12 @@ def build_captionforge_pipeline_plan(
validator = {
"backend": "ollama",
"model_family": str(validator_model_family or "Llama Vision").strip() or "Llama Vision",
"model": str(validator_model or "llama3.2-vision:11b").strip() or "llama3.2-vision:11b",
"base_seed": v_seed,
"model": str(validator_model or "gemma4:26b").strip() or "gemma4:26b",
"system_prompt": str(validator_system_prompt or DEFAULT_VALIDATOR_SYSTEM_PROMPT),
"prompt": str(validator_prompt or DEFAULT_VALIDATOR_INSTRUCTIONS),
"seed": v_seed,
"seed_mode": _normalize_seed_mode(validator_seed_mode),
"image_root": input_path_n,
"num_predict": _coerce_int(validator_num_predict, 2200, 64, 12000),
"num_predict": _coerce_int(validator_num_predict, 2112, 64, 12000),
"temperature": _coerce_float(validator_temperature, 0.0, 0.0, 2.0),
"top_p": _coerce_float(validator_top_p, 0.92, 0.0, 1.0),
"top_k": _coerce_int(validator_top_k, 80, 0, 500),
@@ -516,8 +564,20 @@ def build_captionforge_pipeline_plan(
"preserve_raw_vlm_response": _coerce_bool(validator_preserve_raw_vlm_response, False),
"role": "image_aware_precision_validation",
}
formatter = {
"backend": "ollama",
"model": str(formatter_model or "mistral-small:24b").strip() or "mistral-small:24b",
"prompt": str(formatter_prompt or DEFAULT_TAGGY_FORMATTER_INSTRUCTIONS),
"seed": f_seed,
"num_predict": _coerce_int(formatter_num_predict, 3200, 64, 12000),
"temperature": _coerce_float(formatter_temperature, 0.12, 0.0, 2.0),
"top_p": _coerce_float(formatter_top_p, 0.88, 0.0, 1.0),
"top_k": _coerce_int(formatter_top_k, 50, 0, 500),
"write_prompt_jsonl": _coerce_bool(formatter_write_prompt_jsonl, False),
"preserve_raw_response": _coerce_bool(formatter_preserve_raw_response, False),
"role": "text_only_taggy_formatting",
}
final = {
"caption_style": str(final_caption_style or "narrative").strip() or "narrative",
"write_txt_sidecars": _coerce_bool(final_write_txt_sidecars, True),
"write_jsonl": _coerce_bool(final_write_jsonl, True),
"overwrite_outputs": _coerce_bool(overwrite_outputs, True),
@@ -525,10 +585,15 @@ def build_captionforge_pipeline_plan(
"large_model_passes_after_validator": False,
}
from .captionforge_dataset_export import normalize_export_settings
return {
"dataset_export": normalize_export_settings(dataset_export),
"captionforge_config_type": PIPELINE_PLAN_TYPE,
"captionforge_config_version": PIPELINE_PLAN_VERSION,
"shared": shared,
"cleanup": cleanup,
"ollama": ollama,
"paths": paths,
"pass_a": {
"joy": {"model_key": "joy", "enabled": joy_runs > 0, "runs_per_image": joy_runs, "role": "rich_descriptive_caption_witness"},
@@ -556,9 +621,11 @@ def build_captionforge_pipeline_plan(
},
"distiller": distiller,
"validator": validator,
"formatter": formatter,
"final": final,
"pass_b_distiller": distiller,
"pass_c_vlm_validator": validator,
"pass_d_formatter": formatter,
"final_export": final,
}
@@ -572,7 +639,7 @@ def build_captionforge_run_config(
top_p_schedule: str = "",
top_k_schedule: str = "",
max_size: int = 1024,
max_new_tokens: int = 512,
max_new_tokens: int = 4096,
trigger_word: str = "",
output_dir: str = "",
input_path: str = "",
+78
View File
@@ -0,0 +1,78 @@
"""Frozen CaptionForge 1.0 prompt defaults shared by Planner and Orchestrator.
These strings define the production B/C/D contract: a recall-oriented text
draft, image-grounded LONG validation, and one text-only SHORT/TAGGY format
call. Keeping them here prevents the two control surfaces from drifting.
"""
from __future__ import annotations
DEFAULT_FAT_DRAFT_INSTRUCTIONS = """/no_think
You are a detail-preserving caption merger for LoRA dataset preparation.
You receive multiple captions of the same image. You do NOT see the image.
Task:
Merge all non-contradictory caption details into one deliberately over-complete draft caption.
Rules:
- Do not validate against the image.
- Do not decide that details are false just because they appear once.
- Do not summarize aggressively.
- Preserve concrete details from all captions.
- Split contradictions by choosing cautious wording or listing the alternative only when needed.
- Prefer specific visual language over generic language.
- Keep visible body, clothing, material, accessory, color, pose, lighting, style, and framing details.
- Preserve doll-like, glossy/plastic-like, material, garment-construction, body-shape, and facial-feature details when present.
- Use neutral dataset-caption language, including visible sensual styling or revealing clothing when present.
- Do not add details absent from the captions.
- Treat subject names or trigger-like identity tokens as optional identity labels. Preserve them only when they appear consistently in the captions; do not let them replace visible description.
- Output only one paragraph, no notes, no JSON."""
DEFAULT_VALIDATOR_SYSTEM_PROMPT = (
"/no_think\n"
"You are a direct image validation engine. Inspect the image and answer only with the requested caption."
)
DEFAULT_VALIDATOR_INSTRUCTIONS = """/no_think
Look at the image and validate this draft caption.
Task:
Return a corrected caption paragraph that keeps only image-supported details.
Rules:
- Output only the corrected caption.
- One paragraph.
- No reasoning, no notes, no JSON.
- Keep all true visible details from the draft.
- Delete unsupported details.
- Correct small visible errors.
- Do not add new details unless needed to correct an error already present.
- Preserve useful LoRA details: subject, face, hair, eyes, makeup, lips, skin texture, pose, body shape, outfit, accessories, materials, colors, lighting, background, framing, and visual style.
- Visible sensual styling, revealing clothing, cleavage, thighs, bare skin, swimwear, lingerie, or body-shape details may be described neutrally when present.
- Do not invent hidden anatomy, unseen clothing, explicit acts, or details contradicted by the image."""
DEFAULT_TAGGY_FORMATTER_INSTRUCTIONS = """/no_think
You are a LoRA caption format converter. The validated paragraph is your only source of truth.
Output exactly two labeled lines:
SHORT: <a concise natural-language caption, typically around 100 words, that preserves all LoRA-useful validated details>
TAGGY: <one compact comma-separated caption>
SHORT must preserve the image's distinctive training identity across the whole source:
1. subject, defining face/hair/body traits, and every major outfit piece/material;
2. pose/action and key accessories or unusual visible details;
3. setting, lighting, framing, and visual medium/style.
Omit a category only when absent. Use only source details; never add, infer, euphemize, or correct. Compress wording, not category coverage. Do not copy only the source opening.
Aim for roughly 100 words. Keep it concise, but allow modest variation when needed to preserve important information and finish the caption naturally. Do not cut off a sentence merely to satisfy a word-count target.
TAGGY must preserve all concrete LoRA-useful source details as compact comma-separated phrases.
No markdown, reasoning, notes, or other labels."""
+38
View File
@@ -0,0 +1,38 @@
"""Canonical source-image identities shared by CaptionForge Pass-A witnesses."""
from __future__ import annotations
from pathlib import Path
OPTIONAL_IMAGE_KEY_PREFIX = "captionforge-optional-image://"
def file_source_identity(image_path: Path, input_root: Path) -> tuple[str, str]:
"""Return the readable stem and portable, collision-safe dataset key."""
path = Path(image_path)
root = Path(input_root)
image = path.stem or "image"
if root.is_dir():
try:
return image, path.relative_to(root).as_posix()
except ValueError:
pass
return image, path.name
def optional_image_identity(index: int) -> tuple[str, str]:
"""Return deterministic display/key values for a filename-less IMAGE tensor."""
filename = f"comfy_image_{int(index):04d}.png"
return Path(filename).stem, f"{OPTIONAL_IMAGE_KEY_PREFIX}{filename}"
def optional_image_filename(image_key: object) -> str:
"""Return the synthetic PNG name encoded by an optional-image key, if any."""
text = str(image_key or "")
if not text.startswith(OPTIONAL_IMAGE_KEY_PREFIX):
return ""
filename = text[len(OPTIONAL_IMAGE_KEY_PREFIX):]
if not filename or "/" in filename or "\\" in filename:
return ""
return filename
+11 -14
View File
@@ -8,11 +8,9 @@ CaptionForge VLM Validator Engine
- Repository:
https://github.com/Damkohler/CaptionForge
- CaptionForge focuses on practical dataset-captioning infrastructure for
LoRA dataset preparation, using multi-engine caption generation, JSONL
audit trails, claim extraction and refinement, text-LLM distillation,
image-aware VLM validation, and consensus-oriented caption improvement
to produce grounded, auditable training captions.
- CaptionForge 1.0 uses independent Pass-A witnesses, text-LLM synthesis,
image-aware validation, SHORT/TAGGY formatting, and JSONL audit trails to
produce grounded LoRA dataset captions.
- Engine Purpose
- The **CaptionForge VLM Validator Engine** is the Pass C image-aware
@@ -58,9 +56,8 @@ CaptionForge VLM Validator Engine
- The engine writes structured JSONL records and optional readable sidecars.
- The module also contains an experimental cleaner function used by the
quarantined reversed-pipeline branch; that path is not the recommended
mainline v0.1.0 workflow.
- The module also contains an experimental cleaner retained for historical
diagnostics; it is not part of the production CaptionForge 1.0 workflow.
- Design Philosophy
- The distiller is recall-oriented; the VLM validator is grounding-oriented.
@@ -77,10 +74,10 @@ CaptionForge VLM Validator Engine
- Final candidate captions should remain rich, visually grounded, and
auditable.
- Development Status
- CaptionForge v0.1.0 experimental developer-preview infrastructure.
- Prompt contracts, parser behavior, image-resolution heuristics, and audit
fields may evolve before a stable CaptionForge release.
- Reference Status
- This CLI-oriented engine is retained for reference, diagnostics, and
controlled experiments. It is not imported or registered by the active
ComfyUI pipeline; production Pass C runs inside ``jlc_captionforge_node``.
- Attribution & License
- Concept and implementation by **J. L. Córdova**
@@ -134,7 +131,7 @@ from typing import Any, Iterable, Optional
@dataclass
class VLMValidatorConfig:
engine_name: str = "CaptionForge VLM Validator Engine"
engine_version: str = "0.2.0"
engine_version: str = CAPTIONFORGE_VERSION
vlm_backend: str = "ollama" # ollama, manual_json, prompt_only
vlm_model: str = ""
@@ -2148,7 +2145,7 @@ def _cf_cleaner_record(
return {
"captionforge_pass": CLEANER_PASS,
"engine": "CaptionForge VLM Validator Engine",
"engine_version": getattr(config, "engine_version", "0.2.0"),
"engine_version": getattr(config, "engine_version", CAPTIONFORGE_VERSION),
"contract": "vlm_statement_cleaner_v0.1",
"image_key": _cf_cleaner_record_key(source_record),
"image": str(source_record.get("image") or ""),
@@ -1,477 +0,0 @@
from __future__ import annotations
MANIFEST = {"name":"JLC BLIP Caption Engine","version":(0,1,0),"author":"J. L. Córdova"}
import fnmatch, json, random, re, shutil
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Iterable, Optional
from PIL import Image
import torch
from .captionforge_model_cache import make_cache_key, get_cached_model, register_model, prepare_for_model_load, unload_after_run
@dataclass(frozen=True)
class BlipModelInfo:
repo_id: str
local_folder: str
model_type: str
notes: str = ""
MODEL_REGISTRY: dict[str, BlipModelInfo] = {
"BLIP-2 FLAN-T5 XL": BlipModelInfo("Salesforce/blip2-flan-t5-xl", "blip2-flan-t5-xl", "blip2", "Advanced BLIP-2 checkpoint with FLAN-T5 XL."),
"BLIP-2 OPT 2.7B": BlipModelInfo("Salesforce/blip2-opt-2.7b", "blip2-opt-2.7b", "blip2", "BLIP-2 checkpoint with OPT-2.7B."),
"BLIP Large": BlipModelInfo("Salesforce/blip-image-captioning-large", "blip-image-captioning-large", "blip", "Classic BLIP large captioning baseline."),
"BLIP Base": BlipModelInfo("Salesforce/blip-image-captioning-base", "blip-image-captioning-base", "blip", "Classic BLIP base captioning baseline."),
}
SUPPORTED_EXTENSIONS={".jpg",".jpeg",".png",".webp",".bmp",".tif",".tiff"}
PROMPT_PRESETS={
"caption":"",
"detailed":"Describe this image in detail.",
"lora_literal":"Describe this image as a concise visual caption for image dataset training. Mention only visible subject, pose, clothing, appearance, style, lighting, and background.",
"qa_detail":"Question: Describe this image in detailed visual terms. Answer:",
"qa_lora":"Question: What visible details should be included in a text-to-image training caption for this image? Answer:",
}
DEFAULT_PROMPT_PRESET="detailed"
MEMORY_MODES={"Default":{}, "Balanced (8-bit)":{"load_in_8bit":True}}
@dataclass
class GenerationConfig:
max_new_tokens:int=256
num_beams:int=3
temperature:float=0.0
top_p:float=0.90
top_k:int=50
repetition_penalty:float=1.0
seed:Optional[int]=None
@dataclass
class CleanupConfig:
trigger:str=""
prefix:str=""
suffix:str=""
forbidden_phrases:list[str]=field(default_factory=list)
replacement_rules:list[tuple[str,str]]=field(default_factory=list)
replace_case_insensitive:bool=True
replace_whole_words_only:bool=False
strip_boilerplate_prefixes:bool=True
strip_trailing_period:bool=True
@dataclass
class BlipCaptionConfig:
model_name:str="BLIP-2 FLAN-T5 XL"
model_path:str=""
model_root:str="models/LLM/JLC_BLIPCaption"
memory_mode:str="Default"
dtype:str="bf16"
device:str="auto"
trust_remote_code:bool=False
keep_loaded:bool=True
cache_policy:str=""
quiet_transformers_load:bool=True
max_size:int=768
prompt:str=PROMPT_PRESETS[DEFAULT_PROMPT_PRESET]
allow_download:bool=True
use_comfy_model_management:bool=True
@dataclass
class BatchCaptionConfig:
input_path:str=""
recursive:bool=True
filename_glob:str="*"
extensions:set[str]=field(default_factory=lambda:set(SUPPORTED_EXTENSIONS))
output_dir:str=""
write_txt:bool=True
write_jsonl:bool=False
jsonl_filename:str="captions.jsonl"
also_jsonl_path:str=""
write_run_config:bool=True
run_config_filename:str=""
overwrite:bool=False
backup_existing:bool=True
dry_run:bool=False
limit:int=0
skip_existing_txt:bool=True
skip_existing_jsonl_images:bool=False
@dataclass
class CaptionRecord:
image:str
caption:str
raw_caption:str
model_name:str
model_path:str
prompt:str
seed:Optional[int]
temperature:float
top_p:float
top_k:int
num_beams:int
max_new_tokens:int
max_size:int
timestamp:str
status:str="ok"
error:str=""
captionforge_pass:str="A"
model_family:str="blip"
ensemble_run_index:int=0
image_key:str=""
@dataclass
class BatchCaptionResult:
records:list[CaptionRecord]=field(default_factory=list)
skipped:int=0
failed:int=0
@property
def processed(self): return len([r for r in self.records if r.status=="ok"])
@property
def captions_text(self): return "\n\n".join(r.caption for r in self.records if r.status=="ok")
@property
def jsonl_text(self): return "\n".join(json.dumps(record_to_json(r), ensure_ascii=False) for r in self.records)
def timestamp(): return datetime.now().strftime("%Y%m%d_%H%M%S")
def iso_timestamp(): return datetime.now().isoformat(timespec="seconds")
def safe_mkdir(path:Path): path.mkdir(parents=True, exist_ok=True)
def set_seed(seed:int):
random.seed(seed); torch.manual_seed(seed)
if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
def resolve_prompt(prompt_preset:str=DEFAULT_PROMPT_PRESET, custom_prompt:str="") -> str:
custom=(custom_prompt or "").strip()
return custom if custom else PROMPT_PRESETS.get(prompt_preset, PROMPT_PRESETS[DEFAULT_PROMPT_PRESET])
UNWANTED_PREFIXES=["The image depicts ","The image shows ","This image depicts ","This image shows ","An image of ","A photo of ","A photograph of ","The photo shows ","This photo shows "]
def normalize_caption(caption:str, strip_boilerplate_prefixes=True, strip_trailing_period=True)->str:
text=str(caption or "").strip()
if strip_boilerplate_prefixes:
for p in UNWANTED_PREFIXES:
if text.lower().startswith(p.lower()): text=text[len(p):].strip(); break
text=text.replace("\n"," "); text=re.sub(r"\s+"," ",text).strip(); text=re.sub(r"\s+,",",",text); text=re.sub(r",\s*,+",",",text); text=re.sub(r"\s+\.",".",text)
if strip_trailing_period and text.endswith("."): text=text[:-1].strip()
return text
def apply_replacements(caption, rules, case_insensitive=True, whole_words_only=False):
result=caption; flags=re.IGNORECASE if case_insensitive else 0
for old,new in rules or []:
old=old.strip()
if not old: continue
pat=re.escape(old); pat=(r"\b"+pat+r"\b") if whole_words_only else pat
result=re.sub(pat, new.strip(), result, flags=flags)
return result
def remove_forbidden_phrases(caption, forbidden_phrases):
result=caption
for phrase in forbidden_phrases or []:
phrase=phrase.strip()
if phrase: result=re.sub(re.escape(phrase), "", result, flags=re.IGNORECASE)
result=re.sub(r"\s+,",",",result); result=re.sub(r",\s*,+",",",result); result=re.sub(r"\s+"," ",result)
return result.strip(" ,")
def add_trigger_prefix_suffix(caption, trigger="", prefix="", suffix=""):
parts=[]
for item in (trigger,prefix):
item=(item or "").strip().strip(" ,")
if item: parts.append(item)
if caption: parts.append(caption.lstrip(" ,"))
final=", ".join(parts); final=re.sub(r",\s*,+",",",final).strip(" ,")
suffix=(suffix or "").strip()
return (final+suffix).strip() if suffix else final.strip()
def cleanup_caption(caption, config:CleanupConfig):
text=normalize_caption(caption, config.strip_boilerplate_prefixes, config.strip_trailing_period)
text=apply_replacements(text, config.replacement_rules, config.replace_case_insensitive, config.replace_whole_words_only)
text=remove_forbidden_phrases(text, config.forbidden_phrases)
text=normalize_caption(text, config.strip_boilerplate_prefixes, config.strip_trailing_period)
return add_trigger_prefix_suffix(text, config.trigger, config.prefix, config.suffix)
def iter_image_files(input_path, recursive=True, filename_glob="*", extensions=None):
p=Path(input_path); exts=extensions or set(SUPPORTED_EXTENSIONS); filename_glob=(filename_glob or "*").strip() or "*"
if p.is_file():
if p.suffix.lower() in exts and fnmatch.fnmatch(p.name, filename_glob): yield p
return
if not p.exists(): raise FileNotFoundError(f"input_path does not exist: {p}")
if not p.is_dir(): raise NotADirectoryError(f"input_path is not a file or folder: {p}")
iterator=p.rglob("*") if recursive else p.glob("*")
for child in sorted(iterator):
if child.is_file() and child.suffix.lower() in exts and fnmatch.fnmatch(child.name, filename_glob): yield child
def load_image_file(path): return Image.open(path).convert("RGB")
def resize_for_model(image, max_size:int):
if max_size<=0: return image
w,h=image.size; longest=max(w,h)
if longest<=max_size: return image
scale=max_size/float(longest)
return image.resize((max(1,round(w*scale)), max(1,round(h*scale))), Image.Resampling.LANCZOS)
def sidecar_txt_path(image_path:Path, output_dir=None):
if output_dir:
out=Path(output_dir); safe_mkdir(out); return out/f"{image_path.stem}.txt"
return image_path.with_suffix(".txt")
def backup_existing_file(path:Path, dry_run=False):
if not path.exists(): return None
b=path.with_name(f"{path.name}.bak_{timestamp()}")
if not dry_run: shutil.copy2(path,b)
return b
def write_text_sidecar(path:Path, text:str, overwrite=False, backup_existing=True, dry_run=False):
if path.exists() and not overwrite: return False
if dry_run: return True
safe_mkdir(path.parent)
if path.exists() and overwrite and backup_existing: backup_existing_file(path, False)
path.write_text(text.rstrip()+"\n", encoding="utf-8"); return True
def load_existing_jsonl_images(path:Path):
seen=set()
if not path.exists(): return seen
with path.open("r", encoding="utf-8") as f:
for line in f:
try: rec=json.loads(line)
except json.JSONDecodeError: continue
val=rec.get("image") or rec.get("image_path") or rec.get("source")
if val:
seen.add(str(val))
try: seen.add(str(Path(val)))
except Exception: pass
return seen
def record_to_json(record): return asdict(record)
def append_jsonl_records(path:Path, records, dry_run=False):
if dry_run or not records: return
safe_mkdir(path.parent)
with path.open("a", encoding="utf-8") as f:
for r in records: f.write(json.dumps(record_to_json(r), ensure_ascii=False)+"\n")
def json_safe(v):
if isinstance(v,set): return sorted(v)
if isinstance(v,Path): return str(v)
if isinstance(v,dict): return {k:json_safe(x) for k,x in v.items()}
if isinstance(v,list): return [json_safe(x) for x in v]
if isinstance(v,tuple): return [json_safe(x) for x in v]
return v
def write_run_config_json(path:Path, config, dry_run=False):
if dry_run: return
safe_mkdir(path.parent); path.write_text(json.dumps(json_safe(config), indent=2, ensure_ascii=False)+"\n", encoding="utf-8")
def model_folder_has_weights(local_path:Path):
if not local_path.exists() or not local_path.is_dir(): return False
pats=["*.safetensors","*.bin","*.pt","*.pth","*.gguf","*.ckpt"]
return any(any(local_path.rglob(p)) for p in pats)
def get_model_info(model_name):
if model_name not in MODEL_REGISTRY: raise KeyError(f"Unknown BLIP model_name: {model_name}. Known models: {', '.join(MODEL_REGISTRY)}")
return MODEL_REGISTRY[model_name]
def get_registry_model_path(model_name, model_root): return Path(model_root)/get_model_info(model_name).local_folder
def download_registry_model_if_needed(model_name, model_root, metadata_only=False, allow_download=True):
local_path=get_registry_model_path(model_name, model_root)
if not metadata_only and model_folder_has_weights(local_path): return local_path
if not allow_download:
if metadata_only: safe_mkdir(local_path); return local_path
raise FileNotFoundError(f"Model folder does not contain weights and allow_download=False: {local_path}")
try: from huggingface_hub import snapshot_download
except Exception as exc: raise RuntimeError("Missing dependency: huggingface_hub. Install: pip install huggingface_hub") from exc
info=get_model_info(model_name); safe_mkdir(Path(model_root)); safe_mkdir(local_path)
if metadata_only:
print(f"[JLC BLIP Engine] Probe download for {info.repo_id} -> {local_path}")
snapshot_download(repo_id=info.repo_id, local_dir=str(local_path), ignore_patterns=["*.safetensors","*.bin","*.pt","*.pth","*.gguf","*.onnx","*.ckpt","*.h5","*.msgpack","*.tflite"])
return local_path
print(f"[JLC BLIP Engine] Downloading {info.repo_id} -> {local_path}")
snapshot_download(repo_id=info.repo_id, local_dir=str(local_path)); return local_path
def probe_registry_model_download(model_name, model_root):
local_path=download_registry_model_if_needed(model_name, model_root, metadata_only=True, allow_download=True)
files=[]
try:
for p in sorted(local_path.rglob("*")):
if p.is_file(): files.append(str(p.relative_to(local_path)))
except Exception: files=[]
preview="\n".join(files[:40])
if len(files)>40: preview += f"\n... plus {len(files)-40} more files"
return f"JLC BLIP Caption download probe completed.\n\nModel: {model_name}\nFolder: {local_path}\n\nLarge model weight files were intentionally skipped.\n\nFiles found:\n{preview if preview else '(no files listed)'}"
def _try_get_comfy_model_management():
try:
import comfy.model_management as model_management
return model_management
except Exception: return None
def _unload_blip_bundle(bundle):
model=bundle.get("model") if isinstance(bundle,dict) else None
if model is not None:
try: model.to("cpu")
except Exception: pass
try: del model
except Exception: pass
if isinstance(bundle,dict): bundle.clear()
class BlipCaptionEngine:
def __init__(self, config:BlipCaptionConfig, generation=None, cleanup=None):
self.config=config; self.generation=generation or GenerationConfig(); self.cleanup=cleanup or CleanupConfig()
self.processor=None; self.model=None; self.local_model_path=None; self.model_size_bytes=None
self._comfy_mm=_try_get_comfy_model_management() if config.use_comfy_model_management else None
self.inference_device=self._resolve_inference_device(config.device); self.offload_device=self._resolve_offload_device(config.device)
def resolve_model_path(self):
if self.config.model_path.strip(): return Path(self.config.model_path).expanduser()
return download_registry_model_if_needed(self.config.model_name, self.config.model_root, False, self.config.allow_download)
def _cache_policy(self):
p=getattr(self.config,"cache_policy",""); p=p.strip() if isinstance(p,str) else ""
return p or ("evict_other_caption_models" if getattr(self.config,"keep_loaded",True) else "unload_after_run")
def _cache_key(self, local_path:Path):
info=get_model_info(self.config.model_name)
return make_cache_key(role="caption", family=f"blip-{info.model_type}", model_path=str(local_path.resolve()), device=str(self.inference_device), quantization=self.config.memory_mode, dtype=str(self.config.dtype))
def _resolve_inference_device(self, device):
if self._comfy_mm is not None:
try:
dev=self._comfy_mm.get_torch_device()
if dev.type!="cuda": raise RuntimeError("CaptionForge BLIP requires CUDA for release builds. Silent CPU fallback is disabled.")
return dev
except RuntimeError: raise
except Exception: pass
if device=="auto":
if not torch.cuda.is_available(): raise RuntimeError("CaptionForge BLIP requires CUDA for release builds. Silent CPU fallback is disabled.")
return torch.device("cuda")
dev=torch.device(device)
if dev.type!="cuda": raise RuntimeError(f"CaptionForge BLIP release mode does not allow inference device={device!r}.")
return dev
def _resolve_offload_device(self, device):
if self._comfy_mm is not None:
try: return self._comfy_mm.unet_offload_device()
except Exception: pass
return torch.device("cpu")
@staticmethod
def _resolve_dtype(dtype):
t=dtype.lower().strip()
if t=="auto": return "auto"
if t in {"bf16","bfloat16"}: return torch.bfloat16
if t in {"fp16","float16"}: return torch.float16
if t in {"fp32","float32"}: return torch.float32
raise ValueError(f"Unsupported dtype: {dtype}")
def _module_size(self, module):
if self._comfy_mm is not None:
try: return int(self._comfy_mm.module_size(module))
except Exception: pass
total=0
try:
for p in module.parameters(): total += p.numel()*p.element_size()
for b in module.buffers(): total += b.numel()*b.element_size()
except Exception: total=0
return total
def _free_memory(self, bytes_needed, device):
if bytes_needed is None: return
if self._comfy_mm is not None:
try: self._comfy_mm.free_memory(bytes_needed, device); return
except Exception: pass
if device.type=="cuda" and torch.cuda.is_available(): torch.cuda.empty_cache()
def _soft_empty_cache(self):
if self._comfy_mm is not None:
try: self._comfy_mm.soft_empty_cache(); return
except Exception: pass
if torch.cuda.is_available(): torch.cuda.empty_cache()
def load(self):
local_path=self.resolve_model_path()
if not local_path.exists(): raise FileNotFoundError(f"Model path does not exist: {local_path}")
self.local_model_path=local_path; key=self._cache_key(local_path); cached=get_cached_model(key)
if cached is not None:
self.processor=cached.get("processor"); self.model=cached.get("model"); self.model_size_bytes=cached.get("model_size_bytes")
print(f"[JLC BLIP Engine] Reusing cached model bundle: {local_path}"); return
prepare_for_model_load(key, policy=self._cache_policy(), role="caption")
if self.generation.seed is not None: set_seed(self.generation.seed)
info=get_model_info(self.config.model_name)
try:
if info.model_type=="blip2":
from transformers import Blip2Processor as ProcessorClass, Blip2ForConditionalGeneration as ModelClass
elif info.model_type=="blip":
from transformers import BlipProcessor as ProcessorClass, BlipForConditionalGeneration as ModelClass
else: raise RuntimeError(f"Unsupported BLIP model_type: {info.model_type}")
except Exception as exc: raise RuntimeError("Could not import BLIP Transformers classes. Update transformers/accelerate/pillow/huggingface_hub.") from exc
torch_dtype=self._resolve_dtype(self.config.dtype)
print(f"[JLC BLIP Engine] Loading processor: {local_path}")
self.processor=ProcessorClass.from_pretrained(str(local_path), trust_remote_code=self.config.trust_remote_code)
print(f"[JLC BLIP Engine] Loading model: {local_path} dtype={self.config.dtype} mode={self.config.memory_mode}")
kwargs={"torch_dtype":torch_dtype, "trust_remote_code":self.config.trust_remote_code}
if self.config.memory_mode=="Balanced (8-bit)":
try: from transformers import BitsAndBytesConfig
except Exception as exc: raise RuntimeError("BLIP Balanced (8-bit) requires bitsandbytes. Use Default or install bitsandbytes.") from exc
kwargs["quantization_config"]=BitsAndBytesConfig(load_in_8bit=True); kwargs["device_map"]={"": self.inference_device.index or 0}
elif self.config.memory_mode!="Default": raise ValueError(f"Unsupported memory_mode: {self.config.memory_mode}")
self.model=ModelClass.from_pretrained(str(local_path), **kwargs); self.model.eval(); self.model_size_bytes=self._module_size(self.model)
if self.config.memory_mode=="Default": self._free_memory(self.model_size_bytes, self.inference_device); self.model.to(self.inference_device)
register_model(key, {"processor":self.processor,"model":self.model,"model_size_bytes":self.model_size_bytes,"local_path":str(local_path)}, family="blip", model_path=str(local_path), device=str(self.inference_device), quantization=self.config.memory_mode, role="caption", unload_fn=_unload_blip_bundle, keep=self._cache_policy()=="keep_this_model")
print(f"[JLC BLIP Engine] Loaded model: {self.config.model_name}")
def prepare_for_inference(self):
if self.processor is None or self.model is None: self.load()
if self.model is None: raise RuntimeError("BLIP model failed to load.")
if self.config.memory_mode=="Default": self._free_memory(self.model_size_bytes, self.inference_device); self.model.to(self.inference_device)
def cleanup_after_inference(self):
if self._cache_policy()!="unload_after_run": return
if self.model is not None and self.config.memory_mode=="Default": self.model.to(self.offload_device); self._soft_empty_cache()
def unload(self):
if self.local_model_path is not None: unload_after_run(self._cache_key(self.local_model_path), enabled=True)
self.processor=None; self.model=None; self.model_size_bytes=None; self._soft_empty_cache()
@torch.inference_mode()
def caption_pil(self, image:Image.Image):
if self.generation.seed is not None: set_seed(self.generation.seed)
self.prepare_for_inference()
if self.model is None or self.processor is None: raise RuntimeError("Model is not loaded.")
image_for_model=resize_for_model(image.convert("RGB"), self.config.max_size); prompt=(self.config.prompt or "").strip()
inputs=self.processor(images=image_for_model, text=prompt, return_tensors="pt") if prompt else self.processor(images=image_for_model, return_tensors="pt")
inputs={k:(v.to(self.inference_device) if hasattr(v,"to") else v) for k,v in inputs.items()}
model_dtype=getattr(self.model,"dtype",None)
if isinstance(model_dtype, torch.dtype) and "pixel_values" in inputs and torch.is_floating_point(inputs["pixel_values"]): inputs["pixel_values"]=inputs["pixel_values"].to(dtype=model_dtype)
gen={"max_new_tokens":int(self.generation.max_new_tokens), "num_beams":max(1,int(self.generation.num_beams)), "do_sample":bool(self.generation.temperature>0), "use_cache":True}
if self.generation.temperature>0: gen.update({"temperature":float(self.generation.temperature),"top_p":float(self.generation.top_p),"top_k":int(self.generation.top_k)})
if self.generation.repetition_penalty and self.generation.repetition_penalty!=1.0: gen["repetition_penalty"]=float(self.generation.repetition_penalty)
try: ids=self.model.generate(**inputs, **gen)
finally: self.cleanup_after_inference()
raw=self.processor.batch_decode(ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0].strip()
return cleanup_caption(raw,self.cleanup), raw
def caption_path(self, image_path):
p=Path(image_path)
if not p.exists(): raise FileNotFoundError(f"Image does not exist: {p}")
final,raw=self.caption_pil(load_image_file(p))
return CaptionRecord(str(p), final, raw, self.config.model_name, str(self.local_model_path or self.config.model_path), self.config.prompt, self.generation.seed, self.generation.temperature, self.generation.top_p, self.generation.top_k, self.generation.num_beams, self.generation.max_new_tokens, self.config.max_size, iso_timestamp(), captionforge_pass="A", model_family="blip", ensemble_run_index=0, image_key=str(p.resolve()))
def caption_batch(self, batch:BatchCaptionConfig):
if not batch.input_path.strip(): raise ValueError("BatchCaptionConfig.input_path is required.")
result=BatchCaptionResult(); images=list(iter_image_files(batch.input_path,batch.recursive,batch.filename_glob,batch.extensions))
if batch.limit and batch.limit>0: images=images[:int(batch.limit)]
if not images: print(f"[JLC BLIP Engine] No images found in: {batch.input_path}"); return result
output_dir=Path(batch.output_dir) if batch.output_dir.strip() else None; jsonl_path=None
if batch.write_jsonl: jsonl_path=(output_dir or images[0].parent)/(batch.jsonl_filename.strip() or "captions.jsonl")
also_jsonl_path=Path(batch.also_jsonl_path) if batch.also_jsonl_path.strip() else None
seen=set()
if batch.skip_existing_jsonl_images:
if jsonl_path: seen.update(load_existing_jsonl_images(jsonl_path))
if also_jsonl_path: seen.update(load_existing_jsonl_images(also_jsonl_path))
if batch.write_run_config:
base=output_dir or images[0].parent; name=batch.run_config_filename.strip() or f"jlc_blip_caption_run_config_{timestamp()}.json"; write_run_config_json(base/name, self.build_run_config(batch), batch.dry_run)
print(f"[JLC BLIP Engine] Found {len(images)} image(s).")
try:
for idx,path in enumerate(images,1):
try:
txt=sidecar_txt_path(path, output_dir); src=str(path)
if batch.skip_existing_txt and batch.write_txt and txt.exists() and not batch.overwrite: print(f"[{idx}/{len(images)}] SKIP existing TXT: {txt}"); result.skipped+=1; continue
if batch.skip_existing_jsonl_images and (src in seen or path.name in seen): print(f"[{idx}/{len(images)}] SKIP existing JSONL image: {src}"); result.skipped+=1; continue
print(f"[{idx}/{len(images)}] Captioning: {path}"); rec=self.caption_path(path); result.records.append(rec)
if batch.write_txt and not write_text_sidecar(txt, rec.caption, batch.overwrite, batch.backup_existing, batch.dry_run): result.skipped+=1; print(f"[{idx}/{len(images)}] SKIP existing TXT after caption: {txt}")
if jsonl_path: append_jsonl_records(jsonl_path,[rec],batch.dry_run)
if also_jsonl_path: append_jsonl_records(also_jsonl_path,[rec],batch.dry_run)
except KeyboardInterrupt: raise
except Exception as exc:
result.failed+=1; print(f"[JLC BLIP Engine] ERROR on {path}: {exc}")
err=CaptionRecord(str(path),"","",self.config.model_name,str(self.local_model_path or self.config.model_path),self.config.prompt,self.generation.seed,self.generation.temperature,self.generation.top_p,self.generation.top_k,self.generation.num_beams,self.generation.max_new_tokens,self.config.max_size,iso_timestamp(),status="error",error=str(exc),captionforge_pass="A",model_family="blip",ensemble_run_index=0,image_key=str(path.resolve()))
result.records.append(err)
if jsonl_path: append_jsonl_records(jsonl_path,[err],batch.dry_run)
if also_jsonl_path: append_jsonl_records(also_jsonl_path,[err],batch.dry_run)
finally:
if self._cache_policy()=="unload_after_run": self.unload()
return result
def build_run_config(self,batch=None):
return {"timestamp":iso_timestamp(),"engine":"JLC BLIP Caption Engine","blip_config":asdict(self.config),"generation":asdict(self.generation),"cleanup":{**asdict(self.cleanup),"replacement_rules":[list(r) for r in self.cleanup.replacement_rules]},"batch":json_safe(asdict(batch)) if batch is not None else None}
def caption_one_image(image_path, blip_config, generation=None, cleanup=None):
engine=BlipCaptionEngine(blip_config,generation,cleanup); engine.load()
try: return engine.caption_path(image_path)
finally:
p=getattr(blip_config,"cache_policy",""); p=p.strip() if isinstance(p,str) else ""
if not p: p="evict_other_caption_models" if getattr(blip_config,"keep_loaded",True) else "unload_after_run"
if p=="unload_after_run": engine.unload()
File diff suppressed because it is too large Load Diff
@@ -1,812 +0,0 @@
#!/usr/bin/env python
"""
JLC InternVL Caption Engine
"""
from __future__ import annotations
MANIFEST = {
"name": "JLC InternVL Caption Engine",
"version": (0, 1, 0),
"author": "J. L. Córdova",
"description": "InternVL CaptionForge Pass A engine.",
}
import fnmatch
import json
import random
import re
import shutil
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Iterable, Optional
from PIL import Image
import torch
from torchvision import transforms
from .captionforge_model_cache import (
make_cache_key,
get_cached_model,
register_model,
prepare_for_model_load,
unload_after_run,
)
SUPPORTED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tif", ".tiff"}
PROMPT_PRESETS: dict[str, str] = {
"caption": "Describe this image.",
"detailed": "Describe this image in detail.",
"lora_literal": "Write a concise but informative visual caption for dataset training. Mention only visible details.",
"style_focus": "Describe the subject, appearance, pose, clothing, visual style, and background of this image.",
}
DEFAULT_PROMPT_PRESET = "detailed"
MEMORY_MODES = {"Default": {}, "Balanced (8-bit)": {"load_in_8bit": True}}
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)
@dataclass
class GenerationConfig:
max_new_tokens: int = 256
num_beams: int = 3
temperature: float = 0.0
top_p: float = 0.90
top_k: int = 50
repetition_penalty: float = 1.0
seed: Optional[int] = None
@dataclass
class CleanupConfig:
trigger: str = ""
prefix: str = ""
suffix: str = ""
forbidden_phrases: list[str] = field(default_factory=list)
replacement_rules: list[tuple[str, str]] = field(default_factory=list)
replace_case_insensitive: bool = True
replace_whole_words_only: bool = False
strip_boilerplate_prefixes: bool = True
strip_trailing_period: bool = True
@dataclass
class BatchCaptionConfig:
input_path: str = ""
recursive: bool = True
filename_glob: str = "*"
extensions: set[str] = field(default_factory=lambda: set(SUPPORTED_EXTENSIONS))
output_dir: str = ""
write_txt: bool = True
write_jsonl: bool = False
jsonl_filename: str = "captions.jsonl"
also_jsonl_path: str = ""
write_run_config: bool = True
run_config_filename: str = ""
overwrite: bool = False
backup_existing: bool = True
dry_run: bool = False
limit: int = 0
skip_existing_txt: bool = True
skip_existing_jsonl_images: bool = False
@dataclass
class CaptionRecord:
image: str
caption: str
raw_caption: str
model_name: str
model_path: str
prompt: str
seed: Optional[int]
temperature: float
top_p: float
top_k: int
num_beams: int
max_new_tokens: int
max_size: int
timestamp: str
status: str = "ok"
error: str = ""
captionforge_pass: str = "A"
model_family: str = "internvl"
ensemble_run_index: int = 0
image_key: str = ""
@dataclass
class BatchCaptionResult:
records: list[CaptionRecord] = field(default_factory=list)
skipped: int = 0
failed: int = 0
def timestamp() -> str:
return datetime.now().strftime("%Y%m%d_%H%M%S")
def iso_timestamp() -> str:
return datetime.now().isoformat(timespec="seconds")
def safe_mkdir(path: Path) -> None:
path.mkdir(parents=True, exist_ok=True)
def set_seed(seed: int) -> None:
random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
def resolve_prompt(prompt_preset: str = DEFAULT_PROMPT_PRESET, custom_prompt: str = "") -> str:
custom = (custom_prompt or "").strip()
return custom or PROMPT_PRESETS.get(prompt_preset, PROMPT_PRESETS[DEFAULT_PROMPT_PRESET])
def normalize_caption(text: str, strip_boilerplate_prefixes: bool = True, strip_trailing_period: bool = True) -> str:
value = str(text or "").strip().replace("\n", " ")
prefixes = ["The image shows ", "The image depicts ", "This image shows ", "This image depicts ", "A photo of ", "An image of ", "Assistant: "]
if strip_boilerplate_prefixes:
for prefix in prefixes:
if value.lower().startswith(prefix.lower()):
value = value[len(prefix):].strip()
break
value = re.sub(r"\s+", " ", value).strip()
if strip_trailing_period and value.endswith("."):
value = value[:-1].strip()
return value
def apply_replacements(text: str, rules: list[tuple[str, str]], case_insensitive: bool = True, whole_words_only: bool = False) -> str:
flags = re.IGNORECASE if case_insensitive else 0
out = text
for old, new in rules or []:
old = old.strip()
if not old:
continue
patt = re.escape(old)
if whole_words_only:
patt = r"\b" + patt + r"\b"
out = re.sub(patt, new.strip(), out, flags=flags)
return out
def cleanup_caption(caption: str, config: CleanupConfig) -> str:
text = normalize_caption(caption, config.strip_boilerplate_prefixes, config.strip_trailing_period)
text = apply_replacements(text, config.replacement_rules, config.replace_case_insensitive, config.replace_whole_words_only)
for phrase in config.forbidden_phrases or []:
phrase = phrase.strip()
if phrase:
text = re.sub(re.escape(phrase), "", text, flags=re.IGNORECASE)
text = normalize_caption(text, config.strip_boilerplate_prefixes, config.strip_trailing_period)
parts = []
for x in (config.trigger, config.prefix):
x = (x or "").strip().strip(" ,")
if x:
parts.append(x)
if text:
parts.append(text)
final = ", ".join(parts).strip(" ,")
if config.suffix:
final += config.suffix
return final.strip()
def iter_image_files(input_path: str | Path, recursive: bool = True, filename_glob: str = "*", extensions: Optional[set[str]] = None) -> Iterable[Path]:
p = Path(input_path)
exts = extensions or set(SUPPORTED_EXTENSIONS)
glob_text = (filename_glob or "*").strip() or "*"
if p.is_file():
if p.suffix.lower() in exts and fnmatch.fnmatch(p.name, glob_text):
yield p
return
if not p.exists():
raise FileNotFoundError(f"input_path does not exist: {p}")
it = p.rglob("*") if recursive else p.glob("*")
for child in sorted(it):
if child.is_file() and child.suffix.lower() in exts and fnmatch.fnmatch(child.name, glob_text):
yield child
def load_image_file(path: str | Path) -> Image.Image:
return Image.open(path).convert("RGB")
def resize_for_model(image: Image.Image, max_size: int) -> Image.Image:
if max_size <= 0:
return image
w, h = image.size
longest = max(w, h)
if longest <= max_size:
return image
scale = max_size / float(longest)
return image.resize((max(1, round(w * scale)), max(1, round(h * scale))), Image.Resampling.LANCZOS)
def sidecar_txt_path(image_path: Path, output_dir: str | Path | None = None) -> Path:
if output_dir:
out = Path(output_dir)
safe_mkdir(out)
return out / f"{image_path.stem}.txt"
return image_path.with_suffix(".txt")
def backup_existing_file(path: Path, dry_run: bool = False) -> Optional[Path]:
if not path.exists():
return None
backup = path.with_name(f"{path.name}.bak_{timestamp()}")
if not dry_run:
shutil.copy2(path, backup)
return backup
def write_text_sidecar(path: Path, text: str, overwrite: bool = False, backup_existing: bool = True, dry_run: bool = False) -> bool:
if path.exists() and not overwrite:
return False
if dry_run:
return True
safe_mkdir(path.parent)
if path.exists() and overwrite and backup_existing:
backup_existing_file(path, dry_run=False)
path.write_text(text.rstrip() + "\n", encoding="utf-8")
return True
def load_existing_jsonl_images(path: Path) -> set[str]:
seen = set()
if not path.exists():
return seen
with path.open("r", encoding="utf-8") as f:
for line in f:
try:
obj = json.loads(line)
except Exception:
continue
image_value = obj.get("image") or obj.get("image_path") or obj.get("source")
if image_value:
seen.add(str(image_value))
return seen
def record_to_json(record: CaptionRecord) -> dict[str, Any]:
return asdict(record)
def append_jsonl_records(path: Path, records: list[CaptionRecord], dry_run: bool = False) -> None:
if dry_run or not records:
return
safe_mkdir(path.parent)
with path.open("a", encoding="utf-8") as f:
for rec in records:
f.write(json.dumps(record_to_json(rec), ensure_ascii=False) + "\n")
def json_safe(value):
if isinstance(value, set):
return sorted(value)
if isinstance(value, Path):
return str(value)
if isinstance(value, dict):
return {k: json_safe(v) for k, v in value.items()}
if isinstance(value, list):
return [json_safe(v) for v in value]
if isinstance(value, tuple):
return [json_safe(v) for v in value]
return value
def write_run_config_json(path: Path, config: dict[str, Any], dry_run: bool = False) -> None:
if dry_run:
return
safe_mkdir(path.parent)
path.write_text(json.dumps(json_safe(config), indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
def model_folder_has_weights(local_path: Path) -> bool:
if not local_path.exists() or not local_path.is_dir():
return False
patterns = ["*.safetensors", "*.bin", "*.pt", "*.pth"]
return any(any(local_path.rglob(pat)) for pat in patterns)
def _try_get_comfy_model_management():
try:
import comfy.model_management as mm
return mm
except Exception:
return None
def build_transform(input_size: int):
return transforms.Compose([
transforms.Lambda(lambda img: img.convert("RGB")),
transforms.Resize((input_size, input_size), interpolation=transforms.InterpolationMode.BICUBIC),
transforms.ToTensor(),
transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
])
def internvl_load_image(image: Image.Image, input_size: int = 448) -> torch.Tensor:
return build_transform(input_size)(image).unsqueeze(0)
@dataclass(frozen=True)
class InternVLModelInfo:
repo_id: str
local_folder: str
notes: str = ""
MODEL_REGISTRY = {
"InternVL2.5-2B": InternVLModelInfo(
repo_id="OpenGVLab/InternVL2_5-2B",
local_folder="InternVL2_5-2B",
notes="Default InternVL family entry.",
),
}
@dataclass
class InternVLCaptionConfig:
model_name: str = "InternVL2.5-2B"
model_path: str = ""
model_root: str = "models/LLM/JLC_InternVLCaption"
memory_mode: str = "Default"
dtype: str = "bf16"
device: str = "auto"
trust_remote_code: bool = True
keep_loaded: bool = True
cache_policy: str = ""
quiet_transformers_load: bool = True
max_size: int = 768
prompt: str = PROMPT_PRESETS[DEFAULT_PROMPT_PRESET]
allow_download: bool = True
use_comfy_model_management: bool = True
input_size: int = 448
def get_model_info(model_name: str) -> InternVLModelInfo:
if model_name not in MODEL_REGISTRY:
raise KeyError(f"Unknown InternVL model_name: {model_name}")
return MODEL_REGISTRY[model_name]
def get_registry_model_path(model_name: str, model_root: str | Path) -> Path:
return Path(model_root) / get_model_info(model_name).local_folder
def download_registry_model_if_needed(model_name: str, model_root: str | Path, metadata_only: bool = False, allow_download: bool = True) -> Path:
local_path = get_registry_model_path(model_name, model_root)
if not metadata_only and model_folder_has_weights(local_path):
return local_path
if not allow_download:
if metadata_only:
safe_mkdir(local_path)
return local_path
raise FileNotFoundError(f"Model folder does not contain weights and allow_download=False: {local_path}")
from huggingface_hub import snapshot_download
info = get_model_info(model_name)
safe_mkdir(local_path)
if metadata_only:
snapshot_download(repo_id=info.repo_id, local_dir=str(local_path), ignore_patterns=["*.safetensors", "*.bin", "*.pt", "*.pth", "*.gguf", "*.onnx", "*.ckpt"])
else:
snapshot_download(repo_id=info.repo_id, local_dir=str(local_path))
return local_path
def probe_registry_model_download(model_name: str, model_root: str | Path) -> str:
local_path = download_registry_model_if_needed(model_name, model_root, metadata_only=True, allow_download=True)
return f"JLC InternVL Caption download probe completed.\n\nModel: {model_name}\nFolder: {local_path}\nLarge model weight files were intentionally skipped."
def _unload_bundle(bundle: dict[str, Any]) -> None:
model = bundle.get("model") if isinstance(bundle, dict) else None
if model is not None:
try:
model.to("cpu")
except Exception:
pass
if isinstance(bundle, dict):
bundle.clear()
def patch_internvl_loaded_model_compat(model) -> int:
"""Patch loaded InternVL/InternLM instances for newer Transformers generation."""
patched = 0
try:
from transformers import GenerationConfig
except Exception:
return 0
targets = [model]
language_model = getattr(model, "language_model", None)
if language_model is not None:
targets.append(language_model)
for target in targets:
if target is None:
continue
if not hasattr(target, "generation_config"):
try:
config = getattr(target, "config", None)
target.generation_config = GenerationConfig.from_model_config(config)
patched += 1
except Exception:
try:
target.generation_config = GenerationConfig()
patched += 1
except Exception:
pass
for attr_name in ("config", "generation_config"):
obj = getattr(target, attr_name, None)
if obj is not None and hasattr(obj, "use_cache"):
try:
obj.use_cache = False
patched += 1
except Exception:
pass
if patched:
print(
f"[JLC InternVL Engine] Patched {patched} loaded model "
f"generation_config attribute(s)."
)
return patched
def patch_internvl_transformers_compat() -> int:
"""Patch older InternVL remote-code classes for newer Transformers loaders."""
import inspect
import sys
try:
from transformers.generation import GenerationMixin
except Exception:
from transformers.generation.utils import GenerationMixin
patched = 0
for module_name, module in list(sys.modules.items()):
module_dict = getattr(module, "__dict__", None)
if not isinstance(module_dict, dict):
continue
# Keep this narrow: only touch InternVL / InternLM remote-code modules.
lowered_name = str(module_name).lower()
if (
"internvl" not in lowered_name
and "internlm" not in lowered_name
and "transformers_modules" not in lowered_name
):
continue
for obj in list(module_dict.values()):
if not inspect.isclass(obj):
continue
obj_name = getattr(obj, "__name__", "")
obj_module = getattr(obj, "__module__", "")
combined = f"{obj_module}.{obj_name}".lower()
if "internvl" not in combined and "internlm" not in combined:
continue
# Newer Transformers expects this dict-like attribute during
# _move_missing_keys_from_meta_to_device().
if hasattr(obj, "_tied_weights_keys") and not hasattr(obj, "all_tied_weights_keys"):
try:
setattr(obj, "all_tied_weights_keys", {})
patched += 1
except Exception:
pass
# Transformers >=4.50 no longer gives PreTrainedModel subclasses
# GenerationMixin automatically. InternLM2 remote code defines
# prepare_inputs_for_generation, so binding GenerationMixin.generate
# is the compatibility bridge needed by InternVLChatModel.chat().
if hasattr(obj, "prepare_inputs_for_generation"):
for name, value in GenerationMixin.__dict__.items():
if name.startswith("__"):
continue
# Do not overwrite methods/attrs the remote model already defines.
if hasattr(obj, name):
continue
try:
setattr(obj, name, value)
patched += 1
except Exception:
pass
if patched:
print(
f"[JLC InternVL Engine] Patched {patched} InternVL/InternLM "
f"remote-code compatibility attribute(s)."
)
return patched
def repair_internvl_dynamic_module_cache(local_path: Path) -> int:
"""Copy InternVL remote-code .py files into HF dynamic-module cache dirs.
Some Transformers/HF combinations create the InternVL dynamic module cache with
modeling_internvl_chat.py present but relative import files missing.
"""
try:
import os
import shutil
except Exception:
return 0
source_files = [p for p in Path(local_path).glob("*.py") if p.is_file()]
if not source_files:
return 0
hf_home = Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface"))
modules_root = hf_home / "modules" / "transformers_modules"
if not modules_root.exists():
return 0
candidate_dirs: list[Path] = []
for root in modules_root.glob("*InternVL*"):
if root.is_dir():
candidate_dirs.append(root)
candidate_dirs.extend([p for p in root.rglob("*") if p.is_dir()])
copied = 0
for target_dir in candidate_dirs:
if not (
(target_dir / "modeling_internvl_chat.py").exists()
or (target_dir / "configuration_internvl_chat.py").exists()
or "InternVL" in str(target_dir)
):
continue
for src_file in source_files:
dst = target_dir / src_file.name
if not dst.exists():
try:
shutil.copy2(src_file, dst)
copied += 1
except Exception:
pass
if copied:
print(
f"[JLC InternVL Engine] Repaired HF dynamic-module cache "
f"with {copied} local InternVL Python file(s)."
)
return copied
class InternVLCaptionEngine:
def __init__(self, config: InternVLCaptionConfig, generation: Optional[GenerationConfig] = None, cleanup: Optional[CleanupConfig] = None) -> None:
self.config = config
self.generation = generation or GenerationConfig()
self.cleanup = cleanup or CleanupConfig()
self.tokenizer = None
self.model = None
self.local_model_path: Optional[Path] = None
self.model_size_bytes: Optional[int] = None
self._comfy_mm = _try_get_comfy_model_management() if config.use_comfy_model_management else None
self.inference_device = self._resolve_inference_device(config.device)
def _resolve_inference_device(self, device: str) -> torch.device:
if self._comfy_mm is not None:
try:
dev = self._comfy_mm.get_torch_device()
if dev.type == "cuda":
return dev
except Exception:
pass
if device == "auto":
if not torch.cuda.is_available():
raise RuntimeError("CaptionForge InternVL requires CUDA for release builds.")
return torch.device("cuda")
dev = torch.device(device)
if dev.type != "cuda":
raise RuntimeError(f"CaptionForge InternVL release mode does not allow inference device={device!r}")
return dev
def _resolve_dtype(self, dtype: str):
text = dtype.lower().strip()
if text in {"bf16", "bfloat16"}:
return torch.bfloat16
if text in {"fp16", "float16"}:
return torch.float16
if text in {"fp32", "float32"}:
return torch.float32
return torch.bfloat16
def _cache_policy(self) -> str:
policy = (self.config.cache_policy or "").strip()
return policy or ("evict_other_caption_models" if self.config.keep_loaded else "unload_after_run")
def _cache_key(self, local_path: Path) -> str:
return make_cache_key(role="caption", family="internvl", model_path=str(local_path.resolve()), device=str(self.inference_device), quantization=self.config.memory_mode, dtype=self.config.dtype)
def resolve_model_path(self) -> Path:
if self.config.model_path.strip():
return Path(self.config.model_path).expanduser()
return download_registry_model_if_needed(self.config.model_name, self.config.model_root, metadata_only=False, allow_download=self.config.allow_download)
def _module_size(self, model) -> int:
total = 0
try:
for p in model.parameters():
total += p.numel() * p.element_size()
for b in model.buffers():
total += b.numel() * b.element_size()
except Exception:
pass
return total
patch_internvl_transformers_compat()
def load(self) -> None:
local_path = self.resolve_model_path()
self.local_model_path = local_path
cache_key = self._cache_key(local_path)
cached = get_cached_model(cache_key)
if cached is not None:
self.tokenizer = cached.get("tokenizer")
self.model = cached.get("model")
self.model_size_bytes = cached.get("model_size_bytes")
patch_internvl_transformers_compat()
patch_internvl_loaded_model_compat(self.model)
return
prepare_for_model_load(cache_key, policy=self._cache_policy(), role="caption")
if self.generation.seed is not None:
set_seed(self.generation.seed)
from transformers import AutoModel, AutoTokenizer
self.tokenizer = AutoTokenizer.from_pretrained(str(local_path), trust_remote_code=self.config.trust_remote_code, use_fast=False)
kwargs: dict[str, Any] = {"torch_dtype": self._resolve_dtype(self.config.dtype), "trust_remote_code": self.config.trust_remote_code, "low_cpu_mem_usage": True}
if self.config.memory_mode == "Balanced (8-bit)":
from transformers import BitsAndBytesConfig
kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
kwargs["device_map"] = {"": self.inference_device.index or 0}
try:
kwargs["use_flash_attn"] = False
except Exception:
pass
last_exc = None
for attempt in range(3):
try:
self.model = AutoModel.from_pretrained(str(local_path), **kwargs)
break
except FileNotFoundError as exc:
missing = str(exc)
last_exc = exc
if "transformers_modules" not in missing or "InternVL" not in missing:
raise
print(
"[JLC InternVL Engine] HF dynamic-module cache missing "
f"a relative import: {missing}"
)
repair_internvl_dynamic_module_cache(local_path)
patch_internvl_transformers_compat()
continue
except AttributeError as exc:
message = str(exc)
last_exc = exc
if "all_tied_weights_keys" not in message:
raise
print(
"[JLC InternVL Engine] Patching InternVL remote-code "
"compatibility for all_tied_weights_keys."
)
patched = patch_internvl_transformers_compat()
if not patched:
raise
continue
else:
raise last_exc
patch_internvl_transformers_compat()
patch_internvl_loaded_model_compat(self.model)
self.model.eval()
if self.config.memory_mode == "Default":
self.model.to(self.inference_device)
self.model_size_bytes = self._module_size(self.model)
register_model(cache_key, {"tokenizer": self.tokenizer, "model": self.model, "model_size_bytes": self.model_size_bytes}, family="internvl", model_path=str(local_path), device=str(self.inference_device), quantization=self.config.memory_mode, role="caption", unload_fn=_unload_bundle)
def unload(self) -> None:
if self.local_model_path is not None:
unload_after_run(self._cache_key(self.local_model_path), enabled=True)
self.tokenizer = None
self.model = None
if torch.cuda.is_available():
torch.cuda.empty_cache()
@torch.inference_mode()
def caption_pil(self, image: Image.Image) -> tuple[str, str]:
if self.tokenizer is None or self.model is None:
self.load()
if self.generation.seed is not None:
set_seed(self.generation.seed)
image = resize_for_model(image.convert("RGB"), self.config.max_size)
pixel_values = internvl_load_image(image, self.config.input_size)
model_dtype = getattr(self.model, "dtype", self._resolve_dtype(self.config.dtype))
pixel_values = pixel_values.to(device=self.inference_device, dtype=model_dtype)
prompt = (self.config.prompt or "").strip() or "Describe this image in detail."
question = f"<image>\n{prompt}"
gen_config = {
"max_new_tokens": int(self.generation.max_new_tokens),
"num_beams": max(1, int(self.generation.num_beams)),
"do_sample": bool(self.generation.temperature > 0),
}
if self.generation.temperature > 0:
gen_config.update({"temperature": float(self.generation.temperature), "top_p": float(self.generation.top_p), "top_k": int(self.generation.top_k)})
raw = str(self.model.chat(self.tokenizer, pixel_values, question, generation_config=gen_config)).strip()
final = cleanup_caption(raw, self.cleanup)
if self._cache_policy() == "unload_after_run":
self.unload()
return final, raw
def caption_path(self, image_path: str | Path) -> CaptionRecord:
p = Path(image_path)
final, raw = self.caption_pil(load_image_file(p))
return CaptionRecord(image=str(p), caption=final, raw_caption=raw, model_name=self.config.model_name, model_path=str(self.local_model_path or self.config.model_path), prompt=self.config.prompt, seed=self.generation.seed, temperature=self.generation.temperature, top_p=self.generation.top_p, top_k=self.generation.top_k, num_beams=self.generation.num_beams, max_new_tokens=self.generation.max_new_tokens, max_size=self.config.max_size, timestamp=iso_timestamp(), image_key=str(p.resolve()))
def caption_batch(self, batch: BatchCaptionConfig) -> BatchCaptionResult:
if not batch.input_path.strip():
raise ValueError("BatchCaptionConfig.input_path is required.")
result = BatchCaptionResult()
images = list(iter_image_files(batch.input_path, batch.recursive, batch.filename_glob, batch.extensions))
if batch.limit > 0:
images = images[: int(batch.limit)]
if not images:
return result
output_dir = Path(batch.output_dir) if batch.output_dir.strip() else None
jsonl_path = (output_dir or images[0].parent) / (batch.jsonl_filename.strip() or "captions.jsonl") if batch.write_jsonl else None
also_jsonl_path = Path(batch.also_jsonl_path) if batch.also_jsonl_path.strip() else None
seen_jsonl = set()
if batch.skip_existing_jsonl_images:
if jsonl_path is not None:
seen_jsonl |= load_existing_jsonl_images(jsonl_path)
if also_jsonl_path is not None:
seen_jsonl |= load_existing_jsonl_images(also_jsonl_path)
if batch.write_run_config:
outdir = output_dir or images[0].parent
name = batch.run_config_filename.strip() or f"jlc_internvl_caption_run_config_{timestamp()}.json"
write_run_config_json(outdir / name, self.build_run_config(batch), dry_run=batch.dry_run)
for image_path in images:
try:
txt_path = sidecar_txt_path(image_path, output_dir)
if batch.skip_existing_txt and batch.write_txt and txt_path.exists() and not batch.overwrite:
result.skipped += 1
continue
if batch.skip_existing_jsonl_images and str(image_path) in seen_jsonl:
result.skipped += 1
continue
rec = self.caption_path(image_path)
result.records.append(rec)
if batch.write_txt:
write_text_sidecar(txt_path, rec.caption, batch.overwrite, batch.backup_existing, batch.dry_run)
if jsonl_path is not None:
append_jsonl_records(jsonl_path, [rec], batch.dry_run)
if also_jsonl_path is not None:
append_jsonl_records(also_jsonl_path, [rec], batch.dry_run)
except Exception as exc:
result.failed += 1
err = CaptionRecord(image=str(image_path), caption="", raw_caption="", model_name=self.config.model_name, model_path=str(self.local_model_path or self.config.model_path), prompt=self.config.prompt, seed=self.generation.seed, temperature=self.generation.temperature, top_p=self.generation.top_p, top_k=self.generation.top_k, num_beams=self.generation.num_beams, max_new_tokens=self.generation.max_new_tokens, max_size=self.config.max_size, timestamp=iso_timestamp(), status="error", error=str(exc), image_key=str(Path(image_path).resolve()))
result.records.append(err)
if jsonl_path is not None:
append_jsonl_records(jsonl_path, [err], batch.dry_run)
if also_jsonl_path is not None:
append_jsonl_records(also_jsonl_path, [err], batch.dry_run)
return result
def build_run_config(self, batch: Optional[BatchCaptionConfig] = None) -> dict[str, Any]:
cleanup = asdict(self.cleanup)
cleanup["replacement_rules"] = [list(x) for x in self.cleanup.replacement_rules]
return {"timestamp": iso_timestamp(), "engine": "JLC InternVL Caption Engine", "internvl_config": asdict(self.config), "generation": asdict(self.generation), "cleanup": cleanup, "batch": json_safe(asdict(batch)) if batch is not None else None}
@@ -1,729 +0,0 @@
#!/usr/bin/env python
"""
JLC SmolVLM Caption Engine
"""
from __future__ import annotations
MANIFEST = {
"name": "JLC SmolVLM Caption Engine",
"version": (0, 1, 0),
"author": "J. L. Córdova",
"description": "SmolVLM CaptionForge Pass A engine.",
}
import fnmatch
import json
import random
import re
import shutil
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Iterable, Optional
from PIL import Image
import torch
from .captionforge_model_cache import (
make_cache_key,
get_cached_model,
register_model,
prepare_for_model_load,
unload_after_run,
)
try:
from .captionforge_caption_prompt_kit import (
CAPTION_TYPE_CHOICES,
CAPTION_LENGTH_CHOICES,
EXTRA_OPTIONS as PROMPT_KIT_EXTRA_OPTIONS,
build_caption_prompt_spec,
)
except Exception: # pragma: no cover - direct-file fallback
CAPTION_TYPE_CHOICES = ["Descriptive", "LoRA Literal", "SFW Character Caption"]
CAPTION_LENGTH_CHOICES = ["any", "very short", "short", "medium-length", "long", "very long"]
PROMPT_KIT_EXTRA_OPTIONS = []
def build_caption_prompt_spec(caption_type="Descriptive", caption_length="medium-length", extra_options=None, name_input="", dialect="smolvlm"):
prompt = "Describe this image clearly and factually. Keep it SFW and mention only visible details."
return prompt, {
"caption_type": caption_type,
"caption_length": str(caption_length),
"extra_options": tuple(extra_options or ()),
"name_input": name_input,
"dialect": dialect,
"source": "fallback_smol_prompt",
}
SUPPORTED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tif", ".tiff"}
PROMPT_PRESETS: dict[str, str] = {
"caption": "Describe this image clearly and factually. Keep it SFW.",
"detailed": "Describe this image for dataset captioning. Mention visible subject, pose, expression, clothing, background, lighting, and style. Keep it SFW.",
"lora_literal": "Write a concise SFW LoRA training caption for this image. Mention only visible details.",
"sfw_character": "Write a SFW character-focused caption. Describe visible character design, pose, expression, hair, outfit, scene, lighting, and style.",
"taggy": "Use concise comma-separated visual phrases only. Include visible subject, pose, clothing, style, lighting, and background.",
"style_focus": "Describe the subject, appearance, pose, clothing, visual style, lighting, and background of this image.",
}
DEFAULT_PROMPT_PRESET = "detailed"
MEMORY_MODES = {"Default": {}, "Balanced (8-bit)": {"load_in_8bit": True}}
@dataclass
class GenerationConfig:
max_new_tokens: int = 256
num_beams: int = 3
temperature: float = 0.0
top_p: float = 0.90
top_k: int = 50
repetition_penalty: float = 1.0
seed: Optional[int] = None
@dataclass
class CleanupConfig:
trigger: str = ""
prefix: str = ""
suffix: str = ""
forbidden_phrases: list[str] = field(default_factory=list)
replacement_rules: list[tuple[str, str]] = field(default_factory=list)
replace_case_insensitive: bool = True
replace_whole_words_only: bool = False
strip_boilerplate_prefixes: bool = True
strip_trailing_period: bool = True
@dataclass
class BatchCaptionConfig:
input_path: str = ""
recursive: bool = True
filename_glob: str = "*"
extensions: set[str] = field(default_factory=lambda: set(SUPPORTED_EXTENSIONS))
output_dir: str = ""
write_txt: bool = True
write_jsonl: bool = False
jsonl_filename: str = "captions.jsonl"
also_jsonl_path: str = ""
write_run_config: bool = True
run_config_filename: str = ""
overwrite: bool = False
backup_existing: bool = True
dry_run: bool = False
limit: int = 0
skip_existing_txt: bool = True
skip_existing_jsonl_images: bool = False
@dataclass
class CaptionRecord:
image: str
caption: str
raw_caption: str
model_name: str
model_path: str
prompt: str
seed: Optional[int]
temperature: float
top_p: float
top_k: int
num_beams: int
max_new_tokens: int
max_size: int
timestamp: str
status: str = "ok"
error: str = ""
captionforge_pass: str = "A"
model_family: str = "smolvlm"
ensemble_run_index: int = 0
image_key: str = ""
prompt_metadata: dict[str, Any] = field(default_factory=dict)
@dataclass
class BatchCaptionResult:
records: list[CaptionRecord] = field(default_factory=list)
skipped: int = 0
failed: int = 0
def timestamp() -> str:
return datetime.now().strftime("%Y%m%d_%H%M%S")
def iso_timestamp() -> str:
return datetime.now().isoformat(timespec="seconds")
def safe_mkdir(path: Path) -> None:
path.mkdir(parents=True, exist_ok=True)
def set_seed(seed: int) -> None:
random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
def resolve_prompt(prompt_preset: str = DEFAULT_PROMPT_PRESET, custom_prompt: str = "") -> str:
custom = (custom_prompt or "").strip()
return custom or PROMPT_PRESETS.get(prompt_preset, PROMPT_PRESETS[DEFAULT_PROMPT_PRESET])
def normalize_caption(text: str, strip_boilerplate_prefixes: bool = True, strip_trailing_period: bool = True) -> str:
value = str(text or "").strip().replace("\n", " ")
prefixes = [
"The image shows ", "The image depicts ", "This image shows ", "This image depicts ",
"A photo of ", "An image of ", "Assistant: ",
]
if strip_boilerplate_prefixes:
for prefix in prefixes:
if value.lower().startswith(prefix.lower()):
value = value[len(prefix):].strip()
break
value = re.sub(r"\s+", " ", value).strip()
if strip_trailing_period and value.endswith("."):
value = value[:-1].strip()
return value
def apply_replacements(text: str, rules: list[tuple[str, str]], case_insensitive: bool = True, whole_words_only: bool = False) -> str:
flags = re.IGNORECASE if case_insensitive else 0
out = text
for old, new in rules or []:
old = old.strip()
if not old:
continue
patt = re.escape(old)
if whole_words_only:
patt = r"\b" + patt + r"\b"
out = re.sub(patt, new.strip(), out, flags=flags)
return out
def cleanup_caption(caption: str, config: CleanupConfig) -> str:
text = normalize_caption(caption, config.strip_boilerplate_prefixes, config.strip_trailing_period)
text = apply_replacements(text, config.replacement_rules, config.replace_case_insensitive, config.replace_whole_words_only)
for phrase in config.forbidden_phrases or []:
phrase = phrase.strip()
if phrase:
text = re.sub(re.escape(phrase), "", text, flags=re.IGNORECASE)
text = normalize_caption(text, config.strip_boilerplate_prefixes, config.strip_trailing_period)
parts = []
for x in (config.trigger, config.prefix):
x = (x or "").strip().strip(" ,")
if x:
parts.append(x)
if text:
parts.append(text)
final = ", ".join(parts).strip(" ,")
if config.suffix:
final += config.suffix
return final.strip()
def iter_image_files(input_path: str | Path, recursive: bool = True, filename_glob: str = "*", extensions: Optional[set[str]] = None) -> Iterable[Path]:
p = Path(input_path)
exts = extensions or set(SUPPORTED_EXTENSIONS)
glob_text = (filename_glob or "*").strip() or "*"
if p.is_file():
if p.suffix.lower() in exts and fnmatch.fnmatch(p.name, glob_text):
yield p
return
if not p.exists():
raise FileNotFoundError(f"input_path does not exist: {p}")
it = p.rglob("*") if recursive else p.glob("*")
for child in sorted(it):
if child.is_file() and child.suffix.lower() in exts and fnmatch.fnmatch(child.name, glob_text):
yield child
def load_image_file(path: str | Path) -> Image.Image:
return Image.open(path).convert("RGB")
def resize_for_model(image: Image.Image, max_size: int) -> Image.Image:
if max_size <= 0:
return image
w, h = image.size
longest = max(w, h)
if longest <= max_size:
return image
scale = max_size / float(longest)
return image.resize((max(1, round(w * scale)), max(1, round(h * scale))), Image.Resampling.LANCZOS)
def sidecar_txt_path(image_path: Path, output_dir: str | Path | None = None) -> Path:
if output_dir:
out = Path(output_dir)
safe_mkdir(out)
return out / f"{image_path.stem}.txt"
return image_path.with_suffix(".txt")
def backup_existing_file(path: Path, dry_run: bool = False) -> Optional[Path]:
if not path.exists():
return None
backup = path.with_name(f"{path.name}.bak_{timestamp()}")
if not dry_run:
shutil.copy2(path, backup)
return backup
def write_text_sidecar(path: Path, text: str, overwrite: bool = False, backup_existing: bool = True, dry_run: bool = False) -> bool:
if path.exists() and not overwrite:
return False
if dry_run:
return True
safe_mkdir(path.parent)
if path.exists() and overwrite and backup_existing:
backup_existing_file(path, dry_run=False)
path.write_text(text.rstrip() + "\n", encoding="utf-8")
return True
def load_existing_jsonl_images(path: Path) -> set[str]:
seen = set()
if not path.exists():
return seen
with path.open("r", encoding="utf-8") as f:
for line in f:
try:
obj = json.loads(line)
except Exception:
continue
image_value = obj.get("image") or obj.get("image_path") or obj.get("source")
if image_value:
seen.add(str(image_value))
return seen
def record_to_json(record: CaptionRecord) -> dict[str, Any]:
return asdict(record)
def append_jsonl_records(path: Path, records: list[CaptionRecord], dry_run: bool = False) -> None:
if dry_run or not records:
return
safe_mkdir(path.parent)
with path.open("a", encoding="utf-8") as f:
for rec in records:
f.write(json.dumps(record_to_json(rec), ensure_ascii=False) + "\n")
def json_safe(value):
if isinstance(value, set):
return sorted(value)
if isinstance(value, Path):
return str(value)
if isinstance(value, dict):
return {k: json_safe(v) for k, v in value.items()}
if isinstance(value, list):
return [json_safe(v) for v in value]
if isinstance(value, tuple):
return [json_safe(v) for v in value]
return value
def write_run_config_json(path: Path, config: dict[str, Any], dry_run: bool = False) -> None:
if dry_run:
return
safe_mkdir(path.parent)
path.write_text(json.dumps(json_safe(config), indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
def model_folder_has_weights(local_path: Path) -> bool:
if not local_path.exists() or not local_path.is_dir():
return False
patterns = ["*.safetensors", "*.bin", "*.pt", "*.pth"]
return any(any(local_path.rglob(pat)) for pat in patterns)
def _try_get_comfy_model_management():
try:
import comfy.model_management as mm
return mm
except Exception:
return None
@dataclass(frozen=True)
class SmolVLMModelInfo:
repo_id: str
local_folder: str
notes: str = ""
MODEL_REGISTRY = {
"SmolVLM-Instruct": SmolVLMModelInfo(
repo_id="HuggingFaceTB/SmolVLM-Instruct",
local_folder="SmolVLM-Instruct",
notes="Default SmolVLM family entry.",
),
}
@dataclass
class SmolVLMCaptionConfig:
model_name: str = "SmolVLM-Instruct"
model_path: str = ""
model_root: str = "models/LLM/JLC_SmolVLMCaption"
memory_mode: str = "Default"
dtype: str = "bf16"
device: str = "auto"
trust_remote_code: bool = True
keep_loaded: bool = True
cache_policy: str = ""
quiet_transformers_load: bool = True
max_size: int = 768
# If prompt is non-empty, it wins and preserves legacy/custom behavior.
# If prompt is empty, CaptionForge builds a generic Qwen/Smol prompt-kit prompt.
prompt: str = ""
caption_type: str = "SFW Character Caption"
caption_length: str = "medium-length"
extra_options: tuple[str, ...] = (
"Keep the caption SFW and suitable for dataset training.",
"Describe only visible details; do not infer identity, story, or intent.",
"Use one dense natural-language caption rather than a tag list.",
)
name_input: str = ""
allow_download: bool = True
use_comfy_model_management: bool = True
def get_model_info(model_name: str) -> SmolVLMModelInfo:
if model_name not in MODEL_REGISTRY:
raise KeyError(f"Unknown SmolVLM model_name: {model_name}")
return MODEL_REGISTRY[model_name]
def get_registry_model_path(model_name: str, model_root: str | Path) -> Path:
return Path(model_root) / get_model_info(model_name).local_folder
def download_registry_model_if_needed(model_name: str, model_root: str | Path, metadata_only: bool = False, allow_download: bool = True) -> Path:
local_path = get_registry_model_path(model_name, model_root)
if not metadata_only and model_folder_has_weights(local_path):
return local_path
if not allow_download:
if metadata_only:
safe_mkdir(local_path)
return local_path
raise FileNotFoundError(f"Model folder does not contain weights and allow_download=False: {local_path}")
from huggingface_hub import snapshot_download
info = get_model_info(model_name)
safe_mkdir(local_path)
if metadata_only:
snapshot_download(repo_id=info.repo_id, local_dir=str(local_path), ignore_patterns=["*.safetensors", "*.bin", "*.pt", "*.pth", "*.gguf", "*.onnx", "*.ckpt"])
else:
snapshot_download(repo_id=info.repo_id, local_dir=str(local_path))
return local_path
def probe_registry_model_download(model_name: str, model_root: str | Path) -> str:
local_path = download_registry_model_if_needed(model_name, model_root, metadata_only=True, allow_download=True)
return f"JLC SmolVLM Caption download probe completed.\n\nModel: {model_name}\nFolder: {local_path}\nLarge model weight files were intentionally skipped."
def _unload_bundle(bundle: dict[str, Any]) -> None:
model = bundle.get("model") if isinstance(bundle, dict) else None
if model is not None:
try:
if _model_uses_accelerate_hooks(model):
print("[JLC SmolVLM Engine] Accelerate-dispatched model detected; skipping model.to('cpu') during unload.")
else:
model.to("cpu")
except Exception as exc:
print(f"[JLC SmolVLM Engine] Non-fatal unload warning: {exc}")
if isinstance(bundle, dict):
bundle.clear()
try:
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
except Exception:
pass
def _format_bytes(num_bytes: int | None) -> str:
if not num_bytes:
return "unknown"
value = float(num_bytes)
for unit in ("B", "KB", "MB", "GB", "TB"):
if value < 1024.0 or unit == "TB":
return f"{value:.2f} {unit}"
value /= 1024.0
return f"{num_bytes} B"
def _folder_weight_size_bytes(path: Path) -> int:
total = 0
for pattern in ("*.safetensors", "*.bin", "*.pt", "*.pth", "*.gguf", "*.ckpt"):
try:
total += sum(p.stat().st_size for p in path.rglob(pattern) if p.is_file())
except Exception:
pass
return total
def _cuda_diagnostic_line() -> str:
if not torch.cuda.is_available():
return "CUDA unavailable"
try:
idx = torch.cuda.current_device()
props = torch.cuda.get_device_properties(idx)
allocated = torch.cuda.memory_allocated(idx)
reserved = torch.cuda.memory_reserved(idx)
total = getattr(props, "total_memory", 0)
return (
f"CUDA:{idx} {props.name} | "
f"allocated={_format_bytes(allocated)}, reserved={_format_bytes(reserved)}, total={_format_bytes(total)}"
)
except Exception as exc:
return f"CUDA diagnostics unavailable: {exc}"
def _model_uses_accelerate_hooks(model: Any) -> bool:
if model is None:
return False
if getattr(model, "hf_device_map", None):
return True
if getattr(model, "_hf_hook", None) is not None:
return True
try:
for module in model.modules():
if getattr(module, "_hf_hook", None) is not None:
return True
except Exception:
pass
return False
class SmolVLMCaptionEngine:
def __init__(self, config: SmolVLMCaptionConfig, generation: Optional[GenerationConfig] = None, cleanup: Optional[CleanupConfig] = None) -> None:
self.config = config
self.generation = generation or GenerationConfig()
self.cleanup = cleanup or CleanupConfig()
self.processor = None
self.model = None
self.local_model_path: Optional[Path] = None
self.model_size_bytes: Optional[int] = None
self._last_prompt_metadata: dict[str, Any] = {}
self._comfy_mm = _try_get_comfy_model_management() if config.use_comfy_model_management else None
self.inference_device = self._resolve_inference_device(config.device)
def resolve_caption_prompt(self) -> tuple[str, dict[str, Any]]:
custom = (self.config.prompt or "").replace("\\n", "\n").strip()
if custom:
metadata = {
"source": "custom_or_legacy_engine_prompt",
"dialect": "smolvlm",
"caption_type": getattr(self.config, "caption_type", "custom"),
"caption_length": getattr(self.config, "caption_length", "custom"),
"extra_options": list(getattr(self.config, "extra_options", ())),
"name_input": getattr(self.config, "name_input", ""),
}
return custom, metadata
prompt, spec = build_caption_prompt_spec(
caption_type=getattr(self.config, "caption_type", "SFW Character Caption"),
caption_length=getattr(self.config, "caption_length", "medium-length"),
extra_options=getattr(self.config, "extra_options", ()),
name_input=getattr(self.config, "name_input", ""),
dialect="smolvlm",
)
metadata = spec.to_metadata() if hasattr(spec, "to_metadata") else dict(spec)
return prompt, metadata
def _resolve_inference_device(self, device: str) -> torch.device:
if self._comfy_mm is not None:
try:
dev = self._comfy_mm.get_torch_device()
if dev.type == "cuda":
return dev
except Exception:
pass
if device == "auto":
if not torch.cuda.is_available():
raise RuntimeError("CaptionForge SmolVLM requires CUDA for release builds.")
return torch.device("cuda")
dev = torch.device(device)
if dev.type != "cuda":
raise RuntimeError(f"CaptionForge SmolVLM release mode does not allow inference device={device!r}")
return dev
def _resolve_dtype(self, dtype: str):
text = dtype.lower().strip()
if text in {"bf16", "bfloat16"}:
return torch.bfloat16
if text in {"fp16", "float16"}:
return torch.float16
if text in {"fp32", "float32"}:
return torch.float32
return torch.bfloat16
def _cache_policy(self) -> str:
policy = (self.config.cache_policy or "").strip()
return policy or ("evict_other_caption_models" if self.config.keep_loaded else "unload_after_run")
def _cache_key(self, local_path: Path) -> str:
return make_cache_key(role="caption", family="smolvlm", model_path=str(local_path.resolve()), device=str(self.inference_device), quantization=self.config.memory_mode, dtype=self.config.dtype)
def resolve_model_path(self) -> Path:
if self.config.model_path.strip():
return Path(self.config.model_path).expanduser()
return download_registry_model_if_needed(self.config.model_name, self.config.model_root, metadata_only=False, allow_download=self.config.allow_download)
def _module_size(self, model) -> int:
total = 0
try:
for p in model.parameters():
total += p.numel() * p.element_size()
for b in model.buffers():
total += b.numel() * b.element_size()
except Exception:
pass
return total
def load(self) -> None:
local_path = self.resolve_model_path()
self.local_model_path = local_path
cache_key = self._cache_key(local_path)
cached = get_cached_model(cache_key)
if cached is not None:
self.processor = cached.get("processor")
self.model = cached.get("model")
self.model_size_bytes = cached.get("model_size_bytes")
print(f"[JLC SmolVLM Engine] Reusing cached model: {local_path}")
return
prepare_for_model_load(cache_key, policy=self._cache_policy(), role="caption")
if self.generation.seed is not None:
set_seed(self.generation.seed)
print(f"[JLC SmolVLM Engine] Local model path: {local_path}")
print(f"[JLC SmolVLM Engine] Local weight size: {_format_bytes(_folder_weight_size_bytes(local_path))}")
print(f"[JLC SmolVLM Engine] { _cuda_diagnostic_line() }")
from transformers import AutoProcessor
try:
from transformers import AutoModelForImageTextToText as SmolModel
except Exception:
try:
from transformers import AutoModelForVision2Seq as SmolModel
except Exception:
from transformers import AutoModelForCausalLM as SmolModel
print(f"[JLC SmolVLM Engine] Loading processor: {local_path}")
self.processor = AutoProcessor.from_pretrained(str(local_path), trust_remote_code=self.config.trust_remote_code)
kwargs: dict[str, Any] = {"torch_dtype": self._resolve_dtype(self.config.dtype), "trust_remote_code": self.config.trust_remote_code}
if self.config.memory_mode == "Balanced (8-bit)":
from transformers import BitsAndBytesConfig
kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
kwargs["device_map"] = {"": self.inference_device.index or 0}
print(f"[JLC SmolVLM Engine] Loading model: {local_path}")
print(f"[JLC SmolVLM Engine] dtype={kwargs.get('torch_dtype')}, device={self.inference_device}, memory_mode={self.config.memory_mode}")
self.model = SmolModel.from_pretrained(str(local_path), **kwargs)
device_map = getattr(self.model, "hf_device_map", None)
if device_map:
print(f"[JLC SmolVLM Engine] hf_device_map: {device_map}")
else:
try:
print(f"[JLC SmolVLM Engine] first parameter device: {next(self.model.parameters()).device}")
except Exception:
pass
self.model.eval()
if self.config.memory_mode == "Default":
self.model.to(self.inference_device)
self.model_size_bytes = self._module_size(self.model)
print(f"[JLC SmolVLM Engine] Estimated loaded module size: {_format_bytes(self.model_size_bytes)}")
register_model(cache_key, {"processor": self.processor, "model": self.model, "model_size_bytes": self.model_size_bytes}, family="smolvlm", model_path=str(local_path), device=str(self.inference_device), quantization=self.config.memory_mode, role="caption", unload_fn=_unload_bundle)
print("[JLC SmolVLM Engine] Model loaded.")
def unload(self) -> None:
if self.local_model_path is not None:
unload_after_run(self._cache_key(self.local_model_path), enabled=True)
self.processor = None
self.model = None
if torch.cuda.is_available():
torch.cuda.empty_cache()
@torch.inference_mode()
def caption_pil(self, image: Image.Image) -> tuple[str, str]:
if self.processor is None or self.model is None:
self.load()
if self.generation.seed is not None:
set_seed(self.generation.seed)
image = resize_for_model(image.convert("RGB"), self.config.max_size)
prompt, prompt_metadata = self.resolve_caption_prompt()
self._last_prompt_metadata = dict(prompt_metadata)
messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt}]}]
if hasattr(self.processor, "apply_chat_template"):
text = self.processor.apply_chat_template(messages, add_generation_prompt=True)
inputs = self.processor(text=text, images=[image], return_tensors="pt")
else:
inputs = self.processor(text=prompt, images=image, return_tensors="pt")
inputs = {k: (v.to(self.inference_device) if hasattr(v, "to") else v) for k, v in inputs.items()}
if "pixel_values" in inputs and hasattr(self.model, "dtype") and torch.is_floating_point(inputs["pixel_values"]):
inputs["pixel_values"] = inputs["pixel_values"].to(dtype=self.model.dtype)
gen_kwargs = {"max_new_tokens": int(self.generation.max_new_tokens), "num_beams": max(1, int(self.generation.num_beams)), "do_sample": bool(self.generation.temperature > 0), "use_cache": True}
if self.generation.temperature > 0:
gen_kwargs.update({"temperature": float(self.generation.temperature), "top_p": float(self.generation.top_p), "top_k": int(self.generation.top_k)})
if self.generation.repetition_penalty != 1.0:
gen_kwargs["repetition_penalty"] = float(self.generation.repetition_penalty)
generated = self.model.generate(**inputs, **gen_kwargs)
input_len = inputs["input_ids"].shape[-1] if isinstance(inputs.get("input_ids"), torch.Tensor) else 0
trimmed = generated[:, input_len:] if input_len and generated.shape[-1] > input_len else generated
raw = self.processor.batch_decode(trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0].strip()
final = cleanup_caption(raw, self.cleanup)
if self._cache_policy() == "unload_after_run":
self.unload()
return final, raw
def caption_path(self, image_path: str | Path) -> CaptionRecord:
p = Path(image_path)
final, raw = self.caption_pil(load_image_file(p))
return CaptionRecord(image=str(p), caption=final, raw_caption=raw, model_name=self.config.model_name, model_path=str(self.local_model_path or self.config.model_path), prompt=self.resolve_caption_prompt()[0], seed=self.generation.seed, temperature=self.generation.temperature, top_p=self.generation.top_p, top_k=self.generation.top_k, num_beams=self.generation.num_beams, max_new_tokens=self.generation.max_new_tokens, max_size=self.config.max_size, timestamp=iso_timestamp(), image_key=str(p.resolve()), prompt_metadata=dict(self._last_prompt_metadata))
def caption_batch(self, batch: BatchCaptionConfig) -> BatchCaptionResult:
if not batch.input_path.strip():
raise ValueError("BatchCaptionConfig.input_path is required.")
result = BatchCaptionResult()
images = list(iter_image_files(batch.input_path, batch.recursive, batch.filename_glob, batch.extensions))
if batch.limit > 0:
images = images[: int(batch.limit)]
if not images:
return result
output_dir = Path(batch.output_dir) if batch.output_dir.strip() else None
jsonl_path = (output_dir or images[0].parent) / (batch.jsonl_filename.strip() or "captions.jsonl") if batch.write_jsonl else None
also_jsonl_path = Path(batch.also_jsonl_path) if batch.also_jsonl_path.strip() else None
seen_jsonl = set()
if batch.skip_existing_jsonl_images:
if jsonl_path is not None:
seen_jsonl |= load_existing_jsonl_images(jsonl_path)
if also_jsonl_path is not None:
seen_jsonl |= load_existing_jsonl_images(also_jsonl_path)
if batch.write_run_config:
outdir = output_dir or images[0].parent
name = batch.run_config_filename.strip() or f"jlc_smolvlm_caption_run_config_{timestamp()}.json"
write_run_config_json(outdir / name, self.build_run_config(batch), dry_run=batch.dry_run)
for image_path in images:
try:
txt_path = sidecar_txt_path(image_path, output_dir)
if batch.skip_existing_txt and batch.write_txt and txt_path.exists() and not batch.overwrite:
result.skipped += 1
continue
if batch.skip_existing_jsonl_images and str(image_path) in seen_jsonl:
result.skipped += 1
continue
rec = self.caption_path(image_path)
result.records.append(rec)
if batch.write_txt:
write_text_sidecar(txt_path, rec.caption, batch.overwrite, batch.backup_existing, batch.dry_run)
if jsonl_path is not None:
append_jsonl_records(jsonl_path, [rec], batch.dry_run)
if also_jsonl_path is not None:
append_jsonl_records(also_jsonl_path, [rec], batch.dry_run)
except Exception as exc:
result.failed += 1
err = CaptionRecord(image=str(image_path), caption="", raw_caption="", model_name=self.config.model_name, model_path=str(self.local_model_path or self.config.model_path), prompt=self.resolve_caption_prompt()[0], seed=self.generation.seed, temperature=self.generation.temperature, top_p=self.generation.top_p, top_k=self.generation.top_k, num_beams=self.generation.num_beams, max_new_tokens=self.generation.max_new_tokens, max_size=self.config.max_size, timestamp=iso_timestamp(), status="error", error=str(exc), image_key=str(Path(image_path).resolve()), prompt_metadata=self.resolve_caption_prompt()[1])
result.records.append(err)
if jsonl_path is not None:
append_jsonl_records(jsonl_path, [err], batch.dry_run)
if also_jsonl_path is not None:
append_jsonl_records(also_jsonl_path, [err], batch.dry_run)
return result
def build_run_config(self, batch: Optional[BatchCaptionConfig] = None) -> dict[str, Any]:
cleanup = asdict(self.cleanup)
cleanup["replacement_rules"] = [list(x) for x in self.cleanup.replacement_rules]
return {"timestamp": iso_timestamp(), "engine": "JLC SmolVLM Caption Engine", "smolvlm_config": asdict(self.config), "prompt_metadata": self.resolve_caption_prompt()[1], "generation": asdict(self.generation), "cleanup": cleanup, "batch": json_safe(asdict(batch)) if batch is not None else None}
+28 -46
View File
@@ -9,12 +9,9 @@ JLC Joy Caption Engine
- Repository
https://github.com/Damkohler/CaptionForge
- CaptionForge focuses on practical dataset-captioning infrastructure for:
• LoRA dataset preparation
• multi-engine caption generation
• JSONL audit trails
• claim extraction and refinement
• consensus-oriented caption improvement
- CaptionForge 1.0 uses independent Pass-A witnesses, text-LLM synthesis,
image-aware validation, SHORT/TAGGY formatting, and JSONL audit trails to
produce grounded LoRA dataset captions.
- Module Purpose
- The **JLC Joy Caption Engine** provides the shared importable backend for
@@ -101,12 +98,10 @@ JLC Joy Caption Engine
- The engine therefore prioritizes reproducibility, auditability, and clean
separation between model-specific inference code and ComfyUI node wrappers.
- ⚠️ Development Status
- This is early CaptionForge Pass A engine infrastructure.
- Registry entries are intentionally conservative and should be expanded only
after processor/model compatibility is tested.
- Memory behavior, prompt presets, and audit fields may evolve as the
multi-pass CaptionForge pipeline matures.
- Production Status
- Active CaptionForge 1.0 Joy witness backend. Its registry remains
deliberately conservative because processor/model compatibility and VRAM
behavior are model-specific.
- Attribution & License
- Concept and implementation by **J. L. Córdova**
@@ -150,7 +145,6 @@ from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Iterable, Optional
import warnings
from PIL import Image
import torch
@@ -162,22 +156,10 @@ from .captionforge_model_cache import (
prepare_for_model_load,
unload_after_run,
)
# -------------------------------------------------------------------------
# Eliminate noise from transformers: UserWarning:
# MatMul8bitLt: inputs will be cast from torch.float32 to float16 during quantization
# -------------------------------------------------------------------------
warnings.filterwarnings(
"ignore",
message=r".*MatMul8bitLt: inputs will be cast.*",
category=UserWarning,
)
warnings.filterwarnings(
"ignore",
message=r".*torchvision backend image processor with LANCZOS resample.*",
from .captionforge_joy_warnings import quantized_inference_warnings, warn_if_suspicious_8bit_stack
from .captionforge_cleanup import (
remove_forbidden_phrases as _remove_forbidden_phrases_boundary_safe,
replace_phrases as _replace_phrases_boundary_safe,
)
@@ -389,7 +371,7 @@ class CleanupConfig:
forbidden_phrases: list[str] = field(default_factory=list)
replacement_rules: list[tuple[str, str]] = field(default_factory=list)
replace_case_insensitive: bool = True
replace_whole_words_only: bool = False
replace_whole_words_only: bool = True
strip_boilerplate_prefixes: bool = True
strip_trailing_period: bool = True
@@ -801,7 +783,7 @@ def apply_replacements(
caption: str,
rules: list[tuple[str, str]],
case_insensitive: bool = True,
whole_words_only: bool = False,
whole_words_only: bool = True,
) -> str:
if not rules:
return caption
@@ -815,13 +797,15 @@ def apply_replacements(
if not old:
continue
flags = re.IGNORECASE if case_insensitive else 0
pattern = re.escape(old)
if whole_words_only:
pattern = r"\b" + pattern + r"\b"
result = re.sub(pattern, new, result, flags=flags)
result = _replace_phrases_boundary_safe(
result,
[(old, new)],
case_insensitive=case_insensitive,
)
else:
flags = re.IGNORECASE if case_insensitive else 0
result = re.sub(re.escape(old), new, result, flags=flags)
return result
@@ -830,14 +814,7 @@ def remove_forbidden_phrases(caption: str, forbidden_phrases: list[str]) -> str:
if not forbidden_phrases:
return caption
result = caption
for phrase in forbidden_phrases:
phrase = phrase.strip()
if not phrase:
continue
result = re.sub(re.escape(phrase), "", result, flags=re.IGNORECASE)
result = _remove_forbidden_phrases_boundary_safe(caption, forbidden_phrases)
result = re.sub(r"\s+,", ",", result)
result = re.sub(r",\s*,+", ",", result)
result = re.sub(r"\s+", " ", result)
@@ -1348,6 +1325,7 @@ class JoyCaptionEngine:
self._free_memory(self.model_size_bytes, self.offload_device)
self.model.to(self.offload_device)
else:
warn_if_suspicious_8bit_stack("Joy Caption")
print(f"[JLC Joy Engine] Loading model in {self.config.memory_mode}: {local_path}")
try:
from transformers import BitsAndBytesConfig
@@ -1582,7 +1560,11 @@ class JoyCaptionEngine:
generation_kwargs["eos_token_id"] = eos_token_id
try:
with torch.autocast(
# Keep Joy's native BF16 path. LLM.int8 casts activations to FP16
# internally; silence only that exact diagnostic during 8-bit generation.
with quantized_inference_warnings(
self.config.memory_mode == "Balanced (8-bit)"
), torch.autocast(
device_type=device_type,
dtype=torch.bfloat16,
enabled=autocast_enabled and bf16_supported,
+254 -51
View File
@@ -9,12 +9,9 @@ JLC Qwen Caption Engine
- Repository
https://github.com/Damkohler/CaptionForge
- CaptionForge focuses on practical dataset-captioning infrastructure for:
• LoRA dataset preparation
• multi-engine caption generation
• JSONL audit trails
• claim extraction and refinement
• consensus-oriented caption improvement
- CaptionForge 1.0 uses independent Pass-A witnesses, text-LLM synthesis,
image-aware validation, SHORT/TAGGY formatting, and JSONL audit trails to
produce grounded LoRA dataset captions.
- Module Purpose
- The **JLC Qwen Caption Engine** provides the shared importable backend for
@@ -112,12 +109,10 @@ JLC Qwen Caption Engine
- The engine therefore prioritizes reproducibility, auditability, and clean
separation between model-specific inference code and ComfyUI node wrappers.
- ⚠️ Development Status
- This is early CaptionForge Pass A engine infrastructure.
- Registry entries should be expanded only after model-class, processor, and
memory behavior are tested.
- Quantization, prompt presets, compatibility patches, and audit fields may
evolve as the multi-pass CaptionForge pipeline matures.
- Production Status
- Active CaptionForge 1.0 Qwen witness backend. Registry expansion remains
conservative because processor/model compatibility, quantization, and
VRAM behavior are model-specific.
- Attribution & License
- Concept and implementation by **J. L. Córdova**
@@ -151,9 +146,11 @@ MANIFEST = {
),
}
import ctypes
import fnmatch
import json
import logging
import os
import random
import re
import shutil
@@ -161,7 +158,6 @@ from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Iterable, Optional
import warnings
from PIL import Image
import torch
@@ -173,6 +169,11 @@ from .captionforge_model_cache import (
prepare_for_model_load,
unload_after_run,
)
from .captionforge_joy_warnings import quantized_inference_warnings, warn_if_suspicious_8bit_stack
from .captionforge_cleanup import (
remove_forbidden_phrases as _remove_forbidden_phrases_boundary_safe,
replace_phrases as _replace_phrases_boundary_safe,
)
try:
from .captionforge_caption_prompt_kit import (
@@ -339,7 +340,7 @@ class CleanupConfig:
forbidden_phrases: list[str] = field(default_factory=list)
replacement_rules: list[tuple[str, str]] = field(default_factory=list)
replace_case_insensitive: bool = True
replace_whole_words_only: bool = False
replace_whole_words_only: bool = True
strip_boilerplate_prefixes: bool = True
strip_trailing_period: bool = True
@@ -786,7 +787,7 @@ def apply_replacements(
caption: str,
rules: list[tuple[str, str]],
case_insensitive: bool = True,
whole_words_only: bool = False,
whole_words_only: bool = True,
) -> str:
if not rules:
return caption
@@ -800,13 +801,15 @@ def apply_replacements(
if not old:
continue
flags = re.IGNORECASE if case_insensitive else 0
pattern = re.escape(old)
if whole_words_only:
pattern = r"\b" + pattern + r"\b"
result = re.sub(pattern, new, result, flags=flags)
result = _replace_phrases_boundary_safe(
result,
[(old, new)],
case_insensitive=case_insensitive,
)
else:
flags = re.IGNORECASE if case_insensitive else 0
result = re.sub(re.escape(old), new, result, flags=flags)
return result
@@ -815,14 +818,7 @@ def remove_forbidden_phrases(caption: str, forbidden_phrases: list[str]) -> str:
if not forbidden_phrases:
return caption
result = caption
for phrase in forbidden_phrases:
phrase = phrase.strip()
if not phrase:
continue
result = re.sub(re.escape(phrase), "", result, flags=re.IGNORECASE)
result = _remove_forbidden_phrases_boundary_safe(caption, forbidden_phrases)
result = re.sub(r"\s+,", ",", result)
result = re.sub(r",\s*,+", ",", result)
result = re.sub(r"\s+", " ", result)
@@ -1079,6 +1075,175 @@ def _cuda_diagnostic_line() -> str:
except Exception as exc:
return f"CUDA diagnostics unavailable: {exc}"
def _available_system_memory_bytes() -> int | None:
"""Best-effort available physical RAM without adding a new dependency."""
try:
if os.name == "nt":
class MEMORYSTATUSEX(ctypes.Structure):
_fields_ = [
("dwLength", ctypes.c_ulong),
("dwMemoryLoad", ctypes.c_ulong),
("ullTotalPhys", ctypes.c_ulonglong),
("ullAvailPhys", ctypes.c_ulonglong),
("ullTotalPageFile", ctypes.c_ulonglong),
("ullAvailPageFile", ctypes.c_ulonglong),
("ullTotalVirtual", ctypes.c_ulonglong),
("ullAvailVirtual", ctypes.c_ulonglong),
("ullAvailExtendedVirtual", ctypes.c_ulonglong),
]
status = MEMORYSTATUSEX()
status.dwLength = ctypes.sizeof(MEMORYSTATUSEX)
if ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status)):
return int(status.ullAvailPhys)
page_size = os.sysconf("SC_PAGE_SIZE")
avail_pages = os.sysconf("SC_AVPHYS_PAGES")
return int(page_size * avail_pages)
except Exception:
return None
def _qwen_memory_budget(headroom: float = 0.20) -> dict[Any, int]:
"""Return conservative Accelerate max_memory budgets from currently free memory."""
usable_fraction = max(0.10, min(1.0, 1.0 - float(headroom)))
budgets: dict[Any, int] = {}
if torch.cuda.is_available():
try:
index = torch.cuda.current_device()
free_bytes, _total_bytes = torch.cuda.mem_get_info(index)
budgets[index] = max(1, int(free_bytes * usable_fraction))
except Exception:
pass
available_ram = _available_system_memory_bytes()
if available_ram:
# infer_auto_device_map below estimates the skeleton with dtype=int8.
# bitsandbytes CPU-offloaded modules, however, remain FP32. Reduce the
# CPU capacity by 4x so an int8-sized placement estimate corresponds to
# the actual four-byte-per-parameter CPU residency requirement.
budgets["cpu"] = max(
1,
int((available_ram * usable_fraction) / 4.0),
)
return budgets
def _format_memory_budget(max_memory: dict[Any, int] | None) -> str:
if not max_memory:
return "automatic"
return ", ".join(
f"{device}={_format_bytes(value)}"
for device, value in max_memory.items()
)
def _summarize_device_map(device_map: dict[str, Any] | None) -> str:
if not isinstance(device_map, dict) or not device_map:
return "none"
counts: dict[str, int] = {}
for mapped_device in device_map.values():
if isinstance(mapped_device, int):
label = f"cuda:{mapped_device}"
else:
label = str(mapped_device)
counts[label] = counts.get(label, 0) + 1
return ", ".join(
f"{device}: {count} module(s)"
for device, count in sorted(counts.items())
)
def _resolve_model_execution_device(model: Any) -> torch.device:
"""Choose the device that should receive inference inputs for dispatched models."""
hook = getattr(model, "_hf_hook", None)
hook_device = getattr(hook, "execution_device", None)
if hook_device is not None:
return torch.device(hook_device)
device_map = getattr(model, "hf_device_map", None)
if isinstance(device_map, dict):
for mapped_device in device_map.values():
if mapped_device in {"cpu", "disk", "meta"}:
continue
if isinstance(mapped_device, int):
return torch.device(f"cuda:{mapped_device}")
try:
candidate = torch.device(mapped_device)
except Exception:
continue
if candidate.type not in {"cpu", "meta"}:
return candidate
return next(model.parameters()).device
def _build_qwen_8bit_device_map(
model_cls: Any,
local_path: Path,
trust_remote_code: bool,
) -> tuple[dict[str, Any] | str, dict[Any, int]]:
"""
Build a conservative Accelerate placement map without materializing weights.
GPU placement is estimated as int8. CPU overflow remains FP32 at real load
time through bitsandbytes CPU offload. Twenty percent of currently available
GPU/RAM is intentionally left outside the placement budget.
"""
if not torch.cuda.is_available():
return "auto", {}
max_memory = _qwen_memory_budget(headroom=0.20)
if not max_memory:
return "auto", {}
try:
from accelerate import infer_auto_device_map, init_empty_weights
from transformers import AutoConfig
model_config = AutoConfig.from_pretrained(
str(local_path),
trust_remote_code=trust_remote_code,
)
with init_empty_weights():
empty_model = model_cls(model_config)
no_split_modules = getattr(empty_model, "_no_split_modules", None)
device_map = infer_auto_device_map(
empty_model,
max_memory=max_memory,
no_split_module_classes=no_split_modules,
dtype=torch.int8,
)
if any(str(device).lower() == "disk" for device in device_map.values()):
raise RuntimeError(
"CaptionForge Qwen 8-bit placement would require disk offload. "
"Disk spill is intentionally not enabled because it is extremely slow "
"and can make ComfyUI inference impractical. More GPU VRAM or available "
"system RAM is required for this model."
)
return device_map, max_memory
except RuntimeError:
raise
except Exception as exc:
print(
"[JLC Qwen Engine] Adaptive placement probe was unavailable; "
f"falling back to Accelerate device_map='auto': {exc}"
)
return "auto", max_memory
def json_safe(value):
if isinstance(value, set):
return sorted(value)
@@ -1247,6 +1412,9 @@ class QwenCaptionEngine:
self.model = cached["model"]
print(f"[JLC Qwen Engine] Reusing cached model: {local_path}")
return
if quantization == "bnb_8bit":
warn_if_suspicious_8bit_stack("Qwen Caption")
cache_policy = getattr(
self.config,
@@ -1319,19 +1487,26 @@ class QwenCaptionEngine:
"Install/verify compatible packages before using quantization='bnb_8bit'."
) from exc
warnings.filterwarnings(
"ignore",
message=r"MatMul8bitLt: inputs will be cast from torch\.bfloat16 to float16 during quantization",
category=UserWarning,
module=r"bitsandbytes\.autograd\._functions",
model_kwargs["quantization_config"] = BitsAndBytesConfig(
load_in_8bit=True,
llm_int8_enable_fp32_cpu_offload=True,
)
model_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
# bitsandbytes quantized models should be loaded through Accelerate dispatch.
# Keep this explicit so users do not accidentally request a later .to(device).
if not effective_device_map:
effective_device_map = "auto"
# Balanced 8-bit is intended to remain usable when the complete
# quantized model does not fit in VRAM. Build an adaptive placement
# map with headroom and permit supported FP32 CPU overflow.
if not effective_device_map or effective_device_map == "auto":
effective_device_map, max_memory = _build_qwen_8bit_device_map(
model_cls,
local_path,
self.config.trust_remote_code,
)
if max_memory:
model_kwargs["max_memory"] = max_memory
print(
"[JLC Qwen Engine] Adaptive memory budget: "
f"{_format_memory_budget(max_memory)}"
)
if effective_device_map:
model_kwargs["device_map"] = effective_device_map
@@ -1342,14 +1517,40 @@ class QwenCaptionEngine:
f"device_map={effective_device_map!r}, quantization={quantization}"
)
self.model = model_cls.from_pretrained(
str(local_path),
**model_kwargs,
)
try:
self.model = model_cls.from_pretrained(
str(local_path),
**model_kwargs,
)
except Exception as exc:
message = str(exc)
if quantization == "bnb_8bit" and (
"Some modules are dispatched on the CPU or the disk" in message
or "llm_int8_enable_fp32_cpu_offload" in message
or "device_map" in message and "CPU" in message
):
raise RuntimeError(
"CaptionForge could not place this Qwen model within the "
"available GPU/CPU memory budget using Balanced (8-bit). "
f"{_cuda_diagnostic_line()}. "
f"Memory budget: {_format_memory_budget(model_kwargs.get('max_memory'))}. "
"The 8-bit path supports FP32 CPU offload, but this configuration "
"still could not be loaded. Close other GPU/RAM-heavy applications, "
"use a smaller Qwen model, or free additional system memory."
) from exc
raise
device_map = getattr(self.model, "hf_device_map", None)
if device_map:
print(f"[JLC Qwen Engine] hf_device_map: {device_map}")
if any(str(mapped).lower() == "disk" for mapped in device_map.values()):
raise RuntimeError(
"CaptionForge Qwen loaded with disk-offloaded modules. "
"This configuration is not supported for interactive captioning."
)
print(
"[JLC Qwen Engine] Placement: "
f"{_summarize_device_map(device_map)}"
)
else:
try:
print(f"[JLC Qwen Engine] first parameter device: {next(self.model.parameters()).device}")
@@ -1485,8 +1686,7 @@ class QwenCaptionEngine:
)
try:
device = next(self.model.parameters()).device
inputs = inputs.to(device)
inputs = inputs.to(_resolve_model_execution_device(self.model))
except Exception:
pass
@@ -1519,10 +1719,13 @@ class QwenCaptionEngine:
generation_kwargs["do_sample"] = False
generated_ids = self.model.generate(
**inputs,
**generation_kwargs,
)
with quantized_inference_warnings(
self._resolve_quantization(self.config.quantization) == "bnb_8bit"
):
generated_ids = self.model.generate(
**inputs,
**generation_kwargs,
)
generated_ids_trimmed = [
out_ids[len(in_ids):]
+1 -1
View File
@@ -1 +1 @@
# empty file
"""ComfyUI-facing CaptionForge nodes and production support helpers."""
+1 -1
View File
@@ -1 +1 @@
# empty file
"""Active Joy, Qwen, and Ollama Pass-A witness node wrappers."""
@@ -8,11 +8,9 @@ JLC CaptionForge Joy Caption — ComfyUI Node Wrapper
- Repository
https://github.com/Damkohler/CaptionForge
- CaptionForge focuses on practical dataset-captioning infrastructure for
LoRA dataset preparation, using multi-engine caption generation, JSONL
audit trails, claim extraction and refinement, text-LLM distillation,
image-aware VLM validation, and consensus-oriented caption improvement
to produce grounded, auditable training captions.
- CaptionForge 1.0 uses independent Pass-A witnesses, text-LLM synthesis,
image-aware validation, SHORT/TAGGY formatting, and JSONL audit trails to
produce grounded LoRA dataset captions.
- Node Purpose
- The **JLC CaptionForge Joy Caption** node provides a ComfyUI
@@ -30,7 +28,6 @@ JLC CaptionForge Joy Caption — ComfyUI Node Wrapper
• clear template-vs-custom prompt controls
• CaptionForge Pipeline Planner consumption through `pipeline_plan`
• CaptionForge Template Options consumption through `template_options`
• TXT audit sidecar writing in planned runs
• shared JSONL audit output in planned runs
• direct ComfyUI caption and resolved-prompt string outputs
• IMAGE and pipeline-plan passthrough for clean graph chaining
@@ -89,7 +86,6 @@ JLC CaptionForge Joy Caption — ComfyUI Node Wrapper
• shared LoRA trigger word
- In planned mode, the node can write:
• TXT audit sidecar captions
• JSONL audit records
• run-configuration JSON files
@@ -134,12 +130,10 @@ JLC CaptionForge Joy Caption — ComfyUI Node Wrapper
clean separation between the ComfyUI interface and the backend model
engine.
- ⚠️ Development Status
- This is early CaptionForge raw-caption infrastructure.
- The UI, model registry, prompt behavior, and output audit fields may
evolve as the multi-pass CaptionForge pipeline matures.
- The node is intended for local dataset preparation and controlled caption
audit workflows.
- Production Status
- Active CaptionForge 1.0 Pass-A witness node. In planned mode the Planner
owns run counts and sampling schedules; in standalone mode this node's
visible controls are authoritative.
- Attribution & License
- Concept and implementation by **J. L. Córdova**
@@ -173,7 +167,7 @@ MANIFEST = {
"consumed only through the template_options input from the CaptionForge "
"Template Options sidecar. The UI separates caption_template_mode and "
"custom_prompt_mode for clearer prompt routing while delegating model "
"loading, memory mode, generation, cleanup, TXT sidecars, and JSONL audit "
"loading, memory mode, generation, cleanup, and JSONL audit "
"records to the Joy caption engine."
),
}
@@ -206,9 +200,11 @@ from ...engines.jlc_joy_caption_engine import (
resolve_prompt,
timestamp,
write_run_config_json,
write_text_sidecar,
)
from ...engines.captionforge_pipeline_planner_engine import expand_captionforge_runs
from ...engines.captionforge_cleanup import resolve_cleanup_settings
from ...engines.captionforge_source_identity import file_source_identity, optional_image_identity
from ...engines.captionforge_dataset_export import is_dataset_export
from ..jlc_captionforge_template_options import resolve_effective_extra_options
@@ -250,20 +246,7 @@ def _tensor_to_pil(image_tensor) -> list[Image.Image]:
return images
def _safe_source_name(value: str) -> str:
cleaned = []
for ch in value.replace("\\", "/"):
if ch.isalnum() or ch in {"-", "_", "."}:
cleaned.append(ch)
elif ch == "/":
cleaned.append("__")
else:
cleaned.append("_")
out = "".join(cleaned).strip("._")
return out or "image"
def _iter_input_path_images(input_path: str, recursive: bool, filename_glob: str) -> list[tuple[str, Path]]:
def _iter_input_path_images(input_path: str, recursive: bool, filename_glob: str) -> list[tuple[str, str, Path]]:
root = Path(str(input_path or "").strip())
if not root:
return []
@@ -271,22 +254,26 @@ def _iter_input_path_images(input_path: str, recursive: bool, filename_glob: str
raise RuntimeError(f"CaptionForge input_path does not exist: {root}")
glob_text = (filename_glob or "*").strip() or "*"
explicit_dataset_source = is_dataset_export(root)
if root.is_file():
if root.suffix.lower() not in _SUPPORTED_IMAGE_SUFFIXES:
raise RuntimeError(f"CaptionForge input_path is not a supported image file: {root}")
return [(_safe_source_name(root.stem), root)]
image_name, image_key = file_source_identity(root, root)
return [(image_name, image_key, root)]
pattern_iter = root.rglob(glob_text) if recursive else root.glob(glob_text)
paths = sorted(
p for p in pattern_iter
if p.is_file() and p.suffix.lower() in _SUPPORTED_IMAGE_SUFFIXES
if p.is_file()
and p.suffix.lower() in _SUPPORTED_IMAGE_SUFFIXES
and (explicit_dataset_source or not is_dataset_export(p))
)
items: list[tuple[str, Path]] = []
items: list[tuple[str, str, Path]] = []
for path in paths:
rel = path.relative_to(root).with_suffix("")
items.append((_safe_source_name(str(rel)), path))
image_name, image_key = file_source_identity(path, root)
items.append((image_name, image_key, path))
return items
@@ -323,12 +310,6 @@ def _planned_caption_jsonl_path(pipeline_plan) -> Path | None:
return None
def _run_txt_path(output_dir: Path, source_name: str, run_count: int, run_index: int) -> Path:
if run_count <= 1:
return output_dir / f"{source_name}.txt"
return output_dir / f"{source_name}__cf_run_{run_index:02d}.txt"
def _use_template_mode(caption_template_mode: bool, custom_prompt_mode: bool) -> bool:
if bool(custom_prompt_mode):
return False
@@ -666,7 +647,7 @@ class JLC_CaptionForgeJoy:
"resolved_prompt",
)
FUNCTION = "caption"
CATEGORY = "Captioning/CaptionForge/Captioning Nodes"
CATEGORY = "Caption/CaptionForge/Caption Nodes"
@classmethod
def IS_CHANGED(cls, **kwargs):
@@ -715,7 +696,7 @@ class JLC_CaptionForgeJoy:
if download_probe_only:
result = probe_registry_model_download(model, JLC_JOY_MODEL_ROOT)
return (image, pipeline_plan, result, resolved_prompt)
return (image, pipeline_plan, template_options, result, resolved_prompt)
effective_seed = -1 if seed is None else int(seed)
@@ -739,9 +720,12 @@ class JLC_CaptionForgeJoy:
if run_plan_connected and not run_plan:
status = "[CaptionForge] Joy disabled by Pipeline Planner."
print(status)
return (image, pipeline_plan, status, resolved_prompt)
return (image, pipeline_plan, template_options, status, resolved_prompt)
first_run = run_plan[0]
effective_forbidden, effective_replacements = resolve_cleanup_settings(
_normalize_pipeline_plan(pipeline_plan), forbidden_phrases, replace_pairs
)
generation = GenerationConfig(
max_new_tokens=int(first_run.max_new_tokens),
@@ -756,8 +740,8 @@ class JLC_CaptionForgeJoy:
trigger="",
prefix=(f"{first_run.trigger_word}," if first_run.trigger_word else ""),
suffix="",
forbidden_phrases=_parse_forbidden_lines(forbidden_phrases),
replacement_rules=_parse_replace_pairs(replace_pairs),
forbidden_phrases=effective_forbidden,
replacement_rules=effective_replacements,
)
joy_config = JoyCaptionConfig(
@@ -780,8 +764,8 @@ class JLC_CaptionForgeJoy:
engine = JoyCaptionEngine(config=joy_config, generation=generation, cleanup=cleanup)
direct_images = [(f"comfy_image_{i:04d}", pil) for i, pil in enumerate(_tensor_to_pil(image))]
file_images: list[tuple[str, Path]] = []
direct_images = [(*optional_image_identity(i), pil) for i, pil in enumerate(_tensor_to_pil(image))]
file_images: list[tuple[str, str, Path]] = []
if first_run.input_path:
file_images = _iter_input_path_images(first_run.input_path, first_run.recursive, first_run.filename_glob)
@@ -819,7 +803,7 @@ class JLC_CaptionForgeJoy:
dry_run=False,
)
def process_one(source_name: str, pil: Image.Image):
def process_one(source_name: str, image_key: str, pil: Image.Image):
for run in run_plan:
engine.generation = GenerationConfig(
max_new_tokens=int(run.max_new_tokens),
@@ -833,8 +817,8 @@ class JLC_CaptionForgeJoy:
trigger="",
prefix=(f"{run.trigger_word}," if run.trigger_word else ""),
suffix="",
forbidden_phrases=_parse_forbidden_lines(forbidden_phrases),
replacement_rules=_parse_replace_pairs(replace_pairs),
forbidden_phrases=effective_forbidden,
replacement_rules=effective_replacements,
)
engine.config.max_size = int(run.max_size)
@@ -861,29 +845,22 @@ class JLC_CaptionForgeJoy:
captionforge_pass="A",
model_family="joy",
ensemble_run_index=run.ensemble_run_index,
image_key=source_name,
image_key=image_key,
)
all_records.append(record)
if run_plan_connected and jsonl_path is not None and output_dir is not None:
append_jsonl_records(jsonl_path, [record], dry_run=False)
write_text_sidecar(
_run_txt_path(output_dir, source_name, len(run_plan), run.ensemble_run_index),
record.caption,
overwrite=True,
backup_existing=False,
dry_run=False,
)
print(
f"[JLC CaptionForge Joy] Captioned {source_name} "
f"run {run.ensemble_run_index + 1}/{len(run_plan)}"
)
for source_name, pil in direct_images:
process_one(source_name, pil)
for source_name, image_key, pil in direct_images:
process_one(source_name, image_key, pil)
for source_name, path in file_images:
process_one(source_name, _open_image(path))
for source_name, image_key, path in file_images:
process_one(source_name, image_key, _open_image(path))
if not keep_loaded:
engine.unload()
@@ -8,11 +8,9 @@ JLC CaptionForge Ollama Caption — ComfyUI Node Wrapper
- Repository
https://github.com/Damkohler/CaptionForge
- CaptionForge focuses on practical dataset-captioning infrastructure for
LoRA dataset preparation, using multi-engine caption generation, JSONL
audit trails, claim extraction and refinement, text-LLM distillation,
image-aware VLM validation, and consensus-oriented caption improvement
to produce grounded, auditable training captions.
- CaptionForge 1.0 uses independent Pass-A witnesses, text-LLM synthesis,
image-aware validation, SHORT/TAGGY formatting, and JSONL audit trails to
produce grounded LoRA dataset captions.
- Node Purpose
- The **JLC CaptionForge Ollama Caption** node provides a ComfyUI
@@ -31,7 +29,6 @@ JLC CaptionForge Ollama Caption — ComfyUI Node Wrapper
• clear template-vs-custom prompt controls
• CaptionForge Pipeline Planner consumption through `pipeline_plan`
• CaptionForge Template Options consumption through `template_options`
• TXT audit sidecar writing in planned runs
• shared JSONL audit output in planned runs
• direct ComfyUI caption and resolved-prompt string outputs
• IMAGE and pipeline-plan passthrough for clean graph chaining
@@ -91,7 +88,6 @@ JLC CaptionForge Ollama Caption — ComfyUI Node Wrapper
• shared LoRA trigger word
- In planned mode, the node can write:
• TXT audit sidecar captions
• JSONL audit records
• run-configuration JSON files
@@ -142,12 +138,10 @@ JLC CaptionForge Ollama Caption — ComfyUI Node Wrapper
clean separation between the ComfyUI interface and the backend model
service.
- ⚠️ Development Status
- This is early CaptionForge Ollama-backed raw-caption infrastructure.
- The UI, model-tag config, prompt behavior, and output audit fields may
evolve as the multi-pass CaptionForge pipeline matures.
- The node is intended for local dataset preparation and controlled caption
audit workflows.
- Production Status
- Active CaptionForge 1.0 Pass-A witness node. The selected Ollama VLM tag
remains local to this node; a connected Planner owns shared run counts and
sampling schedules.
- Attribution & License
- Concept and implementation by **J. L. Córdova**
@@ -206,6 +200,10 @@ from PIL import Image
import folder_paths
from ...engines.captionforge_pipeline_planner_engine import expand_captionforge_runs
from ...engines.captionforge_cleanup import resolve_cleanup_settings
from ...engines.captionforge_source_identity import file_source_identity, optional_image_identity
from ...engines.captionforge_dataset_export import is_dataset_export
from ...engines.captionforge_cleanup import contains_forbidden_phrase, replace_phrases
from ...engines.captionforge_caption_prompt_kit import (
CAPTION_LENGTH_CHOICES,
CAPTION_TYPE_CHOICES,
@@ -355,7 +353,7 @@ class OllamaCaptionRecord:
model_path: str = ""
prompt: str = ""
system_prompt: str = ""
seed: int = -1
seed: int | None = None
temperature: float = 0.18
top_p: float = 0.92
top_k: int = 60
@@ -654,7 +652,7 @@ def _ollama_options(
top_p: float,
top_k: int,
repetition_penalty: float,
seed: int,
seed: int | None,
) -> dict[str, Any]:
"""Build Ollama generation options.
@@ -669,7 +667,7 @@ def _ollama_options(
"top_k": int(top_k),
"repeat_penalty": float(repetition_penalty),
}
if int(seed) >= 0:
if seed is not None and int(seed) >= 0:
options["seed"] = int(seed)
return options
@@ -780,7 +778,7 @@ def _ollama_generate_caption(
top_p: float,
top_k: int,
repetition_penalty: float,
seed: int,
seed: int | None,
max_size: int,
keep_loaded: bool,
timeout: float,
@@ -917,20 +915,7 @@ def _pil_to_base64_png(pil: Image.Image, max_size: int) -> str:
return base64.b64encode(buf.getvalue()).decode("ascii")
def _safe_source_name(value: str) -> str:
cleaned = []
for ch in value.replace("\\", "/"):
if ch.isalnum() or ch in {"-", "_", "."}:
cleaned.append(ch)
elif ch == "/":
cleaned.append("__")
else:
cleaned.append("_")
out = "".join(cleaned).strip("._")
return out or "image"
def _iter_input_path_images(input_path: str, recursive: bool, filename_glob: str) -> list[tuple[str, Path]]:
def _iter_input_path_images(input_path: str, recursive: bool, filename_glob: str) -> list[tuple[str, str, Path]]:
root = Path(str(input_path or "").strip())
if not root:
return []
@@ -938,22 +923,26 @@ def _iter_input_path_images(input_path: str, recursive: bool, filename_glob: str
raise RuntimeError(f"CaptionForge input_path does not exist: {root}")
glob_text = (filename_glob or "*").strip() or "*"
explicit_dataset_source = is_dataset_export(root)
if root.is_file():
if root.suffix.lower() not in _SUPPORTED_IMAGE_SUFFIXES:
raise RuntimeError(f"CaptionForge input_path is not a supported image file: {root}")
return [(_safe_source_name(root.stem), root)]
image_name, image_key = file_source_identity(root, root)
return [(image_name, image_key, root)]
pattern_iter = root.rglob(glob_text) if recursive else root.glob(glob_text)
paths = sorted(
p for p in pattern_iter
if p.is_file() and p.suffix.lower() in _SUPPORTED_IMAGE_SUFFIXES
if p.is_file()
and p.suffix.lower() in _SUPPORTED_IMAGE_SUFFIXES
and (explicit_dataset_source or not is_dataset_export(p))
)
items: list[tuple[str, Path]] = []
items: list[tuple[str, str, Path]] = []
for path in paths:
rel = path.relative_to(root).with_suffix("")
items.append((_safe_source_name(str(rel)), path))
image_name, image_key = file_source_identity(path, root)
items.append((image_name, image_key, path))
return items
@@ -990,12 +979,6 @@ def _planned_caption_jsonl_path(pipeline_plan) -> Path | None:
return None
def _run_txt_path(output_dir: Path, source_name: str, run_count: int, run_index: int) -> Path:
if run_count <= 1:
return output_dir / f"{source_name}.txt"
return output_dir / f"{source_name}__cf_run_{run_index:02d}.txt"
def _use_template_mode(caption_template_mode: bool, custom_prompt_mode: bool) -> bool:
if bool(custom_prompt_mode):
return False
@@ -1048,14 +1031,16 @@ def _clean_caption(
replacement_rules: list[tuple[str, str]],
) -> tuple[str, str]:
text = str(raw or "").strip().strip('"').strip()
for old, new in replacement_rules:
text = text.replace(old, new)
text = replace_phrases(
text,
replacement_rules,
case_insensitive=False,
)
if forbidden_phrases:
kept: list[str] = []
for line in text.splitlines() or [text]:
lowered = line.lower()
if any(phrase.lower() in lowered for phrase in forbidden_phrases if phrase):
if contains_forbidden_phrase(line, forbidden_phrases):
continue
kept.append(line)
text = "\n".join(line.strip() for line in kept if line.strip()).strip()
@@ -1074,11 +1059,6 @@ def _append_jsonl_records(path: Path, records: list[OllamaCaptionRecord]) -> Non
f.write(json.dumps(asdict(record), ensure_ascii=False) + "\n")
def _write_text_sidecar(path: Path, text: str) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(str(text or ""), encoding="utf-8")
def _write_run_config_json(path: Path, data: dict[str, Any]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(data, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
@@ -1096,6 +1076,8 @@ def _build_run_config(
top_k: int,
repetition_penalty: float,
max_size: int,
forbidden_phrases: list[str],
replacement_rules: list[tuple[str, str]],
) -> dict[str, Any]:
return {
"backend": "ollama",
@@ -1111,6 +1093,14 @@ def _build_run_config(
"repetition_penalty": float(repetition_penalty),
"max_size": int(max_size),
},
"cleanup": {
"forbidden_phrases": list(forbidden_phrases),
"replacement_rules": [list(rule) for rule in replacement_rules],
"replacement_match_mode": "whole_word_or_phrase_boundary",
"replacement_case_insensitive": False,
"forbidden_match_mode": "whole_word_or_phrase_boundary",
"forbidden_action": "drop_matching_line",
},
"timestamp": datetime.now().isoformat(timespec="seconds"),
}
@@ -1455,7 +1445,7 @@ class JLC_CaptionForgeOllamaCaption:
"resolved_prompt",
)
FUNCTION = "caption"
CATEGORY = "Captioning/CaptionForge/Caption Nodes"
CATEGORY = "Caption/CaptionForge/Caption Nodes"
@classmethod
def IS_CHANGED(cls, **kwargs):
@@ -1510,15 +1500,6 @@ class JLC_CaptionForgeOllamaCaption:
resolved_prompt = _format_resolved_prompt(system_prompt, prompt)
_evict_python_models_before_ollama_if_needed("JLC CaptionForge Ollama Caption")
if download_probe_only:
result = _probe_ollama_model(ollama_url, model_tag, timeout=min(timeout, 60.0))
return (image, pipeline_plan, result, resolved_prompt)
_ensure_ollama_model(ollama_url, model_tag, timeout=timeout)
_persist_caption_model_if_possible(model_tag)
effective_seed = -1 if seed is None else int(seed)
run_plan, planner_key = _expand_ollama_runs_compat(
@@ -1534,7 +1515,16 @@ class JLC_CaptionForgeOllamaCaption:
if run_plan_connected and not run_plan:
status = "[CaptionForge] Ollama Caption disabled by Pipeline Planner or Planner has no Ollama caption count yet."
print(status)
return (image, pipeline_plan, status, resolved_prompt)
return (image, pipeline_plan, template_options, status, resolved_prompt)
_evict_python_models_before_ollama_if_needed("JLC CaptionForge Ollama Caption")
if download_probe_only:
result = _probe_ollama_model(ollama_url, model_tag, timeout=min(timeout, 60.0))
return (image, pipeline_plan, template_options, result, resolved_prompt)
_ensure_ollama_model(ollama_url, model_tag, timeout=timeout)
_persist_caption_model_if_possible(model_tag)
if not run_plan:
# Defensive fallback for standalone mode if expand_captionforge_runs ever changes behavior.
@@ -1557,9 +1547,12 @@ class JLC_CaptionForgeOllamaCaption:
run_plan = [standalone_run]
first_run = run_plan[0]
forbidden, replacements = resolve_cleanup_settings(
_normalize_pipeline_plan(pipeline_plan), forbidden_phrases, replace_pairs
)
direct_images = [(f"comfy_image_{i:04d}", pil) for i, pil in enumerate(_tensor_to_pil(image))]
file_images: list[tuple[str, Path]] = []
direct_images = [(*optional_image_identity(i), pil) for i, pil in enumerate(_tensor_to_pil(image))]
file_images: list[tuple[str, str, Path]] = []
if getattr(first_run, "input_path", ""):
file_images = _iter_input_path_images(first_run.input_path, first_run.recursive, first_run.filename_glob)
@@ -1599,14 +1592,13 @@ class JLC_CaptionForgeOllamaCaption:
top_k=int(first_run.top_k),
repetition_penalty=float(repetition_penalty),
max_size=int(first_run.max_size),
forbidden_phrases=forbidden,
replacement_rules=replacements,
),
)
all_records: list[OllamaCaptionRecord] = []
forbidden = _parse_forbidden_lines(forbidden_phrases)
replacements = _parse_replace_pairs(replace_pairs)
def process_one(source_name: str, pil: Image.Image):
def process_one(source_name: str, image_key: str, pil: Image.Image):
for run in run_plan:
t0 = time.perf_counter()
raw_caption = _ollama_generate_caption(
@@ -1620,7 +1612,7 @@ class JLC_CaptionForgeOllamaCaption:
top_p=float(run.top_p),
top_k=int(run.top_k),
repetition_penalty=float(repetition_penalty),
seed=int(run.seed),
seed=run.seed,
max_size=int(run.max_size),
keep_loaded=bool(keep_loaded),
timeout=timeout,
@@ -1643,7 +1635,7 @@ class JLC_CaptionForgeOllamaCaption:
model_path=model_tag,
prompt=prompt,
system_prompt=system_prompt,
seed=int(run.seed),
seed=run.seed,
temperature=float(run.temperature),
top_p=float(run.top_p),
top_k=int(run.top_k),
@@ -1654,7 +1646,7 @@ class JLC_CaptionForgeOllamaCaption:
captionforge_pass="A",
model_family="ollama",
ensemble_run_index=int(run.ensemble_run_index),
image_key=source_name,
image_key=image_key,
backend="ollama",
status=status,
)
@@ -1662,20 +1654,16 @@ class JLC_CaptionForgeOllamaCaption:
if run_plan_connected and jsonl_path is not None and output_dir is not None:
_append_jsonl_records(jsonl_path, [record])
_write_text_sidecar(
_run_txt_path(output_dir, source_name, len(run_plan), int(run.ensemble_run_index)),
record.caption,
)
print(
f"[JLC CaptionForge Ollama Caption] Captioned {source_name} "
f"run {int(run.ensemble_run_index) + 1}/{len(run_plan)} via Planner key '{planner_key}'"
)
for source_name, pil in direct_images:
process_one(source_name, pil)
for source_name, image_key, pil in direct_images:
process_one(source_name, image_key, pil)
for source_name, path in file_images:
process_one(source_name, _open_image(path))
for source_name, image_key, path in file_images:
process_one(source_name, image_key, _open_image(path))
# Ollama model residency is requested through keep_alive on generation calls;
# Ollama's server/runtime policy ultimately decides residency.
@@ -8,11 +8,9 @@ JLC CaptionForge Qwen Caption — ComfyUI Node Wrapper
- Repository
https://github.com/Damkohler/CaptionForge
- CaptionForge focuses on practical dataset-captioning infrastructure for
LoRA dataset preparation, using multi-engine caption generation, JSONL
audit trails, claim extraction and refinement, text-LLM distillation,
image-aware VLM validation, and consensus-oriented caption improvement
to produce grounded, auditable training captions.
- CaptionForge 1.0 uses independent Pass-A witnesses, text-LLM synthesis,
image-aware validation, SHORT/TAGGY formatting, and JSONL audit trails to
produce grounded LoRA dataset captions.
- Node Purpose
- The **JLC CaptionForge Qwen Caption** node provides a ComfyUI
@@ -30,7 +28,6 @@ JLC CaptionForge Qwen Caption — ComfyUI Node Wrapper
• clear template-vs-custom prompt controls
• CaptionForge Pipeline Planner consumption through `pipeline_plan`
• CaptionForge Template Options consumption through `template_options`
• TXT audit sidecar writing in planned runs
• shared JSONL audit output in planned runs
• direct ComfyUI caption and resolved-prompt string outputs
• IMAGE and pipeline-plan passthrough for clean graph chaining
@@ -89,7 +86,6 @@ JLC CaptionForge Qwen Caption — ComfyUI Node Wrapper
• shared LoRA trigger word
- In planned mode, the node can write:
• TXT audit sidecar captions
• JSONL audit records
• run-configuration JSON files
@@ -133,12 +129,10 @@ JLC CaptionForge Qwen Caption — ComfyUI Node Wrapper
clean separation between the ComfyUI interface and the backend model
engine.
- ⚠️ Development Status
- This is early CaptionForge raw-caption infrastructure.
- The UI, model registry, prompt behavior, and output audit fields may
evolve as the multi-pass CaptionForge pipeline matures.
- The node is intended for local dataset preparation and controlled caption
audit workflows.
- Production Status
- Active CaptionForge 1.0 Pass-A witness node. In planned mode the Planner
owns run counts and sampling schedules; in standalone mode this node's
visible controls are authoritative.
- Attribution & License
- Concept and implementation by **J. L. Córdova**
@@ -167,7 +161,7 @@ MANIFEST = {
"the template_options input from the CaptionForge Template Options sidecar. "
"The UI separates caption_template_mode and custom_prompt_mode for clearer "
"prompt routing while delegating model loading, quantization, generation, "
"cleanup, TXT sidecars, and JSONL audit records to the Qwen caption engine."
"cleanup and JSONL audit records to the Qwen caption engine."
),
}
@@ -194,9 +188,11 @@ from ...engines.jlc_qwen_caption_engine import (
probe_registry_model_download,
timestamp,
write_run_config_json,
write_text_sidecar,
)
from ...engines.captionforge_pipeline_planner_engine import expand_captionforge_runs
from ...engines.captionforge_cleanup import resolve_cleanup_settings
from ...engines.captionforge_source_identity import file_source_identity, optional_image_identity
from ...engines.captionforge_dataset_export import is_dataset_export
from ...engines.captionforge_caption_prompt_kit import (
CAPTION_LENGTH_CHOICES,
CAPTION_TYPE_CHOICES,
@@ -291,20 +287,7 @@ def _tensor_to_pil(image_tensor) -> list[Image.Image]:
return images
def _safe_source_name(value: str) -> str:
cleaned = []
for ch in value.replace("\\", "/"):
if ch.isalnum() or ch in {"-", "_", "."}:
cleaned.append(ch)
elif ch == "/":
cleaned.append("__")
else:
cleaned.append("_")
out = "".join(cleaned).strip("._")
return out or "image"
def _iter_input_path_images(input_path: str, recursive: bool, filename_glob: str) -> list[tuple[str, Path]]:
def _iter_input_path_images(input_path: str, recursive: bool, filename_glob: str) -> list[tuple[str, str, Path]]:
root = Path(str(input_path or "").strip())
if not root:
return []
@@ -312,22 +295,26 @@ def _iter_input_path_images(input_path: str, recursive: bool, filename_glob: str
raise RuntimeError(f"CaptionForge input_path does not exist: {root}")
glob_text = (filename_glob or "*").strip() or "*"
explicit_dataset_source = is_dataset_export(root)
if root.is_file():
if root.suffix.lower() not in _SUPPORTED_IMAGE_SUFFIXES:
raise RuntimeError(f"CaptionForge input_path is not a supported image file: {root}")
return [(_safe_source_name(root.stem), root)]
image_name, image_key = file_source_identity(root, root)
return [(image_name, image_key, root)]
pattern_iter = root.rglob(glob_text) if recursive else root.glob(glob_text)
paths = sorted(
p for p in pattern_iter
if p.is_file() and p.suffix.lower() in _SUPPORTED_IMAGE_SUFFIXES
if p.is_file()
and p.suffix.lower() in _SUPPORTED_IMAGE_SUFFIXES
and (explicit_dataset_source or not is_dataset_export(p))
)
items: list[tuple[str, Path]] = []
items: list[tuple[str, str, Path]] = []
for path in paths:
rel = path.relative_to(root).with_suffix("")
items.append((_safe_source_name(str(rel)), path))
image_name, image_key = file_source_identity(path, root)
items.append((image_name, image_key, path))
return items
@@ -364,12 +351,6 @@ def _planned_caption_jsonl_path(pipeline_plan) -> Path | None:
return None
def _run_txt_path(output_dir: Path, source_name: str, run_count: int, run_index: int) -> Path:
if run_count <= 1:
return output_dir / f"{source_name}.txt"
return output_dir / f"{source_name}__cf_run_{run_index:02d}.txt"
def _use_template_mode(caption_template_mode: bool, custom_prompt_mode: bool) -> bool:
if bool(custom_prompt_mode):
return False
@@ -689,7 +670,7 @@ class JLC_CaptionForgeQwen:
"resolved_prompt",
)
FUNCTION = "caption"
CATEGORY = "Captioning/CaptionForge/Captioning Nodes"
CATEGORY = "Caption/CaptionForge/Caption Nodes"
@classmethod
def IS_CHANGED(cls, **kwargs):
@@ -741,7 +722,7 @@ class JLC_CaptionForgeQwen:
if download_probe_only:
result = probe_registry_model_download(model, JLC_QWEN_MODEL_ROOT)
return (image, pipeline_plan, result, resolved_prompt)
return (image, pipeline_plan, template_options, result, resolved_prompt)
effective_seed = -1 if seed is None else int(seed)
@@ -765,9 +746,12 @@ class JLC_CaptionForgeQwen:
if run_plan_connected and not run_plan:
status = "[CaptionForge] Qwen disabled by Pipeline Planner."
print(status)
return (image, pipeline_plan, status, resolved_prompt)
return (image, pipeline_plan, template_options, status, resolved_prompt)
first_run = run_plan[0]
effective_forbidden, effective_replacements = resolve_cleanup_settings(
_normalize_pipeline_plan(pipeline_plan), forbidden_phrases, replace_pairs
)
qwen_quantization_value = "bnb_8bit" if qwen_quantization == "Balanced (8-bit)" else "none"
generation = GenerationConfig(
@@ -783,8 +767,8 @@ class JLC_CaptionForgeQwen:
trigger="",
prefix=(f"{first_run.trigger_word}," if first_run.trigger_word else ""),
suffix="",
forbidden_phrases=_parse_forbidden_lines(forbidden_phrases),
replacement_rules=_parse_replace_pairs(replace_pairs),
forbidden_phrases=effective_forbidden,
replacement_rules=effective_replacements,
)
qwen_config = QwenCaptionConfig(
@@ -807,8 +791,8 @@ class JLC_CaptionForgeQwen:
engine = QwenCaptionEngine(config=qwen_config, generation=generation, cleanup=cleanup)
direct_images = [(f"comfy_image_{i:04d}", pil) for i, pil in enumerate(_tensor_to_pil(image))]
file_images: list[tuple[str, Path]] = []
direct_images = [(*optional_image_identity(i), pil) for i, pil in enumerate(_tensor_to_pil(image))]
file_images: list[tuple[str, str, Path]] = []
if first_run.input_path:
file_images = _iter_input_path_images(first_run.input_path, first_run.recursive, first_run.filename_glob)
@@ -846,7 +830,7 @@ class JLC_CaptionForgeQwen:
dry_run=False,
)
def process_one(source_name: str, pil: Image.Image):
def process_one(source_name: str, image_key: str, pil: Image.Image):
for run in run_plan:
engine.generation = GenerationConfig(
max_new_tokens=int(run.max_new_tokens),
@@ -860,8 +844,8 @@ class JLC_CaptionForgeQwen:
trigger="",
prefix=(f"{run.trigger_word}," if run.trigger_word else ""),
suffix="",
forbidden_phrases=_parse_forbidden_lines(forbidden_phrases),
replacement_rules=_parse_replace_pairs(replace_pairs),
forbidden_phrases=effective_forbidden,
replacement_rules=effective_replacements,
)
engine.config.max_size = int(run.max_size)
@@ -887,29 +871,22 @@ class JLC_CaptionForgeQwen:
captionforge_pass="A",
model_family="qwen",
ensemble_run_index=run.ensemble_run_index,
image_key=source_name,
image_key=image_key,
)
all_records.append(record)
if run_plan_connected and jsonl_path is not None and output_dir is not None:
append_jsonl_records(jsonl_path, [record], dry_run=False)
write_text_sidecar(
_run_txt_path(output_dir, source_name, len(run_plan), run.ensemble_run_index),
record.caption,
overwrite=True,
backup_existing=False,
dry_run=False,
)
print(
f"[JLC CaptionForge Qwen] Captioned {source_name} "
f"run {run.ensemble_run_index + 1}/{len(run_plan)}"
)
for source_name, pil in direct_images:
process_one(source_name, pil)
for source_name, image_key, pil in direct_images:
process_one(source_name, image_key, pil)
for source_name, path in file_images:
process_one(source_name, _open_image(path))
for source_name, image_key, path in file_images:
process_one(source_name, image_key, _open_image(path))
if not keep_loaded:
engine.unload()
+8 -10
View File
@@ -8,11 +8,9 @@ CaptionForge Ollama Model Dropdown Helper
- Repository:
https://github.com/Damkohler/CaptionForge
- CaptionForge focuses on practical dataset-captioning infrastructure for
LoRA dataset preparation, using multi-engine caption generation, JSONL
audit trails, claim extraction and refinement, text-LLM distillation,
image-aware VLM validation, and consensus-oriented caption improvement
to produce grounded, auditable training captions.
- CaptionForge 1.0 uses independent Pass-A witnesses, text-LLM synthesis,
image-aware validation, SHORT/TAGGY formatting, and JSONL audit trails to
produce grounded LoRA dataset captions.
- Helper Purpose
- The **CaptionForge Ollama Model Dropdown Helper** centralizes loading of
@@ -57,17 +55,17 @@ CaptionForge Ollama Model Dropdown Helper
- CaptionForge should avoid duplicated dropdown constants spread across
multiple node wrappers.
- A single helper makes pre-release cleanup easier, keeps ComfyUI widgets
- A single helper keeps release configuration centralized, keeps ComfyUI widgets
consistent, and allows users to edit one JSON file rather than patching
several Python files.
- The helper favors explicitness, reproducibility, and predictable fallback
behavior over clever model discovery.
- Development Status
- CaptionForge v0.1.0 experimental developer-preview infrastructure.
- This helper is active shared infrastructure, but the exact JSON schema may
evolve before the first stable CaptionForge release.
- Production Status
- Active CaptionForge 1.0 configuration support. It reads concrete Ollama
model tags for Pass-A captioning and the production B/C/D stages without
performing inference itself.
- Attribution & License
- Concept and implementation by **J. L. Córdova**
@@ -1,105 +0,0 @@
from __future__ import annotations
MANIFEST={"name":"JLC BLIP Caption","version":(0,1,0),"author":"J. L. Córdova"}
import json, time
from datetime import datetime
from pathlib import Path
import numpy as np, torch
from PIL import Image
import folder_paths
from ..engines.jlc_blip_caption_engine import BatchCaptionConfig, CaptionRecord, CleanupConfig, BlipCaptionConfig, BlipCaptionEngine, GenerationConfig, MEMORY_MODES, MODEL_REGISTRY, PROMPT_PRESETS, append_jsonl_records, load_existing_jsonl_images, probe_registry_model_download, record_to_json, resolve_prompt, timestamp, write_run_config_json, write_text_sidecar
from ..engines.captionforge_pipeline_planner_engine import expand_captionforge_runs
JLC_BLIP_MODEL_ROOT=Path(folder_paths.models_dir)/"LLM"/"JLC_BLIPCaption"
def _tensor_to_pil(t):
if t is None: return []
if isinstance(t,torch.Tensor): t=t.detach().cpu()
if t.ndim==3: t=t.unsqueeze(0)
out=[]
for img in t:
arr=np.clip(img.numpy()*255.0,0,255).astype(np.uint8); out.append(Image.fromarray(arr).convert("RGB"))
return out
def _lines(v): return [x.strip() for x in (v or "").splitlines() if x.strip()]
def _pairs(v):
out=[]
for line in (v or "").splitlines():
line=line.strip()
if line and not line.startswith("#") and "=>" in line:
a,b=line.split("=>",1); a=a.strip(); b=b.strip()
if a: out.append((a,b))
return out
def _norm_plan(c):
if c is None: return {}
if isinstance(c,dict): return dict(c)
if isinstance(c,str):
try: o=json.loads(c.strip()); return o if isinstance(o,dict) else {}
except Exception: return {}
return {}
def _planned_jsonl(c):
paths=_norm_plan(c).get("paths") if isinstance(_norm_plan(c).get("paths"),dict) else {}
for k in ("caption_jsonl","pass_a_jsonl"):
v=str(paths.get(k) or "").strip()
if v: return Path(v)
return None
def _jsonl(records): return "\n".join(json.dumps(record_to_json(r), ensure_ascii=False) for r in records)
class JLC_BLIPCaption:
@classmethod
def INPUT_TYPES(cls):
return {"required":{
"model":(list(MODEL_REGISTRY.keys()),{"default":"BLIP-2 FLAN-T5 XL"}),
"memory_mode":(list(MEMORY_MODES.keys()),{"default":"Default"}),
"input_path":("STRING",{"default":"","multiline":False}), "recursive":("BOOLEAN",{"default":True}), "filename_glob":("STRING",{"default":"*","multiline":False}),
"prompt_preset":(list(PROMPT_PRESETS.keys()),{"default":"detailed"}), "custom_prompt":("STRING",{"default":"","multiline":True}),
"max_new_tokens":("INT",{"default":256,"min":16,"max":2048,"step":8}), "num_beams":("INT",{"default":3,"min":1,"max":16,"step":1}),
"temperature":("FLOAT",{"default":0.0,"min":0.0,"max":2.0,"step":0.01}), "top_p":("FLOAT",{"default":0.90,"min":0.0,"max":1.0,"step":0.01}), "top_k":("INT",{"default":50,"min":0,"max":500,"step":1}),
"repetition_penalty":("FLOAT",{"default":1.0,"min":1.0,"max":2.0,"step":0.01}), "captions_per_image":("INT",{"default":1,"min":1,"max":100,"step":1}), "seed":("INT",{"default":-1,"min":-1,"max":0xFFFFFFFF,"step":1}),
"max_size":("INT",{"default":768,"min":0,"max":4096,"step":64}), "output_dir":("STRING",{"default":"","multiline":False}),
"write_txt":("BOOLEAN",{"default":True}), "write_jsonl":("BOOLEAN",{"default":False}), "also_jsonl":("BOOLEAN",{"default":False}), "write_run_config":("BOOLEAN",{"default":True}), "jsonl_filename":("STRING",{"default":"captions.jsonl","multiline":False}),
"overwrite":("BOOLEAN",{"default":False}), "backup_existing":("BOOLEAN",{"default":True}), "dry_run":("BOOLEAN",{"default":False}), "limit":("INT",{"default":0,"min":0,"max":100000,"step":1}),
"skip_existing_txt":("BOOLEAN",{"default":True}), "skip_existing_jsonl_images":("BOOLEAN",{"default":False}), "prefix":("STRING",{"default":"","multiline":False}), "suffix":("STRING",{"default":"","multiline":False}),
"forbidden_phrases":("STRING",{"default":"","multiline":True}), "replace_pairs":("STRING",{"default":"","multiline":True}), "keep_loaded":("BOOLEAN",{"default":True}), "download_probe_only":("BOOLEAN",{"default":False}),
},"optional":{"image":("IMAGE",{}),"captionforge_run_config":("CAPTIONFORGE_PIPELINE_PLAN",{})}}
RETURN_TYPES=("CAPTIONFORGE_PIPELINE_PLAN","STRING","STRING","STRING"); RETURN_NAMES=("captionforge_run_config_out","caption","jsonl_records","resolved_prompt"); FUNCTION="caption"; CATEGORY="JLC/Captioning"
@classmethod
def IS_CHANGED(cls, **kwargs): return float("NaN")
def caption(self, model, memory_mode, input_path, recursive, filename_glob, prompt_preset, custom_prompt, max_new_tokens, num_beams, temperature, top_p, top_k, repetition_penalty, captions_per_image, seed, max_size, output_dir, write_txt, write_jsonl, also_jsonl, write_run_config, jsonl_filename, overwrite, backup_existing, dry_run, limit, skip_existing_txt, skip_existing_jsonl_images, prefix, suffix, forbidden_phrases, replace_pairs, keep_loaded, download_probe_only, image=None, captionforge_run_config=None):
if download_probe_only: return (captionforge_run_config, probe_registry_model_download(model,JLC_BLIP_MODEL_ROOT), "", "")
prompt=resolve_prompt(prompt_preset, custom_prompt)
runs=expand_captionforge_runs(captionforge_run_config, model_key="blip", widget_captions_per_image=int(captions_per_image), widget_seed=int(seed), widget_temperature=float(temperature), widget_top_p=float(top_p), widget_top_k=int(top_k), widget_max_new_tokens=int(max_new_tokens), widget_max_size=int(max_size), widget_trigger_word="", widget_output_dir=output_dir, widget_input_path=input_path, widget_recursive=bool(recursive), widget_filename_glob=filename_glob)
planned=bool(captionforge_run_config)
if planned and not runs:
status="[CaptionForge] BLIP Caption disabled by Pipeline Planner."; print(status); return (captionforge_run_config,status,"",prompt)
first=runs[0]
if planned: write_jsonl=True; also_jsonl=False; skip_existing_txt=False; skip_existing_jsonl_images=False; overwrite=True
outdir=(first.output_dir or output_dir or "").strip(); inpath=(first.input_path or input_path or "").strip(); jsonl_filename=(jsonl_filename or "captions.jsonl").strip() or "captions.jsonl"
planned_jsonl=_planned_jsonl(captionforge_run_config) if planned else None
if planned_jsonl: outdir=str(planned_jsonl.parent); jsonl_filename=planned_jsonl.name
eff_prefix=f"{first.trigger_word}, {prefix}".strip(", ") if first.trigger_word else prefix
engine=BlipCaptionEngine(BlipCaptionConfig(model, "", str(JLC_BLIP_MODEL_ROOT), memory_mode, "bf16", "auto", False, bool(keep_loaded), "", True, int(first.max_size), prompt, True, True), GenerationConfig(int(first.max_new_tokens), int(num_beams), float(first.temperature), float(first.top_p), int(first.top_k), float(repetition_penalty), first.seed), CleanupConfig("",eff_prefix,suffix,_lines(forbidden_phrases),_pairs(replace_pairs)))
all_records=[]; pil_images=_tensor_to_pil(image); batch_result=None; use_jsonl=bool(write_jsonl or also_jsonl)
if pil_images:
engine.load(); image_out=Path(outdir) if outdir else Path(folder_paths.get_output_directory())/"jlc_blip_caption"; image_out.mkdir(parents=True,exist_ok=True); jsonl_path=image_out/jsonl_filename
if write_run_config and not inpath: write_run_config_json(image_out/f"jlc_blip_caption_run_config_{timestamp()}.json", engine.build_run_config(), bool(dry_run))
seen=load_existing_jsonl_images(jsonl_path) if use_jsonl and skip_existing_jsonl_images else set()
for idx,pil in enumerate(pil_images):
src=f"comfy_image_{idx:04d}"; txt=image_out/f"{src}.txt"
if len(runs)==1 and skip_existing_txt and write_txt and txt.exists() and not overwrite: print(f"[JLC BLIP Caption] Skipping existing TXT: {txt}"); continue
if use_jsonl and skip_existing_jsonl_images and src in seen: print(f"[JLC BLIP Caption] Skipping existing JSONL image: {src}"); continue
for r in runs:
engine.generation=GenerationConfig(int(r.max_new_tokens),int(num_beams),float(r.temperature),float(r.top_p),int(r.top_k),float(repetition_penalty),r.seed); engine.config.max_size=int(r.max_size); engine.config.prompt=prompt
engine.cleanup=CleanupConfig("", f"{r.trigger_word}, {prefix}".strip(", ") if r.trigger_word else prefix, suffix, _lines(forbidden_phrases), _pairs(replace_pairs))
run_txt=txt if len(runs)<=1 else image_out/f"{src}__cf_run_{r.ensemble_run_index:02d}.txt"
if skip_existing_txt and write_txt and run_txt.exists() and not overwrite: print(f"[JLC BLIP Caption] Skipping existing TXT: {run_txt}"); continue
t0=time.perf_counter(); final,raw=engine.caption_pil(pil); dt=time.perf_counter()-t0; print(f"[JLC BLIP Caption] Generation time run {r.ensemble_run_index}: {dt:.2f}s")
rec=CaptionRecord(src,final,raw,model,str(engine.local_model_path or ""),prompt,r.seed,r.temperature,r.top_p,r.top_k,int(num_beams),r.max_new_tokens,int(r.max_size),datetime.now().isoformat(timespec="seconds"),captionforge_pass="A",model_family="blip",ensemble_run_index=r.ensemble_run_index,image_key=src)
all_records.append(rec)
if use_jsonl: append_jsonl_records(jsonl_path,[rec],bool(dry_run))
if write_txt: write_text_sidecar(run_txt,rec.caption,bool(overwrite),bool(backup_existing),bool(dry_run))
print(f"[JLC BLIP Caption] Captioned IMAGE {idx+1}/{len(pil_images)} run {r.ensemble_run_index+1}/{len(runs)}: {src}")
if inpath:
batch=BatchCaptionConfig(inpath,bool(first.recursive if planned else recursive),first.filename_glob if planned else ((filename_glob or "*").strip() or "*"),output_dir=outdir,write_txt=bool(write_txt),write_jsonl=bool(write_jsonl),jsonl_filename=jsonl_filename,write_run_config=bool(write_run_config),overwrite=bool(overwrite),backup_existing=bool(backup_existing),dry_run=bool(dry_run),limit=int(limit),skip_existing_txt=bool(skip_existing_txt),skip_existing_jsonl_images=bool(skip_existing_jsonl_images))
batch_result=engine.caption_batch(batch); all_records.extend(batch_result.records)
if not pil_images and batch_result is None: raise RuntimeError("No image input found. Connect an IMAGE input or provide input_path pointing to an image file or folder.")
if not keep_loaded: engine.unload()
return (captionforge_run_config,"\n\n".join(r.caption for r in all_records if r.status=="ok"),_jsonl(all_records),prompt)
NODE_CLASS_MAPPINGS={"JLC_BLIPCaption":JLC_BLIPCaption}
NODE_DISPLAY_NAME_MAPPINGS={"JLC_BLIPCaption":"\u2003JLC BLIP Caption"}
@@ -1,84 +0,0 @@
from __future__ import annotations
MANIFEST={"name":"JLC BLIP Caption (Lite)","version":(0,1,0),"author":"J. L. Córdova"}
import json,time
from datetime import datetime
from pathlib import Path
import numpy as np, torch
from PIL import Image
import folder_paths
from ..engines.jlc_blip_caption_engine import CaptionRecord, CleanupConfig, BlipCaptionConfig, BlipCaptionEngine, GenerationConfig, MEMORY_MODES, MODEL_REGISTRY, PROMPT_PRESETS, append_jsonl_records, record_to_json, resolve_prompt, timestamp, write_run_config_json, write_text_sidecar
from ..engines.captionforge_pipeline_planner_engine import expand_captionforge_runs
JLC_BLIP_MODEL_ROOT=Path(folder_paths.models_dir)/"LLM"/"JLC_BLIPCaption"; DEFAULT_JSONL_FILENAME="captions.jsonl"; _SUF={".png",".jpg",".jpeg",".webp",".bmp",".tif",".tiff"}
def _tensor_to_pil(t):
if t is None: return []
if isinstance(t,torch.Tensor): t=t.detach().cpu()
if t.ndim==3: t=t.unsqueeze(0)
return [Image.fromarray(np.clip(img.numpy()*255.0,0,255).astype(np.uint8)).convert("RGB") for img in t]
def _safe(v):
s=[]
for ch in v.replace("\\","/"):
if ch.isalnum() or ch in {"-","_","."}: s.append(ch)
elif ch=="/": s.append("__")
else: s.append("_")
return "".join(s).strip("._") or "image"
def _iter(path, recursive, glob):
root=Path(str(path or "").strip())
if not root: return []
if not root.exists(): raise RuntimeError(f"CaptionForge input_path does not exist: {root}")
if root.is_file(): return [(_safe(root.stem),root)] if root.suffix.lower() in _SUF else []
paths=sorted(p for p in (root.rglob(glob or "*") if recursive else root.glob(glob or "*")) if p.is_file() and p.suffix.lower() in _SUF)
return [(_safe(str(p.relative_to(root).with_suffix(""))),p) for p in paths]
def _open(p):
with Image.open(p) as im: return im.convert("RGB")
def _norm(c):
if c is None: return {}
if isinstance(c,dict): return dict(c)
if isinstance(c,str):
try: o=json.loads(c.strip()); return o if isinstance(o,dict) else {}
except Exception: return {}
return {}
def _planned(c):
paths=_norm(c).get("paths") if isinstance(_norm(c).get("paths"),dict) else {}
for k in ("caption_jsonl","pass_a_jsonl"):
v=str(paths.get(k) or "").strip()
if v: return Path(v)
return None
def _jsonl(records): return "\n".join(json.dumps(record_to_json(r), ensure_ascii=False) for r in records)
def _txt(out,src,n,i): return out/f"{src}.txt" if n<=1 else out/f"{src}__cf_run_{i:02d}.txt"
class JLC_BLIPCaptionLite:
@classmethod
def INPUT_TYPES(cls):
return {"required":{"model":(list(MODEL_REGISTRY.keys()),{"default":"BLIP-2 FLAN-T5 XL"}),"memory_mode":(list(MEMORY_MODES.keys()),{"default":"Default"}),"keep_loaded":("BOOLEAN",{"default":True}),"prompt_preset":(list(PROMPT_PRESETS.keys()),{"default":"detailed"}),"custom_prompt":("STRING",{"default":"","multiline":True}),"max_new_tokens":("INT",{"default":256,"min":16,"max":2048,"step":8}),"num_beams":("INT",{"default":3,"min":1,"max":16,"step":1}),"temperature":("FLOAT",{"default":0.0,"min":0.0,"max":2.0,"step":0.01}),"top_p":("FLOAT",{"default":0.90,"min":0.0,"max":1.0,"step":0.01}),"top_k":("INT",{"default":50,"min":0,"max":500,"step":1})},"optional":{"image":("IMAGE",{}),"captionforge_run_config":("CAPTIONFORGE_PIPELINE_PLAN",{})}}
RETURN_TYPES=("CAPTIONFORGE_PIPELINE_PLAN","STRING","STRING","STRING"); RETURN_NAMES=("captionforge_run_config_out","caption","jsonl_records","resolved_prompt"); FUNCTION="caption"; CATEGORY="JLC/Captioning"
@classmethod
def IS_CHANGED(cls, **kwargs): return float("NaN")
def caption(self, model, memory_mode, keep_loaded, prompt_preset, custom_prompt, max_new_tokens, num_beams, temperature, top_p, top_k, image=None, captionforge_run_config=None):
planned=bool(captionforge_run_config); prompt=resolve_prompt(prompt_preset, custom_prompt)
runs=expand_captionforge_runs(captionforge_run_config, model_key="blip", widget_captions_per_image=1, widget_seed=-1, widget_temperature=float(temperature), widget_top_p=float(top_p), widget_top_k=int(top_k), widget_max_new_tokens=int(max_new_tokens), widget_max_size=768, widget_trigger_word="", widget_output_dir="", widget_input_path="", widget_recursive=True, widget_filename_glob="*")
if planned and not runs:
status="[CaptionForge] BLIP Caption Lite disabled by Pipeline Planner."; print(status); return (captionforge_run_config,status,"",prompt)
first=runs[0]
engine=BlipCaptionEngine(BlipCaptionConfig(model,"",str(JLC_BLIP_MODEL_ROOT),memory_mode,"bf16","auto",False,bool(keep_loaded),"",True,int(first.max_size),prompt,True,True),GenerationConfig(int(first.max_new_tokens),int(num_beams),float(first.temperature),float(first.top_p),int(first.top_k),1.0,first.seed),CleanupConfig("",f"{first.trigger_word}," if first.trigger_word else ""))
direct=[(f"comfy_image_{i:04d}",pil) for i,pil in enumerate(_tensor_to_pil(image))]; files=[]
if first.input_path: files=_iter(first.input_path, first.recursive, first.filename_glob)
if not direct and not files: raise RuntimeError("No image input found. Connect an IMAGE input or provide input_path in the CaptionForge Run Plan.")
pj=_planned(captionforge_run_config) if planned else None; out=Path(first.output_dir) if first.output_dir else Path(folder_paths.get_output_directory())/"jlc_blip_caption_lite"
if pj: out=pj.parent
if planned: out.mkdir(parents=True,exist_ok=True)
jsonl_path=pj or (out/DEFAULT_JSONL_FILENAME); all_records=[]; engine.load()
if planned: write_run_config_json(out/f"jlc_blip_caption_lite_run_config_{timestamp()}.json", engine.build_run_config(), False)
def process(src,pil):
for r in runs:
engine.generation=GenerationConfig(int(r.max_new_tokens),int(num_beams),float(r.temperature),float(r.top_p),int(r.top_k),1.0,r.seed); engine.cleanup=CleanupConfig("",f"{r.trigger_word}," if r.trigger_word else ""); engine.config.max_size=int(r.max_size); engine.config.prompt=prompt
t0=time.perf_counter(); final,raw=engine.caption_pil(pil); dt=time.perf_counter()-t0; print(f"[JLC BLIP Caption Lite] Generation time run {r.ensemble_run_index}: {dt:.2f}s")
rec=CaptionRecord(src,final,raw,model,str(engine.local_model_path or ""),prompt,r.seed,r.temperature,r.top_p,r.top_k,int(num_beams),r.max_new_tokens,int(r.max_size),datetime.now().isoformat(timespec="seconds"),captionforge_pass="A",model_family="blip",ensemble_run_index=r.ensemble_run_index,image_key=src)
all_records.append(rec)
if planned: append_jsonl_records(jsonl_path,[rec],False); write_text_sidecar(_txt(out,src,len(runs),r.ensemble_run_index),rec.caption,True,False,False)
print(f"[JLC BLIP Caption Lite] Captioned {src} run {r.ensemble_run_index+1}/{len(runs)}")
for src,pil in direct: process(src,pil)
for src,path in files: process(src,_open(path))
if not keep_loaded: engine.unload()
return (captionforge_run_config,"\n\n".join(r.caption for r in all_records if r.status=="ok"),_jsonl(all_records),prompt)
NODE_CLASS_MAPPINGS={"JLC_BLIPCaptionLite":JLC_BLIPCaptionLite}
NODE_DISPLAY_NAME_MAPPINGS={"JLC_BLIPCaptionLite":"\u2003JLC BLIP Caption (Lite)"}
@@ -1,565 +0,0 @@
"""
JLC Florence Caption — ComfyUI Node Wrapper
Drop target:
CaptionForge/nodes/jlc_florence_caption_CUI_node.py
Full Florence-2 CaptionForge Pass A node.
"""
from __future__ import annotations
MANIFEST = {
"name": "JLC Florence Caption",
"version": (0, 1, 0),
"author": "J. L. Córdova",
"description": (
"ComfyUI frontend for Florence-2-family CaptionForge Pass A captioning. Provides "
"model/task controls, direct IMAGE and file/folder routing, TXT/JSONL/run-config "
"outputs, cleanup rules, download probe, Pipeline Planner consumption, and delegates "
"model loading/generation/cache behavior to jlc_florence_caption_engine.py."
),
}
import json
from datetime import datetime
from pathlib import Path
import time
import numpy as np
import torch
from PIL import Image
import folder_paths
from ..engines.jlc_florence_caption_engine import (
BatchCaptionConfig,
CaptionRecord,
CleanupConfig,
FlorenceCaptionConfig,
FlorenceCaptionEngine,
GenerationConfig,
MEMORY_MODES,
MODEL_REGISTRY,
TASK_PROMPTS,
append_jsonl_records,
load_existing_jsonl_images,
probe_registry_model_download,
record_to_json,
resolve_task_prompt,
timestamp,
write_run_config_json,
write_text_sidecar,
)
from ..engines.captionforge_pipeline_planner_engine import expand_captionforge_runs
JLC_FLORENCE_MODEL_ROOT = Path(folder_paths.models_dir) / "LLM" / "JLC_FlorenceCaption"
def _tensor_to_pil(image_tensor) -> list[Image.Image]:
if image_tensor is None:
return []
if isinstance(image_tensor, torch.Tensor):
image_tensor = image_tensor.detach().cpu()
if image_tensor.ndim == 3:
image_tensor = image_tensor.unsqueeze(0)
images: list[Image.Image] = []
for img in image_tensor:
arr = img.numpy()
arr = np.clip(arr * 255.0, 0, 255).astype(np.uint8)
images.append(Image.fromarray(arr).convert("RGB"))
return images
def _parse_forbidden_lines(value: str) -> list[str]:
value = (value or "").strip()
if not value:
return []
return [line.strip() for line in value.splitlines() if line.strip()]
def _parse_replace_pairs(value: str) -> list[tuple[str, str]]:
value = (value or "").strip()
if not value:
return []
rules: list[tuple[str, str]] = []
for line in value.splitlines():
line = line.strip()
if not line or line.startswith("#") or "=>" not in line:
continue
old, new = line.split("=>", 1)
old = old.strip()
new = new.strip()
if old:
rules.append((old, new))
return rules
def _normalize_pipeline_plan(config):
if config is None:
return {}
if isinstance(config, dict):
return dict(config)
if isinstance(config, str):
text = config.strip()
if not text:
return {}
try:
obj = json.loads(text)
return obj if isinstance(obj, dict) else {}
except Exception:
return {}
return {}
def _planned_caption_jsonl_path(captionforge_run_config) -> Path | None:
cfg = _normalize_pipeline_plan(captionforge_run_config)
paths = cfg.get("paths") if isinstance(cfg.get("paths"), dict) else {}
for key in ("caption_jsonl", "pass_a_jsonl"):
value = str(paths.get(key) or "").strip()
if value:
return Path(value)
return None
def _make_jsonl_string(records: list[CaptionRecord]) -> str:
return "\n".join(json.dumps(record_to_json(r), ensure_ascii=False) for r in records)
class JLC_FlorenceCaption:
"""Full Florence-2 captioner for CaptionForge."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": (
list(MODEL_REGISTRY.keys()),
{
"default": "Florence-2-base-ft",
"tooltip": "Select Florence-2 model. Models live under ComfyUI/models/LLM/JLC_FlorenceCaption/.",
},
),
"memory_mode": (
list(MEMORY_MODES.keys()),
{
"default": "Default",
"tooltip": "Florence v1 uses normal Transformers loading; no bitsandbytes path is exposed.",
},
),
"input_path": (
"STRING",
{"default": "", "multiline": False, "tooltip": "Optional standalone image file or folder path."},
),
"recursive": (
"BOOLEAN",
{"default": True, "tooltip": "Standalone folder mode: search subfolders."},
),
"filename_glob": (
"STRING",
{"default": "*", "multiline": False, "tooltip": "Standalone folder mode filename filter."},
),
"task": (
list(TASK_PROMPTS.keys()),
{"default": "Detailed Caption", "tooltip": "Florence caption task."},
),
"custom_task_prompt": (
"STRING",
{
"default": "",
"multiline": False,
"tooltip": "Optional raw Florence task prompt override, e.g. <MORE_DETAILED_CAPTION>.",
},
),
"max_new_tokens": (
"INT",
{"default": 384, "min": 16, "max": 2048, "step": 8, "tooltip": "Max generated tokens. Planner overrides when connected."},
),
"num_beams": (
"INT",
{"default": 3, "min": 1, "max": 16, "step": 1, "tooltip": "Beam count for Florence generation."},
),
"temperature": (
"FLOAT",
{"default": 0.0, "min": 0.0, "max": 2.0, "step": 0.01, "tooltip": "Sampling temperature. 0.0 is deterministic."},
),
"top_p": (
"FLOAT",
{"default": 0.90, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Top-p when sampling."},
),
"top_k": (
"INT",
{"default": 50, "min": 0, "max": 500, "step": 1, "tooltip": "Top-k when sampling."},
),
"captions_per_image": (
"INT",
{"default": 1, "min": 1, "max": 100, "step": 1, "tooltip": "Standalone captions per image. Planner overrides when connected."},
),
"seed": (
"INT",
{"default": -1, "min": -1, "max": 0xFFFFFFFF, "step": 1, "tooltip": "Standalone seed. -1 means unseeded/planner-style behavior."},
),
"max_size": (
"INT",
{"default": 1024, "min": 0, "max": 4096, "step": 64, "tooltip": "Maximum longest-side image size."},
),
"output_dir": (
"STRING",
{"default": "", "multiline": False, "tooltip": "Standalone output folder. Planner overrides when connected."},
),
"write_txt": (
"BOOLEAN",
{"default": True, "tooltip": "Write TXT sidecar captions."},
),
"write_jsonl": (
"BOOLEAN",
{"default": False, "tooltip": "Append JSONL evidence records."},
),
"also_jsonl": (
"BOOLEAN",
{"default": False, "tooltip": "Standalone compatibility toggle for JSONL in addition to TXT."},
),
"write_run_config": (
"BOOLEAN",
{"default": True, "tooltip": "Write run-config JSON."},
),
"jsonl_filename": (
"STRING",
{"default": "captions.jsonl", "multiline": False, "tooltip": "JSONL evidence filename."},
),
"overwrite": (
"BOOLEAN",
{"default": False, "tooltip": "Allow overwriting TXT sidecars."},
),
"backup_existing": (
"BOOLEAN",
{"default": True, "tooltip": "Backup existing TXT before overwrite."},
),
"dry_run": (
"BOOLEAN",
{"default": False, "tooltip": "Run logic without writing files."},
),
"limit": (
"INT",
{"default": 0, "min": 0, "max": 100000, "step": 1, "tooltip": "Standalone folder image limit; 0 means no limit."},
),
"skip_existing_txt": (
"BOOLEAN",
{"default": True, "tooltip": "Skip images whose TXT exists unless overwrite is enabled."},
),
"skip_existing_jsonl_images": (
"BOOLEAN",
{"default": False, "tooltip": "Skip images already present in JSONL."},
),
"prefix": (
"STRING",
{"default": "", "multiline": False, "tooltip": "Optional final caption prefix."},
),
"suffix": (
"STRING",
{"default": "", "multiline": False, "tooltip": "Optional final caption suffix."},
),
"forbidden_phrases": (
"STRING",
{"default": "", "multiline": True, "tooltip": "Phrases to remove, one per line."},
),
"replace_pairs": (
"STRING",
{"default": "", "multiline": True, "tooltip": "Search-replace rules, one per line, old=>new."},
),
"keep_loaded": (
"BOOLEAN",
{"default": True, "tooltip": "Keep Florence cached after run."},
),
"download_probe_only": (
"BOOLEAN",
{"default": False, "tooltip": "Download metadata/config only; skip large weights and do not caption."},
),
},
"optional": {
"image": ("IMAGE", {"tooltip": "Optional ComfyUI IMAGE input."}),
"captionforge_run_config": (
"CAPTIONFORGE_PIPELINE_PLAN",
{"tooltip": "Optional shared CaptionForge Pipeline Planner."},
),
},
}
RETURN_TYPES = ("CAPTIONFORGE_PIPELINE_PLAN", "STRING", "STRING", "STRING")
RETURN_NAMES = ("captionforge_run_config_out", "caption", "jsonl_records", "resolved_prompt")
FUNCTION = "caption"
CATEGORY = "JLC/Captioning"
@classmethod
def IS_CHANGED(cls, **kwargs):
return float("NaN")
def caption(
self,
model,
memory_mode,
input_path,
recursive,
filename_glob,
task,
custom_task_prompt,
max_new_tokens,
num_beams,
temperature,
top_p,
top_k,
captions_per_image,
seed,
max_size,
output_dir,
write_txt,
write_jsonl,
also_jsonl,
write_run_config,
jsonl_filename,
overwrite,
backup_existing,
dry_run,
limit,
skip_existing_txt,
skip_existing_jsonl_images,
prefix,
suffix,
forbidden_phrases,
replace_pairs,
keep_loaded,
download_probe_only,
image=None,
captionforge_run_config=None,
):
if download_probe_only:
result = probe_registry_model_download(model, JLC_FLORENCE_MODEL_ROOT)
return (captionforge_run_config, result, "", "")
task_prompt = resolve_task_prompt(task, custom_task_prompt)
run_plan = expand_captionforge_runs(
captionforge_run_config,
model_key="florence",
widget_captions_per_image=int(captions_per_image),
widget_seed=int(seed),
widget_temperature=float(temperature),
widget_top_p=float(top_p),
widget_top_k=int(top_k),
widget_max_new_tokens=int(max_new_tokens),
widget_max_size=int(max_size),
widget_trigger_word="",
widget_output_dir=output_dir,
widget_input_path=input_path,
widget_recursive=bool(recursive),
widget_filename_glob=filename_glob,
)
run_plan_connected = bool(captionforge_run_config)
if run_plan_connected and not run_plan:
status = "[CaptionForge] Florence Caption disabled by Pipeline Planner."
print(status)
return (captionforge_run_config, status, "", task_prompt)
first_run = run_plan[0]
if run_plan_connected:
write_jsonl = True
also_jsonl = False
skip_existing_txt = False
skip_existing_jsonl_images = False
overwrite = True
effective_prefix = prefix
if first_run.trigger_word:
effective_prefix = f"{first_run.trigger_word}, {prefix}".strip(", ")
generation = GenerationConfig(
max_new_tokens=int(first_run.max_new_tokens),
num_beams=int(num_beams),
temperature=float(first_run.temperature),
top_p=float(first_run.top_p),
top_k=int(first_run.top_k),
seed=first_run.seed,
)
cleanup = CleanupConfig(
trigger="",
prefix=effective_prefix,
suffix=suffix,
forbidden_phrases=_parse_forbidden_lines(forbidden_phrases),
replacement_rules=_parse_replace_pairs(replace_pairs),
)
florence_config = FlorenceCaptionConfig(
model_name=model,
model_path="",
model_root=str(JLC_FLORENCE_MODEL_ROOT),
memory_mode=memory_mode,
dtype="bf16",
device="auto",
trust_remote_code=True,
keep_loaded=bool(keep_loaded),
quiet_transformers_load=True,
max_size=int(first_run.max_size),
task_prompt=task_prompt,
allow_download=True,
use_comfy_model_management=True,
)
engine = FlorenceCaptionEngine(config=florence_config, generation=generation, cleanup=cleanup)
input_path = (input_path or "").strip()
output_dir = (output_dir or "").strip()
jsonl_filename = (jsonl_filename or "captions.jsonl").strip() or "captions.jsonl"
planned_caption_jsonl_path = _planned_caption_jsonl_path(captionforge_run_config) if run_plan_connected else None
if first_run.output_dir:
output_dir = first_run.output_dir
if first_run.input_path:
input_path = first_run.input_path
if run_plan_connected:
recursive = bool(first_run.recursive)
filename_glob = first_run.filename_glob or "*"
if planned_caption_jsonl_path is not None:
output_dir = str(planned_caption_jsonl_path.parent)
jsonl_filename = planned_caption_jsonl_path.name
use_jsonl = bool(write_jsonl or also_jsonl)
all_records: list[CaptionRecord] = []
# Direct IMAGE input
pil_images = _tensor_to_pil(image)
if pil_images and input_path:
print("[CaptionForge] IMAGE input and input_path are both active; captioning both sources.")
if pil_images:
engine.load()
fallback_dir = Path(folder_paths.get_output_directory()) / "jlc_florence_caption"
image_output_dir = Path(output_dir) if output_dir else fallback_dir
image_output_dir.mkdir(parents=True, exist_ok=True)
jsonl_path = image_output_dir / jsonl_filename
seen_jsonl_images: set[str] = set()
if use_jsonl and skip_existing_jsonl_images:
seen_jsonl_images = load_existing_jsonl_images(jsonl_path)
if write_run_config and not input_path:
write_run_config_json(
image_output_dir / f"jlc_florence_caption_run_config_{timestamp()}.json",
engine.build_run_config(),
dry_run=bool(dry_run),
)
for index, pil in enumerate(pil_images):
source_name = f"comfy_image_{index:04d}"
txt_path = image_output_dir / f"{source_name}.txt"
if len(run_plan) == 1 and skip_existing_txt and write_txt and txt_path.exists() and not overwrite:
print(f"[JLC Florence Caption] Skipping existing TXT: {txt_path}")
continue
if use_jsonl and skip_existing_jsonl_images and source_name in seen_jsonl_images:
print(f"[JLC Florence Caption] Skipping existing JSONL image: {source_name}")
continue
for run in run_plan:
engine.generation = GenerationConfig(
max_new_tokens=int(run.max_new_tokens),
num_beams=int(num_beams),
temperature=float(run.temperature),
top_p=float(run.top_p),
top_k=int(run.top_k),
seed=run.seed,
)
effective_prefix = prefix
if run.trigger_word:
effective_prefix = f"{run.trigger_word}, {prefix}".strip(", ")
engine.cleanup = CleanupConfig(
trigger="",
prefix=effective_prefix,
suffix=suffix,
forbidden_phrases=_parse_forbidden_lines(forbidden_phrases),
replacement_rules=_parse_replace_pairs(replace_pairs),
)
engine.config.max_size = int(run.max_size)
engine.config.task_prompt = task_prompt
run_txt_path = txt_path if len(run_plan) <= 1 else image_output_dir / f"{source_name}__cf_run_{run.ensemble_run_index:02d}.txt"
if skip_existing_txt and write_txt and run_txt_path.exists() and not overwrite:
print(f"[JLC Florence Caption] Skipping existing TXT: {run_txt_path}")
continue
t0 = time.perf_counter()
final_caption, raw_caption = engine.caption_pil(pil)
dt = time.perf_counter() - t0
print(f"[JLC Florence Caption] Generation time run {run.ensemble_run_index}: {dt:.2f}s")
record = CaptionRecord(
image=source_name,
caption=final_caption,
raw_caption=raw_caption,
model_name=model,
model_path=str(engine.local_model_path or ""),
task_prompt=task_prompt,
seed=run.seed,
temperature=run.temperature,
top_p=run.top_p,
top_k=run.top_k,
num_beams=int(num_beams),
max_new_tokens=run.max_new_tokens,
max_size=int(run.max_size),
timestamp=datetime.now().isoformat(timespec="seconds"),
captionforge_pass="A",
model_family="florence",
ensemble_run_index=run.ensemble_run_index,
image_key=source_name,
)
all_records.append(record)
if use_jsonl:
append_jsonl_records(jsonl_path, [record], dry_run=bool(dry_run))
if write_txt:
write_text_sidecar(run_txt_path, record.caption, overwrite=bool(overwrite), backup_existing=bool(backup_existing), dry_run=bool(dry_run))
print(f"[JLC Florence Caption] Captioned IMAGE {index + 1}/{len(pil_images)} run {run.ensemble_run_index + 1}/{len(run_plan)}: {source_name}")
# File/folder input path. Note: engine.caption_batch is single-caption-per-image.
# For planned multi-run file mode, use direct iteration logic in a later patch if needed.
batch_result = None
if input_path:
also_jsonl_path = ""
if also_jsonl and not write_jsonl:
base_dir = Path(output_dir) if output_dir else (Path(input_path).parent if Path(input_path).is_file() else Path(input_path))
also_jsonl_path = str(base_dir / jsonl_filename)
batch = BatchCaptionConfig(
input_path=input_path,
recursive=bool(recursive),
filename_glob=(filename_glob or "*").strip() or "*",
output_dir=output_dir,
write_txt=bool(write_txt),
write_jsonl=bool(write_jsonl),
jsonl_filename=jsonl_filename,
also_jsonl_path=also_jsonl_path,
write_run_config=bool(write_run_config),
overwrite=bool(overwrite),
backup_existing=bool(backup_existing),
dry_run=bool(dry_run),
limit=int(limit),
skip_existing_txt=bool(skip_existing_txt),
skip_existing_jsonl_images=bool(skip_existing_jsonl_images),
)
batch_result = engine.caption_batch(batch)
all_records.extend(batch_result.records)
if not pil_images and batch_result is None:
raise RuntimeError("No image input found. Connect an IMAGE input or provide input_path pointing to an image file or folder.")
if not keep_loaded:
engine.unload()
caption_string = "\n\n".join(r.caption for r in all_records if r.status == "ok")
jsonl_string = _make_jsonl_string(all_records)
return (captionforge_run_config, caption_string, jsonl_string, task_prompt)
NODE_CLASS_MAPPINGS = {
"JLC_FlorenceCaption": JLC_FlorenceCaption,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"JLC_FlorenceCaption": "\u2003JLC Florence Caption",
}
@@ -1,438 +0,0 @@
"""
JLC Florence Caption (Lite) — ComfyUI Node Wrapper
Drop target:
CaptionForge/nodes/jlc_florence_caption_lite_CUI_node.py
Compact Florence-2 raw-caption witness node for CaptionForge.
"""
from __future__ import annotations
MANIFEST = {
"name": "JLC Florence Caption (Lite)",
"version": (0, 1, 0),
"author": "J. L. Córdova",
"description": (
"Compact ComfyUI frontend for Florence-2-family CaptionForge Pass A captioning. "
"Consumes CAPTIONFORGE_PIPELINE_PLAN objects, writes TXT audit sidecars and shared "
"JSONL evidence in planned runs, and delegates model loading/generation/cache behavior "
"to jlc_florence_caption_engine.py."
),
}
import json
from datetime import datetime
from pathlib import Path
import time
import numpy as np
import torch
from PIL import Image
import folder_paths
from ..engines.jlc_florence_caption_engine import (
CaptionRecord,
CleanupConfig,
FlorenceCaptionConfig,
FlorenceCaptionEngine,
GenerationConfig,
MEMORY_MODES,
MODEL_REGISTRY,
TASK_PROMPTS,
append_jsonl_records,
record_to_json,
resolve_task_prompt,
timestamp,
write_run_config_json,
write_text_sidecar,
)
from ..engines.captionforge_pipeline_planner_engine import expand_captionforge_runs
JLC_FLORENCE_MODEL_ROOT = Path(folder_paths.models_dir) / "LLM" / "JLC_FlorenceCaption"
DEFAULT_JSONL_FILENAME = "captions.jsonl"
_SUPPORTED_IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff"}
def _tensor_to_pil(image_tensor) -> list[Image.Image]:
if image_tensor is None:
return []
if isinstance(image_tensor, torch.Tensor):
image_tensor = image_tensor.detach().cpu()
if image_tensor.ndim == 3:
image_tensor = image_tensor.unsqueeze(0)
images: list[Image.Image] = []
for img in image_tensor:
arr = img.numpy()
arr = np.clip(arr * 255.0, 0, 255).astype(np.uint8)
images.append(Image.fromarray(arr).convert("RGB"))
return images
def _safe_source_name(value: str) -> str:
cleaned = []
for ch in value.replace("\\", "/"):
if ch.isalnum() or ch in {"-", "_", "."}:
cleaned.append(ch)
elif ch == "/":
cleaned.append("__")
else:
cleaned.append("_")
return "".join(cleaned).strip("._") or "image"
def _iter_input_path_images(input_path: str, recursive: bool, filename_glob: str) -> list[tuple[str, Path]]:
text = str(input_path or "").strip()
if not text:
return []
root = Path(text)
if not root.exists():
raise RuntimeError(f"CaptionForge input_path does not exist: {root}")
glob_text = (filename_glob or "*").strip() or "*"
if root.is_file():
if root.suffix.lower() not in _SUPPORTED_IMAGE_SUFFIXES:
raise RuntimeError(f"CaptionForge input_path is not a supported image file: {root}")
return [(_safe_source_name(root.stem), root)]
pattern_iter = root.rglob(glob_text) if recursive else root.glob(glob_text)
paths = sorted(p for p in pattern_iter if p.is_file() and p.suffix.lower() in _SUPPORTED_IMAGE_SUFFIXES)
return [(_safe_source_name(str(p.relative_to(root).with_suffix(""))), p) for p in paths]
def _open_image(path: Path) -> Image.Image:
with Image.open(path) as im:
return im.convert("RGB")
def _normalize_pipeline_plan(config):
if config is None:
return {}
if isinstance(config, dict):
return dict(config)
if isinstance(config, str):
text = config.strip()
if not text:
return {}
try:
obj = json.loads(text)
return obj if isinstance(obj, dict) else {}
except Exception:
return {}
return {}
def _planned_caption_jsonl_path(captionforge_run_config) -> Path | None:
cfg = _normalize_pipeline_plan(captionforge_run_config)
paths = cfg.get("paths") if isinstance(cfg.get("paths"), dict) else {}
for key in ("caption_jsonl", "pass_a_jsonl"):
value = str(paths.get(key) or "").strip()
if value:
return Path(value)
return None
def _make_jsonl_string(records: list[CaptionRecord]) -> str:
return "\n".join(json.dumps(record_to_json(r), ensure_ascii=False) for r in records)
def _run_txt_path(output_dir: Path, source_name: str, run_count: int, run_index: int) -> Path:
if run_count <= 1:
return output_dir / f"{source_name}.txt"
return output_dir / f"{source_name}__cf_run_{run_index:02d}.txt"
class JLC_FlorenceCaptionLite:
"""Minimal Florence caption node; Run Plan turns it into a CaptionForge evidence node."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": (
list(MODEL_REGISTRY.keys()),
{
"default": "Florence-2-base-ft",
"tooltip": "Select Florence-2 model. Models are loaded from ComfyUI/models/LLM/JLC_FlorenceCaption/.",
},
),
"memory_mode": (
list(MEMORY_MODES.keys()),
{
"default": "Default",
"tooltip": "Florence v1 uses normal Transformers loading; no bitsandbytes path is exposed.",
},
),
"keep_loaded": (
"BOOLEAN",
{
"default": True,
"tooltip": "Keep Florence cached after captioning. The shared CaptionForge cache may still evict it when another caption model loads.",
},
),
"task": (
list(TASK_PROMPTS.keys()),
{
"default": "Detailed Caption",
"tooltip": "Florence caption task. Detailed Caption is the suggested CaptionForge Lite default.",
},
),
"custom_task_prompt": (
"STRING",
{
"default": "",
"multiline": False,
"tooltip": "Optional raw Florence task prompt override, such as <MORE_DETAILED_CAPTION>. Leave blank to use the task dropdown.",
"advanced": True,
},
),
"max_new_tokens": (
"INT",
{
"default": 256,
"min": 16,
"max": 2048,
"step": 8,
"tooltip": "Standalone token budget. Run Plan overrides this when connected.",
},
),
"num_beams": (
"INT",
{
"default": 3,
"min": 1,
"max": 16,
"step": 1,
"tooltip": "Beam count for deterministic Florence generation. 3 is the common Florence example value.",
},
),
"temperature": (
"FLOAT",
{
"default": 0.0,
"min": 0.0,
"max": 2.0,
"step": 0.01,
"tooltip": "Standalone sampling temperature. 0.0 keeps Florence deterministic. Run Plan overrides this when connected.",
},
),
"top_p": (
"FLOAT",
{
"default": 0.90,
"min": 0.0,
"max": 1.0,
"step": 0.01,
"tooltip": "Top-p used only when temperature > 0. Run Plan overrides this when connected.",
},
),
"top_k": (
"INT",
{
"default": 50,
"min": 0,
"max": 500,
"step": 1,
"tooltip": "Top-k used only when temperature > 0. Run Plan overrides this when connected.",
},
),
},
"optional": {
"image": ("IMAGE", {"tooltip": "Image or batch of images to caption."}),
"captionforge_run_config": (
"CAPTIONFORGE_PIPELINE_PLAN",
{
"tooltip": (
"Optional shared CaptionForge Run Plan. When connected, this Lite node uses "
"the florence model-family plan for witness count, routing, seeds/schedules, "
"size, token budget, trigger word, and shared JSONL evidence."
),
},
),
},
}
RETURN_TYPES = ("CAPTIONFORGE_PIPELINE_PLAN", "STRING", "STRING", "STRING")
RETURN_NAMES = ("captionforge_run_config_out", "caption", "jsonl_records", "resolved_prompt")
FUNCTION = "caption"
CATEGORY = "JLC/Captioning"
@classmethod
def IS_CHANGED(cls, **kwargs):
return float("NaN")
def caption(
self,
model,
memory_mode,
keep_loaded,
task,
custom_task_prompt,
max_new_tokens,
num_beams,
temperature,
top_p,
top_k,
image=None,
captionforge_run_config=None,
):
run_plan_connected = bool(captionforge_run_config)
task_prompt = resolve_task_prompt(task, custom_task_prompt)
run_plan = expand_captionforge_runs(
captionforge_run_config,
model_key="florence",
widget_captions_per_image=1,
widget_seed=-1,
widget_temperature=float(temperature),
widget_top_p=float(top_p),
widget_top_k=int(top_k),
widget_max_new_tokens=int(max_new_tokens),
widget_max_size=1024,
widget_trigger_word="",
widget_output_dir="",
widget_input_path="",
widget_recursive=True,
widget_filename_glob="*",
)
if run_plan_connected and not run_plan:
status = "[CaptionForge] Florence Caption Lite disabled by Pipeline Planner."
print(status)
return (captionforge_run_config, status, "", task_prompt)
first_run = run_plan[0]
generation = GenerationConfig(
max_new_tokens=int(first_run.max_new_tokens),
num_beams=int(num_beams),
temperature=float(first_run.temperature),
top_p=float(first_run.top_p),
top_k=int(first_run.top_k),
seed=first_run.seed,
)
cleanup = CleanupConfig(
trigger="",
prefix=(f"{first_run.trigger_word}," if first_run.trigger_word else ""),
suffix="",
forbidden_phrases=[],
replacement_rules=[],
)
florence_config = FlorenceCaptionConfig(
model_name=model,
model_path="",
model_root=str(JLC_FLORENCE_MODEL_ROOT),
memory_mode=memory_mode,
dtype="bf16",
device="auto",
trust_remote_code=True,
keep_loaded=bool(keep_loaded),
quiet_transformers_load=True,
max_size=int(first_run.max_size),
task_prompt=task_prompt,
allow_download=True,
use_comfy_model_management=True,
)
engine = FlorenceCaptionEngine(config=florence_config, generation=generation, cleanup=cleanup)
direct_images = [(f"comfy_image_{i:04d}", pil) for i, pil in enumerate(_tensor_to_pil(image))]
file_images: list[tuple[str, Path]] = []
if first_run.input_path:
file_images = _iter_input_path_images(first_run.input_path, first_run.recursive, first_run.filename_glob)
if direct_images and file_images:
print("[CaptionForge] IMAGE input and Run Plan input_path are both active; captioning both sources.")
if not direct_images and not file_images:
raise RuntimeError("No image input found. Connect an IMAGE input or provide input_path in the CaptionForge Run Plan.")
planned_caption_jsonl_path = _planned_caption_jsonl_path(captionforge_run_config) if run_plan_connected else None
output_dir = Path(first_run.output_dir) if first_run.output_dir else Path(folder_paths.get_output_directory()) / "jlc_florence_caption_lite"
if planned_caption_jsonl_path is not None:
output_dir = planned_caption_jsonl_path.parent
if run_plan_connected:
output_dir.mkdir(parents=True, exist_ok=True)
jsonl_path = planned_caption_jsonl_path or (output_dir / DEFAULT_JSONL_FILENAME)
all_records: list[CaptionRecord] = []
engine.load()
if run_plan_connected:
write_run_config_json(output_dir / f"jlc_florence_caption_lite_run_config_{timestamp()}.json", engine.build_run_config(), dry_run=False)
def process_one(source_name: str, pil: Image.Image):
for run in run_plan:
engine.generation = GenerationConfig(
max_new_tokens=int(run.max_new_tokens),
num_beams=int(num_beams),
temperature=float(run.temperature),
top_p=float(run.top_p),
top_k=int(run.top_k),
seed=run.seed,
)
engine.cleanup = CleanupConfig(
trigger="",
prefix=(f"{run.trigger_word}," if run.trigger_word else ""),
suffix="",
forbidden_phrases=[],
replacement_rules=[],
)
engine.config.max_size = int(run.max_size)
engine.config.task_prompt = task_prompt
t0 = time.perf_counter()
final_caption, raw_caption = engine.caption_pil(pil)
dt = time.perf_counter() - t0
print(f"[JLC Florence Caption Lite] Generation time run {run.ensemble_run_index}: {dt:.2f}s")
record = CaptionRecord(
image=source_name,
caption=final_caption,
raw_caption=raw_caption,
model_name=model,
model_path=str(engine.local_model_path or ""),
task_prompt=task_prompt,
seed=run.seed,
temperature=run.temperature,
top_p=run.top_p,
top_k=run.top_k,
num_beams=int(num_beams),
max_new_tokens=run.max_new_tokens,
max_size=int(run.max_size),
timestamp=datetime.now().isoformat(timespec="seconds"),
captionforge_pass="A",
model_family="florence",
ensemble_run_index=run.ensemble_run_index,
image_key=source_name,
)
all_records.append(record)
if run_plan_connected:
append_jsonl_records(jsonl_path, [record], dry_run=False)
write_text_sidecar(
_run_txt_path(output_dir, source_name, len(run_plan), run.ensemble_run_index),
record.caption,
overwrite=True,
backup_existing=False,
dry_run=False,
)
print(f"[JLC Florence Caption Lite] Captioned {source_name} run {run.ensemble_run_index + 1}/{len(run_plan)}")
for source_name, pil in direct_images:
process_one(source_name, pil)
for source_name, path in file_images:
process_one(source_name, _open_image(path))
if not keep_loaded:
engine.unload()
caption_string = "\n\n".join(r.caption for r in all_records if r.status == "ok")
jsonl_string = _make_jsonl_string(all_records)
return (captionforge_run_config, caption_string, jsonl_string, task_prompt)
NODE_CLASS_MAPPINGS = {
"JLC_FlorenceCaptionLite": JLC_FlorenceCaptionLite,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"JLC_FlorenceCaptionLite": "\u2003JLC Florence Caption (Lite)",
}
@@ -1,243 +0,0 @@
"""
JLC InternVL Caption — ComfyUI Node Wrapper
"""
from __future__ import annotations
MANIFEST = {
"name": "JLC InternVL Caption",
"version": (0, 1, 0),
"author": "J. L. Córdova",
"description": "ComfyUI frontend for InternVL-family CaptionForge Pass A captioning.",
}
import json
from datetime import datetime
from pathlib import Path
import time
import numpy as np
import torch
from PIL import Image
import folder_paths
from ..engines.jlc_internvl_caption_engine import (
BatchCaptionConfig,
CaptionRecord,
CleanupConfig,
InternVLCaptionConfig,
InternVLCaptionEngine,
GenerationConfig,
MEMORY_MODES,
MODEL_REGISTRY,
PROMPT_PRESETS,
append_jsonl_records,
load_existing_jsonl_images,
probe_registry_model_download,
record_to_json,
resolve_prompt,
timestamp,
write_run_config_json,
write_text_sidecar,
)
from ..engines.captionforge_pipeline_planner_engine import expand_captionforge_runs
JLC_INTERVLVL_MODEL_ROOT = Path(folder_paths.models_dir) / "LLM" / "JLC_InternVLCaption"
def _tensor_to_pil(image_tensor) -> list[Image.Image]:
if image_tensor is None:
return []
if isinstance(image_tensor, torch.Tensor):
image_tensor = image_tensor.detach().cpu()
if image_tensor.ndim == 3:
image_tensor = image_tensor.unsqueeze(0)
images: list[Image.Image] = []
for img in image_tensor:
arr = img.numpy()
arr = np.clip(arr * 255.0, 0, 255).astype(np.uint8)
images.append(Image.fromarray(arr).convert("RGB"))
return images
def _parse_forbidden_lines(value: str) -> list[str]:
value = (value or "").strip()
if not value:
return []
return [line.strip() for line in value.splitlines() if line.strip()]
def _parse_replace_pairs(value: str) -> list[tuple[str, str]]:
value = (value or "").strip()
if not value:
return []
rules: list[tuple[str, str]] = []
for line in value.splitlines():
line = line.strip()
if not line or line.startswith("#") or "=>" not in line:
continue
old, new = line.split("=>", 1)
old = old.strip(); new = new.strip()
if old:
rules.append((old, new))
return rules
def _normalize_pipeline_plan(config):
if config is None:
return {}
if isinstance(config, dict):
return dict(config)
if isinstance(config, str):
text = config.strip()
if not text:
return {}
try:
obj = json.loads(text)
return obj if isinstance(obj, dict) else {}
except Exception:
return {}
return {}
def _planned_caption_jsonl_path(captionforge_run_config) -> Path | None:
cfg = _normalize_pipeline_plan(captionforge_run_config)
paths = cfg.get("paths") if isinstance(cfg.get("paths"), dict) else {}
for key in ("caption_jsonl", "pass_a_jsonl"):
value = str(paths.get(key) or "").strip()
if value:
return Path(value)
return None
def _make_jsonl_string(records: list[CaptionRecord]) -> str:
return "\n".join(json.dumps(record_to_json(r), ensure_ascii=False) for r in records)
class JLC_InternVLCaption:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": (list(MODEL_REGISTRY.keys()), {"default": "InternVL2.5-2B", "tooltip": "Select model."}),
"memory_mode": (list(MEMORY_MODES.keys()), {"default": "Default", "tooltip": "Default uses regular Transformers loading. Balanced (8-bit) uses bitsandbytes when available."}),
"input_path": ("STRING", {"default": "", "multiline": False}),
"recursive": ("BOOLEAN", {"default": True}),
"filename_glob": ("STRING", {"default": "*", "multiline": False}),
"prompt_preset": (list(PROMPT_PRESETS.keys()), {"default": "detailed"}),
"custom_prompt": ("STRING", {"default": "", "multiline": True}),
"max_new_tokens": ("INT", {"default": 256, "min": 16, "max": 2048, "step": 8}),
"num_beams": ("INT", {"default": 3, "min": 1, "max": 16, "step": 1}),
"temperature": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 2.0, "step": 0.01}),
"top_p": ("FLOAT", {"default": 0.90, "min": 0.0, "max": 1.0, "step": 0.01}),
"top_k": ("INT", {"default": 50, "min": 0, "max": 500, "step": 1}),
"repetition_penalty": ("FLOAT", {"default": 1.0, "min": 1.0, "max": 2.0, "step": 0.01}),
"captions_per_image": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}),
"seed": ("INT", {"default": -1, "min": -1, "max": 0xFFFFFFFF, "step": 1}),
"max_size": ("INT", {"default": 768, "min": 0, "max": 4096, "step": 64}),
"output_dir": ("STRING", {"default": "", "multiline": False}),
"write_txt": ("BOOLEAN", {"default": True}),
"write_jsonl": ("BOOLEAN", {"default": False}),
"also_jsonl": ("BOOLEAN", {"default": False}),
"write_run_config": ("BOOLEAN", {"default": True}),
"jsonl_filename": ("STRING", {"default": "captions.jsonl", "multiline": False}),
"overwrite": ("BOOLEAN", {"default": False}),
"backup_existing": ("BOOLEAN", {"default": True}),
"dry_run": ("BOOLEAN", {"default": False}),
"limit": ("INT", {"default": 0, "min": 0, "max": 100000, "step": 1}),
"skip_existing_txt": ("BOOLEAN", {"default": True}),
"skip_existing_jsonl_images": ("BOOLEAN", {"default": False}),
"prefix": ("STRING", {"default": "", "multiline": False}),
"suffix": ("STRING", {"default": "", "multiline": False}),
"forbidden_phrases": ("STRING", {"default": "", "multiline": True}),
"replace_pairs": ("STRING", {"default": "", "multiline": True}),
"keep_loaded": ("BOOLEAN", {"default": True}),
"download_probe_only": ("BOOLEAN", {"default": False}),
},
"optional": {
"image": ("IMAGE", {}),
"captionforge_run_config": ("CAPTIONFORGE_PIPELINE_PLAN", {}),
},
}
RETURN_TYPES = ("CAPTIONFORGE_PIPELINE_PLAN", "STRING", "STRING", "STRING")
RETURN_NAMES = ("captionforge_run_config_out", "caption", "jsonl_records", "resolved_prompt")
FUNCTION = "caption"
CATEGORY = "JLC/Captioning"
@classmethod
def IS_CHANGED(cls, **kwargs):
return float("NaN")
def caption(self, model, memory_mode, input_path, recursive, filename_glob, prompt_preset, custom_prompt, max_new_tokens, num_beams, temperature, top_p, top_k, repetition_penalty, captions_per_image, seed, max_size, output_dir, write_txt, write_jsonl, also_jsonl, write_run_config, jsonl_filename, overwrite, backup_existing, dry_run, limit, skip_existing_txt, skip_existing_jsonl_images, prefix, suffix, forbidden_phrases, replace_pairs, keep_loaded, download_probe_only, image=None, captionforge_run_config=None):
if download_probe_only:
result = probe_registry_model_download(model, JLC_INTERVLVL_MODEL_ROOT)
return (captionforge_run_config, result, "", "")
prompt = resolve_prompt(prompt_preset, custom_prompt)
run_plan = expand_captionforge_runs(captionforge_run_config, model_key="internvl", widget_captions_per_image=int(captions_per_image), widget_seed=int(seed), widget_temperature=float(temperature), widget_top_p=float(top_p), widget_top_k=int(top_k), widget_max_new_tokens=int(max_new_tokens), widget_max_size=int(max_size), widget_trigger_word="", widget_output_dir=output_dir, widget_input_path=input_path, widget_recursive=bool(recursive), widget_filename_glob=filename_glob)
run_plan_connected = bool(captionforge_run_config)
if run_plan_connected and not run_plan:
status = "[CaptionForge] JLC InternVL Caption disabled by Pipeline Planner."
return (captionforge_run_config, status, "", prompt)
first_run = run_plan[0]
if run_plan_connected:
write_jsonl = True; also_jsonl = False; skip_existing_txt = False; skip_existing_jsonl_images = False; overwrite = True
effective_prefix = prefix
if first_run.trigger_word:
effective_prefix = f"{first_run.trigger_word}, {prefix}".strip(", ")
generation = GenerationConfig(max_new_tokens=int(first_run.max_new_tokens), num_beams=int(num_beams), temperature=float(first_run.temperature), top_p=float(first_run.top_p), top_k=int(first_run.top_k), repetition_penalty=float(repetition_penalty), seed=first_run.seed)
cleanup = CleanupConfig(trigger="", prefix=effective_prefix, suffix=suffix, forbidden_phrases=_parse_forbidden_lines(forbidden_phrases), replacement_rules=_parse_replace_pairs(replace_pairs))
cfg = InternVLCaptionConfig(model_name=model, model_path="", model_root=str(JLC_INTERVLVL_MODEL_ROOT), memory_mode=memory_mode, dtype="bf16", device="auto", trust_remote_code=True, keep_loaded=bool(keep_loaded), quiet_transformers_load=True, max_size=int(first_run.max_size), prompt=prompt, allow_download=True, use_comfy_model_management=True)
engine = InternVLCaptionEngine(config=cfg, generation=generation, cleanup=cleanup)
input_path = (input_path or "").strip(); output_dir = (output_dir or "").strip(); jsonl_filename = (jsonl_filename or "captions.jsonl").strip() or "captions.jsonl"
planned_caption_jsonl_path = _planned_caption_jsonl_path(captionforge_run_config) if run_plan_connected else None
if first_run.output_dir: output_dir = first_run.output_dir
if first_run.input_path: input_path = first_run.input_path
if run_plan_connected:
recursive = bool(first_run.recursive); filename_glob = first_run.filename_glob or "*"
if planned_caption_jsonl_path is not None:
output_dir = str(planned_caption_jsonl_path.parent); jsonl_filename = planned_caption_jsonl_path.name
use_jsonl = bool(write_jsonl or also_jsonl)
all_records: list[CaptionRecord] = []
pil_images = _tensor_to_pil(image)
if pil_images:
engine.load()
fallback_dir = Path(folder_paths.get_output_directory()) / "jlc_internvl_caption"
image_output_dir = Path(output_dir) if output_dir else fallback_dir
image_output_dir.mkdir(parents=True, exist_ok=True)
jsonl_path = image_output_dir / jsonl_filename
seen_jsonl = load_existing_jsonl_images(jsonl_path) if use_jsonl and skip_existing_jsonl_images else set()
if write_run_config and not input_path:
write_run_config_json(image_output_dir / f"jlc_internvl_caption_run_config_{timestamp()}.json", engine.build_run_config(), dry_run=bool(dry_run))
for idx, pil in enumerate(pil_images):
source_name = f"comfy_image_{idx:04d}"
txt_path = image_output_dir / f"{source_name}.txt"
if len(run_plan) == 1 and skip_existing_txt and write_txt and txt_path.exists() and not overwrite:
continue
if use_jsonl and skip_existing_jsonl_images and source_name in seen_jsonl:
continue
for run in run_plan:
engine.generation = GenerationConfig(max_new_tokens=int(run.max_new_tokens), num_beams=int(num_beams), temperature=float(run.temperature), top_p=float(run.top_p), top_k=int(run.top_k), repetition_penalty=float(repetition_penalty), seed=run.seed)
eff_prefix = prefix
if run.trigger_word:
eff_prefix = f"{run.trigger_word}, {prefix}".strip(", ")
engine.cleanup = CleanupConfig(trigger="", prefix=eff_prefix, suffix=suffix, forbidden_phrases=_parse_forbidden_lines(forbidden_phrases), replacement_rules=_parse_replace_pairs(replace_pairs))
engine.config.max_size = int(run.max_size); engine.config.prompt = prompt
run_txt_path = txt_path if len(run_plan) <= 1 else image_output_dir / f"{source_name}__cf_run_{run.ensemble_run_index:02d}.txt"
t0 = time.perf_counter(); final_caption, raw_caption = engine.caption_pil(pil); _ = time.perf_counter() - t0
rec = CaptionRecord(image=source_name, caption=final_caption, raw_caption=raw_caption, model_name=model, model_path=str(engine.local_model_path or ""), prompt=prompt, seed=run.seed, temperature=run.temperature, top_p=run.top_p, top_k=run.top_k, num_beams=int(num_beams), max_new_tokens=run.max_new_tokens, max_size=int(run.max_size), timestamp=datetime.now().isoformat(timespec="seconds"), captionforge_pass="A", model_family="internvl", ensemble_run_index=run.ensemble_run_index, image_key=source_name)
all_records.append(rec)
if use_jsonl:
append_jsonl_records(jsonl_path, [rec], dry_run=bool(dry_run))
if write_txt:
write_text_sidecar(run_txt_path, rec.caption, overwrite=bool(overwrite), backup_existing=bool(backup_existing), dry_run=bool(dry_run))
batch_result = None
if input_path:
also_jsonl_path = ""
if also_jsonl and not write_jsonl:
base_dir = Path(output_dir) if output_dir else (Path(input_path).parent if Path(input_path).is_file() else Path(input_path))
also_jsonl_path = str(base_dir / jsonl_filename)
batch = BatchCaptionConfig(input_path=input_path, recursive=bool(recursive), filename_glob=(filename_glob or "*").strip() or "*", output_dir=output_dir, write_txt=bool(write_txt), write_jsonl=bool(write_jsonl), jsonl_filename=jsonl_filename, also_jsonl_path=also_jsonl_path, write_run_config=bool(write_run_config), overwrite=bool(overwrite), backup_existing=bool(backup_existing), dry_run=bool(dry_run), limit=int(limit), skip_existing_txt=bool(skip_existing_txt), skip_existing_jsonl_images=bool(skip_existing_jsonl_images))
batch_result = engine.caption_batch(batch)
all_records.extend(batch_result.records)
if not pil_images and batch_result is None:
raise RuntimeError("No image input found. Connect an IMAGE input or provide input_path.")
if not keep_loaded:
engine.unload()
caption_string = "\n\n".join(r.caption for r in all_records if r.status == "ok")
jsonl_string = _make_jsonl_string(all_records)
return (captionforge_run_config, caption_string, jsonl_string, prompt)
NODE_CLASS_MAPPINGS = {"JLC_InternVLCaption": JLC_InternVLCaption}
NODE_DISPLAY_NAME_MAPPINGS = {"JLC_InternVLCaption": " JLC InternVL Caption"}
@@ -1,205 +0,0 @@
"""
JLC InternVL Caption (Lite) — ComfyUI Node Wrapper
"""
from __future__ import annotations
MANIFEST = {
"name": "JLC InternVL Caption (Lite)",
"version": (0, 1, 0),
"author": "J. L. Córdova",
"description": "Compact ComfyUI frontend for InternVL-family CaptionForge Pass A captioning.",
}
import json
from datetime import datetime
from pathlib import Path
import time
import numpy as np
import torch
from PIL import Image
import folder_paths
from ..engines.jlc_internvl_caption_engine import (
CaptionRecord,
CleanupConfig,
InternVLCaptionConfig,
InternVLCaptionEngine,
GenerationConfig,
MEMORY_MODES,
MODEL_REGISTRY,
PROMPT_PRESETS,
append_jsonl_records,
record_to_json,
resolve_prompt,
timestamp,
write_run_config_json,
write_text_sidecar,
)
from ..engines.captionforge_pipeline_planner_engine import expand_captionforge_runs
JLC_INTERVLVL_MODEL_ROOT = Path(folder_paths.models_dir) / "LLM" / "JLC_InternVLCaption"
DEFAULT_JSONL_FILENAME = "captions.jsonl"
_SUPPORTED_IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff"}
def _tensor_to_pil(image_tensor) -> list[Image.Image]:
if image_tensor is None:
return []
if isinstance(image_tensor, torch.Tensor):
image_tensor = image_tensor.detach().cpu()
if image_tensor.ndim == 3:
image_tensor = image_tensor.unsqueeze(0)
images: list[Image.Image] = []
for img in image_tensor:
arr = img.numpy(); arr = np.clip(arr * 255.0, 0, 255).astype(np.uint8)
images.append(Image.fromarray(arr).convert("RGB"))
return images
def _safe_source_name(value: str) -> str:
cleaned = []
for ch in value.replace("\\", "/"):
if ch.isalnum() or ch in {"-", "_", "."}:
cleaned.append(ch)
elif ch == "/":
cleaned.append("__")
else:
cleaned.append("_")
out = "".join(cleaned).strip("._")
return out or "image"
def _iter_input_path_images(input_path: str, recursive: bool, filename_glob: str) -> list[tuple[str, Path]]:
root = Path(str(input_path or "").strip())
if not root:
return []
if not root.exists():
raise RuntimeError(f"CaptionForge input_path does not exist: {root}")
glob_text = (filename_glob or "*").strip() or "*"
if root.is_file():
if root.suffix.lower() not in _SUPPORTED_IMAGE_SUFFIXES:
raise RuntimeError(f"CaptionForge input_path is not a supported image file: {root}")
return [(_safe_source_name(root.stem), root)]
pattern_iter = root.rglob(glob_text) if recursive else root.glob(glob_text)
paths = sorted(p for p in pattern_iter if p.is_file() and p.suffix.lower() in _SUPPORTED_IMAGE_SUFFIXES)
items = []
for path in paths:
rel = path.relative_to(root).with_suffix("")
items.append((_safe_source_name(str(rel)), path))
return items
def _open_image(path: Path) -> Image.Image:
with Image.open(path) as im:
return im.convert("RGB")
def _normalize_pipeline_plan(config):
if config is None:
return {}
if isinstance(config, dict):
return dict(config)
if isinstance(config, str):
text = config.strip()
if not text:
return {}
try:
obj = json.loads(text)
return obj if isinstance(obj, dict) else {}
except Exception:
return {}
return {}
def _planned_caption_jsonl_path(captionforge_run_config) -> Path | None:
cfg = _normalize_pipeline_plan(captionforge_run_config)
paths = cfg.get("paths") if isinstance(cfg.get("paths"), dict) else {}
for key in ("caption_jsonl", "pass_a_jsonl"):
value = str(paths.get(key) or "").strip()
if value:
return Path(value)
return None
def _make_jsonl_string(records: list[CaptionRecord]) -> str:
return "\n".join(json.dumps(record_to_json(r), ensure_ascii=False) for r in records)
def _run_txt_path(output_dir: Path, source_name: str, run_count: int, run_index: int) -> Path:
if run_count <= 1:
return output_dir / f"{source_name}.txt"
return output_dir / f"{source_name}__cf_run_{run_index:02d}.txt"
class JLC_InternVLCaptionLite:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": (list(MODEL_REGISTRY.keys()), {"default": "InternVL2.5-2B"}),
"memory_mode": (list(MEMORY_MODES.keys()), {"default": "Default"}),
"keep_loaded": ("BOOLEAN", {"default": True}),
"prompt_preset": (list(PROMPT_PRESETS.keys()), {"default": "detailed"}),
"custom_prompt": ("STRING", {"default": "", "multiline": True}),
"max_new_tokens": ("INT", {"default": 256, "min": 16, "max": 2048, "step": 8}),
"num_beams": ("INT", {"default": 3, "min": 1, "max": 16, "step": 1}),
"temperature": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 2.0, "step": 0.01}),
"top_p": ("FLOAT", {"default": 0.90, "min": 0.0, "max": 1.0, "step": 0.01}),
"top_k": ("INT", {"default": 50, "min": 0, "max": 500, "step": 1}),
},
"optional": {
"image": ("IMAGE", {}),
"captionforge_run_config": ("CAPTIONFORGE_PIPELINE_PLAN", {}),
},
}
RETURN_TYPES = ("CAPTIONFORGE_PIPELINE_PLAN", "STRING", "STRING", "STRING")
RETURN_NAMES = ("captionforge_run_config_out", "caption", "jsonl_records", "resolved_prompt")
FUNCTION = "caption"
CATEGORY = "JLC/Captioning"
@classmethod
def IS_CHANGED(cls, **kwargs):
return float("NaN")
def caption(self, model, memory_mode, keep_loaded, prompt_preset, custom_prompt, max_new_tokens, num_beams, temperature, top_p, top_k, image=None, captionforge_run_config=None):
run_plan_connected = bool(captionforge_run_config)
prompt = resolve_prompt(prompt_preset, custom_prompt)
run_plan = expand_captionforge_runs(captionforge_run_config, model_key="internvl", widget_captions_per_image=1, widget_seed=-1, widget_temperature=float(temperature), widget_top_p=float(top_p), widget_top_k=int(top_k), widget_max_new_tokens=int(max_new_tokens), widget_max_size=768, widget_trigger_word="", widget_output_dir="", widget_input_path="", widget_recursive=True, widget_filename_glob="*")
if run_plan_connected and not run_plan:
status = "[CaptionForge] JLC InternVL Caption (Lite) disabled by Pipeline Planner."
return (captionforge_run_config, status, "", prompt)
first_run = run_plan[0]
generation = GenerationConfig(max_new_tokens=int(first_run.max_new_tokens), num_beams=int(num_beams), temperature=float(first_run.temperature), top_p=float(first_run.top_p), top_k=int(first_run.top_k), repetition_penalty=1.0, seed=first_run.seed)
cleanup = CleanupConfig(trigger="", prefix=(f"{first_run.trigger_word}," if first_run.trigger_word else ""), suffix="", forbidden_phrases=[], replacement_rules=[])
cfg = InternVLCaptionConfig(model_name=model, model_path="", model_root=str(JLC_INTERVLVL_MODEL_ROOT), memory_mode=memory_mode, dtype="bf16", device="auto", trust_remote_code=True, keep_loaded=bool(keep_loaded), quiet_transformers_load=True, max_size=int(first_run.max_size), prompt=prompt, allow_download=True, use_comfy_model_management=True)
engine = InternVLCaptionEngine(config=cfg, generation=generation, cleanup=cleanup)
direct_images = [(f"comfy_image_{i:04d}", pil) for i, pil in enumerate(_tensor_to_pil(image))]
file_images = _iter_input_path_images(first_run.input_path, first_run.recursive, first_run.filename_glob) if first_run.input_path else []
if not direct_images and not file_images:
raise RuntimeError("No image input found. Connect an IMAGE input or provide input_path in the CaptionForge Run Plan.")
planned_caption_jsonl_path = _planned_caption_jsonl_path(captionforge_run_config) if run_plan_connected else None
output_dir = Path(first_run.output_dir) if first_run.output_dir else Path(folder_paths.get_output_directory()) / "jlc_internvl_caption_lite"
if planned_caption_jsonl_path is not None:
output_dir = planned_caption_jsonl_path.parent
if run_plan_connected:
output_dir.mkdir(parents=True, exist_ok=True)
jsonl_path = planned_caption_jsonl_path or (output_dir / DEFAULT_JSONL_FILENAME)
all_records: list[CaptionRecord] = []
engine.load()
if run_plan_connected:
write_run_config_json(output_dir / f"jlc_internvl_caption_lite_run_config_{timestamp()}.json", engine.build_run_config(), dry_run=False)
def process_one(source_name: str, pil: Image.Image):
for run in run_plan:
engine.generation = GenerationConfig(max_new_tokens=int(run.max_new_tokens), num_beams=int(num_beams), temperature=float(run.temperature), top_p=float(run.top_p), top_k=int(run.top_k), repetition_penalty=1.0, seed=run.seed)
engine.cleanup = CleanupConfig(trigger="", prefix=(f"{run.trigger_word}," if run.trigger_word else ""), suffix="", forbidden_phrases=[], replacement_rules=[])
engine.config.max_size = int(run.max_size); engine.config.prompt = prompt
t0 = time.perf_counter(); final_caption, raw_caption = engine.caption_pil(pil); _ = time.perf_counter() - t0
rec = CaptionRecord(image=source_name, caption=final_caption, raw_caption=raw_caption, model_name=model, model_path=str(engine.local_model_path or ""), prompt=prompt, seed=run.seed, temperature=run.temperature, top_p=run.top_p, top_k=run.top_k, num_beams=int(num_beams), max_new_tokens=run.max_new_tokens, max_size=int(run.max_size), timestamp=datetime.now().isoformat(timespec="seconds"), captionforge_pass="A", model_family="internvl", ensemble_run_index=run.ensemble_run_index, image_key=source_name)
all_records.append(rec)
if run_plan_connected:
append_jsonl_records(jsonl_path, [rec], dry_run=False)
write_text_sidecar(_run_txt_path(output_dir, source_name, len(run_plan), run.ensemble_run_index), rec.caption, overwrite=True, backup_existing=False, dry_run=False)
for source_name, pil in direct_images:
process_one(source_name, pil)
for source_name, path in file_images:
process_one(source_name, _open_image(path))
if not keep_loaded:
engine.unload()
caption_string = "\n\n".join(r.caption for r in all_records if r.status == "ok")
jsonl_string = _make_jsonl_string(all_records)
return (captionforge_run_config, caption_string, jsonl_string, prompt)
NODE_CLASS_MAPPINGS = {"JLC_InternVLCaptionLite": JLC_InternVLCaptionLite}
NODE_DISPLAY_NAME_MAPPINGS = {"JLC_InternVLCaptionLite": " JLC InternVL Caption (Lite)"}
@@ -1,277 +0,0 @@
"""
JLC SmolVLM Caption — ComfyUI Node Wrapper
"""
from __future__ import annotations
MANIFEST = {
"name": "JLC SmolVLM Caption",
"version": (0, 2, 0),
"author": "J. L. Córdova",
"description": "ComfyUI frontend for SmolVLM-family CaptionForge Pass A captioning.",
}
import json
from datetime import datetime
from pathlib import Path
import time
import numpy as np
import torch
from PIL import Image
import folder_paths
from ..engines.jlc_smolvlm_caption_engine import (
BatchCaptionConfig,
CaptionRecord,
CleanupConfig,
SmolVLMCaptionConfig,
SmolVLMCaptionEngine,
GenerationConfig,
MEMORY_MODES,
MODEL_REGISTRY,
PROMPT_PRESETS,
append_jsonl_records,
load_existing_jsonl_images,
probe_registry_model_download,
record_to_json,
resolve_prompt,
timestamp,
write_run_config_json,
write_text_sidecar,
)
from ..engines.captionforge_pipeline_planner_engine import expand_captionforge_runs
from ..engines.captionforge_caption_prompt_kit import (
CAPTION_LENGTH_CHOICES,
CAPTION_TYPE_CHOICES,
build_caption_prompt,
)
from .jlc_captionforge_template_options import (
CAPTIONFORGE_EXTRA_OPTIONS_CHOICES,
resolve_effective_extra_options,
)
JLC_SMOLVLM_MODEL_ROOT = Path(folder_paths.models_dir) / "LLM" / "JLC_SmolVLMCaption"
def _tensor_to_pil(image_tensor) -> list[Image.Image]:
if image_tensor is None:
return []
if isinstance(image_tensor, torch.Tensor):
image_tensor = image_tensor.detach().cpu()
if image_tensor.ndim == 3:
image_tensor = image_tensor.unsqueeze(0)
images: list[Image.Image] = []
for img in image_tensor:
arr = img.numpy()
arr = np.clip(arr * 255.0, 0, 255).astype(np.uint8)
images.append(Image.fromarray(arr).convert("RGB"))
return images
def _parse_forbidden_lines(value: str) -> list[str]:
value = (value or "").strip()
if not value:
return []
return [line.strip() for line in value.splitlines() if line.strip()]
def _parse_replace_pairs(value: str) -> list[tuple[str, str]]:
value = (value or "").strip()
if not value:
return []
rules: list[tuple[str, str]] = []
for line in value.splitlines():
line = line.strip()
if not line or line.startswith("#") or "=>" not in line:
continue
old, new = line.split("=>", 1)
old = old.strip(); new = new.strip()
if old:
rules.append((old, new))
return rules
def _normalize_pipeline_plan(config):
if config is None:
return {}
if isinstance(config, dict):
return dict(config)
if isinstance(config, str):
text = config.strip()
if not text:
return {}
try:
obj = json.loads(text)
return obj if isinstance(obj, dict) else {}
except Exception:
return {}
return {}
def _planned_caption_jsonl_path(captionforge_run_config) -> Path | None:
cfg = _normalize_pipeline_plan(captionforge_run_config)
paths = cfg.get("paths") if isinstance(cfg.get("paths"), dict) else {}
for key in ("caption_jsonl", "pass_a_jsonl"):
value = str(paths.get(key) or "").strip()
if value:
return Path(value)
return None
def _make_jsonl_string(records: list[CaptionRecord]) -> str:
return "\n".join(json.dumps(record_to_json(r), ensure_ascii=False) for r in records)
def _extra_options_from_widgets(*items: str) -> list[str]:
return [str(item).strip() for item in items if str(item).strip()]
class JLC_SmolVLMCaption:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": (list(MODEL_REGISTRY.keys()), {"default": "SmolVLM-Instruct", "tooltip": "Select model."}),
"memory_mode": (list(MEMORY_MODES.keys()), {"default": "Default", "tooltip": "Default uses regular Transformers loading. Balanced (8-bit) uses bitsandbytes when available."}),
"input_path": ("STRING", {"default": "", "multiline": False}),
"recursive": ("BOOLEAN", {"default": True}),
"filename_glob": ("STRING", {"default": "*", "multiline": False}),
"prompt_mode": (["custom_or_preset", "caption_template"], {"default": "caption_template", "tooltip": "custom_or_preset uses custom_prompt or prompt_preset. caption_template uses caption_type, caption_length, and extra options."}),
"prompt_preset": (list(PROMPT_PRESETS.keys()), {"default": "detailed"}),
"caption_type": (CAPTION_TYPE_CHOICES, {"default": "SFW Character Caption" if "SFW Character Caption" in CAPTION_TYPE_CHOICES else CAPTION_TYPE_CHOICES[0]}),
"caption_length": (CAPTION_LENGTH_CHOICES, {"default": "any"}),
"extra_option1": (CAPTIONFORGE_EXTRA_OPTIONS_CHOICES, {"default": ""}),
"extra_option2": (CAPTIONFORGE_EXTRA_OPTIONS_CHOICES, {"default": "", "advanced": True}),
"extra_option3": (CAPTIONFORGE_EXTRA_OPTIONS_CHOICES, {"default": "", "advanced": True}),
"extra_option4": (CAPTIONFORGE_EXTRA_OPTIONS_CHOICES, {"default": "", "advanced": True}),
"extra_option5": (CAPTIONFORGE_EXTRA_OPTIONS_CHOICES, {"default": "", "advanced": True}),
"extra_option6": (CAPTIONFORGE_EXTRA_OPTIONS_CHOICES, {"default": "", "advanced": True}),
"person_name": ("STRING", {"default": "", "multiline": False, "advanced": True}),
"custom_prompt": ("STRING", {"default": "", "multiline": True}),
"max_new_tokens": ("INT", {"default": 256, "min": 16, "max": 2048, "step": 8}),
"num_beams": ("INT", {"default": 3, "min": 1, "max": 16, "step": 1}),
"temperature": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 2.0, "step": 0.01}),
"top_p": ("FLOAT", {"default": 0.90, "min": 0.0, "max": 1.0, "step": 0.01}),
"top_k": ("INT", {"default": 50, "min": 0, "max": 500, "step": 1}),
"repetition_penalty": ("FLOAT", {"default": 1.0, "min": 1.0, "max": 2.0, "step": 0.01}),
"captions_per_image": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}),
"seed": ("INT", {"default": -1, "min": -1, "max": 0xFFFFFFFF, "step": 1}),
"max_size": ("INT", {"default": 768, "min": 0, "max": 4096, "step": 64}),
"output_dir": ("STRING", {"default": "", "multiline": False}),
"write_txt": ("BOOLEAN", {"default": True}),
"write_jsonl": ("BOOLEAN", {"default": False}),
"also_jsonl": ("BOOLEAN", {"default": False}),
"write_run_config": ("BOOLEAN", {"default": True}),
"jsonl_filename": ("STRING", {"default": "captions.jsonl", "multiline": False}),
"overwrite": ("BOOLEAN", {"default": False}),
"backup_existing": ("BOOLEAN", {"default": True}),
"dry_run": ("BOOLEAN", {"default": False}),
"limit": ("INT", {"default": 0, "min": 0, "max": 100000, "step": 1}),
"skip_existing_txt": ("BOOLEAN", {"default": True}),
"skip_existing_jsonl_images": ("BOOLEAN", {"default": False}),
"prefix": ("STRING", {"default": "", "multiline": False}),
"suffix": ("STRING", {"default": "", "multiline": False}),
"forbidden_phrases": ("STRING", {"default": "", "multiline": True}),
"replace_pairs": ("STRING", {"default": "", "multiline": True}),
"keep_loaded": ("BOOLEAN", {"default": True}),
"download_probe_only": ("BOOLEAN", {"default": False}),
},
"optional": {
"image": ("IMAGE", {}),
"captionforge_run_config": ("CAPTIONFORGE_PIPELINE_PLAN", {}),
"extra_options": ("CAPTIONFORGE_EXTRA_OPTIONS", {"tooltip": "Optional shared CaptionForge Extra Options payload. Overrides local extra_option widgets when non-empty."}),
},
}
RETURN_TYPES = ("CAPTIONFORGE_PIPELINE_PLAN", "STRING", "STRING", "STRING")
RETURN_NAMES = ("captionforge_run_config_out", "caption", "jsonl_records", "resolved_prompt")
FUNCTION = "caption"
CATEGORY = "JLC/Captioning"
@classmethod
def IS_CHANGED(cls, **kwargs):
return float("NaN")
def caption(self, model, memory_mode, input_path, recursive, filename_glob, prompt_mode, prompt_preset, caption_type, caption_length, extra_option1, extra_option2, extra_option3, extra_option4, extra_option5, extra_option6, person_name, custom_prompt, max_new_tokens, num_beams, temperature, top_p, top_k, repetition_penalty, captions_per_image, seed, max_size, output_dir, write_txt, write_jsonl, also_jsonl, write_run_config, jsonl_filename, overwrite, backup_existing, dry_run, limit, skip_existing_txt, skip_existing_jsonl_images, prefix, suffix, forbidden_phrases, replace_pairs, keep_loaded, download_probe_only, image=None, captionforge_run_config=None, extra_options=None):
if download_probe_only:
result = probe_registry_model_download(model, JLC_SMOLVLM_MODEL_ROOT)
return (captionforge_run_config, result, "", "")
local_extra_options = _extra_options_from_widgets(extra_option1, extra_option2, extra_option3, extra_option4, extra_option5, extra_option6)
effective_extra_options, effective_person_name, extra_options_metadata = resolve_effective_extra_options(
payload=extra_options,
local_options=local_extra_options,
local_name=person_name,
)
if prompt_mode == "caption_template":
prompt = build_caption_prompt(caption_type, caption_length, effective_extra_options, effective_person_name, dialect="smolvlm")
else:
prompt = resolve_prompt(prompt_preset, custom_prompt)
run_plan = expand_captionforge_runs(captionforge_run_config, model_key="smolvlm", widget_captions_per_image=int(captions_per_image), widget_seed=int(seed), widget_temperature=float(temperature), widget_top_p=float(top_p), widget_top_k=int(top_k), widget_max_new_tokens=int(max_new_tokens), widget_max_size=int(max_size), widget_trigger_word="", widget_output_dir=output_dir, widget_input_path=input_path, widget_recursive=bool(recursive), widget_filename_glob=filename_glob)
run_plan_connected = bool(captionforge_run_config)
if run_plan_connected and not run_plan:
status = "[CaptionForge] JLC SmolVLM Caption disabled by Pipeline Planner."
return (captionforge_run_config, status, "", prompt)
first_run = run_plan[0]
if run_plan_connected:
write_jsonl = True; also_jsonl = False; skip_existing_txt = False; skip_existing_jsonl_images = False; overwrite = True
effective_prefix = prefix
if first_run.trigger_word:
effective_prefix = f"{first_run.trigger_word}, {prefix}".strip(", ")
generation = GenerationConfig(max_new_tokens=int(first_run.max_new_tokens), num_beams=int(num_beams), temperature=float(first_run.temperature), top_p=float(first_run.top_p), top_k=int(first_run.top_k), repetition_penalty=float(repetition_penalty), seed=first_run.seed)
cleanup = CleanupConfig(trigger="", prefix=effective_prefix, suffix=suffix, forbidden_phrases=_parse_forbidden_lines(forbidden_phrases), replacement_rules=_parse_replace_pairs(replace_pairs))
cfg = SmolVLMCaptionConfig(model_name=model, model_path="", model_root=str(JLC_SMOLVLM_MODEL_ROOT), memory_mode=memory_mode, dtype="bf16", device="auto", trust_remote_code=True, keep_loaded=bool(keep_loaded), quiet_transformers_load=True, max_size=int(first_run.max_size), prompt=prompt, allow_download=True, use_comfy_model_management=True)
engine = SmolVLMCaptionEngine(config=cfg, generation=generation, cleanup=cleanup)
input_path = (input_path or "").strip(); output_dir = (output_dir or "").strip(); jsonl_filename = (jsonl_filename or "captions.jsonl").strip() or "captions.jsonl"
planned_caption_jsonl_path = _planned_caption_jsonl_path(captionforge_run_config) if run_plan_connected else None
if first_run.output_dir: output_dir = first_run.output_dir
if first_run.input_path: input_path = first_run.input_path
if run_plan_connected:
recursive = bool(first_run.recursive); filename_glob = first_run.filename_glob or "*"
if planned_caption_jsonl_path is not None:
output_dir = str(planned_caption_jsonl_path.parent); jsonl_filename = planned_caption_jsonl_path.name
use_jsonl = bool(write_jsonl or also_jsonl)
all_records: list[CaptionRecord] = []
pil_images = _tensor_to_pil(image)
if pil_images:
engine.load()
fallback_dir = Path(folder_paths.get_output_directory()) / "jlc_smolvlm_caption"
image_output_dir = Path(output_dir) if output_dir else fallback_dir
image_output_dir.mkdir(parents=True, exist_ok=True)
jsonl_path = image_output_dir / jsonl_filename
seen_jsonl = load_existing_jsonl_images(jsonl_path) if use_jsonl and skip_existing_jsonl_images else set()
if write_run_config and not input_path:
write_run_config_json(image_output_dir / f"jlc_smolvlm_caption_run_config_{timestamp()}.json", engine.build_run_config(), dry_run=bool(dry_run))
for idx, pil in enumerate(pil_images):
source_name = f"comfy_image_{idx:04d}"
txt_path = image_output_dir / f"{source_name}.txt"
if len(run_plan) == 1 and skip_existing_txt and write_txt and txt_path.exists() and not overwrite:
continue
if use_jsonl and skip_existing_jsonl_images and source_name in seen_jsonl:
continue
for run in run_plan:
engine.generation = GenerationConfig(max_new_tokens=int(run.max_new_tokens), num_beams=int(num_beams), temperature=float(run.temperature), top_p=float(run.top_p), top_k=int(run.top_k), repetition_penalty=float(repetition_penalty), seed=run.seed)
eff_prefix = prefix
if run.trigger_word:
eff_prefix = f"{run.trigger_word}, {prefix}".strip(", ")
engine.cleanup = CleanupConfig(trigger="", prefix=eff_prefix, suffix=suffix, forbidden_phrases=_parse_forbidden_lines(forbidden_phrases), replacement_rules=_parse_replace_pairs(replace_pairs))
engine.config.max_size = int(run.max_size); engine.config.prompt = prompt
run_txt_path = txt_path if len(run_plan) <= 1 else image_output_dir / f"{source_name}__cf_run_{run.ensemble_run_index:02d}.txt"
t0 = time.perf_counter(); final_caption, raw_caption = engine.caption_pil(pil); _ = time.perf_counter() - t0
rec = CaptionRecord(image=source_name, caption=final_caption, raw_caption=raw_caption, model_name=model, model_path=str(engine.local_model_path or ""), prompt=prompt, seed=run.seed, temperature=run.temperature, top_p=run.top_p, top_k=run.top_k, num_beams=int(num_beams), max_new_tokens=run.max_new_tokens, max_size=int(run.max_size), timestamp=datetime.now().isoformat(timespec="seconds"), captionforge_pass="A", model_family="smolvlm", ensemble_run_index=run.ensemble_run_index, image_key=source_name)
all_records.append(rec)
if use_jsonl:
append_jsonl_records(jsonl_path, [rec], dry_run=bool(dry_run))
if write_txt:
write_text_sidecar(run_txt_path, rec.caption, overwrite=bool(overwrite), backup_existing=bool(backup_existing), dry_run=bool(dry_run))
batch_result = None
if input_path:
also_jsonl_path = ""
if also_jsonl and not write_jsonl:
base_dir = Path(output_dir) if output_dir else (Path(input_path).parent if Path(input_path).is_file() else Path(input_path))
also_jsonl_path = str(base_dir / jsonl_filename)
batch = BatchCaptionConfig(input_path=input_path, recursive=bool(recursive), filename_glob=(filename_glob or "*").strip() or "*", output_dir=output_dir, write_txt=bool(write_txt), write_jsonl=bool(write_jsonl), jsonl_filename=jsonl_filename, also_jsonl_path=also_jsonl_path, write_run_config=bool(write_run_config), overwrite=bool(overwrite), backup_existing=bool(backup_existing), dry_run=bool(dry_run), limit=int(limit), skip_existing_txt=bool(skip_existing_txt), skip_existing_jsonl_images=bool(skip_existing_jsonl_images))
batch_result = engine.caption_batch(batch)
all_records.extend(batch_result.records)
if not pil_images and batch_result is None:
raise RuntimeError("No image input found. Connect an IMAGE input or provide input_path.")
if not keep_loaded:
engine.unload()
caption_string = "\n\n".join(r.caption for r in all_records if r.status == "ok")
jsonl_string = _make_jsonl_string(all_records)
return (captionforge_run_config, caption_string, jsonl_string, prompt)
NODE_CLASS_MAPPINGS = {"JLC_SmolVLMCaption": JLC_SmolVLMCaption}
NODE_DISPLAY_NAME_MAPPINGS = {"JLC_SmolVLMCaption": " JLC SmolVLM Caption"}
@@ -1,770 +0,0 @@
"""
JLC CaptionForge SmolVLM (Lite) — ComfyUI Node Wrapper
- CaptionForge
- This node is part of **CaptionForge**, a model-agnostic captioning
framework for ComfyUI developed by **J. L. Córdova**.
- Repository
https://github.com/Damkohler/CaptionForge
- CaptionForge focuses on practical dataset-captioning infrastructure for:
• LoRA dataset preparation
• multi-engine caption generation
• JSONL audit trails
• claim extraction and refinement
• consensus-oriented caption improvement
- Node Purpose
- The **JLC CaptionForge SmolVLM (Lite)** node provides a compact ComfyUI
frontend for SmolVLM-family image captioning inside CaptionForge.
- This file is the **ComfyUI-facing wrapper**, not the reusable captioning
engine. It is responsible for:
• compact ComfyUI INPUT_TYPES / widget definitions
• optional direct IMAGE tensor input
• IMAGE tensor conversion to PIL images
• optional Pipeline Planner file/folder `input_path` routing
• fixed model-root integration under `models/LLM/`
• minimal model and memory-mode controls
• clear template-vs-custom prompt controls
• CaptionForge Pipeline Planner consumption through `pipeline_plan`
• CaptionForge Template Options consumption through `template_options`
• TXT audit sidecar writing in planned runs
• shared JSONL audit output in planned runs
• ComfyUI output strings for captions and JSONL records
• node display name, category, and mapping registration
• passing user settings into the shared SmolVLM caption engine
- The actual reusable captioning implementation lives in:
jlc_smolvlm_caption_engine.py
That engine handles:
• SmolVLM model registry lookup
• Hugging Face model probing/downloading
• processor and model loading
• CaptionForge shared model-cache integration
• ComfyUI model-management bridge behavior
• generation settings
• caption cleanup
• TXT sidecar writing
• JSONL audit output
• run-configuration export
- CaptionForge Pipeline Role
- This node participates as a **raw-caption witness** in the CaptionForge
pipeline.
- Raw-caption witnesses generate auditable caption evidence records from one
or more captioning engines.
- The SmolVLM Lite node contributes compact SFW-oriented caption evidence
compatible with downstream CaptionForge claim extraction, semantic
synthesis, and final caption construction.
- CaptionForge audit fields include:
• captionforge_pass
• model_family
• ensemble_run_index
• image_key
• raw_caption
• final cleaned caption
• generation parameters
• prompt metadata
• model metadata
- Node Workflow Model
- The node supports two image sources:
• direct ComfyUI IMAGE input
• file/folder routing supplied by the CaptionForge Pipeline Planner
- These may be used independently or together.
- In standalone mode, the node behaves as a compact direct-IMAGE captioning
node and writes to:
ComfyUI/output/jlc_captionforge_smolvlm_lite/
- When connected to the CaptionForge Pipeline Planner, this node consumes
the shared CAPTIONFORGE_PIPELINE_PLAN through the `pipeline_plan` pin and
uses the SmolVLM model-family plan to determine:
• whether SmolVLM captioning participates in the run
• how many SmolVLM raw-caption witnesses to generate per image
• shared output directory
• optional shared input path
• recursive folder traversal
• filename filtering
• per-caption-instance seed values
• per-caption-instance sampling values
• shared max image size
• shared max token budget
• shared LoRA trigger word
- The node can write:
• TXT audit sidecar captions
• JSONL audit records
• run-configuration JSON files
• direct ComfyUI string outputs
- Prompting Model
- The Lite node supports two explicit prompt modes:
• caption_template_mode
• custom_prompt_mode
- `caption_template_mode` uses caption_type, caption_length, and optional
Template Options supplied by the `template_options` pin.
- `custom_prompt_mode` uses custom_prompt when provided, otherwise a local
prompt_preset fallback.
- If both booleans are enabled, custom_prompt_mode intentionally takes
precedence. If both are disabled, the node falls back to template mode.
- Local extra-option widgets have been removed. Shared template modifiers
now live only in the CaptionForge Template Options sidecar node.
- Model and Dependency Notes
- SmolVLM repositories are multi-file Hugging Face model directories, not
single checkpoint files.
- Runtime behavior depends on the active ComfyUI Python environment,
PyTorch, Transformers, Accelerate, Pillow, and Hugging Face Hub
dependencies.
- Design Philosophy
- The Lite node is intended to be a compact production-friendly raw-caption
witness node.
- It keeps model-family configuration local to the caption node while taking
shared workflow policy from the CaptionForge Pipeline Planner.
- CaptionForge is engine-democratic: SmolVLM captions contribute evidence to
the broader audit and consensus pipeline alongside Joy, Qwen, future local
VLMs, cleanup LLMs, validators, or other robustness engines.
- The node prioritizes reproducibility, auditability, low UI clutter, and
clean separation between the ComfyUI interface and the shared
model-specific captioning engine.
- ⚠️ Development Status
- This is early CaptionForge raw-caption witness infrastructure.
- The UI, model registry, prompt behavior, and output audit fields may
evolve as the multi-pass CaptionForge pipeline matures.
- The node is intended for local dataset preparation and controlled caption
audit workflows.
- Attribution & License
- Concept and implementation by **J. L. Córdova**
with development assistance from **ChatGPT (OpenAI)**.
- CaptionForge's template-option workflow is locally adapted and was inspired
in part by the practical template interface pattern used by the public
JoyCaption Beta One Hugging Face Space.
- SmolVLM-family model loading is designed around compatible Hugging Face
Transformers interfaces and publicly available SmolVLM-family checkpoints.
- Designed for use with:
https://github.com/comfyanonymous/ComfyUI
- Copyright (c) 2026 J. L. Córdova
- Released under the **MIT License**.
"""
from __future__ import annotations
MANIFEST = {
"name": "JLC CaptionForge SmolVLM (Lite)",
"version": (0, 3, 0),
"author": "J. L. Córdova",
"description": (
"Compact CaptionForge frontend for SmolVLM-family image captioning. Provides "
"direct IMAGE captioning while also consuming CAPTIONFORGE_PIPELINE_PLAN objects "
"from the CaptionForge Pipeline Planner through the pipeline_plan input. Template "
"extras are no longer duplicated as local widgets; they are consumed only through "
"the template_options input from the CaptionForge Template Options sidecar. The UI "
"now separates caption_template_mode and custom_prompt_mode for clearer prompt routing "
"while delegating reusable SmolVLM-family model loading, generation, cleanup, cache "
"integration, and audit-record creation to jlc_smolvlm_caption_engine.py."
),
}
import json
from datetime import datetime
from pathlib import Path
import time
import numpy as np
import torch
from PIL import Image
import folder_paths
from ...engines.jlc_smolvlm_caption_engine import (
CaptionRecord,
CleanupConfig,
SmolVLMCaptionConfig,
SmolVLMCaptionEngine,
GenerationConfig,
MEMORY_MODES,
MODEL_REGISTRY,
PROMPT_PRESETS,
append_jsonl_records,
record_to_json,
resolve_prompt,
timestamp,
write_run_config_json,
write_text_sidecar,
)
from ...engines.captionforge_pipeline_planner_engine import expand_captionforge_runs
from ...engines.captionforge_caption_prompt_kit import (
CAPTION_LENGTH_CHOICES,
CAPTION_TYPE_CHOICES,
build_caption_prompt,
)
from ..jlc_captionforge_template_options import resolve_effective_extra_options
JLC_SMOLVLM_MODEL_ROOT = Path(folder_paths.models_dir) / "LLM" / "JLC_SmolVLMCaption"
DEFAULT_JSONL_FILENAME = "captions.jsonl"
_SUPPORTED_IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff"}
def _tensor_to_pil(image_tensor) -> list[Image.Image]:
"""Convert ComfyUI IMAGE tensor [B,H,W,C] float 0..1 into PIL RGB images."""
if image_tensor is None:
return []
if isinstance(image_tensor, torch.Tensor):
image_tensor = image_tensor.detach().cpu()
if image_tensor.ndim == 3:
image_tensor = image_tensor.unsqueeze(0)
images: list[Image.Image] = []
for img in image_tensor:
arr = img.numpy()
arr = np.clip(arr * 255.0, 0, 255).astype(np.uint8)
images.append(Image.fromarray(arr).convert("RGB"))
return images
def _safe_source_name(value: str) -> str:
cleaned = []
for ch in value.replace("\\", "/"):
if ch.isalnum() or ch in {"-", "_", "."}:
cleaned.append(ch)
elif ch == "/":
cleaned.append("__")
else:
cleaned.append("_")
out = "".join(cleaned).strip("._")
return out or "image"
def _iter_input_path_images(input_path: str, recursive: bool, filename_glob: str) -> list[tuple[str, Path]]:
root = Path(str(input_path or "").strip())
if not root:
return []
if not root.exists():
raise RuntimeError(f"CaptionForge input_path does not exist: {root}")
glob_text = (filename_glob or "*").strip() or "*"
if root.is_file():
if root.suffix.lower() not in _SUPPORTED_IMAGE_SUFFIXES:
raise RuntimeError(f"CaptionForge input_path is not a supported image file: {root}")
return [(_safe_source_name(root.stem), root)]
pattern_iter = root.rglob(glob_text) if recursive else root.glob(glob_text)
paths = sorted(
p for p in pattern_iter
if p.is_file() and p.suffix.lower() in _SUPPORTED_IMAGE_SUFFIXES
)
items: list[tuple[str, Path]] = []
for path in paths:
rel = path.relative_to(root).with_suffix("")
items.append((_safe_source_name(str(rel)), path))
return items
def _open_image(path: Path) -> Image.Image:
with Image.open(path) as im:
return im.convert("RGB")
def _normalize_pipeline_plan(config):
if config is None:
return {}
if isinstance(config, dict):
return dict(config)
if isinstance(config, str):
text = config.strip()
if not text:
return {}
try:
obj = json.loads(text)
return obj if isinstance(obj, dict) else {}
except Exception:
return {}
return {}
def _planned_caption_jsonl_path(pipeline_plan) -> Path | None:
"""Return Planner-owned shared Pass A JSONL path, if present."""
cfg = _normalize_pipeline_plan(pipeline_plan)
paths = cfg.get("paths") if isinstance(cfg.get("paths"), dict) else {}
for key in ("caption_jsonl", "pass_a_jsonl"):
value = str(paths.get(key) or "").strip()
if value:
return Path(value)
return None
def _make_jsonl_string(records: list[CaptionRecord]) -> str:
return "\n".join(json.dumps(record_to_json(r), ensure_ascii=False) for r in records)
def _run_txt_path(output_dir: Path, source_name: str, run_count: int, run_index: int) -> Path:
if run_count <= 1:
return output_dir / f"{source_name}.txt"
return output_dir / f"{source_name}__cf_run_{run_index:02d}.txt"
def _use_template_mode(caption_template_mode: bool, custom_prompt_mode: bool) -> bool:
if bool(custom_prompt_mode):
return False
if bool(caption_template_mode):
return True
print("[CaptionForge] No caption mode was selected; falling back to caption_template_mode.")
return True
def _resolve_template_options(template_options):
options, name_input, metadata = resolve_effective_extra_options(
payload=template_options,
local_options=[],
local_name="",
)
return options, name_input, metadata
class JLC_SmolVLMCaptionLite:
"""Compact SmolVLM witness node for CaptionForge."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": (
list(MODEL_REGISTRY.keys()),
{
"default": "SmolVLM-Instruct",
"tooltip": (
"Select the SmolVLM-family vision-language model. Models are loaded from "
"ComfyUI/models/LLM/JLC_SmolVLMCaption/."
),
},
),
"memory_mode": (
list(MEMORY_MODES.keys()),
{
"default": "Default",
"tooltip": "SmolVLM model memory mode.",
},
),
"keep_loaded": (
"BOOLEAN",
{
"default": True,
"tooltip": (
"Keep the model cached after captioning for faster repeated runs. "
"CaptionForge cache policy may still evict it when another caption "
"model must load."
),
},
),
"caption_template_mode": (
"BOOLEAN",
{
"default": True,
"tooltip": (
"Use the structured CaptionForge template path: caption_type, "
"caption_length, and optional Template Options from the template_options pin. "
"If custom_prompt_mode is also enabled, custom_prompt_mode takes precedence."
),
},
),
"caption_type": (
CAPTION_TYPE_CHOICES,
{
"default": "SFW Character Caption" if "SFW Character Caption" in CAPTION_TYPE_CHOICES else CAPTION_TYPE_CHOICES[0],
"tooltip": "Caption template style used when caption_template_mode is active.",
},
),
"caption_length": (
CAPTION_LENGTH_CHOICES,
{
"default": "any",
"tooltip": "Target caption length used when caption_template_mode is active.",
},
),
"custom_prompt_mode": (
"BOOLEAN",
{
"default": False,
"tooltip": (
"Use custom_prompt when non-empty, otherwise use prompt_preset. "
"This overrides caption_template_mode when both toggles are enabled."
),
},
),
"prompt_preset": (
list(PROMPT_PRESETS.keys()),
{
"default": "detailed" if "detailed" in PROMPT_PRESETS else list(PROMPT_PRESETS.keys())[0],
"tooltip": "Built-in prompt preset used only in custom_prompt_mode when custom_prompt is blank.",
},
),
"custom_prompt": (
"STRING",
{
"default": "",
"multiline": True,
"tooltip": "Custom prompt used only when custom_prompt_mode is enabled. Overrides prompt_preset when non-empty.",
},
),
"max_new_tokens": (
"INT",
{
"default": 256,
"min": 16,
"max": 2048,
"step": 8,
"tooltip": (
"Standalone token budget. When a CaptionForge Pipeline Planner is connected, "
"this is overridden by the Pipeline Planner's shared max_new_tokens."
),
},
),
"num_beams": (
"INT",
{
"default": 3,
"min": 1,
"max": 16,
"step": 1,
"tooltip": "Standalone beam count for SmolVLM generation.",
},
),
"temperature": (
"FLOAT",
{
"default": 0.0,
"min": 0.0,
"max": 2.0,
"step": 0.01,
"tooltip": (
"Standalone sampling temperature. When a CaptionForge Pipeline Planner is "
"connected, this is overridden by the Pipeline Planner temperature schedule."
),
},
),
"top_p": (
"FLOAT",
{
"default": 0.90,
"min": 0.0,
"max": 1.0,
"step": 0.01,
"tooltip": (
"Standalone top-p sampling value. When a CaptionForge Pipeline Planner is "
"connected, this is overridden by the Pipeline Planner top-p schedule."
),
},
),
"top_k": (
"INT",
{
"default": 50,
"min": 0,
"max": 500,
"step": 1,
"tooltip": (
"Standalone top-k sampling limit. When a CaptionForge Pipeline Planner is "
"connected, this is overridden by the Pipeline Planner top-k schedule."
),
},
),
},
"optional": {
"image": (
"IMAGE",
{
"tooltip": (
"Image or batch of images to caption. When a CaptionForge Pipeline Planner "
"also supplies input_path, both sources are captioned."
),
},
),
"pipeline_plan": (
"CAPTIONFORGE_PIPELINE_PLAN",
{
"tooltip": (
"Connect the CaptionForge Pipeline Planner output here. When connected, "
"this node switches into Pass A evidence mode: shared input_path/output_dir, "
"per-run seeds and schedules, TXT audit sidecars, common captions.jsonl, "
"and run config export."
),
},
),
"template_options": (
"CAPTIONFORGE_EXTRA_OPTIONS",
{
"tooltip": (
"Connect the CaptionForge Template Options node here. This supplies shared "
"caption-template modifiers and optional name text. Local extra-option widgets "
"were removed so template modifiers have one source of truth."
),
},
),
"seed": (
"INT",
{
"forceInput": True,
"tooltip": "Optional standalone seed input. Ignored when a Pipeline Planner supplies a seed schedule.",
},
),
},
}
RETURN_TYPES = ("CAPTIONFORGE_PIPELINE_PLAN", "STRING", "STRING", "STRING")
RETURN_NAMES = ("pipeline_plan_out", "caption", "jsonl_records", "resolved_prompt")
FUNCTION = "caption"
CATEGORY = "JLC/CaptionForge/Caption Nodes"
@classmethod
def IS_CHANGED(cls, **kwargs):
return float("NaN")
def caption(
self,
model,
memory_mode,
keep_loaded,
caption_template_mode,
caption_type,
caption_length,
custom_prompt_mode,
prompt_preset,
custom_prompt,
max_new_tokens,
num_beams,
temperature,
top_p,
top_k,
image=None,
pipeline_plan=None,
template_options=None,
seed=None,
):
run_plan_connected = bool(pipeline_plan)
use_caption_template = _use_template_mode(caption_template_mode, custom_prompt_mode)
effective_extra_options, effective_person_name, _ = _resolve_template_options(template_options)
if use_caption_template:
prompt = build_caption_prompt(
caption_type,
caption_length,
effective_extra_options,
effective_person_name,
dialect="smolvlm",
)
else:
prompt = resolve_prompt(prompt_preset, custom_prompt)
effective_seed = -1 if seed is None else int(seed)
run_plan = expand_captionforge_runs(
pipeline_plan,
model_key="smolvlm",
widget_captions_per_image=1,
widget_seed=effective_seed,
widget_temperature=float(temperature),
widget_top_p=float(top_p),
widget_top_k=int(top_k),
widget_max_new_tokens=int(max_new_tokens),
widget_max_size=768,
widget_trigger_word="",
widget_output_dir="",
widget_input_path="",
widget_recursive=True,
widget_filename_glob="*",
)
if run_plan_connected and not run_plan:
status = "[CaptionForge] SmolVLM Caption Lite disabled by Pipeline Planner."
print(status)
return (pipeline_plan, status, "", prompt)
first_run = run_plan[0]
generation = GenerationConfig(
max_new_tokens=int(first_run.max_new_tokens),
num_beams=int(num_beams),
temperature=float(first_run.temperature),
top_p=float(first_run.top_p),
top_k=int(first_run.top_k),
repetition_penalty=1.0,
seed=first_run.seed,
)
cleanup = CleanupConfig(
trigger="",
prefix=(f"{first_run.trigger_word}," if first_run.trigger_word else ""),
suffix="",
forbidden_phrases=[],
replacement_rules=[],
)
cfg = SmolVLMCaptionConfig(
model_name=model,
model_path="",
model_root=str(JLC_SMOLVLM_MODEL_ROOT),
memory_mode=memory_mode,
dtype="bf16",
device="auto",
trust_remote_code=True,
keep_loaded=bool(keep_loaded),
quiet_transformers_load=True,
max_size=int(first_run.max_size),
prompt=prompt,
allow_download=True,
use_comfy_model_management=True,
)
engine = SmolVLMCaptionEngine(config=cfg, generation=generation, cleanup=cleanup)
direct_images = [(f"comfy_image_{i:04d}", pil) for i, pil in enumerate(_tensor_to_pil(image))]
file_images = (
_iter_input_path_images(first_run.input_path, first_run.recursive, first_run.filename_glob)
if first_run.input_path
else []
)
if direct_images and file_images:
print(
"[CaptionForge] IMAGE input and Pipeline Planner input_path are both active; "
"captioning both sources."
)
if not direct_images and not file_images:
raise RuntimeError(
"No image input found. Connect an IMAGE input or provide input_path in the "
"CaptionForge Pipeline Planner."
)
planned_caption_jsonl_path = _planned_caption_jsonl_path(pipeline_plan) if run_plan_connected else None
output_dir = Path(first_run.output_dir) if first_run.output_dir else Path(folder_paths.get_output_directory()) / "jlc_captionforge_smolvlm_lite"
if planned_caption_jsonl_path is not None:
output_dir = planned_caption_jsonl_path.parent
if run_plan_connected:
output_dir.mkdir(parents=True, exist_ok=True)
jsonl_path = planned_caption_jsonl_path or (output_dir / DEFAULT_JSONL_FILENAME)
all_records: list[CaptionRecord] = []
engine.load()
if run_plan_connected:
write_run_config_json(
output_dir / f"jlc_captionforge_smolvlm_lite_run_config_{timestamp()}.json",
engine.build_run_config(),
dry_run=False,
)
def process_one(source_name: str, pil: Image.Image):
for run in run_plan:
engine.generation = GenerationConfig(
max_new_tokens=int(run.max_new_tokens),
num_beams=int(num_beams),
temperature=float(run.temperature),
top_p=float(run.top_p),
top_k=int(run.top_k),
repetition_penalty=1.0,
seed=run.seed,
)
engine.cleanup = CleanupConfig(
trigger="",
prefix=(f"{run.trigger_word}," if run.trigger_word else ""),
suffix="",
forbidden_phrases=[],
replacement_rules=[],
)
engine.config.max_size = int(run.max_size)
engine.config.prompt = prompt
t0 = time.perf_counter()
final_caption, raw_caption = engine.caption_pil(pil)
dt = time.perf_counter() - t0
print(f"[JLC CaptionForge SmolVLM Lite] Generation time run {run.ensemble_run_index}: {dt:.2f}s")
rec = CaptionRecord(
image=source_name,
caption=final_caption,
raw_caption=raw_caption,
model_name=model,
model_path=str(engine.local_model_path or ""),
prompt=prompt,
seed=run.seed,
temperature=run.temperature,
top_p=run.top_p,
top_k=run.top_k,
num_beams=int(num_beams),
max_new_tokens=run.max_new_tokens,
max_size=int(run.max_size),
timestamp=datetime.now().isoformat(timespec="seconds"),
captionforge_pass="A",
model_family="smolvlm",
ensemble_run_index=run.ensemble_run_index,
image_key=source_name,
)
all_records.append(rec)
if run_plan_connected:
append_jsonl_records(jsonl_path, [rec], dry_run=False)
write_text_sidecar(
_run_txt_path(output_dir, source_name, len(run_plan), run.ensemble_run_index),
rec.caption,
overwrite=True,
backup_existing=False,
dry_run=False,
)
print(
f"[JLC CaptionForge SmolVLM Lite] Captioned {source_name} "
f"run {run.ensemble_run_index + 1}/{len(run_plan)}"
)
for source_name, pil in direct_images:
process_one(source_name, pil)
for source_name, path in file_images:
process_one(source_name, _open_image(path))
if not keep_loaded:
engine.unload()
caption_string = "\n\n".join(r.caption for r in all_records if r.status == "ok")
jsonl_string = _make_jsonl_string(all_records)
return (pipeline_plan, caption_string, jsonl_string, prompt)
NODE_CLASS_MAPPINGS = {
"JLC_SmolVLMCaptionLite": JLC_SmolVLMCaptionLite,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"JLC_SmolVLMCaptionLite": "\u2003JLC CaptionForge SmolVLM (Lite)",
}
File diff suppressed because it is too large Load Diff
+395 -101
View File
@@ -8,11 +8,9 @@ JLC CaptionForge Pipeline Planner — ComfyUI Node Wrapper
- Repository
https://github.com/Damkohler/CaptionForge
- CaptionForge focuses on practical dataset-captioning infrastructure for
LoRA dataset preparation, using multi-engine caption generation, JSONL
audit trails, claim extraction and refinement, text-LLM distillation,
image-aware VLM validation, and consensus-oriented caption improvement
to produce grounded, auditable training captions.
- CaptionForge 1.0 uses independent Pass-A witnesses, text-LLM synthesis,
image-aware validation, SHORT/TAGGY formatting, and JSONL audit trails to
produce grounded LoRA dataset captions.
- Node Purpose
- The **JLC CaptionForge Pipeline Planner** is the ordinary-run control
@@ -24,11 +22,12 @@ JLC CaptionForge Pipeline Planner — ComfyUI Node Wrapper
• optional IMAGE passthrough for quick single-image workflows
• shared input path, recursion, and filename-glob routing
• output folder and run-name policy
• LoRA trigger word and user caption anchor routing
• LoRA trigger word and persistent caption/training anchor routing
• raw-caption run counts for Joy, Qwen, and generic Ollama Caption nodes
• caption seed, sampling, image-size, and token policy
• Distiller model/settings selection
• Validator model/settings selection
• SHORT/TAGGY Formatter model/settings selection
• final export policy
• JSON serialization of the plan for audit/debugging
@@ -36,32 +35,38 @@ JLC CaptionForge Pipeline Planner — ComfyUI Node Wrapper
captionforge_pipeline_planner_engine.py
- Ollama Model Dropdowns
- Distiller and Validator dropdown values are explicit Ollama model tags.
- Distiller, Validator, and Formatter dropdown values are explicit Ollama
model tags.
- Caption-stage Ollama models are intentionally selected on each
**JLC CaptionForge Ollama Caption** node, not in this Planner.
- Dropdown choices for Distiller/Validator are loaded at node-import time
- Dropdown choices for Distiller/Validator/Formatter are loaded at node-import time
from:
config/captionforge_ollama_models.json
- If the JSON file is missing or malformed, the node falls back to:
Distiller: llama3.1:8b
Validator: gemma4:e4b
Distiller: mistral-small:24b
Validator: gemma4:26b
Formatter: mistral-small:24b
- The optional custom choice lets users enter any installed Ollama model tag
without editing Python.
- CaptionForge Pipeline Role
- The planner emits the CAPTIONFORGE_PIPELINE_PLAN consumed by caption
nodes and the JLC CaptionForge capstone node.
nodes and the JLC CaptionForge Orchestrator.
- The canonical graph flow is:
Pipeline Planner
-> Joy/Qwen/Ollama raw-caption nodes
-> JLC CaptionForge capstone node
-> Distiller Engine
-> VLM Validator Engine
-> final deterministic TXT/JSONL export
-> Joy/Qwen/Ollama Pass A caption witness nodes
-> JLC CaptionForge Orchestrator
-> Pass B fat draft (text Ollama call)
-> Pass C image-aware validator (Ollama VLM call)
-> Pass D SHORT + TAGGY formatter (text Ollama call)
-> final deterministic TXT/JSONL export
- SmolVLM is not exposed in the current mainline Planner UI. It may remain
available as a standalone/experimental node and can be revisited later.
- The standalone CLI-oriented distiller and validator engines are
prototype/reference implementations, not the production runtime path.
- Historical SmolVLM experiments are not registered or exposed by the
CaptionForge 1.0 production package.
- Design Philosophy
- CaptionForge is an original concept and implementation, not derived from
@@ -75,10 +80,10 @@ JLC CaptionForge Pipeline Planner — ComfyUI Node Wrapper
low UI ambiguity, and clean separation between ComfyUI UI and reusable
pipeline logic.
- ⚠️ Development Status
- This is release-candidate CaptionForge infrastructure.
- Widget names, output schema details, and downstream validation strategy may
evolve as CaptionForge matures.
- Production Status
- This is the active CaptionForge 1.0 project-level control surface. When
connected, Planner values override corresponding Orchestrator controls; the
frozen production defaults are aligned between both nodes.
- Attribution & License
- Concept and implementation by **J. L. Córdova**
@@ -99,10 +104,10 @@ MANIFEST = {
"description": (
"ComfyUI-facing Pipeline Planner node for CaptionForge. Builds the "
"CAPTIONFORGE_PIPELINE_PLAN consumed through the pipeline_plan pin by "
"caption nodes and the JLC CaptionForge capstone node. Exposes Joy, Qwen, "
"caption nodes and the JLC CaptionForge Orchestrator. Exposes Joy, Qwen, "
"and generic Ollama Caption run counts for the current supported Pass A set. "
"Caption-stage Ollama model tags are selected directly on each Ollama Caption "
"node. Loads explicit Ollama Distiller/Validator dropdown tags from "
"node. Loads explicit Ollama Distiller/Validator/Formatter dropdown tags from "
"config/captionforge_ollama_models.json, with no family aliases or shorthand "
"model substitutions. The selected output folder is treated as an output root; "
"the planner derives a run-specific working directory for JSON/JSONL artifacts. "
@@ -117,11 +122,34 @@ import re
from pathlib import Path
from typing import Any
from ..engines.captionforge_dataset_export import (
dataset_export_inputs,
dataset_root,
export_settings_from_widgets,
normalize_export_settings,
prepare_dataset_root,
)
try:
from .captionforge_ollama_model_dropdowns import load_ollama_model_dropdowns
except Exception: # pragma: no cover - useful for direct local smoke tests
from captionforge_ollama_model_dropdowns import load_ollama_model_dropdowns
try:
from ..engines.captionforge_prompt_defaults import (
DEFAULT_FAT_DRAFT_INSTRUCTIONS,
DEFAULT_TAGGY_FORMATTER_INSTRUCTIONS,
DEFAULT_VALIDATOR_INSTRUCTIONS,
DEFAULT_VALIDATOR_SYSTEM_PROMPT,
)
except Exception: # pragma: no cover - useful for direct local smoke tests
from engines.captionforge_prompt_defaults import (
DEFAULT_FAT_DRAFT_INSTRUCTIONS,
DEFAULT_TAGGY_FORMATTER_INSTRUCTIONS,
DEFAULT_VALIDATOR_INSTRUCTIONS,
DEFAULT_VALIDATOR_SYSTEM_PROMPT,
)
try:
import folder_paths
except Exception:
@@ -142,17 +170,16 @@ DEFAULT_CAPTION_TEMPERATURE_SCHEDULE = "0.90"
DEFAULT_CAPTION_TOP_P_SCHEDULE = "0.60"
DEFAULT_CAPTION_TOP_K_SCHEDULE = "80"
DEFAULT_CAPTION_MAX_IMAGE_SIZE = 1024
DEFAULT_CAPTION_MAX_NEW_TOKENS = 6000
DEFAULT_CAPTION_MAX_NEW_TOKENS = 4096
_MODEL_DROPDOWNS = load_ollama_model_dropdowns(__file__)
DISTILLER_MODEL_CHOICES = _MODEL_DROPDOWNS["distiller_models"]
VALIDATOR_MODEL_CHOICES = _MODEL_DROPDOWNS["validator_models"]
FORMAT_MODEL_CHOICES = _MODEL_DROPDOWNS["format_models"]
DEFAULT_DISTILLER_MODEL = _MODEL_DROPDOWNS["distiller_default"]
DEFAULT_VALIDATOR_MODEL = _MODEL_DROPDOWNS["validator_default"]
DISTILLER_STRATEGIES = ["single_pass", "by_model_then_global"]
FINAL_CAPTION_STYLES = ["narrative", "comma", "both"]
DEFAULT_FORMAT_MODEL = _MODEL_DROPDOWNS["format_default"]
DEFAULT_OLLAMA_URL = "http://127.0.0.1:11434"
def _default_output_dir() -> str:
if folder_paths is not None:
try:
@@ -174,6 +201,11 @@ def _as_bool(value: Any) -> bool:
return str(value).strip().lower() in {"1", "true", "yes", "on"}
def _value_or_default(value: Any, default: Any) -> Any:
"""Return a default only for an absent/empty value, preserving numeric zero."""
return default if value is None or value == "" else value
def _runs_per_image(value: Any, default: str = "Disabled") -> int:
"""Normalize caption witness count widgets. Disabled means exactly 0 runs."""
text = str(value if value is not None else default).strip()
@@ -222,10 +254,19 @@ def _call_build_captionforge_pipeline_plan_compat(**kwargs) -> dict[str, Any]:
shared.setdefault("single_image_connected", bool(kwargs.get("single_image_connected", False)))
shared.setdefault("overwrite_outputs", bool(kwargs.get("overwrite_outputs", True)))
plan["ollama"] = {
"url": kwargs.get("ollama_url", DEFAULT_OLLAMA_URL),
"keep_loaded": kwargs.get("ollama_keep_loaded", True),
"request_timeout_seconds": kwargs.get("ollama_request_timeout_seconds", 1800),
}
# Concrete model names are stored in several compatible locations because
# older and newer capstone nodes look in slightly different namespaces.
distiller_model = str(kwargs.get("distiller_model") or DEFAULT_DISTILLER_MODEL).strip() or DEFAULT_DISTILLER_MODEL
validator_model = str(kwargs.get("validator_model") or DEFAULT_VALIDATOR_MODEL).strip() or DEFAULT_VALIDATOR_MODEL
distiller_seed = kwargs.get("distiller_seed", kwargs.get("distiller_base_seed", -1))
validator_seed = kwargs.get("validator_seed", kwargs.get("validator_base_seed", -1))
formatter_seed = kwargs.get("formatter_seed", -1)
caption_common = {
"base_seed": kwargs.get("base_seed", -1),
@@ -249,7 +290,8 @@ def _call_build_captionforge_pipeline_plan_compat(**kwargs) -> dict[str, Any]:
"backend": "ollama",
"model": distiller_model,
"ollama_model": distiller_model,
"seed": kwargs.get("distiller_base_seed", -1),
"prompt": kwargs.get("distiller_prompt", DEFAULT_FAT_DRAFT_INSTRUCTIONS),
"seed": distiller_seed,
})
pass_c_vlm_validator = plan.setdefault("pass_c_vlm_validator", {})
@@ -258,16 +300,17 @@ def _call_build_captionforge_pipeline_plan_compat(**kwargs) -> dict[str, Any]:
"backend": "ollama",
"model": validator_model,
"ollama_model": validator_model,
"seed": kwargs.get("validator_base_seed", -1),
"system_prompt": kwargs.get("validator_system_prompt", DEFAULT_VALIDATOR_SYSTEM_PROMPT),
"prompt": kwargs.get("validator_prompt", DEFAULT_VALIDATOR_INSTRUCTIONS),
"seed": validator_seed,
})
distiller_common = {
"model": distiller_model,
"ollama_model": distiller_model,
"model_family": distiller_model,
"base_seed": kwargs.get("distiller_base_seed", -1),
"seed_mode": kwargs.get("distiller_seed_mode", "fixed"),
"strategy": kwargs.get("distiller_strategy", "single_pass"),
"prompt": kwargs.get("distiller_prompt", DEFAULT_FAT_DRAFT_INSTRUCTIONS),
"seed": distiller_seed,
"max_caption_chars_for_llm": kwargs.get("distiller_max_caption_chars_for_llm", 1536),
"num_predict": kwargs.get("distiller_num_predict", 3096),
"temperature": kwargs.get("distiller_temperature", 0.24),
@@ -283,9 +326,10 @@ def _call_build_captionforge_pipeline_plan_compat(**kwargs) -> dict[str, Any]:
"model": validator_model,
"ollama_model": validator_model,
"model_family": validator_model,
"base_seed": kwargs.get("validator_base_seed", -1),
"seed_mode": kwargs.get("validator_seed_mode", "fixed"),
"num_predict": kwargs.get("validator_num_predict", 2200),
"system_prompt": kwargs.get("validator_system_prompt", DEFAULT_VALIDATOR_SYSTEM_PROMPT),
"prompt": kwargs.get("validator_prompt", DEFAULT_VALIDATOR_INSTRUCTIONS),
"seed": validator_seed,
"num_predict": kwargs.get("validator_num_predict", 2112),
"temperature": kwargs.get("validator_temperature", 0.0),
"top_p": kwargs.get("validator_top_p", 0.92),
"top_k": kwargs.get("validator_top_k", 80),
@@ -295,13 +339,31 @@ def _call_build_captionforge_pipeline_plan_compat(**kwargs) -> dict[str, Any]:
plan["validator"] = dict(validator_common)
plan["pass_c"] = dict(validator_common)
formatter_model = str(kwargs.get("formatter_model") or DEFAULT_FORMAT_MODEL).strip() or DEFAULT_FORMAT_MODEL
formatter_common = {
"model": formatter_model,
"ollama_model": formatter_model,
"model_family": formatter_model,
"prompt": kwargs.get("formatter_prompt", DEFAULT_TAGGY_FORMATTER_INSTRUCTIONS),
"seed": formatter_seed,
"num_predict": kwargs.get("formatter_num_predict", 3200),
"temperature": kwargs.get("formatter_temperature", 0.12),
"top_p": kwargs.get("formatter_top_p", 0.88),
"top_k": kwargs.get("formatter_top_k", 50),
"write_prompt_jsonl": kwargs.get("formatter_write_prompt_jsonl", False),
"preserve_raw_response": kwargs.get("formatter_preserve_raw_response", False),
}
plan["formatter"] = dict(formatter_common)
plan["pass_d"] = dict(formatter_common)
plan["pass_d_formatter"] = dict(formatter_common)
plan["final"] = {
"caption_style": kwargs.get("final_caption_style", "narrative"),
"write_txt_sidecars": kwargs.get("final_write_txt_sidecars", True),
"write_jsonl": kwargs.get("final_write_jsonl", True),
"overwrite_outputs": kwargs.get("overwrite_outputs", True),
}
plan["output"] = {"overwrite_outputs": kwargs.get("overwrite_outputs", True)}
plan["dataset_export"] = normalize_export_settings(kwargs.get("dataset_export"))
return plan
@@ -419,7 +481,7 @@ def _patch_supported_caption_witnesses(
def _patch_v010_working_image_paths(plan: dict[str, Any]) -> dict[str, Any]:
"""Normalize v0.1.x optional-image directories in the emitted plan.
"""Normalize legacy optional-image directories in the emitted plan.
The planner engine owns most path derivation. This wrapper patch keeps
JSON/JSONL/audit artifacts inside the run working directory while promoting
@@ -608,6 +670,35 @@ class JLC_CaptionForge_Pipeline_Planner:
},
),
# -----------------------------------------------------------------
# Shared Ollama runtime controls for Passes B/C/D.
# -----------------------------------------------------------------
"Ollama - URL": (
"STRING",
{
"default": DEFAULT_OLLAMA_URL,
"multiline": False,
"tooltip": "Planner-owned local Ollama server URL for Passes B, C, and D.",
},
),
"Ollama - keep loaded": (
"BOOLEAN",
{
"default": True,
"tooltip": "Planner-owned keep_alive policy for Passes B, C, and D.",
},
),
"Ollama - request timeout seconds": (
"INT",
{
"default": 1800,
"min": 10,
"max": 7200,
"step": 10,
"tooltip": "Planner-owned HTTP patience for Passes B, C, and D; does not affect caption quality.",
},
),
# -----------------------------------------------------------------
# LoRA metadata immediately after outputs.
# -----------------------------------------------------------------
@@ -624,7 +715,27 @@ class JLC_CaptionForge_Pipeline_Planner:
{
"default": "",
"multiline": False,
"tooltip": "Optional user style/identity anchor passed to distiller and validator.",
"tooltip": (
"Optional persistent caption/training anchor. In CaptionForge 1.x, a non-empty "
"anchor is preserved in the final caption variants rather than treated as image "
"evidence that the Validator may remove."
),
},
),
"Cleanup - forbidden phrases": (
"STRING",
{
"default": "",
"multiline": True,
"tooltip": "One forbidden word or phrase per line. Planner values override caption-node and Orchestrator cleanup controls.",
},
),
"Cleanup - replace pairs": (
"STRING",
{
"default": "",
"multiline": True,
"tooltip": "One boundary-safe old=>new replacement per line. Planner values override standalone cleanup controls.",
},
),
@@ -641,14 +752,14 @@ class JLC_CaptionForge_Pipeline_Planner:
"Caption - Qwen runs/image": (
CAPTION_RUNS,
{
"default": "2",
"default": "1",
"tooltip": "Qwen Caption runs per image. Set to Disabled to omit Qwen from this run. Dropdown is capped at 5 to prevent accidental giant runs.",
},
),
"Caption - Ollama runs/image": (
CAPTION_RUNS,
{
"default": "Disabled",
"default": "1",
"tooltip": (
"Ollama Caption runs per image for each connected JLC CaptionForge Ollama Caption node. "
"The actual Ollama model tag is selected in each Ollama Caption node. Connecting multiple "
@@ -663,12 +774,19 @@ class JLC_CaptionForge_Pipeline_Planner:
"min": -1,
"max": MAX_SEED_32,
"step": 1,
"tooltip": "Base seed for caption generation. -1 means unseeded when supported.",
"tooltip": "Base seed for the reusable Pass-A run schedule. -1 means intentionally unseeded.",
},
),
"Caption - seed mode": (
SEED_MODES,
{"default": "fixed"},
{
"default": "fixed",
"tooltip": (
"How the base seed changes across witness runs: fixed reuses it, "
"increment/decrement step by one, and random creates a repeatable "
"hash-derived schedule. A base seed of -1 remains unseeded in every mode."
),
},
),
"Caption - temperature schedule": (
"STRING",
@@ -680,19 +798,48 @@ class JLC_CaptionForge_Pipeline_Planner:
),
"Caption - top p schedule": (
"STRING",
{"default": DEFAULT_CAPTION_TOP_P_SCHEDULE, "multiline": False},
{
"default": DEFAULT_CAPTION_TOP_P_SCHEDULE,
"multiline": False,
"tooltip": (
"Comma-separated nucleus-sampling values for Pass-A runs. Lower values "
"limit choices to more likely tokens; the final value repeats as needed."
),
},
),
"Caption - top k schedule": (
"STRING",
{"default": DEFAULT_CAPTION_TOP_K_SCHEDULE, "multiline": False},
{
"default": DEFAULT_CAPTION_TOP_K_SCHEDULE,
"multiline": False,
"tooltip": (
"Comma-separated token-choice limits for Pass-A runs. Lower values are "
"more restrictive; the final value repeats as needed."
),
},
),
"Caption - max image size": (
"INT",
{"default": DEFAULT_CAPTION_MAX_IMAGE_SIZE, "min": 0, "max": 4096, "step": 64},
{
"default": DEFAULT_CAPTION_MAX_IMAGE_SIZE,
"min": 0,
"max": 4096,
"step": 64,
"tooltip": (
"Maximum longest image side sent to Pass-A captioners. Larger images are "
"resized proportionally; 0 keeps their original size."
),
},
),
"Caption - max new tokens": (
"INT",
{"default": DEFAULT_CAPTION_MAX_NEW_TOKENS, "min": 16, "max": 12000, "step": 64},
{
"default": DEFAULT_CAPTION_MAX_NEW_TOKENS,
"min": 16,
"max": 4096,
"step": 64,
"tooltip": "Maximum Pass-A generation budget. CaptionForge supports up to 4096 tokens.",
},
),
# -----------------------------------------------------------------
@@ -716,51 +863,69 @@ class JLC_CaptionForge_Pipeline_Planner:
"tooltip": "Used only when Distiller - model is custom, e.g. my-model:latest.",
},
),
"Distiller - base seed": (
"Distiller - prompt": (
"STRING",
{
"default": DEFAULT_FAT_DRAFT_INSTRUCTIONS,
"multiline": True,
"tooltip": (
"Planner-owned instructions for the text-only fat draft LLM. "
"Pass A captions are appended automatically."
),
},
),
"Distiller - seed": (
"INT",
{
"default": -1,
"min": -1,
"max": MAX_SEED_32,
"step": 1,
"tooltip": "Base seed for the distiller. -1 means omit seed.",
"tooltip": "Fixed seed used by the distiller for every image. -1 means omit seed.",
},
),
"Distiller - seed mode": (
SEED_MODES,
{"default": "fixed"},
),
"Distiller - strategy": (
DISTILLER_STRATEGIES,
{"default": "single_pass"},
),
"Distiller - max caption chars for LLM": (
"INT",
{"default": 1536, "min": 0, "max": 12000, "step": 64},
{
"default": 1536,
"min": 0,
"max": 12000,
"step": 64,
"tooltip": (
"Maximum characters retained from each Pass-A source caption before "
"building the Pass-B prompt. 0 keeps the complete caption."
),
},
),
"Distiller - num predict": (
"INT",
{"default": 3096, "min": 64, "max": 12000, "step": 64},
{
"default": 3096,
"min": 64,
"max": 12000,
"step": 64,
"tooltip": "Planner token budget for Pass B; maps to Ollama num_predict.",
},
),
"Distiller - temperature": (
"FLOAT",
{"default": 0.24, "min": 0.0, "max": 2.0, "step": 0.01},
{"default": 0.24, "min": 0.0, "max": 2.0, "step": 0.01, "tooltip": "Pass-B variation level. Lower values are steadier; higher values permit more varied wording."},
),
"Distiller - top p": (
"FLOAT",
{"default": 0.90, "min": 0.0, "max": 1.0, "step": 0.01},
{"default": 0.90, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Pass-B nucleus-sampling limit. Lower values restrict the model to more likely tokens."},
),
"Distiller - top k": (
"INT",
{"default": 60, "min": 0, "max": 500, "step": 1},
{"default": 60, "min": 0, "max": 500, "step": 1, "tooltip": "Pass-B token-choice limit. Lower values are more restrictive; 0 lets the backend disable top-k filtering."},
),
"Distiller - write prompt JSONL": (
"BOOLEAN",
{"default": False},
{"default": False, "tooltip": "Write the complete Pass-B prompt to a separate JSONL audit file. This can substantially increase output size."},
),
"Distiller - preserve raw response": (
"BOOLEAN",
{"default": False},
{"default": False, "tooltip": "Keep the unparsed Pass-B model response in audit records for troubleshooting."},
),
# -----------------------------------------------------------------
@@ -784,52 +949,142 @@ class JLC_CaptionForge_Pipeline_Planner:
"tooltip": "Used only when Validator - model is custom, e.g. gemma4:e4b or another installed VLM tag.",
},
),
"Validator - base seed": (
"Validator - system prompt": (
"STRING",
{
"default": DEFAULT_VALIDATOR_SYSTEM_PROMPT,
"multiline": True,
"tooltip": "Planner-owned system prompt for the image-aware VLM validator.",
},
),
"Validator - prompt": (
"STRING",
{
"default": DEFAULT_VALIDATOR_INSTRUCTIONS,
"multiline": True,
"tooltip": (
"Planner-owned instructions for the image-aware VLM validator. "
"The fat draft is appended automatically."
),
},
),
"Validator - seed": (
"INT",
{
"default": -1,
"min": -1,
"max": MAX_SEED_32,
"step": 1,
"tooltip": "Base seed for the VLM validator. -1 means omit seed.",
"tooltip": "Fixed seed used by the VLM validator for every image. -1 means omit seed.",
},
),
"Validator - seed mode": (
SEED_MODES,
{"default": "fixed"},
),
"Validator - num predict": (
"INT",
{"default": 2200, "min": 64, "max": 12000, "step": 64},
{
"default": 2112,
"min": 64,
"max": 12000,
"step": 64,
"tooltip": "Planner token budget for Pass C; maps to Ollama num_predict.",
},
),
"Validator - temperature": (
"FLOAT",
{"default": 0.0, "min": 0.0, "max": 2.0, "step": 0.01},
{"default": 0.0, "min": 0.0, "max": 2.0, "step": 0.01, "tooltip": "Pass-C variation level. Zero requests the most deterministic image-validation result."},
),
"Validator - top p": (
"FLOAT",
{"default": 0.92, "min": 0.0, "max": 1.0, "step": 0.01},
{"default": 0.92, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Pass-C nucleus-sampling limit. Lower values restrict the validator to more likely tokens."},
),
"Validator - top k": (
"INT",
{"default": 80, "min": 0, "max": 500, "step": 1},
{"default": 80, "min": 0, "max": 500, "step": 1, "tooltip": "Pass-C token-choice limit. Lower values are more restrictive; 0 lets the backend disable top-k filtering."},
),
"Validator - write prompt JSONL": (
"BOOLEAN",
{"default": False},
{"default": False, "tooltip": "Write the complete image-validation prompt to a separate JSONL audit file."},
),
"Validator - preserve raw VLM response": (
"BOOLEAN",
{"default": False},
{"default": False, "tooltip": "Keep the unparsed Pass-C VLM response in audit records for troubleshooting."},
),
# -----------------------------------------------------------------
# Formatter controls.
# -----------------------------------------------------------------
"Formatter - model": (
FORMAT_MODEL_CHOICES,
{
"default": DEFAULT_FORMAT_MODEL,
"tooltip": (
"Concrete Ollama text model tag for the SHORT + TAGGY formatter. "
"Use custom to enter any other installed Ollama text model."
),
},
),
"Formatter - custom Ollama model": (
"STRING",
{
"default": "",
"multiline": False,
"tooltip": "Used only when Formatter - model is custom.",
},
),
"Formatter - prompt": (
"STRING",
{
"default": DEFAULT_TAGGY_FORMATTER_INSTRUCTIONS,
"multiline": True,
"tooltip": (
"Planner-owned SHORT + TAGGY instructions for the text-only formatter. "
"The validated LONG caption is appended automatically."
),
},
),
"Formatter - seed": (
"INT",
{
"default": -1,
"min": -1,
"max": MAX_SEED_32,
"step": 1,
"tooltip": "Fixed seed used by the SHORT/TAGGY formatter for every image. -1 means omit seed.",
},
),
"Formatter - num predict": (
"INT",
{
"default": 3200,
"min": 64,
"max": 12000,
"step": 64,
"tooltip": "Planner token budget for Pass D; maps to Ollama num_predict.",
},
),
"Formatter - temperature": (
"FLOAT",
{"default": 0.12, "min": 0.0, "max": 2.0, "step": 0.01, "tooltip": "Pass-D variation level. Lower values make SHORT/TAGGY formatting more consistent."},
),
"Formatter - top p": (
"FLOAT",
{"default": 0.88, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Pass-D nucleus-sampling limit. Lower values restrict the formatter to more likely tokens."},
),
"Formatter - top k": (
"INT",
{"default": 50, "min": 0, "max": 500, "step": 1, "tooltip": "Pass-D token-choice limit. Lower values are more restrictive; 0 lets the backend disable top-k filtering."},
),
"Formatter - write prompt JSONL": (
"BOOLEAN",
{"default": False, "tooltip": "Write the complete Pass-D SHORT/TAGGY prompt to a separate JSONL audit file."},
),
"Formatter - preserve raw response": (
"BOOLEAN",
{"default": False, "tooltip": "Keep the unparsed Pass-D model response in audit records for troubleshooting."},
),
# -----------------------------------------------------------------
# Final export controls.
# -----------------------------------------------------------------
"Final - caption style": (
FINAL_CAPTION_STYLES,
{"default": "narrative"},
),
"Final - write TXT sidecars": (
"BOOLEAN",
{
@@ -841,10 +1096,11 @@ class JLC_CaptionForge_Pipeline_Planner:
),
"Final - write JSONL": (
"BOOLEAN",
{"default": True},
{"default": True, "tooltip": "Write the final run-level JSONL containing LONG, SHORT, and TAGGY captions for every processed image."},
),
},
"optional": {
**dataset_export_inputs(),
"Input - single image": (
"IMAGE",
{
@@ -860,7 +1116,7 @@ class JLC_CaptionForge_Pipeline_Planner:
RETURN_TYPES = ("IMAGE", "CAPTIONFORGE_PIPELINE_PLAN", "STRING")
RETURN_NAMES = ("single_image", "pipeline_plan", "pipeline_plan_json")
FUNCTION = "plan"
CATEGORY = "Captioning/CaptionForge"
CATEGORY = "Caption/CaptionForge"
@classmethod
def IS_CHANGED(cls, **kwargs):
@@ -882,12 +1138,16 @@ class JLC_CaptionForge_Pipeline_Planner:
kwargs.get("Validator - custom Ollama model", ""),
DEFAULT_VALIDATOR_MODEL,
)
formatter_model = _resolve_ollama_model_name(
kwargs.get("Formatter - model", DEFAULT_FORMAT_MODEL),
kwargs.get("Formatter - custom Ollama model", ""),
DEFAULT_FORMAT_MODEL,
)
planner_enabled = _as_bool(kwargs.get("Planner - enabled", True))
joy_runs = _runs_per_image(kwargs.get("Caption - Joy runs/image", "2"), "2")
qwen_runs = _runs_per_image(kwargs.get("Caption - Qwen runs/image", "2"), "2")
ollama_runs = _runs_per_image(kwargs.get("Caption - Ollama runs/image", "Disabled"), "Disabled")
smolvlm_runs = 0
qwen_runs = _runs_per_image(kwargs.get("Caption - Qwen runs/image", "1"), "1")
ollama_runs = _runs_per_image(kwargs.get("Caption - Ollama runs/image", "1"), "1")
if not planner_enabled:
plan: dict[str, Any] = {}
@@ -912,40 +1172,68 @@ class JLC_CaptionForge_Pipeline_Planner:
smol_runs_per_image=0,
florence_runs_per_image=0,
llama_vision_runs_per_image=0,
base_seed=int(kwargs.get("Caption - base seed", -1) or -1),
base_seed=int(_value_or_default(kwargs.get("Caption - base seed", -1), -1)),
seed_mode=str(kwargs.get("Caption - seed mode", "fixed") or "fixed"),
temperature_schedule=str(kwargs.get("Caption - temperature schedule", DEFAULT_CAPTION_TEMPERATURE_SCHEDULE) or DEFAULT_CAPTION_TEMPERATURE_SCHEDULE),
top_p_schedule=str(kwargs.get("Caption - top p schedule", DEFAULT_CAPTION_TOP_P_SCHEDULE) or DEFAULT_CAPTION_TOP_P_SCHEDULE),
top_k_schedule=str(kwargs.get("Caption - top k schedule", DEFAULT_CAPTION_TOP_K_SCHEDULE) or DEFAULT_CAPTION_TOP_K_SCHEDULE),
max_size=int(kwargs.get("Caption - max image size", DEFAULT_CAPTION_MAX_IMAGE_SIZE) or DEFAULT_CAPTION_MAX_IMAGE_SIZE),
max_new_tokens=int(kwargs.get("Caption - max new tokens", DEFAULT_CAPTION_MAX_NEW_TOKENS) or DEFAULT_CAPTION_MAX_NEW_TOKENS),
max_size=int(_value_or_default(kwargs.get("Caption - max image size", DEFAULT_CAPTION_MAX_IMAGE_SIZE), DEFAULT_CAPTION_MAX_IMAGE_SIZE)),
max_new_tokens=int(_value_or_default(kwargs.get("Caption - max new tokens", DEFAULT_CAPTION_MAX_NEW_TOKENS), DEFAULT_CAPTION_MAX_NEW_TOKENS)),
trigger_word=str(kwargs.get("LoRA - trigger word", "") or "").strip(),
user_caption_anchor=str(kwargs.get("LoRA - user caption anchor", "") or "").strip(),
forbidden_phrases=str(kwargs.get("Cleanup - forbidden phrases", "") or ""),
replace_pairs=str(kwargs.get("Cleanup - replace pairs", "") or ""),
ollama_url=str(kwargs.get("Ollama - URL", DEFAULT_OLLAMA_URL) or DEFAULT_OLLAMA_URL),
ollama_keep_loaded=_as_bool(kwargs.get("Ollama - keep loaded", True)),
ollama_request_timeout_seconds=int(
_value_or_default(kwargs.get("Ollama - request timeout seconds", 1800), 1800)
),
distiller_model=distiller_model,
distiller_model_family=distiller_model,
distiller_base_seed=int(kwargs.get("Distiller - base seed", -1) or -1),
distiller_seed_mode=str(kwargs.get("Distiller - seed mode", "fixed") or "fixed"),
distiller_strategy=str(kwargs.get("Distiller - strategy", "single_pass") or "single_pass"),
distiller_max_caption_chars_for_llm=int(kwargs.get("Distiller - max caption chars for LLM", 1536) or 1536),
distiller_prompt=str(
kwargs.get("Distiller - prompt", DEFAULT_FAT_DRAFT_INSTRUCTIONS)
or DEFAULT_FAT_DRAFT_INSTRUCTIONS
),
distiller_seed=int(_value_or_default(kwargs.get("Distiller - seed", -1), -1)),
distiller_max_caption_chars_for_llm=int(_value_or_default(kwargs.get("Distiller - max caption chars for LLM", 1536), 1536)),
distiller_num_predict=int(kwargs.get("Distiller - num predict", 3096) or 3096),
distiller_temperature=float(kwargs.get("Distiller - temperature", 0.24) or 0.0),
distiller_top_p=float(kwargs.get("Distiller - top p", 0.90) or 0.90),
distiller_top_k=int(kwargs.get("Distiller - top k", 60) or 60),
distiller_top_p=float(_value_or_default(kwargs.get("Distiller - top p", 0.90), 0.90)),
distiller_top_k=int(_value_or_default(kwargs.get("Distiller - top k", 60), 60)),
distiller_write_prompt_jsonl=_as_bool(kwargs.get("Distiller - write prompt JSONL", False)),
distiller_preserve_raw_response=_as_bool(kwargs.get("Distiller - preserve raw response", False)),
validator_model=validator_model,
validator_model_family=validator_model,
validator_base_seed=int(kwargs.get("Validator - base seed", -1) or -1),
validator_seed_mode=str(kwargs.get("Validator - seed mode", "fixed") or "fixed"),
validator_num_predict=int(kwargs.get("Validator - num predict", 2200) or 2200),
validator_system_prompt=str(
kwargs.get("Validator - system prompt", DEFAULT_VALIDATOR_SYSTEM_PROMPT)
or DEFAULT_VALIDATOR_SYSTEM_PROMPT
),
validator_prompt=str(
kwargs.get("Validator - prompt", DEFAULT_VALIDATOR_INSTRUCTIONS)
or DEFAULT_VALIDATOR_INSTRUCTIONS
),
validator_seed=int(_value_or_default(kwargs.get("Validator - seed", -1), -1)),
validator_num_predict=int(kwargs.get("Validator - num predict", 2112) or 2112),
validator_temperature=float(kwargs.get("Validator - temperature", 0.0) or 0.0),
validator_top_p=float(kwargs.get("Validator - top p", 0.92) or 0.92),
validator_top_k=int(kwargs.get("Validator - top k", 80) or 80),
validator_top_p=float(_value_or_default(kwargs.get("Validator - top p", 0.92), 0.92)),
validator_top_k=int(_value_or_default(kwargs.get("Validator - top k", 80), 80)),
validator_write_prompt_jsonl=_as_bool(kwargs.get("Validator - write prompt JSONL", False)),
validator_preserve_raw_vlm_response=_as_bool(kwargs.get("Validator - preserve raw VLM response", False)),
final_caption_style=str(kwargs.get("Final - caption style", "narrative") or "narrative"),
formatter_model=formatter_model,
formatter_prompt=str(
kwargs.get("Formatter - prompt", DEFAULT_TAGGY_FORMATTER_INSTRUCTIONS)
or DEFAULT_TAGGY_FORMATTER_INSTRUCTIONS
),
formatter_seed=int(_value_or_default(kwargs.get("Formatter - seed", -1), -1)),
formatter_num_predict=int(kwargs.get("Formatter - num predict", 3200) or 3200),
formatter_temperature=float(kwargs.get("Formatter - temperature", 0.12) or 0.0),
formatter_top_p=float(_value_or_default(kwargs.get("Formatter - top p", 0.88), 0.88)),
formatter_top_k=int(_value_or_default(kwargs.get("Formatter - top k", 50), 50)),
formatter_write_prompt_jsonl=_as_bool(kwargs.get("Formatter - write prompt JSONL", False)),
formatter_preserve_raw_response=_as_bool(kwargs.get("Formatter - preserve raw response", False)),
final_write_txt_sidecars=_as_bool(kwargs.get("Final - write TXT sidecars", True)),
final_write_jsonl=_as_bool(kwargs.get("Final - write JSONL", True)),
dataset_export=export_settings_from_widgets(kwargs),
overwrite_outputs=_as_bool(kwargs.get("Output - overwrite outputs", True)),
)
plan = _patch_supported_caption_witnesses(
@@ -956,6 +1244,12 @@ class JLC_CaptionForge_Pipeline_Planner:
)
plan = _patch_v010_working_image_paths(plan)
export = plan["dataset_export"]
if export["enabled"]:
root = dataset_root(export, output_dir)
prepare_dataset_root(root, str(kwargs.get("Input - image path", "") or ""))
plan["paths"]["training_dataset_dir"] = str(root)
overwrite_outputs = _as_bool(kwargs.get("Output - overwrite outputs", True))
_reset_pass_a_jsonl_for_overwrite(plan, overwrite_outputs=overwrite_outputs)
+8 -12
View File
@@ -8,11 +8,9 @@ JLC CaptionForge Template Options — ComfyUI Node Wrapper
- Repository
https://github.com/Damkohler/CaptionForge
- CaptionForge focuses on practical dataset-captioning infrastructure for
LoRA dataset preparation, using multi-engine caption generation, JSONL
audit trails, claim extraction and refinement, text-LLM distillation,
image-aware VLM validation, and consensus-oriented caption improvement
to produce grounded, auditable training captions.
- CaptionForge 1.0 uses independent Pass-A witnesses, text-LLM synthesis,
image-aware validation, SHORT/TAGGY formatting, and JSONL audit trails to
produce grounded LoRA dataset captions.
- Node Purpose
- The **JLC CaptionForge Template Options** node provides a shared
@@ -85,12 +83,10 @@ JLC CaptionForge Template Options — ComfyUI Node Wrapper
Joy, Qwen, Ollama VLMs, or future caption models without duplicating
every checkbox in every caption node.
- ⚠️ Development Status
- This is active CaptionForge raw-caption infrastructure.
- The UI, option taxonomy, payload schema, and future configuration-file
support may evolve as CaptionForge matures.
- The node is intended for local dataset preparation and controlled caption
audit workflows.
- Production Status
- Active CaptionForge 1.0 Pass-A support node. It supplies shared prompt
modifiers only; each witness backend retains its own prompt dialect and
model behavior.
- Attribution & License
- Concept and implementation by **J. L. Córdova**
@@ -339,7 +335,7 @@ class JLC_CaptionForgeExtraOptions:
RETURN_TYPES = ("CAPTIONFORGE_EXTRA_OPTIONS", "STRING")
RETURN_NAMES = ("template_options", "template_options_json")
FUNCTION = "build"
CATEGORY = "Captioning/CaptionForge"
CATEGORY = "Caption/CaptionForge"
def build(self, **kwargs):
# Accept the old name_input key as a courtesy during hot-swaps, but the
+42 -24
View File
@@ -4,50 +4,55 @@ build-backend = "setuptools.build_meta"
[project]
name = "captionforge"
version = "0.1.0"
description = "Experimental ComfyUI captioning framework for LoRA dataset preparation using multi-engine caption witnesses, JSONL audit trails, text-LLM distillation, image-aware VLM validation, and grounded final-caption export."
version = "1.0.2"
description = "Local, auditable ComfyUI captioning for LoRA datasets: multi-witness evidence, text-LLM synthesis, image-aware VLM validation, and LONG/SHORT/TAGGY exports."
readme = "README.md"
requires-python = ">=3.10,<3.13"
license = { text = "MIT" }
authors = [
{ name = "J. L. Córdova" }
{ name = "J. L. Córdova", email = "jose.l.cordova@gmail.com" }
]
maintainers = [
{ name = "J. L. Córdova" }
{ name = "J. L. Córdova", email = "jose.l.cordova@gmail.com" }
]
keywords = [
"comfyui",
"comfyui-nodes",
"comfyui-custom-nodes",
"captioning",
"lora",
"caption-generation",
"dataset-captioning",
"dataset-preparation",
"image-captioning",
"image-dataset",
"lora",
"lora-training",
"vision-language-models",
"vlm",
"vlm-validation",
"llm",
"multimodal-ai",
"ollama",
"huggingface",
"qwen",
"joycaption",
"jsonl",
"audit-trail",
"local-ai",
]
classifiers = [
"Development Status :: 2 - Pre-Alpha",
"Intended Audience :: Developers",
"Intended Audience :: Science/Research",
"Programming Language :: Python :: 3",
"License :: OSI Approved :: MIT License",
"Topic :: Multimedia :: Graphics",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
"Topic :: Software Development :: Libraries :: Python Modules",
"License :: OSI Approved :: MIT License",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.10",
"Operating System :: OS Independent",
]
dependencies = [
"torch",
"transformers",
"accelerate",
"bitsandbytes>=0.46.1",
"huggingface-hub",
"pillow",
"numpy",
@@ -57,7 +62,7 @@ dependencies = [
[project.optional-dependencies]
quantization = [
"bitsandbytes",
"bitsandbytes>=0.46.1",
]
dev = [
@@ -70,21 +75,34 @@ Homepage = "https://github.com/Damkohler/CaptionForge"
Repository = "https://github.com/Damkohler/CaptionForge"
Issues = "https://github.com/Damkohler/CaptionForge/issues"
[tool.setuptools]
include-package-data = true
[tool.comfy]
supported_comfyui_version = ">=0.2"
PublisherId = "damkohler"
DisplayName = "CaptionForge"
Icon = "https://raw.githubusercontent.com/Damkohler/CaptionForge/main/assets/icons/jlc-comfyui-nodes_Logo-0512.png"
Category = "nodes"
Homepage = "https://github.com/Damkohler/CaptionForge"
[tool.setuptools.packages.find]
where = ["."]
include = [
"engines*",
"nodes*",
[tool.setuptools]
include-package-data = false
packages = [
"CaptionForge",
"CaptionForge.engines",
"CaptionForge.nodes",
"CaptionForge.nodes.caption_nodes",
]
[tool.setuptools.package-dir]
CaptionForge = "."
"CaptionForge.engines" = "engines"
"CaptionForge.nodes" = "nodes"
"CaptionForge.nodes.caption_nodes" = "nodes/caption_nodes"
[tool.setuptools.package-data]
"*" = [
CaptionForge = [
"config/*.json",
"config/**/*.json",
"assets/**/*",
"assets/icons/*",
"assets/workflows/*",
"web/**/*",
]
+11
View File
@@ -0,0 +1,11 @@
# ComfyUI Manager installs this file. Keep aligned with project.dependencies
# in pyproject.toml; tests/test_joy_hardening.py checks both install paths.
torch
transformers
accelerate
bitsandbytes>=0.46.1
huggingface-hub
pillow
numpy
safetensors
qwen-vl-utils
+231
View File
@@ -0,0 +1,231 @@
"""Regression tests for boundary-safe Pass-A forbidden-phrase cleanup."""
from __future__ import annotations
import importlib
import sys
import types
import unittest
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
def _install_namespace(name: str, path: Path) -> None:
if name in sys.modules:
return
module = types.ModuleType(name)
module.__path__ = [str(path)]
sys.modules[name] = module
_install_namespace("CaptionForge", ROOT)
_install_namespace("CaptionForge.engines", ROOT / "engines")
_install_namespace("CaptionForge.nodes", ROOT / "nodes")
_install_namespace("CaptionForge.nodes.caption_nodes", ROOT / "nodes" / "caption_nodes")
if "folder_paths" not in sys.modules:
folder_paths = types.ModuleType("folder_paths")
folder_paths.models_dir = str(ROOT / "models")
folder_paths.get_output_directory = lambda: str(ROOT / "output")
sys.modules["folder_paths"] = folder_paths
cleanup = importlib.import_module("CaptionForge.engines.captionforge_cleanup")
planner_engine = importlib.import_module("CaptionForge.engines.captionforge_pipeline_planner_engine")
joy = importlib.import_module("CaptionForge.engines.jlc_joy_caption_engine")
qwen = importlib.import_module("CaptionForge.engines.jlc_qwen_caption_engine")
ollama = importlib.import_module(
"CaptionForge.nodes.caption_nodes.jlc_captionforge_ollama_caption_node"
)
class SharedForbiddenPhraseContractTests(unittest.TestCase):
def test_pipeline_cleanup_settings_propagate_and_override_standalone_values(self) -> None:
plan = planner_engine.build_captionforge_pipeline_plan(
forbidden_phrases="old\nsafety disclaimer",
replace_pairs="former=>current\nred car=>blue car",
)
self.assertEqual(plan["cleanup"]["forbidden_phrases"], ["old", "safety disclaimer"])
self.assertEqual(
plan["cleanup"]["replace_pairs"],
[{"old": "former", "new": "current"}, {"old": "red car", "new": "blue car"}],
)
forbidden, pairs = cleanup.resolve_cleanup_settings(
plan, "standalone forbidden", "standalone old=>standalone new"
)
self.assertEqual(forbidden, ["old", "safety disclaimer"])
self.assertEqual(pairs, [("former", "current"), ("red car", "blue car")])
def test_standalone_cleanup_settings_survive_without_planner(self) -> None:
forbidden, pairs = cleanup.resolve_cleanup_settings({}, "old", "former=>current")
self.assertEqual(forbidden, ["old"])
self.assertEqual(pairs, [("former", "current")])
def test_end_to_end_cleanup_keeps_boundaries_and_normalizes_spacing(self) -> None:
value = cleanup.apply_cleanup_contract(
"bold holding gold, old, former wording.", ["old"], [("former", "current")]
)
self.assertEqual(value, "bold holding gold, current wording.")
def test_substrings_inside_legitimate_words_do_not_match(self) -> None:
text = "A bold subject is holding a gold accessory."
self.assertFalse(cleanup.contains_forbidden_phrase(text, ["old"]))
self.assertEqual(cleanup.remove_forbidden_phrases(text, ["old"]), text)
def test_boundary_safe_replacements_preserve_containing_words(self) -> None:
source = "bold pose, holding a gold prop, old stone wall"
expected = "bold pose, holding a gold prop, young stone wall"
self.assertEqual(
cleanup.replace_phrases(source, [("old", "young")]),
expected,
)
def test_true_word_and_phrase_matches_are_boundary_aware(self) -> None:
self.assertTrue(cleanup.contains_forbidden_phrase("an old stone wall", ["old"]))
self.assertTrue(
cleanup.contains_forbidden_phrase(
"caption includes safety disclaimer text",
["safety disclaimer"],
)
)
self.assertFalse(
cleanup.contains_forbidden_phrase(
"caption includes safety disclaimers",
["safety disclaimer"],
)
)
def test_joy_and_qwen_replace_pairs_are_boundary_safe_by_default(self) -> None:
source = "bold pose, holding a gold prop, old stone wall"
expected = "bold pose, holding a gold prop, young stone wall"
self.assertTrue(joy.CleanupConfig().replace_whole_words_only)
self.assertTrue(qwen.CleanupConfig().replace_whole_words_only)
self.assertEqual(
joy.apply_replacements(source, [("old", "young")]),
expected,
)
self.assertEqual(
qwen.apply_replacements(source, [("old", "young")]),
expected,
)
def test_joy_and_qwen_preserve_containing_words_but_remove_true_match(self) -> None:
source = "bold pose, holding a gold prop, old stone wall"
expected = "bold pose, holding a gold prop, stone wall"
self.assertEqual(joy.remove_forbidden_phrases(source, ["old"]), expected)
self.assertEqual(qwen.remove_forbidden_phrases(source, ["old"]), expected)
def test_ollama_does_not_drop_paragraph_for_substring_false_positive(self) -> None:
raw = "A bold figure is holding a gold accessory in dramatic light."
cleaned, status = ollama._clean_caption(
raw,
trigger_word="",
forbidden_phrases=["old"],
replacement_rules=[],
)
self.assertEqual(status, "ok")
self.assertEqual(cleaned, raw)
def test_ollama_replace_pairs_are_boundary_safe(self) -> None:
raw = "A bold figure is holding a gold prop near an old wall."
cleaned, status = ollama._clean_caption(
raw,
trigger_word="",
forbidden_phrases=[],
replacement_rules=[("old", "young")],
)
self.assertEqual(status, "ok")
self.assertEqual(
cleaned,
"A bold figure is holding a gold prop near an young wall.",
)
def test_ollama_replacement_case_sensitivity_is_preserved(self) -> None:
raw = "Old wall beside an old wall."
cleaned, status = ollama._clean_caption(
raw,
trigger_word="",
forbidden_phrases=[],
replacement_rules=[("old", "young")],
)
self.assertEqual(status, "ok")
self.assertEqual(cleaned, "Old wall beside an young wall.")
def test_ollama_still_drops_line_for_true_forbidden_match(self) -> None:
raw = "First safe line.\nAn old line.\nFinal safe line."
cleaned, status = ollama._clean_caption(
raw,
trigger_word="",
forbidden_phrases=["old"],
replacement_rules=[],
)
self.assertEqual(status, "ok")
self.assertEqual(cleaned, "First safe line.\nFinal safe line.")
def test_ollama_run_config_audits_cleanup_settings(self) -> None:
config = ollama._build_run_config(
model_tag="example:model",
ollama_url="http://127.0.0.1:11434",
system_prompt="system",
prompt="prompt",
max_new_tokens=100,
temperature=0.2,
top_p=0.9,
top_k=40,
repetition_penalty=1.03,
max_size=1024,
forbidden_phrases=["old", "safety disclaimer"],
replacement_rules=[("foo", "bar")],
)
self.assertEqual(
config["cleanup"]["forbidden_phrases"],
["old", "safety disclaimer"],
)
self.assertEqual(config["cleanup"]["replacement_rules"], [["foo", "bar"]])
self.assertEqual(
config["cleanup"]["replacement_match_mode"],
"whole_word_or_phrase_boundary",
)
self.assertFalse(config["cleanup"]["replacement_case_insensitive"])
self.assertEqual(
config["cleanup"]["forbidden_match_mode"],
"whole_word_or_phrase_boundary",
)
self.assertEqual(config["cleanup"]["forbidden_action"], "drop_matching_line")
def test_test01_style_seven_witnesses_all_remain_eligible(self) -> None:
"""Reproduce the Test_01 false-positive signature across seven Pass-A witnesses."""
witnesses = [
("joy", "A subject is holding a prop with both hands."),
("joy", "A bold composition with detailed clothing."),
("qwen", "The subject is holding an accessory against a plain background."),
("ollama", "A bold figure is holding a gold accessory in dramatic light."),
("ollama", "The subject is holding a pose with bold styling."),
("ollama", "A gold ornament is visible while the subject is holding the garment."),
("ollama", "A detailed portrait with a confident pose and studio lighting."),
]
survivors = []
for family, caption in witnesses:
if family == "joy":
cleaned = joy.remove_forbidden_phrases(caption, ["old"])
status = "ok" if cleaned else "filtered"
elif family == "qwen":
cleaned = qwen.remove_forbidden_phrases(caption, ["old"])
status = "ok" if cleaned else "filtered"
else:
cleaned, status = ollama._clean_caption(
caption,
trigger_word="",
forbidden_phrases=["old"],
replacement_rules=[],
)
if status == "ok" and cleaned:
survivors.append((family, cleaned))
self.assertEqual(len(survivors), 7)
if __name__ == "__main__":
unittest.main()
File diff suppressed because it is too large Load Diff
+305
View File
@@ -0,0 +1,305 @@
"""CPU checks for paired exports, source protection, and Planner ownership."""
from __future__ import annotations
import importlib
import io
import json
import sys
import tempfile
import types
import unittest
import urllib.error
from pathlib import Path
from unittest import mock
from PIL import Image
ROOT = Path(__file__).resolve().parents[1]
for name, path in (("CaptionForge", ROOT), ("CaptionForge.engines", ROOT / "engines"),
("CaptionForge.nodes", ROOT / "nodes"),
("CaptionForge.nodes.caption_nodes", ROOT / "nodes" / "caption_nodes")):
if name not in sys.modules:
module = types.ModuleType(name)
module.__path__ = [str(path)]
sys.modules[name] = module
if "folder_paths" not in sys.modules:
sys.modules["folder_paths"] = types.ModuleType("folder_paths")
sys.modules["folder_paths"].models_dir = str(ROOT / "models")
sys.modules["folder_paths"].get_output_directory = lambda: str(ROOT / "output")
exporter = importlib.import_module("CaptionForge.engines.captionforge_dataset_export")
planner = importlib.import_module("CaptionForge.nodes.jlc_captionforge_pipeline_planner_node")
orchestrator = importlib.import_module("CaptionForge.nodes.jlc_captionforge_node")
class DatasetExportTests(unittest.TestCase):
def setUp(self):
temporary = tempfile.TemporaryDirectory()
self.addCleanup(temporary.cleanup)
self.root = Path(temporary.name)
self.source = self.root / "portraits" / "photo.jpg"
self.source.parent.mkdir()
Image.new("RGB", (900, 600), "orange").save(self.source)
self.original_bytes = self.source.read_bytes()
self.settings = exporter.normalize_export_settings({"enabled": True, "max_size": 256})
def pair(self, **overrides):
arguments = dict(source=self.source, image_key="portraits/photo.jpg", input_root=self.root,
root=self.root / "training_dataset", settings=self.settings,
validator_max_size=1536, captions={"short": "A portrait.", "long": "Long portrait.", "taggy": "portrait"},
overwrite=False)
arguments.update(overrides)
return exporter.export_pair(**arguments)
def test_downscale_and_divisibility_never_enlarge(self):
result = self.pair()
with Image.open(result["image"]) as image:
self.assertEqual(image.size, (256, 160))
self.assertEqual(Path(result["caption"]).read_text().strip(), "A portrait.")
self.assertEqual(self.source.read_bytes(), self.original_bytes)
self.assertFalse(self.source.with_suffix(".txt").exists())
def test_small_image_and_divisor_one(self):
image = Image.new("RGB", (99, 65))
self.assertEqual(exporter.resize_for_export(image, 1536, 1).size, (99, 65))
self.assertEqual(exporter.resize_for_export(image, 1536, 16).size, (96, 64))
with self.assertRaises(ValueError):
exporter.resize_for_export(Image.new("RGB", (10, 9)), 1536, 16)
def test_fractional_divisor_rejected(self):
with self.assertRaises(ValueError):
exporter.normalize_export_settings({"divisor": 16.5})
def test_original_extensions_and_folders_cannot_collide(self):
first = self.pair()
second_source = self.source.with_suffix(".png")
Image.new("RGB", (900, 600)).save(second_source)
second = self.pair(source=second_source, image_key="portraits/photo.png")
self.assertNotEqual(first["image"], second["image"])
self.assertNotEqual(first["caption"], second["caption"])
self.assertEqual(Path(first["image"]).relative_to(self.root).as_posix(), "training_dataset/files/portraits/photo.jpg.png")
def test_resume_and_changed_settings(self):
first = self.pair()
original_mtime = Path(first["image"]).stat().st_mtime_ns
self.assertTrue(self.pair()["resumed"])
self.assertEqual(Path(first["image"]).stat().st_mtime_ns, original_mtime)
with self.assertRaises(FileExistsError):
self.pair(settings={**self.settings, "divisor": 32})
regenerated = self.pair(settings={**self.settings, "divisor": 32}, overwrite=True)
self.assertFalse(regenerated["resumed"])
def test_foreign_folder_and_foreign_file_are_protected(self):
destination = self.root / "training_dataset"
destination.mkdir()
original = destination / "archive.png"
original.write_bytes(b"archive")
with self.assertRaises(ValueError):
self.pair(overwrite=True)
self.assertEqual(original.read_bytes(), b"archive")
original.unlink()
exporter.prepare_dataset_root(destination)
target, _, _ = exporter.export_paths(destination, self.source, self.root, "portraits/photo.jpg", "PNG")
target.parent.mkdir(parents=True)
target.write_bytes(b"untracked")
with self.assertRaises(FileExistsError):
self.pair(overwrite=True)
self.assertEqual(target.read_bytes(), b"untracked")
def test_destination_containing_input_is_rejected(self):
with self.assertRaises(ValueError):
exporter.prepare_dataset_root(self.root, self.source)
def test_generated_images_excluded_from_every_witness_scan(self):
self.pair()
for family in ("joy", "qwen", "ollama"):
module = importlib.import_module(f"CaptionForge.nodes.caption_nodes.jlc_captionforge_{family}_caption_node")
paths = module._iter_input_path_images(str(self.root), True, "*")
self.assertEqual([entry[2] for entry in paths], [self.source])
def test_explicit_generated_dataset_root_is_valid_witness_input(self):
exported = self.pair()
dataset_root = self.root / "training_dataset"
exported_image = Path(exported["image"])
for family in ("joy", "qwen", "ollama"):
module = importlib.import_module(f"CaptionForge.nodes.caption_nodes.jlc_captionforge_{family}_caption_node")
# A generated dataset discovered under a broader source root remains excluded.
parent_scan = module._iter_input_path_images(str(self.root), True, "*")
self.assertEqual([entry[2] for entry in parent_scan], [self.source])
# But explicitly selecting that generated dataset is an intentional user action
# and must make its images available for recaptioning/reprocessing.
explicit_scan = module._iter_input_path_images(str(dataset_root), True, "*")
self.assertEqual([entry[2] for entry in explicit_scan], [exported_image])
def test_jpeg_and_optional_image_namespace(self):
result = self.pair(image_key="captionforge-optional-image://comfy_image_0000.png",
settings={**self.settings, "image_format": "JPEG", "caption": "taggy"})
self.assertIn("optional", Path(result["image"]).parts)
with Image.open(result["image"]) as image:
self.assertEqual(image.format, "JPEG")
self.assertEqual(Path(result["caption"]).read_text().strip(), "portrait")
def run_pipeline(self, *, plan=None, overwrite=True, validator_side_effect=None, **extra):
caption_path = self.root / "witness.jsonl"
caption_path.write_text(json.dumps({"image": str(self.source), "image_key": "portraits/photo.jpg",
"caption": "A portrait.", "model_family": "joy", "status": "ok"}) + "\n")
kwargs = {"Input - captions JSONL": str(caption_path), "Input - image path": str(self.root),
"Output - folder": str(self.root), "Output - run name": "prototype",
"Output - overwrite outputs": overwrite, "Validator - max image size": 256,
"Dataset - export image and caption": True, "pipeline_plan": plan}
kwargs.update(extra)
with mock.patch.object(orchestrator, "_evict_python_models_before_ollama_if_needed"), \
mock.patch.object(orchestrator, "_ollama_generate_text", return_value=("SHORT: A portrait.\nTAGGY: portrait", {})) as text, \
mock.patch.object(orchestrator, "_ollama_chat_image", return_value=("A long portrait.", {}), side_effect=validator_side_effect) as validator:
result = orchestrator.JLC_CaptionForge().forge(**kwargs)
return json.loads(result[3])["records"][0], text.call_count, validator.call_count
def test_standalone_export_and_resume_without_model_calls(self):
record, _, _ = self.run_pipeline()
self.assertEqual(record["status"], "ok", record)
self.assertEqual(record["dataset_export"]["width"], 256)
self.assertTrue(Path(record["outputs"]["short"]).is_relative_to(self.root / "training_dataset"))
self.assertFalse(self.source.with_name("photo_short.txt").exists())
resumed, text_calls, validator_calls = self.run_pipeline(overwrite=False)
self.assertEqual((text_calls, validator_calls), (0, 0))
self.assertTrue(resumed["dataset_export"]["resumed"])
def test_planner_controls_override_every_local_setting_including_blank_folder(self):
_, plan, _ = planner.JLC_CaptionForge_Pipeline_Planner().plan(**{
"Input - image path": str(self.root), "Output - folder": str(self.root),
"Output - run name": "prototype", "Caption - max image size": 384,
"Dataset - export image and caption": True, "Dataset - dimension divisor": 32,
"Dataset - caption": "long", "Dataset - image format": "JPEG", "Dataset - JPEG quality": 88,
})
plan["paths"]["a_raw_captions_jsonl"] = str(self.root / "witness.jsonl")
# The canonical plan may expose more than one compatible ledger key.
for key in plan["paths"]:
if "caption" in key.lower() and "jsonl" in key.lower():
plan["paths"][key] = str(self.root / "witness.jsonl")
record, _, _ = self.run_pipeline(plan=plan, **{
"Dataset - export image and caption": False,
"Dataset - output folder": str(self.root / "wrong"), "Dataset - max image size": 128,
"Dataset - dimension divisor": 1, "Dataset - caption": "taggy",
"Dataset - image format": "PNG", "Dataset - JPEG quality": 1,
})
self.assertEqual(record["status"], "ok", record)
exported = record["dataset_export"]
self.assertTrue(Path(exported["image"]).is_relative_to(self.root / "training_dataset"))
self.assertEqual((exported["width"], exported["height"]), (384, 256))
self.assertEqual(exported["caption_style"], "long")
self.assertTrue(exported["image"].endswith(".jpg"))
self.assertFalse((self.root / "wrong").exists())
receipt = json.loads(Path(exported["image"]).with_suffix(".captionforge.json").read_text())
self.assertEqual(receipt["signature"]["jpeg_quality"], 88)
def test_old_or_disabled_planner_cannot_enable_local_export(self):
record, _, _ = self.run_pipeline(plan={"captionforge_config_type": "captionforge_pipeline_plan"})
self.assertEqual(record["status"], "ok", record)
self.assertNotIn("dataset_export", record)
self.assertFalse((self.root / "training_dataset").exists())
def test_export_backfills_completed_captions_without_model_calls(self):
first, _, _ = self.run_pipeline(**{"Dataset - export image and caption": False})
self.assertNotIn("dataset_export", first)
record, text_calls, validator_calls = self.run_pipeline(overwrite=False)
self.assertEqual(record["dataset_export"]["status"], "ok")
self.assertEqual((text_calls, validator_calls), (0, 0))
def test_failed_export_preserves_captions_and_can_retry_without_models(self):
record, _, _ = self.run_pipeline(**{"Dataset - dimension divisor": 512})
self.assertEqual(record["status"], "error")
self.assertEqual(record["error_stage"], "dataset_export")
self.assertEqual(record["long"], "A long portrait.")
recovered, text_calls, validator_calls = self.run_pipeline(overwrite=False)
self.assertEqual(recovered["status"], "ok", recovered)
self.assertEqual((text_calls, validator_calls), (0, 0))
def test_validator_retry_does_not_reduce_training_export(self):
error = urllib.error.HTTPError("http://localhost/api/chat", 413, "Too large", {}, io.BytesIO())
record, _, calls = self.run_pipeline(
validator_side_effect=[error, ("A long portrait.", {})],
**{"Validator - max image size": 768},
)
self.assertEqual(calls, 2)
self.assertEqual((record["dataset_export"]["width"], record["dataset_export"]["height"]), (768, 512))
def test_interrupted_publication_requires_overwrite_and_recovers(self):
replace = exporter.os.replace
counter = 0
def interrupted(source, target):
nonlocal counter
counter += 1
if counter == 2:
raise OSError("Interrupted caption publication")
return replace(source, target)
with mock.patch.object(exporter.os, "replace", side_effect=interrupted):
with self.assertRaises(OSError):
self.pair()
with self.assertRaises(FileExistsError):
self.pair()
self.assertEqual(self.pair(overwrite=True)["status"], "ok")
self.assertEqual(self.source.read_bytes(), self.original_bytes)
def test_hardlink_to_source_cannot_be_overwritten(self):
destination = self.root / "training_dataset"
exporter.prepare_dataset_root(destination)
target, _, _ = exporter.export_paths(destination, self.source, self.root, "portraits/photo.jpg", "PNG")
target.parent.mkdir(parents=True)
target.hardlink_to(self.source)
with self.assertRaises(ValueError):
self.pair(overwrite=True)
self.assertEqual(self.source.read_bytes(), self.original_bytes)
def test_caption_variant_hardlinks_and_untracked_files_are_protected(self):
destination = self.root / "training_dataset"
exporter.prepare_dataset_root(destination)
target, _, _ = exporter.export_paths(destination, self.source, self.root, "portraits/photo.jpg", "PNG")
target.parent.mkdir(parents=True)
variant = target.with_name(target.stem + "_long.txt")
variant.hardlink_to(self.source)
with self.assertRaises(ValueError):
self.pair(overwrite=True, write_variants=True)
self.assertEqual(self.source.read_bytes(), self.original_bytes)
variant.unlink()
variant.write_text("My existing caption")
with self.assertRaises(FileExistsError):
self.pair(overwrite=True, write_variants=True)
self.assertEqual(variant.read_text(), "My existing caption")
def test_format_change_cannot_leave_duplicate_training_images(self):
original = self.pair()
with self.assertRaises(FileExistsError):
self.pair(settings={**self.settings, "image_format": "JPEG"}, overwrite=True)
self.assertTrue(Path(original["image"]).exists())
self.assertFalse(Path(original["image"]).with_suffix(".jpg").exists())
def test_changed_variant_cannot_silently_resume(self):
first = self.pair(write_variants=True)
Path(first["variants"]["long"]).write_text("Changed caption")
with self.assertRaises(FileExistsError):
self.pair(write_variants=True)
self.assertEqual(self.pair(write_variants=True, overwrite=True)["status"], "ok")
def test_canonical_workflows_serialize_dataset_widgets(self):
directory = ROOT / "assets" / "workflows"
ui = json.loads((directory / "CaptionForge_FullWorkflow_Rel_v1.0.2.json").read_text(encoding="utf-8"))
api = json.loads((directory / "CaptionForge_FullWorkflow_API_Rel_v1.0.2.json").read_text(encoding="utf-8"))
for name, cls in (("JLC_CaptionForge", orchestrator.JLC_CaptionForge),
("JLC_CaptionForge_Pipeline_Planner", planner.JLC_CaptionForge_Pipeline_Planner)):
node = next(item for item in ui["nodes"] if item["type"] == name)
api_node = next(item for item in api.values() if item["class_type"] == name)
names = list(cls.INPUT_TYPES()["required"]) + list(exporter.dataset_export_inputs())
self.assertEqual(node["widgets_values"], [node["widgets_values_named"][key] for key in names])
for key in names:
self.assertEqual(api_node["inputs"][key], node["widgets_values_named"][key])
if __name__ == "__main__":
unittest.main()
+258
View File
@@ -0,0 +1,258 @@
"""CPU-only warning and installation-contract tests; no model downloads required."""
import ast
from contextlib import nullcontext
import importlib.util
import logging
import threading
from pathlib import Path
import types
import unittest
from unittest import mock
from unittest.mock import Mock
import warnings
try:
import tomllib
except ModuleNotFoundError:
import tomli as tomllib
ROOT = Path(__file__).resolve().parents[1]
SPEC = importlib.util.spec_from_file_location(
'joy_warnings', ROOT / 'engines/captionforge_joy_warnings.py'
)
joy_warnings = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(joy_warnings)
scope = joy_warnings.quantized_inference_warnings
MESSAGE = 'MatMul8bitLt: inputs will be cast from torch.bfloat16 to float16 during quantization'
FP32_MESSAGE = MESSAGE.replace('bfloat16', 'float32')
MODULE = 'bitsandbytes.autograd._functions'
def emit(message=MESSAGE, category=UserWarning, module=MODULE):
warnings.warn_explicit(message, category, filename='mock_bnb.py', lineno=1, module=module)
class WarningScopeTests(unittest.TestCase):
def test_old_8bit_stack_warns_for_joy_and_qwen_without_blocking(self):
with mock.patch.object(joy_warnings.metadata, "version", return_value="0.45.5"):
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always")
joy_warnings.warn_if_suspicious_8bit_stack("Joy Caption")
joy_warnings.warn_if_suspicious_8bit_stack("Qwen Caption")
messages = [str(item.message) for item in caught]
self.assertTrue(any("CaptionForge's Joy Caption node detected" in item for item in messages))
self.assertTrue(any("CaptionForge's Qwen Caption node detected" in item for item in messages))
self.assertTrue(all("continue without modifying packages" in item for item in messages))
def test_current_8bit_stack_does_not_warn(self):
with mock.patch.object(joy_warnings.metadata, "version", return_value="0.46.1"):
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always")
joy_warnings.warn_if_suspicious_8bit_stack("Joy Caption")
joy_warnings.warn_if_suspicious_8bit_stack("Qwen Caption")
self.assertEqual(caught, [])
def test_import_does_not_change_filters(self):
before = list(warnings.filters)
SPEC.loader.exec_module(joy_warnings)
self.assertEqual(warnings.filters, before)
def test_six_caption_bursts_and_useful_diagnostics(self):
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter('always')
before = list(warnings.filters)
for _ in range(6):
with scope(True):
for _ in range(100):
emit()
emit(FP32_MESSAGE)
emit('Unrelated bitsandbytes warning')
self.assertEqual(warnings.filters, before)
self.assertEqual(len(caught), 6)
self.assertTrue(all(str(w.message) == 'Unrelated bitsandbytes warning' for w in caught))
def test_near_matches_remain_visible(self):
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter('always')
with scope(True):
emit(MESSAGE.replace('bfloat16', 'float64'))
emit(MESSAGE + ' extra context')
emit('prefix ' + MESSAGE)
emit(module='another_engine')
emit(FP32_MESSAGE, module='another_engine')
emit(category=RuntimeWarning)
emit(module=MODULE + '.other')
self.assertEqual(len(caught), 7)
def test_default_mode_and_after_scope_remain_visible(self):
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter('always')
with scope(False):
emit()
with scope(True):
emit()
emit()
self.assertEqual(len(caught), 2)
def test_exception_restores_callers_error_policy(self):
with warnings.catch_warnings():
warnings.simplefilter('error')
before = list(warnings.filters)
with self.assertRaisesRegex(ValueError, 'generation failed'):
with scope(True):
emit()
raise ValueError('generation failed')
self.assertEqual(warnings.filters, before)
with self.assertRaises(UserWarning):
emit()
with scope(True):
with self.assertRaisesRegex(UserWarning, 'unrelated'):
emit('unrelated')
class LoggingScopeTests(unittest.TestCase):
def setUp(self):
self.logger = logging.getLogger(MODULE)
self.before = (list(self.logger.filters), self.logger.level, list(self.logger.handlers), self.logger.propagate)
def tearDown(self):
self.assertEqual((list(self.logger.filters), self.logger.level, list(self.logger.handlers), self.logger.propagate), self.before)
def test_six_logging_bursts_keep_other_diagnostics(self):
with self.assertLogs(MODULE, level='WARNING') as caught:
for _ in range(6):
with scope(True):
for _ in range(100):
self.logger.warning('MatMul8bitLt: inputs will be cast from %s to float16 during quantization', 'torch.bfloat16')
self.logger.warning('MatMul8bitLt: inputs will be cast from %s to float16 during quantization', 'torch.float32')
self.logger.warning('Useful diagnostic')
self.assertEqual([r.getMessage() for r in caught.records], ['Useful diagnostic'] * 6)
def test_near_matches_levels_child_logger_and_other_thread(self):
with self.assertLogs(MODULE, level='INFO') as caught:
with scope(True):
self.logger.warning(MESSAGE.replace('bfloat16', 'float64'))
self.logger.warning(MESSAGE + ' extra')
self.logger.error(MESSAGE)
self.logger.info(MESSAGE)
logging.getLogger(MODULE + '.other').warning(MESSAGE)
thread = threading.Thread(target=lambda: self.logger.warning(MESSAGE))
thread.start()
thread.join()
self.assertEqual(len(caught.records), 6)
def test_default_after_scope_and_failure(self):
with self.assertLogs(MODULE, level='WARNING') as caught:
with scope(False):
self.logger.warning(MESSAGE)
with self.assertRaises(ValueError):
with scope(True):
self.logger.warning(MESSAGE)
raise ValueError('generation failed')
self.logger.warning(MESSAGE)
self.assertEqual(len(caught.records), 2)
def test_existing_filters_preserved_in_nested_scope(self):
existing = logging.Filter()
self.logger.addFilter(existing)
try:
with self.assertLogs(MODULE, level='WARNING') as caught:
with scope(True):
with scope(True):
self.logger.warning(MESSAGE)
self.logger.warning(MESSAGE)
self.assertIn(existing, self.logger.filters)
self.logger.warning(MESSAGE)
self.assertEqual(len(caught.records), 1)
self.assertEqual(self.logger.filters, self.before[0] + [existing])
finally:
self.logger.removeFilter(existing)
class GenerationScopeTests(unittest.TestCase):
"""Execute the real caption_pil method with processor/torch/model test doubles."""
def run_generation(self, mode, fail=False):
tree = ast.parse((ROOT / 'engines/jlc_joy_caption_engine.py').read_text(encoding='utf-8'))
cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == 'JoyCaptionEngine')
method = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == 'caption_pil')
method.decorator_list = []
namespace = {
'quantized_inference_warnings': scope,
'torch': types.SimpleNamespace(
bfloat16='bf16', dtype=type(None), autocast=Mock(return_value=nullcontext())
),
'resize_for_model': lambda image, size: image,
'cleanup_caption': lambda text, config: text,
}
# Postponed annotations avoid importing PIL/torch just to test control flow.
module = ast.Module(body=[ast.ImportFrom(module='__future__', names=[ast.alias(name='annotations')], level=0), method], type_ignores=[])
exec(compile(ast.fix_missing_locations(module), '<caption_pil>', 'exec'), namespace)
class Inputs(dict):
def to(self, device):
return self
processor = Mock(return_value=Inputs(input_ids=types.SimpleNamespace(shape=(1, 2))))
processor.tokenizer.eos_token_id = 2
processor.tokenizer.pad_token_id = 2
processor.tokenizer.decode.return_value = 'caption'
def generate(**kwargs):
emit()
emit(FP32_MESSAGE)
logging.getLogger(MODULE).warning(
'MatMul8bitLt: inputs will be cast from %s to float16 during quantization',
'torch.bfloat16',
)
logging.getLogger(MODULE).warning(FP32_MESSAGE)
emit('Useful generation diagnostic')
if fail:
raise ValueError('generation failed')
return [[1, 2, 3]]
engine = types.SimpleNamespace(
generation=types.SimpleNamespace(seed=None, max_new_tokens=10, temperature=0, repetition_penalty=1),
config=types.SimpleNamespace(system_prompt='system', prompt='prompt', max_size=1024, memory_mode=mode),
cleanup=None, inference_device='cpu', processor=processor,
model=types.SimpleNamespace(dtype=None, generate=generate),
prepare_for_inference=Mock(), cleanup_after_inference=Mock(side_effect=lambda: emit('Cleanup diagnostic')),
_autocast_device_type=lambda: 'cpu', _autocast_enabled=lambda device: False,
)
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter('always')
before = list(warnings.filters)
if fail:
with self.assertRaisesRegex(ValueError, 'generation failed'):
namespace['caption_pil'](engine, Mock())
else:
self.assertEqual(namespace['caption_pil'](engine, Mock()), ('caption', 'caption'))
self.assertEqual(warnings.filters, before)
engine.cleanup_after_inference.assert_called_once()
self.assertEqual(namespace['torch'].autocast.call_args.kwargs['dtype'], 'bf16')
return [str(w.message) for w in caught]
def test_balanced_generation(self):
self.assertEqual(self.run_generation('Balanced (8-bit)'), ['Useful generation diagnostic', 'Cleanup diagnostic'])
def test_default_generation(self):
self.assertEqual(self.run_generation('Default'), [MESSAGE, FP32_MESSAGE, 'Useful generation diagnostic', 'Cleanup diagnostic'])
def test_failed_generation(self):
self.assertEqual(self.run_generation('Balanced (8-bit)', fail=True), ['Useful generation diagnostic', 'Cleanup diagnostic'])
class DependencyContractTests(unittest.TestCase):
def test_manager_and_package_dependencies_match(self):
project = tomllib.loads((ROOT / 'pyproject.toml').read_text(encoding='utf-8'))['project']
requirements = {
line.strip() for line in (ROOT / 'requirements.txt').read_text().splitlines()
if line.strip() and not line.lstrip().startswith('#')
}
self.assertEqual(requirements, set(project['dependencies']))
self.assertIn('accelerate', requirements)
self.assertIn('bitsandbytes>=0.46.1', requirements)
self.assertEqual(project['optional-dependencies']['quantization'], ['bitsandbytes>=0.46.1'])
if __name__ == '__main__':
unittest.main()
+406
View File
@@ -0,0 +1,406 @@
"""Focused tests for the CaptionForge Orchestrator output contract."""
from __future__ import annotations
import base64
import importlib
import io
import json
import sys
import tempfile
import types
import unittest
from contextlib import redirect_stdout
from pathlib import Path
from unittest import mock
from PIL import Image
ROOT = Path(__file__).resolve().parents[1]
def _install_namespace(name: str, path: Path) -> None:
if name in sys.modules:
return
module = types.ModuleType(name)
module.__path__ = [str(path)]
sys.modules[name] = module
_install_namespace("CaptionForge", ROOT)
_install_namespace("CaptionForge.engines", ROOT / "engines")
_install_namespace("CaptionForge.nodes", ROOT / "nodes")
if "folder_paths" not in sys.modules:
folder_paths = types.ModuleType("folder_paths")
folder_paths.get_output_directory = lambda: str(ROOT / "output")
sys.modules["folder_paths"] = folder_paths
planner_engine = importlib.import_module("CaptionForge.engines.captionforge_pipeline_planner_engine")
orchestrator = importlib.import_module("CaptionForge.nodes.jlc_captionforge_node")
class ValidatorImagePreparationTests(unittest.TestCase):
def _prepare(self, size: tuple[int, int], max_size: int) -> tuple[tuple[int, int], str, Path]:
temp_dir = tempfile.TemporaryDirectory()
self.addCleanup(temp_dir.cleanup)
image_path = Path(temp_dir.name) / "source.jpg"
Image.new("RGB", size, "white").save(image_path, format="JPEG", quality=95)
output = io.StringIO()
with redirect_stdout(output):
encoded = orchestrator._pil_to_base64_png(image_path, max_size=max_size)
with Image.open(io.BytesIO(base64.b64decode(encoded))) as transmitted:
transmitted_size = transmitted.size
transmitted_format = transmitted.format
self.assertEqual(transmitted_format, "PNG")
return transmitted_size, output.getvalue(), image_path
def test_large_image_is_reduced_to_longest_side_maximum(self) -> None:
transmitted_size, log, image_path = self._prepare((400, 200), max_size=100)
self.assertEqual(transmitted_size, (100, 50))
with Image.open(image_path) as source:
self.assertEqual(source.size, (400, 200))
self.assertIn("source=400x200", log)
self.assertIn("validator=100x50", log)
self.assertIn("encoded_payload=", log)
self.assertIn("MiB", log)
def test_resize_preserves_aspect_ratio(self) -> None:
transmitted_size, _, _ = self._prepare((300, 500), max_size=100)
self.assertEqual(transmitted_size, (60, 100))
self.assertEqual(transmitted_size[0] / transmitted_size[1], 300 / 500)
def test_image_below_maximum_is_not_enlarged(self) -> None:
transmitted_size, _, _ = self._prepare((80, 40), max_size=100)
self.assertEqual(transmitted_size, (80, 40))
def test_zero_maximum_disables_resizing(self) -> None:
transmitted_size, _, _ = self._prepare((400, 200), max_size=0)
self.assertEqual(transmitted_size, (400, 200))
class OrchestratorOutputContractTests(unittest.TestCase):
def test_public_output_surface_is_exactly_five_named_strings(self) -> None:
self.assertEqual(
orchestrator.JLC_CaptionForge.RETURN_NAMES,
("long_captions", "short_captions", "taggy_captions", "final_records", "status"),
)
self.assertEqual(orchestrator.JLC_CaptionForge.RETURN_TYPES, ("STRING",) * 5)
def _run(self, *, planner_connected: bool) -> tuple[tuple[str, ...], Path]:
temp_dir = tempfile.TemporaryDirectory()
self.addCleanup(temp_dir.cleanup)
root = Path(temp_dir.name)
image_names = ("one.png", "two.png")
for name in image_names:
Image.new("RGB", (2, 2), "white").save(root / name)
caption_path = root / "captions.jsonl"
caption_path.write_text(
"\n".join(
json.dumps(
{
"image": name,
"image_key": name,
"caption": f"Pass-A caption for {name}",
"model_family": "joy",
"status": "ok",
}
)
for name in image_names
)
+ "\n",
encoding="utf-8",
)
kwargs = {
"Input - captions JSONL": str(caption_path),
"Input - image path": str(root),
"Output - folder": str(root / "standalone-output"),
"Output - run name": "output-contract",
"Output - overwrite outputs": True,
"Final - write TXT sidecars": True,
"Final - write JSONL": True,
}
if planner_connected:
plan = planner_engine.build_captionforge_pipeline_plan(
output_dir=str(root / "planned-output"),
input_path=str(root),
run_name="output-contract-planned",
final_write_txt_sidecars=True,
final_write_jsonl=True,
)
plan["paths"]["caption_jsonl"] = str(caption_path)
plan["paths"]["pass_a_jsonl"] = str(caption_path)
kwargs["pipeline_plan"] = plan
text_results = iter(
(
"fat draft one",
"SHORT: Short one.\nTAGGY: tag one, detail one",
"fat draft two",
"SHORT: Short two.\nTAGGY: tag two, detail two",
)
)
validator_results = iter(("Long one.", "Long two."))
def fake_text(**_kwargs):
return next(text_results), {"ok": True}
def fake_image(**_kwargs):
return next(validator_results), {"ok": True}
with mock.patch.object(
orchestrator, "_evict_python_models_before_ollama_if_needed"
), mock.patch.object(
orchestrator, "_ollama_generate_text", side_effect=fake_text
), mock.patch.object(
orchestrator, "_ollama_chat_image", side_effect=fake_image
):
result = orchestrator.JLC_CaptionForge().forge(**kwargs)
return result, root
def _assert_semantic_outputs(self, *, planner_connected: bool) -> None:
result, root = self._run(planner_connected=planner_connected)
self.assertEqual(len(result), 5)
long_captions, short_captions, taggy_captions, final_records_json, status = result
self.assertEqual(long_captions.split("\n\n"), ["Long one.", "Long two."])
self.assertEqual(short_captions.split("\n\n"), ["Short one.", "Short two."])
self.assertEqual(
taggy_captions.split("\n\n"),
["tag one, detail one", "tag two, detail two"],
)
payload = json.loads(final_records_json)
self.assertEqual(set(payload), {"records", "run_outputs"})
self.assertEqual(len(payload["records"]), 2)
self.assertEqual(
[(record["long"], record["short"], record["taggy"]) for record in payload["records"]],
[
("Long one.", "Short one.", "tag one, detail one"),
("Long two.", "Short two.", "tag two, detail two"),
],
)
for image_name, record in zip(("one.png", "two.png"), payload["records"], strict=True):
self.assertEqual(record["image_key"], image_name)
self.assertEqual(set(record["outputs"]), {"long", "short", "taggy"})
self.assertEqual(len(record["sidecar_paths"]), 3)
for path in record["outputs"].values():
self.assertTrue(Path(path).is_file())
self.assertEqual(Path(path).parent, root)
run_outputs = payload["run_outputs"]
for key in (
"caption_jsonl",
"pass_a_jsonl",
"fat_draft_jsonl",
"validator_jsonl",
"taggy_jsonl",
"final_jsonl",
"output_paths_json",
"image_root",
"image_roots",
):
self.assertIn(key, run_outputs)
self.assertEqual(run_outputs["planner_connected"], planner_connected)
self.assertIsInstance(status, str)
self.assertIn("CaptionForge Orchestrator", status)
self.assertIn(f"planner_connected={planner_connected}", status)
self.assertIn("images=2", status)
self.assertIn("final_ok=2 final_failed=0", status)
def test_standalone_returns_aligned_semantic_products_and_records(self) -> None:
self._assert_semantic_outputs(planner_connected=False)
def test_planner_connected_returns_the_same_semantic_products(self) -> None:
self._assert_semantic_outputs(planner_connected=True)
def test_failed_item_keeps_an_empty_aligned_semantic_position(self) -> None:
temp_dir = tempfile.TemporaryDirectory()
self.addCleanup(temp_dir.cleanup)
root = Path(temp_dir.name)
image_names = ("one.png", "two.png", "three.png")
for name in image_names:
Image.new("RGB", (2, 2), "white").save(root / name)
caption_path = root / "captions.jsonl"
source_records = (
{
"image": "one.png",
"image_key": "one.png",
"caption": "usable caption one",
"model_family": "joy",
"status": "ok",
},
{
"image": "two.png",
"image_key": "two.png",
"caption": "FAILED: exception text must never become a caption",
"model_family": "joy",
"status": "error",
},
{
"image": "three.png",
"image_key": "three.png",
"caption": "usable caption three",
"model_family": "joy",
"status": "ok",
},
)
caption_path.write_text(
"\n".join(json.dumps(record) for record in source_records) + "\n",
encoding="utf-8",
)
text_results = iter(
(
"fat draft one",
"SHORT: Short one.\nTAGGY: tag one",
"fat draft three",
"SHORT: Short three.\nTAGGY: tag three",
)
)
validator_results = iter(("Long one.", "Long three."))
with mock.patch.object(
orchestrator, "_evict_python_models_before_ollama_if_needed"
), mock.patch.object(
orchestrator,
"_ollama_generate_text",
side_effect=lambda **_kwargs: (next(text_results), {"ok": True}),
), mock.patch.object(
orchestrator,
"_ollama_chat_image",
side_effect=lambda **_kwargs: (next(validator_results), {"ok": True}),
):
result = orchestrator.JLC_CaptionForge().forge(
**{
"Input - captions JSONL": str(caption_path),
"Input - image path": str(root),
"Output - folder": str(root / "output"),
"Output - run name": "failure-output-contract",
"Output - overwrite outputs": True,
"Final - write TXT sidecars": True,
"Final - write JSONL": True,
}
)
long_captions, short_captions, taggy_captions, final_records_json, status = result
semantic_entries = (
long_captions.split("\n\n"),
short_captions.split("\n\n"),
taggy_captions.split("\n\n"),
)
self.assertEqual(semantic_entries[0], ["Long one.", "", "Long three."])
self.assertEqual(semantic_entries[1], ["Short one.", "", "Short three."])
self.assertEqual(semantic_entries[2], ["tag one", "", "tag three"])
for entries in semantic_entries:
self.assertEqual(len(entries), 3)
self.assertEqual(entries[1], "")
joined = " ".join(entries).lower()
for diagnostic in ("failed", "failure", "error", "exception", "traceback"):
self.assertNotIn(diagnostic, joined)
records = json.loads(final_records_json)["records"]
self.assertEqual(len(records), 3)
self.assertEqual(records[1]["status"], "error")
self.assertEqual(records[1]["error"], "no_usable_captions_selected")
self.assertEqual(
(records[1]["long"], records[1]["short"], records[1]["taggy"]),
("", "", ""),
)
self.assertIn("final_ok=2 final_failed=1", status)
def test_downstream_reintroduction_is_cleaned_with_planner_precedence_and_audited(self) -> None:
temp_dir = tempfile.TemporaryDirectory()
self.addCleanup(temp_dir.cleanup)
root = Path(temp_dir.name)
Image.new("RGB", (2, 2), "white").save(root / "one.png")
captions = root / "captions.jsonl"
captions.write_text(json.dumps({
"image": "one.png", "image_key": "one.png", "caption": "clean witness",
"model_family": "joy", "status": "ok",
}) + "\n", encoding="utf-8")
plan = planner_engine.build_captionforge_pipeline_plan(
output_dir=str(root / "out"), input_path=str(root), run_name="cleanup",
forbidden_phrases="old", replace_pairs="former=>current",
)
plan["paths"]["caption_jsonl"] = str(captions)
plan["paths"]["pass_a_jsonl"] = str(captions)
generated = iter((
"former draft with old but bold holding gold",
"SHORT: former short old bold.\nTAGGY: former, old, gold",
))
with mock.patch.object(orchestrator, "_evict_python_models_before_ollama_if_needed"), \
mock.patch.object(orchestrator, "_ollama_generate_text", side_effect=lambda **_: (next(generated), {})), \
mock.patch.object(orchestrator, "_ollama_chat_image", return_value=("former long old bold holding gold.", {})):
result = orchestrator.JLC_CaptionForge().forge(**{
"Input - captions JSONL": str(captions), "Input - image path": str(root),
"Output - folder": str(root / "standalone"), "Output - run name": "ignored",
"Output - overwrite outputs": True, "Cleanup - forbidden phrases": "gold",
"Cleanup - replace pairs": "former=>wrong", "Final - write TXT sidecars": False,
"Final - write JSONL": True, "pipeline_plan": plan,
})
long_text, short_text, taggy_text, payload_text, _ = result
for text_value in (long_text, short_text, taggy_text):
self.assertNotIn(" old", f" {text_value.lower()}")
self.assertNotIn("former", text_value.lower())
self.assertIn("current", text_value.lower())
self.assertIn("bold", long_text)
self.assertIn("holding", long_text)
self.assertIn("gold", long_text)
payload = json.loads(payload_text)
self.assertEqual(payload["run_outputs"]["cleanup"], plan["cleanup"])
self.assertEqual(payload["records"][0]["cleanup"], plan["cleanup"])
run_config = json.loads(Path(plan["paths"]["run_config_json"]).read_text(encoding="utf-8"))
self.assertEqual(run_config["cleanup"], plan["cleanup"])
def test_standalone_orchestrator_cleanup_values_are_effective_and_audited(self) -> None:
temp_dir = tempfile.TemporaryDirectory()
self.addCleanup(temp_dir.cleanup)
root = Path(temp_dir.name)
Image.new("RGB", (2, 2), "white").save(root / "one.png")
captions = root / "captions.jsonl"
captions.write_text(json.dumps({
"image": "one.png", "image_key": "one.png", "caption": "clean witness",
"model_family": "joy", "status": "ok",
}) + "\n", encoding="utf-8")
generated = iter(("former draft old", "SHORT: former short old.\nTAGGY: former, old"))
with mock.patch.object(orchestrator, "_evict_python_models_before_ollama_if_needed"), \
mock.patch.object(orchestrator, "_ollama_generate_text", side_effect=lambda **_: (next(generated), {})), \
mock.patch.object(orchestrator, "_ollama_chat_image", return_value=("former long old.", {})):
result = orchestrator.JLC_CaptionForge().forge(**{
"Input - captions JSONL": str(captions), "Input - image path": str(root),
"Output - folder": str(root / "out"), "Output - run name": "standalone-cleanup",
"Output - overwrite outputs": True, "Cleanup - forbidden phrases": "old",
"Cleanup - replace pairs": "former=>current", "Final - write TXT sidecars": False,
"Final - write JSONL": True,
})
for text_value in result[:3]:
self.assertNotIn("old", text_value.lower())
self.assertIn("current", text_value.lower())
payload = json.loads(result[3])
expected = {
"forbidden_phrases": ["old"],
"replace_pairs": [{"old": "former", "new": "current"}],
"matching": "boundary_safe_case_insensitive",
"order": ["replace_pairs", "forbidden_phrases", "normalize_whitespace_punctuation"],
}
self.assertEqual(payload["records"][0]["cleanup"], expected)
run_config_path = Path(payload["run_outputs"]["run_config_json"])
self.assertEqual(json.loads(run_config_path.read_text(encoding="utf-8"))["cleanup"], expected)
if __name__ == "__main__":
unittest.main()
+297
View File
@@ -0,0 +1,297 @@
"""Focused resilience tests for the CaptionForge Orchestrator."""
from __future__ import annotations
import importlib
import io
import json
import sys
import tempfile
import types
import unittest
import urllib.error
from pathlib import Path
from unittest import mock
from PIL import Image
ROOT = Path(__file__).resolve().parents[1]
def _install_namespace(name: str, path: Path) -> None:
if name in sys.modules:
return
module = types.ModuleType(name)
module.__path__ = [str(path)]
sys.modules[name] = module
_install_namespace("CaptionForge", ROOT)
_install_namespace("CaptionForge.engines", ROOT / "engines")
_install_namespace("CaptionForge.nodes", ROOT / "nodes")
if "folder_paths" not in sys.modules:
folder_paths = types.ModuleType("folder_paths")
folder_paths.get_output_directory = lambda: str(ROOT / "output")
sys.modules["folder_paths"] = folder_paths
orchestrator = importlib.import_module("CaptionForge.nodes.jlc_captionforge_node")
def _http_error(status: int) -> urllib.error.HTTPError:
return urllib.error.HTTPError(
"http://127.0.0.1:11434/api/chat",
status,
"Request Entity Too Large" if status == 413 else "HTTP failure",
{},
io.BytesIO(b"error"),
)
def _wrapped_http_error(status: int) -> RuntimeError:
error = RuntimeError(f"HTTP {status} from Ollama")
error.__cause__ = _http_error(status)
return error
def _completed_record(image_key: str, *, status: str = "ok") -> dict:
complete = status == "ok"
return {
"captionforge_pass": "D_FINAL_EXPORT",
"image_key": image_key,
"status": status,
"final_caption": "Existing long." if complete else "",
"long": "Existing long." if complete else "",
"short": "Existing short." if complete else "",
"taggy": "existing, tags" if complete else "",
}
class OrchestratorResilienceTests(unittest.TestCase):
def _run(
self,
image_keys: tuple[str, ...],
*,
overwrite: bool = True,
max_size: int | None = None,
existing_final_records: tuple[dict, ...] = (),
text_side_effect=None,
validator_side_effect=None,
) -> tuple[tuple[str, ...], Path, mock.Mock, mock.Mock, mock.Mock]:
temp_dir = tempfile.TemporaryDirectory()
self.addCleanup(temp_dir.cleanup)
root = Path(temp_dir.name)
for image_key in image_keys:
Image.new("RGB", (4, 2), "white").save(root / f"{image_key}.png")
caption_path = root / "captions.jsonl"
caption_path.write_text(
"\n".join(
json.dumps(
{
"image": f"{image_key}.png",
"image_key": image_key,
"caption": f"source caption for {image_key}",
"model_family": "joy",
"status": "ok",
}
)
for image_key in image_keys
)
+ "\n",
encoding="utf-8",
)
output_dir = root / "output"
run_name = "resilience"
plan: dict = {}
if max_size is not None:
plan = {
"captionforge_config_type": "captionforge_pipeline_plan",
"caption_settings": {"max_size": max_size},
}
paths = orchestrator._derive_paths(output_dir, run_name, plan)
final_path = Path(paths["final_jsonl"])
if existing_final_records:
orchestrator._write_jsonl(final_path, list(existing_final_records), append=True)
if text_side_effect is None:
text_side_effect = lambda **_kwargs: (
"SHORT: Generated short.\nTAGGY: generated, tags",
{"ok": True},
)
if validator_side_effect is None:
validator_side_effect = lambda **_kwargs: ("Generated long caption.", {"ok": True})
kwargs = {
"Input - captions JSONL": str(caption_path),
"Input - image path": str(root),
"Output - folder": str(output_dir),
"Output - run name": run_name,
"Output - overwrite outputs": overwrite,
"Final - write TXT sidecars": False,
"Final - write JSONL": True,
}
if plan:
kwargs["pipeline_plan"] = plan
original_prepare = orchestrator._pil_to_base64_png
with mock.patch.object(
orchestrator, "_evict_python_models_before_ollama_if_needed"
), mock.patch.object(
orchestrator, "_ollama_generate_text", side_effect=text_side_effect
) as text_mock, mock.patch.object(
orchestrator, "_ollama_chat_image", side_effect=validator_side_effect
) as validator_mock, mock.patch.object(
orchestrator, "_pil_to_base64_png", wraps=original_prepare
) as prepare_mock:
result = orchestrator.JLC_CaptionForge().forge(**kwargs)
return result, final_path, text_mock, validator_mock, prepare_mock
def test_middle_image_failure_is_isolated_and_later_image_completes(self) -> None:
def text_response(**kwargs):
if "source caption for middle" in kwargs["prompt"]:
raise RuntimeError("synthetic distiller failure")
return "SHORT: Generated short.\nTAGGY: generated, tags", {"ok": True}
result, _, _, validator_mock, _ = self._run(
("first", "middle", "third"),
text_side_effect=text_response,
)
payload = json.loads(result[3])
records = payload["records"]
self.assertEqual([record["status"] for record in records], ["ok", "error", "ok"])
self.assertEqual(records[1]["image_key"], "middle")
self.assertEqual(records[1]["error_stage"], "distiller")
self.assertEqual(records[1]["error_type"], "RuntimeError")
self.assertEqual(records[1]["error_message"], "synthetic distiller failure")
self.assertEqual(validator_mock.call_count, 2)
self.assertIn("final_ok=2 final_failed=1", result[4])
def test_resume_skips_only_latest_complete_final_records(self) -> None:
existing = (
_completed_record("a", status="error"),
_completed_record("a", status="ok"),
_completed_record("b", status="ok"),
_completed_record("b", status="error"),
)
result, final_path, text_mock, validator_mock, _ = self._run(
("a", "b", "c"),
overwrite=False,
existing_final_records=existing,
)
self.assertEqual(text_mock.call_count, 4)
self.assertEqual(validator_mock.call_count, 2)
self.assertIn("resume_skipped=1", result[4])
self.assertEqual(json.loads(result[3])["run_outputs"]["resume_skipped"], 1)
ledger = orchestrator._read_jsonl(final_path)
successful_a = [
record
for record in ledger
if record.get("image_key") == "a" and record.get("status") == "ok"
]
self.assertEqual(len(successful_a), 1)
self.assertEqual(ledger[-2]["image_key"], "b")
self.assertEqual(ledger[-1]["image_key"], "c")
def test_resume_completion_uses_latest_usable_record_and_ignores_bad_line(self) -> None:
temp_dir = tempfile.TemporaryDirectory()
self.addCleanup(temp_dir.cleanup)
ledger_path = Path(temp_dir.name) / "final.jsonl"
partial_b = _completed_record("b")
partial_b["short"] = ""
records = (
_completed_record("a", status="error"),
_completed_record("a"),
_completed_record("b"),
partial_b,
_completed_record("c"),
_completed_record("c", status="error"),
)
orchestrator._write_jsonl(ledger_path, list(records), append=True)
with ledger_path.open("a", encoding="utf-8") as ledger:
ledger.write("{interrupted")
completed = orchestrator._completed_final_records(ledger_path)
self.assertEqual(set(completed), {"a"})
def test_overwrite_true_ignores_existing_completion_and_runs_fresh(self) -> None:
result, final_path, text_mock, validator_mock, _ = self._run(
("a", "b", "c"),
overwrite=True,
existing_final_records=(_completed_record("a"),),
)
self.assertEqual(text_mock.call_count, 6)
self.assertEqual(validator_mock.call_count, 3)
self.assertIn("resume_skipped=0", result[4])
self.assertEqual(len(orchestrator._read_jsonl(final_path)), 3)
def test_validator_413_retries_1024_then_succeeds_at_768(self) -> None:
result, _, _, validator_mock, prepare_mock = self._run(
("image",),
max_size=1024,
validator_side_effect=(_wrapped_http_error(413), ("Recovered long caption.", {"ok": True})),
)
self.assertEqual(result[4].count("final_ok=1"), 1)
self.assertEqual(validator_mock.call_count, 2)
self.assertEqual(
[call.kwargs["max_size"] for call in prepare_mock.call_args_list],
[1024, 768],
)
def test_unrelated_validator_http_error_is_not_retried(self) -> None:
result, _, _, validator_mock, prepare_mock = self._run(
("image",),
max_size=1024,
validator_side_effect=_wrapped_http_error(500),
)
record = json.loads(result[3])["records"][0]
self.assertEqual(record["status"], "error")
self.assertEqual(record["error_stage"], "validator")
self.assertEqual(validator_mock.call_count, 1)
self.assertEqual(len(prepare_mock.call_args_list), 1)
def test_exhausted_413_retries_fail_one_image_and_continue(self) -> None:
result, _, _, validator_mock, prepare_mock = self._run(
("first", "second"),
max_size=768,
validator_side_effect=(
_wrapped_http_error(413),
_wrapped_http_error(413),
("Second image completed.", {"ok": True}),
),
)
records = json.loads(result[3])["records"]
self.assertEqual([record["status"] for record in records], ["error", "ok"])
self.assertEqual(records[0]["error_stage"], "validator")
self.assertEqual(records[1]["image_key"], "second")
self.assertEqual(validator_mock.call_count, 3)
self.assertEqual(
[call.kwargs["max_size"] for call in prepare_mock.call_args_list],
[768, 512, 768],
)
self.assertIn("final_ok=1 final_failed=1", result[4])
def test_validator_retry_ladders_are_strictly_smaller(self) -> None:
self.assertEqual(orchestrator._validator_image_retry_caps(0), (0, 1024, 768, 512))
self.assertEqual(orchestrator._validator_image_retry_caps(1536), (1536, 1024, 768, 512))
self.assertEqual(orchestrator._validator_image_retry_caps(1024), (1024, 768, 512))
self.assertEqual(orchestrator._validator_image_retry_caps(768), (768, 512))
self.assertEqual(orchestrator._validator_image_retry_caps(512), (512,))
def test_explicit_user_abort_is_not_isolated(self) -> None:
with self.assertRaises(KeyboardInterrupt):
self._run(("image",), validator_side_effect=KeyboardInterrupt())
if __name__ == "__main__":
unittest.main()
+203
View File
@@ -0,0 +1,203 @@
"""Focused tests for Pass-A source identity and artifact persistence."""
from __future__ import annotations
import importlib
import json
import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest import mock
from PIL import Image
ROOT = Path(__file__).resolve().parents[1]
def _install_namespace(name: str, path: Path) -> None:
if name in sys.modules:
return
module = types.ModuleType(name)
module.__path__ = [str(path)]
sys.modules[name] = module
_install_namespace("CaptionForge", ROOT)
_install_namespace("CaptionForge.engines", ROOT / "engines")
_install_namespace("CaptionForge.nodes", ROOT / "nodes")
_install_namespace("CaptionForge.nodes.caption_nodes", ROOT / "nodes" / "caption_nodes")
if "folder_paths" not in sys.modules:
folder_paths = types.ModuleType("folder_paths")
folder_paths.models_dir = str(ROOT / "models")
folder_paths.get_output_directory = lambda: str(ROOT / "output")
sys.modules["folder_paths"] = folder_paths
elif not hasattr(sys.modules["folder_paths"], "models_dir"):
sys.modules["folder_paths"].models_dir = str(ROOT / "models")
planner_engine = importlib.import_module("CaptionForge.engines.captionforge_pipeline_planner_engine")
source_identity = importlib.import_module("CaptionForge.engines.captionforge_source_identity")
orchestrator = importlib.import_module("CaptionForge.nodes.jlc_captionforge_node")
joy = importlib.import_module("CaptionForge.nodes.caption_nodes.jlc_captionforge_joy_caption_node")
qwen = importlib.import_module("CaptionForge.nodes.caption_nodes.jlc_captionforge_qwen_caption_node")
ollama = importlib.import_module("CaptionForge.nodes.caption_nodes.jlc_captionforge_ollama_caption_node")
FILE_SOURCE_IDENTITIES = {
"_DSC2094-Edit-2.jpg": "_DSC2094-Edit-2",
"ChatGPT Image May 26, 2026, 08_21_31 AM.png": "ChatGPT Image May 26, 2026, 08_21_31 AM",
"Set A/Sub Folder/_DSC2094-Edit-2.jpg": "_DSC2094-Edit-2",
"Set A/photo.jpg": "photo",
"Set A/photo.png": "photo",
"Set B/photo.jpg": "photo",
"comfy_image_0000.png": "comfy_image_0000",
}
SOURCE_IDENTITIES = {
**FILE_SOURCE_IDENTITIES,
"captionforge-optional-image://comfy_image_0000.png": "comfy_image_0000",
}
def _required_defaults(node_class) -> dict[str, object]:
return {
name: spec[1]["default"]
for name, spec in node_class.INPUT_TYPES()["required"].items()
}
class PassAArtifactIdentityTests(unittest.TestCase):
def test_each_witness_writes_collision_safe_raw_identity_without_txt_sidecars(self) -> None:
with tempfile.TemporaryDirectory() as temp_dir:
root = Path(temp_dir)
input_dir = root / "inputs"
output_dir = root / "outputs"
input_dir.mkdir()
for relative_name in FILE_SOURCE_IDENTITIES:
image_path = input_dir / Path(*relative_name.split("/"))
image_path.parent.mkdir(parents=True, exist_ok=True)
Image.new("RGB", (2, 2), "white").save(image_path)
optional_pil = Image.new("RGB", (2, 2), "black")
plan = planner_engine.build_captionforge_pipeline_plan(
output_dir=str(output_dir),
input_path=str(input_dir),
run_name="captionforge_run",
joy_runs_per_image=1,
qwen_runs_per_image=1,
ollama_runs_per_image=1,
)
joy_engine = mock.Mock()
joy_engine.config = types.SimpleNamespace(max_size=0)
joy_engine.local_model_path = ""
joy_engine.caption_pil.return_value = ("joy caption", "joy raw")
joy_engine.build_run_config.return_value = {}
joy_kwargs = _required_defaults(joy.JLC_CaptionForgeJoy)
joy_kwargs["pipeline_plan"] = plan
with mock.patch.object(joy, "JoyCaptionEngine", return_value=joy_engine), mock.patch.object(
joy, "_tensor_to_pil", return_value=[optional_pil]
):
joy.JLC_CaptionForgeJoy().caption(**joy_kwargs)
qwen_engine = mock.Mock()
qwen_engine.config = types.SimpleNamespace(max_size=0)
qwen_engine.local_model_path = ""
qwen_engine.caption_pil.return_value = ("qwen caption", "qwen raw")
qwen_engine.build_run_config.return_value = {}
qwen_kwargs = _required_defaults(qwen.JLC_CaptionForgeQwen)
qwen_kwargs["pipeline_plan"] = plan
with mock.patch.object(qwen, "QwenCaptionEngine", return_value=qwen_engine), mock.patch.object(
qwen, "_tensor_to_pil", return_value=[optional_pil]
):
qwen.JLC_CaptionForgeQwen().caption(**qwen_kwargs)
ollama_kwargs = _required_defaults(ollama.JLC_CaptionForgeOllamaCaption)
ollama_kwargs["pipeline_plan"] = plan
with mock.patch.object(ollama, "_ensure_ollama_model"), mock.patch.object(
ollama, "_persist_caption_model_if_possible"
), mock.patch.object(
ollama, "_evict_python_models_before_ollama_if_needed"
), mock.patch.object(
ollama, "_ollama_generate_caption", return_value="ollama caption"
), mock.patch.object(
ollama, "_tensor_to_pil", return_value=[optional_pil]
):
ollama.JLC_CaptionForgeOllamaCaption().caption(**ollama_kwargs)
raw_path = Path(plan["paths"]["caption_jsonl"])
records = [json.loads(line) for line in raw_path.read_text(encoding="utf-8").splitlines()]
self.assertEqual(len(records), len(SOURCE_IDENTITIES) * 3)
for family in ("joy", "qwen", "ollama"):
family_records = [record for record in records if record["model_family"] == family]
actual = {record["image_key"]: record["image"] for record in family_records}
self.assertEqual(actual, SOURCE_IDENTITIES)
grouped = orchestrator._group_records_by_image(records)
self.assertEqual(set(grouped), set(SOURCE_IDENTITIES))
self.assertTrue(all(len(group) == 3 for group in grouped.values()))
self.assertEqual(list(root.rglob("*.txt")), [])
opt_images_dir = root / "opt_images"
opt_images_dir.mkdir()
optional_path = opt_images_dir / "comfy_image_0000.png"
optional_pil.save(optional_path)
optional_key = source_identity.optional_image_identity(0)[1]
self.assertEqual(optional_key, "captionforge-optional-image://comfy_image_0000.png")
self.assertEqual(
orchestrator._resolve_image_path_for_group(
grouped[optional_key],
[input_dir, opt_images_dir],
optional_image_root=opt_images_dir,
),
optional_path,
)
self.assertEqual(
orchestrator._resolve_image_path_for_group(
grouped["Set A/Sub Folder/_DSC2094-Edit-2.jpg"],
[input_dir, opt_images_dir],
optional_image_root=opt_images_dir,
),
input_dir / "Set A" / "Sub Folder" / "_DSC2094-Edit-2.jpg",
)
self.assertEqual(
orchestrator._resolve_image_path_for_group(
grouped["Set A/photo.jpg"],
[input_dir, opt_images_dir],
optional_image_root=opt_images_dir,
),
input_dir / "Set A" / "photo.jpg",
)
empty_dataset = root / "empty-inputs"
empty_dataset.mkdir()
self.assertEqual(
orchestrator._resolve_image_path_for_group(
[{"image": "comfy_image_0000", "image_key": "comfy_image_0000"}],
[empty_dataset, opt_images_dir],
optional_image_root=opt_images_dir,
),
optional_path,
)
self.assertEqual(
orchestrator._resolve_image_path_for_group(
[{"image": "photo", "image_key": r"Set A\photo.jpg"}],
[input_dir, opt_images_dir],
optional_image_root=opt_images_dir,
),
input_dir / "Set A" / "photo.jpg",
)
def test_raw_response_filenames_do_not_collapse_relative_keys(self) -> None:
with tempfile.TemporaryDirectory() as temp_dir:
raw_dir = Path(temp_dir)
first = orchestrator._save_raw_response(raw_dir, "Set A/photo.jpg", "stage", {"ok": 1})
second = orchestrator._save_raw_response(raw_dir, "Set B/photo.jpg", "stage", {"ok": 2})
extension_variant = orchestrator._save_raw_response(raw_dir, "Set A/photo.png", "stage", {"ok": 3})
self.assertEqual(len({first, second, extension_variant}), 3)
self.assertTrue(all(Path(path).exists() for path in (first, second, extension_variant)))
if __name__ == "__main__":
unittest.main()
+264
View File
@@ -0,0 +1,264 @@
"""Regression tests for Qwen 8-bit low-VRAM CPU offload behavior."""
from __future__ import annotations
import ast
import importlib
import sys
import tempfile
import types
import unittest
from types import SimpleNamespace
from pathlib import Path
from unittest import mock
import torch
ROOT = Path(__file__).resolve().parents[1]
def _install_namespace(name: str, path: Path) -> None:
if name in sys.modules:
return
module = types.ModuleType(name)
module.__path__ = [str(path)]
sys.modules[name] = module
_install_namespace("CaptionForge", ROOT)
_install_namespace("CaptionForge.engines", ROOT / "engines")
qwen = importlib.import_module("CaptionForge.engines.jlc_qwen_caption_engine")
class QwenOffloadTests(unittest.TestCase):
def test_generation_uses_scoped_bnb_diagnostic_handling(self):
tree = ast.parse(
(ROOT / "engines/jlc_qwen_caption_engine.py").read_text(encoding="utf-8")
)
engine_class = next(
node
for node in tree.body
if isinstance(node, ast.ClassDef) and node.name == "QwenCaptionEngine"
)
caption_method = next(
node
for node in engine_class.body
if isinstance(node, ast.FunctionDef) and node.name == "caption_pil"
)
scoped_generate_calls = []
for node in ast.walk(caption_method):
if not isinstance(node, ast.With):
continue
uses_scope = any(
isinstance(item.context_expr, ast.Call)
and isinstance(item.context_expr.func, ast.Name)
and item.context_expr.func.id == "quantized_inference_warnings"
for item in node.items
)
if uses_scope:
scoped_generate_calls.extend(
child
for child in ast.walk(node)
if isinstance(child, ast.Call)
and isinstance(child.func, ast.Attribute)
and child.func.attr == "generate"
)
self.assertEqual(len(scoped_generate_calls), 1)
self.assertNotIn("filterwarnings", ast.unparse(caption_method))
def test_memory_budget_reserves_headroom_and_fp32_cpu_cost(self):
gib = 1024 ** 3
with mock.patch.object(qwen.torch.cuda, "is_available", return_value=True), \
mock.patch.object(qwen.torch.cuda, "current_device", return_value=0), \
mock.patch.object(qwen.torch.cuda, "mem_get_info", return_value=(10 * gib, 12 * gib)), \
mock.patch.object(qwen, "_available_system_memory_bytes", return_value=40 * gib):
budget = qwen._qwen_memory_budget(headroom=0.20)
# GPU receives 80% of currently free VRAM.
self.assertEqual(budget[0], int(10 * gib * 0.80))
# CPU placement is estimated with dtype=int8 but actual bitsandbytes
# CPU overflow remains FP32, so CPU capacity is divided by four.
self.assertEqual(budget["cpu"], int(40 * gib * 0.80 / 4.0))
def test_execution_device_prefers_accelerate_hook(self):
model = SimpleNamespace(
_hf_hook=SimpleNamespace(execution_device="cuda:0"),
hf_device_map={"": "cpu"},
)
device = qwen._resolve_model_execution_device(model)
self.assertEqual(device, torch.device("cuda:0"))
def test_execution_device_prefers_gpu_from_device_map(self):
model = SimpleNamespace(
_hf_hook=None,
hf_device_map={
"visual": "cpu",
"model.layers.0": 0,
"model.layers.1": "cpu",
},
)
device = qwen._resolve_model_execution_device(model)
self.assertEqual(device, torch.device("cuda:0"))
def test_execution_device_falls_back_to_parameter_device(self):
parameter = torch.nn.Parameter(torch.zeros(1))
model = torch.nn.Linear(1, 1)
model.weight = parameter
device = qwen._resolve_model_execution_device(model)
self.assertEqual(device, parameter.device)
def test_device_map_summary(self):
summary = qwen._summarize_device_map(
{
"visual": 0,
"layer0": 0,
"layer1": "cpu",
}
)
self.assertIn("cuda:0: 2 module(s)", summary)
self.assertIn("cpu: 1 module(s)", summary)
def test_load_uses_supported_8bit_cpu_offload_and_preserves_preload_eviction(self):
events = []
captured = {}
class FakeBitsAndBytesConfig:
def __init__(self, **kwargs):
captured["bnb"] = dict(kwargs)
class FakeProcessor:
@classmethod
def from_pretrained(cls, *_args, **_kwargs):
return cls()
class FakeModel:
hf_device_map = {
"visual": 0,
"model.layers.0": 0,
"model.layers.1": "cpu",
}
def eval(self):
return self
class FakeModelClass:
@classmethod
def from_pretrained(cls, *_args, **kwargs):
events.append("load")
captured["model_kwargs"] = dict(kwargs)
return FakeModel()
fake_transformers = types.ModuleType("transformers")
fake_transformers.AutoProcessor = FakeProcessor
fake_transformers.BitsAndBytesConfig = FakeBitsAndBytesConfig
with tempfile.TemporaryDirectory() as temp_dir:
local_path = Path(temp_dir)
config = qwen.QwenCaptionConfig(
model_path=str(local_path),
quantization="bnb_8bit",
device_map="auto",
keep_loaded=True,
allow_download=False,
)
engine = qwen.QwenCaptionEngine(config)
def fake_prepare(*_args, **_kwargs):
events.append("prepare")
with mock.patch.dict(sys.modules, {"transformers": fake_transformers}), \
mock.patch.object(qwen, "get_cached_model", return_value=None), \
mock.patch.object(qwen, "prepare_for_model_load", side_effect=fake_prepare), \
mock.patch.object(qwen, "register_model"), \
mock.patch.object(engine, "_load_model_class", return_value=FakeModelClass), \
mock.patch.object(engine, "_detect_model_type", return_value="qwen2_5_vl"), \
mock.patch.object(
qwen,
"_build_qwen_8bit_device_map",
return_value=(
{
"visual": 0,
"model.layers.0": 0,
"model.layers.1": "cpu",
},
{0: 8 * 1024**3, "cpu": 8 * 1024**3},
),
):
engine.load()
self.assertEqual(events[:2], ["prepare", "load"])
self.assertTrue(captured["bnb"]["load_in_8bit"])
self.assertTrue(
captured["bnb"]["llm_int8_enable_fp32_cpu_offload"]
)
self.assertEqual(
captured["model_kwargs"]["device_map"]["model.layers.1"],
"cpu",
)
self.assertIn("max_memory", captured["model_kwargs"])
def test_adaptive_mapper_rejects_disk_spill(self):
class FakeModel:
_no_split_modules = []
class FakeModelClass:
def __new__(cls, *_args, **_kwargs):
return FakeModel()
fake_accelerate = SimpleNamespace(
init_empty_weights=mock.MagicMock(),
infer_auto_device_map=mock.MagicMock(
return_value={
"visual": 0,
"model.layers.0": "cpu",
"model.layers.1": "disk",
}
),
)
class FakeContext:
def __enter__(self):
return None
def __exit__(self, exc_type, exc, tb):
return False
fake_accelerate.init_empty_weights.return_value = FakeContext()
fake_transformers = SimpleNamespace(
AutoConfig=SimpleNamespace(
from_pretrained=mock.MagicMock(return_value=object())
)
)
modules = {
"accelerate": fake_accelerate,
"transformers": fake_transformers,
}
with mock.patch.dict("sys.modules", modules), \
mock.patch.object(qwen.torch.cuda, "is_available", return_value=True), \
mock.patch.object(qwen, "_qwen_memory_budget", return_value={0: 1, "cpu": 1}):
with self.assertRaisesRegex(RuntimeError, "disk offload"):
qwen._build_qwen_8bit_device_map(
FakeModelClass,
qwen.Path("."),
True,
)
if __name__ == "__main__":
unittest.main()
+126
View File
@@ -0,0 +1,126 @@
"""CPU-only release metadata, module documentation, and node-help checks."""
from __future__ import annotations
import ast
import importlib
import json
import re
import sys
import types
import unittest
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
def _install_namespace(name: str, path: Path) -> None:
if name in sys.modules:
return
module = types.ModuleType(name)
module.__path__ = [str(path)]
sys.modules[name] = module
_install_namespace("CaptionForge", ROOT)
_install_namespace("CaptionForge.engines", ROOT / "engines")
_install_namespace("CaptionForge.nodes", ROOT / "nodes")
_install_namespace("CaptionForge.nodes.caption_nodes", ROOT / "nodes" / "caption_nodes")
if "folder_paths" not in sys.modules:
folder_paths = types.ModuleType("folder_paths")
folder_paths.models_dir = str(ROOT / "models")
folder_paths.get_output_directory = lambda: str(ROOT / "output")
sys.modules["folder_paths"] = folder_paths
elif not hasattr(sys.modules["folder_paths"], "models_dir"):
sys.modules["folder_paths"].models_dir = str(ROOT / "models")
version_module = importlib.import_module("CaptionForge.captionforge_version")
planner = importlib.import_module("CaptionForge.nodes.jlc_captionforge_pipeline_planner_node")
capstone = importlib.import_module("CaptionForge.nodes.jlc_captionforge_node")
template_options = importlib.import_module("CaptionForge.nodes.jlc_captionforge_template_options")
joy = importlib.import_module("CaptionForge.nodes.caption_nodes.jlc_captionforge_joy_caption_node")
qwen = importlib.import_module("CaptionForge.nodes.caption_nodes.jlc_captionforge_qwen_caption_node")
ollama = importlib.import_module("CaptionForge.nodes.caption_nodes.jlc_captionforge_ollama_caption_node")
class ReleaseMetadataTests(unittest.TestCase):
def test_current_release_versions_are_1_0_2(self) -> None:
self.assertEqual(version_module.CAPTIONFORGE_VERSION, "1.0.2")
pyproject = (ROOT / "pyproject.toml").read_text(encoding="utf-8")
match = re.search(r'^version\s*=\s*"([^"]+)"', pyproject, re.MULTILINE)
self.assertIsNotNone(match)
self.assertEqual(match.group(1), version_module.CAPTIONFORGE_VERSION)
config = json.loads(
(ROOT / "config" / "captionforge_ollama_models.json").read_text(encoding="utf-8")
)
self.assertEqual(config["_meta"]["version"], version_module.CAPTIONFORGE_VERSION)
self.assertEqual(capstone.CAPTIONFORGE_NODE_VERSION, version_module.CAPTIONFORGE_VERSION)
def test_anchor_tooltips_describe_frozen_1_x_behavior(self) -> None:
expected = (
"Optional persistent caption/training anchor. In CaptionForge 1.x, a non-empty "
"anchor is preserved in the final caption variants rather than treated as image "
"evidence that the Validator may remove."
)
self.assertEqual(
planner.JLC_CaptionForge_Pipeline_Planner.INPUT_TYPES()["required"]
["LoRA - user caption anchor"][1]["tooltip"],
expected,
)
self.assertEqual(
capstone.JLC_CaptionForge.INPUT_TYPES()["required"]
["LoRA - user caption anchor"][1]["tooltip"],
expected,
)
def test_package_discovery_is_explicit_and_production_only(self) -> None:
pyproject = (ROOT / "pyproject.toml").read_text(encoding="utf-8")
self.assertNotIn("[tool.setuptools.packages.find]", pyproject)
for package in (
"CaptionForge",
"CaptionForge.engines",
"CaptionForge.nodes",
"CaptionForge.nodes.caption_nodes",
):
self.assertIn(f'"{package}"', pyproject)
for excluded in (".backups", ".tests", ".bak", ".deprecated", "experimental"):
self.assertNotIn(f'"{excluded}', pyproject)
def test_active_python_modules_have_module_docstrings(self) -> None:
files = [ROOT / "__init__.py", ROOT / "captionforge_version.py"]
files.extend((ROOT / "engines").glob("*.py"))
files.extend((ROOT / "nodes").glob("*.py"))
files.extend((ROOT / "nodes" / "caption_nodes").glob("*.py"))
files.extend((ROOT / "tests").glob("*.py"))
for path in files:
with self.subTest(path=path.relative_to(ROOT)):
tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
self.assertTrue(ast.get_docstring(tree))
class ActiveNodeHelpTests(unittest.TestCase):
def test_every_active_input_has_plain_help_text(self) -> None:
node_classes = (
planner.JLC_CaptionForge_Pipeline_Planner,
capstone.JLC_CaptionForge,
template_options.JLC_CaptionForgeExtraOptions,
joy.JLC_CaptionForgeJoy,
qwen.JLC_CaptionForgeQwen,
ollama.JLC_CaptionForgeOllamaCaption,
)
for node_class in node_classes:
inputs = node_class.INPUT_TYPES()
for section in ("required", "optional"):
for name, spec in inputs.get(section, {}).items():
with self.subTest(node=node_class.__name__, input=name):
self.assertGreaterEqual(len(spec), 2)
metadata = spec[1]
self.assertIsInstance(metadata, dict)
self.assertTrue(str(metadata.get("tooltip", "")).strip())
if __name__ == "__main__":
unittest.main()
+300
View File
@@ -0,0 +1,300 @@
"""CPU tests for the shared Pass-A and standalone seed contracts."""
from __future__ import annotations
import importlib
import json
import random
import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest import mock
from PIL import Image
ROOT = Path(__file__).resolve().parents[1]
def _install_namespace(name: str, path: Path) -> None:
if name in sys.modules:
return
module = types.ModuleType(name)
module.__path__ = [str(path)]
sys.modules[name] = module
_install_namespace("CaptionForge", ROOT)
_install_namespace("CaptionForge.engines", ROOT / "engines")
_install_namespace("CaptionForge.nodes", ROOT / "nodes")
_install_namespace("CaptionForge.nodes.caption_nodes", ROOT / "nodes" / "caption_nodes")
if "folder_paths" not in sys.modules:
folder_paths = types.ModuleType("folder_paths")
folder_paths.get_output_directory = lambda: str(ROOT / "output")
sys.modules["folder_paths"] = folder_paths
planner_engine = importlib.import_module("CaptionForge.engines.captionforge_pipeline_planner_engine")
planner_node = importlib.import_module("CaptionForge.nodes.jlc_captionforge_pipeline_planner_node")
capstone = importlib.import_module("CaptionForge.nodes.jlc_captionforge_node")
ollama_caption = importlib.import_module("CaptionForge.nodes.caption_nodes.jlc_captionforge_ollama_caption_node")
class PassASeedContractTests(unittest.TestCase):
def _schedule(self, base_seed: int, seed_mode: str, runs: int = 3) -> list[int | None]:
plan = planner_engine.build_captionforge_pipeline_plan(
output_dir="unused",
joy_runs_per_image=runs,
qwen_runs_per_image=0,
ollama_runs_per_image=0,
base_seed=base_seed,
seed_mode=seed_mode,
)
return [run.seed for run in planner_engine.expand_captionforge_runs(plan, model_key="joy")]
def test_fixed_increment_and_decrement(self) -> None:
self.assertEqual(self._schedule(100, "fixed"), [100, 100, 100])
self.assertEqual(self._schedule(100, "increment"), [100, 101, 102])
self.assertEqual(self._schedule(100, "decrement"), [100, 99, 98])
def test_uint32_wraparound(self) -> None:
maximum = planner_engine.MAX_SEED_32
self.assertEqual(self._schedule(maximum - 1, "increment"), [maximum - 1, maximum, 0])
self.assertEqual(self._schedule(0, "decrement"), [0, maximum, maximum - 1])
def test_literal_zero_survives(self) -> None:
self.assertEqual(self._schedule(0, "fixed"), [0, 0, 0])
self.assertEqual(self._schedule(0, "increment"), [0, 1, 2])
self.assertEqual(self._schedule(0, "random"), self._schedule(0, "random"))
self.assertTrue(all(seed is not None for seed in self._schedule(0, "random")))
def test_hash_random_has_a_stable_golden_schedule(self) -> None:
expected = [90425488, 1231494799, 4241276224, 2828971892]
self.assertEqual(self._schedule(123, "random", 4), expected)
self.assertEqual(self._schedule(123, "random", 4), expected)
self.assertEqual(len(set(expected)), len(expected))
def test_hash_random_is_independent_of_global_random_state(self) -> None:
state = random.getstate()
try:
random.seed(1)
first = self._schedule(321, "random", 5)
for _ in range(100):
random.random()
random.seed(999999)
second = self._schedule(321, "random", 5)
finally:
random.setstate(state)
self.assertEqual(first, second)
def test_schedule_is_image_ordinal_independent_and_reusable(self) -> None:
expected = self._schedule(100, "increment", 3)
per_image = [self._schedule(100, "increment", 3) for _ in range(4)]
self.assertEqual(per_image, [expected, expected, expected, expected])
def test_minus_one_is_unseeded_for_every_mode(self) -> None:
for mode in ("fixed", "increment", "decrement", "random"):
with self.subTest(mode=mode):
self.assertEqual(self._schedule(-1, mode), [None, None, None])
def test_joy_qwen_and_ollama_share_the_planner_schedule(self) -> None:
plan = planner_engine.build_captionforge_pipeline_plan(
output_dir="unused",
joy_runs_per_image=3,
qwen_runs_per_image=3,
ollama_runs_per_image=3,
base_seed=77,
seed_mode="random",
)
schedules = [
[run.seed for run in planner_engine.expand_captionforge_runs(plan, model_key=family)]
for family in ("joy", "qwen", "ollama")
]
self.assertEqual(schedules[0], schedules[1])
self.assertEqual(schedules[1], schedules[2])
class OllamaPassAUnsetSeedTests(unittest.TestCase):
def test_unset_seed_is_omitted_from_ollama_options(self) -> None:
options = ollama_caption._ollama_options(
max_new_tokens=10,
temperature=0.5,
top_p=0.9,
top_k=20,
repetition_penalty=1.0,
seed=None,
)
self.assertNotIn("seed", options)
def test_planned_unset_seed_reaches_caption_call_without_int_none(self) -> None:
with tempfile.TemporaryDirectory() as temp_dir:
root = Path(temp_dir)
Image.new("RGB", (2, 2), "white").save(root / "image.png")
plan = planner_engine.build_captionforge_pipeline_plan(
output_dir=temp_dir,
input_path=temp_dir,
joy_runs_per_image=0,
qwen_runs_per_image=0,
ollama_runs_per_image=1,
base_seed=-1,
seed_mode="random",
)
required = ollama_caption.JLC_CaptionForgeOllamaCaption.INPUT_TYPES()["required"]
kwargs = {name: spec[1]["default"] for name, spec in required.items()}
kwargs["pipeline_plan"] = plan
with mock.patch.object(ollama_caption, "_ensure_ollama_model"), mock.patch.object(
ollama_caption, "_persist_caption_model_if_possible"
), mock.patch.object(
ollama_caption, "_evict_python_models_before_ollama_if_needed"
), mock.patch.object(
ollama_caption, "_ollama_generate_caption", return_value="caption"
) as generate:
ollama_caption.JLC_CaptionForgeOllamaCaption().caption(**kwargs)
self.assertIsNone(generate.call_args.kwargs["seed"])
class StageSeedContractTests(unittest.TestCase):
def _run_capstone(
self,
*,
standalone_seeds: tuple[int | None, int | None, int | None],
plan_seeds: tuple[int, int, int] | None = None,
) -> tuple[list[int | None], list[int | None]]:
with tempfile.TemporaryDirectory() as temp_dir:
root = Path(temp_dir)
for name in ("one.png", "two.png"):
Image.new("RGB", (2, 2), "white").save(root / name)
caption_path = root / "captions.jsonl"
caption_path.write_text(
"\n".join(
json.dumps(
{
"image": name,
"image_key": Path(name).stem,
"caption": f"caption for {name}",
"model_family": "joy",
"status": "ok",
}
)
for name in ("one.png", "two.png")
)
+ "\n",
encoding="utf-8",
)
kwargs = {
"Input - captions JSONL": str(caption_path),
"Input - image path": temp_dir,
"Output - folder": str(root / "output"),
"Output - run name": "seed-test",
"Output - overwrite outputs": True,
"Distiller seed": standalone_seeds[0],
"Validator seed": standalone_seeds[1],
"Formatter seed": standalone_seeds[2],
}
if plan_seeds is not None:
kwargs["pipeline_plan"] = {
"captionforge_config_type": "captionforge_pipeline_plan",
"distiller": {"seed": plan_seeds[0]},
"validator": {"seed": plan_seeds[1]},
"formatter": {"seed": plan_seeds[2]},
}
text_seeds: list[int | None] = []
image_seeds: list[int | None] = []
def fake_text(**call_kwargs):
text_seeds.append(call_kwargs["seed"])
return "generated text", {"ok": True}
def fake_image(**call_kwargs):
image_seeds.append(call_kwargs["seed"])
return "validated natural caption", {"ok": True}
with mock.patch.object(capstone, "_evict_python_models_before_ollama_if_needed"), mock.patch.object(
capstone, "_ollama_generate_text", side_effect=fake_text
), mock.patch.object(capstone, "_ollama_chat_image", side_effect=fake_image):
capstone.JLC_CaptionForge().forge(**kwargs)
return text_seeds, image_seeds
def test_standalone_exact_stage_seeds_repeat_for_every_image(self) -> None:
text_seeds, image_seeds = self._run_capstone(standalone_seeds=(200, 300, 400))
self.assertEqual(text_seeds, [200, 400, 200, 400])
self.assertEqual(image_seeds, [300, 300])
def test_planner_overrides_all_standalone_stage_seed_inputs(self) -> None:
text_seeds, image_seeds = self._run_capstone(
standalone_seeds=(1, 2, 3),
plan_seeds=(200, 300, 400),
)
self.assertEqual(text_seeds, [200, 400, 200, 400])
self.assertEqual(image_seeds, [300, 300])
def test_literal_zero_survives_all_stage_paths(self) -> None:
text_seeds, image_seeds = self._run_capstone(
standalone_seeds=(9, 9, 9),
plan_seeds=(0, 0, 0),
)
self.assertEqual(text_seeds, [0, 0, 0, 0])
self.assertEqual(image_seeds, [0, 0])
def test_omitted_standalone_seeds_are_unseeded_and_do_not_crash(self) -> None:
text_seeds, image_seeds = self._run_capstone(standalone_seeds=(None, None, None))
self.assertEqual(text_seeds, [None, None, None, None])
self.assertEqual(image_seeds, [None, None])
class SeedOwnershipSchemaTests(unittest.TestCase):
def test_planner_preserves_three_exact_stage_seeds_without_modes(self) -> None:
plan = planner_engine.build_captionforge_pipeline_plan(
output_dir="unused",
joy_runs_per_image=1,
qwen_runs_per_image=0,
distiller_seed=200,
validator_seed=300,
formatter_seed=400,
)
self.assertEqual(plan["distiller"]["seed"], 200)
self.assertEqual(plan["validator"]["seed"], 300)
self.assertEqual(plan["formatter"]["seed"], 400)
self.assertNotIn("seed_mode", plan["distiller"])
self.assertNotIn("seed_mode", plan["validator"])
self.assertNotIn("seed_mode", plan["formatter"])
def test_planner_input_schema_has_fixed_stage_seeds_only(self) -> None:
required = planner_node.JLC_CaptionForge_Pipeline_Planner.INPUT_TYPES()["required"]
self.assertIn("Distiller - seed", required)
self.assertIn("Validator - seed", required)
self.assertIn("Formatter - seed", required)
self.assertNotIn("Distiller - seed mode", required)
self.assertNotIn("Validator - seed mode", required)
self.assertNotIn("Formatter - seed mode", required)
def test_capstone_exposes_only_optional_stage_seed_inputs(self) -> None:
input_types = capstone.JLC_CaptionForge.INPUT_TYPES()
required = input_types["required"]
optional = input_types["optional"]
self.assertIn("Distiller seed", optional)
self.assertIn("Validator seed", optional)
self.assertIn("Formatter seed", optional)
for name in ("Distiller seed", "Validator seed", "Formatter seed"):
self.assertTrue(optional[name][1]["forceInput"])
for old_name in (
"Fat Draft - base seed",
"Fat Draft - seed mode",
"Validator - base seed",
"Validator - seed mode",
"Formatter - base seed",
"Formatter - seed mode",
):
self.assertNotIn(old_name, required)
self.assertNotIn(old_name, optional)
if __name__ == "__main__":
unittest.main()
+7 -6
View File
@@ -3,11 +3,12 @@ import { app } from "/scripts/app.js";
const ICON_SIZE = 12;
const CAPTIONFORGE_NODE_NAMES = new Set([
"JLC_QwenCaption",
"JLC_JoyCaption",
"JLC_QwenCaptionLite",
"JLC_JoyCaptionLite",
"JLC_CaptionForgeClaimExtractor",
"JLC_CaptionForge_Pipeline_Planner",
"JLC_CaptionForgeExtraOptions",
"JLC_CaptionForgeJoy",
"JLC_CaptionForgeQwen",
"JLC_CaptionForgeOllamaCaption",
"JLC_CaptionForge",
]);
const iconImage = new Image();
@@ -69,4 +70,4 @@ app.registerExtension({
}
};
},
});
});