github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

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

* Base training code done

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

* Fix issue with quantizing ace step

* Add album artwork to samples with waveform.

* Cleanup logs

* Add album art endpoint to speed up album art loading

* Made an make video with artwork script

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

* Add prompt tagging system for special tagged models.

* prompt tagging processing for ui working.

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

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

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

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

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

* Added qwen3 vl as a captioner for images

* Have an indicator when a dataset is being captioned.

* Adjust the way caption jobs look in the queue

* Fix a few issues. Adjust defaults.

* Version bump

* Added ace step to the readme.
This commit is contained in:
Jaret Burkett
2026-04-09 15:02:03 -06:00
committed by GitHub
parent 9ca58e9aa2
commit 78cf049c29
54 changed files with 5589 additions and 536 deletions
Download Patch File Download Diff File Expand all files Collapse all files

119
toolkit/audio/album_artwork.py Normal file
View File

@@ -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)

149
toolkit/audio/make_video.py Normal file
View File

@@ -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}")

BIN
toolkit/audio/ostris_logo.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 22 KiB

6
toolkit/config_modules.py
View File

@@ -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))

11
toolkit/data_loader.py
View File

@@ -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:

19
toolkit/data_transfer_object/data_loader.py
View File

@@ -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

32
toolkit/dataloader_mixins.py
View File

@@ -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

4
toolkit/models/base_model.py
View File

@@ -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