Add support for ACE-Step 1.5 and ACE-Step 1.5 XL. Also added dataset captioning through the UI. (#785)

* Base ace step 1.5 xl added. Generating, still wip on training and ui

* Base training code done

* Fix some issues with caching text embeddings. Update sample cards to show audio

* Fix issue with quantizing ace step

* Add album artwork to samples with waveform.

* Cleanup logs

* Add album art endpoint to speed up album art loading

* Made an make video with artwork script

* Make ui handle basic audio models. Make multi line adjustments to the editor and better syntax hilighting.

* Add prompt tagging system for special tagged models.

* prompt tagging processing for ui working.

* Moved default samples to a special file so we can add more when needed and they can be adjusted for a specific model

* Add a captioner job with music captioner that is prepped for use with the ui

* Add basit ui setup for captioning modal and handeling captioning jobs

* Starting captioning job from ui working. Still better management for it.

* Better filtering of job options in the job view for captioning jobs

* Added qwen3 vl as a captioner for images

* Have an indicator when a dataset is being captioned.

* Adjust the way caption jobs look in the queue

* Fix a few issues. Adjust defaults.

* Version bump

* Added ace step to the readme.

extensions_built_in/audio_models/__init__.py

@@ -0,0 +1,7 @@
+from .ace_step import AceStep15Model, AceStep15XLModel
+
+AI_TOOLKIT_MODELS = [
+    # put a list of models here
+    AceStep15Model,
+    AceStep15XLModel,
+]

extensions_built_in/audio_models/ace_step/__init__.py

@@ -0,0 +1 @@
+from .ace_step_15_model import AceStep15Model, AceStep15XLModel

extensions_built_in/audio_models/ace_step/ace_step_15_model.py

@@ -0,0 +1,333 @@
+import json
+import os
+from typing import List, Optional
+import huggingface_hub
+import torch
+from safetensors.torch import load_file, save_file
+from extensions_built_in.audio_models.base_audio_model import BaseAudioModel
+from toolkit.basic import flush
+from toolkit.config_modules import GenerateImageConfig
+from toolkit.prompt_utils import PromptEmbeds, concat_prompt_embeds
+from toolkit.samplers.custom_flowmatch_sampler import (
+    CustomFlowMatchEulerDiscreteScheduler,
+)
+from toolkit.util.quantize import get_qtype, quantize, quantize_model
+
+from optimum.quanto import freeze
+from .src.model import (
+    AceStep15,
+    OobleckVAE,
+    TextEncoder,
+    get_silence_latent,
+    load_models,
+)
+from transformers import AutoTokenizer
+from .src.pipeline import AceStep15Pipeline
+
+scheduler_config = {
+    "num_train_timesteps": 1000,
+    "shift": 3.0,
+    "use_dynamic_shifting": False,
+}
+
+def to_number(str_or_number, default):
+    if isinstance(str_or_number, (int, float)):
+        return str_or_number
+    if str_or_number is None:
+        return default
+    if str_or_number == "":
+        return default
+    try:
+        return float(str_or_number)
+    except ValueError:
+        try:
+            return int(str_or_number)
+        except ValueError as e:
+            raise ValueError(f"Could not convert {str_or_number} to a number") from e
+
+
+def parse_ace_step_caption(text):
+    """Parse a tagged caption file back into a dict."""
+    import re
+
+    def tag(name):
+        m = re.search(rf"<{name}>(.*?)</{name}>", text, re.DOTALL)
+        return m.group(1).strip() if m else ""
+
+    return {
+        "caption": tag("CAPTION"),
+        "lyrics": tag("LYRICS"),
+        "bpm": to_number(tag("BPM"), 120),
+        "keyscale": tag("KEYSCALE"),
+        "timesignature": tag("TIMESIGNATURE"),
+        "duration": to_number(tag("DURATION"), 1.0),
+        "language": tag("LANGUAGE"),
+    }
+
+
+class AceStep15Model(BaseAudioModel):
+    arch = "ace_step_15"
+    sample_rate = 48000
+
+    def __init__(
+        self,
+        device,
+        model_config,
+        dtype="bf16",
+        custom_pipeline=None,
+        noise_scheduler=None,
+        **kwargs,
+    ):
+        super().__init__(
+            device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs
+        )
+        self.is_flow_matching = True
+        self.is_transformer = True
+        # self.target_lora_modules = ['AceStep15']
+        self.target_lora_modules = ["DiTModel"]
+
+    # static method to get the noise scheduler
+    @staticmethod
+    def get_train_scheduler():
+        return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config)
+
+    def load_model(self):
+        dtype = self.torch_dtype
+        device = self.device_torch
+
+        model_path = self.model_config.name_or_path
+
+        if not os.path.exists(model_path):
+            # assume it is a hf repo like org/repo/filename.safetensors
+            path_parts = model_path.split("/")
+            if len(path_parts) != 3:
+                raise ValueError(
+                    f"Model path {model_path} does not exist and is not a valid Hugging Face repo path"
+                )
+            model_path = huggingface_hub.hf_hub_download(
+                repo_id=f"{path_parts[0]}/{path_parts[1]}",
+                filename=path_parts[2],
+            )
+        # load the models from the single safetensors file
+        load_device = device
+        if self.model_config.low_vram:
+            load_device = "cpu"
+            
+        models = load_models(model_path, device=load_device, dtype=dtype)
+
+        self.model = models["model"]
+        
+        if self.model_config.quantize:
+            self.print_and_status_update("Quantizing Transformer")
+            # quantize_model(self, self.model.decoder)
+            quantize(self.model, weights=get_qtype(self.model_config.qtype))
+            freeze(self.model)
+            flush()
+        
+        if self.model_config.low_vram:
+            self.print_and_status_update("Moving transformer to CPU")
+            self.model.to("cpu")
+            
+        
+        if (
+            self.model_config.layer_offloading
+            and self.model_config.layer_offloading_transformer_percent > 0
+        ):
+            raise NotImplementedError("Layer offloading not yet implemented for AceStep15Model")
+        
+        self.text_encoder = models["text_encoder"]
+        
+        if self.model_config.quantize_te:
+            self.print_and_status_update("Quantizing Text Encoder")
+            quantize(self.text_encoder, weights=get_qtype(self.model_config.qtype_te))
+            freeze(self.text_encoder)
+            flush()
+        
+        self.vae = models["vae"]
+        
+        # move back to device
+        self.model.to(device)
+        self.text_encoder.to(device)
+        self.vae.to(device)
+        self.tokenizer = models["tokenizer"]
+        
+        self.pipeline = AceStep15Pipeline(
+            transformer=self.model,
+            vae=self.vae,
+            text_encoder=self.text_encoder,
+            tokenizer=self.tokenizer,
+            scheduler=self.get_train_scheduler(),
+        )
+
+    def get_prompt_embeds(self, prompt: str) -> PromptEmbeds:
+        if isinstance(prompt, str):
+            prompts = [prompt]
+        else:
+            prompts = prompt
+        
+        if self.text_encoder.device == torch.device("cpu"):
+            self.text_encoder.to(self.device_torch)
+        # we need the encoder from the model
+        if self.model.encoder.device == torch.device("cpu"):
+            self.model.encoder.to(self.device_torch)
+
+        # the prompt should be json as a string. Try to parse it.
+        json_prompts = []
+        for p in prompts:
+            try:
+                json_prompts.append(parse_ace_step_caption(p))
+            except json.JSONDecodeError:
+                raise ValueError(
+                    f"Prompt {p} is not a valid JSON string. Prompts must be JSON for this model"
+                )
+
+        if self.pipeline.text_encoder.device == torch.device("cpu"):
+            self.pipeline.text_encoder.to(self.device_torch)
+
+        device = self.text_encoder.device
+        dtype = self.text_encoder.dtype
+
+        batch_pe = None
+        # TODO not sure this will allow for proper batching
+
+        for json_prompt in json_prompts:
+            prompt = json_prompt.get("caption", "")
+            lyrics = json_prompt.get("lyrics", "")
+            bpm = json_prompt.get("bpm", 120)
+            key = json_prompt.get("key", "C")
+            time_sig = json_prompt.get("time_sig", "4/4")
+            duration = json_prompt.get("duration", 10)
+            duration = int(duration) if isinstance(duration, (int, float)) else 10
+            language = json_prompt.get("language", "en")
+
+            text_embeddings, text_mask, lyric_embeddings, lyric_mask = (
+                self.pipeline.get_text_embedings(
+                    prompt, lyrics, bpm, key, time_sig, duration, language
+                )
+            )
+            latent_len = int(duration * self.pipeline.LATENT_RATE)
+            # Silence as source latent [1, 64, T] -> [1, T, 64] for DiT
+            sil = get_silence_latent(latent_len, device, dtype)  # [1, 64, T]
+            src = sil.transpose(1, 2)  # [1, T, 64]
+            chunk_masks = torch.ones_like(src)
+
+            # Reference audio (silence)
+            ref = sil[:, :, :750].transpose(1, 2)  # [1, 750, 64]
+            ref_order = torch.zeros(1, device=device, dtype=torch.long)
+            enc_h, enc_m, _ = self.pipeline.transformer.prepare_condition(
+                text_embeddings,
+                text_mask,
+                lyric_embeddings,
+                lyric_mask,
+                ref,
+                ref_order,
+                src,
+                chunk_masks,
+            )
+
+            pe = PromptEmbeds(enc_h, attention_mask=enc_m)
+            if batch_pe is None:
+                batch_pe = pe
+            else:
+                batch_pe = concat_prompt_embeds(batch_pe, pe)
+        return batch_pe
+
+    def get_transformer_block_names(self) -> Optional[List[str]]:
+        return ["layers"]
+    
+    def get_generation_pipeline(self):
+        return self.pipeline
+
+    def generate_single_audio(
+        self,
+        pipeline,
+        gen_config: GenerateImageConfig,
+        conditional_embeds: PromptEmbeds,
+        unconditional_embeds: PromptEmbeds,
+        generator: torch.Generator,
+        extra: dict,
+    ):
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+        # make sure gen config is setup for audio
+        if gen_config.output_ext not in ['mp3', 'wav']:
+            gen_config.output_ext = 'mp3'
+        prompt = gen_config.prompt
+        json_prompt = parse_ace_step_caption(prompt)
+        prompt = json_prompt.get("caption", "")
+        lyrics = json_prompt.get("lyrics", "")
+        bpm = json_prompt.get("bpm", 120)
+        key = json_prompt.get("key", "C")
+        time_sig = json_prompt.get("time_sig", "4/4")
+        duration = json_prompt.get("duration", 0)
+        language = json_prompt.get("language", "en")
+
+        output = self.pipeline(
+            prompt=None,  # we are passing in the embeds directly, so no need for a prompt
+            encoder_embeddings=conditional_embeds.text_embeds.to(self.device_torch, dtype=self.torch_dtype),
+            encoder_mask=conditional_embeds.attention_mask.to(self.device_torch, dtype=torch.bool),
+            num_inference_steps=gen_config.num_inference_steps,
+            duration=duration,
+            generator=generator,
+            bpm=bpm,
+            key=key,
+            time_sig=time_sig,
+            language=language,
+            guidance_scale=gen_config.guidance_scale,
+        )
+        return output
+
+    def get_noise_prediction(
+        self,
+        latent_model_input: torch.Tensor, #(1, 300, 64)
+        timestep: torch.Tensor,  # 0 to 1000 scale
+        text_embeddings: PromptEmbeds,
+        **kwargs,
+    ):
+        if self.model.decoder.device == torch.device("cpu"):
+            self.model.decoder.to(self.device_torch)
+        with torch.no_grad():
+            model: AceStep15 = self.model
+            tt = timestep.to(self.device_torch, dtype=torch.long) / 1000
+            latent_len = latent_model_input.shape[1]
+            device = self.device_torch
+            dtype = self.torch_dtype
+            attn = torch.ones(1, latent_len, device=device, dtype=dtype)
+
+            # build context from silence latent matching the actual input length
+            sil = get_silence_latent(latent_len, device, dtype)  # [1, 64, T]
+            src = sil.transpose(1, 2)  # [1, T, 64]
+            chunk_masks = torch.ones_like(src)
+            context = torch.cat([src, chunk_masks], dim=-1)  # [1, T, 128]
+
+        pred = model.decoder(
+            x=latent_model_input.detach(),
+            timestep=tt.detach(),
+            timestep_r=tt.detach(),
+            attention_mask=attn.detach(),
+            enc_h=text_embeddings.text_embeds.to(self.device_torch, dtype=self.torch_dtype).detach(),
+            enc_m=text_embeddings.attention_mask.to(self.device_torch, dtype=torch.bool).detach(),
+            context=context.detach(),
+        )
+        return pred
+    
+    def get_loss_target(self, *args, **kwargs):
+        noise = kwargs.get("noise")
+        batch = kwargs.get("batch")
+        return (noise - batch.latents).detach()
+    
+    def encode_audio(self, audio_tensor: torch.Tensor, device=None, dtype=None):
+        if device is None:
+            device = self.device_torch
+        if dtype is None:
+            dtype = self.torch_dtype
+        if self.vae.device == torch.device("cpu"):
+            self.vae.to(device)
+        output = self.vae.encode(audio_tensor.to(device=device, dtype=dtype))
+        # transpose from [B, 64, T] to [B, T, 64] for DiT
+        output = output.transpose(1, 2)
+        return output
+
+
+class AceStep15XLModel(AceStep15Model):
+    arch = "ace_step_15_xl"

extensions_built_in/audio_models/ace_step/src/pipeline.py

@@ -0,0 +1,163 @@
+from typing import List, Optional
+
+import torch
+import time
+import os
+from .model import (
+    SAMPLE_RATE,
+    AceStep15,
+    OobleckVAE,
+    TextEncoder,
+    get_silence_latent,
+    compute_timesteps,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from transformers import AutoTokenizer
+
+SFT_PROMPT = """# Instruction
+{instruction}
+
+# Caption
+{caption}
+
+# Metas
+{metas}<|endoftext|>
+"""
+
+
+class AceStep15Pipeline:
+    SAMPLE_RATE = 48000
+    LATENT_RATE = 25  # 48000 / 1920
+    SFT_PROMPT = SFT_PROMPT
+
+    def __init__(self, transformer, vae, text_encoder, tokenizer, scheduler):
+        self.transformer: AceStep15 = transformer
+        self.vae: OobleckVAE = vae
+        self.text_encoder: TextEncoder = text_encoder
+        self.tokenizer: AutoTokenizer = tokenizer
+        self.scheduler = scheduler
+
+    def to(self, *args, **kwargs):
+        self.transformer.to(*args, **kwargs)
+        self.vae.to(*args, **kwargs)
+        self.text_encoder.to(*args, **kwargs)
+
+    def get_text_embedings(
+        self, prompt, lyrics, bpm, key, time_sig, duration, language
+    ):
+        metas = f"- bpm: {bpm}\n- timesignature: {time_sig}\n- keyscale: {key}\n- duration: {int(duration)} seconds\n"
+        caption = self.SFT_PROMPT.format(
+            instruction="Fill the audio semantic mask based on the given conditions:",
+            caption=prompt,
+            metas=metas,
+        )
+        lyrics_text = f"# Languages\n{language}\n\n# Lyric\n{lyrics}<|endoftext|>"
+
+        cap_tok = self.tokenizer(
+            caption, truncation=True, max_length=256, return_tensors="pt"
+        )
+        lyr_tok = self.tokenizer(
+            lyrics_text, truncation=True, max_length=2048, return_tensors="pt"
+        )
+
+        text_embeddings = self.text_encoder.encode_text(
+            cap_tok.input_ids.to(self.text_encoder.device)
+        ).to(self.transformer.dtype)
+        text_mask = cap_tok.attention_mask.to(self.text_encoder.device).bool()
+        lyric_embeddings = self.text_encoder.encode_lyrics(
+            lyr_tok.input_ids.to(self.text_encoder.device)
+        ).to(self.transformer.dtype)
+        lyric_mask = lyr_tok.attention_mask.to(self.text_encoder.device).bool()
+
+        return text_embeddings, text_mask, lyric_embeddings, lyric_mask
+
+    def __call__(
+        self,
+        prompt="",
+        lyrics="",
+        encoder_embeddings: Optional[List[torch.Tensor]] = None,
+        encoder_mask: Optional[List[torch.Tensor]] = None,
+        # uses a null conditional for unconditional if not provided, which is what we want for CFG
+        num_inference_steps=50,
+        duration=30.0,
+        generator: torch.Generator = None,
+        bpm="N/A",
+        key="N/A",
+        time_sig="N/A",
+        language="en",
+        guidance_scale=1.0,
+    ):
+        t_sched = compute_timesteps(num_inference_steps, 3.0)
+        latent_len = int(duration * self.LATENT_RATE)
+        device = self.transformer.device
+        dtype = self.transformer.dtype
+
+        # Text encoding
+        if encoder_embeddings is not None and encoder_mask is not None:
+            enc_h = encoder_embeddings
+            enc_m = encoder_mask
+            sil = get_silence_latent(latent_len, device, dtype)  # [1, 64, T]
+            src = sil.transpose(1, 2)  # [1, T, 64]
+            chunk_masks = torch.ones_like(src)
+            ctx = torch.cat([src, chunk_masks.to(src.dtype)], dim=-1)
+        else:
+            text_h, text_m, lyric_h, lyric_m = self.get_text_embedings(
+                prompt, lyrics, bpm, key, time_sig, duration, language
+            )
+
+            # Silence as source latent [1, 64, T] -> [1, T, 64] for DiT
+            sil = get_silence_latent(latent_len, device, dtype)  # [1, 64, T]
+            src = sil.transpose(1, 2)  # [1, T, 64]
+            chunk_masks = torch.ones_like(src)
+
+            # Reference audio (silence)
+            ref = sil[:, :, :750].transpose(1, 2)  # [1, 750, 64]
+            ref_order = torch.zeros(1, device=device, dtype=torch.long)
+
+            # Prepare conditions (conditional)
+            enc_h, enc_m, ctx = self.transformer.prepare_condition(
+                text_h, text_m, lyric_h, lyric_m, ref, ref_order, src, chunk_masks
+            )
+
+        # Prepare unconditional conditions for CFG
+        use_cfg = guidance_scale > 1.0
+        enc_h_uncond = None
+        if use_cfg:
+            enc_h_uncond = self.transformer.null_condition_emb.expand_as(enc_h)
+
+        # Noise
+        if generator is None:
+            generator = torch.Generator(device=device)
+        noise_ch = ctx.shape[-1] // 2
+        xt = randn_tensor(
+            (1, latent_len, noise_ch), generator=generator, device=device, dtype=dtype
+        )
+        # xt = torch.randn(1, latent_len, noise_ch, generator=generator, device=device, dtype=dtype)
+
+        # Diffusion
+        t_sched_t = torch.tensor(t_sched, device=device, dtype=dtype)
+        attn = torch.ones(1, latent_len, device=device, dtype=dtype)
+
+        for i in range(len(t_sched_t)):
+            tv = t_sched_t[i].item()
+            tt = torch.full((1,), tv, device=device, dtype=dtype)
+
+            vt_cond = self.transformer.decoder(xt, tt, tt, attn, enc_h, enc_m, ctx)
+
+            if use_cfg:
+                vt_uncond = self.transformer.decoder(
+                    xt, tt, tt, attn, enc_h_uncond, enc_m, ctx
+                )
+                vt = vt_uncond + guidance_scale * (vt_cond - vt_uncond)
+            else:
+                vt = vt_cond
+
+            if i == len(t_sched_t) - 1:
+                xt = xt - vt * tv
+            else:
+                xt = xt - vt * (tv - t_sched_t[i + 1].item())
+
+        # VAE decode
+        wav = self.vae.decode(xt.transpose(1, 2))  # [1, 2, samples]
+        wav = wav[0, :, : int(duration * SAMPLE_RATE)]
+        return wav

extensions_built_in/audio_models/base_audio_model.py

@@ -0,0 +1,99 @@
+import json
+
+import torch
+
+from toolkit.config_modules import GenerateImageConfig, ModelConfig
+from toolkit.models.base_model import BaseModel
+from toolkit.prompt_utils import PromptEmbeds
+
+
+class BaseAudioModel(BaseModel):
+    sample_rate = 48000
+
+    def __init__(
+        self,
+        device,
+        model_config: ModelConfig,
+        dtype="bf16",
+        custom_pipeline=None,
+        noise_scheduler=None,
+        **kwargs,
+    ):
+        super().__init__(
+            device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs
+        )
+        self.is_audio_model = True
+
+    def generate_single_image(
+        self,
+        pipeline,
+        gen_config: GenerateImageConfig,
+        conditional_embeds: PromptEmbeds,
+        unconditional_embeds: PromptEmbeds,
+        generator: torch.Generator,
+        extra: dict,
+    ):
+        # This is called on the base model. We override it to make it make more sense for audio models.
+        return self.generate_single_audio(
+            pipeline,
+            gen_config,
+            conditional_embeds,
+            unconditional_embeds,
+            generator,
+            extra,
+        )
+
+    def generate_single_audio(
+        self,
+        pipeline,
+        gen_config: GenerateImageConfig,
+        conditional_embeds: PromptEmbeds,
+        unconditional_embeds: PromptEmbeds,
+        generator: torch.Generator,
+        extra: dict,
+    ):
+        # This is called on the base model. We override it to make it make more sense for audio models.
+        raise NotImplementedError(
+            "generate_single_audio is not implemented for this model"
+        )
+
+    def get_model_has_grad(self):
+        return False
+
+    def get_te_has_grad(self):
+        return False
+
+    def save_model(self, output_path, meta, save_dtype):
+        # we need to save the model, vae, text encoder, and tokenizer together since they are all trained together and depend on each other
+        raise NotImplementedError(
+            "save_model is not implemented for this model. Use the pipeline directly instead."
+        )
+
+    def convert_lora_weights_before_save(self, state_dict):
+        # currently starte with transformer. but needs to start with diffusion_model. for comfyui
+        new_sd = {}
+        for key, value in state_dict.items():
+            new_key = key.replace("transformer.", "diffusion_model.")
+            new_sd[new_key] = value
+        return new_sd
+
+    def convert_lora_weights_before_load(self, state_dict):
+        # saved as diffusion_model. but needs to be transformer. for ai-toolkit
+        new_sd = {}
+        for key, value in state_dict.items():
+            new_key = key.replace("diffusion_model.", "transformer.")
+            new_sd[new_key] = value
+        return new_sd
+
+    def encode_images(self, image_list: torch.Tensor, device=None, dtype=None):
+        # make it more obvious for audio models
+        return self.encode_audio(image_list, device=device, dtype=dtype)
+
+    def encode_audio(self, audio_tensor: torch.Tensor, device=None, dtype=None):
+        if device is None:
+            device = self.device_torch
+        if dtype is None:
+            dtype = self.torch_dtype
+        if self.vae.device == torch.device("cpu"):
+            self.vae.to(device)
+        return self.vae.encode(audio_tensor.to(device=device, dtype=dtype))

extensions_built_in/captioner/AceStepCaptioner.py

@@ -0,0 +1,243 @@
+import librosa
+import numpy as np
+import torch
+import torchaudio
+from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor
+from collections import OrderedDict
+
+from optimum.quanto import freeze
+from toolkit.basic import flush
+from toolkit.util.quantize import quantize, get_qtype
+
+from .BaseCaptioner import BaseCaptioner
+import transformers
+import logging
+import warnings
+
+transformers.logging.set_verbosity_error()
+warnings.filterwarnings("ignore")
+logging.disable(logging.WARNING)
+
+TARGET_SAMPLE_RATE = 16000
+CAPTIONER_ID = "ACE-Step/acestep-captioner"
+TRANSCRIBER_ID = "ACE-Step/acestep-transcriber"
+
+# Key profiles for Krumhansl-Schmuckler key detection
+MAJOR_PROFILE = np.array(
+    [6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88]
+)
+MINOR_PROFILE = np.array(
+    [6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17]
+)
+KEY_NAMES = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Audio analysis (BPM, key, time signature) via librosa
+# ═══════════════════════════════════════════════════════════════════════════════
+
+
+def analyze_audio(audio_path):
+    """Extract BPM, key, and time signature from audio using librosa."""
+    y, sr = librosa.load(audio_path, sr=22050, mono=True)
+    duration = librosa.get_duration(y=y, sr=sr)
+
+    # BPM
+    tempo, _ = librosa.beat.beat_track(y=y, sr=sr)
+    if hasattr(tempo, "__len__"):
+        tempo = tempo[0]
+    bpm = int(round(float(tempo)))
+
+    # Key detection via chroma correlation with key profiles
+    chroma = librosa.feature.chroma_cqt(y=y, sr=sr)
+    chroma_avg = chroma.mean(axis=1)
+    major_corrs = np.array(
+        [np.corrcoef(np.roll(MAJOR_PROFILE, i), chroma_avg)[0, 1] for i in range(12)]
+    )
+    minor_corrs = np.array(
+        [np.corrcoef(np.roll(MINOR_PROFILE, i), chroma_avg)[0, 1] for i in range(12)]
+    )
+
+    best_major_idx = major_corrs.argmax()
+    best_minor_idx = minor_corrs.argmax()
+    if major_corrs[best_major_idx] >= minor_corrs[best_minor_idx]:
+        keyscale = f"{KEY_NAMES[best_major_idx]} major"
+    else:
+        keyscale = f"{KEY_NAMES[best_minor_idx]} minor"
+
+    # Time signature estimation from beat strength pattern
+    onset_env = librosa.onset.onset_strength(y=y, sr=sr)
+    tempo_est, beats = librosa.beat.beat_track(onset_envelope=onset_env, sr=sr)
+    if len(beats) >= 8:
+        beat_strengths = onset_env[beats]
+        # Check 3/4 vs 4/4 by looking at periodicity of strong beats
+        acf = np.correlate(
+            beat_strengths - beat_strengths.mean(),
+            beat_strengths - beat_strengths.mean(),
+            mode="full",
+        )
+        acf = acf[len(acf) // 2 :]
+        if len(acf) > 6:
+            # Look at autocorrelation peaks at lag 3 vs lag 4
+            score_3 = acf[3] if len(acf) > 3 else 0
+            score_4 = acf[4] if len(acf) > 4 else 0
+            timesig = "3" if score_3 > score_4 * 1.2 else "4"
+        else:
+            timesig = "4"
+    else:
+        timesig = "4"
+
+    return {
+        "bpm": bpm,
+        "keyscale": keyscale,
+        "timesignature": timesig,
+        "duration": int(round(duration)),
+    }
+
+
+class AceStepCaptioner(BaseCaptioner):
+    def __init__(self, process_id: int, job, config: OrderedDict, **kwargs):
+        super(AceStepCaptioner, self).__init__(process_id, job, config, **kwargs)
+
+    def load_model(self):
+        self.print_and_status_update("Loading transcriber model")
+        self.model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
+            self.caption_config.model_name_or_path,
+            dtype=self.torch_dtype,
+            device_map="cpu",
+        )
+        self.model.to(self.device_torch)
+        self.model.disable_talker()
+        if self.caption_config.quantize:
+            self.print_and_status_update("Quantizing transcriber model")
+            quantize(self.model, weights=get_qtype(self.caption_config.qtype))
+            freeze(self.model)
+            flush()
+        self.processor = Qwen2_5OmniProcessor.from_pretrained(
+            self.caption_config.model_name_or_path
+        )
+        if self.caption_config.low_vram:
+            self.model.to("cpu")
+
+        # 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")
+        flush()
+
+    def run_qwen_audio(self, model, processor, audio_data, sr, prompt_text):
+        """Run a Qwen2.5-Omni model on audio with a text prompt."""
+        conversation = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "audio", "audio": "<|audio_bos|><|AUDIO|><|audio_eos|>"},
+                    {"type": "text", "text": prompt_text},
+                ],
+            }
+        ]
+        text = processor.apply_chat_template(
+            conversation, add_generation_prompt=True, tokenize=False
+        )
+        inputs = processor(
+            text=text,
+            audio=[audio_data],
+            images=None,
+            videos=None,
+            return_tensors="pt",
+            padding=True,
+            sampling_rate=sr,
+        )
+        inputs = inputs.to(model.device).to(model.dtype)
+        text_ids = model.generate(**inputs, return_audio=False)
+        output = processor.batch_decode(
+            text_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )
+        result = output[0]
+        marker = "assistant\n"
+        if marker in result:
+            result = result[result.rfind(marker) + len(marker) :]
+        return result.strip()
+
+    def get_audio_lyrics(self, audio_data: torch.Tensor) -> str:
+        if self.caption_config.low_vram and self.model2.device != torch.device("cpu"):
+            # move captioner to cpu
+            self.model2.to("cpu")
+        # move lyric model if needed
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+
+        prompt_text = "*Task* Transcribe this audio in detail"
+        return self.run_qwen_audio(
+            self.model, self.processor, audio_data, TARGET_SAMPLE_RATE, prompt_text
+        )
+
+    def get_audio_caption(self, audio_data: torch.Tensor) -> str:
+        if self.caption_config.low_vram and self.model.device != torch.device("cpu"):
+            # move lyricmodel to cpu
+            self.model.to("cpu")
+        # move captioner model if needed
+        if self.model2.device == torch.device("cpu"):
+            self.model2.to(self.device_torch)
+        prompt_text = "*Task* Describe this music in detail. Include genre, mood, instrumentation, tempo feel, and vocal style if present."
+        return self.run_qwen_audio(
+            self.model2, self.processor2, audio_data, TARGET_SAMPLE_RATE, prompt_text
+        )
+
+    def get_caption_for_file(self, file_path: str) -> str:
+        try:
+            # analyze audio with librosa
+            analysis = analyze_audio(file_path)
+
+            # load audio with torchaudio for transcription
+            waveform, sr = torchaudio.load(file_path)
+            waveform = waveform.to(self.device_torch)
+            if waveform.shape[0] > 1:
+                waveform = waveform.mean(dim=0, keepdim=True)
+            if sr != TARGET_SAMPLE_RATE:
+                waveform = torchaudio.functional.resample(
+                    waveform, sr, TARGET_SAMPLE_RATE
+                )
+            audio_data = waveform.squeeze(0).cpu().numpy()
+
+            # get the lyrics from the audio
+            lyrics = self.get_audio_lyrics(audio_data)
+
+            language = "en"
+
+            if "# Languages" in lyrics and "# Lyrics" in lyrics:
+                language = lyrics.split("# Languages")[1].split("# Lyrics")[0]
+                # remove newlines and extra spaces from language
+                language = language.replace("\n", "").strip()
+                lyrics = lyrics.split("# Lyrics")[1].strip()
+
+            # get the caption from the audio
+            caption = self.get_audio_caption(audio_data)
+
+            output = f"<CAPTION>\n{caption}\n</CAPTION>\n"
+            output += f"<LYRICS>\n{lyrics}\n</LYRICS>\n"
+            output += f"<BPM>{analysis['bpm']}</BPM>\n"
+            output += f"<KEYSCALE>{analysis['keyscale']}</KEYSCALE>\n"
+            output += f"<TIMESIGNATURE>{analysis['timesignature']}</TIMESIGNATURE>\n"
+            output += f"<DURATION>{analysis['duration']}</DURATION>\n"
+            output += f"<LANGUAGE>{language}</LANGUAGE>"
+            return output
+        except Exception as e:
+            print(f"Error processing {file_path}: {e}")
+            return None

extensions_built_in/captioner/BaseCaptioner.py

@@ -0,0 +1,399 @@
+import asyncio
+from collections import OrderedDict
+
+import sqlite3
+import os
+from typing import Literal, Optional
+import threading
+import time
+import signal
+import concurrent.futures
+from PIL import Image
+
+import torch
+from jobs.process import BaseExtensionProcess
+import tqdm
+
+from toolkit.train_tools import get_torch_dtype
+
+AITK_Status = Literal["running", "stopped", "error", "completed"]
+
+
+class CaptionConfig:
+    def __init__(self, **kwargs):
+        self.model_name_or_path = kwargs.get("model_name_or_path", None)
+        if self.model_name_or_path is None:
+            raise ValueError("model_name_or_path is required in config")
+        self.model_name_or_path2 = kwargs.get("model_name_or_path2", None)
+        self.extensions = kwargs.get("extensions", [])
+        if self.extensions is None or len(self.extensions) == 0:
+            raise ValueError("At least one extension is required in config")
+        self.path_to_caption = kwargs.get("path_to_caption", None)
+        if self.path_to_caption is None:
+            raise ValueError("path_to_caption is required in config")
+        self.dtype = kwargs.get("dtype", "bf16")
+        self.device = kwargs.get("device", "cuda")
+        self.quantize = kwargs.get("quantize", False)
+        self.qtype = kwargs.get("qtype", "float8")
+        self.low_vram = kwargs.get("low_vram", False)
+        self.caption_extension = kwargs.get("caption_extension", "txt")
+        self.recaption = kwargs.get("recaption", False)
+        self.max_res = kwargs.get("max_res", 512)
+        self.max_new_tokens = kwargs.get("max_new_tokens", 128)
+        self.caption_prompt = kwargs.get(
+            "caption_prompt", "Describe this image in detail."
+        )
+
+
+class BaseCaptioner(BaseExtensionProcess):
+    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")
+        self.job_id = os.environ.get("AITK_JOB_ID", None)
+        self.job_id = self.job_id.strip() if self.job_id is not None else None
+        self.is_ui_captioner = True
+        if not os.path.exists(self.sqlite_db_path):
+            self.is_ui_captioner = False
+        else:
+            print(f"Using SQLite database at {self.sqlite_db_path}")
+        if self.job_id is None:
+            self.is_ui_captioner = False
+        else:
+            print(f'Job ID: "{self.job_id}"')
+
+        self.is_stopping = False
+
+        if self.is_ui_captioner:
+            self.is_stopping = False
+            # Create a thread pool for database operations
+            self.thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+            # Track all async tasks
+            self._async_tasks = []
+            # Initialize the status
+            self._run_async_operation(self._update_status("running", "Starting"))
+            self._stop_watcher_started = False
+            # self.start_stop_watcher(interval_sec=2.0)
+
+        self.caption_config = CaptionConfig(**self.get_conf("caption", {}))
+        self.model = None
+        self.processor = None
+        self.model2 = None
+        self.processor2 = None
+        self.file_paths = []
+        self.device_torch = torch.device(self.caption_config.device)
+        self.torch_dtype = get_torch_dtype(self.caption_config.dtype)
+
+    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("")
+
+        print("****************************************************")
+        print("Captioning complete")
+        print("****************************************************")
+
+    def run_caption_loop(self):
+        for file_path in tqdm.tqdm(
+            self.file_paths, desc="Captioning files", unit="file"
+        ):
+            if self.is_ui_captioner:
+                self.maybe_stop()
+                if self.is_stopping:
+                    break
+            try:
+                file_caption = self.get_caption_for_file(file_path)
+                if file_caption is not None:
+                    self.save_caption_for_file(file_path, file_caption)
+            except Exception as e:
+                print(f"Error captioning file {file_path}: {e}")
+                continue
+
+    def load_pil_image(self, file_path: str, max_res: Optional[int] = None) -> Image:
+        image = Image.open(file_path).convert("RGB")
+        if max_res is not None:
+            max_pixels = max_res * max_res
+            image_pixels = image.width * image.height
+            if image_pixels > max_pixels:
+                scale_factor = (max_pixels / image_pixels) ** 0.5
+                new_width = int(image.width * scale_factor)
+                new_height = int(image.height * scale_factor)
+                image = image.resize((new_width, new_height), resample=Image.BICUBIC)
+        return image
+
+    def save_caption_for_file(self, file_path: str, caption: str):
+        filename_no_ext = os.path.splitext(file_path)[0]
+        caption_file_path = f"{filename_no_ext}.{self.caption_config.caption_extension}"
+        # delete it if it already exists
+        if os.path.exists(caption_file_path):
+            os.remove(caption_file_path)
+        with open(caption_file_path, "w") as f:
+            f.write(caption)
+
+    def get_caption_for_file(self, file_path: str) -> str:
+        raise NotImplementedError("Captioning not implemented for this captioner")
+
+    def print_and_status_update(self, status: str):
+        print(status)
+        self.update_status("running", status)
+
+    def find_files(self):
+        # recursivly find all the files in the path_to_caption with the specified extensions and save the paths to self.file_paths
+        for root, dirs, files in os.walk(self.caption_config.path_to_caption):
+            for file in files:
+                if any(
+                    file.lower().endswith(f".{ext}")
+                    for ext in self.caption_config.extensions
+                ):
+                    full_path = os.path.join(root, file)
+                    self.file_paths.append(full_path)
+        # sort
+        self.file_paths.sort()
+        # it not recaption, remove the ones with captions
+        if not self.caption_config.recaption:
+            filtered_file_paths = []
+            for file_path in self.file_paths:
+                filename_no_ext = os.path.splitext(file_path)[0]
+                caption_file_path = (
+                    f"{filename_no_ext}.{self.caption_config.caption_extension}"
+                )
+                if not os.path.exists(caption_file_path):
+                    filtered_file_paths.append(file_path)
+            print(
+                f"Found {len(self.file_paths)} files. {len(filtered_file_paths)} need captioning."
+            )
+            self.file_paths = filtered_file_paths
+        else:
+            print(f"Found {len(self.file_paths)} files to caption")
+
+    def load_model(self):
+        raise NotImplementedError("Model loading not implemented for this captioner")
+
+    def start_stop_watcher(self, interval_sec: float = 5.0):
+        """
+        Start a daemon thread that periodically checks should_stop()
+        and terminates the process immediately when triggered.
+        """
+        if not self.is_ui_captioner:
+            return
+        if getattr(self, "_stop_watcher_started", False):
+            return
+        self._stop_watcher_started = True
+        t = threading.Thread(
+            target=self._stop_watcher_thread, args=(interval_sec,), daemon=True
+        )
+        t.start()
+
+    def _stop_watcher_thread(self, interval_sec: float):
+        while True:
+            try:
+                if self.should_stop():
+                    # Mark and update status (non-blocking; uses existing infra)
+                    self.is_stopping = True
+                    self._run_async_operation(
+                        self._update_status("stopped", "Job stopped (remote)")
+                    )
+                    # Best-effort flush pending async ops
+                    try:
+                        asyncio.run(self.wait_for_all_async())
+                    except RuntimeError:
+                        pass
+                    # Try to stop DB thread pool quickly
+                    try:
+                        self.thread_pool.shutdown(wait=False, cancel_futures=True)
+                    except TypeError:
+                        self.thread_pool.shutdown(wait=False)
+                    print("")
+                    print("****************************************************")
+                    print("    Stop signal received; terminating process.      ")
+                    print("****************************************************")
+                    os.kill(os.getpid(), signal.SIGINT)
+                time.sleep(interval_sec)
+            except Exception:
+                time.sleep(interval_sec)
+
+    def _run_async_operation(self, coro):
+        """Helper method to run an async coroutine and track the task."""
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            # No event loop exists, create a new one
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        # Create a task and track it
+        if loop.is_running():
+            task = asyncio.run_coroutine_threadsafe(coro, loop)
+            self._async_tasks.append(asyncio.wrap_future(task))
+        else:
+            task = loop.create_task(coro)
+            self._async_tasks.append(task)
+            loop.run_until_complete(task)
+
+    async def _execute_db_operation(self, operation_func):
+        """Execute a database operation in a separate thread with retry on lock."""
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(
+            self.thread_pool, lambda: self._retry_db_operation(operation_func)
+        )
+
+    def _db_connect(self):
+        """Create a new connection for each operation to avoid locking."""
+        conn = sqlite3.connect(self.sqlite_db_path, timeout=30.0)
+        conn.isolation_level = None  # Enable autocommit mode
+        return conn
+
+    def _retry_db_operation(self, operation_func, max_retries=3, base_delay=2.0):
+        """Retry a database operation with exponential backoff on lock errors."""
+        last_error = None
+        for attempt in range(max_retries + 1):
+            try:
+                return operation_func()
+            except sqlite3.OperationalError as e:
+                if "database is locked" in str(e):
+                    last_error = e
+                    if attempt < max_retries:
+                        delay = base_delay * (2**attempt)  # 2s, 4s, 8s
+                        print(
+                            f"[AITK] Database locked (attempt {attempt + 1}/{max_retries + 1}), retrying in {delay:.1f}s..."
+                        )
+                        time.sleep(delay)
+                    else:
+                        print(
+                            f"[AITK] Database locked after {max_retries + 1} attempts, giving up."
+                        )
+                else:
+                    raise
+        raise last_error
+
+    def should_stop(self):
+        if not self.is_ui_captioner:
+            return False
+
+        def _check_stop():
+            with self._db_connect() as conn:
+                cursor = conn.cursor()
+                cursor.execute("SELECT stop FROM Job WHERE id = ?", (self.job_id,))
+                stop = cursor.fetchone()
+                return False if stop is None else stop[0] == 1
+
+        return self._retry_db_operation(_check_stop)
+
+    def should_return_to_queue(self):
+        if not self.is_ui_captioner:
+            return False
+
+        def _check_return_to_queue():
+            with self._db_connect() as conn:
+                cursor = conn.cursor()
+                cursor.execute(
+                    "SELECT return_to_queue FROM Job WHERE id = ?", (self.job_id,)
+                )
+                return_to_queue = cursor.fetchone()
+                return False if return_to_queue is None else return_to_queue[0] == 1
+
+        return self._retry_db_operation(_check_return_to_queue)
+
+    def maybe_stop(self):
+        if not self.is_ui_captioner:
+            return
+        if self.should_stop():
+            self._run_async_operation(self._update_status("stopped", "Job stopped"))
+            self.is_stopping = True
+            raise Exception("Job stopped")
+        if self.should_return_to_queue():
+            self._run_async_operation(self._update_status("queued", "Job queued"))
+            self.is_stopping = True
+            raise Exception("Job returning to queue")
+
+    async def _update_key(self, key, value):
+        def _do_update():
+            with self._db_connect() as conn:
+                cursor = conn.cursor()
+                cursor.execute("BEGIN IMMEDIATE")
+                try:
+                    # Convert the value to string if it's not already
+                    if isinstance(value, str):
+                        value_to_insert = value
+                    else:
+                        value_to_insert = str(value)
+
+                    # Use parameterized query for both the column name and value
+                    update_query = f"UPDATE Job SET {key} = ? WHERE id = ?"
+                    cursor.execute(update_query, (value_to_insert, self.job_id))
+                finally:
+                    cursor.execute("COMMIT")
+
+        await self._execute_db_operation(_do_update)
+
+    def update_step(self):
+        """Non-blocking update of the step count."""
+        if self.is_ui_captioner:
+            self._run_async_operation(self._update_key("step", self.step_num))
+
+    def update_db_key(self, key, value):
+        """Non-blocking update a key in the database."""
+        if self.is_ui_captioner:
+            self._run_async_operation(self._update_key(key, value))
+
+    async def _update_status(self, status: AITK_Status, info: Optional[str] = None):
+        if not self.is_ui_captioner:
+            return
+
+        def _do_update():
+            with self._db_connect() as conn:
+                cursor = conn.cursor()
+                cursor.execute("BEGIN IMMEDIATE")
+                try:
+                    if info is not None:
+                        cursor.execute(
+                            "UPDATE Job SET status = ?, info = ? WHERE id = ?",
+                            (status, info, self.job_id),
+                        )
+                    else:
+                        cursor.execute(
+                            "UPDATE Job SET status = ? WHERE id = ?",
+                            (status, self.job_id),
+                        )
+                finally:
+                    cursor.execute("COMMIT")
+
+        await self._execute_db_operation(_do_update)
+
+    def update_status(self, status: AITK_Status, info: Optional[str] = None):
+        if self.is_ui_captioner:
+            """Non-blocking update of status."""
+            self._run_async_operation(self._update_status(status, info))
+
+    def on_error(self, e: Exception):
+        super(BaseCaptioner, self).on_error(e)
+        if self.is_ui_captioner:
+            try:
+                if not self.is_stopping:
+                    self.update_status("error", str(e))
+                asyncio.run(self.wait_for_all_async())
+            except Exception as db_err:
+                print(
+                    f"[AITK] Warning: failed to update DB during error handling: {db_err}"
+                )
+            finally:
+                self.thread_pool.shutdown(wait=True)
+
+    async def wait_for_all_async(self):
+        """Wait for all tracked async operations to complete."""
+        if not self._async_tasks:
+            return
+
+        try:
+            await asyncio.gather(*self._async_tasks)
+        except Exception as e:
+            pass
+        finally:
+            # Clear the task list after completion
+            self._async_tasks.clear()

extensions_built_in/captioner/Qwen3VLCaptioner.py

@@ -0,0 +1,86 @@
+from transformers import Qwen3VLForConditionalGeneration, AutoProcessor
+from collections import OrderedDict
+
+from optimum.quanto import freeze
+from toolkit.basic import flush
+from toolkit.util.quantize import quantize, get_qtype
+
+from .BaseCaptioner import BaseCaptioner
+import transformers
+import logging
+import warnings
+
+transformers.logging.set_verbosity_error()
+warnings.filterwarnings("ignore")
+logging.disable(logging.WARNING)
+
+
+class Qwen3VLCaptioner(BaseCaptioner):
+    def __init__(self, process_id: int, job, config: OrderedDict, **kwargs):
+        super(Qwen3VLCaptioner, self).__init__(process_id, job, config, **kwargs)
+
+    def load_model(self):
+        self.print_and_status_update("Loading Qwen3VL model")
+        self.model = Qwen3VLForConditionalGeneration.from_pretrained(
+            self.caption_config.model_name_or_path,
+            dtype=self.torch_dtype,
+            device_map="cpu",
+        )
+        if not self.caption_config.low_vram:
+            self.model.to(self.device_torch)
+        if self.caption_config.quantize:
+            self.print_and_status_update("Quantizing Qwen3VL model")
+            quantize(self.model, weights=get_qtype(self.caption_config.qtype))
+            freeze(self.model)
+            flush()
+        self.processor = AutoProcessor.from_pretrained(
+            self.caption_config.model_name_or_path
+        )
+        if self.caption_config.low_vram:
+            self.model.to(self.device_torch)
+        flush()
+
+    def get_caption_for_file(self, file_path: str) -> str:
+        img = self.load_pil_image(file_path, max_res=self.caption_config.max_res)
+        try:
+            messages = [
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image",
+                            "image": img,
+                        },
+                        {"type": "text", "text": self.caption_config.caption_prompt},
+                    ],
+                }
+            ]
+
+            # Preparation for inference
+            inputs = self.processor.apply_chat_template(
+                messages,
+                tokenize=True,
+                add_generation_prompt=True,
+                return_dict=True,
+                return_tensors="pt",
+            )
+            inputs = inputs.to(self.device_torch)
+
+            # Inference: Generation of the output
+            generated_ids = self.model.generate(
+                **inputs, max_new_tokens=self.caption_config.max_new_tokens
+            )
+            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,
+            )
+
+            return output_text[0].strip()
+        except Exception as e:
+            print(f"Error processing {file_path}: {e}")
+            return None

extensions_built_in/captioner/__init__.py

@@ -0,0 +1,31 @@
+from toolkit.extension import Extension
+
+
+class AceStepCaptionerExtension(Extension):
+    uid = "AceStepCaptioner"
+    name = "Ace Step Captioner"
+
+    @classmethod
+    def get_process(cls):
+        # import your process class here so it is only loaded when needed and return it
+        from .AceStepCaptioner import AceStepCaptioner
+
+        return AceStepCaptioner
+
+
+class Qwen3VLCaptionerExtension(Extension):
+    uid = "Qwen3VLCaptioner"
+    name = "Qwen 3VL Captioner"
+
+    @classmethod
+    def get_process(cls):
+        # import your process class here so it is only loaded when needed and return it
+        from .Qwen3VLCaptioner import Qwen3VLCaptioner
+
+        return Qwen3VLCaptioner
+
+
+AI_TOOLKIT_EXTENSIONS = [
+    AceStepCaptionerExtension,
+    Qwen3VLCaptionerExtension,
+]