diff --git a/README.md b/README.md index 00b37984..0b40d300 100644 --- a/README.md +++ b/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) diff --git a/extensions_built_in/audio_models/__init__.py b/extensions_built_in/audio_models/__init__.py new file mode 100644 index 00000000..b01e2655 --- /dev/null +++ b/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, +] diff --git a/extensions_built_in/audio_models/ace_step/__init__.py b/extensions_built_in/audio_models/ace_step/__init__.py new file mode 100644 index 00000000..5b24a377 --- /dev/null +++ b/extensions_built_in/audio_models/ace_step/__init__.py @@ -0,0 +1 @@ +from .ace_step_15_model import AceStep15Model, AceStep15XLModel \ No newline at end of file diff --git a/extensions_built_in/audio_models/ace_step/ace_step_15_model.py b/extensions_built_in/audio_models/ace_step/ace_step_15_model.py new file mode 100644 index 00000000..4dc089c9 --- /dev/null +++ b/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}>(.*?)", 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" diff --git a/extensions_built_in/audio_models/ace_step/src/__init__.py b/extensions_built_in/audio_models/ace_step/src/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/extensions_built_in/audio_models/ace_step/src/model.py b/extensions_built_in/audio_models/ace_step/src/model.py new file mode 100644 index 00000000..60a364f7 --- /dev/null +++ b/extensions_built_in/audio_models/ace_step/src/model.py @@ -0,0 +1,1520 @@ +#!/usr/bin/env python3 +""" +ACE-Step v1.5 — Standalone single-file inference. + +Generates music from text + lyrics. All model code inlined — no project imports, +no trust_remote_code. Uses ComfyUI-style architecture for AIO checkpoint compat. + +Requirements: + pip install torch torchaudio transformers safetensors + +Usage: + python simple_inference.py --prompt "indie folk, warm female vocal, 100 bpm" \ + --lyrics "[Verse]\\nSunlight through the window pane" --duration 30 +""" + +import argparse +import math +import os +import time + +import torch +import torch.nn.functional as F +import torchaudio +from safetensors.torch import load_file +from torch import nn +from transformers import AutoTokenizer +import torch.utils.checkpoint as ckpt + +# ═══════════════════════════════════════════════════════════════════════════════ +# Constants +# ═══════════════════════════════════════════════════════════════════════════════ + +MODELS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models") +MODEL_PATHS = { + "base": os.path.join(MODELS_DIR, "ace_step_1.5_xl_base_aio.safetensors"), + "turbo": os.path.join(MODELS_DIR, "ace_step_1.5_turbo_aio.safetensors"), +} +SAMPLE_RATE = 48000 +LATENT_RATE = 25 # 48000 / 1920 + +SFT_PROMPT = """# Instruction +{instruction} + +# Caption +{caption} + +# Metas +{metas}<|endoftext|> +""" + +TURBO_TIMESTEPS = { + 1.0: [1.0, 0.875, 0.75, 0.625, 0.5, 0.375, 0.25, 0.125], + 2.0: [1.0, 0.933, 0.857, 0.769, 0.667, 0.545, 0.4, 0.222], + 3.0: [ + 1.0, + 0.9545454545454546, + 0.9, + 0.8333333333333334, + 0.75, + 0.6428571428571429, + 0.5, + 0.3, + ], +} + + +def compute_timesteps(num_steps, shift=3.0): + """Compute flow-matching timestep schedule with shifting.""" + import numpy as np + + sigmas = np.linspace(1.0, 0.0, num_steps + 1)[:-1] # exclude final 0 + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + return sigmas.tolist() + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Silence latent (hardcoded, from ComfyUI) +# ═══════════════════════════════════════════════════════════════════════════════ + + +def get_silence_latent(length, device, dtype=torch.bfloat16): + head = torch.tensor( + [ + [ + [ + 0.5707, + 0.0982, + 0.6909, + -0.5658, + 0.6266, + 0.6996, + -0.1365, + -0.1291, + -0.0776, + -0.1171, + -0.2743, + -0.8422, + -0.1168, + 1.5539, + -4.6936, + 0.7436, + -1.1846, + -0.2637, + 0.6933, + -6.7266, + 0.0966, + -0.1187, + -0.3501, + -1.1736, + 0.0587, + -2.0517, + -1.3651, + 0.7508, + -0.2490, + -1.3548, + -0.1290, + -0.7261, + 1.1132, + -0.3249, + 0.2337, + 0.3004, + 0.6605, + -0.0298, + -0.1989, + -0.4041, + 0.2843, + -1.0963, + -0.5519, + 0.2639, + -1.0436, + -0.1183, + 0.0640, + 0.4460, + -1.1001, + -0.6172, + -1.3241, + 1.1379, + 0.5623, + -0.1507, + -0.1963, + -0.4742, + -2.4697, + 0.5302, + 0.5381, + 0.4636, + -0.1782, + -0.0687, + 1.0333, + 0.4202, + ], + [ + 0.3040, + -0.1367, + 0.6200, + 0.0665, + -0.0642, + 0.4655, + -0.1187, + -0.0440, + 0.2941, + -0.2753, + 0.0173, + -0.2421, + -0.0147, + 1.5603, + -2.7025, + 0.7907, + -0.9736, + -0.0682, + 0.1294, + -5.0707, + -0.2167, + 0.3302, + -0.1513, + -0.8100, + -0.3894, + -0.2884, + -0.3149, + 0.8660, + -0.3817, + -1.7061, + 0.5824, + -0.4840, + 0.6938, + 0.1859, + 0.1753, + 0.3081, + 0.0195, + 0.1403, + -0.0754, + -0.2091, + 0.1251, + -0.1578, + -0.4968, + -0.1052, + -0.4554, + -0.0320, + 0.1284, + 0.4974, + -1.1889, + -0.0344, + -0.8313, + 0.2953, + 0.5445, + -0.6249, + -0.1595, + -0.0682, + -3.1412, + 0.0484, + 0.4153, + 0.8260, + -0.1526, + -0.0625, + 0.5366, + 0.8473, + ], + [ + 5.3524e-02, + -1.7534e-01, + 5.4443e-01, + -4.3501e-01, + -2.1317e-03, + 3.7200e-01, + -4.0143e-03, + -1.5516e-01, + -1.2968e-01, + -1.5375e-01, + -7.7107e-02, + -2.0593e-01, + -3.2780e-01, + 1.5142e00, + -2.6101e00, + 5.8698e-01, + -1.2716e00, + -2.4773e-01, + -2.7933e-02, + -5.0799e00, + 1.1601e-01, + 4.0987e-01, + -2.2030e-02, + -6.6495e-01, + -2.0995e-01, + -6.3474e-01, + -1.5893e-01, + 8.2745e-01, + -2.2992e-01, + -1.6816e00, + 5.4440e-01, + -4.9579e-01, + 5.5128e-01, + 3.0477e-01, + 8.3052e-02, + -6.1782e-02, + 5.9036e-03, + 2.9553e-01, + -8.0645e-02, + -1.0060e-01, + 1.9144e-01, + -3.8124e-01, + -7.2949e-01, + 2.4520e-02, + -5.0814e-01, + 2.3977e-01, + 9.2943e-02, + 3.9256e-01, + -1.1993e00, + -3.2752e-01, + -7.2707e-01, + 2.9476e-01, + 4.3542e-01, + -8.8597e-01, + -4.1686e-01, + -8.5390e-02, + -2.9018e00, + 6.4988e-02, + 5.3945e-01, + 9.1988e-01, + 5.8762e-02, + -7.0098e-02, + 6.4772e-01, + 8.9118e-01, + ], + [ + -3.2225e-02, + -1.3195e-01, + 5.6411e-01, + -5.4766e-01, + -5.2170e-03, + 3.1425e-01, + -5.4367e-02, + -1.9419e-01, + -1.3059e-01, + -1.3660e-01, + -9.0984e-02, + -1.9540e-01, + -2.5590e-01, + 1.5440e00, + -2.6349e00, + 6.8273e-01, + -1.2532e00, + -1.9810e-01, + -2.2793e-02, + -5.0506e00, + 1.8818e-01, + 5.0109e-01, + 7.3546e-03, + -6.8771e-01, + -3.0676e-01, + -7.3257e-01, + -1.6687e-01, + 9.2232e-01, + -1.8987e-01, + -1.7267e00, + 5.3355e-01, + -5.3179e-01, + 4.4953e-01, + 2.8820e-01, + 1.3012e-01, + -2.0943e-01, + -1.1348e-01, + 3.3929e-01, + -1.5069e-01, + -1.2919e-01, + 1.8929e-01, + -3.6166e-01, + -8.0756e-01, + 6.6387e-02, + -5.8867e-01, + 1.6978e-01, + 1.0134e-01, + 3.3877e-01, + -1.2133e00, + -3.2492e-01, + -8.1237e-01, + 3.8101e-01, + 4.3765e-01, + -8.0596e-01, + -4.4531e-01, + -4.7513e-02, + -2.9266e00, + 1.1741e-03, + 4.5123e-01, + 9.3075e-01, + 5.3688e-02, + -1.9621e-01, + 6.4530e-01, + 9.3870e-01, + ], + ] + ], + device=device, + ).movedim(-1, 1) + body = ( + torch.tensor( + [ + [ + [ + -1.3672e-01, + -1.5820e-01, + 5.8594e-01, + -5.7422e-01, + 3.0273e-02, + 2.7930e-01, + -2.5940e-03, + -2.0703e-01, + -1.6113e-01, + -1.4746e-01, + -2.7710e-02, + -1.8066e-01, + -2.9688e-01, + 1.6016e00, + -2.6719e00, + 7.7734e-01, + -1.3516e00, + -1.9434e-01, + -7.1289e-02, + -5.0938e00, + 2.4316e-01, + 4.7266e-01, + 4.6387e-02, + -6.6406e-01, + -2.1973e-01, + -6.7578e-01, + -1.5723e-01, + 9.5312e-01, + -2.0020e-01, + -1.7109e00, + 5.8984e-01, + -5.7422e-01, + 5.1562e-01, + 2.8320e-01, + 1.4551e-01, + -1.8750e-01, + -5.9814e-02, + 3.6719e-01, + -1.0059e-01, + -1.5723e-01, + 2.0605e-01, + -4.3359e-01, + -8.2812e-01, + 4.5654e-02, + -6.6016e-01, + 1.4844e-01, + 9.4727e-02, + 3.8477e-01, + -1.2578e00, + -3.3203e-01, + -8.5547e-01, + 4.3359e-01, + 4.2383e-01, + -8.9453e-01, + -5.0391e-01, + -5.6152e-02, + -2.9219e00, + -2.4658e-02, + 5.0391e-01, + 9.8438e-01, + 7.2754e-02, + -2.1582e-01, + 6.3672e-01, + 1.0000e00, + ] + ] + ], + device=device, + ) + .movedim(-1, 1) + .repeat(1, 1, length) + ) + body[:, :, : head.shape[-1]] = head + return body.to(dtype) # [1, 64, T] + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Helpers +# ═══════════════════════════════════════════════════════════════════════════════ + + +class RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.eps = eps + + def forward(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight + + +class RotaryEmbedding(nn.Module): + def __init__(self, dim, base=1000000.0): + super().__init__() + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self._cos = None + self._sin = None + self._cached_len = 0 + + def _build_cache(self, seq_len, device, dtype): + if ( + seq_len <= self._cached_len + and self._cos is not None + and self._cos.device == device + ): + return + t = torch.arange(seq_len, device=device, dtype=torch.float32) + freqs = torch.outer(t, self.inv_freq.to(device)) + emb = torch.cat((freqs, freqs), dim=-1) + self._cos = emb.cos().to(dtype) + self._sin = emb.sin().to(dtype) + self._cached_len = seq_len + + def forward(self, x, seq_len): + self._build_cache(seq_len, x.device, x.dtype) + return self._cos[:seq_len], self._sin[:seq_len] + + +def rotate_half(x): + x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary(q, k, cos, sin): + cos, sin = cos.unsqueeze(0).unsqueeze(0), sin.unsqueeze(0).unsqueeze(0) + return (q * cos + rotate_half(q) * sin), (k * cos + rotate_half(k) * sin) + + +class MLP(nn.Module): + def __init__(self, hidden, inter): + super().__init__() + self.gate_proj = nn.Linear(hidden, inter, bias=False) + self.up_proj = nn.Linear(hidden, inter, bias=False) + self.down_proj = nn.Linear(inter, hidden, bias=False) + + def forward(self, x): + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + +def pack_sequences(h1, h2, m1, m2): + h = torch.cat([h1, h2], dim=1) + if m1 is not None and m2 is not None: + m = torch.cat([m1, m2], dim=1) + B, L, D = h.shape + idx = m.argsort(dim=1, descending=True, stable=True) + h = torch.gather(h, 1, idx.unsqueeze(-1).expand(B, L, D)) + lengths = m.sum(dim=1) + m = torch.arange(L, device=h.device).unsqueeze(0) < lengths.unsqueeze(1) + else: + m = None + return h, m + + +def timestep_embedding(t, dim, scale=1000, max_period=10000): + t = t * scale + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(half, dtype=torch.float32, device=t.device) + / half + ) + args = t[:, None].float() * freqs[None] + return torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + + +# ═══════════════════════════════════════════════════════════════════════════════ +# DiT model components (ComfyUI-style, matches AIO weight keys) +# ═══════════════════════════════════════════════════════════════════════════════ + + +class TimestepEmbed(nn.Module): + def __init__(self, hidden): + super().__init__() + self.linear_1 = nn.Linear(256, hidden) + self.act1 = nn.SiLU() + self.linear_2 = nn.Linear(hidden, hidden) + self.act2 = nn.SiLU() + self.time_proj = nn.Linear(hidden, hidden * 6) + self.scale = 1000 + + def forward(self, t, dtype=None): + emb = timestep_embedding(t, 256, self.scale) + temb = self.act1(self.linear_1(emb.to(dtype=dtype))) + temb = self.linear_2(temb) + proj = self.time_proj(self.act2(temb)).view(-1, 6, temb.shape[-1]) + return temb, proj + + +class Attention(nn.Module): + def __init__( + self, + hidden, + num_heads, + num_kv, + head_dim, + eps=1e-6, + is_cross=False, + sliding_window=None, + ): + super().__init__() + self.num_heads = num_heads + self.num_kv = num_kv + self.head_dim = head_dim + self.is_cross = is_cross + self.sliding_window = sliding_window + self.q_proj = nn.Linear(hidden, num_heads * head_dim, bias=False) + self.k_proj = nn.Linear(hidden, num_kv * head_dim, bias=False) + self.v_proj = nn.Linear(hidden, num_kv * head_dim, bias=False) + self.o_proj = nn.Linear(num_heads * head_dim, hidden, bias=False) + self.q_norm = RMSNorm(head_dim, eps) + self.k_norm = RMSNorm(head_dim, eps) + + def forward(self, x, encoder_hidden_states=None, position_embeddings=None): + B, L, _ = x.shape + q = self.q_norm( + self.q_proj(x).view(B, L, self.num_heads, self.head_dim) + ).transpose(1, 2) + + src = ( + encoder_hidden_states + if (self.is_cross and encoder_hidden_states is not None) + else x + ) + sL = src.shape[1] + k = self.k_norm( + self.k_proj(src).view(B, sL, self.num_kv, self.head_dim) + ).transpose(1, 2) + v = self.v_proj(src).view(B, sL, self.num_kv, self.head_dim).transpose(1, 2) + + if position_embeddings is not None and not ( + self.is_cross and encoder_hidden_states is not None + ): + q, k = apply_rotary(q, k, *position_embeddings) + + n_rep = self.num_heads // self.num_kv + if n_rep > 1: + k = k.repeat_interleave(n_rep, dim=1) + v = v.repeat_interleave(n_rep, dim=1) + + attn_bias = None + if self.sliding_window is not None and not self.is_cross: + idx = torch.arange(L, device=q.device) + in_win = ( + torch.abs(idx.unsqueeze(1) - idx.unsqueeze(0)) <= self.sliding_window + ) + attn_bias = torch.zeros(L, sL, device=q.device, dtype=q.dtype) + attn_bias.masked_fill_(~in_win, torch.finfo(q.dtype).min) + attn_bias = attn_bias.unsqueeze(0).unsqueeze(0) + + out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_bias) + return self.o_proj(out.transpose(1, 2).reshape(B, L, -1)) + + +class EncoderLayer(nn.Module): + def __init__(self, hidden, heads, kv, head_dim, inter, eps=1e-6): + super().__init__() + self.self_attn = Attention(hidden, heads, kv, head_dim, eps) + self.input_layernorm = RMSNorm(hidden, eps) + self.post_attention_layernorm = RMSNorm(hidden, eps) + self.mlp = MLP(hidden, inter) + + def forward(self, x, position_embeddings): + x = x + self.self_attn( + self.input_layernorm(x), position_embeddings=position_embeddings + ) + x = x + self.mlp(self.post_attention_layernorm(x)) + return x + + +class DiTLayer(nn.Module): + def __init__( + self, hidden, heads, kv, head_dim, inter, eps=1e-6, sliding_window=None + ): + super().__init__() + self.self_attn_norm = RMSNorm(hidden, eps) + self.self_attn = Attention( + hidden, heads, kv, head_dim, eps, sliding_window=sliding_window + ) + self.cross_attn_norm = RMSNorm(hidden, eps) + self.cross_attn = Attention(hidden, heads, kv, head_dim, eps, is_cross=True) + self.mlp_norm = RMSNorm(hidden, eps) + self.mlp = MLP(hidden, inter) + self.scale_shift_table = nn.Parameter(torch.empty(1, 6, hidden)) + + def forward(self, x, temb, enc, position_embeddings): + s_msa, sc_msa, g_msa, s_mlp, sc_mlp, g_mlp = ( + self.scale_shift_table.to(temb) + temb + ).chunk(6, dim=1) + x = ( + x + + self.self_attn( + self.self_attn_norm(x) * (1 + sc_msa) + s_msa, + position_embeddings=position_embeddings, + ) + * g_msa + ) + x = x + self.cross_attn(self.cross_attn_norm(x), encoder_hidden_states=enc) + x = x + self.mlp(self.mlp_norm(x) * (1 + sc_mlp) + s_mlp) * g_mlp + return x + + +# ── Encoders ── + + +class LyricEncoder(nn.Module): + def __init__( + self, text_dim, hidden, n_layers, heads, kv, head_dim, inter, eps=1e-6 + ): + super().__init__() + self.embed_tokens = nn.Linear(text_dim, hidden) + self.norm = RMSNorm(hidden, eps) + self.rotary_emb = RotaryEmbedding(head_dim) + self.layers = nn.ModuleList( + [ + EncoderLayer(hidden, heads, kv, head_dim, inter, eps) + for _ in range(n_layers) + ] + ) + + def forward(self, embeds): + x = self.embed_tokens(embeds) + cos, sin = self.rotary_emb(x, x.shape[1]) + for layer in self.layers: + x = layer(x, (cos, sin)) + return self.norm(x) + + +class TimbreEncoder(nn.Module): + def __init__( + self, timbre_dim, hidden, n_layers, heads, kv, head_dim, inter, eps=1e-6 + ): + super().__init__() + self.embed_tokens = nn.Linear(timbre_dim, hidden) + self.norm = RMSNorm(hidden, eps) + self.rotary_emb = RotaryEmbedding(head_dim) + self.layers = nn.ModuleList( + [ + EncoderLayer(hidden, heads, kv, head_dim, inter, eps) + for _ in range(n_layers) + ] + ) + self.special_token = nn.Parameter(torch.empty(1, 1, hidden)) + + def forward(self, packed, order_mask): + x = self.embed_tokens(packed) + cos, sin = self.rotary_emb(x, x.shape[1]) + for layer in self.layers: + x = layer(x, (cos, sin)) + x = self.norm(x) + cls = x[:, 0, :] + # Unpack to batch + N, D = cls.shape + B = int(order_mask.max().item() + 1) + counts = torch.bincount(order_mask, minlength=B) + mc = counts.max().item() + result = torch.zeros(B, mc, D, device=cls.device, dtype=cls.dtype) + mask = torch.zeros(B, mc, device=cls.device, dtype=torch.long) + for i in range(N): + b = order_mask[i].item() + pos = (order_mask[:i] == b).sum().item() + result[b, pos] = cls[i] + mask[b, pos] = 1 + return result, mask + + +class ConditionEncoder(nn.Module): + def __init__( + self, + text_dim, + timbre_dim, + hidden, + n_lyric, + n_timbre, + heads, + kv, + head_dim, + inter, + eps=1e-6, + ): + super().__init__() + self.text_projector = nn.Linear(text_dim, hidden, bias=False) + self.lyric_encoder = LyricEncoder( + text_dim, hidden, n_lyric, heads, kv, head_dim, inter, eps + ) + self.timbre_encoder = TimbreEncoder( + timbre_dim, hidden, n_timbre, heads, kv, head_dim, inter, eps + ) + + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + + def forward(self, text_h, text_m, lyric_h, lyric_m, refer_packed, refer_order): + text_proj = self.text_projector(text_h) + lyric_enc = self.lyric_encoder(lyric_h) + timbre_enc, timbre_mask = self.timbre_encoder(refer_packed, refer_order) + merged, merged_m = pack_sequences(lyric_enc, timbre_enc, lyric_m, timbre_mask) + final, final_m = pack_sequences(merged, text_proj, merged_m, text_m) + return final, final_m + + +# ── DiT ── + + +class DiTModel(nn.Module): + def __init__( + self, + in_ch, + hidden, + n_layers, + heads, + kv, + head_dim, + inter, + patch, + out_ch, + layer_types=None, + sliding_window=128, + eps=1e-6, + cond_dim=None, + ): + super().__init__() + self.patch_size = patch + self.rotary_emb = RotaryEmbedding(head_dim) + self.proj_in = nn.Sequential( + nn.Identity(), nn.Conv1d(in_ch, hidden, kernel_size=patch, stride=patch) + ) + self.time_embed = TimestepEmbed(hidden) + self.time_embed_r = TimestepEmbed(hidden) + self.condition_embedder = nn.Linear(cond_dim or hidden, hidden) + lt = layer_types or [ + "sliding_attention" if i % 2 == 0 else "full_attention" + for i in range(n_layers) + ] + self.layers = nn.ModuleList( + [ + DiTLayer( + hidden, + heads, + kv, + head_dim, + inter, + eps, + sliding_window=sliding_window + if lt[i] == "sliding_attention" + else None, + ) + for i in range(n_layers) + ] + ) + self.norm_out = RMSNorm(hidden, eps) + self.proj_out = nn.Sequential( + nn.Identity(), + nn.ConvTranspose1d(hidden, out_ch, kernel_size=patch, stride=patch), + ) + self.scale_shift_table = nn.Parameter(torch.empty(1, 2, hidden)) + self.gradient_checkpointing = False + + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + + def forward(self, x, timestep, timestep_r, attention_mask, enc_h, enc_m, context): + temb_t, proj_t = self.time_embed(timestep, dtype=x.dtype) + temb_r, proj_r = self.time_embed_r(timestep - timestep_r, dtype=x.dtype) + temb = temb_t + temb_r + tproj = proj_t + proj_r + + h = torch.cat([context, x], dim=-1) + orig_len = h.shape[1] + if h.shape[1] % self.patch_size != 0: + h = F.pad(h, (0, 0, 0, self.patch_size - h.shape[1] % self.patch_size)) + h = self.proj_in(h.transpose(1, 2)).transpose(1, 2) + enc = self.condition_embedder(enc_h) + cos, sin = self.rotary_emb(h, h.shape[1]) + for layer in self.layers: + if torch.is_grad_enabled() and self.gradient_checkpointing: + h = ckpt.checkpoint( + layer, h, tproj, enc, (cos, sin), use_reentrant=False + ) + else: + h = layer(h, tproj, enc, (cos, sin)) + shift, scale = (self.scale_shift_table.to(temb) + temb.unsqueeze(1)).chunk( + 2, dim=1 + ) + h = self.norm_out(h) * (1 + scale) + shift + h = self.proj_out(h.transpose(1, 2)).transpose(1, 2) + return h[:, :orig_len, :] + + +# ── Top-level model ── + + +class AceStep15(nn.Module): + def __init__( + self, + hidden=2048, + text_dim=1024, + timbre_dim=64, + out_ch=64, + n_dit=24, + n_lyric=8, + n_timbre=4, + heads=16, + kv=8, + head_dim=128, + inter=6144, + patch=2, + in_ch=192, + sliding_window=128, + eps=1e-6, + layer_types=None, + # Encoder can have different size than decoder (XL models) + enc_hidden=None, + enc_heads=None, + enc_kv=None, + enc_inter=None, + ): + super().__init__() + eh = enc_hidden or hidden + eheads = enc_heads or heads + ekv = enc_kv or kv + einter = enc_inter or inter + + self.decoder = DiTModel( + in_ch, + hidden, + n_dit, + heads, + kv, + head_dim, + inter, + patch, + out_ch, + layer_types, + sliding_window, + eps, + cond_dim=eh, + ) + self.encoder = ConditionEncoder( + text_dim, + timbre_dim, + eh, + n_lyric, + n_timbre, + eheads, + ekv, + head_dim, + einter, + eps, + ) + self.null_condition_emb = nn.Parameter(torch.empty(1, 1, eh)) + self._gradient_checkpointing = False + + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + + def enable_gradient_checkpointing(self): + self.gradient_checkpointing = True + + @property + def gradient_checkpointing(self): + return self._gradient_checkpointing + + @gradient_checkpointing.setter + def gradient_checkpointing(self, value): + self._gradient_checkpointing = value + self.decoder.gradient_checkpointing = value + + def prepare_condition( + self, + text_h, + text_m, + lyric_h, + lyric_m, + refer_packed, + refer_order, + src_latents, + chunk_masks, + ): + enc_h, enc_m = self.encoder( + text_h, text_m, lyric_h, lyric_m, refer_packed, refer_order + ) + context = torch.cat([src_latents, chunk_masks.to(src_latents.dtype)], dim=-1) + return enc_h, enc_m, context + + +# ═══════════════════════════════════════════════════════════════════════════════ +# VAE (ComfyUI Oobleck style — uses parametrizations.weight_norm) +# ═══════════════════════════════════════════════════════════════════════════════ + + +def WNConv1d(*args, **kwargs): + return torch.nn.utils.parametrizations.weight_norm(nn.Conv1d(*args, **kwargs)) + + +def WNConvT1d(*args, **kwargs): + return torch.nn.utils.parametrizations.weight_norm( + nn.ConvTranspose1d(*args, **kwargs) + ) + + +class SnakeBeta(nn.Module): + def __init__(self, channels): + super().__init__() + self.alpha = nn.Parameter(torch.zeros(channels)) + self.beta = nn.Parameter(torch.zeros(channels)) + + def forward(self, x): + a = self.alpha.unsqueeze(0).unsqueeze(-1).exp().to(x.device) + b = self.beta.unsqueeze(0).unsqueeze(-1).exp().to(x.device) + return x + (1.0 / (b + 1e-9)) * torch.sin(x * a).pow(2) + + +class ResUnit(nn.Module): + def __init__(self, ch, dilation): + super().__init__() + self.layers = nn.Sequential( + SnakeBeta(ch), + WNConv1d(ch, ch, 7, dilation=dilation, padding=(dilation * 6) // 2), + SnakeBeta(ch), + WNConv1d(ch, ch, 1), + ) + + def forward(self, x): + return x + self.layers(x) + + +class EncBlock(nn.Module): + def __init__(self, in_ch, out_ch, stride): + super().__init__() + self.layers = nn.Sequential( + ResUnit(in_ch, 1), + ResUnit(in_ch, 3), + ResUnit(in_ch, 9), + SnakeBeta(in_ch), + WNConv1d( + in_ch, out_ch, 2 * stride, stride=stride, padding=math.ceil(stride / 2) + ), + ) + + def forward(self, x): + return self.layers(x) + + +class DecBlock(nn.Module): + def __init__(self, in_ch, out_ch, stride): + super().__init__() + self.layers = nn.Sequential( + SnakeBeta(in_ch), + WNConvT1d( + in_ch, out_ch, 2 * stride, stride=stride, padding=math.ceil(stride / 2) + ), + ResUnit(out_ch, 1), + ResUnit(out_ch, 3), + ResUnit(out_ch, 9), + ) + + def forward(self, x): + return self.layers(x) + + +class VAEBottleneck(nn.Module): + def encode(self, x): + mean, scale = x.chunk(2, dim=1) + return mean + + def decode(self, x): + return x + + +class _SeqWrap(nn.Module): + """Wraps Sequential as .layers so state_dict keys match AIO format.""" + + def __init__(self, *modules): + super().__init__() + self.layers = nn.Sequential(*modules) + + def forward(self, x): + return self.layers(x) + + +class OobleckVAE(nn.Module): + def __init__( + self, + in_ch=2, + channels=128, + latent_dim=64, + c_mults=(1, 2, 4, 8, 16), + strides=(2, 4, 4, 6, 10), + ): + super().__init__() + cm = [1] + list(c_mults) + # Encoder + enc = [WNConv1d(in_ch, cm[0] * channels, 7, padding=3)] + for i in range(len(cm) - 1): + enc.append(EncBlock(cm[i] * channels, cm[i + 1] * channels, strides[i])) + enc += [ + SnakeBeta(cm[-1] * channels), + WNConv1d(cm[-1] * channels, latent_dim * 2, 3, padding=1), + ] + self.encoder = _SeqWrap(*enc) + # Decoder + dec = [WNConv1d(latent_dim, cm[-1] * channels, 7, padding=3)] + for i in range(len(cm) - 1, 0, -1): + dec.append(DecBlock(cm[i] * channels, cm[i - 1] * channels, strides[i - 1])) + dec += [ + SnakeBeta(cm[0] * channels), + WNConv1d(cm[0] * channels, in_ch, 7, padding=3, bias=False), + ] + self.decoder = _SeqWrap(*dec) + self.bottleneck = VAEBottleneck() + + def encode(self, x): + return self.bottleneck.encode(self.encoder(x)) + + def decode(self, x): + return self.decoder(self.bottleneck.decode(x)) + + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Text encoder (Qwen3-Embedding, just need embed_tokens + model) +# ═══════════════════════════════════════════════════════════════════════════════ + + +class TextEncoder(nn.Module): + """Wraps Qwen3 weights loaded from AIO. Forward returns last_hidden_state.""" + + def __init__(self, qwen_model): + super().__init__() + self.model = qwen_model # the inner model (layers, norm, embed_tokens) + + def encode_text(self, input_ids): + return self.model(input_ids=input_ids).last_hidden_state + + def encode_lyrics(self, input_ids): + return self.model.embed_tokens(input_ids) + + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Loading +# ═══════════════════════════════════════════════════════════════════════════════ + + +def infer_dit_config(dit_sd): + """Infer model config from DiT state dict tensor shapes.""" + # hidden_size from decoder norm + hidden = dit_sd["decoder.scale_shift_table"].shape[2] + # intermediate_size from MLP gate_proj + inter = dit_sd["decoder.layers.0.mlp.gate_proj.weight"].shape[0] + # num_heads from q_proj: q_proj.weight is [num_heads * head_dim, hidden] + q_size = dit_sd["decoder.layers.0.self_attn.q_proj.weight"].shape[0] + # head_dim from q_norm + head_dim = dit_sd["decoder.layers.0.self_attn.q_norm.weight"].shape[0] + heads = q_size // head_dim + # num_kv_heads from k_proj + k_size = dit_sd["decoder.layers.0.self_attn.k_proj.weight"].shape[0] + kv = k_size // head_dim + # num_dit_layers: count unique layer indices + n_dit = ( + max(int(k.split(".")[2]) for k in dit_sd if k.startswith("decoder.layers.")) + 1 + ) + # encoder hidden (may differ from decoder hidden for XL models) + enc_hidden = dit_sd["encoder.text_projector.weight"].shape[0] + # encoder layers + n_lyric = ( + max( + int(k.split(".")[3]) + for k in dit_sd + if k.startswith("encoder.lyric_encoder.layers.") + ) + + 1 + ) + n_timbre = ( + max( + int(k.split(".")[3]) + for k in dit_sd + if k.startswith("encoder.timbre_encoder.layers.") + ) + + 1 + ) + # encoder attention config + enc_heads = ( + dit_sd["encoder.lyric_encoder.layers.0.self_attn.q_proj.weight"].shape[0] + // head_dim + ) + enc_kv = ( + dit_sd["encoder.lyric_encoder.layers.0.self_attn.k_proj.weight"].shape[0] + // head_dim + ) + enc_inter = dit_sd["encoder.lyric_encoder.layers.0.mlp.gate_proj.weight"].shape[0] + config = dict( + hidden=hidden, + inter=inter, + heads=heads, + kv=kv, + head_dim=head_dim, + n_dit=n_dit, + n_lyric=n_lyric, + n_timbre=n_timbre, + enc_hidden=enc_hidden, + enc_heads=enc_heads, + enc_kv=enc_kv, + enc_inter=enc_inter, + ) + print( + f" Detected config: hidden={hidden}, inter={inter}, heads={heads}, kv={kv}, " + f"n_dit={n_dit}, enc_hidden={enc_hidden}" + ) + return config + + +def load_models(checkpoint_path, device="cuda", dtype=torch.bfloat16): + if not os.path.isfile(checkpoint_path): + raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}") + print(f"Loading from: {checkpoint_path}") + + sd = load_file(checkpoint_path) + + # --- DiT --- + print(" Loading DiT...") + dit_sd = { + k.removeprefix("model.diffusion_model."): v + for k, v in sd.items() + if k.startswith("model.diffusion_model.") + } + cfg = infer_dit_config(dit_sd) + model = AceStep15( + hidden=cfg["hidden"], + inter=cfg["inter"], + heads=cfg["heads"], + kv=cfg["kv"], + head_dim=cfg["head_dim"], + n_dit=cfg["n_dit"], + n_lyric=cfg["n_lyric"], + n_timbre=cfg["n_timbre"], + enc_hidden=cfg["enc_hidden"], + enc_heads=cfg["enc_heads"], + enc_kv=cfg["enc_kv"], + enc_inter=cfg["enc_inter"], + ) + missing, unexpected = model.load_state_dict(dit_sd, strict=False) + # tokenizer/detokenizer keys are expected to be unused (cover mode only) + unexpected = [ + k for k in unexpected if not k.startswith(("tokenizer.", "detokenizer.")) + ] + if missing: + print(f" DiT missing: {len(missing)} (first 3: {missing[:3]})") + if unexpected: + print(f" DiT unexpected: {len(unexpected)} (first 3: {unexpected[:3]})") + model = model.to(device).to(dtype).eval() + + # --- VAE --- + print(" Loading VAE...") + vae_sd = {k.removeprefix("vae."): v for k, v in sd.items() if k.startswith("vae.")} + vae = OobleckVAE() + m, u = vae.load_state_dict(vae_sd, strict=False) + if m: + print(f" VAE missing: {len(m)} (first 3: {m[:3]})") + if u: + print(f" VAE unexpected: {len(u)}") + vae = vae.to(device).to(dtype).eval() + + # --- Text encoder (Qwen3-Embedding from AIO) --- + print(" Loading text encoder...") + te_sd = { + k.removeprefix("text_encoders.qwen3_06b.transformer.model."): v + for k, v in sd.items() + if k.startswith("text_encoders.qwen3_06b.transformer.model.") + } + # Load Qwen3 model structure from transformers, then override weights + from transformers import Qwen3Model, Qwen3Config + + qwen_cfg = Qwen3Config( + vocab_size=151669, + hidden_size=1024, + intermediate_size=3072, + num_hidden_layers=28, + num_attention_heads=16, + num_key_value_heads=8, + head_dim=128, + max_position_embeddings=32768, + rms_norm_eps=1e-6, + ) + qwen = Qwen3Model(qwen_cfg) + m2, u2 = qwen.load_state_dict(te_sd, strict=False) + if m2: + print(f" TE missing: {len(m2)} (first 3: {m2[:3]})") + te = TextEncoder(qwen).to(device).to(dtype).eval() + + # Tokenizer — download from HF + print(" Loading tokenizer...") + tok = AutoTokenizer.from_pretrained( + "Qwen/Qwen3-Embedding-0.6B", trust_remote_code=False + ) + + del sd # free memory + torch.cuda.empty_cache() if torch.cuda.is_available() else None + + print(" Done.\n") + return dict( + model=model, vae=vae, text_encoder=te, tokenizer=tok, device=device, dtype=dtype + ) + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Inference +# ═══════════════════════════════════════════════════════════════════════════════ + + +@torch.inference_mode() +def get_latent(audio_path, models): + """Encode audio file to VAE latent. Returns [1, 64, T] tensor.""" + vae, device, dtype = models["vae"], models["device"], models["dtype"] + wav, sr = torchaudio.load(audio_path) + if sr != SAMPLE_RATE: + wav = torchaudio.functional.resample(wav, sr, SAMPLE_RATE) + if wav.shape[0] == 1: + wav = wav.repeat(2, 1) + elif wav.shape[0] > 2: + wav = wav[:2] + return vae.encode(wav.unsqueeze(0).to(device, dtype)) # [1, 64, T] + + +@torch.inference_mode() +def generate( + models, + prompt, + lyrics="", + duration=30.0, + seed=42, + bpm="N/A", + key="N/A", + time_sig="N/A", + language="en", + timesteps=None, + guidance_scale=1.0, +): + model = models["model"] + vae = models["vae"] + te = models["text_encoder"] + tok = models["tokenizer"] + device = models["device"] + dtype = models["dtype"] + + t_sched = timesteps + latent_len = int(duration * LATENT_RATE) + print( + f"Duration: {duration}s -> {latent_len} latent frames, {len(t_sched)} steps" + + (f", CFG={guidance_scale}" if guidance_scale > 1.0 else "") + ) + + # 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) + + # Text encoding + metas = f"- bpm: {bpm}\n- timesignature: {time_sig}\n- keyscale: {key}\n- duration: {int(duration)} seconds\n" + caption = 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 = tok(caption, truncation=True, max_length=256, return_tensors="pt") + lyr_tok = tok(lyrics_text, truncation=True, max_length=2048, return_tensors="pt") + + text_h = te.encode_text(cap_tok.input_ids.to(device)).to(dtype) + text_m = cap_tok.attention_mask.to(device).bool() + lyric_h = te.encode_lyrics(lyr_tok.input_ids.to(device)).to(dtype) + lyric_m = lyr_tok.attention_mask.to(device).bool() + + # 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) + print("Preparing conditions...") + enc_h, enc_m, ctx = model.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 = model.null_condition_emb.expand_as(enc_h) + + # Noise + gen = torch.Generator(device=device).manual_seed(seed) + noise_ch = ctx.shape[-1] // 2 + xt = torch.randn(1, latent_len, noise_ch, generator=gen, device=device, dtype=dtype) + + # Diffusion + print("Running diffusion...") + t0 = time.time() + 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 = model.decoder(xt, tt, tt, attn, enc_h, enc_m, ctx) + + if use_cfg: + vt_uncond = model.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()) + + print(f"Diffusion: {time.time() - t0:.2f}s") + + # VAE decode + print("Decoding audio...") + t0 = time.time() + wav = vae.decode(xt.transpose(1, 2)) # [1, 2, samples] + wav = wav[0, :, : int(duration * SAMPLE_RATE)] + print(f"VAE decode: {time.time() - t0:.2f}s") + return wav.cpu().float() + + +# ═══════════════════════════════════════════════════════════════════════════════ +# CLI +# ═══════════════════════════════════════════════════════════════════════════════ + + +def main(): + p = argparse.ArgumentParser(description="ACE-Step v1.5 standalone inference") + p.add_argument("--prompt", required=True) + p.add_argument("--lyrics", default="") + p.add_argument("--duration", type=float, default=30.0) + p.add_argument("--output", default="output.wav") + p.add_argument("--seed", type=int, default=42) + p.add_argument( + "--model", + default="base", + choices=["base", "turbo"], + help="Model variant (default: base)", + ) + p.add_argument( + "--checkpoint", default=None, help="Override path to AIO .safetensors" + ) + p.add_argument("--device", default=None) + p.add_argument( + "--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"] + ) + p.add_argument("--bpm", default="N/A") + p.add_argument("--key", default="N/A") + p.add_argument("--time-sig", default="N/A") + p.add_argument("--language", default="en") + p.add_argument( + "--steps", + type=int, + default=None, + help="Diffusion steps (default: 30 for base, 8 for turbo)", + ) + p.add_argument( + "--shift", type=float, default=3.0, help="Timestep shift (default: 3.0)" + ) + p.add_argument( + "--cfg", + type=float, + default=None, + help="CFG guidance scale (default: 3.5 for base, 1.0 for turbo)", + ) + args = p.parse_args() + + device = args.device or ( + "cuda" + if torch.cuda.is_available() + else "mps" + if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available() + else "cpu" + ) + dtype = { + "bfloat16": torch.bfloat16, + "float16": torch.float16, + "float32": torch.float32, + }[args.dtype] + if device == "mps": + dtype = torch.float32 + + lyrics = args.lyrics + if lyrics.startswith("@") and os.path.isfile(lyrics[1:]): + lyrics = open(lyrics[1:]).read() + else: + lyrics = lyrics.replace("\\n", "\n") + + # Model-specific defaults + is_turbo = args.model == "turbo" + ckpt = args.checkpoint or MODEL_PATHS[args.model] + steps = args.steps or (8 if is_turbo else 30) + cfg = args.cfg if args.cfg is not None else (1.0 if is_turbo else 3.5) + + # Timestep schedule + if is_turbo and steps == 8: + ts = TURBO_TIMESTEPS.get(args.shift, TURBO_TIMESTEPS[3.0]) + else: + ts = compute_timesteps(steps, args.shift) + + print( + f"ACE-Step v1.5 ({args.model}) | {device} ({dtype}) | seed={args.seed} | {args.duration}s | {steps} steps | CFG={cfg}" + ) + models = load_models(ckpt, device, dtype) + wav = generate( + models, + args.prompt, + lyrics, + args.duration, + args.seed, + args.bpm, + args.key, + args.time_sig, + args.language, + ts, + cfg, + ) + + os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True) + torchaudio.save(args.output, wav, SAMPLE_RATE) + print(f"Saved: {args.output} ({wav.shape[1] / SAMPLE_RATE:.1f}s stereo)") + + +if __name__ == "__main__": + main() diff --git a/extensions_built_in/audio_models/ace_step/src/pipeline.py b/extensions_built_in/audio_models/ace_step/src/pipeline.py new file mode 100644 index 00000000..def452a4 --- /dev/null +++ b/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 diff --git a/extensions_built_in/audio_models/base_audio_model.py b/extensions_built_in/audio_models/base_audio_model.py new file mode 100644 index 00000000..6860a1aa --- /dev/null +++ b/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)) diff --git a/extensions_built_in/captioner/AceStepCaptioner.py b/extensions_built_in/captioner/AceStepCaptioner.py new file mode 100644 index 00000000..8ecb1ff8 --- /dev/null +++ b/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"\n{caption}\n\n" + output += f"\n{lyrics}\n\n" + output += f"{analysis['bpm']}\n" + output += f"{analysis['keyscale']}\n" + output += f"{analysis['timesignature']}\n" + output += f"{analysis['duration']}\n" + output += f"{language}" + return output + except Exception as e: + print(f"Error processing {file_path}: {e}") + return None diff --git a/extensions_built_in/captioner/BaseCaptioner.py b/extensions_built_in/captioner/BaseCaptioner.py new file mode 100644 index 00000000..8216eca7 --- /dev/null +++ b/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() diff --git a/extensions_built_in/captioner/Qwen3VLCaptioner.py b/extensions_built_in/captioner/Qwen3VLCaptioner.py new file mode 100644 index 00000000..9c76f223 --- /dev/null +++ b/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 diff --git a/extensions_built_in/captioner/__init__.py b/extensions_built_in/captioner/__init__.py new file mode 100644 index 00000000..78b93a23 --- /dev/null +++ b/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, +] diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index 0ef7077c..295f0e5d 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/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( diff --git a/requirements_base.txt b/requirements_base.txt index 3d8d11f4..4e7611e9 100644 --- a/requirements_base.txt +++ b/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 diff --git a/scripts/caption_audio_dataset.py b/scripts/caption_audio_dataset.py new file mode 100644 index 00000000..8cddf713 --- /dev/null +++ b/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"\n{caption}\n\n" + f"\n{lyrics}\n\n" + f"{analysis['bpm']}\n" + f"{analysis['keyscale']}\n" + f"{analysis['timesignature']}\n" + f"{analysis['duration']}\n" + f"{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}>(.*?)", 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() diff --git a/toolkit/audio/album_artwork.py b/toolkit/audio/album_artwork.py new file mode 100644 index 00000000..c1ef4b39 --- /dev/null +++ b/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) diff --git a/toolkit/audio/make_video.py b/toolkit/audio/make_video.py new file mode 100644 index 00000000..6fc95561 --- /dev/null +++ b/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}") diff --git a/toolkit/audio/ostris_logo.jpg b/toolkit/audio/ostris_logo.jpg new file mode 100644 index 00000000..dc34ea71 Binary files /dev/null and b/toolkit/audio/ostris_logo.jpg differ diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 8e69d171..a5bdeabf 100644 --- a/toolkit/config_modules.py +++ b/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)) diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py index 2d96f1ee..1d5536e0 100644 --- a/toolkit/data_loader.py +++ b/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: diff --git a/toolkit/data_transfer_object/data_loader.py b/toolkit/data_transfer_object/data_loader.py index 5768ba16..36ed1766 100644 --- a/toolkit/data_transfer_object/data_loader.py +++ b/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 diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index 55bbff88..cb03d8d4 100644 --- a/toolkit/dataloader_mixins.py +++ b/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 diff --git a/toolkit/models/base_model.py b/toolkit/models/base_model.py index 9e4e8d1a..d75c5ebf 100644 --- a/toolkit/models/base_model.py +++ b/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 diff --git a/ui/prisma/schema.prisma b/ui/prisma/schema.prisma index 95c2ea19..f183e589 100644 --- a/ui/prisma/schema.prisma +++ b/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]) + } diff --git a/ui/src/app/api/audio/art/[...audioPath]/route.ts b/ui/src/app/api/audio/art/[...audioPath]/route.ts new file mode 100644 index 00000000..de187f9f --- /dev/null +++ b/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 }); + } +} diff --git a/ui/src/app/api/jobs/route.ts b/ui/src/app/api/jobs/route.ts index 9a03965c..11dd9527 100644 --- a/ui/src/app/api/jobs/route.ts +++ b/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); diff --git a/ui/src/app/datasets/[datasetName]/page.tsx b/ui/src/app/datasets/[datasetName]/page.tsx index a3409afb..85a4c649 100644 --- a/ui/src/app/datasets/[datasetName]/page.tsx +++ b/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
+
-

Job: {job?.name}

+

{title}

{job && ( @@ -103,16 +113,21 @@ export default function JobPage({ params }: { params: { jobID: string } }) { )}
- {pages.map(page => ( - - ))} + {pages.map(page => { + if (page.jobTypes && !page.jobTypes.includes(jobType)) { + return null; + } + return ( + + ); + })} {page?.menuItem && ( <>
diff --git a/ui/src/app/jobs/new/AdvancedJob.tsx b/ui/src/app/jobs/new/AdvancedJob.tsx index 7ea39f08..b893e62d 100644 --- a/ui/src/app/jobs/new/AdvancedJob.tsx +++ b/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(''); + const [hasError, setHasError] = useState(false); const lastJobConfigUpdateStringRef = useRef(''); const editorRef = useRef(null); + const monacoRef = useRef(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 ( - <> +
+ {hasError && ( +
+ )} - +
); } diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx index 4abd5f4c..4442c81b 100644 --- a/ui/src/app/jobs/new/SimpleJob.tsx +++ b/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[] | 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 ( <>

Dataset {i + 1}

-
+
)} - - - Flip X - - } - checked={dataset.flip_x || false} - onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].flip_x`)} - /> - - Flip Y - - } - checked={dataset.flip_y || false} - onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].flip_y`)} - /> - -
-
- -
- {[ - [256, 512, 768, 1024], - [1280, 1328, 1536], - ].map(resGroup => ( -
- {resGroup.map(res => ( - { - const resolutions = dataset.resolution.includes(res) - ? dataset.resolution.filter(r => r !== res) - : [...dataset.resolution, res]; - setJobConfig(resolutions, `config.process[0].datasets[${i}].resolution`); - }} - /> - ))} -
- ))} -
- + {!isAudioModel && ( + + + Flip X + + } + checked={dataset.flip_x || false} + onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].flip_x`)} + /> + + Flip Y + + } + checked={dataset.flip_y || false} + onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].flip_y`)} + /> + + )}
+ {!isAudioModel && ( +
+ +
+ {[ + [256, 512, 768, 1024], + [1280, 1328, 1536], + ].map(resGroup => ( +
+ {resGroup.map(res => ( + { + const resolutions = dataset.resolution.includes(res) + ? dataset.resolution.filter(r => r !== res) + : [...dataset.resolution, res]; + setJobConfig(resolutions, `config.process[0].datasets[${i}].resolution`); + }} + /> + ))} +
+ ))} +
+ +
+ )}
))} @@ -1039,13 +1110,7 @@ export default function SimpleJob({
-
+
-
- setJobConfig(value, 'config.process[0].sample.width')} - placeholder="eg. 1024" - min={0} - required - /> - setJobConfig(value, 'config.process[0].sample.height')} - placeholder="eg. 1024" - className="pt-2" - min={0} - required - /> - {isVideoModel && ( -
- setJobConfig(value, 'config.process[0].sample.num_frames')} - placeholder="eg. 0" - className="pt-2" - min={0} - required - /> - setJobConfig(value, 'config.process[0].sample.fps')} - placeholder="eg. 0" - className="pt-2" - min={0} - required - /> -
- )} -
+ + {!isAudioModel && ( +
+ setJobConfig(value, 'config.process[0].sample.width')} + placeholder="eg. 1024" + min={0} + required + /> + setJobConfig(value, 'config.process[0].sample.height')} + placeholder="eg. 1024" + className="pt-2" + min={0} + required + /> + {isVideoModel && ( +
+ setJobConfig(value, 'config.process[0].sample.num_frames')} + placeholder="eg. 0" + className="pt-2" + min={0} + required + /> + setJobConfig(value, 'config.process[0].sample.fps')} + placeholder="eg. 0" + className="pt-2" + min={0} + required + /> +
+ )} +
+ )}
- setJobConfig(value, `config.process[0].sample.samples[${i}].prompt`)} - placeholder="Enter prompt" - required - /> -
- { - // 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) => ( +
+ {Object.entries(sampleTagSection).map(([tagKey, tag]) => ( +
+ {tag.type === 'text' && ( + { + 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' && ( + { + 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' && ( + { + let taggedSample = { ...taggedSampleArr[i] }; + taggedSample[tagKey] = value; + setJobConfig( + objToTags(taggedSample), + `config.process[0].sample.samples[${i}].prompt`, + ); + }} + placeholder={`Enter ${tag.title.toLowerCase()}`} + /> + )} +
+ ))} +
+ ))} + + ) : ( + <> + {modelArch?.hasMultiLinePrompts ? ( + setJobConfig(value, `config.process[0].sample.samples[${i}].prompt`)} + placeholder="Enter prompt" + required + /> + ) : ( + setJobConfig(value, `config.process[0].sample.samples[${i}].prompt`)} + placeholder="Enter prompt" + required + /> + )} + + )} + +
+ {!isAudioModel && ( + { + // 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)`} - /> - { - // 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 && ( + { + // 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)`} + /> + )} { // Sort by label, case-insensitive return a.label.localeCompare(b.label, undefined, { sensitivity: 'base' }); diff --git a/ui/src/app/jobs/page.tsx b/ui/src/app/jobs/page.tsx index 8dfa0c70..817ad92a 100644 --- a/ui/src/app/jobs/page.tsx +++ b/ui/src/app/jobs/page.tsx @@ -9,7 +9,7 @@ export default function Dashboard() { <>
-

Training Queue

+

Queue

diff --git a/ui/src/components/AudioPlayer.tsx b/ui/src/components/AudioPlayer.tsx index 541716d2..93f1a588 100644 --- a/ui/src/components/AudioPlayer.tsx +++ b/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 { - 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(null); - const [metaArtist, setMetaArtist] = useState(null); - const [metaAlbum, setMetaAlbum] = useState(null); + // Album art: served by /api/audio/art endpoint (fast, server-side extraction) const [albumArtUrl, setAlbumArtUrl] = useState(null); - const albumArtBlobUrlRef = useRef(null); // WebAudio analyser const audioCtxRef = useRef(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}
- {effectiveSubtitle ? ( + {subtitle ? (
- {effectiveSubtitle} + {subtitle}
) : null} {err ? ( diff --git a/ui/src/components/AutoCaptionButton.tsx b/ui/src/components/AutoCaptionButton.tsx new file mode 100644 index 00000000..207ba5c1 --- /dev/null +++ b/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 ( + + + Auto Captioning... + + ); + } + return ( + + ); +} diff --git a/ui/src/components/CaptionDatasetModal.tsx b/ui/src/components/CaptionDatasetModal.tsx new file mode 100644 index 00000000..2701eaa8 --- /dev/null +++ b/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(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(objectCopy(defaultCaptionJobConfig)); + const [gpuIDs, setGpuIDs] = useState(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 ( + +
+ { + e.preventDefault(); + saveJob(); + }} + > +
+
+
+ { + handleCaptionerTypeChange(jobConfig.config.process[0].type, value, jobConfig, setJobConfig); + }} + options={groupedCaptionerTypes} + /> +
+ {showGPUSelect && ( +
+ setGpuIDs(value)} + options={gpuList.map((gpu: any) => ({ value: `${gpu.index}`, label: `GPU #${gpu.index}` }))} + /> +
+ )} +
+
+ { + if (value?.trim() === '') { + value = null; + } + setJobConfig(value, 'config.process[0].caption.model_name_or_path'); + }} + placeholder="" + options={selectedCaptionOption?.name_or_path_options || []} + required + /> +
+ {selectedCaptionOption?.additionalSections?.includes('caption.model_name_or_path2') && ( +
+ { + if (value?.trim() === '') { + value = null; + } + setJobConfig(value, 'config.process[0].caption.model_name_or_path2'); + }} + placeholder="" + options={selectedCaptionOption?.name_or_path2_options || []} + /> +
+ )} +
+
+ { + 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') && ( +
+ { + const intVal = parseInt(value); + if (!isNaN(intVal)) { + setJobConfig(intVal, 'config.process[0].caption.max_res'); + } + }} + options={maxResOptions} + /> +
+ )} + {selectedCaptionOption?.additionalSections?.includes('caption.max_new_tokens') && ( +
+ { + const intVal = parseInt(value); + if (!isNaN(intVal)) { + setJobConfig(intVal, 'config.process[0].caption.max_new_tokens'); + } + }} + options={maxNewTokensOptions} + /> +
+ )} +
+
+ + setJobConfig(value, 'config.process[0].caption.low_vram')} + /> + setJobConfig(value, 'config.process[0].caption.recaption')} + /> + +
+
+ {selectedCaptionOption?.additionalSections?.includes('caption.caption_prompt') && ( +
+ { + setJobConfig(value, 'config.process[0].caption.caption_prompt'); + }} + placeholder="Enter caption prompt" + /> +
+ )} +
+ +
+ + +
+ +
+ + ); +}; diff --git a/ui/src/components/DatasetImageCard.tsx b/ui/src/components/DatasetImageCard.tsx index 7f90c302..101baffa 100644 --- a/ui/src/components/DatasetImageCard.tsx +++ b/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 = ({ imageUrl, alt, + isAutoCaptioning, children, className = '', onDelete = () => {}, @@ -80,6 +82,16 @@ const DatasetImageCard: React.FC = ({ } }, [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 = ({ 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 = ({ controls /> )} - {isItAudio && ( + {isItAudio && !showAudioPlayer && ( +
setShowAudioPlayer(true)} + > + {alt} { + (e.target as HTMLImageElement).style.display = 'none'; + }} + /> +
+ )} + {isItAudio && showAudioPlayer && ( = ({ onBlur={saveCaption} >