github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity
1,235 Commits 9 Branches 0 Tags
df6ea4263dc004db0fb0cfcebf318640d2edec05
Commit Graph

140 Commits

Author SHA1 Message Date
Jaret Burkett
4a99ddabad Fix breaking change with diffusers qwen image 2026-06-26 10:04:35 -06:00
Jaret Burkett
5f04ae7ad5 Rework merge_network_on_save to handle dequantization on merging and saving for much more efficient full finetuning. 2026-06-25 13:19:38 -06:00
Jaret Burkett
7a089fd0d7 Add support for training directly on Krea2 Turbo with a training adapter 2026-06-23 20:32:02 -06:00
Jaret Burkett
a803611ec1 Enable tiling on vae when decoding with low_vram flag on krea2 2026-06-23 18:38:52 -06:00
Jaret Burkett
724e67d634 Set krea 2 to use new lokr format 2026-06-23 13:31:50 -06:00
Jaret Burkett
e20b42e84a Add offloading support for krea2 2026-06-23 10:43:36 -06:00
Jaret Burkett (Ostris)
99be3d96a2 Add support for Krea2 (#906) 
* Add support for krea2

* Update repo pointer to actual repo
 2026-06-23 09:18:14 -06:00
Jaret Burkett
60c1ac6a50 Add support for Boogu Image and Boogu Image Edit 2026-06-18 15:05:49 -06:00
Jaret Burkett
e886745051 Added additional information on addine new models and some additional gotchas 2026-06-18 15:04:50 -06:00
Jaret Burkett
ad49d4ef25 Update example to cover some common issues 2026-06-18 07:54:37 -06:00
Jaret Burkett
92814f9e6d Apply ideogram dynamic shifting to sampling 2026-06-16 18:52:16 -06:00
Jaret Burkett
178eb5fbbe Add unconditional lora support so Ideogram 4 inference will more closely resemble the full pipeline results. I pushed a finetuned unconditional lora to the hub as an adapter. 2026-06-16 13:27:43 -06:00
Jaret Burkett
f6c0104f25 Handle unconditional conditioning for ideogram 4 more in line with example code. 2026-06-16 10:23:08 -06:00
Jaret Burkett
86b19589a0 Update the Ideogram 4 prompt generation/parsing/ui to handle the updated format notes better. 2026-06-16 09:44:38 -06:00
Jaret Burkett
fcccc0fbd2 Add gradient checkpointing to ideogram4 vae 2026-06-15 10:00:40 -06:00
Jaret Burkett
c0d600b5d6 Allow using flash backend for ideogram 2026-06-13 15:17:20 -06:00
Jaret Burkett
c8cd78b1a4 Allow nested transformer block names for quantization, lora targeting, quantizing 2026-06-13 14:33:16 -06:00
Jaret Burkett
6c3b82696e Add support for PRX Pixel T2I 2026-06-13 09:47:53 -06:00
Jaret Burkett
a01c83073a Added an example model with docs so people and agents can add models easier. 2026-06-13 08:17:06 -06:00
Jaret Burkett
35588726de Fixed issue where a buffer was stuck on cpu when offloading ideogram4 2026-06-07 13:31:03 -06:00
Jaret Burkett
8968e41234 Set ideogram to use new loke saving format 2026-06-07 09:49:57 -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
Jaret Burkett
63b3181262 Added experimental support for Ideogram 4 2026-06-04 09:03:00 -06:00
Jaret Burkett
5e84bf0d0b Fixed issue with hidream-01 that could cause a weird nan state. Took forever to track down as it was 1 in 10 starts. 2026-05-28 12:48:24 -06:00
Jaret Burkett
ef57c1077c Change z image divisibility 2026-05-28 09:09:36 -06:00
Jaret Burkett
7354def271 Added decode latent method to qwen image model 2026-05-25 10:02:09 -06:00
Jaret Burkett
fbac1cb7f5 Dont force flash attention on hidream 01. Causes random issues and is slower. 2026-05-24 16:05:22 -06:00
Jaret Burkett
644a6f9246 Fix device casting for zimage in some instances 2026-05-23 10:37:33 -06:00
Jaret Burkett
c6ecc03ccd Fixed saving full model of z_image l2p for finetuning 2026-05-23 07:36:24 -06:00
Jaret Burkett
6102370df9 Add support for ZImage L2P 2026-05-22 14:58:28 -06:00
Jaret Burkett
e03c6e4dc9 Fix potential inconsistency with different attention mentods in hidream01 2026-05-13 09:08:11 -06:00
Jaret Burkett (Ostris)
fc4d6ebf39 Add support for fine-tuning Hidream O1 (#831) 
* Initial support for hidream. Lora keys likely need work

* Fix saving for hidream-o1

* Remove dependence on flash attention for hidream o1

* Fix gradient checkpointing for hidream o1

* A lot of fixes for hidream. Handle loading and saving as comfy model.

* Omit layers not used in comfy. Fix issue with lora loading keys in comfy

* Version bumpo
 2026-05-12 11:15:16 -06:00
Jaret Burkett
e9ab387dfd Fixed issue with qwen image edit models when using multiple control images when not caching text embeddings. 2026-04-30 11:59:12 +00:00
Jaret Burkett
deb409085a Made it possible to use the flux 2 small decoder VAE when setting the vae_path manually for flux2 models 2026-04-30 05:42:31 -06:00
Jaret Burkett
7ccec8ec2c Add checkpointing and a proper decode for flux 2 VAEs so they can be used with DFE 2026-04-30 04:27:13 -06:00
Jaret Burkett
7c4f18ce51 Fix ernie unpatchify 2026-04-17 12:03:39 -06:00
Jaret Burkett
8cb9649382 Add decode latent to ernie pipe 2026-04-17 12:01:09 -06:00
Jaret Burkett
ab1ee4df34 Hotfix some issues with Wan models caused by diffusers and transformers updates 2026-04-16 20:53:50 +00:00
Jaret Burkett
afb62b1fa5 Add support for Nucleus-Image 2026-04-16 13:09:10 -06:00
Jaret Burkett
dd7074a21f Fix issue with layer offloading on ernie 2026-04-14 19:19:02 -06:00
Jaret Burkett (Ostris)
3e0c904054 Add support for Baidu's ERNIE-Image (#793) 
* Add support for ERNIE Image

* change float64 to float32

* Version bump

* Update ERNIE defaults
 2026-04-14 09:45:12 -06:00
Jaret Burkett
233e292256 Added some experimental low step things for zeta 2026-04-13 09:37:34 -06:00
Jaret Burkett
9ca58e9aa2 Fixed offload and quantize order of ltx 2.3 text encoder. 2026-04-07 15:11:50 -06:00
M. Hofer
f213e3b1e5 Fix FLUX2 Klein load-time VRAM spikes on low-memory GPUs. (#726) 
Keep the transformer and Qwen text encoder off CUDA during initial load/quantization in low-VRAM mode so model startup avoids full-model OOM before offloading and quantization can take effect.

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Jaret Burkett <jaretburkett@gmail.com>
 2026-04-01 09:36:55 -06:00
Jaret Burkett
affa411edc Fixed an issue where Flux.2 model VAE can be left offloaded to CPU when encoding control images while caching latents 2026-03-29 09:49:10 -06:00
abionda-sc
4ef5cbe5bc Fixing bug where width and height are inverted for control image resizing (#707) 2026-03-28 13:00:32 -06:00
科林 KELIN
489b194231 Fix CPU/CUDA device mismatch in Klein edit control image encoding (#742) 
When training Klein models with a `control_path` (edit/kontext-style
paired datasets), `encode_image_refs()` returns tensors that reside on
the VAE's device (CPU, since the VAE weights are loaded via
`load_file(..., device="cpu")` and are never explicitly moved to the
training device).  Concatenating those CPU tensors with the training
latents (`packed_latents`) that live on CUDA raises:

    RuntimeError: Expected all tensors to be on the same device

Fix: move `img_cond_seq` and `img_cond_seq_ids` to the same device
(and dtype) as `img_input` / `img_input_ids` before concatenation.

Co-authored-by: HuangYuChuh <HuangYuChuh@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
 2026-03-25 11:45:38 -06:00
Jaret Burkett
dfde30f231 Fix issue with ltx2 custom te repo path 2026-03-25 09:50:18 -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
Rayane
99a4a5887b Fix Qwen attention mask crash with diffusers >=0.37 (#748) 
* Fix Qwen Image mask handling

* Fix Qwen attention mask crash with diffusers >=0.37

diffusers v0.37 (PR #12987) optimizes all-ones attention masks to None
in encode_prompt() when there is no padding. This breaks ai-toolkit's
Qwen extensions which call .to() on the mask unconditionally.

Fix: reconstruct the all-ones mask at the boundary (get_prompt_embeds)
right after encode_prompt() returns. This keeps the rest of the code
unchanged and works with both old and new diffusers versions.

Also removes redundant duplicate mask assignments in qwen_image_edit
and qwen_image_edit_plus.

Fixes #740
 2026-03-23 14:43:08 -06:00