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.

README.md

@@ -43,6 +43,10 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t
 - [Lightricks/LTX-2](https://huggingface.co/Lightricks/LTX-2) (LTX-2)
 - [Lightricks/LTX-2.3](https://huggingface.co/Lightricks/LTX-2.3) (LTX-2.3)
 
+### Audio
+- [ACE-Step/Ace-Step1.5](https://huggingface.co/ACE-Step/Ace-Step1.5) (Ace Step 1.5)
+- [ACE-Step/acestep-v15-xl-base](https://huggingface.co/ACE-Step/acestep-v15-xl-base) (Ace Step 1.5 XL)
+
 ### Experimental
 - [lodestones/Zeta-Chroma](https://huggingface.co/lodestones/Zeta-Chroma) (Zeta Chroma)
 

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,
+]

jobs/process/BaseSDTrainProcess.py

@@ -1188,7 +1188,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
                     if self.sd.is_flux or 'flex' in self.sd.arch:
                         # flux is a patch size of 1, but latents are divided by 2, so we need to double it
                         patch_size = 2
-                    elif hasattr(self.sd.unet.config, 'patch_size'):
+                    elif hasattr(self.sd.unet, 'config') and hasattr(self.sd.unet.config, 'patch_size'):
                         patch_size = self.sd.unet.config.patch_size
                     
                     self.sd.noise_scheduler.set_train_timesteps(

requirements_base.txt

@@ -38,3 +38,5 @@ matplotlib==3.10.1
 setuptools==69.5.1
 av==16.0.1
 torchcodec==0.9.1
+librosa==0.11.0
+mutagen==1.47.0

scripts/caption_audio_dataset.py

@@ -0,0 +1,309 @@
+#!/usr/bin/env python3
+"""
+Caption audio files for ACE-Step v1.5 training.
+
+Produces .txt files containing all training metadata:
+  - caption (from acestep-captioner)
+  - lyrics (from acestep-transcriber)
+  - bpm, keyscale, timesignature (from librosa)
+  - duration, language
+
+Requirements:
+    pip install torch torchaudio transformers librosa numpy
+
+Usage:
+    python caption_dir.py input_dir/
+    python caption_dir.py input_dir/ --low_vram --skip_existing
+"""
+
+import argparse
+import gc
+import os
+import glob
+import logging
+import warnings
+
+import librosa
+import numpy as np
+import torch
+import torchaudio
+from tqdm import tqdm
+from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor
+
+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"]
+
+
+def get_audio_files(input_dir):
+    extensions = ["*.wav", "*.mp3", "*.flac", "*.ogg", "*.WAV", "*.MP3", "*.FLAC"]
+    files = []
+    for ext in extensions:
+        files.extend(glob.glob(os.path.join(input_dir, ext)))
+    return sorted(set(files))
+
+
+def load_audio_mono_16k(audio_path):
+    waveform, sr = torchaudio.load(audio_path)
+    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)
+    return waveform.squeeze(0).numpy(), TARGET_SAMPLE_RATE
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# 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)),
+    }
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Model management
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def offload_to_cpu(model):
+    """Move model to CPU and free GPU memory."""
+    if model is not None:
+        model.to("cpu")
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+def load_qwen_model(model_id, device="cuda", dtype=torch.bfloat16):
+    """Load a Qwen2.5-Omni model."""
+    model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
+        model_id, torch_dtype=dtype, device_map=device,
+    )
+    model.disable_talker()
+    processor = Qwen2_5OmniProcessor.from_pretrained(model_id)
+    return model, processor
+
+
+def run_qwen_audio(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()
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Output formatting
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def format_output(caption, lyrics, analysis, language="en"):
+    """Format all metadata into tagged format for easy parsing."""
+    return (
+        f"<CAPTION>\n{caption}\n</CAPTION>\n"
+        f"<LYRICS>\n{lyrics}\n</LYRICS>\n"
+        f"<BPM>{analysis['bpm']}</BPM>\n"
+        f"<KEYSCALE>{analysis['keyscale']}</KEYSCALE>\n"
+        f"<TIMESIGNATURE>{analysis['timesignature']}</TIMESIGNATURE>\n"
+        f"<DURATION>{analysis['duration']}</DURATION>\n"
+        f"<LANGUAGE>{language}</LANGUAGE>"
+    )
+
+
+def parse_caption_file(path):
+    """Parse a tagged caption file back into a dict."""
+    import re
+    text = open(path, "r", encoding="utf-8").read()
+    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": tag("BPM"),
+        "keyscale": tag("KEYSCALE"),
+        "timesignature": tag("TIMESIGNATURE"),
+        "duration": tag("DURATION"),
+        "language": tag("LANGUAGE"),
+    }
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Main
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def main():
+    parser = argparse.ArgumentParser(description="Caption audio files for ACE-Step training")
+    parser.add_argument("input_dir", type=str, help="Directory containing audio files")
+    parser.add_argument("--skip_existing", action="store_true", help="Skip files that already have captions")
+    parser.add_argument("--low_vram", action="store_true", help="Offload models to CPU between stages")
+    parser.add_argument("--language", default="en", help="Default language code (default: en)")
+    args = parser.parse_args()
+
+    if not os.path.isdir(args.input_dir):
+        print(f"Error: {args.input_dir} is not a valid directory")
+        return
+
+    audio_files = get_audio_files(args.input_dir)
+    if not audio_files:
+        print("No audio files found in the directory")
+        return
+
+    print(f"Found {len(audio_files)} audio files")
+
+    # ── Stage 1: Audio analysis (BPM, key, time sig) — no GPU needed ─────
+    print("\n[Stage 1/3] Analyzing audio (BPM, key, time signature)...")
+    analyses = {}
+    for audio_path in tqdm(audio_files, desc="Analyzing"):
+        base_name = os.path.splitext(audio_path)[0]
+        if args.skip_existing and os.path.exists(base_name + ".txt"):
+            continue
+        try:
+            analyses[audio_path] = analyze_audio(audio_path)
+        except Exception as e:
+            print(f"\n  Error analyzing {os.path.basename(audio_path)}: {e}")
+            analyses[audio_path] = {"bpm": 120, "keyscale": "C major", "timesignature": "4",
+                                    "duration": 30}
+
+    # Filter to only files that need processing
+    files_to_process = [f for f in audio_files if f in analyses]
+    if not files_to_process:
+        print("All files already captioned (use without --skip_existing to overwrite)")
+        return
+
+    # ── Stage 2: Captioning ──────────────────────────────────────────────
+    print(f"\n[Stage 2/3] Captioning {len(files_to_process)} files...")
+    print("  Loading captioner model...")
+    captioner, cap_processor = load_qwen_model(CAPTIONER_ID)
+
+    captions = {}
+    for audio_path in tqdm(files_to_process, desc="Captioning"):
+        try:
+            audio_data, sr = load_audio_mono_16k(audio_path)
+            caption = run_qwen_audio(
+                captioner, cap_processor, audio_data, sr,
+                "*Task* Describe this music in detail. Include genre, mood, instrumentation, tempo feel, and vocal style if present."
+            )
+            captions[audio_path] = caption
+        except Exception as e:
+            print(f"\n  Error captioning {os.path.basename(audio_path)}: {e}")
+            captions[audio_path] = ""
+
+    if args.low_vram:
+        print("  Offloading captioner to CPU...")
+        offload_to_cpu(captioner)
+        del captioner, cap_processor
+
+    # ── Stage 3: Lyrics transcription ────────────────────────────────────
+    print(f"\n[Stage 3/3] Transcribing lyrics for {len(files_to_process)} files...")
+    print("  Loading transcriber model...")
+    transcriber, trans_processor = load_qwen_model(TRANSCRIBER_ID)
+
+    lyrics_map = {}
+    for audio_path in tqdm(files_to_process, desc="Transcribing"):
+        try:
+            audio_data, sr = load_audio_mono_16k(audio_path)
+            lyrics = run_qwen_audio(
+                transcriber, trans_processor, audio_data, sr,
+                "*Task* Transcribe this audio in detail"
+            )
+            lyrics_map[audio_path] = lyrics
+        except Exception as e:
+            print(f"\n  Error transcribing {os.path.basename(audio_path)}: {e}")
+            lyrics_map[audio_path] = "[Instrumental]"
+
+    if args.low_vram:
+        print("  Offloading transcriber to CPU...")
+        offload_to_cpu(transcriber)
+        del transcriber, trans_processor
+
+    # ── Write output files ───────────────────────────────────────────────
+    print("\nWriting output files...")
+    for audio_path in files_to_process:
+        base_name = os.path.splitext(audio_path)[0]
+        output_path = base_name + ".txt"
+
+        caption = captions.get(audio_path, "")
+        lyrics = lyrics_map.get(audio_path, "[Instrumental]")
+        analysis = analyses[audio_path]
+
+        output = format_output(caption, lyrics, analysis, args.language)
+        with open(output_path, "w", encoding="utf-8") as f:
+            f.write(output)
+
+    print(f"Done! Processed {len(files_to_process)} files.")
+
+
+if __name__ == "__main__":
+    main()

toolkit/audio/album_artwork.py

@@ -0,0 +1,119 @@
+import io
+import os
+import numpy as np
+import av
+from PIL import Image, ImageDraw
+
+
+ARTWORK_DIR = os.path.dirname(os.path.abspath(__file__))
+BACKGROUND_PATH = os.path.join(ARTWORK_DIR, "ostris_logo.jpg")
+WAVEFORM_COLOR = (0xFB, 0xBF, 0x24, 230)  # #fbbf24 at 90% opacity
+ARTWORK_SIZE = 1024
+
+
+def load_waveform(audio_path: str, num_samples: int = 512) -> np.ndarray:
+    """Load audio and return a downsampled waveform envelope using PyAV."""
+    container = av.open(audio_path)
+    stream = container.streams.audio[0]
+    stream.codec_context.thread_type = "AUTO"
+
+    frames = []
+    for frame in container.decode(stream):
+        arr = frame.to_ndarray()
+        # mix down to mono
+        if arr.ndim > 1:
+            arr = arr.mean(axis=0)
+        frames.append(arr)
+    container.close()
+
+    audio = np.concatenate(frames)
+
+    # downsample to num_samples bins by taking max absolute value per bin
+    bin_size = len(audio) // num_samples
+    if bin_size == 0:
+        bin_size = 1
+    trimmed = audio[: bin_size * num_samples]
+    bins = trimmed.reshape(num_samples, bin_size)
+    envelope = np.max(np.abs(bins), axis=1)
+
+    # normalize to 0-1
+    peak = envelope.max()
+    if peak > 0:
+        envelope = envelope / peak
+    return envelope
+
+
+def create_artwork(waveform: np.ndarray, size: int = ARTWORK_SIZE) -> Image.Image:
+    """Create album artwork with logo background and waveform overlay."""
+    bg = Image.open(BACKGROUND_PATH).convert("RGBA").resize((size, size), Image.LANCZOS)
+
+    # draw waveform on separate overlay for alpha compositing
+    wave_overlay = Image.new("RGBA", (size, size), (0, 0, 0, 0))
+    draw = ImageDraw.Draw(wave_overlay)
+
+    num_bars = len(waveform)
+    padding = int(size * 0.02)
+    draw_w = size - 2 * padding
+    bar_width = max(1, draw_w / num_bars)
+    center_y = size // 2
+
+    max_amp = (size // 2) * 0.85  # leave a little margin
+
+    for i, amp in enumerate(waveform):
+        x = padding + i * bar_width
+        h = amp * max_amp
+        y_top = center_y - h
+        y_bot = center_y + h
+        draw.rectangle(
+            [x, y_top, x + bar_width - 1, y_bot],
+            fill=WAVEFORM_COLOR,
+        )
+
+    bg = Image.alpha_composite(bg, wave_overlay)
+    return bg.convert("RGB")
+
+
+def add_album_artwork(song_path: str) -> None:
+    """Add album artwork with waveform visualization to an MP3 file."""
+    from mutagen.id3 import ID3, APIC, ID3NoHeaderError
+
+    if not os.path.isfile(song_path):
+        raise FileNotFoundError(f"Audio file not found: {song_path}")
+
+    waveform = load_waveform(song_path)
+    artwork = create_artwork(waveform)
+
+    # encode artwork to JPEG bytes in memory
+    buf = io.BytesIO()
+    artwork.save(buf, format="JPEG", quality=85)
+    artwork_data = buf.getvalue()
+
+    # embed into MP3 via mutagen ID3 tags
+    try:
+        tags = ID3(song_path)
+    except ID3NoHeaderError:
+        tags = ID3()
+
+    tags.delall("APIC")
+    tags.add(
+        APIC(
+            encoding=3,  # UTF-8
+            mime="image/jpeg",
+            type=3,  # front cover
+            desc="Cover",
+            data=artwork_data,
+        )
+    )
+    tags.save(song_path, v2_version=3)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Add album artwork with waveform to an MP3 file"
+    )
+    parser.add_argument("mp3", help="Path to the MP3 file")
+    args = parser.parse_args()
+
+    add_album_artwork(args.mp3)

toolkit/audio/make_video.py

@@ -0,0 +1,149 @@
+import os
+import numpy as np
+import av
+from PIL import Image, ImageDraw
+
+
+ARTWORK_DIR = os.path.dirname(os.path.abspath(__file__))
+BACKGROUND_PATH = os.path.join(ARTWORK_DIR, "ostris_logo.jpg")
+WAVEFORM_COLOR = (0xFB, 0xBF, 0x24, 230)  # #fbbf24 at 90% opacity
+ARTWORK_SIZE = 1024
+
+
+def load_waveform(audio_path: str, num_samples: int = 512) -> np.ndarray:
+    """Load audio and return a downsampled waveform envelope using PyAV."""
+    container = av.open(audio_path)
+    stream = container.streams.audio[0]
+    stream.codec_context.thread_type = "AUTO"
+
+    frames = []
+    for frame in container.decode(stream):
+        arr = frame.to_ndarray()
+        # mix down to mono
+        if arr.ndim > 1:
+            arr = arr.mean(axis=0)
+        frames.append(arr)
+    container.close()
+
+    audio = np.concatenate(frames)
+
+    # downsample to num_samples bins by taking max absolute value per bin
+    bin_size = len(audio) // num_samples
+    if bin_size == 0:
+        bin_size = 1
+    trimmed = audio[: bin_size * num_samples]
+    bins = trimmed.reshape(num_samples, bin_size)
+    envelope = np.max(np.abs(bins), axis=1)
+
+    # normalize to 0-1
+    peak = envelope.max()
+    if peak > 0:
+        envelope = envelope / peak
+    return envelope
+
+
+def create_artwork(waveform: np.ndarray, size: int = ARTWORK_SIZE) -> Image.Image:
+    """Create album artwork with logo background and waveform overlay."""
+    bg = Image.open(BACKGROUND_PATH).convert("RGBA").resize((size, size), Image.LANCZOS)
+
+    # draw waveform on separate overlay for alpha compositing
+    wave_overlay = Image.new("RGBA", (size, size), (0, 0, 0, 0))
+    draw = ImageDraw.Draw(wave_overlay)
+
+    num_bars = len(waveform)
+    padding = int(size * 0.02)
+    draw_w = size - 2 * padding
+    bar_width = max(1, draw_w / num_bars)
+    center_y = size // 2
+
+    max_amp = (size // 2) * 0.85  # leave a little margin
+
+    for i, amp in enumerate(waveform):
+        x = padding + i * bar_width
+        h = amp * max_amp
+        y_top = center_y - h
+        y_bot = center_y + h
+        draw.rectangle(
+            [x, y_top, x + bar_width - 1, y_bot],
+            fill=WAVEFORM_COLOR,
+        )
+
+    bg = Image.alpha_composite(bg, wave_overlay)
+    return bg.convert("RGB")
+
+
+def make_video(song_path: str, video_size: int = 512) -> str:
+    """Create an MP4 video with album artwork as a static image for the duration of the audio."""
+    if not os.path.isfile(song_path):
+        raise FileNotFoundError(f"Audio file not found: {song_path}")
+
+    waveform = load_waveform(song_path)
+    artwork = create_artwork(waveform)
+    artwork = artwork.resize((video_size, video_size), Image.LANCZOS)
+
+    # get audio duration
+    container = av.open(song_path)
+    duration = float(container.duration) / av.time_base
+    container.close()
+
+    # output path: same name as input but .mp4, in the same directory
+    base, _ = os.path.splitext(song_path)
+    output_path = base + ".mp4"
+
+    fps = 1  # static image, 1 fps is enough
+    total_frames = max(1, int(duration * fps))
+
+    # convert artwork to numpy array for video encoding
+    frame_data = np.array(artwork)
+
+    out_container = av.open(output_path, mode="w")
+    video_stream = out_container.add_stream("libx264", rate=fps)
+    video_stream.width = video_size
+    video_stream.height = video_size
+    video_stream.pix_fmt = "yuv420p"
+
+    for _ in range(total_frames):
+        frame = av.VideoFrame.from_ndarray(frame_data, format="rgb24")
+        for packet in video_stream.encode(frame):
+            out_container.mux(packet)
+
+    # flush
+    for packet in video_stream.encode():
+        out_container.mux(packet)
+
+    out_container.close()
+
+    # mux audio into the video using ffmpeg via subprocess
+    import subprocess
+    final_path = base + "_final.mp4"
+    subprocess.run(
+        [
+            "ffmpeg", "-y",
+            "-i", output_path,
+            "-i", song_path,
+            "-c:v", "copy",
+            "-c:a", "aac",
+            "-shortest",
+            final_path,
+        ],
+        check=True,
+        capture_output=True,
+    )
+
+    # replace silent video with final muxed version
+    os.replace(final_path, output_path)
+
+    return output_path
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Create an MP4 video with album artwork from an audio file"
+    )
+    parser.add_argument("audio", help="Path to the audio file")
+    args = parser.parse_args()
+
+    out = make_video(args.audio)
+    print(f"Created video: {out}")

toolkit/config_modules.py

@@ -6,6 +6,7 @@ import random
 import torch
 import torchaudio
 
+from toolkit.audio.album_artwork import add_album_artwork
 from toolkit.prompt_utils import PromptEmbeds
 from torchao.quantization.quant_primitives import _DTYPE_TO_BIT_WIDTH
 
@@ -1201,13 +1202,16 @@ class GenerateImageConfig:
                 raise ValueError(f"Unsupported video format {self.output_ext}")
         elif self.output_ext in ['wav', 'mp3']:
             # save audio file
+            audio_path = self.get_image_path(count, max_count)
             torchaudio.save(
-                self.get_image_path(count, max_count), 
+                audio_path, 
                 image[0].to('cpu'),
                 sample_rate=48000, 
                 format=None, 
                 backend=None
             )
+            if self.output_ext == 'mp3':
+                add_album_artwork(audio_path)
         else:
             # TODO save image gen header info for A1111 and us, our seeds probably wont match
             image.save(self.get_image_path(count, max_count))

toolkit/data_loader.py

@@ -38,6 +38,7 @@ if TYPE_CHECKING:
 
 image_extensions = ['.jpg', '.jpeg', '.png', '.webp']
 video_extensions = ['.mp4', '.avi', '.mov', '.webm', '.mkv', '.wmv', '.m4v', '.flv']
+audio_extensions = ['.mp3', '.wav', '.flac', '.aac', '.ogg', '.m4a']
 
 
 class RescaleTransform:
@@ -393,6 +394,7 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
         # update bucket divisibility
         self.dataset_config.bucket_tolerance = sd.get_bucket_divisibility()
         self.is_video = dataset_config.num_frames > 1 or dataset_config.auto_frame_count
+        self.is_audio_model = hasattr(sd, 'is_audio_model') and sd.is_audio_model if sd is not None else False
         super().__init__()
         folder_path = dataset_config.folder_path
         self.dataset_path = dataset_config.dataset_path
@@ -425,10 +427,13 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
         # check if dataset_path is a folder or json
         if os.path.isdir(self.dataset_path):
             extensions = image_extensions
-            if self.is_video:
+            if self.is_audio_model:
+                # only look for audio files
+                extensions = audio_extensions
+            elif self.is_video:
                 # only look for videos
                 extensions = video_extensions
-            file_list = [os.path.join(root, file) for root, _, files in os.walk(self.dataset_path) for file in files if file.lower().endswith(tuple(extensions))]
+            file_list = [os.path.join(root, file) for root, _, files in os.walk(self.dataset_path) for file in files if file.lower().endswith(tuple(extensions)) and not file.startswith('.')]
         else:
             # assume json
             with open(self.dataset_path, 'r') as f:
@@ -522,6 +527,7 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
                 file_item = FileItemDTO(
                     sd=self.sd,
                     path=file,
+                    is_audio_model=self.is_audio_model,
                     dataset_config=dataset_config,
                     dataloader_transforms=self.transform,
                     size_database=self.size_database,
@@ -531,6 +537,7 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
                     te_padding_side=self.sd.te_padding_side if self.sd else "right",
                     latent_space_version=latent_space_version,
                     temporal_compression=temporal_compression,
+                    sample_rate=self.sd.sample_rate if self.is_audio_model and self.sd is not None else 48000,
                 )
                 self.file_list.append(file_item)
             except Exception as e:

toolkit/data_transfer_object/data_loader.py

@@ -5,7 +5,8 @@ import torch
 
 from PIL import Image
 from PIL.ImageOps import exif_transpose
-
+import av
+            
 from toolkit import image_utils
 from toolkit.basic import get_quick_signature_string
 from toolkit.dataloader_mixins import (
@@ -21,6 +22,7 @@ from toolkit.dataloader_mixins import (
     ClipImageFileItemDTOMixin,
     InpaintControlFileItemDTOMixin,
     TextEmbeddingFileItemDTOMixin,
+    AudioProcessingDTOMixin,
 )
 from toolkit.prompt_utils import PromptEmbeds, concat_prompt_embeds
 
@@ -42,6 +44,7 @@ class FileItemDTO(
     TextEmbeddingFileItemDTOMixin,
     CaptionProcessingDTOMixin,
     ImageProcessingDTOMixin,
+    AudioProcessingDTOMixin,
     ControlFileItemDTOMixin,
     InpaintControlFileItemDTOMixin,
     ClipImageFileItemDTOMixin,
@@ -55,6 +58,8 @@ class FileItemDTO(
         self.path = kwargs.get("path", "")
         self.dataset_config: "DatasetConfig" = kwargs.get("dataset_config", None)
         self.is_video = self.dataset_config.num_frames > 1 or self.dataset_config.auto_frame_count
+        self.is_audio_model = kwargs.get("is_audio_model", False)
+        self.sample_rate = kwargs.get("sample_rate", 48000)
         self.num_frames = self.dataset_config.num_frames
         self.temporal_compression = kwargs.get("temporal_compression", 8)
         size_database = kwargs.get("size_database", {})
@@ -84,8 +89,16 @@ class FileItemDTO(
                 and db_entry[2] == file_signature
             ):
                 use_db_entry = True
-
-        if use_db_entry:
+        if self.is_audio_model:
+            # get the length of the audio file in ms
+            with av.open(self.path) as c:
+                if c.duration is not None:
+                    w =  int(c.duration / 1_000)
+                else:
+                    s = c.streams.audio[0]
+                    w = int(float(s.duration * s.time_base) * 1_000)
+            h = 1
+        elif use_db_entry:
             w, h, _ = size_database[file_key]
         elif self.is_video:
             # Open the video file

toolkit/dataloader_mixins.py

@@ -231,6 +231,12 @@ class BucketsMixin:
         # for file_item in enumerate(file_list):
         for idx, file_item in enumerate(file_list):
             file_item: 'FileItemDTO' = file_item
+            if self.is_audio_model:
+                bucket_key = f"{file_item.width}ms"
+                if bucket_key not in self.buckets:
+                    self.buckets[bucket_key] = Bucket(file_item.width, 1)
+                self.buckets[bucket_key].file_list_idx.append(idx)
+                continue
             width = int(file_item.width * file_item.dataset_config.scale)
             height = int(file_item.height * file_item.dataset_config.scale)
 
@@ -465,6 +471,27 @@ class CaptionProcessingDTOMixin:
             pass
         return caption
 
+class AudioProcessingDTOMixin:
+    def load_and_process_audio(self: 'FileItemDTO'):
+        # Default to "no audio" unless we successfully extract it
+        self.audio_data = None
+        self.audio_tensor = None
+        self.tensor = None
+        try:
+            import torchaudio
+
+            waveform, sample_rate = torchaudio.load(self.path)  # [channels, samples]
+            waveform = waveform_to_stereo(waveform)  # Convert to stereo if not already
+            if sample_rate != self.sample_rate:
+                waveform = torchaudio.functional.resample(waveform, sample_rate, self.sample_rate)
+            self.tensor = waveform
+            self.audio_tensor = waveform
+            self.audio_data = {"waveform": waveform, "sample_rate": int(self.sample_rate)}
+
+        except Exception as e:
+            # if issue with libtorchcodec "Could not load libtorchcodec"
+            raise Exception(f"** WARNING ** - Error Processing audio for {self.path}. Error: {e}")
+        
 
 class ImageProcessingDTOMixin:
     def load_and_process_video(
@@ -772,6 +799,9 @@ class ImageProcessingDTOMixin:
             if self.has_unconditional:
                 self.load_unconditional_image()
             return
+        if self.is_audio_model:
+            self.load_and_process_audio()
+            return
         if self.dataset_config.num_frames > 1 or self.dataset_config.auto_frame_count:
             self.load_and_process_video(transform, only_load_latents)
             return
@@ -1852,6 +1882,8 @@ class LatentCachingMixin:
         self.latent_cache = {}
 
     def cache_latents_all_latents(self: 'AiToolkitDataset'):
+        if self.is_audio_model:
+            raise Exception("Audio models are not supported for latent caching yet")
         with accelerator.main_process_first():
             print_acc(f"Caching latents for {self.dataset_path}")
             # cache all latents to disk

toolkit/models/base_model.py

@@ -427,7 +427,7 @@ class BaseModel:
                 if network is not None:
                     assert network.is_active
 
-                for i in tqdm(range(len(image_configs)), desc=f"Generating Images", leave=False):
+                for i in tqdm(range(len(image_configs)), desc=f"Generating Samples", leave=False):
                     gen_config = image_configs[i]
 
                     extra = {}
@@ -1590,7 +1590,7 @@ class BaseModel:
     
     def get_base_model_version(self) -> str:
         # override in child classes to get the base model version
-        return "unknown"
+        return self.arch if self.arch is not None else 'unknown'
 
     def get_model_to_train(self):
         # called to get model to attach LoRAs to. Can be overridden in child classes

ui/prisma/schema.prisma

@@ -36,7 +36,12 @@ model Job {
   speed_string    String   @default("")
   queue_position  Int      @default(0)
   pid             Int?
+  job_type        String   @default("train") // 'train', 'caption'
+  job_ref         String? // can be used for anything for special jobs, like dataset path for caption jobs
 
   @@index([status])
   @@index([gpu_ids])
+  @@index([job_type])
+  @@index([job_ref])
+
 }

ui/src/app/api/audio/art/[...audioPath]/route.ts

@@ -0,0 +1,184 @@
+/* eslint-disable */
+import { NextRequest, NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot, getTrainingFolder, getDataRoot } from '@/server/settings';
+
+/**
+ * Serves embedded album art from an MP3 file's ID3v2 tag.
+ * Reads only the tag header from disk (no full-file buffering).
+ * Returns the raw image bytes with correct Content-Type.
+ */
+
+function synchsafeToInt(b0: number, b1: number, b2: number, b3: number) {
+  return ((b0 & 0x7f) << 21) | ((b1 & 0x7f) << 14) | ((b2 & 0x7f) << 7) | (b3 & 0x7f);
+}
+
+function deUnsync(bytes: Buffer) {
+  const out: number[] = [];
+  for (let i = 0; i < bytes.length; i++) {
+    out.push(bytes[i]);
+    if (bytes[i] === 0xff && i + 1 < bytes.length && bytes[i + 1] === 0x00) i += 1;
+  }
+  return Buffer.from(out);
+}
+
+function readNullTerminated(buf: Buffer, start: number, wide: boolean): { text: string; next: number } {
+  if (wide) {
+    let i = start;
+    while (i + 1 < buf.length && !(buf[i] === 0 && buf[i + 1] === 0)) i += 2;
+    return { text: buf.slice(start, i).toString('utf16le'), next: i + 2 };
+  }
+  let i = start;
+  while (i < buf.length && buf[i] !== 0) i++;
+  return { text: buf.slice(start, i).toString('latin1'), next: i + 1 };
+}
+
+type ArtResult = { mime: string; data: Buffer } | null;
+
+function extractArtFromTag(buf: Buffer): ArtResult {
+  if (buf.length < 10) return null;
+  if (buf[0] !== 0x49 || buf[1] !== 0x44 || buf[2] !== 0x33) return null; // "ID3"
+
+  const verMajor = buf[3]; // 2, 3, or 4
+  const flags = buf[5];
+  const tagSize = synchsafeToInt(buf[6], buf[7], buf[8], buf[9]);
+  const tagEnd = Math.min(10 + tagSize, buf.length);
+
+  let tagData = buf.slice(10, tagEnd);
+  if ((flags & 0x80) !== 0) tagData = deUnsync(tagData);
+
+  let offset = 0;
+
+  // Skip extended header
+  if ((verMajor === 3 || verMajor === 4) && (flags & 0x40) !== 0 && tagData.length >= 4) {
+    const extSize =
+      verMajor === 4
+        ? synchsafeToInt(tagData[0], tagData[1], tagData[2], tagData[3])
+        : (tagData[0] << 24) | (tagData[1] << 16) | (tagData[2] << 8) | tagData[3];
+    offset += 4 + Math.max(0, extSize);
+  }
+
+  while (offset < tagData.length) {
+    if (tagData[offset] === 0x00) break;
+
+    if (verMajor === 2) {
+      // ID3v2.2: 3-byte frame ID, 3-byte size
+      if (offset + 6 > tagData.length) break;
+      const id = tagData.slice(offset, offset + 3).toString('latin1');
+      const size = (tagData[offset + 3] << 16) | (tagData[offset + 4] << 8) | tagData[offset + 5];
+      offset += 6;
+      if (!id.trim() || size <= 0 || offset + size > tagData.length) break;
+
+      if (id === 'PIC' && size > 6) {
+        const frame = tagData.slice(offset, offset + size);
+        const fmt = frame.slice(1, 4).toString('latin1').toLowerCase();
+        const mime = fmt === 'png' ? 'image/png' : 'image/jpeg';
+        // skip: encoding(1) + format(3) + pictureType(1) = 5, then null-terminated description
+        let p = 5;
+        const enc = frame[0];
+        const wide = enc === 1 || enc === 2;
+        const desc = readNullTerminated(frame as any, p, wide);
+        p = desc.next;
+        if (p < frame.length) {
+          const img = frame.slice(p);
+          if (img.length > 64) return { mime, data: Buffer.from(img) };
+        }
+      }
+      offset += size;
+    } else {
+      // ID3v2.3/v2.4: 4-byte frame ID, 4-byte size, 2-byte flags
+      if (offset + 10 > tagData.length) break;
+      const id = tagData.slice(offset, offset + 4).toString('latin1');
+      let size =
+        verMajor === 4
+          ? synchsafeToInt(tagData[offset + 4], tagData[offset + 5], tagData[offset + 6], tagData[offset + 7])
+          : (tagData[offset + 4] << 24) |
+            (tagData[offset + 5] << 16) |
+            (tagData[offset + 6] << 8) |
+            tagData[offset + 7];
+      const flag2 = tagData[offset + 9];
+      offset += 10;
+      if (!id.trim() || size <= 0 || offset + size > tagData.length) break;
+
+      if (id === 'APIC') {
+        let frame = tagData.slice(offset, offset + size);
+        if (verMajor === 4 && (flag2 & 0x02) !== 0) frame = deUnsync(frame);
+
+        const enc = frame[0];
+        // mime type: null-terminated latin1
+        const mimeZ = readNullTerminated(frame as any, 1, false);
+        const mime = mimeZ.text || 'image/jpeg';
+        let p = mimeZ.next;
+        if (p < frame.length) p += 1; // picture type byte
+        const wide = enc === 1 || enc === 2;
+        const desc = readNullTerminated(frame as any, p, wide);
+        p = desc.next;
+        if (p < frame.length) {
+          const img = frame.slice(p);
+          if (img.length > 64) return { mime, data: Buffer.from(img) };
+        }
+      }
+      offset += size;
+    }
+  }
+  return null;
+}
+
+export async function GET(request: NextRequest, { params }: { params: { audioPath: string } }) {
+  const { audioPath } = await params;
+  try {
+    const filepath = decodeURIComponent(audioPath);
+
+    // Security check
+    const datasetRoot = await getDatasetsRoot();
+    const trainingRoot = await getTrainingFolder();
+    const dataRoot = await getDataRoot();
+    const allowedDirs = [datasetRoot, trainingRoot, dataRoot];
+    const isAllowed = allowedDirs.some(d => filepath.startsWith(d)) && !filepath.includes('..');
+    if (!isAllowed) {
+      return new NextResponse('Access denied', { status: 403 });
+    }
+
+    const stat = await fs.promises.stat(filepath).catch(() => null);
+    if (!stat || !stat.isFile()) {
+      return new NextResponse('File not found', { status: 404 });
+    }
+
+    // Read only the ID3 tag (first min(tagSize, 4MB) bytes).
+    // First read 10 bytes to get tag size, then read the full tag.
+    const fd = await fs.promises.open(filepath, 'r');
+    try {
+      const headerBuf = Buffer.alloc(10);
+      await fd.read(headerBuf, 0, 10, 0);
+
+      if (headerBuf[0] !== 0x49 || headerBuf[1] !== 0x44 || headerBuf[2] !== 0x33) {
+        return new NextResponse('No ID3 tag', { status: 404 });
+      }
+
+      const tagSize = synchsafeToInt(headerBuf[6], headerBuf[7], headerBuf[8], headerBuf[9]);
+      const totalRead = Math.min(10 + tagSize, 4_000_000);
+
+      const tagBuf = Buffer.alloc(totalRead);
+      await fd.read(tagBuf, 0, totalRead, 0);
+
+      const art = extractArtFromTag(tagBuf);
+      if (!art) {
+        return new NextResponse('No album art found', { status: 404 });
+      }
+
+      return new NextResponse(art.data, {
+        headers: {
+          'Content-Type': art.mime,
+          'Content-Length': String(art.data.length),
+          'Cache-Control': 'public, max-age=604800, immutable',
+        },
+      });
+    } finally {
+      await fd.close();
+    }
+  } catch (error) {
+    console.error('Error extracting album art:', error);
+    return new NextResponse('Internal Server Error', { status: 500 });
+  }
+}

ui/src/app/api/jobs/route.ts

@@ -7,6 +7,8 @@ const prisma = new PrismaClient();
 export async function GET(request: Request) {
   const { searchParams } = new URL(request.url);
   const id = searchParams.get('id');
+  const job_ref = searchParams.get('job_ref');
+  const job_type = searchParams.get('job_type');
 
   try {
     if (id) {
@@ -15,8 +17,16 @@ export async function GET(request: Request) {
       });
       return NextResponse.json(job);
     }
+    if (job_ref) {
+      const job = await prisma.job.findFirst({
+        where: { job_ref },
+        orderBy: { updated_at: 'desc' },
+      });
+      return NextResponse.json(job);
+    }
 
     const jobs = await prisma.job.findMany({
+      where: job_type ? { job_type } : undefined,
       orderBy: { created_at: 'desc' },
     });
     return NextResponse.json({ jobs: jobs });
@@ -36,6 +46,15 @@ export async function POST(request: Request) {
       gpu_ids = "mps";
     }
 
+    const extra: any = {};
+    if ("job_ref" in body) {
+      extra["job_ref"] = body.job_ref;
+    }
+
+    if ("job_type" in body) {
+      extra["job_type"] = body.job_type;
+    }
+
     if (id) {
       // Update existing training
       const training = await prisma.job.update({
@@ -44,6 +63,7 @@ export async function POST(request: Request) {
           name,
           gpu_ids,
           job_config: JSON.stringify(job_config),
+          ...extra,
         },
       });
       return NextResponse.json(training);
@@ -63,6 +83,7 @@ export async function POST(request: Request) {
           gpu_ids,
           job_config: JSON.stringify(job_config),
           queue_position: newQueuePosition,
+          ...extra,
         },
       });
       return NextResponse.json(training);

ui/src/app/datasets/[datasetName]/page.tsx

@@ -8,12 +8,18 @@ import { Button } from '@headlessui/react';
 import AddImagesModal, { openImagesModal, useOpenImagesModalOnDrag } from '@/components/AddImagesModal';
 import { TopBar, MainContent } from '@/components/layout';
 import { apiClient } from '@/utils/api';
+import { CaptionDatasetModal, openCaptionDatasetModal } from '@/components/CaptionDatasetModal';
+import useSettings from '@/hooks/useSettings';
+import { pathJoin } from '@/utils/basic';
+import AutoCaptionButton from '@/components/AutoCaptionButton';
 
 export default function DatasetPage({ params }: { params: { datasetName: string } }) {
   const [imgList, setImgList] = useState<{ img_path: string }[]>([]);
+  const [isAutoCaptioning, setIsAutoCaptioning] = useState(false);
   const usableParams = use(params as any) as { datasetName: string };
   const datasetName = usableParams.datasetName;
   const [status, setStatus] = useState<'idle' | 'loading' | 'success' | 'error'>('idle');
+  const { settings, isSettingsLoaded } = useSettings();
 
   const refreshImageList = (dbName: string) => {
     setStatus('loading');
@@ -105,6 +111,10 @@ export default function DatasetPage({ params }: { params: { datasetName: string
         </div>
         <div className="flex-1"></div>
         <div>
+          <AutoCaptionButton
+            datasetPath={`${pathJoin(settings.DATASETS_FOLDER, datasetName)}`}
+            setIsAutoCaptioning={setIsAutoCaptioning}
+          />
           <Button
             className="text-white bg-slate-600 px-3 py-1 rounded-md"
             onClick={() => openImagesModal(datasetName, () => refreshImageList(datasetName))}
@@ -121,6 +131,7 @@ export default function DatasetPage({ params }: { params: { datasetName: string
               <DatasetImageCard
                 key={img.img_path}
                 alt="image"
+                isAutoCaptioning={isAutoCaptioning}
                 imageUrl={img.img_path}
                 onDelete={() => refreshImageList(datasetName)}
               />
@@ -129,6 +140,7 @@ export default function DatasetPage({ params }: { params: { datasetName: string
         )}
       </MainContent>
       <AddImagesModal />
+      <CaptionDatasetModal />
     </>
   );
 }

ui/src/app/jobs/[jobID]/page.tsx

@@ -23,6 +23,7 @@ interface Page {
   component: React.ComponentType<{ job: Job }>;
   menuItem?: React.ComponentType<{ job?: Job | null }> | null;
   mainCss?: string;
+  jobTypes?: string[]; // if specified, only show this page for these job types
 }
 
 const pages: Page[] = [
@@ -40,6 +41,7 @@ const pages: Page[] = [
     component: SampleImages,
     menuItem: SampleImagesMenu,
     mainCss: 'pt-24',
+    jobTypes: ['train'],
   },
   {
     name: 'Loss Graph',
@@ -47,6 +49,7 @@ const pages: Page[] = [
     icon: MdShowChart,
     component: JobLossGraph,
     mainCss: 'pt-24',
+    jobTypes: ['train'],
   },
   {
     name: 'Config File',
@@ -65,6 +68,13 @@ export default function JobPage({ params }: { params: { jobID: string } }) {
 
   const page = pages.find(p => p.value === pageKey);
 
+  const jobType = job?.job_type || 'unknown';
+
+  let title = `Job: ${job?.name || 'Loading...'}`;
+  if (jobType === 'caption') {
+    title = `Captioning: ${job?.job_ref || 'Loading...'}`;
+  }
+
   return (
     <>
       {/* Fixed top bar */}
@@ -75,7 +85,7 @@ export default function JobPage({ params }: { params: { jobID: string } }) {
           </Button>
         </div>
         <div>
-          <h1 className="text-lg">Job: {job?.name}</h1>
+          <h1 className="text-lg">{title}</h1>
         </div>
         <div className="flex-1"></div>
         {job && (
@@ -103,16 +113,21 @@ export default function JobPage({ params }: { params: { jobID: string } }) {
         )}
       </MainContent>
       <div className="bg-gray-800 absolute top-12 left-0 w-full h-8 flex items-center px-2 text-sm">
-        {pages.map(page => (
-          <Button
-            key={page.value}
-            onClick={() => setPageKey(page.value)}
-            className={`px-4 py-1 h-8 flex items-center gap-1.5 ${page.value === pageKey ? 'bg-gray-300 dark:bg-gray-700 text-white' : ''}`}
-          >
-            <page.icon className="text-sm" />
-            {page.name}
-          </Button>
-        ))}
+        {pages.map(page => {
+          if (page.jobTypes && !page.jobTypes.includes(jobType)) {
+            return null;
+          }
+          return (
+            <Button
+              key={page.value}
+              onClick={() => setPageKey(page.value)}
+              className={`px-4 py-1 h-8 flex items-center gap-1.5 ${page.value === pageKey ? 'bg-gray-300 dark:bg-gray-700 text-white' : ''}`}
+            >
+              <page.icon className="text-sm" />
+              {page.name}
+            </Button>
+          );
+        })}
         {page?.menuItem && (
           <>
             <div className="flex-grow"></div>

ui/src/app/jobs/new/AdvancedJob.tsx

@@ -35,23 +35,38 @@ const yamlConfig: YAML.DocumentOptions &
   directives: true,
 };
 
+function toYaml(obj: any): string {
+  const doc = new YAML.Document(obj, yamlConfig);
+  YAML.visit(doc, {
+    Scalar(_key, node) {
+      if (typeof node.value === 'string' && node.value.includes('\n')) {
+        node.type = YAML.Scalar.BLOCK_LITERAL;
+      }
+    },
+  });
+  return doc.toString(yamlConfig);
+}
+
 export default function AdvancedJob({ jobConfig, setJobConfig, settings }: Props) {
   const { theme } = useTheme();
   const [editorValue, setEditorValue] = useState<string>('');
+  const [hasError, setHasError] = useState(false);
   const lastJobConfigUpdateStringRef = useRef('');
   const editorRef = useRef<editor.IStandaloneCodeEditor | null>(null);
+  const monacoRef = useRef<any>(null);
 
   // Track if the editor has been mounted
   const isEditorMounted = useRef(false);
 
   // Handler for editor mounting
-  const handleEditorDidMount: OnMount = editor => {
+  const handleEditorDidMount: OnMount = (editor, monaco) => {
     editorRef.current = editor;
+    monacoRef.current = monaco;
     isEditorMounted.current = true;
 
     // Initial content setup
     try {
-      const yamlContent = YAML.stringify(jobConfig, yamlConfig);
+      const yamlContent = toYaml(jobConfig);
       setEditorValue(yamlContent);
       lastJobConfigUpdateStringRef.current = JSON.stringify(jobConfig);
     } catch (e) {
@@ -78,7 +93,7 @@ export default function AdvancedJob({ jobConfig, setJobConfig, settings }: Props
         const scrollTop = editor.getScrollTop();
 
         // Update content
-        const yamlContent = YAML.stringify(jobConfig, yamlConfig);
+        const yamlContent = toYaml(jobConfig);
 
         // Only update if the content is actually different
         if (yamlContent !== editor.getValue()) {
@@ -98,11 +113,29 @@ export default function AdvancedJob({ jobConfig, setJobConfig, settings }: Props
     }
   }, [jobConfig]);
 
+  const setMarkers = (errors: { message: string; line: number }[]) => {
+    const monaco = monacoRef.current;
+    const model = editorRef.current?.getModel();
+    if (!monaco || !model) return;
+    const markers = errors.map(err => ({
+      severity: monaco.MarkerSeverity.Error,
+      message: err.message,
+      startLineNumber: err.line,
+      startColumn: 1,
+      endLineNumber: err.line,
+      endColumn: model.getLineMaxColumn(err.line),
+    }));
+    monaco.editor.setModelMarkers(model, 'yaml', markers);
+  };
+
   const handleChange = (value: string | undefined) => {
     if (value === undefined) return;
 
     try {
       const parsed = YAML.parse(value);
+      setHasError(false);
+      setMarkers([]);
+
       // Don't update jobConfig if the change came from the editor itself
       // to avoid a circular update loop
       if (JSON.stringify(parsed) !== lastJobConfigUpdateStringRef.current) {
@@ -121,14 +154,21 @@ export default function AdvancedJob({ jobConfig, setJobConfig, settings }: Props
         migrateJobConfig(parsed);
         setJobConfig(parsed);
       }
-    } catch (e) {
-      // Don't update on parsing errors
-      console.warn(e);
+    } catch (e: any) {
+      setHasError(true);
+      const line = e?.linePos?.[0]?.line ?? e?.linePos?.line ?? 1;
+      setMarkers([{ message: e?.message ?? 'Invalid YAML', line }]);
     }
   };
 
   return (
-    <>
+    <div className="relative h-full w-full">
+      {hasError && (
+        <div
+          className="absolute inset-0 z-10 pointer-events-none rounded-sm"
+          style={{ boxShadow: 'inset 0 0 12px 2px rgba(239, 68, 68, 0.5)' }}
+        />
+      )}
       <Editor
         height="100%"
         width="100%"
@@ -143,6 +183,6 @@ export default function AdvancedJob({ jobConfig, setJobConfig, settings }: Props
           automaticLayout: true,
         }}
       />
-    </>
+    </div>
   );
 }

ui/src/app/jobs/new/SimpleJob.tsx

@@ -7,11 +7,20 @@ import {
   quantizationOptions,
   defaultQtype,
   jobTypeOptions,
+  SampleTags,
 } from './options';
 import { defaultDatasetConfig } from './jobConfig';
 import { GroupedSelectOption, JobConfig, SelectOption } from '@/types';
-import { objectCopy } from '@/utils/basic';
-import { TextInput, SelectInput, Checkbox, FormGroup, NumberInput, SliderInput } from '@/components/formInputs';
+import { objectCopy, tagsToObj, objToTags } from '@/utils/basic';
+import {
+  TextInput,
+  TextAreaInput,
+  SelectInput,
+  Checkbox,
+  FormGroup,
+  NumberInput,
+  SliderInput,
+} from '@/components/formInputs';
 import Card from '@/components/Card';
 import { X, Copy } from 'lucide-react';
 import AddSingleImageModal, { openAddImageModal } from '@/components/AddSingleImageModal';
@@ -68,6 +77,44 @@ export default function SimpleJob({
   }, [modelArch, jobType]);
 
   const isVideoModel = !!(modelArch?.group === 'video');
+  const isAudioModel = !!(modelArch?.group === 'audio');
+
+  const taggedSampleArr: Record<string, any>[] | null = useMemo(() => {
+    if (!modelArch) return null;
+    if (!modelArch.sampleTags) return null;
+    if (!jobConfig.config.process[0].sample.samples) return null;
+    let sampleArr: any[] = [];
+    for (let i = 0; i < jobConfig.config.process[0].sample.samples.length; i++) {
+      const taggedPrompt = jobConfig.config.process[0].sample.samples[i].prompt;
+      const tagsObj = tagsToObj(taggedPrompt);
+      sampleArr.push(tagsObj);
+    }
+    return sampleArr;
+  }, [modelArch, jobConfig.config.process[0].sample.samples]);
+
+  const modelArchTagSections: SampleTags[] | null = useMemo(() => {
+    if (!modelArch?.sampleTags) return null;
+    const maxPerGroup = 5;
+    let sections: SampleTags[] = [];
+    let subSection: SampleTags = {};
+    for (const [tagKey, tag] of Object.entries(modelArch.sampleTags)) {
+      if ((tag.full && Object.keys(subSection).length > 0) || Object.keys(subSection).length >= maxPerGroup) {
+        // reset the sub section build if the next tag is full or max per group is reached
+        sections.push(subSection);
+        subSection = {};
+      }
+      subSection[tagKey] = tag;
+      if (tag.full) {
+        // if the tag is full, push the section immediately and reset the sub section build
+        sections.push(subSection);
+        subSection = {};
+      }
+    }
+    if (Object.keys(subSection).length > 0) {
+      sections.push(subSection);
+    }
+    return sections.length > 0 ? sections : null;
+  }, [modelArch]);
 
   const numTopCards = useMemo(() => {
     let count = 4; // job settings, model config, target config, save config
@@ -149,6 +196,26 @@ export default function SimpleJob({
 
   const showGPUSelect = !isMac();
 
+  let numDatasetCols = 4;
+  let numSampleTopCols = 4;
+  let datasetStyleClass = 'grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 gap-6';
+  let sampleTopStyleClass = 'grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 gap-6';
+  if (isVideoModel) {
+    numSampleTopCols += 1;
+  }
+  if (isAudioModel) {
+    numDatasetCols -= 1;
+    numSampleTopCols -= 1;
+  }
+  if (numDatasetCols == 3) {
+    datasetStyleClass = 'grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 gap-6';
+  }
+  if (numSampleTopCols == 5) {
+    sampleTopStyleClass = 'grid grid-cols-1 md:grid-cols-3 lg:grid-cols-5 gap-6';
+  }
+  if (numSampleTopCols == 3) {
+    sampleTopStyleClass = 'grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 gap-6';
+  }
   return (
     <>
       <form
@@ -790,7 +857,7 @@ export default function SimpleJob({
                     </button>
                   </div>
                   <h2 className="text-lg font-bold mb-4">Dataset {i + 1}</h2>
-                  <div className="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 gap-6">
+                  <div className={datasetStyleClass}>
                     <div>
                       <SelectInput
                         label="Target Dataset"
@@ -971,53 +1038,57 @@ export default function SimpleJob({
                           />
                         )}
                       </FormGroup>
-                      <FormGroup label="Flipping" docKey={'datasets.flip'} className="mt-2">
-                        <Checkbox
-                          label={
-                            <>
-                              Flip X <FlipHorizontal2 className="inline-block w-4 h-4 ml-1" />
-                            </>
-                          }
-                          checked={dataset.flip_x || false}
-                          onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].flip_x`)}
-                        />
-                        <Checkbox
-                          label={
-                            <>
-                              Flip Y <FlipVertical2 className="inline-block w-4 h-4 ml-1" />
-                            </>
-                          }
-                          checked={dataset.flip_y || false}
-                          onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].flip_y`)}
-                        />
-                      </FormGroup>
-                    </div>
-                    <div>
-                      <FormGroup label="Resolutions" className="pt-2">
-                        <div className="grid grid-cols-2 gap-2">
-                          {[
-                            [256, 512, 768, 1024],
-                            [1280, 1328, 1536],
-                          ].map(resGroup => (
-                            <div key={resGroup[0]} className="space-y-2">
-                              {resGroup.map(res => (
-                                <Checkbox
-                                  key={res}
-                                  label={res.toString()}
-                                  checked={dataset.resolution.includes(res)}
-                                  onChange={value => {
-                                    const resolutions = dataset.resolution.includes(res)
-                                      ? dataset.resolution.filter(r => r !== res)
-                                      : [...dataset.resolution, res];
-                                    setJobConfig(resolutions, `config.process[0].datasets[${i}].resolution`);
-                                  }}
-                                />
-                              ))}
-                            </div>
-                          ))}
-                        </div>
-                      </FormGroup>
+                      {!isAudioModel && (
+                        <FormGroup label="Flipping" docKey={'datasets.flip'} className="mt-2">
+                          <Checkbox
+                            label={
+                              <>
+                                Flip X <FlipHorizontal2 className="inline-block w-4 h-4 ml-1" />
+                              </>
+                            }
+                            checked={dataset.flip_x || false}
+                            onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].flip_x`)}
+                          />
+                          <Checkbox
+                            label={
+                              <>
+                                Flip Y <FlipVertical2 className="inline-block w-4 h-4 ml-1" />
+                              </>
+                            }
+                            checked={dataset.flip_y || false}
+                            onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].flip_y`)}
+                          />
+                        </FormGroup>
+                      )}
                     </div>
+                    {!isAudioModel && (
+                      <div>
+                        <FormGroup label="Resolutions" className="pt-2">
+                          <div className="grid grid-cols-2 gap-2">
+                            {[
+                              [256, 512, 768, 1024],
+                              [1280, 1328, 1536],
+                            ].map(resGroup => (
+                              <div key={resGroup[0]} className="space-y-2">
+                                {resGroup.map(res => (
+                                  <Checkbox
+                                    key={res}
+                                    label={res.toString()}
+                                    checked={dataset.resolution.includes(res)}
+                                    onChange={value => {
+                                      const resolutions = dataset.resolution.includes(res)
+                                        ? dataset.resolution.filter(r => r !== res)
+                                        : [...dataset.resolution, res];
+                                      setJobConfig(resolutions, `config.process[0].datasets[${i}].resolution`);
+                                    }}
+                                  />
+                                ))}
+                              </div>
+                            ))}
+                          </div>
+                        </FormGroup>
+                      </div>
+                    )}
                   </div>
                 </div>
               ))}
@@ -1039,13 +1110,7 @@ export default function SimpleJob({
         </div>
         <div>
           <Card title="Sample">
-            <div
-              className={
-                isVideoModel
-                  ? 'grid grid-cols-1 md:grid-cols-3 lg:grid-cols-5 gap-6'
-                  : 'grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 gap-6'
-              }
-            >
+            <div className={sampleTopStyleClass}>
               <div>
                 <NumberInput
                   label="Sample Every"
@@ -1084,47 +1149,50 @@ export default function SimpleJob({
                   required
                 />
               </div>
-              <div>
-                <NumberInput
-                  label="Width"
-                  value={jobConfig.config.process[0].sample.width}
-                  onChange={value => setJobConfig(value, 'config.process[0].sample.width')}
-                  placeholder="eg. 1024"
-                  min={0}
-                  required
-                />
-                <NumberInput
-                  label="Height"
-                  value={jobConfig.config.process[0].sample.height}
-                  onChange={value => setJobConfig(value, 'config.process[0].sample.height')}
-                  placeholder="eg. 1024"
-                  className="pt-2"
-                  min={0}
-                  required
-                />
-                {isVideoModel && (
-                  <div>
-                    <NumberInput
-                      label="Num Frames"
-                      value={jobConfig.config.process[0].sample.num_frames}
-                      onChange={value => setJobConfig(value, 'config.process[0].sample.num_frames')}
-                      placeholder="eg. 0"
-                      className="pt-2"
-                      min={0}
-                      required
-                    />
-                    <NumberInput
-                      label="FPS"
-                      value={jobConfig.config.process[0].sample.fps}
-                      onChange={value => setJobConfig(value, 'config.process[0].sample.fps')}
-                      placeholder="eg. 0"
-                      className="pt-2"
-                      min={0}
-                      required
-                    />
-                  </div>
-                )}
-              </div>
+
+              {!isAudioModel && (
+                <div>
+                  <NumberInput
+                    label="Width"
+                    value={jobConfig.config.process[0].sample.width}
+                    onChange={value => setJobConfig(value, 'config.process[0].sample.width')}
+                    placeholder="eg. 1024"
+                    min={0}
+                    required
+                  />
+                  <NumberInput
+                    label="Height"
+                    value={jobConfig.config.process[0].sample.height}
+                    onChange={value => setJobConfig(value, 'config.process[0].sample.height')}
+                    placeholder="eg. 1024"
+                    className="pt-2"
+                    min={0}
+                    required
+                  />
+                  {isVideoModel && (
+                    <div>
+                      <NumberInput
+                        label="Num Frames"
+                        value={jobConfig.config.process[0].sample.num_frames}
+                        onChange={value => setJobConfig(value, 'config.process[0].sample.num_frames')}
+                        placeholder="eg. 0"
+                        className="pt-2"
+                        min={0}
+                        required
+                      />
+                      <NumberInput
+                        label="FPS"
+                        value={jobConfig.config.process[0].sample.fps}
+                        onChange={value => setJobConfig(value, 'config.process[0].sample.fps')}
+                        placeholder="eg. 0"
+                        className="pt-2"
+                        min={0}
+                        required
+                      />
+                    </div>
+                  )}
+                </div>
+              )}
 
               <div>
                 <NumberInput
@@ -1199,68 +1267,143 @@ export default function SimpleJob({
                   <div className="flex-1">
                     <div className="flex">
                       <div className="flex-1">
-                        <TextInput
-                          label={`Prompt`}
-                          value={sample.prompt}
-                          onChange={value => setJobConfig(value, `config.process[0].sample.samples[${i}].prompt`)}
-                          placeholder="Enter prompt"
-                          required
-                        />
-                        <div className="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 gap-4 mt-2">
-                          <TextInput
-                            label={`Width`}
-                            value={sample.width ? `${sample.width}` : ''}
-                            onChange={value => {
-                              // remove any non-numeric characters
-                              value = value.replace(/\D/g, '');
-                              if (value === '') {
-                                // remove the key from the config if empty
-                                let newConfig = objectCopy(jobConfig);
-                                if (newConfig.config.process[0].sample.samples[i]) {
-                                  delete newConfig.config.process[0].sample.samples[i].width;
-                                  setJobConfig(
-                                    newConfig.config.process[0].sample.samples,
-                                    'config.process[0].sample.samples',
-                                  );
-                                }
-                              } else {
-                                const intValue = parseInt(value);
-                                if (!isNaN(intValue)) {
-                                  setJobConfig(intValue, `config.process[0].sample.samples[${i}].width`);
+                        {modelArch?.sampleTags && taggedSampleArr && modelArchTagSections ? (
+                          <>
+                            {modelArchTagSections.map((sampleTagSection, sti) => (
+                              <div key={sti} className="grid w-full lg:grid-flow-col lg:auto-cols-fr gap-4 mt-2">
+                                {Object.entries(sampleTagSection).map(([tagKey, tag]) => (
+                                  <div key={tagKey} className="mb-2">
+                                    {tag.type === 'text' && (
+                                      <TextInput
+                                        label={tag.title}
+                                        value={taggedSampleArr[i][tagKey] ?? ''}
+                                        onChange={value => {
+                                          let taggedSample = { ...taggedSampleArr[i] };
+                                          taggedSample[tagKey] = value;
+                                          setJobConfig(
+                                            objToTags(taggedSample),
+                                            `config.process[0].sample.samples[${i}].prompt`,
+                                          );
+                                        }}
+                                        placeholder={`Enter ${tag.title.toLowerCase()}`}
+                                      />
+                                    )}
+                                    {tag.type === 'multiline' && (
+                                      <TextAreaInput
+                                        label={tag.title}
+                                        value={taggedSampleArr[i][tagKey] ?? ''}
+                                        onChange={value => {
+                                          let taggedSample = { ...taggedSampleArr[i] };
+                                          taggedSample[tagKey] = value;
+                                          setJobConfig(
+                                            objToTags(taggedSample),
+                                            `config.process[0].sample.samples[${i}].prompt`,
+                                          );
+                                        }}
+                                        placeholder={`Enter ${tag.title.toLowerCase()}`}
+                                      />
+                                    )}
+                                    {tag.type === 'number' && (
+                                      <NumberInput
+                                        label={tag.title}
+                                        value={taggedSampleArr[i][tagKey] ?? ''}
+                                        onChange={value => {
+                                          let taggedSample = { ...taggedSampleArr[i] };
+                                          taggedSample[tagKey] = value;
+                                          setJobConfig(
+                                            objToTags(taggedSample),
+                                            `config.process[0].sample.samples[${i}].prompt`,
+                                          );
+                                        }}
+                                        placeholder={`Enter ${tag.title.toLowerCase()}`}
+                                      />
+                                    )}
+                                  </div>
+                                ))}
+                              </div>
+                            ))}
+                          </>
+                        ) : (
+                          <>
+                            {modelArch?.hasMultiLinePrompts ? (
+                              <TextAreaInput
+                                label={`Prompt`}
+                                value={sample.prompt}
+                                onChange={value => setJobConfig(value, `config.process[0].sample.samples[${i}].prompt`)}
+                                placeholder="Enter prompt"
+                                required
+                              />
+                            ) : (
+                              <TextInput
+                                label={`Prompt`}
+                                value={sample.prompt}
+                                onChange={value => setJobConfig(value, `config.process[0].sample.samples[${i}].prompt`)}
+                                placeholder="Enter prompt"
+                                required
+                              />
+                            )}
+                          </>
+                        )}
+
+                        <div className="grid w-full lg:grid-flow-col lg:auto-cols-fr gap-4 mt-2">
+                          {!isAudioModel && (
+                            <TextInput
+                              label={`Width`}
+                              value={sample.width ? `${sample.width}` : ''}
+                              onChange={value => {
+                                // remove any non-numeric characters
+                                value = value.replace(/\D/g, '');
+                                if (value === '') {
+                                  // remove the key from the config if empty
+                                  let newConfig = objectCopy(jobConfig);
+                                  if (newConfig.config.process[0].sample.samples[i]) {
+                                    delete newConfig.config.process[0].sample.samples[i].width;
+                                    setJobConfig(
+                                      newConfig.config.process[0].sample.samples,
+                                      'config.process[0].sample.samples',
+                                    );
+                                  }
                                 } else {
-                                  console.warn('Invalid width value:', value);
+                                  const intValue = parseInt(value);
+                                  if (!isNaN(intValue)) {
+                                    setJobConfig(intValue, `config.process[0].sample.samples[${i}].width`);
+                                  } else {
+                                    console.warn('Invalid width value:', value);
+                                  }
                                 }
-                              }
-                            }}
-                            placeholder={`${jobConfig.config.process[0].sample.width} (default)`}
-                          />
-                          <TextInput
-                            label={`Height`}
-                            value={sample.height ? `${sample.height}` : ''}
-                            onChange={value => {
-                              // remove any non-numeric characters
-                              value = value.replace(/\D/g, '');
-                              if (value === '') {
-                                // remove the key from the config if empty
-                                let newConfig = objectCopy(jobConfig);
-                                if (newConfig.config.process[0].sample.samples[i]) {
-                                  delete newConfig.config.process[0].sample.samples[i].height;
-                                  setJobConfig(
-                                    newConfig.config.process[0].sample.samples,
-                                    'config.process[0].sample.samples',
-                                  );
-                                }
-                              } else {
-                                const intValue = parseInt(value);
-                                if (!isNaN(intValue)) {
-                                  setJobConfig(intValue, `config.process[0].sample.samples[${i}].height`);
+                              }}
+                              placeholder={`${jobConfig.config.process[0].sample.width} (default)`}
+                            />
+                          )}
+                          {!isAudioModel && (
+                            <TextInput
+                              label={`Height`}
+                              value={sample.height ? `${sample.height}` : ''}
+                              onChange={value => {
+                                // remove any non-numeric characters
+                                value = value.replace(/\D/g, '');
+                                if (value === '') {
+                                  // remove the key from the config if empty
+                                  let newConfig = objectCopy(jobConfig);
+                                  if (newConfig.config.process[0].sample.samples[i]) {
+                                    delete newConfig.config.process[0].sample.samples[i].height;
+                                    setJobConfig(
+                                      newConfig.config.process[0].sample.samples,
+                                      'config.process[0].sample.samples',
+                                    );
+                                  }
                                 } else {
-                                  console.warn('Invalid height value:', value);
+                                  const intValue = parseInt(value);
+                                  if (!isNaN(intValue)) {
+                                    setJobConfig(intValue, `config.process[0].sample.samples[${i}].height`);
+                                  } else {
+                                    console.warn('Invalid height value:', value);
+                                  }
                                 }
-                              }
-                            }}
-                            placeholder={`${jobConfig.config.process[0].sample.height} (default)`}
-                          />
+                              }}
+                              placeholder={`${jobConfig.config.process[0].sample.height} (default)`}
+                            />
+                          )}
                           <TextInput
                             label={`Seed`}
                             value={sample.seed ? `${sample.seed}` : ''}

ui/src/app/jobs/new/jobConfig.ts

@@ -1,6 +1,7 @@
 'use client';
 import { isMac } from '@/helpers/basic';
-import { JobConfig, DatasetConfig, SliderConfig } from '@/types';
+import { defaultSampleConfig } from '@/helpers/defaultSamples';
+import { JobConfig, SampleConfig, DatasetConfig, SliderConfig } from '@/types';
 
 export const defaultDatasetConfig: DatasetConfig = {
   folder_path: '/path/to/images/folder',
@@ -108,54 +109,7 @@ export const defaultJobConfig: JobConfig = {
           low_vram: false,
           model_kwargs: {},
         },
-        sample: {
-          sampler: 'flowmatch',
-          sample_every: 250,
-          width: 1024,
-          height: 1024,
-          samples: [
-            {
-              prompt: 'woman with red hair, playing chess at the park, bomb going off in the background',
-            },
-            {
-              prompt: 'a woman holding a coffee cup, in a beanie, sitting at a cafe',
-            },
-            {
-              prompt: 'a horse is a DJ at a night club, fish eye lens, smoke machine, lazer lights, holding a martini',
-            },
-            {
-              prompt:
-                'a man showing off his cool new t shirt at the beach, a shark is jumping out of the water in the background',
-            },
-            {
-              prompt: 'a bear building a log cabin in the snow covered mountains',
-            },
-            {
-              prompt: 'woman playing the guitar, on stage, singing a song, laser lights, punk rocker',
-            },
-            {
-              prompt: 'hipster man with a beard, building a chair, in a wood shop',
-            },
-            {
-              prompt:
-                'photo of a man, white background, medium shot, modeling clothing, studio lighting, white backdrop',
-            },
-            {
-              prompt: "a man holding a sign that says, 'this is a sign'",
-            },
-            {
-              prompt:
-                'a bulldog, in a post apocalyptic world, with a shotgun, in a leather jacket, in a desert, with a motorcycle',
-            },
-          ],
-          neg: '',
-          seed: 42,
-          walk_seed: true,
-          guidance_scale: 4,
-          sample_steps: 25,
-          num_frames: 1,
-          fps: 1,
-        },
+        sample: defaultSampleConfig,
       },
     ],
   },

ui/src/app/jobs/new/options.ts

@@ -1,5 +1,6 @@
 import { GroupedSelectOption, SelectOption, JobConfig } from '@/types';
 import { defaultSliderConfig } from './jobConfig';
+import { defaultAudioSampleConfig, defaultSampleConfig } from '@/helpers/defaultSamples';
 
 type Control = 'depth' | 'line' | 'pose' | 'inpaint';
 
@@ -31,7 +32,17 @@ type AdditionalSections =
   | 'model.qie.match_target_res'
   | 'model.assistant_lora_path';
 
-type ModelGroup = 'image' | 'instruction' | 'video' | 'experimental';
+type ModelGroup = 'image' | 'instruction' | 'video' | 'experimental' | 'audio';
+
+export type SampleTag = {
+  title: string;
+  type: 'text' | 'multiline' | 'number'
+  full?: boolean;
+}
+
+export interface SampleTags {
+  [key: string]: SampleTag;
+}
 
 export interface ModelArch {
   name: string;
@@ -39,10 +50,12 @@ export interface ModelArch {
   group: ModelGroup;
   controls?: Control[];
   isVideoModel?: boolean;
+  hasMultiLinePrompts?: boolean;
   defaults?: { [key: string]: any };
   disableSections?: DisableableSections[];
   additionalSections?: AdditionalSections[];
   accuracyRecoveryAdapters?: { [key: string]: string };
+  sampleTags?: SampleTags;
 }
 
 const defaultNameOrPath = '';
@@ -760,6 +773,116 @@ export const modelArchs: ModelArch[] = [
       'model.qie.match_target_res',
     ],
   },
+  {
+    name: 'ace_step_15_xl',
+    label: 'ACE-Step 1.5 XL',
+    group: 'audio',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_xl_base_aio.safetensors', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].model.low_vram': [true, false],
+      'config.process[0].train.unload_text_encoder': [false, false],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
+      'config.process[0].model.qtype': ['qfloat8', 'qfloat8'],
+      'config.process[0].sample': [defaultAudioSampleConfig, defaultSampleConfig],
+    },
+    sampleTags: {
+      "CAPTION": {
+        title: "Audio Prompt",
+        type: "text",
+        full: true,
+      },
+      "LYRICS": {
+        title: "Lyrics",
+        type: "multiline",
+        full: true,
+      },
+      "BPM": {
+        title: "BPM",
+        type: "number",
+      },
+      "KEYSCALE": {
+        title: "Key Scale",
+        type: "text",
+      },
+      "TIMESIGNATURE": {
+        title: "Time Signature",
+        type: "text",
+      },
+      "DURATION": {
+        title: "Duration (sec)",
+        type: "number",
+      },
+      "LANGUAGE": {
+        title: "Language",
+        type: "text",
+      },
+    },
+    disableSections: ['network.conv'],
+    additionalSections: [
+      'sample.multi_ctrl_imgs',
+      'model.low_vram',
+      'model.layer_offloading',
+    ],
+  },
+  {
+    name: 'ace_step_15',
+    label: 'ACE-Step 1.5',
+    group: 'audio',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_base_aio.safetensors', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].model.low_vram': [true, false],
+      'config.process[0].train.unload_text_encoder': [false, false],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
+      'config.process[0].model.qtype': ['qfloat8', 'qfloat8'],
+      'config.process[0].sample': [defaultAudioSampleConfig, defaultSampleConfig],
+    },
+    sampleTags: {
+      "CAPTION": {
+        title: "Audio Prompt",
+        type: "text",
+        full: true,
+      },
+      "LYRICS": {
+        title: "Lyrics",
+        type: "multiline",
+        full: true,
+      },
+      "BPM": {
+        title: "BPM",
+        type: "number",
+      },
+      "KEYSCALE": {
+        title: "Key Scale",
+        type: "text",
+      },
+      "TIMESIGNATURE": {
+        title: "Time Signature",
+        type: "text",
+      },
+      "DURATION": {
+        title: "Duration (sec)",
+        type: "number",
+      },
+      "LANGUAGE": {
+        title: "Language",
+        type: "text",
+      },
+    },
+    disableSections: ['network.conv'],
+    additionalSections: [
+      'sample.multi_ctrl_imgs',
+      'model.low_vram',
+      'model.layer_offloading',
+    ],
+  },
 ].sort((a, b) => {
   // Sort by label, case-insensitive
   return a.label.localeCompare(b.label, undefined, { sensitivity: 'base' });

ui/src/app/jobs/page.tsx

@@ -9,7 +9,7 @@ export default function Dashboard() {
     <>
       <TopBar>
         <div>
-          <h1 className="text-lg">Training Queue</h1>
+          <h1 className="text-lg">Queue</h1>
         </div>
         <div className="flex-1"></div>
         <div>

ui/src/components/AudioPlayer.tsx

@@ -1,6 +1,7 @@
 'use client';
 
 import React, { useEffect, useMemo, useRef, useState } from 'react';
+import { apiClient } from '@/utils/api';
 
 type AudioPlayerProps = {
   src: string;
@@ -38,214 +39,17 @@ function broadcastExclusivePlay(token: string) {
 }
 
 /**
- * ID3 helpers (v2.2/v2.3/v2.4):
- * - robust album art extraction: APIC (v2.3/2.4) + PIC (v2.2)
- * - basic text frames: title/artist/album
- * - handles tag-level unsynchronisation
- *
- * Requires fetch() byte access; if CORS blocks, it will fall back gracefully.
+ * Build the server-side album-art URL from the audio src.
+ * The audio src is `/api/img/{encodedPath}` — we extract the path
+ * and point to `/api/audio/art/{encodedPath}` instead.
  */
-type Id3Meta = {
-  title?: string;
-  artist?: string;
-  album?: string;
-  pictureUrl?: string; // object URL
-};
-
-function synchsafeToInt(b0: number, b1: number, b2: number, b3: number) {
-  return ((b0 & 0x7f) << 21) | ((b1 & 0x7f) << 14) | ((b2 & 0x7f) << 7) | (b3 & 0x7f);
-}
-
-function deUnsync(bytes: Uint8Array) {
-  const out: number[] = [];
-  for (let i = 0; i < bytes.length; i++) {
-    const cur = bytes[i];
-    out.push(cur);
-    if (cur === 0xff && i + 1 < bytes.length && bytes[i + 1] === 0x00) i += 1;
+function albumArtUrlFromSrc(src: string): string {
+  const prefix = '/api/img/';
+  if (src.startsWith(prefix)) {
+    return `/api/audio/art/${src.slice(prefix.length)}`;
   }
-  return new Uint8Array(out);
-}
-
-function decodeText(encoding: number, bytes: Uint8Array) {
-  let end = bytes.length;
-  while (end > 0 && bytes[end - 1] === 0) end--;
-  const b = bytes.slice(0, end);
-
-  try {
-    if (encoding === 0) return new TextDecoder('latin1').decode(b);
-    if (encoding === 1) return new TextDecoder('utf-16').decode(b);
-    if (encoding === 2) return new TextDecoder('utf-16be').decode(b);
-    if (encoding === 3) return new TextDecoder('utf-8').decode(b);
-  } catch {
-    // ignore
-  }
-  return new TextDecoder('latin1').decode(b);
-}
-
-function readNullTerminated(bytes: Uint8Array, start: number, encoding: number) {
-  if (encoding === 1 || encoding === 2) {
-    let i = start;
-    while (i + 1 < bytes.length && !(bytes[i] === 0 && bytes[i + 1] === 0)) i += 2;
-    const textBytes = bytes.slice(start, i);
-    return { text: decodeText(encoding, textBytes), next: i + 2 };
-  } else {
-    let i = start;
-    while (i < bytes.length && bytes[i] !== 0) i++;
-    const textBytes = bytes.slice(start, i);
-    return { text: decodeText(encoding, textBytes), next: i + 1 };
-  }
-}
-
-async function fetchBytes(src: string, start: number, endInclusive: number) {
-  const wantLen = endInclusive - start + 1;
-
-  try {
-    const r = await fetch(src, { headers: { Range: `bytes=${start}-${endInclusive}` } });
-    if (!r.ok) throw new Error('range not ok');
-    const buf = await r.arrayBuffer();
-    return new Uint8Array(buf);
-  } catch {
-    const r = await fetch(src);
-    if (!r.ok) throw new Error('fetch not ok');
-    const buf = await r.arrayBuffer();
-    const u8 = new Uint8Array(buf);
-    if (start === 0 && u8.length >= wantLen) return u8.slice(0, wantLen);
-    return u8.slice(start, Math.min(u8.length, endInclusive + 1));
-  }
-}
-
-async function extractId3MetaAndArt(src: string, maxTagBytes = 4_000_000): Promise<Id3Meta> {
-  const head = await fetchBytes(src, 0, 64 * 1024 - 1).catch(() => null);
-  if (!head || head.length < 10) return {};
-  if (head[0] !== 0x49 || head[1] !== 0x44 || head[2] !== 0x33) return {};
-
-  const verMajor = head[3]; // 2,3,4
-  const flags = head[5];
-  const tagSize = synchsafeToInt(head[6], head[7], head[8], head[9]);
-  const tagEnd = 10 + tagSize;
-
-  const need = Math.min(tagEnd, maxTagBytes);
-  let tagBytes = head;
-  if (head.length < need) {
-    const more = await fetchBytes(src, 0, need - 1).catch(() => null);
-    if (!more) return {};
-    tagBytes = more;
-  } else {
-    tagBytes = head.slice(0, need);
-  }
-
-  const tagUnsync = (flags & 0x80) !== 0;
-  const tagDataRaw = tagBytes.slice(10, Math.min(tagBytes.length, tagEnd));
-  const tagData = tagUnsync ? deUnsync(tagDataRaw) : tagDataRaw;
-
-  let offset = 0;
-
-  // Extended header (v2.3/v2.4)
-  if (verMajor === 3 || verMajor === 4) {
-    const hasExt = (flags & 0x40) !== 0;
-    if (hasExt && tagData.length >= 4) {
-      let extSize = 0;
-      if (verMajor === 4) extSize = synchsafeToInt(tagData[0], tagData[1], tagData[2], tagData[3]);
-      else extSize = (tagData[0] << 24) | (tagData[1] << 16) | (tagData[2] << 8) | tagData[3];
-      offset += 4 + Math.max(0, extSize);
-    }
-  }
-
-  const meta: Id3Meta = {};
-  const setIfEmpty = (k: keyof Id3Meta, v?: string) => {
-    if (!v) return;
-    if (!meta[k]) meta[k] = v;
-  };
-
-  while (offset < tagData.length) {
-    if (tagData[offset] === 0x00) break;
-
-    if (verMajor === 2) {
-      if (offset + 6 > tagData.length) break;
-      const id = new TextDecoder('latin1').decode(tagData.slice(offset, offset + 3));
-      const size = (tagData[offset + 3] << 16) | (tagData[offset + 4] << 8) | tagData[offset + 5];
-      offset += 6;
-      if (!id.trim() || size <= 0 || offset + size > tagData.length) break;
-
-      const frame = tagData.slice(offset, offset + size);
-
-      if (id === 'TT2' || id === 'TP1' || id === 'TAL') {
-        const enc = frame[0];
-        const txt = decodeText(enc, frame.slice(1));
-        if (id === 'TT2') setIfEmpty('title', txt);
-        if (id === 'TP1') setIfEmpty('artist', txt);
-        if (id === 'TAL') setIfEmpty('album', txt);
-      }
-
-      if (id === 'PIC' && frame.length > 6) {
-        const enc = frame[0];
-        const fmt = new TextDecoder('latin1').decode(frame.slice(1, 4)).toLowerCase();
-        const imgType = fmt === 'png' ? 'image/png' : 'image/jpeg';
-        let p = 5;
-        const desc = readNullTerminated(frame, p, enc);
-        p = desc.next;
-        if (p < frame.length) {
-          const img = frame.slice(p);
-          if (img.length > 64) {
-            const blob = new Blob([img], { type: imgType });
-            meta.pictureUrl = URL.createObjectURL(blob);
-          }
-        }
-      }
-
-      offset += size;
-    } else {
-      if (offset + 10 > tagData.length) break;
-
-      const id = new TextDecoder('latin1').decode(tagData.slice(offset, offset + 4));
-      let size = 0;
-      if (verMajor === 4)
-        size = synchsafeToInt(tagData[offset + 4], tagData[offset + 5], tagData[offset + 6], tagData[offset + 7]);
-      else
-        size =
-          (tagData[offset + 4] << 24) | (tagData[offset + 5] << 16) | (tagData[offset + 6] << 8) | tagData[offset + 7];
-
-      const flag2 = tagData[offset + 9];
-      offset += 10;
-
-      if (!id.trim() || size <= 0 || offset + size > tagData.length) break;
-
-      let frame = tagData.slice(offset, offset + size);
-
-      const frameUnsync = verMajor === 4 && (flag2 & 0x02) !== 0;
-      if (frameUnsync) frame = deUnsync(frame);
-
-      if (id === 'TIT2' || id === 'TPE1' || id === 'TALB') {
-        const enc = frame[0];
-        const txt = decodeText(enc, frame.slice(1));
-        if (id === 'TIT2') setIfEmpty('title', txt);
-        if (id === 'TPE1') setIfEmpty('artist', txt);
-        if (id === 'TALB') setIfEmpty('album', txt);
-      }
-
-      if (id === 'APIC' && frame.length > 10) {
-        const enc = frame[0];
-        const mimeZ = readNullTerminated(frame, 1, 0);
-        const mime = mimeZ.text || 'image/jpeg';
-        let p = mimeZ.next;
-        if (p < frame.length) p += 1; // pictureType
-        const desc = readNullTerminated(frame, p, enc);
-        p = desc.next;
-
-        if (p < frame.length) {
-          const img = frame.slice(p);
-          if (img.length > 64) {
-            const blob = new Blob([img], { type: mime });
-            if (!meta.pictureUrl) meta.pictureUrl = URL.createObjectURL(blob);
-          }
-        }
-      }
-
-      offset += size;
-    }
-  }
-
-  return meta;
+  // Fallback: assume src is already an encoded path
+  return `/api/audio/art/${encodeURIComponent(src)}`;
 }
 
 export default function AudioPlayer({
@@ -278,12 +82,8 @@ export default function AudioPlayer({
   const [dragging, setDragging] = useState(false);
   const [dragValue, setDragValue] = useState(0);
 
-  // Meta + artwork
-  const [metaTitle, setMetaTitle] = useState<string | null>(null);
-  const [metaArtist, setMetaArtist] = useState<string | null>(null);
-  const [metaAlbum, setMetaAlbum] = useState<string | null>(null);
+  // Album art: served by /api/audio/art endpoint (fast, server-side extraction)
   const [albumArtUrl, setAlbumArtUrl] = useState<string | null>(null);
-  const albumArtBlobUrlRef = useRef<string | null>(null);
 
   // WebAudio analyser
   const audioCtxRef = useRef<AudioContext | null>(null);
@@ -318,57 +118,23 @@ export default function AudioPlayer({
     c.height = Math.max(1, Math.floor(cssH * dpr));
   }, [size]);
 
-  // Extract meta + art
+  // Set album art URL from server endpoint
   useEffect(() => {
+    const artUrl = albumArtUrlFromSrc(src);
     let cancelled = false;
-
-    (async () => {
-      if (albumArtBlobUrlRef.current) {
-        URL.revokeObjectURL(albumArtBlobUrlRef.current);
-        albumArtBlobUrlRef.current = null;
-      }
-      setAlbumArtUrl(null);
-      setMetaTitle(null);
-      setMetaArtist(null);
-      setMetaAlbum(null);
-
-      try {
-        const meta = await extractId3MetaAndArt(src);
-        if (cancelled) return;
-
-        if (meta.title) setMetaTitle(meta.title);
-        if (meta.artist) setMetaArtist(meta.artist);
-        if (meta.album) setMetaAlbum(meta.album);
-
-        if (meta.pictureUrl) {
-          albumArtBlobUrlRef.current = meta.pictureUrl;
-          setAlbumArtUrl(meta.pictureUrl);
-        } else if (defaultAlbumArtUrl) {
-          setAlbumArtUrl(defaultAlbumArtUrl);
-        }
-      } catch {
-        if (!cancelled && defaultAlbumArtUrl) setAlbumArtUrl(defaultAlbumArtUrl);
-      }
-    })();
-
+    apiClient
+      .head(artUrl)
+      .then(() => {
+        if (!cancelled) setAlbumArtUrl(artUrl);
+      })
+      .catch(() => {
+        if (!cancelled) setAlbumArtUrl(defaultAlbumArtUrl ?? null);
+      });
     return () => {
       cancelled = true;
     };
   }, [src, defaultAlbumArtUrl]);
 
-  // Cleanup artwork blob URL on unmount
-  useEffect(() => {
-    return () => {
-      if (albumArtBlobUrlRef.current) {
-        URL.revokeObjectURL(albumArtBlobUrlRef.current);
-        albumArtBlobUrlRef.current = null;
-      }
-    };
-  }, []);
-
-  const effectiveTitle = metaTitle || title;
-  const effectiveSubtitle = metaArtist || subtitle || metaAlbum || '';
-
   const progress = useMemo(() => {
     const cur = dragging ? dragValue : t;
     if (!duration) return 0;
@@ -677,11 +443,11 @@ export default function AudioPlayer({
             className="truncate text-gray-200"
             style={{ fontSize: titleSize, lineHeight: 1.1, letterSpacing: '0.01em' }}
           >
-            {effectiveTitle}
+            {title}
           </div>
-          {effectiveSubtitle ? (
+          {subtitle ? (
             <div className="mt-1 truncate text-gray-400" style={{ fontSize: subSize, lineHeight: 1.15 }}>
-              {effectiveSubtitle}
+              {subtitle}
             </div>
           ) : null}
           {err ? (

ui/src/components/AutoCaptionButton.tsx

@@ -0,0 +1,39 @@
+import React, {use, useEffect} from 'react';
+import { Button } from '@headlessui/react';
+import { CaptionDatasetModal, openCaptionDatasetModal } from '@/components/CaptionDatasetModal';
+import useJobByRef from '@/hooks/useJobByRef';
+import Link from 'next/link';
+import { Loader2 } from 'lucide-react';
+
+type AutoCaptionButtonProps = {
+  datasetPath: string;
+  setIsAutoCaptioning?: (isAutoCaptioning: boolean) => void;
+};
+
+export default function AutoCaptionButton({ datasetPath, setIsAutoCaptioning }: AutoCaptionButtonProps) {
+  const { job, status, refreshJob } = useJobByRef(datasetPath, 5000);
+  useEffect(() => {
+    if (setIsAutoCaptioning) {
+      setIsAutoCaptioning(!!(job && (job.status === 'running')));
+    }
+  }, [job, setIsAutoCaptioning]);
+  
+  if (job && (job.status === 'running' || job.status === 'queued')) {
+    return (
+      <Link href={`/jobs/${job.id}`} className="text-white bg-gray-400 px-3 py-1 rounded-md mr-2 inline-flex items-center gap-1.5">
+        <Loader2 className="w-4 h-4 animate-spin" />
+        Auto Captioning...
+      </Link>
+    );
+  }
+  return (
+    <Button
+      className="text-white bg-blue-600 px-3 py-1 rounded-md mr-2"
+      onClick={() => openCaptionDatasetModal(datasetPath, () => {
+        refreshJob();
+      })}
+    >
+      Auto Caption
+    </Button>
+  );
+}

ui/src/components/CaptionDatasetModal.tsx

@@ -0,0 +1,273 @@
+import React, { useState, useEffect, useRef } from 'react';
+import { Modal } from '@/components/Modal';
+import { createGlobalState } from 'react-global-hooks';
+import { useFromNull } from '@/hooks/useFromNull';
+import {
+  Checkbox,
+  CreatableSelectInput,
+  FormGroup,
+  SelectInput,
+  TextAreaInput,
+  TextInput,
+} from '@/components/formInputs';
+import { CaptionJobConfig } from '@/types';
+import { defaultCaptionJobConfig, handleCaptionerTypeChange } from '@/helpers/captionJobConfig';
+import { objectCopy } from '@/utils/basic';
+import { useNestedState } from '@/utils/hooks';
+import {
+  captionerTypes,
+  defaultQtype,
+  groupedCaptionerTypes,
+  maxNewTokensOptions,
+  maxResOptions,
+  quantizationOptions,
+} from '@/helpers/captionOptions';
+import { isMac } from '@/helpers/basic';
+import useGPUInfo from '@/hooks/useGPUInfo';
+import { apiClient } from '@/utils/api';
+import { v4 as uuidv4 } from 'uuid';
+import { startJob } from '@/utils/jobs';
+import { startQueue } from '@/utils/queue';
+
+export interface CaptionDatasetModalState {
+  datasetPath: string;
+  onClose?: () => void;
+}
+
+export const captionDatasetModalState = createGlobalState<CaptionDatasetModalState | null>(null);
+
+export const openCaptionDatasetModal = (datasetPath: string, onClose?: () => void) => {
+  captionDatasetModalState.set({ datasetPath, onClose });
+};
+
+export const CaptionDatasetModal: React.FC = () => {
+  const [modalInfo, setModalInfo] = captionDatasetModalState.use();
+  const [jobConfig, setJobConfig] = useNestedState<CaptionJobConfig>(objectCopy(defaultCaptionJobConfig));
+  const [gpuIDs, setGpuIDs] = useState<string | null>(null);
+  const { gpuList, isGPUInfoLoaded } = useGPUInfo();
+  const open = modalInfo !== null;
+  const isSavingRef = useRef(false);
+  const showGPUSelect = !isMac();
+
+  useFromNull(() => {
+    // reset the state
+    setJobConfig(objectCopy(defaultCaptionJobConfig));
+    // set the path_to_caption
+    if (modalInfo?.datasetPath) {
+      setJobConfig(modalInfo.datasetPath, 'config.process[0].caption.path_to_caption');
+    }
+  }, [modalInfo]);
+
+  useEffect(() => {
+    if (isGPUInfoLoaded) {
+      if (gpuIDs === null && gpuList.length > 0) {
+        setGpuIDs(`${gpuList[0].index}`);
+      }
+    }
+  }, [gpuList, isGPUInfoLoaded]);
+
+  const handleClose = () => {
+    if (modalInfo?.onClose) {
+      modalInfo.onClose();
+    }
+    setModalInfo(null);
+  };
+
+  const selectedCaptionOption = captionerTypes.find(option => option.name === jobConfig.config.process[0].type);
+
+  const saveJob = async () => {
+    if (isSavingRef.current) return;
+    if (!modalInfo?.datasetPath) {
+      alert('Dataset path is missing. Please try again.');
+      return;
+    }
+    isSavingRef.current = true;
+
+    apiClient
+      .post('/api/jobs', {
+        id: null,
+        name: uuidv4(),
+        gpu_ids: gpuIDs,
+        job_config: jobConfig,
+        job_type: 'caption',
+        job_ref: modalInfo.datasetPath,
+      })
+      .then(async res => {
+        const jobId = res.data.id;
+        await startJob(jobId);
+        // start the queue as well
+        await startQueue(gpuIDs || '');
+        isSavingRef.current = false;
+        handleClose();
+      })
+      .catch(error => {
+        if (error.response?.status === 409) {
+          alert('A caption job for this dataset already exists. Please check your jobs list.');
+        } else {
+          alert('Failed to save job. Please try again.');
+        }
+        console.log('Error saving training:', error);
+        isSavingRef.current = false;
+      });
+  };
+
+  return (
+    <Modal isOpen={open} onClose={handleClose} title="Caption Dataset" size="lg">
+      <div className="space-y-4 text-gray-200">
+        <form
+          onSubmit={e => {
+            e.preventDefault();
+            saveJob();
+          }}
+        >
+          <div className="text-sm text-gray-400">
+            <div className="grid grid-cols-1 md:grid-cols-2 gap-4 mt-4">
+              <div>
+                <SelectInput
+                  label="Captioner Type"
+                  value={jobConfig.config.process[0].type}
+                  onChange={value => {
+                    handleCaptionerTypeChange(jobConfig.config.process[0].type, value, jobConfig, setJobConfig);
+                  }}
+                  options={groupedCaptionerTypes}
+                />
+              </div>
+              {showGPUSelect && (
+                <div>
+                  <SelectInput
+                    label="GPU ID"
+                    value={`${gpuIDs}`}
+                    onChange={value => setGpuIDs(value)}
+                    options={gpuList.map((gpu: any) => ({ value: `${gpu.index}`, label: `GPU #${gpu.index}` }))}
+                  />
+                </div>
+              )}
+            </div>
+            <div className="mt-4">
+              <CreatableSelectInput
+                label="Name or Path"
+                value={jobConfig.config.process[0].caption.model_name_or_path}
+                docKey="config.process[0].caption.model_name_or_path"
+                onChange={(value: string | null) => {
+                  if (value?.trim() === '') {
+                    value = null;
+                  }
+                  setJobConfig(value, 'config.process[0].caption.model_name_or_path');
+                }}
+                placeholder=""
+                options={selectedCaptionOption?.name_or_path_options || []}
+                required
+              />
+            </div>
+            {selectedCaptionOption?.additionalSections?.includes('caption.model_name_or_path2') && (
+              <div className="mt-4">
+                <CreatableSelectInput
+                  label="Name or Path 2"
+                  value={jobConfig.config.process[0].caption.model_name_or_path2 || ''}
+                  onChange={(value: string | null) => {
+                    if (value?.trim() === '') {
+                      value = null;
+                    }
+                    setJobConfig(value, 'config.process[0].caption.model_name_or_path2');
+                  }}
+                  placeholder=""
+                  options={selectedCaptionOption?.name_or_path2_options || []}
+                />
+              </div>
+            )}
+            <div className="grid grid-cols-1 md:grid-cols-2 gap-4 mt-4">
+              <div>
+                <SelectInput
+                  label="Quantize"
+                  value={jobConfig.config.process[0].caption.quantize ? jobConfig.config.process[0].caption.qtype : ''}
+                  onChange={value => {
+                    if (value === '') {
+                      setJobConfig(false, 'config.process[0].caption.quantize');
+                      value = defaultQtype;
+                    } else {
+                      setJobConfig(true, 'config.process[0].caption.quantize');
+                    }
+                    setJobConfig(value, 'config.process[0].caption.qtype');
+                  }}
+                  options={quantizationOptions}
+                />
+                {selectedCaptionOption?.additionalSections?.includes('caption.max_res') && (
+                  <div className="mt-4">
+                    <SelectInput
+                      label="Max Resolution"
+                      value={`${jobConfig.config.process[0].caption.max_res || ''}`}
+                      onChange={value => {
+                        const intVal = parseInt(value);
+                        if (!isNaN(intVal)) {
+                          setJobConfig(intVal, 'config.process[0].caption.max_res');
+                        }
+                      }}
+                      options={maxResOptions}
+                    />
+                  </div>
+                )}
+                {selectedCaptionOption?.additionalSections?.includes('caption.max_new_tokens') && (
+                  <div className="mt-4">
+                    <SelectInput
+                      label="Max New Tokens"
+                      value={`${jobConfig.config.process[0].caption.max_new_tokens || ''}`}
+                      onChange={value => {
+                        const intVal = parseInt(value);
+                        if (!isNaN(intVal)) {
+                          setJobConfig(intVal, 'config.process[0].caption.max_new_tokens');
+                        }
+                      }}
+                      options={maxNewTokensOptions}
+                    />
+                  </div>
+                )}
+              </div>
+              <div>
+                <FormGroup label="Options">
+                  <Checkbox
+                    label="Low VRAM"
+                    checked={jobConfig.config.process[0].caption.low_vram}
+                    onChange={value => setJobConfig(value, 'config.process[0].caption.low_vram')}
+                  />
+                  <Checkbox
+                    label="Recaption"
+                    checked={jobConfig.config.process[0].caption.recaption}
+                    onChange={value => setJobConfig(value, 'config.process[0].caption.recaption')}
+                  />
+                </FormGroup>
+              </div>
+            </div>
+            {selectedCaptionOption?.additionalSections?.includes('caption.caption_prompt') && (
+              <div className="mt-4">
+                <TextAreaInput
+                  label="Caption Prompt"
+                  value={jobConfig.config.process[0].caption.caption_prompt || ''}
+                  onChange={value => {
+                    setJobConfig(value, 'config.process[0].caption.caption_prompt');
+                  }}
+                  placeholder="Enter caption prompt"
+                />
+              </div>
+            )}
+          </div>
+
+          <div className="mt-6 flex justify-end space-x-3">
+            <button
+              type="button"
+              className="rounded-md bg-gray-700 px-4 py-2 text-gray-200 hover:bg-gray-600 focus:outline-none focus:ring-2 focus:ring-gray-500"
+              onClick={handleClose}
+            >
+              Cancel
+            </button>
+            <button
+              type="submit"
+              className="rounded-md bg-blue-600 px-4 py-2 text-white hover:bg-blue-700 focus:outline-none focus:ring-2 focus:ring-blue-500"
+            >
+              Add to Queue
+            </button>
+          </div>
+        </form>
+      </div>
+    </Modal>
+  );
+};

ui/src/components/DatasetImageCard.tsx

@@ -9,6 +9,7 @@ import { isVideo, isAudio } from '@/utils/basic';
 interface DatasetImageCardProps {
   imageUrl: string;
   alt: string;
+  isAutoCaptioning: boolean;
   children?: ReactNode;
   className?: string;
   onDelete?: () => void;
@@ -17,6 +18,7 @@ interface DatasetImageCardProps {
 const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
   imageUrl,
   alt,
+  isAutoCaptioning,
   children,
   className = '',
   onDelete = () => {},
@@ -80,6 +82,16 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
     }
   }, [inViewport, isVisible]);
 
+  // Poll for caption updates every 5 seconds while auto-captioning
+  useEffect(() => {
+    if (!isAutoCaptioning || !inViewport || !isVisible) return;
+    const interval = setInterval(() => {
+      // Reset so fetchCaption will re-fetch
+      setIsCaptionLoaded(false);
+    }, 5000);
+    return () => clearInterval(interval);
+  }, [isAutoCaptioning, inViewport, isVisible]);
+
   useEffect(() => {
     // Create intersection observer to check viewport visibility
     const observer = new IntersectionObserver(
@@ -129,6 +141,8 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
 
   const isCaptionCurrent = caption.trim() === savedCaption;
 
+  const [showAudioPlayer, setShowAudioPlayer] = useState(true);
+
   const isItAVideo = isVideo(imageUrl);
   const isItAudio = isAudio(imageUrl);
   const isItImage = !isItAVideo && !isItAudio;
@@ -154,7 +168,22 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
                   controls
                 />
               )}
-              {isItAudio && (
+              {isItAudio && !showAudioPlayer && (
+                <div
+                  className="w-full h-full cursor-pointer flex items-center justify-center bg-gray-900"
+                  onClick={() => setShowAudioPlayer(true)}
+                >
+                  <img
+                    src={`/api/audio/art/${encodeURIComponent(imageUrl)}`}
+                    alt={alt}
+                    className="w-full h-full object-contain"
+                    onError={e => {
+                      (e.target as HTMLImageElement).style.display = 'none';
+                    }}
+                  />
+                </div>
+              )}
+              {isItAudio && showAudioPlayer && (
                 <AudioPlayer
                   src={`/api/img/${encodeURIComponent(imageUrl)}`}
                   title={imageUrl.replace(/^.*[\\/]/, '')}
@@ -221,9 +250,12 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
             onBlur={saveCaption}
           >
             <textarea
-              className="w-full bg-transparent resize-none outline-none focus:ring-0 focus:outline-none"
+              className={classNames("w-full bg-transparent resize-none outline-none focus:ring-0 focus:outline-none", {
+                'opacity-50 cursor-not-allowed': isAutoCaptioning,
+              })}
               value={caption}
               rows={3}
+              readOnly={isAutoCaptioning}
               onChange={e => setCaption(e.target.value)}
               onKeyDown={handleKeyDown}
             />

ui/src/components/JobActionBar.tsx

@@ -123,11 +123,16 @@ export default function JobActionBar({
           <Cog />
         </MenuButton>
         <MenuItems anchor="bottom" className="bg-gray-900 border border-gray-700 rounded shadow-lg w-48 px-2 py-2 mt-4">
-          <MenuItem>
-            <Link href={`/jobs/new?cloneId=${job.id}`} className="cursor-pointer px-4 py-1 hover:bg-gray-800 rounded block">
-              Clone Job
-            </Link>
-          </MenuItem>
+          {job.job_type === 'train' && (
+            <MenuItem>
+              <Link
+                href={`/jobs/new?cloneId=${job.id}`}
+                className="cursor-pointer px-4 py-1 hover:bg-gray-800 rounded block"
+              >
+                Clone Job
+              </Link>
+            </MenuItem>
+          )}
           <MenuItem>
             <div
               className="cursor-pointer px-4 py-1 hover:bg-gray-800 rounded"

ui/src/components/JobConfigViewer.tsx

@@ -21,11 +21,23 @@ const yamlConfig: YAML.DocumentOptions &
   directives: true,
 };
 
+function toYaml(obj: any): string {
+  const doc = new YAML.Document(obj, yamlConfig);
+  YAML.visit(doc, {
+    Scalar(_key, node) {
+      if (typeof node.value === 'string' && node.value.includes('\n')) {
+        node.type = YAML.Scalar.BLOCK_LITERAL;
+      }
+    },
+  });
+  return doc.toString(yamlConfig);
+}
+
 export default function JobConfigViewer({ job }: Props) {
   const [editorValue, setEditorValue] = useState<string>('');
   useEffect(() => {
     if (job?.job_config) {
-      const yamlContent = YAML.stringify(JSON.parse(job.job_config), yamlConfig);
+      const yamlContent = toYaml(JSON.parse(job.job_config));
       setEditorValue(yamlContent);
     }
   }, [job]);

ui/src/components/JobOverview.tsx

@@ -82,6 +82,8 @@ export default function JobOverview({ job }: JobOverviewProps) {
     }
   };
 
+  const jobType = job?.job_type || 'unknown';
+
   let status = job.status;
   if (isStopping) {
     status = 'stopping';
@@ -100,17 +102,19 @@ export default function JobOverview({ job }: JobOverviewProps) {
 
         <div className="p-4 space-y-6 flex flex-col flex-grow">
           {/* Progress Bar */}
-          <div className="space-y-2">
-            <div className="flex items-center justify-between text-sm">
-              <span className="text-gray-400">Progress</span>
-              <span className="text-gray-200">
-                Step {job.step} of {totalSteps}
-              </span>
+          {job.job_type === 'train' && (
+            <div className="space-y-2">
+              <div className="flex items-center justify-between text-sm">
+                <span className="text-gray-400">Progress</span>
+                <span className="text-gray-200">
+                  Step {job.step} of {totalSteps}
+                </span>
+              </div>
+              <div className="w-full bg-gray-800 rounded-full h-2">
+                <div className="h-2 rounded-full bg-blue-500 transition-all" style={{ width: `${progress}%` }} />
+              </div>
             </div>
-            <div className="w-full bg-gray-800 rounded-full h-2">
-              <div className="h-2 rounded-full bg-blue-500 transition-all" style={{ width: `${progress}%` }} />
-            </div>
-          </div>
+          )}
 
           {/* Job Info Grid */}
           <div className="grid gap-4 grid-cols-1 md:grid-cols-3">
@@ -164,9 +168,11 @@ export default function JobOverview({ job }: JobOverviewProps) {
       <div className="col-span-1">
         <div>{isCPUInfoLoaded && cpuInfo && <CPUWidget cpu={cpuInfo} />}</div>
         <div className="mt-4">{isGPUInfoLoaded && gpuList.length > 0 && <GPUWidget gpu={gpuList[0]} />}</div>
-        <div className="mt-4">
-          <FilesWidget jobID={job.id} />
-        </div>
+        {jobType === 'train' && (
+          <div className="mt-4">
+            <FilesWidget jobID={job.id} />
+          </div>
+        )}
       </div>
     </div>
   );

ui/src/components/JobsTable.tsx

@@ -14,10 +14,11 @@ import useGPUInfo from '@/hooks/useGPUInfo';
 interface JobsTableProps {
   autoStartQueue?: boolean;
   onlyActive?: boolean;
+  job_type?: string | null;
 }
 
-export default function JobsTable({ onlyActive = false }: JobsTableProps) {
-  const { jobs, status, refreshJobs } = useJobsList(onlyActive, 5000);
+export default function JobsTable({ onlyActive = false, job_type = null }: JobsTableProps) {
+  const { jobs, status, refreshJobs } = useJobsList({ onlyActive, reloadInterval: 5000, job_type });
   const { queues, status: queueStatus, refreshQueues } = useQueueList();
   const { gpuList, isGPUInfoLoaded } = useGPUInfo();
 
@@ -30,21 +31,37 @@ export default function JobsTable({ onlyActive = false }: JobsTableProps) {
     {
       title: 'Name',
       key: 'name',
-      render: row => (
-        <Link href={`/jobs/${row.id}`} className="font-medium whitespace-nowrap">
-          {['running', 'stopping'].includes(row.status) ? (
-            <CgSpinner className="inline animate-spin mr-2 text-blue-400" />
-          ) : null}
-          {row.name}
-        </Link>
-      ),
+      render: row => {
+        let title = row.name;
+        // if (row.job_type === 'train') title = `Train: ${title}`;
+        if (row.job_type === 'caption') {
+          let splits = row.job_ref.split(/[/\\]/);
+          const datasetPath = `${splits[splits.length - 1]}`;
+          title = (
+            <>
+              <small className="opacity-50">CAPTION: </small> {datasetPath}
+            </>
+          );
+        }
+        return (
+          <Link href={`/jobs/${row.id}`} className="font-medium whitespace-nowrap">
+            {['running', 'stopping'].includes(row.status) ? (
+              <CgSpinner className="inline animate-spin mr-2 text-blue-400" />
+            ) : null}
+            {title}
+          </Link>
+        );
+      },
     },
     {
       title: 'Steps',
       key: 'steps',
       render: row => {
         const jobConfig: JobConfig = JSON.parse(row.job_config);
-        const totalSteps = jobConfig.config.process[0].train.steps;
+        if (row.job_type !== 'train') {
+          return <></>;
+        }
+        const totalSteps = jobConfig.config.process[0].train?.steps;
 
         return (
           <div>

ui/src/components/SampleImageCard.tsx

@@ -1,5 +1,5 @@
 import React, { useRef, useEffect, useState, ReactNode } from 'react';
-import { isVideo } from '@/utils/basic';
+import { isVideo, isAudio } from '@/utils/basic';
 
 interface SampleImageCardProps {
   imageUrl: string;
@@ -63,7 +63,18 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({
       <div ref={cardRef} className="relative w-full cursor-pointer" style={{ paddingBottom: '100%' }} onClick={onClick}>
         <div className="absolute inset-0 rounded-t-lg shadow-md">
           {isVisible ? (
-            isVideo(imageUrl) ? (
+            isAudio(imageUrl) ? (
+              <div className="w-full h-full flex items-center justify-center bg-gray-900">
+                <img
+                  src={`/api/audio/art/${encodeURIComponent(imageUrl)}`}
+                  alt={alt}
+                  className="w-full h-full object-cover"
+                  onError={e => {
+                    (e.target as HTMLImageElement).style.display = 'none';
+                  }}
+                />
+              </div>
+            ) : isVideo(imageUrl) ? (
               <video
                 ref={videoRef}
                 src={`/api/img/${encodeURIComponent(imageUrl)}`}

ui/src/components/SampleImageViewer.tsx

@@ -7,7 +7,8 @@ import { Cog } from 'lucide-react';
 import { Menu, MenuButton, MenuItem, MenuItems } from '@headlessui/react';
 import { openConfirm } from './ConfirmModal';
 import { apiClient } from '@/utils/api';
-import { isVideo } from '@/utils/basic';
+import { isVideo, isAudio } from '@/utils/basic';
+import AudioPlayer from './AudioPlayer';
 
 interface Props {
   imgPath: string | null; // current image path
@@ -214,7 +215,15 @@ export default function SampleImageViewer({
           >
             <div className="overflow-hidden flex items-center justify-center">
               {displayedImgPath &&
-                (isVideo(displayedImgPath) ? (
+                (isAudio(displayedImgPath) ? (
+                  <div className="w-[500px] h-[500px] max-w-[95vw] max-h-[82vh]">
+                    <AudioPlayer
+                      src={`/api/img/${encodeURIComponent(displayedImgPath)}`}
+                      title={displayedImgPath.replace(/^.*[\\/]/, '')}
+                      autoPlay
+                    />
+                  </div>
+                ) : isVideo(displayedImgPath) ? (
                   <video
                     src={`/api/img/${encodeURIComponent(displayedImgPath)}`}
                     className="w-auto h-auto max-w-[95vw] max-h-[82vh] object-contain"
@@ -282,7 +291,7 @@ export default function SampleImageViewer({
                 </div>
               </div>
             </div>
-            <div className="absolute top-2 right-2 bg-gray-900 rounded-full p-1 leading-[0px] opacity-50 hover:opacity-100">
+            <div className="absolute top-2 right-2 bg-gray-900 rounded-full p-1 leading-[0px] opacity-50 hover:opacity-100 z-20">
               <Menu>
                 <MenuButton>
                   <Cog />

ui/src/components/Sidebar.tsx

@@ -8,7 +8,7 @@ const Sidebar = () => {
   const navigation = [
     { name: 'Dashboard', href: '/dashboard', icon: Home },
     { name: 'New Job', href: '/jobs/new', icon: Plus },
-    { name: 'Training Queue', href: '/jobs', icon: BrainCircuit },
+    { name: 'Queue', href: '/jobs', icon: BrainCircuit },
     { name: 'Datasets', href: '/datasets', icon: Images },
     { name: 'Settings', href: '/settings', icon: Settings },
   ];

ui/src/components/formInputs.tsx

@@ -67,6 +67,49 @@ export const TextInput = forwardRef<HTMLInputElement, TextInputProps>((props: Te
 // 👇 Helpful for debugging
 TextInput.displayName = 'TextInput';
 
+export interface TextAreaInputProps extends InputProps {
+  value: string;
+  onChange: (value: string) => void;
+  disabled?: boolean;
+  rows?: number;
+}
+
+export const TextAreaInput = forwardRef<HTMLTextAreaElement, TextAreaInputProps>((props: TextAreaInputProps, ref) => {
+  const { label, value, onChange, placeholder, required, disabled, rows = 4, className, docKey = null } = props;
+  let { doc } = props;
+  if (!doc && docKey) {
+    doc = getDoc(docKey);
+  }
+  return (
+    <div className={classNames(className)}>
+      {label && (
+        <label className={labelClasses}>
+          {label}{' '}
+          {doc && (
+            <div className="inline-block ml-1 text-xs text-gray-500 cursor-pointer" onClick={() => openDoc(doc)}>
+              <CircleHelp className="inline-block w-4 h-4 cursor-pointer" />
+            </div>
+          )}
+        </label>
+      )}
+      <textarea
+        ref={ref}
+        value={value}
+        onChange={e => {
+          if (!disabled) onChange(e.target.value);
+        }}
+        className={`${inputClasses} ${disabled ? 'opacity-30 cursor-not-allowed' : ''}`}
+        placeholder={placeholder}
+        required={required}
+        disabled={disabled}
+        rows={rows}
+      />
+    </div>
+  );
+});
+
+TextAreaInput.displayName = 'TextAreaInput';
+
 export interface NumberInputProps extends InputProps {
   value: number | null;
   onChange: (value: number | null) => void;
@@ -200,6 +243,115 @@ export const SelectInput = (props: SelectInputProps) => {
   );
 };
 
+export interface CreatableSelectInputProps extends InputProps {
+  value: string;
+  disabled?: boolean;
+  onChange: (value: string) => void;
+  options: GroupedSelectOption[] | SelectOption[];
+}
+
+const CUSTOM_SELECT_VALUE = '__custom__';
+
+export const CreatableSelectInput = (props: CreatableSelectInputProps) => {
+  const { label, value, onChange, options, docKey = null } = props;
+  let { doc } = props;
+  if (!doc && docKey) {
+    doc = getDoc(docKey);
+  }
+
+  // Check if current value matches any predefined option
+  let isInOptions = false;
+  if (options && options.length > 0) {
+    if ('options' in options[0]) {
+      isInOptions = (options as GroupedSelectOption[]).flatMap(g => g.options).some(opt => opt.value === value);
+    } else {
+      isInOptions = (options as SelectOption[]).some(opt => opt.value === value);
+    }
+  }
+
+  const [isCustom, setIsCustom] = React.useState(!isInOptions && !!value);
+
+  // Build select options with "Custom" at the top
+  const customOption: SelectOption = { value: CUSTOM_SELECT_VALUE, label: 'Custom' };
+  const selectOptions = React.useMemo(() => {
+    if (options && options.length > 0 && 'options' in options[0]) {
+      return [{ label: '', options: [customOption] }, ...(options as GroupedSelectOption[])];
+    }
+    return [customOption, ...(options as SelectOption[])];
+  }, [options]);
+
+  const selectedOption = isCustom
+    ? customOption
+    : (() => {
+        if (!options || options.length === 0) return undefined;
+        if ('options' in options[0]) {
+          return (options as GroupedSelectOption[]).flatMap(g => g.options).find(opt => opt.value === value);
+        }
+        return (options as SelectOption[]).find(opt => opt.value === value);
+      })();
+
+  return (
+    <div
+      className={classNames(props.className, {
+        'opacity-30 cursor-not-allowed': props.disabled,
+      })}
+    >
+      {label && (
+        <label className={labelClasses}>
+          {label}{' '}
+          {doc && (
+            <div className="inline-block ml-1 text-xs text-gray-500 cursor-pointer" onClick={() => openDoc(doc)}>
+              <CircleHelp className="inline-block w-4 h-4 cursor-pointer" />
+            </div>
+          )}
+        </label>
+      )}
+      <div className="flex gap-2">
+        <div className={isCustom ? 'w-1/3' : 'w-full'}>
+          <Select
+            value={selectedOption}
+            options={selectOptions}
+            isDisabled={props.disabled}
+            className="aitk-react-select-container"
+            classNamePrefix="aitk-react-select"
+            formatOptionLabel={(option: unknown) => {
+              const opt = option as SelectOption;
+              return opt.value === CUSTOM_SELECT_VALUE ? (
+                <span className="opacity-50 italic">~ Custom ~</span>
+              ) : (
+                opt.label
+              );
+            }}
+            onChange={selected => {
+              if (selected) {
+                const val = (selected as { value: string }).value;
+                if (val === CUSTOM_SELECT_VALUE) {
+                  setIsCustom(true);
+                  onChange('');
+                } else {
+                  setIsCustom(false);
+                  onChange(val);
+                }
+              }
+            }}
+          />
+        </div>
+        {isCustom && (
+          <input
+            type="text"
+            value={value}
+            onChange={e => onChange(e.target.value)}
+            className={`${inputClasses} w-2/3`}
+            placeholder={props.placeholder ?? 'Enter custom value'}
+            disabled={props.disabled}
+            autoFocus
+          />
+        )}
+      </div>
+    </div>
+  );
+};
+
 export interface CheckboxProps {
   label?: string | React.ReactNode;
   checked: boolean;

ui/src/helpers/captionJobConfig.ts

@@ -0,0 +1,74 @@
+import { CaptionJobConfig } from "@/types";
+import { captionerTypes } from "./captionOptions";
+
+
+export const defaultCaptionJobConfig: CaptionJobConfig = {
+  job: 'extension',
+  config: {
+    name: 'Caption Directory',
+    process: [
+      {
+        type: 'AceStepCaptioner',
+        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",
+          dtype: 'bf16',
+          quantize: true,
+          qtype: 'float8',
+          low_vram: true,
+          extensions: ['mp3', 'wav'],
+          path_to_caption: '',
+          recaption: false,
+        },
+      },
+    ],
+  },
+};
+
+
+const repairDefaults = (defaults: { [key: string]: any }) => {
+  let newDefaults: { [key: string]: any } = {};
+  // if the key doesnt start with config.process[0]., then add it
+  for (const key in defaults) {
+    if (!key.startsWith('config.process[0].')) {
+      newDefaults[`config.process[0].${key}`] = defaults[key];
+    } else {
+      newDefaults[key] = defaults[key];
+    }
+  }
+  return newDefaults;
+}
+
+
+
+export const handleCaptionerTypeChange = (
+  currentTypeName: string,
+  newTypeName: string,
+  jobConfig: CaptionJobConfig,
+  setJobConfig: (value: any, key: string) => void,
+) => {
+  const currentType = captionerTypes.find(a => a.name === currentTypeName);
+  if (!currentType || currentType.name === newTypeName) {
+    return;
+  }
+
+  // update the defaults when a model is selected
+  const newType = captionerTypes.find(model => model.name === newTypeName);
+
+  let currentDefaults = repairDefaults(currentType.defaults || {});
+  let newDefaults = repairDefaults(newType?.defaults || {});
+
+  // set new model
+  setJobConfig(newTypeName, 'config.process[0].type');
+
+  // revert defaults from previous model
+  for (const key in currentDefaults) {
+    setJobConfig(currentDefaults[key][1], key);
+  }
+
+  for (const key in newDefaults) {
+    setJobConfig(newDefaults[key][0], key);
+  }
+};

ui/src/helpers/captionOptions.ts

@@ -0,0 +1,114 @@
+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';
+
+export interface CaptionOption {
+    name: string;
+    label: string;
+    group: CaptionGroup;
+    hasMultiLinePrompts?: boolean;
+    defaults?: { [key: string]: any };
+    additionalSections?: AdditionalSections[];
+    name_or_path_options?: SelectOption[];
+    name_or_path2_options?: SelectOption[];
+}
+
+const defaultNameOrPath = '';
+
+const extensionsAudio = ['mp3', 'wav'];
+const extensionsImage = ['jpg', 'jpeg', 'png', 'bmp', 'webp'];
+
+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.";
+
+export const captionerTypes: CaptionOption[] = [
+    {
+        name: 'AceStepCaptioner',
+        label: 'Ace Step',
+        group: 'music',
+        defaults: {
+            'config.process[0].caption.model_name_or_path': ['ACE-Step/acestep-transcriber', defaultNameOrPath],
+            'config.process[0].caption.model_name_or_path2': ['ACE-Step/acestep-captioner', undefined],
+            'config.process[0].caption.extensions': [extensionsAudio, defaultExtensions],
+        },
+        name_or_path_options: [
+            { value: 'ACE-Step/acestep-transcriber', label: 'ACE-Step/acestep-transcriber' },
+        ],
+        name_or_path2_options: [
+            { value: 'ACE-Step/acestep-captioner', label: 'ACE-Step/acestep-captioner' },
+        ],
+        additionalSections: [
+            'caption.model_name_or_path2',
+        ],
+    },
+    {
+        name: 'Qwen3VLCaptioner',
+        label: 'Qwen3-VL',
+        group: 'image',
+        defaults: {
+            'config.process[0].caption.model_name_or_path': ['Qwen/Qwen3-VL-8B-Instruct', defaultNameOrPath],
+            'config.process[0].caption.extensions': [extensionsImage, defaultExtensions],
+            'config.process[0].caption.caption_prompt': [defaultImageCaptionPrompt, undefined],
+            'config.process[0].caption.max_res': [512, undefined],
+            'config.process[0].caption.max_new_tokens': [128, undefined],
+
+        },
+        name_or_path_options: [
+            { value: 'Qwen/Qwen3-VL-2B-Instruct', label: 'Qwen/Qwen3-VL-2B-Instruct' },
+            { value: 'Qwen/Qwen3-VL-4B-Instruct', label: 'Qwen/Qwen3-VL-4B-Instruct' },
+            { value: 'Qwen/Qwen3-VL-8B-Instruct', label: 'Qwen/Qwen3-VL-8B-Instruct' },
+            { value: 'Qwen/Qwen3-VL-30B-A3B-Instruct', label: 'Qwen/Qwen3-VL-30B-A3B-Instruct' },
+        ],
+        additionalSections: [
+            'caption.caption_prompt',
+            'caption.max_res',
+            'caption.max_new_tokens',
+        ],
+    },
+
+].sort((a, b) => {
+    // Sort by label, case-insensitive
+    return a.label.localeCompare(b.label, undefined, { sensitivity: 'base' });
+}) as any;
+
+export const groupedCaptionerTypes: GroupedSelectOption[] = captionerTypes.reduce((acc, arch) => {
+    const group = acc.find(g => g.label === arch.group);
+    if (group) {
+        group.options.push({ value: arch.name, label: arch.label });
+    } else {
+        acc.push({
+            label: arch.group,
+            options: [{ value: arch.name, label: arch.label }],
+        });
+    }
+    return acc;
+}, [] as GroupedSelectOption[]);
+
+export const quantizationOptions: SelectOption[] = [
+    { value: '', label: '- NONE -' },
+    { value: 'float8', label: 'float8 (default)' },
+    { value: 'uint7', label: '7 bit' },
+    { value: 'uint6', label: '6 bit' },
+    { value: 'uint5', label: '5 bit' },
+    { value: 'uint4', label: '4 bit' },
+    { value: 'uint3', label: '3 bit' },
+    { value: 'uint2', label: '2 bit' },
+];
+
+export const maxResOptions: SelectOption[] = [
+    { value: '256', label: '256' },
+    { value: '512', label: '512 (default)' },
+    { value: '768', label: '768' },
+    { value: '1024', label: '1024' },
+];
+export const maxNewTokensOptions: SelectOption[] = [
+    { value: '64', label: '64' },
+    { value: '128', label: '128 (default)' },
+    { value: '256', label: '256' },
+    { value: '512', label: '512' },
+    { value: '1024', label: '1024' },
+];
+
+export const defaultQtype = 'float8';

ui/src/helpers/defaultSamples.ts

@@ -0,0 +1,403 @@
+import { SampleConfig } from "@/types";
+
+export const defaultSampleConfig: SampleConfig = {
+  sampler: 'flowmatch',
+  sample_every: 250,
+  width: 1024,
+  height: 1024,
+  samples: [
+    {
+      prompt: 'woman with red hair, playing chess at the park, bomb going off in the background',
+    },
+    {
+      prompt: 'a woman holding a coffee cup, in a beanie, sitting at a cafe',
+    },
+    {
+      prompt: 'a horse is a DJ at a night club, fish eye lens, smoke machine, lazer lights, holding a martini',
+    },
+    {
+      prompt:
+        'a man showing off his cool new t shirt at the beach, a shark is jumping out of the water in the background',
+    },
+    {
+      prompt: 'a bear building a log cabin in the snow covered mountains',
+    },
+    {
+      prompt: 'woman playing the guitar, on stage, singing a song, laser lights, punk rocker',
+    },
+    {
+      prompt: 'hipster man with a beard, building a chair, in a wood shop',
+    },
+    {
+      prompt:
+        'photo of a man, white background, medium shot, modeling clothing, studio lighting, white backdrop',
+    },
+    {
+      prompt: "a man holding a sign that says, 'this is a sign'",
+    },
+    {
+      prompt:
+        'a bulldog, in a post apocalyptic world, with a shotgun, in a leather jacket, in a desert, with a motorcycle',
+    },
+  ],
+  neg: '',
+  seed: 42,
+  walk_seed: true,
+  guidance_scale: 4,
+  sample_steps: 50,
+  num_frames: 1,
+  fps: 1,
+}
+
+export const defaultAudioSampleConfig: SampleConfig = {
+  sampler: 'flowmatch',
+  sample_every: 250,
+  width: 1024,
+  height: 1024,
+  samples: [
+    {
+      prompt: `
+<CAPTION>my style song</CAPTION>
+<LYRICS>
+[Intro choir]
+Laura
+Laura
+Laura
+Laura Training
+
+[Verse 1]
+A new open model, she been training it nightly
+AI tool kit, she configures it tightly,
+Loss curves dropping down to the floor
+Wondering if she's done or she should train it some more.
+
+[Chorus]
+Laura training
+She trains on what she pleases
+Laura training
+No paying corporate sleazes
+Laura training
+This could be her best one
+Why go outside, Laura training is too fun
+
+[Instrumental Break]
+
+[Verse 4]
+She's caching all the latents
+now she doesn't need a vay
+Training on some voices
+What will she make them say
+
+[Chorus]
+Laura training
+She trains on what she pleases
+Laura training
+No paying corporate sleazes
+Laura training
+This could be her best one
+Why go outside, Laura training is too fun
+
+[Guitar Solo]
+
+[Chorus]
+Laura training
+She trains on what she pleases
+Laura training
+No paying corporate sleazes
+Laura training
+This could be her best one
+Why go outside, Laura training is too fun
+
+[Chorus]
+Laura training
+She trains on what she pleases
+Laura training
+No paying corporate sleazes
+Laura training
+This could be her best one
+Why go outside, Laura training is too fun
+
+[Instrumental Break]
+
+[Outro]
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+</LYRICS>
+<BPM>112</BPM>
+<KEYSCALE>A minor</KEYSCALE>
+<TIMESIGNATURE>4</TIMESIGNATURE>
+<DURATION>180</DURATION>
+<LANGUAGE>en</LANGUAGE>
+`,
+    }, {
+      prompt: `
+<CAPTION>my style song</CAPTION>
+<LYRICS>
+[Intro choir]
+Laura
+Laura
+Laura
+Laura Training
+
+[Verse 1]
+A new open model, she been training it nightly
+AI tool kit, she configures it tightly,
+Loss curves dropping down to the floor
+Wondering if she's done or she should train it some more.
+
+[Chorus]
+Laura training
+She trains on what she pleases
+Laura training
+No paying corporate sleazes
+Laura training
+This could be her best one
+Why go outside, Laura training is too fun
+
+[Instrumental Break]
+
+[Verse 4]
+She's caching all the latents
+now she doesn't need a vay
+Training on some voices
+What will she make them say
+
+[Chorus]
+Laura training
+She trains on what she pleases
+Laura training
+No paying corporate sleazes
+Laura training
+This could be her best one
+Why go outside, Laura training is too fun
+
+[Guitar Solo]
+
+[Chorus]
+Laura training
+She trains on what she pleases
+Laura training
+No paying corporate sleazes
+Laura training
+This could be her best one
+Why go outside, Laura training is too fun
+
+[Chorus]
+Laura training
+She trains on what she pleases
+Laura training
+No paying corporate sleazes
+Laura training
+This could be her best one
+Why go outside, Laura training is too fun
+
+[Instrumental Break]
+
+[Outro]
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+</LYRICS>
+<BPM>112</BPM>
+<KEYSCALE>A minor</KEYSCALE>
+<TIMESIGNATURE>4</TIMESIGNATURE>
+<DURATION>180</DURATION>
+<LANGUAGE>en</LANGUAGE>
+`,
+    }, {
+      prompt: `
+<CAPTION>my style song</CAPTION>
+<LYRICS>
+[Intro choir]
+Laura
+Laura
+Laura
+Laura Training
+
+[Verse 1]
+A new open model, she been training it nightly
+AI tool kit, she configures it tightly,
+Loss curves dropping down to the floor
+Wondering if she's done or she should train it some more.
+
+[Chorus]
+Laura training
+She trains on what she pleases
+Laura training
+No paying corporate sleazes
+Laura training
+This could be her best one
+Why go outside, Laura training is too fun
+
+[Instrumental Break]
+
+[Verse 4]
+She's caching all the latents
+now she doesn't need a vay
+Training on some voices
+What will she make them say
+
+[Chorus]
+Laura training
+She trains on what she pleases
+Laura training
+No paying corporate sleazes
+Laura training
+This could be her best one
+Why go outside, Laura training is too fun
+
+[Guitar Solo]
+
+[Chorus]
+Laura training
+She trains on what she pleases
+Laura training
+No paying corporate sleazes
+Laura training
+This could be her best one
+Why go outside, Laura training is too fun
+
+[Chorus]
+Laura training
+She trains on what she pleases
+Laura training
+No paying corporate sleazes
+Laura training
+This could be her best one
+Why go outside, Laura training is too fun
+
+[Instrumental Break]
+
+[Outro]
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+</LYRICS>
+<BPM>112</BPM>
+<KEYSCALE>A minor</KEYSCALE>
+<TIMESIGNATURE>4</TIMESIGNATURE>
+<DURATION>180</DURATION>
+<LANGUAGE>en</LANGUAGE>
+`,
+    }, {
+      prompt: `
+<CAPTION>my style song</CAPTION>
+<LYRICS>
+[Intro choir]
+Laura
+Laura
+Laura
+Laura Training
+
+[Verse 1]
+A new open model, she been training it nightly
+AI tool kit, she configures it tightly,
+Loss curves dropping down to the floor
+Wondering if she's done or she should train it some more.
+
+[Chorus]
+Laura training
+She trains on what she pleases
+Laura training
+No paying corporate sleazes
+Laura training
+This could be her best one
+Why go outside, Laura training is too fun
+
+[Instrumental Break]
+
+[Verse 4]
+She's caching all the latents
+now she doesn't need a vay
+Training on some voices
+What will she make them say
+
+[Chorus]
+Laura training
+She trains on what she pleases
+Laura training
+No paying corporate sleazes
+Laura training
+This could be her best one
+Why go outside, Laura training is too fun
+
+[Guitar Solo]
+
+[Chorus]
+Laura training
+She trains on what she pleases
+Laura training
+No paying corporate sleazes
+Laura training
+This could be her best one
+Why go outside, Laura training is too fun
+
+[Chorus]
+Laura training
+She trains on what she pleases
+Laura training
+No paying corporate sleazes
+Laura training
+This could be her best one
+Why go outside, Laura training is too fun
+
+[Instrumental Break]
+
+[Outro]
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+Ah yeah!
+It's Converging!
+</LYRICS>
+<BPM>112</BPM>
+<KEYSCALE>A minor</KEYSCALE>
+<TIMESIGNATURE>4</TIMESIGNATURE>
+<DURATION>180</DURATION>
+<LANGUAGE>en</LANGUAGE>
+`,
+    },
+  ],
+  neg: '',
+  seed: 42,
+  walk_seed: true,
+  guidance_scale: 4,
+  sample_steps: 25,
+  num_frames: 1,
+  fps: 1,
+};

ui/src/hooks/useJobByRef.tsx

@@ -0,0 +1,42 @@
+'use client';
+
+import { useEffect, useState } from 'react';
+import { Job } from '@prisma/client';
+import { apiClient } from '@/utils/api';
+
+export default function useJobByRef(jobRef: string | null, reloadInterval: null | number = null) {
+  const [job, setJob] = useState<Job | null>(null);
+  const [status, setStatus] = useState<'idle' | 'loading' | 'success' | 'error'>('idle');
+
+  const refreshJob = () => {
+    setStatus('loading');
+    apiClient
+      .get(`/api/jobs?job_ref=${jobRef}`)
+      .then(res => res.data)
+      .then(data => {
+        console.log('Job:', data);
+        setJob(data);
+        setStatus('success');
+      })
+      .catch(error => {
+        console.error('Error fetching job:', error);
+        setStatus('error');
+      });
+  };
+
+  useEffect(() => {
+    refreshJob();
+
+    if (reloadInterval) {
+      const interval = setInterval(() => {
+        refreshJob();
+      }, reloadInterval);
+
+      return () => {
+        clearInterval(interval);
+      };
+    }
+  }, [jobRef, reloadInterval]);
+
+  return { job, setJob, status, refreshJob };
+}

ui/src/hooks/useJobsList.tsx

@@ -4,14 +4,20 @@ import { useEffect, useState } from 'react';
 import { Job } from '@prisma/client';
 import { apiClient } from '@/utils/api';
 
-export default function useJobsList(onlyActive = false, reloadInterval: null | number = null) {
+type UseJobsListProps = {
+  onlyActive?: boolean;
+  reloadInterval?: number | null;
+  job_type?: string | null;
+};
+
+export default function useJobsList({ onlyActive = false, reloadInterval = null, job_type=null }: UseJobsListProps = {}) {
   const [jobs, setJobs] = useState<Job[]>([]);
   const [status, setStatus] = useState<'idle' | 'loading' | 'success' | 'error'>('idle');
 
   const refreshJobs = () => {
     setStatus('loading');
     apiClient
-      .get('/api/jobs')
+      .get('/api/jobs', { params: job_type ? { job_type } : undefined })
       .then(res => res.data)
       .then(data => {
         console.log('Jobs:', data);

ui/src/types.ts

@@ -252,6 +252,36 @@ export interface JobConfig {
   meta: MetaConfig;
 }
 
+export interface CaptionProcessConfig {
+  type: string;
+  sqlite_db_path?: string;
+  device: string;
+  caption: {
+    model_name_or_path: string;
+    model_name_or_path2?: string;
+    dtype: string;
+    quantize: boolean;
+    qtype: string;
+    low_vram: boolean;
+    extensions: string[];
+    path_to_caption: string;
+    recaption: boolean;
+    caption_prompt?: string;
+    max_res?: number;
+    max_new_tokens?: number;
+  }
+}
+
+export interface CaptionConfigObject {
+  name: string;
+  process: CaptionProcessConfig[];
+}
+
+export interface CaptionJobConfig {
+  job: string;
+  config: CaptionConfigObject;
+}
+
 export interface ConfigDoc {
   title: string | React.ReactNode;
   description: React.ReactNode;

ui/src/utils/basic.ts

@@ -11,3 +11,43 @@ export const audioExtensions = ['.mp3', '.wav'];
 export const isVideo = (filePath: string) => videoExtensions.includes(filePath.toLowerCase().slice(-4));
 export const isImage = (filePath: string) => imgExtensions.includes(filePath.toLowerCase().slice(-4));
 export const isAudio = (filePath: string) => audioExtensions.includes(filePath.toLowerCase().slice(-4));
+
+export const tagsToObj = (tagStr: string): Record<string, any> => {
+  const result: Record<string, any> = {};
+  const regex = /<([A-Z_][A-Z0-9_]*)>([\s\S]*?)<\/\1>/g;
+  let match;
+  while ((match = regex.exec(tagStr)) !== null) {
+    const value = match[2].trim();
+    try {
+      result[match[1]] = JSON.parse(value);
+    } catch {
+      result[match[1]] = value;
+    }
+  }
+  return result;
+};
+
+export const objToTags = (obj: Record<string, any>): string => {
+  return Object.entries(obj)
+    .map(([key, value]) => {
+      const content = typeof value === 'string' ? value : JSON.stringify(value);
+      return `<${key}>${content}</${key}>`;
+    })
+    .join('\n');
+};
+
+export const pathJoin = (...parts: string[]) => {
+  const sep = parts.length > 0 && parts[0].includes('\\') ? '\\' : '/';
+  const leadingTrailing = sep === '\\' ? /^\\+|\\+$/g : /^\/+|\/+$/g;
+  const trailing = sep === '\\' ? /\\+$/ : /\/+$/;
+  return parts
+    .map((part, index) => {
+      if (index === 0) {
+        return part.replace(trailing, '');
+      } else {
+        return part.replace(leadingTrailing, '');
+      }
+    })
+    .filter(part => part.length > 0)
+    .join(sep);
+}

ui/src/utils/jobs.ts

@@ -74,14 +74,18 @@ export const getAvaliableJobActions = (job: Job) => {
   const jobConfig = getJobConfig(job);
   const isStopping = job.stop && job.status === 'running';
   const canDelete = ['queued', 'completed', 'stopped', 'error'].includes(job.status) && !isStopping;
-  const canEdit = ['queued','completed', 'stopped', 'error'].includes(job.status) && !isStopping;
+  let canEdit = ['queued', 'completed', 'stopped', 'error'].includes(job.status) && !isStopping;
   const canRemoveFromQueue = job.status === 'queued';
   const canStop = job.status === 'running' && !isStopping;
   let canStart = ['stopped', 'error'].includes(job.status) && !isStopping;
   // can resume if more steps were added
-  if (job.status === 'completed' && jobConfig.config.process[0].train.steps > job.step && !isStopping) {
+  if (job.status === 'completed' && (jobConfig.config.process[0].train?.steps || 0) > job.step && !isStopping) {
     canStart = true;
   }
+  if (job.job_type !== 'train') {
+    // for non-train jobs, allow editing unless it's currently running
+    canEdit = false;
+  }
   return { canDelete, canEdit, canStop, canStart, canRemoveFromQueue };
 };
 
@@ -92,5 +96,5 @@ export const getNumberOfSamples = (job: Job) => {
 
 export const getTotalSteps = (job: Job) => {
   const jobConfig = getJobConfig(job);
-  return jobConfig.config.process[0].train.steps;
+  return jobConfig.config.process[0].train?.steps || 0;
 };

version.py

@@ -1 +1 @@
-VERSION = "0.8.2"
+VERSION = "0.9.0"