diff --git a/extensions_built_in/captioner/BaseCaptioner.py b/extensions_built_in/captioner/BaseCaptioner.py index a89a6ac..bf9a575 100644 --- a/extensions_built_in/captioner/BaseCaptioner.py +++ b/extensions_built_in/captioner/BaseCaptioner.py @@ -3,6 +3,7 @@ from collections import OrderedDict import sqlite3 import os +import re from typing import Literal, Optional import threading import time @@ -18,6 +19,9 @@ from toolkit.train_tools import get_torch_dtype AITK_Status = Literal["running", "stopped", "error", "completed"] +# prompt special token: -> contents of the media file's sibling .ext +_CAPTION_TOKEN_RE = re.compile(r"") + class CaptionConfig: def __init__(self, **kwargs): @@ -152,6 +156,23 @@ class BaseCaptioner(BaseExtensionProcess): image = image.resize((new_width, new_height), resample=Image.BICUBIC) return image + def resolve_prompt(self, prompt: Optional[str], file_path: str) -> Optional[str]: + """Expand special tokens in a caption prompt for one media file. + becomes the text of .ext next to the file (json included, inserted as-is), + or blank when that file does not exist.""" + if not prompt or " str: + def build_prompt(self, aspect_ratio: str, file_path: Optional[str] = None) -> str: # caption_prompt is the user-editable ADDITIONAL INSTRUCTIONS block, # injected into the fixed system prompt (not the whole prompt). - user_instructions = (self.caption_config.caption_prompt or "").strip() + user_instructions = self.caption_config.caption_prompt or "" + if file_path is not None: + user_instructions = self.resolve_prompt(user_instructions, file_path) + user_instructions = user_instructions.strip() if not user_instructions: user_instructions = "None." prompt = ideogram4_caption_prompt.replace("{{aspect_ratio}}", aspect_ratio) @@ -131,7 +134,7 @@ class Ideogram4Captioner(Qwen3VLCaptioner): aspect_ratio = self.compute_aspect_ratio(width, height) img = self.load_pil_image(file_path, max_res=self.caption_config.max_res) - prompt = self.build_prompt(aspect_ratio) + prompt = self.build_prompt(aspect_ratio, file_path) messages = [ { diff --git a/extensions_built_in/captioner/MossMusicCaptioner.py b/extensions_built_in/captioner/MossMusicCaptioner.py index d7559f3..a9dd148 100644 --- a/extensions_built_in/captioner/MossMusicCaptioner.py +++ b/extensions_built_in/captioner/MossMusicCaptioner.py @@ -117,7 +117,7 @@ class MossMusicCaptioner(AceStepCaptioner): item["language_inputs"] = self._make_inputs(audio_data, LANGUAGE_PROMPT) if self.caption_config.fixed_caption is None: item["caption_inputs"] = self._make_inputs( - audio_data, self.caption_config.caption_prompt + audio_data, self.resolve_prompt(self.caption_config.caption_prompt, file_path) ) return item diff --git a/extensions_built_in/captioner/Qwen25OmniCaptioner.py b/extensions_built_in/captioner/Qwen25OmniCaptioner.py index d1d218a..f8d7efc 100644 --- a/extensions_built_in/captioner/Qwen25OmniCaptioner.py +++ b/extensions_built_in/captioner/Qwen25OmniCaptioner.py @@ -171,7 +171,7 @@ class Qwen25OmniCaptioner(BaseCaptioner): "role": "user", "content": [ media, - {"type": "text", "text": self.caption_config.caption_prompt}, + {"type": "text", "text": self.resolve_prompt(self.caption_config.caption_prompt, _file_path)}, ], } ] diff --git a/extensions_built_in/captioner/Qwen38Captioner.py b/extensions_built_in/captioner/Qwen38Captioner.py index 67cb36f..7ae55c9 100644 --- a/extensions_built_in/captioner/Qwen38Captioner.py +++ b/extensions_built_in/captioner/Qwen38Captioner.py @@ -93,7 +93,7 @@ class Qwen38Captioner(BaseCaptioner): def _prep_media(self, file_path: str) -> dict: """CPU side of one file (image decode / frame extraction, processor); runs in a worker thread ahead of the GPU.""" - item = self.llm.prepare_caption_inputs(file_path, self.caption_config.caption_prompt) + item = self.llm.prepare_caption_inputs(file_path, self.resolve_prompt(self.caption_config.caption_prompt, file_path)) item["file"] = file_path return item diff --git a/extensions_built_in/captioner/Qwen3OmniCaptioner.py b/extensions_built_in/captioner/Qwen3OmniCaptioner.py index 6c72731..a9e96d9 100644 --- a/extensions_built_in/captioner/Qwen3OmniCaptioner.py +++ b/extensions_built_in/captioner/Qwen3OmniCaptioner.py @@ -678,7 +678,7 @@ class Qwen3OmniCaptioner(BaseCaptioner): "role": "user", "content": [ media, - {"type": "text", "text": self.caption_config.caption_prompt}, + {"type": "text", "text": self.resolve_prompt(self.caption_config.caption_prompt, _file_path)}, ], } ] diff --git a/extensions_built_in/captioner/Qwen3VLCaptioner.py b/extensions_built_in/captioner/Qwen3VLCaptioner.py index b4ec37e..c385160 100644 --- a/extensions_built_in/captioner/Qwen3VLCaptioner.py +++ b/extensions_built_in/captioner/Qwen3VLCaptioner.py @@ -98,7 +98,7 @@ class Qwen3VLCaptioner(BaseCaptioner): "type": "image", "image": img, }, - {"type": "text", "text": self.caption_config.caption_prompt}, + {"type": "text", "text": self.resolve_prompt(self.caption_config.caption_prompt, file_path)}, ], } ] diff --git a/ui/src/components/CaptionSimpleJob.tsx b/ui/src/components/CaptionSimpleJob.tsx index 13ee888..3805e08 100644 --- a/ui/src/components/CaptionSimpleJob.tsx +++ b/ui/src/components/CaptionSimpleJob.tsx @@ -415,6 +415,7 @@ const CaptionSimpleJob: React.FC = ({ jobConfig, setJobConfig, gpuIDs, se )} { setJobConfig(value, 'config.process[0].caption.caption_prompt'); diff --git a/ui/src/components/DocModal.tsx b/ui/src/components/DocModal.tsx index bfdd6bf..4e940db 100644 --- a/ui/src/components/DocModal.tsx +++ b/ui/src/components/DocModal.tsx @@ -19,13 +19,13 @@ export default function DocModal() { }; return ( - + -
+
), }, + 'config.process[0].caption.caption_prompt': { + title: 'Caption Prompt', + description: ( + <> + The instruction sent to the captioning model for every file. It can contain special tokens that are replaced + per file before the prompt is sent: +
+
+ <caption.ext> — inserts the text of the file next to the media file that has the same name + and the given extension. For /data/subfolder/img01.webp, <caption.json> inserts + the contents of /data/subfolder/img01.json and <caption.txt> inserts{' '} + img01.txt. The file is inserted as plain text exactly as it is on disk (json included). If the file + does not exist, the token is replaced with nothing. +
+
+ Example: Rewrite this caption as a single sentence: <caption.txt> + + ), + }, }; export const getDoc = (key: string | null | undefined): ConfigDoc | null => { diff --git a/ui/src/helpers/captionOptions.ts b/ui/src/helpers/captionOptions.ts index 1d5055d..d36c603 100644 --- a/ui/src/helpers/captionOptions.ts +++ b/ui/src/helpers/captionOptions.ts @@ -34,7 +34,7 @@ const extensionsVideo = ['mp4', 'mov', 'webm', 'mkv', 'avi']; const defaultExtensions = [...extensionsImage]; -const defaultImageCaptionPrompt = "Caption this image as if you were going to try to generate it with an image generator. Be thurough and describe everything in the image. Be decisive by stating things as they are. Do not say things like \"It appears that\" Or \"possibly\". Start out with things like \"A person on the beach\" or \"A black dragon\". No preamble. Just get to the point."; +const defaultImageCaptionPrompt = "Caption this image as if you were going to try to generate it with an image generator. Be thorough and describe everything in the image. Be decisive by stating things as they are. Do not say things like \"It appears that\" Or \"possibly\". Start out with things like \"A person on the beach\" or \"A black dragon\". No preamble. Just get to the point."; const defaultVideoCaptionPrompt = "Caption this video as if you were going to try to generate it with a video generator. Describe the visual content, how it moves and changes over time, and the camera work. Also describe the audio, including any speech, music, or sound effects, and transcribe spoken dialogue verbatim in quotes. Be decisive by stating things as they are. Do not say things like \"It appears that\" Or \"possibly\". No preamble. Just get to the point.";