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

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: