Commit Graph
76 Commits
Author SHA1 Message Date
Jaret Burkett a949f21f8e Add ability to merge dataset buckets. Useful when using numerous different datasets. 2026-10-08 14:04:30 -06:00
Jaret Burkett 0cbac2c1f3 Add threadding to bucket resizing of a batch. Helpful to speed up large batch gathering. 2026-10-08 11:26:48 -06:00
Jaret Burkett 07abdbedc7 Add match_target_res for Qwen Image 2.1 so reference images are scaled to the target's pixel area with their own aspect ratio, consistently across text-embedding caching, training, and sampling, and key the embedding cache on it. 2026-09-22 13:57:49 -06:00
Jaret Burkett cc98c3b690 Add ability to train with just music or vocals on yue2. Doesnt work as expected though 2026-09-17 20:19:32 -06:00
Jaret Burkett db8dbd68c2 Add support for training VLLMs, omni models, and LLMs 2026-09-15 17:47:02 -06:00
Jaret Burkett (Ostris) 2a2ca26ee8 Merge branch 'main' into fix/dataloader-pin-memory 2026-09-05 08:08:39 -06:00
Jaret Burkett da79ebce99 Add D-OPSD as a distillation handeling option for MiniMax H3 ref2va 2026-08-26 11:06:33 -06:00
Jaret Burkett b96476a841 Allow overriding the batch size on a per dataset level 2026-08-21 12:56:04 -06:00
Jaret Burkett 7f9a142dfd Generate thumbnails for dataset items so things load faster and the ui is more stable. Especially good for videos. 2026-08-12 20:33:42 -06:00
Jaret Burkett 8c1a4082fd Allow images to work with auto frame count, and include images in video datasets if they exist. 2026-08-08 20:26:28 -06:00
Jaret Burkett 9021caa723 Merge branch 'dev' of github.com:ostris/ai-toolkit into dev 2026-07-29 17:50:50 -06:00
Jaret Burkett 3f8afcac7e Allow dataloader to encode first frame with the text embeddings if the model needs it. 2026-07-29 17:50:45 -06:00
Jaret Burkett 3d472de2f1 Reworked spawning on dataloader so windows and mac can use multiple data loader workers now 2026-07-29 17:47:04 -06:00
Jaret Burkett d51463ca52 Add ability to recover from a truncated image file in the dataset 2026-06-22 05:07:30 -06:00
Jaret Burkett 41157b460c Added ability to set the caption extention in dataset viewer, captioner, and trainer so one dataset can have multiple caption styles in different files with different extensions. Added dataset caption template for a blank ideogram 4 formatted template. 2026-06-06 08:32:24 -06:00
Jaret Burkett 6a8d9333b6 Improved the prompt handeling of ideogram4 model. Now used advanced prompts class to store them smaller and allow longer prompts 2026-06-04 13:00:23 -06:00
Booyaka101 0c8793ff47 fix: make pin_memory opt-in via dataset config (default off)
@zappazack found a real regression: pin_memory page-locks RAM that
NVIDIA's Windows driver wants for shared-memory VRAM-overflow fallback
when VRAM saturates. Page-locked pages can't be relocated, so the driver
falls back to PCIe thrashing - step times went from 1.5s to 20+s when
VRAM saturated on RTX 5090 + Qwen LoRA + 1024x1024.

The original "strict-improvement" claim was wrong. Convert to opt-in:

- Add `pin_memory: bool` field on DatasetConfig (default False)
- Only pin when both CUDA is available AND user opts in via config

Users with stable VRAM headroom can opt in for the transfer speedup.
Users at the VRAM ceiling stay safe with the default.

Refs #758.
2026-05-02 11:02:41 +08:00
Booyaka101 67e223b502 perf: enable pin_memory on CUDA DataLoaders for faster CPU->GPU transfer
torch DataLoader's pin_memory was never set, defaulting to False. Setting
pin_memory=True allocates batch tensors in page-locked memory, which lets
PyTorch issue async cuda host->device copies via DMA - typically 10-30%
faster transfer for image/video training tensors and can reduce per-step
overhead noticeably on Windows where num_workers is forced to 0 (see
6d31c6db / #770) and the main process pays the full transfer cost.

pin_memory is orthogonal to num_workers and is safe with num_workers=0,
so this benefits Windows/macOS users (who currently can't use multi-worker
loading due to pickling constraints) without changing any other behavior.
Gated on torch.cuda.is_available() so it's a no-op on CPU/MPS/ROCm-without
-CUDA-API setups.

Refs #758.
2026-05-01 23:07:43 +08:00
Jaret Burkett 78cf049c29 Add support for ACE-Step 1.5 and ACE-Step 1.5 XL. Also added dataset captioning through the UI. (#785)
* Base ace step 1.5 xl added. Generating, still wip on training and ui

* Base training code done

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

* Fix issue with quantizing ace step

* Add album artwork to samples with waveform.

* Cleanup logs

* Add album art endpoint to speed up album art loading

* Made an make video with artwork script

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

* Add prompt tagging system for special tagged models.

* prompt tagging processing for ui working.

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

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

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

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

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

* Added qwen3 vl as a captioner for images

* Have an indicator when a dataset is being captioned.

* Adjust the way caption jobs look in the queue

* Fix a few issues. Adjust defaults.

* Version bump

* Added ace step to the readme.
2026-04-09 15:02:03 -06:00
Jaret Burkett 171535833a Add Mac OS support for Apple Silicon (#770)
* Made an install script and auto updates env for mac

* GPU sensors and initial training working for MAC. Still WIP.

* Switch dataloader to single threaded until I can work around some mac pickeling issues.

* Get quantization working on mac

* Fix mac exclusive imports so they don't break other builds.

* Add mac instructions to the UI
2026-03-30 09:37:47 -06:00
Rodrigo Reis 54f4732c9b Fix the bug in temporal_compression data loader (#754) 2026-03-25 08:16:44 -06:00
Jaret Burkett 7f3309b291 Add support for audo frame count so datasets can have varrying length videos. Varous ltx 2.3 VAE optimizations such as removing tiling articacts, and doing frame split encoding to reduce vram on encoding/decoding. 2026-03-24 12:20:09 -06:00
Jaret Burkett 5642b656b9 Fix audio issues with ltx2 models. Silent codec fails now raised. Auto convert surround sound audio to stereo. Invalidate old caches just to be safe so they recache now. 2026-03-23 20:08:33 +00:00
Jaret Burkett 1ce2428722 Shrink text embeds to max token length for LTX-2. Drastically reduces cached text embedding sizes 2026-01-28 12:54:49 -07:00
Jaret Burkett bf2700f7be Initial support for finetuning qwen image. Will only work with caching for now, need to add controls everywhere. 2025-08-21 16:41:17 -06:00
Jaret Burkett bb6db3d635 Added support for caching text embeddings. This is just initial support and will probably fail for some models. Still needs to be ompimized 2025-08-07 10:27:55 -06:00
Jaret Burkett 77dc38a574 Some work on caching text embeddings 2025-07-26 09:22:04 -06:00
Jaret Burkett e5ed450dc7 Allow finetuning tiny autoencoder in vae trainer 2025-07-16 07:13:30 -06:00
Jaret Burkett 96ba2fd129 Added methods to the dataloader to automatically generate controls for line, mask, inpainting, depth, and pose. 2025-04-09 13:35:04 -06:00
Jaret Burkett 3d131fb27a Added a file signature check on the dataset size caching system to invalidate cached dimensions if the file changes. 2025-04-01 07:39:36 -06:00
Jaret Burkett 860d892214 Pixel shuffle adapter. Some bug fixes thrown in 2025-03-29 21:15:01 -06:00
Jaret Burkett b829983b16 Added ability to load video datasets and train with them 2025-03-19 09:54:26 -06:00
Jaret Burkett 5e663746b8 Working multi gpu training. Still need a lot of tweaks and testing. 2025-01-25 16:46:20 -07:00
Jaret Burkett bedb8197a2 Fixed issue with sizes for some images being loaded sideways resulting in squished images. 2024-10-20 11:51:29 -06:00
Jaret Burkett 0640cdf569 Handle errors in loading size database 2024-10-20 07:04:19 -06:00
Jaret Burkett af108bb964 Bug fix with dataloader. Added a flag to completly disable sampling 2024-08-12 09:19:40 -06:00
Jaret Burkett 6d31c6db73 Added a fix for windows dataloader 2024-08-11 10:48:24 -06:00
Jaret Burkett 22d2f6e28f Fixed issue with grad scaling 2024-07-20 08:21:57 -06:00
Jaret Burkett 58dffd43a8 Added caching to image sizes so we dont do it every time. 2024-07-15 19:07:41 -06:00
Jaret Burkett b96913d73c Improvements to dataloader 2024-04-27 09:28:28 -06:00
Jaret Burkett 5da3613e0b Bug fixes and minor features 2024-04-25 06:14:31 -06:00
Jaret Burkett b0a0f28191 Bug fixes 2024-02-25 08:28:29 -07:00
Jaret Burkett 1ae1017748 Bug fixes. added ability to use l1 loss. varous other tests and improvements 2024-01-31 06:30:54 -07:00
Jaret Burkett 645b27f97a Bug fixes with ip adapter training. Made a clip pre processor that can be trained with ip adapter to help augment the clip input to squeeze in more detail from a larget input. moved clip processing to the dataloader for speed. 2024-01-04 12:59:38 -07:00
Jaret Burkett 0f8daa5612 Bug fixes, work on maing IP adapters more customizable. 2023-12-24 08:32:39 -07:00
Jaret Burkett 1ee62562a4 diffirential guidance is WORKING (from what I can tell) 2023-11-07 19:24:12 -07:00
Jaret Burkett 07bf7bd7de Allow augmentations and targeting different loss types fron the config file 2023-10-18 03:04:57 -06:00
Jaret Burkett 38e441a29c allow flipping for point of interesting autocropping. allow num repeats. Fixed some bugs with new free u 2023-10-12 21:02:47 -06:00
Jaret Burkett 63ceffae24 Massive speed increases and ram optimizations 2023-10-10 06:07:55 -06:00
Jaret Burkett f73402473b Bug fixes. Added some functionality to help with private extensions 2023-10-05 07:09:34 -06:00