Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
33c21baff8 | ||
|
|
b116b08033 | ||
|
|
1bc70ccb2a | ||
|
|
9d5b131b08 | ||
|
|
fd5f519aae | ||
|
|
90753aa7c6 | ||
|
|
a84724a38a | ||
|
|
7fa2871810 | ||
|
|
49080cbae2 | ||
|
|
9f76aae7e7 | ||
|
|
49290e8812 | ||
|
|
f7f770f953 | ||
|
|
a1e613a6ac | ||
|
|
25425d41b5 | ||
|
|
b438fddbc7 | ||
|
|
32a8d4ecc7 | ||
|
|
af785e3d36 | ||
|
|
a373a7e8d9 | ||
|
|
b234e8020c | ||
|
|
6c111930d5 | ||
|
|
31a10619d9 | ||
|
|
5ed231eaad | ||
|
|
435ad73ad6 | ||
|
|
3cf6b3536d | ||
|
|
cb5741b604 | ||
|
|
59c4eb31e9 |
@@ -11,7 +11,7 @@
|
||||
<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.0-blue">
|
||||
<img alt="Version" src="https://img.shields.io/badge/version-1.0.2-blue">
|
||||
</p>
|
||||
|
||||
---
|
||||
@@ -28,16 +28,19 @@ The core idea is simple: a single image captioner can be useful, but it should n
|
||||
|
||||
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.0.**
|
||||
> **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">
|
||||
<img src="assets/workflows/CaptionForge_FullWorkflow.png"
|
||||
<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>
|
||||
@@ -45,9 +48,9 @@ CaptionForge also writes structured JSONL audit records so intermediate evidence
|
||||
|
||||
Canonical workflow files:
|
||||
|
||||
- [`CaptionForge_FullWorkflow.json`](assets/workflows/CaptionForge_FullWorkflow.json) — editable ComfyUI workflow
|
||||
- [`CaptionForge_FullWorkflow_API.json`](assets/workflows/CaptionForge_FullWorkflow_API.json) — API-format workflow
|
||||
- [`CaptionForge_FullWorkflow.png`](assets/workflows/CaptionForge_FullWorkflow.png) — workflow PNG with embedded metadata
|
||||
- [`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.
|
||||
|
||||
@@ -229,6 +232,8 @@ It coordinates:
|
||||
|
||||
In the full workflow, **the Planner is authoritative**.
|
||||
|
||||
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.
|
||||
|
||||
### CaptionForge Joy Caption
|
||||
|
||||
Python/Hugging Face JoyCaption-family Pass A witness.
|
||||
@@ -277,6 +282,8 @@ 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
|
||||
@@ -430,7 +437,7 @@ Use explicit seeds when reproducibility is required.
|
||||
|
||||
## Running the full workflow
|
||||
|
||||
1. Load [`CaptionForge_FullWorkflow.json`](assets/workflows/CaptionForge_FullWorkflow.json).
|
||||
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:
|
||||
@@ -442,11 +449,13 @@ Use explicit seeds when reproducibility is required.
|
||||
|
||||
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.0
|
||||
## Validation for 1.0.1
|
||||
|
||||
The final 1.0.0 release-preparation pass included:
|
||||
The final 1.0.1 release-preparation pass included:
|
||||
|
||||
- Planner ownership/default tests
|
||||
- Pass A artifact and source-identity tests
|
||||
@@ -482,3 +491,4 @@ For large LoRA datasets, the goal is to improve the bulk quality and auditabilit
|
||||
CaptionForge is released under the **MIT License**.
|
||||
|
||||
See [`LICENSE`](LICENSE) for details.
|
||||
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 2.1 MiB |
+169
-181
@@ -1,13 +1,4 @@
|
||||
{
|
||||
"9": {
|
||||
"inputs": {
|
||||
"image": "ComfyUI_00081_.png"
|
||||
},
|
||||
"class_type": "LoadImage",
|
||||
"_meta": {
|
||||
"title": "Load Image"
|
||||
}
|
||||
},
|
||||
"51": {
|
||||
"inputs": {
|
||||
"model": "llama-joycaption-beta-one-hf-llava",
|
||||
@@ -30,11 +21,11 @@
|
||||
"replace_pairs": "",
|
||||
"download_probe_only": false,
|
||||
"image": [
|
||||
"85",
|
||||
"148",
|
||||
0
|
||||
],
|
||||
"pipeline_plan": [
|
||||
"85",
|
||||
"148",
|
||||
1
|
||||
],
|
||||
"template_options": [
|
||||
@@ -49,7 +40,7 @@
|
||||
},
|
||||
"52": {
|
||||
"inputs": {
|
||||
"model": "Qwen2.5-VL-7B-Instruct",
|
||||
"model": "Qwen2.5-VL-7B-NSFW-Caption-V3-abliterated",
|
||||
"qwen_quantization": "Balanced (8-bit)",
|
||||
"keep_loaded": true,
|
||||
"caption_template_mode": true,
|
||||
@@ -128,8 +119,8 @@
|
||||
},
|
||||
"56": {
|
||||
"inputs": {
|
||||
"model": "gemma4:26b",
|
||||
"custom_model_tag": "",
|
||||
"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,
|
||||
@@ -137,49 +128,8 @@
|
||||
"caption_length": "any",
|
||||
"custom_prompt_mode": false,
|
||||
"prompt_preset": "female_character_lora",
|
||||
"system_prompt": "You are a multimodal image captioning engine for female character LoRA dataset preparation.\n\nYour job is to inspect the image and write one strong standalone caption.\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: 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": [
|
||||
"58",
|
||||
0
|
||||
],
|
||||
"pipeline_plan": [
|
||||
"58",
|
||||
1
|
||||
],
|
||||
"template_options": [
|
||||
"58",
|
||||
2
|
||||
]
|
||||
},
|
||||
"class_type": "JLC_CaptionForgeOllamaCaption",
|
||||
"_meta": {
|
||||
"title": " JLC CaptionForge Ollama Caption"
|
||||
}
|
||||
},
|
||||
"58": {
|
||||
"inputs": {
|
||||
"model": "qwen3.6:35B-A3B",
|
||||
"custom_model_tag": "",
|
||||
"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 LoRA dataset preparation.\n\nYour job is to inspect the image and write one strong standalone caption.\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: 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:",
|
||||
"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,
|
||||
@@ -208,73 +158,115 @@
|
||||
"title": " JLC CaptionForge Ollama Caption"
|
||||
}
|
||||
},
|
||||
"67": {
|
||||
"90": {
|
||||
"inputs": {
|
||||
"inputcount": 3,
|
||||
"Update inputs": null,
|
||||
"image_1": [
|
||||
"68",
|
||||
"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
|
||||
],
|
||||
"image_2": [
|
||||
"70",
|
||||
"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
|
||||
],
|
||||
"image_3": [
|
||||
"9",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "ImageBatchMulti",
|
||||
"_meta": {
|
||||
"title": "Image Batch Multi"
|
||||
}
|
||||
},
|
||||
"68": {
|
||||
"inputs": {
|
||||
"image": "Flux.2__00297_.png"
|
||||
},
|
||||
"class_type": "LoadImage",
|
||||
"_meta": {
|
||||
"title": "Load Image"
|
||||
}
|
||||
},
|
||||
"70": {
|
||||
"inputs": {
|
||||
"image": "jlc_Orchestrate_Sample.jpg"
|
||||
},
|
||||
"class_type": "LoadImage",
|
||||
"_meta": {
|
||||
"title": "Load Image"
|
||||
}
|
||||
},
|
||||
"81": {
|
||||
"inputs": {
|
||||
"output": "",
|
||||
"source": [
|
||||
"88",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "Display Any (rgthree)",
|
||||
"_meta": {
|
||||
"title": "Long Caption - Display Any (rgthree)"
|
||||
}
|
||||
},
|
||||
"82": {
|
||||
"inputs": {
|
||||
"output": "",
|
||||
"source": [
|
||||
"88",
|
||||
"pipeline_plan": [
|
||||
"90",
|
||||
1
|
||||
]
|
||||
},
|
||||
"class_type": "Display Any (rgthree)",
|
||||
"class_type": "JLC_CaptionForge",
|
||||
"_meta": {
|
||||
"title": "Short Captions - Display Any (rgthree)"
|
||||
"title": " JLC CaptionForge Orchestrator"
|
||||
}
|
||||
},
|
||||
"85": {
|
||||
"148": {
|
||||
"inputs": {
|
||||
"Planner - enabled": true,
|
||||
"Input - image path": "",
|
||||
@@ -286,16 +278,18 @@
|
||||
"Ollama - URL": "http://127.0.0.1:11434",
|
||||
"Ollama - keep loaded": true,
|
||||
"Ollama - request timeout seconds": 1800,
|
||||
"LoRA - trigger word": "",
|
||||
"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": 3,
|
||||
"Caption - base seed": 1,
|
||||
"Caption - seed mode": "increment",
|
||||
"Caption - temperature schedule": "0.7,0.8,0.9",
|
||||
"Caption - top p schedule": "0.6,0.7,0.8",
|
||||
"Caption - top k schedule": "60,70,80",
|
||||
"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",
|
||||
@@ -304,9 +298,9 @@
|
||||
"Distiller - seed": -1,
|
||||
"Distiller - max caption chars for LLM": 1536,
|
||||
"Distiller - num predict": 3096,
|
||||
"Distiller - temperature": 0.24,
|
||||
"Distiller - top p": 0.9,
|
||||
"Distiller - top k": 60,
|
||||
"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",
|
||||
@@ -325,15 +319,22 @@
|
||||
"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.12,
|
||||
"Formatter - top p": 0.88,
|
||||
"Formatter - top k": 50,
|
||||
"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": [
|
||||
"67",
|
||||
"151",
|
||||
0
|
||||
]
|
||||
},
|
||||
@@ -342,74 +343,61 @@
|
||||
"title": " JLC CaptionForge Pipeline Planner"
|
||||
}
|
||||
},
|
||||
"88": {
|
||||
"151": {
|
||||
"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": "",
|
||||
"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 - 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,
|
||||
"Input - single image": [
|
||||
"56",
|
||||
"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
|
||||
],
|
||||
"pipeline_plan": [
|
||||
"56",
|
||||
]
|
||||
},
|
||||
"class_type": "DisplayAny",
|
||||
"_meta": {
|
||||
"title": "🔧 Display Any"
|
||||
}
|
||||
},
|
||||
"157": {
|
||||
"inputs": {
|
||||
"mode": "raw value",
|
||||
"input": [
|
||||
"147",
|
||||
1
|
||||
]
|
||||
},
|
||||
"class_type": "JLC_CaptionForge",
|
||||
"class_type": "DisplayAny",
|
||||
"_meta": {
|
||||
"title": " JLC CaptionForge Orchestrator"
|
||||
"title": "🔧 Display Any"
|
||||
}
|
||||
},
|
||||
"89": {
|
||||
"158": {
|
||||
"inputs": {
|
||||
"output": "",
|
||||
"source": [
|
||||
"88",
|
||||
"mode": "raw value",
|
||||
"input": [
|
||||
"147",
|
||||
2
|
||||
]
|
||||
},
|
||||
"class_type": "Display Any (rgthree)",
|
||||
"class_type": "DisplayAny",
|
||||
"_meta": {
|
||||
"title": "Taggy Captions - Display Any (rgthree)"
|
||||
"title": "🔧 Display Any"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
+755
-828
File diff suppressed because one or more lines are too long
Binary file not shown.
|
After Width: | Height: | Size: 2.3 MiB |
@@ -1,3 +1,3 @@
|
||||
"""Single authoritative package/release version for CaptionForge."""
|
||||
|
||||
CAPTIONFORGE_VERSION = "1.0.0"
|
||||
CAPTIONFORGE_VERSION = "1.0.2"
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"_meta": {
|
||||
"name": "CaptionForge Ollama Model Dropdowns",
|
||||
"version": "1.0.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",
|
||||
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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(" ,")
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -117,6 +117,7 @@ from .captionforge_prompt_defaults import (
|
||||
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"
|
||||
@@ -391,6 +392,8 @@ def build_captionforge_pipeline_plan(
|
||||
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,
|
||||
@@ -434,6 +437,7 @@ def build_captionforge_pipeline_plan(
|
||||
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,
|
||||
@@ -502,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
|
||||
@@ -573,10 +585,14 @@ 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": {
|
||||
|
||||
@@ -145,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
|
||||
@@ -157,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,
|
||||
)
|
||||
|
||||
|
||||
@@ -384,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
|
||||
|
||||
@@ -796,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
|
||||
@@ -810,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
|
||||
|
||||
@@ -825,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)
|
||||
@@ -1343,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
|
||||
@@ -1577,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,
|
||||
|
||||
@@ -146,9 +146,11 @@ MANIFEST = {
|
||||
),
|
||||
}
|
||||
|
||||
import ctypes
|
||||
import fnmatch
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
import shutil
|
||||
@@ -156,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
|
||||
@@ -168,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 (
|
||||
@@ -334,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
|
||||
|
||||
@@ -781,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
|
||||
@@ -795,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
|
||||
|
||||
@@ -810,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)
|
||||
@@ -1074,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)
|
||||
@@ -1242,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,
|
||||
@@ -1314,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
|
||||
@@ -1337,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}")
|
||||
@@ -1480,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
|
||||
|
||||
@@ -1514,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):]
|
||||
|
||||
@@ -202,7 +202,9 @@ from ...engines.jlc_joy_caption_engine import (
|
||||
write_run_config_json,
|
||||
)
|
||||
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
|
||||
|
||||
|
||||
@@ -252,6 +254,7 @@ 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:
|
||||
@@ -262,7 +265,9 @@ def _iter_input_path_images(input_path: str, recursive: bool, filename_glob: str
|
||||
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, str, Path]] = []
|
||||
@@ -718,6 +723,9 @@ class JLC_CaptionForgeJoy:
|
||||
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),
|
||||
@@ -732,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(
|
||||
@@ -809,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)
|
||||
|
||||
|
||||
@@ -200,7 +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,
|
||||
@@ -920,6 +923,7 @@ 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:
|
||||
@@ -930,7 +934,9 @@ def _iter_input_path_images(input_path: str, recursive: bool, filename_glob: str
|
||||
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, str, Path]] = []
|
||||
@@ -1025,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()
|
||||
@@ -1068,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",
|
||||
@@ -1083,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"),
|
||||
}
|
||||
|
||||
@@ -1529,6 +1547,9 @@ 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 = [(*optional_image_identity(i), pil) for i, pil in enumerate(_tensor_to_pil(image))]
|
||||
file_images: list[tuple[str, str, Path]] = []
|
||||
@@ -1571,13 +1592,12 @@ 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, image_key: str, pil: Image.Image):
|
||||
for run in run_plan:
|
||||
t0 = time.perf_counter()
|
||||
|
||||
@@ -190,7 +190,9 @@ from ...engines.jlc_qwen_caption_engine import (
|
||||
write_run_config_json,
|
||||
)
|
||||
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,
|
||||
@@ -293,6 +295,7 @@ 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:
|
||||
@@ -303,7 +306,9 @@ def _iter_input_path_images(input_path: str, recursive: bool, filename_glob: str
|
||||
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, str, Path]] = []
|
||||
@@ -744,6 +749,9 @@ class JLC_CaptionForgeQwen:
|
||||
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(
|
||||
@@ -759,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(
|
||||
@@ -836,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)
|
||||
|
||||
|
||||
+658
-256
File diff suppressed because it is too large
Load Diff
@@ -22,7 +22,7 @@ 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 persistent semantic 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
|
||||
@@ -122,6 +122,14 @@ 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
|
||||
@@ -355,6 +363,7 @@ def _call_build_captionforge_pipeline_plan_compat(**kwargs) -> dict[str, Any]:
|
||||
"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
|
||||
|
||||
|
||||
@@ -707,12 +716,28 @@ class JLC_CaptionForge_Pipeline_Planner:
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"tooltip": (
|
||||
"Optional persistent semantic caption anchor supplied by the user, for example "
|
||||
"'doll-like quasi-3D render'. CaptionForge encourages compatible anchor content "
|
||||
"to persist through refinement and final caption generation."
|
||||
"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.",
|
||||
},
|
||||
),
|
||||
|
||||
# -----------------------------------------------------------------
|
||||
# Caption-node controls. User-facing label avoids Pass A.
|
||||
@@ -1075,6 +1100,7 @@ class JLC_CaptionForge_Pipeline_Planner:
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
**dataset_export_inputs(),
|
||||
"Input - single image": (
|
||||
"IMAGE",
|
||||
{
|
||||
@@ -1155,6 +1181,8 @@ class JLC_CaptionForge_Pipeline_Planner:
|
||||
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(
|
||||
@@ -1205,6 +1233,7 @@ class JLC_CaptionForge_Pipeline_Planner:
|
||||
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(
|
||||
@@ -1215,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)
|
||||
|
||||
|
||||
+3
-2
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "captionforge"
|
||||
version = "1.0.0"
|
||||
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"
|
||||
@@ -52,6 +52,7 @@ dependencies = [
|
||||
"torch",
|
||||
"transformers",
|
||||
"accelerate",
|
||||
"bitsandbytes>=0.46.1",
|
||||
"huggingface-hub",
|
||||
"pillow",
|
||||
"numpy",
|
||||
@@ -61,7 +62,7 @@ dependencies = [
|
||||
|
||||
[project.optional-dependencies]
|
||||
quantization = [
|
||||
"bitsandbytes",
|
||||
"bitsandbytes>=0.46.1",
|
||||
]
|
||||
|
||||
dev = [
|
||||
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -2,7 +2,9 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import importlib
|
||||
import io
|
||||
import json
|
||||
import struct
|
||||
import sys
|
||||
@@ -375,10 +377,33 @@ class PlannerContractTests(unittest.TestCase):
|
||||
|
||||
|
||||
class WorkflowAssetContractTests(unittest.TestCase):
|
||||
UI_WORKFLOW = ROOT / "assets" / "workflows" / "CaptionForge_FullWorkflow_Rel_v1.0.2.json"
|
||||
API_WORKFLOW = ROOT / "assets" / "workflows" / "CaptionForge_FullWorkflow_API_Rel_v1.0.2.json"
|
||||
PNG_WORKFLOW = ROOT / "assets" / "workflows" / "CaptionForge_FullWorkflow_Rel_v1.0.2.png"
|
||||
|
||||
def test_canonical_workflow_uses_one_optional_image_source_without_batcher(self) -> None:
|
||||
ui = json.loads(self.UI_WORKFLOW.read_text(encoding="utf-8"))
|
||||
api = json.loads(self.API_WORKFLOW.read_text(encoding="utf-8"))
|
||||
png = _embedded_png_workflow(self.PNG_WORKFLOW)
|
||||
self.assertEqual(png, ui)
|
||||
for name, artifact in (("ui", ui), ("api", api), ("png", png)):
|
||||
serialized = json.dumps(artifact)
|
||||
with self.subTest(artifact=name):
|
||||
self.assertNotIn("ImageBatchMulti", serialized)
|
||||
nodes = artifact["nodes"] if "nodes" in artifact else artifact.values()
|
||||
node_types = [node.get("type", node.get("class_type")) for node in nodes]
|
||||
self.assertEqual(node_types.count("JLC_LoadAndResizeImage"), 1)
|
||||
note = next(node for node in ui["nodes"] if node["type"] == "MarkdownNote")
|
||||
text = note["widgets_values_named"]["text"]
|
||||
self.assertIn("quick single-image workflows", text)
|
||||
self.assertIn("Input - image path", text)
|
||||
self.assertIn("does not support heterogeneous IMAGE lists", text)
|
||||
self.assertIn("future/v2 work", text)
|
||||
|
||||
def test_planner_downstream_controls_match_ui_api_and_png_workflows(self) -> None:
|
||||
ui = json.loads((ROOT / "assets" / "workflows" / "CaptionForge_FullWorkflow.json").read_text(encoding="utf-8"))
|
||||
api = json.loads((ROOT / "assets" / "workflows" / "CaptionForge_FullWorkflow_API.json").read_text(encoding="utf-8"))
|
||||
png = _embedded_png_workflow(ROOT / "assets" / "workflows" / "CaptionForge_FullWorkflow.png")
|
||||
ui = json.loads(self.UI_WORKFLOW.read_text(encoding="utf-8"))
|
||||
api = json.loads(self.API_WORKFLOW.read_text(encoding="utf-8"))
|
||||
png = _embedded_png_workflow(self.PNG_WORKFLOW)
|
||||
|
||||
prompt_defaults = {
|
||||
"Distiller - prompt": capstone.DEFAULT_FAT_DRAFT_INSTRUCTIONS,
|
||||
@@ -386,9 +411,8 @@ class WorkflowAssetContractTests(unittest.TestCase):
|
||||
"Validator - prompt": capstone.DEFAULT_VALIDATOR_INSTRUCTIONS,
|
||||
"Formatter - prompt": capstone.DEFAULT_TAGGY_FORMATTER_INSTRUCTIONS,
|
||||
}
|
||||
required_names = list(
|
||||
planner_node.JLC_CaptionForge_Pipeline_Planner.INPUT_TYPES()["required"]
|
||||
)
|
||||
input_types = planner_node.JLC_CaptionForge_Pipeline_Planner.INPUT_TYPES()
|
||||
required_names = list(input_types["required"])
|
||||
ui_planner = next(
|
||||
node for node in ui["nodes"]
|
||||
if node["type"] == "JLC_CaptionForge_Pipeline_Planner"
|
||||
@@ -397,7 +421,12 @@ class WorkflowAssetContractTests(unittest.TestCase):
|
||||
node for node in api.values()
|
||||
if node["class_type"] == "JLC_CaptionForge_Pipeline_Planner"
|
||||
)
|
||||
for name in required_names:
|
||||
widget_names = required_names + [
|
||||
name for name in input_types.get("optional", {})
|
||||
if name in ui_planner["widgets_values_named"]
|
||||
]
|
||||
self.assertEqual(list(ui_planner["widgets_values_named"]), widget_names)
|
||||
for name in widget_names:
|
||||
self.assertEqual(api_planner["inputs"][name], ui_planner["widgets_values_named"][name])
|
||||
|
||||
for workflow in (ui, png):
|
||||
@@ -405,8 +434,8 @@ class WorkflowAssetContractTests(unittest.TestCase):
|
||||
node for node in workflow["nodes"]
|
||||
if node["type"] == "JLC_CaptionForge_Pipeline_Planner"
|
||||
)
|
||||
self.assertEqual(len(node["widgets_values"]), len(required_names))
|
||||
for index, name in enumerate(required_names):
|
||||
self.assertEqual(len(node["widgets_values"]), len(widget_names))
|
||||
for index, name in enumerate(widget_names):
|
||||
self.assertEqual(node["widgets_values"][index], node["widgets_values_named"][name])
|
||||
for name, value in prompt_defaults.items():
|
||||
self.assertEqual(node["widgets_values_named"][name], value)
|
||||
@@ -415,9 +444,9 @@ class WorkflowAssetContractTests(unittest.TestCase):
|
||||
self.assertIn("TAGGY:", formatter_prompt)
|
||||
|
||||
def test_formatter_prompt_matches_ui_api_and_png_workflows(self) -> None:
|
||||
ui = json.loads((ROOT / "assets" / "workflows" / "CaptionForge_FullWorkflow.json").read_text(encoding="utf-8"))
|
||||
api = json.loads((ROOT / "assets" / "workflows" / "CaptionForge_FullWorkflow_API.json").read_text(encoding="utf-8"))
|
||||
png = _embedded_png_workflow(ROOT / "assets" / "workflows" / "CaptionForge_FullWorkflow.png")
|
||||
ui = json.loads(self.UI_WORKFLOW.read_text(encoding="utf-8"))
|
||||
api = json.loads(self.API_WORKFLOW.read_text(encoding="utf-8"))
|
||||
png = _embedded_png_workflow(self.PNG_WORKFLOW)
|
||||
|
||||
expected = capstone.DEFAULT_TAGGY_FORMATTER_INSTRUCTIONS
|
||||
api_capstone = next(
|
||||
@@ -430,10 +459,10 @@ class WorkflowAssetContractTests(unittest.TestCase):
|
||||
node = next(node for node in workflow["nodes"] if node["type"] == "JLC_CaptionForge")
|
||||
self.assertEqual(node["widgets_values_named"]["Formatter - prompt"], expected)
|
||||
|
||||
def test_canonical_nodes_match_source_defaults_and_current_widget_order(self) -> None:
|
||||
ui = json.loads((ROOT / "assets" / "workflows" / "CaptionForge_FullWorkflow.json").read_text(encoding="utf-8"))
|
||||
api = json.loads((ROOT / "assets" / "workflows" / "CaptionForge_FullWorkflow_API.json").read_text(encoding="utf-8"))
|
||||
png = _embedded_png_workflow(ROOT / "assets" / "workflows" / "CaptionForge_FullWorkflow.png")
|
||||
def test_canonical_nodes_match_stable_defaults_and_current_widget_order(self) -> None:
|
||||
ui = json.loads(self.UI_WORKFLOW.read_text(encoding="utf-8"))
|
||||
api = json.loads(self.API_WORKFLOW.read_text(encoding="utf-8"))
|
||||
png = _embedded_png_workflow(self.PNG_WORKFLOW)
|
||||
self.assertEqual(png, ui)
|
||||
|
||||
contracts = (
|
||||
@@ -443,7 +472,6 @@ class WorkflowAssetContractTests(unittest.TestCase):
|
||||
(
|
||||
"Caption - Joy runs/image",
|
||||
"Caption - Qwen runs/image",
|
||||
"Caption - Ollama runs/image",
|
||||
"Ollama - URL",
|
||||
"Ollama - keep loaded",
|
||||
"Ollama - request timeout seconds",
|
||||
@@ -453,9 +481,6 @@ class WorkflowAssetContractTests(unittest.TestCase):
|
||||
"Distiller - seed",
|
||||
"Distiller - max caption chars for LLM",
|
||||
"Distiller - num predict",
|
||||
"Distiller - temperature",
|
||||
"Distiller - top p",
|
||||
"Distiller - top k",
|
||||
"Distiller - write prompt JSONL",
|
||||
"Distiller - preserve raw response",
|
||||
"Validator - model",
|
||||
@@ -474,9 +499,6 @@ class WorkflowAssetContractTests(unittest.TestCase):
|
||||
"Formatter - prompt",
|
||||
"Formatter - seed",
|
||||
"Formatter - num predict",
|
||||
"Formatter - temperature",
|
||||
"Formatter - top p",
|
||||
"Formatter - top k",
|
||||
"Formatter - write prompt JSONL",
|
||||
"Formatter - preserve raw response",
|
||||
),
|
||||
@@ -501,6 +523,7 @@ class WorkflowAssetContractTests(unittest.TestCase):
|
||||
"Validator - system prompt",
|
||||
"Validator - prompt",
|
||||
"Validator - max new tokens",
|
||||
"Validator - max image size",
|
||||
"Validator - temperature",
|
||||
"Validator - top p",
|
||||
"Validator - top k",
|
||||
@@ -517,10 +540,15 @@ class WorkflowAssetContractTests(unittest.TestCase):
|
||||
),
|
||||
)
|
||||
for class_type, node_class, default_names in contracts:
|
||||
required = node_class.INPUT_TYPES()["required"]
|
||||
names = list(required)
|
||||
input_types = node_class.INPUT_TYPES()
|
||||
required = input_types["required"]
|
||||
ui_node = next(node for node in ui["nodes"] if node["type"] == class_type)
|
||||
api_node = next(node for node in api.values() if node["class_type"] == class_type)
|
||||
names = list(required) + [
|
||||
name for name in input_types.get("optional", {})
|
||||
if name in ui_node["widgets_values_named"]
|
||||
]
|
||||
self.assertEqual(list(ui_node["widgets_values_named"]), names)
|
||||
self.assertEqual(len(ui_node["widgets_values"]), len(names))
|
||||
for index, name in enumerate(names):
|
||||
with self.subTest(class_type=class_type, index=index, control=name):
|
||||
@@ -540,9 +568,9 @@ class WorkflowAssetContractTests(unittest.TestCase):
|
||||
)
|
||||
|
||||
def test_canonical_workflows_have_only_public_safe_paths(self) -> None:
|
||||
ui = json.loads((ROOT / "assets" / "workflows" / "CaptionForge_FullWorkflow.json").read_text(encoding="utf-8"))
|
||||
api = json.loads((ROOT / "assets" / "workflows" / "CaptionForge_FullWorkflow_API.json").read_text(encoding="utf-8"))
|
||||
png = _embedded_png_workflow(ROOT / "assets" / "workflows" / "CaptionForge_FullWorkflow.png")
|
||||
ui = json.loads(self.UI_WORKFLOW.read_text(encoding="utf-8"))
|
||||
api = json.loads(self.API_WORKFLOW.read_text(encoding="utf-8"))
|
||||
png = _embedded_png_workflow(self.PNG_WORKFLOW)
|
||||
for name, artifact in (("ui", ui), ("api", api), ("png", png)):
|
||||
serialized = json.dumps(artifact).lower()
|
||||
with self.subTest(artifact=name):
|
||||
@@ -559,10 +587,17 @@ class WorkflowAssetContractTests(unittest.TestCase):
|
||||
|
||||
|
||||
class CapstoneResolutionTests(unittest.TestCase):
|
||||
def _capture_downstream_calls(self, *, plan: dict | None = None) -> tuple[list[dict], dict]:
|
||||
def _capture_downstream_calls(
|
||||
self,
|
||||
*,
|
||||
plan: dict | None = None,
|
||||
planner_max_size: int | None = None,
|
||||
validator_widget_max_size: int | None = None,
|
||||
image_size: tuple[int, int] = (2, 2),
|
||||
) -> tuple[list[dict], dict]:
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
root = Path(temp_dir)
|
||||
Image.new("RGB", (2, 2), "white").save(root / "image.png")
|
||||
Image.new("RGB", image_size, "white").save(root / "image.png")
|
||||
caption_path = root / "captions.jsonl"
|
||||
caption_path.write_text(
|
||||
json.dumps(
|
||||
@@ -577,6 +612,15 @@ class CapstoneResolutionTests(unittest.TestCase):
|
||||
+ "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
if planner_max_size is not None:
|
||||
plan = planner_node._call_build_captionforge_pipeline_plan_compat(
|
||||
output_dir=str(root / "output"),
|
||||
input_path=temp_dir,
|
||||
run_name="validator-contract",
|
||||
max_size=planner_max_size,
|
||||
)
|
||||
plan["paths"]["caption_jsonl"] = str(caption_path)
|
||||
plan["paths"]["pass_a_jsonl"] = str(caption_path)
|
||||
kwargs = {
|
||||
"Input - captions JSONL": str(caption_path),
|
||||
"Input - image path": temp_dir,
|
||||
@@ -612,6 +656,8 @@ class CapstoneResolutionTests(unittest.TestCase):
|
||||
"Formatter - top k": 19,
|
||||
"Formatter seed": 20,
|
||||
}
|
||||
if validator_widget_max_size is not None:
|
||||
kwargs["Validator - max image size"] = validator_widget_max_size
|
||||
if plan is not None:
|
||||
kwargs["pipeline_plan"] = plan
|
||||
|
||||
@@ -743,6 +789,42 @@ class CapstoneResolutionTests(unittest.TestCase):
|
||||
self.assertFalse(downstream_call["keep_loaded"])
|
||||
self.assertEqual(downstream_call["timeout"], 222.0)
|
||||
|
||||
def test_standalone_validator_max_image_size_widget_contract(self) -> None:
|
||||
widget = capstone.JLC_CaptionForge.INPUT_TYPES()["required"]["Validator - max image size"]
|
||||
|
||||
self.assertEqual(widget[0], "INT")
|
||||
self.assertEqual(widget[1]["default"], 1024)
|
||||
self.assertEqual(widget[1]["min"], 0)
|
||||
self.assertEqual(widget[1]["max"], 8192)
|
||||
self.assertEqual(widget[1]["step"], 64)
|
||||
|
||||
def test_standalone_default_max_size_reaches_validator_image_preparation(self) -> None:
|
||||
_, call = self._capture_downstream_calls(image_size=(1200, 600))
|
||||
|
||||
image_bytes = base64.b64decode(call["image_b64"])
|
||||
with Image.open(io.BytesIO(image_bytes)) as transmitted:
|
||||
self.assertEqual(transmitted.size, (1024, 512))
|
||||
|
||||
def test_standalone_custom_max_size_reaches_validator_image_preparation(self) -> None:
|
||||
_, call = self._capture_downstream_calls(
|
||||
validator_widget_max_size=768,
|
||||
image_size=(1200, 600),
|
||||
)
|
||||
|
||||
image_bytes = base64.b64decode(call["image_b64"])
|
||||
with Image.open(io.BytesIO(image_bytes)) as transmitted:
|
||||
self.assertEqual(transmitted.size, (768, 384))
|
||||
|
||||
def test_real_planner_max_size_overrides_standalone_widget(self) -> None:
|
||||
_, call = self._capture_downstream_calls(
|
||||
planner_max_size=1,
|
||||
validator_widget_max_size=768,
|
||||
)
|
||||
|
||||
image_bytes = base64.b64decode(call["image_b64"])
|
||||
with Image.open(io.BytesIO(image_bytes)) as transmitted:
|
||||
self.assertEqual(transmitted.size, (1, 1))
|
||||
|
||||
def test_formatter_derivatives_parse_dual_and_legacy_responses(self) -> None:
|
||||
short, taggy = capstone._parse_formatter_derivatives(
|
||||
"SHORT: A concise natural caption.\n"
|
||||
@@ -962,3 +1044,4 @@ class OverwriteContractTests(unittest.TestCase):
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
@@ -2,12 +2,15 @@
|
||||
|
||||
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
|
||||
|
||||
@@ -38,6 +41,51 @@ planner_engine = importlib.import_module("CaptionForge.engines.captionforge_pipe
|
||||
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(
|
||||
@@ -274,6 +322,85 @@ class OrchestratorOutputContractTests(unittest.TestCase):
|
||||
)
|
||||
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()
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -46,8 +46,8 @@ ollama = importlib.import_module("CaptionForge.nodes.caption_nodes.jlc_captionfo
|
||||
|
||||
|
||||
class ReleaseMetadataTests(unittest.TestCase):
|
||||
def test_current_release_versions_are_1_0_0(self) -> None:
|
||||
self.assertEqual(version_module.CAPTIONFORGE_VERSION, "1.0.0")
|
||||
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)
|
||||
@@ -59,6 +59,23 @@ class ReleaseMetadataTests(unittest.TestCase):
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user