github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Let omni captioner handle just images as well. Set it as the new default captioner.

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-13 10:30:17 -06:00
parent 6b7fb60a22
commit ab18528fdb
3 changed files with 97 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files

174
extensions_built_in/captioner/Qwen3OmniCaptioner.py
View File

@@ -27,6 +27,9 @@ logging.disable(logging.WARNING)
# frame sampling rate for video captioning
VIDEO_FPS = 2
# still-image files caption through the image pipeline (no audio, no frames)
IMAGE_EXTENSIONS = {"jpg", "jpeg", "png", "bmp", "webp"}
# fixed generation ceiling under compiled decode: a constant max_length keeps
# the static kv cache (and so the compiled decode graph) at one shape for
# every video; the real per-caption budget is enforced by a stopping criterion
@@ -501,7 +504,9 @@ class Qwen3OmniCaptioner(BaseCaptioner):
self._model_info = info_for_filename(os.path.basename(name_or_path))
return name_or_path
model_info = CONVROT_MODELS.get(name_or_path, CONVROT_MODELS[DEFAULT_CONVROT_MODEL])
model_info = CONVROT_MODELS.get(
name_or_path, CONVROT_MODELS[DEFAULT_CONVROT_MODEL]
)
filename = model_info["filename"]
if os.path.isdir(name_or_path):
@@ -651,12 +656,20 @@ class Qwen3OmniCaptioner(BaseCaptioner):
self.processor = AutoProcessor.from_pretrained(self._model_info["base_repo"])
flush()
@staticmethod
def _is_image_file(file_path: str) -> bool:
return os.path.splitext(file_path)[1].lower().lstrip(".") in IMAGE_EXTENSIONS
def _build_messages(self, _file_path: str):
if self._is_image_file(_file_path):
media = {"type": "image", "image": _file_path}
else:
media = {"type": "video", "video": _file_path}
return [
{
"role": "user",
"content": [
{"type": "video", "video": _file_path},
media,
{"type": "text", "text": self.caption_config.caption_prompt},
],
}
@@ -672,39 +685,60 @@ class Qwen3OmniCaptioner(BaseCaptioner):
}
def _prep_media(self, file_path: str):
"""CPU side of one video, safe to run in a worker thread: decode +
subsample frames, extract the audio track, render the chat text. At
batch size 1 the full processor (tokenize, resize, mel) runs here too,
so the main thread only moves tensors and generates."""
from transformers.video_utils import load_video
from transformers.audio_utils import load_audio
"""CPU side of one file, safe to run in a worker thread: decode +
subsample frames (or load the image), extract the audio track, render
the chat text. At batch size 1 the full processor (tokenize, resize,
mel) runs here too, so the main thread only moves tensors and
generates."""
if self._is_image_file(file_path):
from PIL import Image
frames = load_video(file_path, fps=VIDEO_FPS)
if isinstance(frames, tuple):
frames = frames[0]
audio = None
try:
a = load_audio(file_path, sampling_rate=16000)
if a is not None and a.size > 0:
audio = a
except Exception:
pass
image = Image.open(file_path).convert("RGB")
item = {"file": file_path, "kind": "image", "image": image, "audio": None}
else:
from transformers.video_utils import load_video
from transformers.audio_utils import load_audio
frames = load_video(file_path, fps=VIDEO_FPS)
if isinstance(frames, tuple):
frames = frames[0]
audio = None
try:
a = load_audio(file_path, sampling_rate=16000)
if a is not None and a.size > 0:
audio = a
except Exception:
pass
item = {
"file": file_path,
"kind": "video_audio" if audio is not None else "video_silent",
"frames": frames,
"audio": audio,
}
template_kwargs = {}
if self.is_thinking_model and not self.thinking_enabled:
template_kwargs["enable_thinking"] = False
text = self.processor.apply_chat_template(
item["text"] = self.processor.apply_chat_template(
self._build_messages(file_path),
tokenize=False,
add_generation_prompt=True,
**template_kwargs,
)
item = {"file": file_path, "frames": frames, "audio": audio, "text": text}
if self.caption_config.batch_size <= 1:
item["inputs"] = self._process_items([item])
return item
def _process_items(self, items):
use_audio = items[0]["audio"] is not None
kind = items[0]["kind"]
if kind == "image":
return self.processor(
text=[it["text"] for it in items],
images=[it["image"] for it in items],
return_tensors="pt",
padding=True,
size=self._size_kwargs(),
)
use_audio = kind == "video_audio"
return self.processor(
text=[it["text"] for it in items],
audio=[it["audio"] for it in items] if use_audio else None,
@@ -718,9 +752,9 @@ class Qwen3OmniCaptioner(BaseCaptioner):
)
def _caption_batch(self, items):
"""Batched generate over preprocessed items (all with audio, or all
silent). Returns captions in item order."""
use_audio = items[0]["audio"] is not None
"""Batched generate over preprocessed items (all the same kind: image,
video with audio, or silent video). Returns captions in item order."""
use_audio = items[0]["kind"] == "video_audio"
if len(items) == 1 and "inputs" in items[0]:
inputs = items[0]["inputs"]
else:
@@ -834,7 +868,8 @@ class Qwen3OmniCaptioner(BaseCaptioner):
break
futures.append((path, executor.submit(self._prep_media, path)))
with_audio, silent = [], []
# batches must be homogeneous: the processor call differs per kind
buckets = {"image": [], "video_audio": [], "video_silent": []}
while futures:
if self.is_ui_captioner:
self.maybe_stop()
@@ -850,12 +885,12 @@ class Qwen3OmniCaptioner(BaseCaptioner):
print(f"Error preprocessing {path}: {e}")
finish(path, None)
continue
bucket = with_audio if item["audio"] is not None else silent
bucket = buckets[item["kind"]]
bucket.append(item)
if len(bucket) >= batch_size:
flush(bucket)
flush(with_audio)
flush(silent)
for bucket in buckets.values():
flush(bucket)
finally:
executor.shutdown(wait=False, cancel_futures=True)
pbar.close()
@@ -891,87 +926,10 @@ class Qwen3OmniCaptioner(BaseCaptioner):
)
def get_caption_for_file(self, file_path: str) -> str:
# single-file path (and the per-file fallback when a batch fails):
# same prep + generate flow as the batched loop, for one item
try:
messages = [
{
"role": "user",
"content": [
{
"type": "video",
"video": file_path,
},
{"type": "text", "text": self.caption_config.caption_prompt},
],
}
]
max_pixels = self.caption_config.max_res * self.caption_config.max_res
# render the chat text only; the media goes to the processor
# directly so the audio track is interleaved INTO the video block
# (use_audio_in_video) instead of forming a separate audio segment
template_kwargs = {}
if self.is_thinking_model and not self.thinking_enabled:
template_kwargs["enable_thinking"] = False
text = self.processor.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
**template_kwargs,
)
# pull the audio track out of the video file; silent videos fall
# back to frames only
from transformers.audio_utils import load_audio
use_audio = True
try:
audio = load_audio(file_path, sampling_rate=16000)
if audio.size == 0:
use_audio = False
except Exception as audio_err:
print(
f"No audio track for {file_path} ({audio_err}); captioning frames only"
)
use_audio = False
inputs = self.processor(
text=text,
audio=[audio] if use_audio else None,
videos=[file_path],
return_tensors="pt",
padding=True,
use_audio_in_video=use_audio,
fps=VIDEO_FPS,
do_sample_frames=True,
# shortest_edge/longest_edge are total pixel counts
# (min_pixels/max_pixels), not edge lengths
size={
"shortest_edge": min(131072, max_pixels),
"longest_edge": max_pixels,
},
)
inputs = inputs.to(self.device_torch).to(self.torch_dtype)
self.model._pad_mask_2d = inputs.get("attention_mask", None)
generated_ids = self.model.generate(
**inputs,
use_audio_in_video=use_audio,
**self._gen_kwargs(inputs["input_ids"].shape[1]),
)
generated_ids_trimmed = [
out_ids[len(in_ids) :]
for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
]
output_text = self.processor.batch_decode(
generated_ids_trimmed,
skip_special_tokens=True,
clean_up_tokenization_spaces=False,
)
caption = output_text[0]
if "</think>" in caption:
caption = caption.split("</think>")[-1]
return caption.strip()
return self._caption_batch([self._prep_media(file_path)])[0]
except Exception as e:
print(f"Error processing {file_path}: {e}")
traceback.print_exc()

16
ui/src/helpers/captionJobConfig.ts
View File

@@ -8,21 +8,25 @@ export const defaultCaptionJobConfig: CaptionJobConfig = {
name: 'Caption Directory',
process: [
{
type: 'AceStepCaptioner',
type: 'Qwen3OmniCaptioner',
sqlite_db_path: './aitk_db.db',
device: 'cuda',
caption: {
model_name_or_path: "ACE-Step/acestep-transcriber",
model_name_or_path2: "ACE-Step/acestep-captioner",
model_name_or_path: 'ai-toolkit/Qwen3-Omni-30B-A3B-Thinking',
dtype: 'bf16',
quantize: true,
qtype: 'convrot8',
low_vram: true,
extensions: ['mp3', 'wav', 'flac', 'ogg'],
low_vram: false,
extensions: ['mp4', 'mov', 'webm', 'mkv', 'avi', 'jpg', 'jpeg', 'png', 'bmp', 'webp'],
path_to_caption: '',
recaption: false,
compile: false,
compile: true,
caption_extension: 'txt',
caption_prompt:
"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.",
max_res: 512,
max_new_tokens: 512,
batch_size: 1,
},
},
],

25
ui/src/helpers/captionOptions.ts
View File

@@ -1,6 +1,6 @@
import { GroupedSelectOption, SelectOption } from "@/types";
type CaptionGroup = 'image' | 'music' | 'video';
type CaptionGroup = 'image' | 'music' | 'video' | 'image/video/sound';
type AdditionalSections = 'caption.model_name_or_path2' | 'caption.caption_prompt' | 'caption.max_res' | 'caption.max_new_tokens' | 'caption.fixed_caption' | 'caption.thinking' | 'caption.batch_size' | 'caption.layer_offloading';
export interface CaptionOption {
@@ -51,6 +51,22 @@ non_diegetic_music: 1-3 sentences describing background music the characters can
Describe only what is actually seen and heard. Be decisive. No preamble and no extra text - output only the three fields.`;
// The T2VA format applied to a still image: one static [Shot 1], no timeline,
// no audio fields beyond the required N/A placeholders.
const minimaxImageCaptionPrompt = `Caption this image as a MiniMax training prompt for a single still frame. Output exactly three fields in this order, each starting on its own line with these exact field names:
integrated_multimodal_description: [Shot 1] ...
overall_soundscape: N/A
non_diegetic_music: N/A
Rules for integrated_multimodal_description: Write a single [Shot 1] with no timestamp, no cuts, and no camera motion - the camera holds a static shot on a still frame. State the overall visual style first (Live-action, cinematic, 2D-animated, 3D CG, claymation, watercolor, or vintage film) and the framing (extreme wide shot, wide shot, medium shot, medium close-up, close-up, or extreme close-up, plus the viewpoint). Then describe everything visible decisively and specifically: subject appearance, pose, and position; clothing, colors, and materials; the scene, lighting, and key props; and the spatial relationships between them. There is no motion, no dialogue, and no sound - do not invent any. Put any legible text in the image (signs, labels, packaging) in double quotation marks verbatim without translating it.
overall_soundscape and non_diegetic_music are always exactly N/A for a still image.
Describe only what is actually visible. Be decisive. No preamble and no extra text - output only the three fields.`;
// Editable ADDITIONAL INSTRUCTIONS block injected into the Ideogram system prompt.
// Users can tweak this for dataset-specific guidance without altering the fixed
// output contract, element/background rules, or bbox format.
@@ -108,10 +124,10 @@ export const captionerTypes: CaptionOption[] = [
{
name: 'Qwen3OmniCaptioner',
label: 'Qwen3-Omni',
group: 'video',
group: 'image/video/sound',
defaults: {
'config.process[0].caption.model_name_or_path': ['ai-toolkit/Qwen3-Omni-30B-A3B-Instruct', defaultNameOrPath],
'config.process[0].caption.extensions': [extensionsVideo, defaultExtensions],
'config.process[0].caption.model_name_or_path': ['ai-toolkit/Qwen3-Omni-30B-A3B-Thinking', defaultNameOrPath],
'config.process[0].caption.extensions': [[...extensionsVideo, ...extensionsImage], defaultExtensions],
'config.process[0].caption.caption_prompt': [defaultVideoCaptionPrompt, undefined],
'config.process[0].caption.max_res': [512, undefined],
'config.process[0].caption.max_new_tokens': [512, undefined],
@@ -126,6 +142,7 @@ export const captionerTypes: CaptionOption[] = [
captionPrompts: {
'General': defaultVideoCaptionPrompt,
'MiniMax H4 T2V': minimaxT2VCaptionPrompt,
'MiniMax H4 Image': minimaxImageCaptionPrompt,
},
additionalSections: [
'caption.caption_prompt',