RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Performance improvements to captioner. Add ability to set a default caption for ace captioner to avoid needing to caption the audio so we only transcribe

This commit is contained in:
Jaret Burkett 2026-04-28 10:15:11 -06:00
parent acc6a36214
commit f972b750e6
8 changed files with 86 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

60
extensions_built_in/captioner/AceStepCaptioner.py
View File

@ -1,3 +1,5 @@
from typing import Optional
import librosa
import numpy as np
import torch
@ -9,7 +11,7 @@ from optimum.quanto import freeze
from toolkit.basic import flush
from toolkit.util.quantize import quantize, get_qtype
from .BaseCaptioner import BaseCaptioner
from .BaseCaptioner import BaseCaptioner, CaptionConfig
import transformers
import logging
import warnings
@ -95,7 +97,16 @@ def analyze_audio(audio_path):
}
class AceStepCaptionConfig(CaptionConfig):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.fixed_caption: Optional[str] = kwargs.get("fixed_caption", None)
class AceStepCaptioner(BaseCaptioner):
caption_config_class = AceStepCaptionConfig
caption_config: AceStepCaptionConfig
def __init__(self, process_id: int, job, config: OrderedDict, **kwargs):
super(AceStepCaptioner, self).__init__(process_id, job, config, **kwargs)
@ -118,27 +129,31 @@ class AceStepCaptioner(BaseCaptioner):
)
if self.caption_config.low_vram:
self.model.to("cpu")
self.model2 = None
self.processor2 = None
# load captioner model
self.print_and_status_update("Loading captioner model")
self.model2 = Qwen2_5OmniForConditionalGeneration.from_pretrained(
self.caption_config.model_name_or_path2,
dtype=self.torch_dtype,
device_map="cpu",
)
self.model2.to(self.device_torch)
self.model2.disable_talker()
if self.caption_config.quantize:
self.print_and_status_update("Quantizing captioner model")
quantize(self.model2, weights=get_qtype(self.caption_config.qtype))
freeze(self.model2)
flush()
self.processor2 = Qwen2_5OmniProcessor.from_pretrained(
self.caption_config.model_name_or_path2,
)
if self.caption_config.fixed_caption is not None:
# load captioner model
self.print_and_status_update("Loading captioner model")
self.model2 = Qwen2_5OmniForConditionalGeneration.from_pretrained(
self.caption_config.model_name_or_path2,
dtype=self.torch_dtype,
device_map="cpu",
)
self.model2.to(self.device_torch)
self.model2.disable_talker()
if self.caption_config.quantize:
self.print_and_status_update("Quantizing captioner model")
quantize(self.model2, weights=get_qtype(self.caption_config.qtype))
freeze(self.model2)
flush()
self.processor2 = Qwen2_5OmniProcessor.from_pretrained(
self.caption_config.model_name_or_path2,
)
if self.caption_config.low_vram:
self.model2.to("cpu")
if self.caption_config.low_vram:
self.model2.to("cpu")
flush()
def run_qwen_audio(self, model, processor, audio_data, sr, prompt_text):
@ -228,7 +243,10 @@ class AceStepCaptioner(BaseCaptioner):
lyrics = lyrics.split("# Lyrics")[1].strip()
# get the caption from the audio
caption = self.get_audio_caption(audio_data)
if self.caption_config.fixed_caption is not None:
caption = self.caption_config.fixed_caption
else:
caption = self.get_audio_caption(audio_data)
output = f"<CAPTION>\n{caption}\n</CAPTION>\n"
output += f"<LYRICS>\n{lyrics}\n</LYRICS>\n"

29
extensions_built_in/captioner/BaseCaptioner.py
View File

@ -46,6 +46,8 @@ class CaptionConfig:
class BaseCaptioner(BaseExtensionProcess):
caption_config_class = CaptionConfig
def __init__(self, process_id: int, job, config: OrderedDict, **kwargs):
super(BaseCaptioner, self).__init__(process_id, job, config, **kwargs)
self.sqlite_db_path = self.config.get("sqlite_db_path", "./aitk_db.db")
@ -74,7 +76,7 @@ class BaseCaptioner(BaseExtensionProcess):
self._stop_watcher_started = False
# self.start_stop_watcher(interval_sec=2.0)
self.caption_config = CaptionConfig(**self.get_conf("caption", {}))
self.caption_config = self.caption_config_class(**self.get_conf("caption", {}))
self.model = None
self.processor = None
self.model2 = None
@ -85,19 +87,20 @@ class BaseCaptioner(BaseExtensionProcess):
def run(self):
super(BaseCaptioner, self).run()
self.start_stop_watcher()
self.update_status("running", "Loading Model")
self.load_model()
self.update_status("running", "Looking for files")
self.find_files()
self.update_status("running", f"Captioning {len(self.file_paths)} files")
self.run_caption_loop()
self.update_status("completed", "Captioning completed")
print("")
with torch.no_grad():
self.start_stop_watcher()
self.update_status("running", "Loading Model")
self.load_model()
self.update_status("running", "Looking for files")
self.find_files()
self.update_status("running", f"Captioning {len(self.file_paths)} files")
self.run_caption_loop()
self.update_status("completed", "Captioning completed")
print("")
print("****************************************************")
print("Captioning complete")
print("****************************************************")
print("****************************************************")
print("Captioning complete")
print("****************************************************")
def run_caption_loop(self):
for file_path in tqdm.tqdm(

6
ui/src/app/globals.css
View File

@ -46,7 +46,7 @@ body {
@layer components {
/* control */
.aitk-react-select-container .aitk-react-select__control {
@apply flex w-full h-8 min-h-0 px-0 text-sm bg-gray-800 border border-gray-700 rounded-sm hover:border-gray-600 items-center;
@apply flex w-full h-8 min-h-0 px-0 text-sm bg-gray-950 dark:bg-gray-800 border border-gray-700 rounded-sm hover:border-gray-600 items-center;
}
/* selected label */
@ -66,12 +66,12 @@ body {
/* menu */
.aitk-react-select-container .aitk-react-select__menu {
@apply bg-gray-800 border border-gray-700;
@apply bg-gray-950 dark:bg-gray-800 border border-gray-700;
}
/* options */
.aitk-react-select-container .aitk-react-select__option {
@apply text-sm text-gray-200 bg-gray-800 hover:bg-gray-700;
@apply text-sm text-gray-200 bg-gray-950 dark:bg-gray-800 hover:bg-gray-700;
}
/* indicator separator */

26
ui/src/components/CaptionDatasetModal.tsx
View File

@ -111,6 +111,8 @@ export const CaptionDatasetModal: React.FC = () => {
});
};
const additionalSections = selectedCaptionOption?.additionalSections || [];
return (
<Modal isOpen={open} onClose={handleClose} title="Caption Dataset" size="lg">
<div className="space-y-4 text-gray-200">
@ -159,7 +161,7 @@ export const CaptionDatasetModal: React.FC = () => {
required
/>
</div>
{selectedCaptionOption?.additionalSections?.includes('caption.model_name_or_path2') && (
{additionalSections.includes('caption.model_name_or_path2') && (
<div className="mt-4">
<CreatableSelectInput
label="Name or Path 2"
@ -175,6 +177,22 @@ export const CaptionDatasetModal: React.FC = () => {
/>
</div>
)}
{additionalSections.includes('caption.fixed_caption') && (
<div className="mt-4">
<TextInput
label="Fixed Caption"
value={jobConfig.config.process[0].caption.fixed_caption || ''}
onChange={value => {
if (value?.trim() === '') {
//@ts-ignore
value = undefined;
}
setJobConfig(value, 'config.process[0].caption.fixed_caption');
}}
placeholder="Enter fixed caption (if you want the same caption for all audio files)"
/>
</div>
)}
<div className="grid grid-cols-1 md:grid-cols-2 gap-4 mt-4">
<div>
<SelectInput
@ -191,7 +209,7 @@ export const CaptionDatasetModal: React.FC = () => {
}}
options={quantizationOptions}
/>
{selectedCaptionOption?.additionalSections?.includes('caption.max_res') && (
{additionalSections.includes('caption.max_res') && (
<div className="mt-4">
<SelectInput
label="Max Resolution"
@ -206,7 +224,7 @@ export const CaptionDatasetModal: React.FC = () => {
/>
</div>
)}
{selectedCaptionOption?.additionalSections?.includes('caption.max_new_tokens') && (
{additionalSections.includes('caption.max_new_tokens') && (
<div className="mt-4">
<SelectInput
label="Max New Tokens"
@ -237,7 +255,7 @@ export const CaptionDatasetModal: React.FC = () => {
</FormGroup>
</div>
</div>
{selectedCaptionOption?.additionalSections?.includes('caption.caption_prompt') && (
{additionalSections.includes('caption.caption_prompt') && (
<div className="mt-4">
<TextAreaInput
label="Caption Prompt"

2
ui/src/components/Modal.tsx
View File

@ -68,7 +68,7 @@ export const Modal: React.FC<ModalProps> = ({
>
{/* Modal panel */}
<div
className={`relative mx-auto w-full ${sizeClasses[size]} rounded-lg bg-gray-800 border border-gray-700 shadow-xl transition-all`}
className={`relative mx-auto w-full ${sizeClasses[size]} rounded-lg bg-white dark:bg-gray-900 border border-gray-700 shadow-xl transition-all`}
onClick={e => e.stopPropagation()}
>
{/* Modal header */}

4
ui/src/components/formInputs.tsx
View File

@ -12,7 +12,7 @@ const Select = dynamic(() => import('react-select'), { ssr: false });
const labelClasses = 'block text-xs mb-1 mt-2 text-gray-300';
const inputClasses =
'w-full text-sm px-3 py-1 bg-gray-800 border border-gray-700 rounded-sm focus:ring-2 focus:ring-gray-600 focus:border-transparent';
'w-full text-sm px-3 py-1 bg-gray-950 dark:bg-gray-800 border border-gray-700 rounded-sm text-gray-100 placeholder:text-gray-500 focus:ring-2 focus:ring-gray-600 focus:border-transparent';
export interface InputProps {
label?: string;
@ -584,7 +584,7 @@ export const SliderInput: React.FC<SliderInputProps> = props => {
</div>
{showValue && (
<div className="min-w-[3.5rem] text-right text-sm px-3 py-1 bg-gray-800 border border-gray-700 rounded-sm">
<div className="min-w-[3.5rem] text-right text-sm px-3 py-1 bg-gray-950 dark:bg-gray-800 border border-gray-700 rounded-sm">
{Number.isFinite(value) ? value : ''}
</div>
)}

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

@ -1,7 +1,7 @@
import { GroupedSelectOption, SelectOption } from "@/types";
type CaptionGroup = 'image' | 'music';
type AdditionalSections = 'caption.model_name_or_path2' | 'caption.caption_prompt' | 'caption.max_res' | 'caption.max_new_tokens';
type AdditionalSections = 'caption.model_name_or_path2' | 'caption.caption_prompt' | 'caption.max_res' | 'caption.max_new_tokens' | 'caption.fixed_caption';
export interface CaptionOption {
name: string;
@ -41,6 +41,7 @@ export const captionerTypes: CaptionOption[] = [
],
additionalSections: [
'caption.model_name_or_path2',
'caption.fixed_caption',
],
},
{

1
ui/src/types.ts
View File

@ -269,6 +269,7 @@ export interface CaptionProcessConfig {
caption_prompt?: string;
max_res?: number;
max_new_tokens?: number;
fixed_caption?: string;
}
}