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

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

* Base training code done

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

* Fix issue with quantizing ace step

* Add album artwork to samples with waveform.

* Cleanup logs

* Add album art endpoint to speed up album art loading

* Made an make video with artwork script

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

* Add prompt tagging system for special tagged models.

* prompt tagging processing for ui working.

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

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

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

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

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

* Added qwen3 vl as a captioner for images

* Have an indicator when a dataset is being captioned.

* Adjust the way caption jobs look in the queue

* Fix a few issues. Adjust defaults.

* Version bump

* Added ace step to the readme.

extensions_built_in/audio_models/__init__.py

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

extensions_built_in/audio_models/ace_step/__init__.py

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

extensions_built_in/audio_models/ace_step/ace_step_15_model.py

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

extensions_built_in/audio_models/ace_step/src/pipeline.py

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

extensions_built_in/audio_models/base_audio_model.py

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