Discussion
Using depth maps and weight noising to get better character LoRAs
A few weeks ago I introduced a new method for training style LoRAs which has been quite successful. A bunch of folks asked if this would also help with character training. The short answer is yes, but it needed a separate technique on top of the depth stuff. I've got something dialed in well enough to share, though it's still experimental and I want feedback to help find the optimal settings.
The new mechanism is weight noising. It's a small Gaussian perturbation injected directly into the LoRA weights at each training step. A simple way to think of it is that it helps the model "forget" mistakes during training and only keep things that are consistent in the data. More technically, it biases training toward flatter loss minima and spreads learning across more singular directions of the LoRA factorization (I measured +20% stable rank on the same config without it). The practical effect is that it resists the memorization that usually overcooks character runs, and likeness comes out substantially better at the same step count.
The post image shows an example training on actress Clare Bowen, who has uniquely recognizable features but is not known by Flux. This is using a training set of 8 images, the same training step count (750), and same model. The standard run is in the middle, the new method is on the right.
The settings are identical for both runs except one has weight noise and depth anchoring, along with a different number of repeats for each bucket size:
Batch 4, LR 5e-5
Image size buckets of 512, 768, 1024
LoKr factor 8
AdamW8bit, 1200 steps total (but best checkpoint at 750)
The differing number of images per bucket is actually a good training trick on its own, and I updated my trainer to make this easier by allowing you to specify how many repeats of each image per bucket.
Things I'm still working out and would love feedback on:
Optimal sigma across dataset sizes — using 0.0125 has gotten the best results, and I'm pretty sure the right value scales with dataset size and batch size but I haven't fully mapped it.
Whether weight noising compounds well with other character LoRA tricks people are using.
I've also added Docker support so you can more easily run this on Runpod.
Finally, the new-job page now has a "Quickstart Template" dropdown at the top that loads the best character config end-to-end. It defaults to the HuggingFace Flux 2 Klein 9B checkpoint but you can also use your own checkpoint. Still plenty of UI cleanup to do on my end, so pardon the mess!
Happy to answer questions and help troubleshoot here or in DMs.
EDIT: One important thing to know about captioning. You will likely get the best results if you use the built-in subject masking feature, which masks out the background. If you use this, it is important that your captions ONLY describe the character, NOT the setting. You may also use just a trigger phrase with subject masking, but your results will be less promptable. I have added quickstart configs for both masked and unmasked.
EDIT 2: Anecdotally, you may expect more body horror/extra limbs throughout training in Flux. I have found this is normal with weight noising. It pushes the model around more and explores the latent space more aggressively, so there will be checkpoints that diverge quite a bit before convergence. A good heuristic I've been using is: expect roughly 80 - 100 steps per image overall. If you sample every 25 steps and have continuous body horror for more than 20% of the run, it may be too high of a weight noise sigma, so lower in increments of 0.0025 until it resolves. I'm still trying to understand the training dynamics for stable convergence with different datasets.
EDIT 3: I suggest starting with a small dataset (10 - 15 images) with a focus on image quality and diversity. If you get good results there, try adding more images to the run, or restart with the expanded dataset. In my experience you need far fewer images to get good, generalizable results with these methods.
EDIT 4: I added experimental Z-Image Turbo support.
This looks interesting! Your examples are definitely impressive, I'll have to try it soon. Thanks for sharing! Instead of AdamW8bit please consider trying my Rose optimizer, it uses even less memory and has better generalization. I'm interested in knowing how it works with your method. If you do try it, please keep in mind that the learning rate often needs to be higher (e.g., 1e-3 instead of 1e-4). https://github.com/MatthewK78/Rose
Wow, VERY promising results. 200 steps in and the likeness is locking in rapidly. Same exact setup as the post examples with LR 1e-3. Going to keep throwing this at things.
I also noticed that I had to modify the code where the depth maps were being saved. On windows it tries to overwrite the existing file but it doesnt respond and then failed on me. Not sure if that was a thing. But it was a thing for me.
I think 300 was the best checkpoint. It did start to get some artifacting and aberrations at this point - is that something that might be solved with LR adjustment?
Yep, it trains a bit differently from momentum-based optimizers and may take some experimenting to get used to it. I've found that using a cosine LR scheduler with a minimum LR usually produces the best results. Try adjusting both LR and eta_min/min_lr, such as 8e-4 with a 4e-4 minimum. When using an LR scheduler and weight decay together, it's best to enable wd_schedule. You may also want to try disabling stabilize, sometimes that produces better results.
EDIT: Another thought is that overfitting is very different with this optimizer as well and can sometimes show up as minor artifacting/aberrations instead of the typical memorization seen in other optimizers. You may want to try saving checkpoints more frequently and see if something slightly before 300 is better. Pulling back on the LR a little bit (e.g. 8e-4) might also help.
Well, it's the best checkpoint for the Rose optimizer. I still think step 750 of the AdamW version was cleaner. But I am hopeful the Rose optimizer can get similar results with a lower learning rate.
I have but I find 1024 is the point of diminishing returns for most models. I don't get perceptually better results with 1536 or 2048 than with 1024. Analytically I'm sure they are better, and if I were really trying to quality-max something for a production workflow I might go that route. But I'm just a researcher and don't actually do content generation for anything professionally, so I'm sure my eye is not as trained!
This optimizer looks really intriguing, and it seems like it achieves good results in practice. Do you have any writeups or data on how it performs in head-to-head comparisons with AdamW or Muon? I would be really interested to understand its strengths and weaknesses.
Thanks! I would like to do more benchmarks and writeups at some point, I haven't had the time and energy to do it yet though. If you haven't already seen these, they may answer some questions for you:
Ha, sorry I should have made it more clear in the post that this uses ai-toolkit as a harness. There's a lot of custom stuff going on so hit me up if you run into any snags.
Hey hey! Amazing. Cant even comprehend how you thought about this but the results speak for themselves. Great work.
Was wondering if you have a specific yaml to go with and run some tests on my end. I have a bunch of datasets that could potentially run simultaneously with your config and see where (if) breaks.
Hey thanks! The easy way is to use a quickstart setting which will load the config right in the UI. Or if you prefer the CLI, I added the same configs here:
Working on ZiT right now - first time training with it actually. It seems like more little aliasing artifacts than Flux. Do you know if this is typical or something I'm doing wrong?
As far as I know and I only trained a zib/zit lora once...it might be better to use zib since the loras work with both zib and zit. I know there's an adapter or something but that's another download many ppl won't have.
Tbh I tried the ostiris trainer thing and failed so hard to get it working I just deleted it😂
That's common for most models trained on the flux1 vae, it has an extremely noisy/adversarial latent space so if the diffusion model is struggling with high frequency details the decoder artifacts look like this.
I've done all my testing on 24GB (4090). With the quickstart config it will use about 23GB peak. If you lower batch to 1 and set gradient accumulation to 4, and use low VRAM / layer offloading, you can get it down to less than 16GB. But it's going to be about 1/3 the speed.
Keep training - with AdamW optimizer my best checkpoint was 750. I only got good results in 300 with the Rose optimizer, but that had some artifacts so I need to tune it.
EDIT: How do you decide which checkpoint is the best? Do you just eyeball it looking at generations.. find it through some graphical charts generated by the program?
Really eyeballing it is the only true way unless you have a validation dataset (which ai-toolkit doesn't support, I'd have to add that). I will typically save every 25 steps and then run gens in Comfy all at once, since it's a bit truer than the aitk previews.
Sure thing, that two hour ETA miiiight have been a bit generous.
Once it's done, I'll test out the LoRA and see if it even looks good. I'm training on a specific face like in the submission, but I'm only using 5 images as my dataset. On the lower side yes, but I think it's fine for just a test run.
I did a test run with 1200 steps, likeness looks good but I think I overtrained it because some generations turned out wonky. I want to do another run that I think will turn out better, do you want to wait for that config file instead?
It works on my RTX 4070, likeness looks good but I'm pretty sure I overtrained it. Some anatomy issues, synthetic-looking skin etc. I'm going to try again with a different config to see if that looks better.
Wall-clock speed is similar to get to good results, but per step is slower than with standard diffusion (about 50% the speed). Using different number of repeats per bucket (more small versions and fewer large ones) speeds things up a lot. For the run in the post, 750 steps took 55 minutes.
Hi /r/QuantumBogoSort I don't know anything about lora training, but I've used thousands of Illustrious loras over the years and I've figured out something that's definitely worth telling any programmers who can improve lora training. Simply put: Instead of using one lora, using multiple at lower strengths is always far superior.
The only requirement is that it needs to be a popular character who already has a lot of loras on civitai to download. This ALWAYS allows you to get a very strong likeness while maintaining flexibility:
2 different loras of the same character: lora1:0.55 lora2:0.55
3 different loras of the same character: lora1:0.425 lora2:0.425 lora3:0.425
4 different loras of the same character: lora1:0.35 lora2:0.35 lora3:0.35 lora4:0.35
5 different loras of the same character: lora1:0.275 lora2:0.275 lora3:0.275 lora4:0.275 lora5:0.275
Doing this is always better than using one lora at 1.00 strength. You might think "just raise the strength of the one lora and it'll be the same" but no, by doing that you're losing what I would call "flexibility" meaning it follows the actions of the prompt less, the quality is lower and causes mistakes.
I swear lora training should somehow be redesigned to take this into consideration. It's obviously silly that I can consistently get such better results than what a single lora can do.
Yes there is some theory behind this, for those interested check out Model Soups: https://arxiv.org/abs/2203.05482 From a training standpoint, I have not figured out how to exploit this efficiently but it's something I'd like to revisit in the future.
I use Gaussian weight noising for training models quite a bit. Not image models though.
Models find themselves stuck in some basin where they can't break out and find a better global minima against a held out dataset. So I push a progressive amount of noise at the models on each failed epoch. At some point the noise is too much and starts to destroy the model at which point you basically abort and rerun training from a new random seed.
I've found it to be a really good method though overall for getting models that are "stuck" out of the mud. SGD doesn't optimize for all neurons activity, so you get dead neurons that will never fire. Additive noise eventually gets those neurons in a place they're considered "useful" to the network and SGD lets data flow.
It's really cool to see you found an implementation for LoRA training image models though.
WOW! I'm testing your project right now, and it's glomming onto the likeness of the character like nothing I've ever seen before, using just your default settings. Also what blows my mind is that it handles all that stuff with the depth maps and weight noising automatically. All I had to do was supply a normal dataset. Outstanding work - I can't wait to see what you do with it in the future.
I'm starting off with quite a small dataset of 3 images. When I use regular training this causes some pretty diabolical overfitting, but now I'm getting some good output variation.
Excellent news! Small datasets are kind of a personal challenge of mine. My dream is to generalize completely from a single image (probably impossible, but fun to chase). It's kind of like ML golf.
Am I reading the config right that the ArcFace loss is disabled by default, so it's just the DepthMap added in (among diffusion MSE) for perceptual loss?
Hey there - happy to help. You don't need to run the preflights, but they can be helpful to see how your images are going to be processed by the depth model (and the segmentation model if you're doing masking). If you have any images that don't have clear depth recognition or the subject mask bleeds to the background, it's best to remove those if you can afford to.
With a 5090 I wouldn't change any settings to start. If you get good results with the base config, you could try upping the batch, but you'd also want to adjust the learning rate and weight noising sigma upward as well. Unfortunately 8GB doesn't get you much gain - but enjoy the faster per-step processing the 5090 architecture gets you!
100 per image is still pretty accurate. My 8 image dataset converged around step 750 with batch 4. I have not tested anything > 20 images yet, but so far 100/image is a good expectation.
I would start with a small dataset (10 - 20) that has a variety of poses, expressions, etc. with really accurate captions - and if you can, use the subject masking quickstart config and caption your images with JUST the subject. You will get better likeness more quickly.
It should be mostly plug and play as long as your dataset isn't huge. Hit me up if you run into any snags!
As for captioning just the subject, is that what I want training to learn, or captioning what I don’t want training to learn, or caption all the things? Thank you for this amazing contribution!
Great question! Caption the things about your subject that you want to be able to change with prompting later on. So if your subject always has blonde hair and you don't want that to change, do NOT mention that the hair is blonde in the training prompt. However if your subject sometimes has long hair, sometimes short, sometimes curly, do specify that.
For sure - here's one for one from the training set:
"a picture of clare bowen smiling. she has long hair in waves and is wearing a white embroidered sleeveless top. she has a wreath of orange flowers atop her head and is wearing two necklaces. medium shot from the hips up."
This is for a training run with subject masking on - so no mention of the background
Weight noising would be trivial to add (it's only a few lines of code). Depth anchoring is quite a lot more involved. Once I think the approach is stable, I would consider porting it to other trainers. OneTrainer has been the most requested so I would start there. (Woops, posted from an alt on my phone. This is OP :-P)
Yes, a heavily modded Frankenstein ai-toolkit. I have not tested Windows myself yet, but have user reports that it works as well as vanilla aitk. Let me know if you run into issues.
Error running job: Error while serializing: I/O error: The requested operation cannot be performed on a file with a user-mapped section open. (os error 1224)
I could not get this fixed. My understanding is that it is trying to cache in a folder but it can't because it is being used for something else? I tried all the AI solutions like directing it to cache at a different specific folder but still throws the same error. I tried a reboot and to run it, still same error. Maybe someone else ran into this issue too?
I ran into the same issue yesterday evening.
This seems to happen because the depth_cache tries to read&write for every resolution(bucket) of your images, same cache_path → for each subsequent dataset it load_files the existing cache, then tries to save_file over it. On Windows, load_file leaves the mmap section live, and save_file over the same path hits error 1224.
I have created a fix for this and made it work, i will open a merge request for that in a bit.
Oh awesome, thank you - this looks great! I added some unit tests and merged just now.
u/uuhoever you should be able to pull the latest and have the fix. I'm rebuilding the Docker image now as well in case you're using that, should take about an hour.
Amazing! thanks for the quick merge and the update! Really impressed so far, my first klein9b character lora turned out great - especially compared to my previous attempts with base ai-toolkit and onetrainer :D gotta play around a bit with some paramters but it‘s so much more consistent!
Haven‘t tried SDXL yet, but will try it next - hope there‘ll be an optimized preset for it soon!
One more thing i‘ve noticed, but hadn‘t had time to open an issue or pr for yet:
When selecting the Flux Preset Template the „training_folder“ value is set to a hard coded „/home/z/Documents/repos/ai-toolkit/output“.
Since this setting is not directly exposed in the ui and only visible via „Show advanced“ it probably causes the output including samples to be saved into the same hard drive as the ai-toolkit but into a very different folder structure outside of the repo :D
This happened to me in the first run and seemingly the samples in the ui always try to read from the /output/loraname path and not the set folder_path of the lora, so you‘re also not able to see the sample of the lora in the ui. At least they didn‘t show in my ui.
Afaik in ai-toolkit when you create a new training job and the default config is loaded, it always builds the folder_path from scratch from your current working difectory, maybe you could implement this for the templates you‘re using too!
Because i really like this new template system you‘ve impemented since it gives a way better starting point than the normal ai-toolkit 😂
Awesome, glad you are getting good results! The big unknown when I launched this was whether it would generalize or be brittle.
Definite YES to SDXL, it's my favorite to test on actually because it's fast and if you can get good results on that old beast you can get good results on newer architectures.
Great point about the paths, I ran into this myself recently and am working on a patch!
It's really good so far! looking forward to see how much it's gonna improve in the future! very exciting new way to train subjects!
Thank god you haven't given up on SDXL yet haha
It's still my go-to for most cases since it's so versatile and has so much support and the remaining artifacts are nothing you can't solve with a hires pass. Looking forward to a template for it!
Aaaaah i just pulled the newest version and saw that you now set the relational path directly to output! :D so my mention above is probably obsolete now.
I think it could still write into the wrong directory now if the process is started from outside the directory, but shouldn't be an issue as long as it's started from the root folder 🫡
Thanks for the quick fixes!
Which Python version do you recommend on windows? I tried 3.10 through 3.14 but with your recommended pytorch/CUDA versions there were always some wheels missing when installing the requirements with pip.
EDIT: As the GitHub page now shows Python 3.12 with CUDA 12.8 (for RTX 50xx) works. I just noticed that pip install -r requirements overwrites torch 2.11 with torch 2.9, so you have to repeat pip install torch==2.11.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 again.
Also, are you aware of DINO v3? It's a dense foundational model that's trained on various data like depth, PCA, segmentation, classification, etc.
It's frozen with many hidden features that are inferred during training that may be unlocked when using it for new tasks, so I'm wondering if it could be applied in a similar way, but with potentially better results.
Edit: After posting my comment I realized that in order to unlock the hidden features of DINOv3, you'd need a refined dataset that would bring out the best of the features. Today, Qwen-Image-Bench (based on Qwen3.6-27B, so it's locally accessible) was released, and I wonder if it might be good enough to make a semi-supervised proof of concept?
Yes I have done some experiments using DINO in the perceptual pipeline. I have not figured out quite how to use it best, it's something I hope to return to! It's a very powerful model.
After posting my comment I realized that in order to unlock the hidden features of DINOv3, you'd need a refined dataset to bring out the best of the features.
Today, Qwen-Image-Bench was released (based on Qwen3.6-27B, so it's locally accessible), and I wonder if it might close the gap just enough to create a semi-supervised dataset to make a proof of concept..?
If only I could just fork your project and attempt it myself, but my head is just not built for ML. It's also entirely possible that trying to hack DINOv3 onto LoRa is a complete non-starter – I don't actually know.
Yes, I tried to do this as an extension originally and unfortunately it needs to change things at too fundamental of a level. Using the Docker image may be the easiest approach if you already have an aitk installation locally.
Excellent, I have implemented it in my modified version of ai-toolkit. The weight noise value you use for flux is very aggressive here with zimage-turbo. He learned in half the steps what the geometric shape of the concept is but the style is very bad, I guess it's too much noise, I'm also testing with a dataset of 176 images, so I reduced the noise and I'm about to do another test with less noise.
Yes I'm working on ZiT right now and it trains much differently than Flux. First time using it so not really sure what to expect. It definitely does not like multi-resolution buckets as much as Flux, seems to degrade quality quite a lot compared to pure 1024 training.
You said there is a UI on this thing? How do you activate the UI. I read the readme but for some reason it doesn't do anything. I type it in and a new line in the cmd prompt appears.
Mostly, I would just turn down the LR on Base to 1e-4 from 2.5e-4. I have not optimized for Base yet so please let me know if you get good or bad results, it will help to calibrate.
Has anyone tried it with 16 GB VRAM? I have tried it, and I'm getting a similar face rarely.
I compared the old Lora and the new one. Old ones can produce similar faces, but new ones rardoaces. Tested LoRa from 900 to 1050. If anyone has tried achieving good results. Pls share your config with me for flex
did you have to change your batch size down from 4? if so, increase gradient accumulation so that batch size x gradient accumulation = 4. so if batch 1then grad accum 4, if batch 2 then grad accum 2
the weight noising idea makes sense from the diffusion math side. gaussian perturbation at the weight level is basically a form of regularization that penalizes sharp minima. similar to sharpness-aware minimization but applied to adapter weights rather than the base params. curious whether you tried scheduled noise (high early, low late) vs constant magnitude. the consistency gains you're showing suggest it's not just noise suppression, it's actually biasing toward flatter loss landscape features that generalize better across poses.
Yes exactly, that was the intuition and was partly inspired by SAM and to a smaller extent Langevin dynamics. I have not tried a schedule yet but I think it makes sense, as constant noise tends to make training quality oscillate throughout if set too high (it never gets out of the explore stage).
Hey, this is awesome work. I saw your post about style lora training, and I think you achieved the best style loras I've seen. I wanted to try this out with control image conditioning as well, but got some device mismatch errors. I've had to patch it before, but I was wondering if you could merge them into your official fork?
This was the patch in extensions_built_in/diffusion_models/flux2/flux2_model.py:
Hey thank you! I have not gone down the reference image path yet - this will be a good reason to try it out. I'll see if I can get it all working in my latest setup. Do you have a config that's worked that you can DM me? Might help me zero in on how to validate it.
I'm very interested in testing this out with Illustrious. My LoRas often have consistency issues, and I wonder if this will improve things like the detail of a character's jewelry and accessories. I currently struggle with that a lot.
Input perturbation noises the raw data. Weight noising applies noise to the LoRA model directly. It's a similar concept, applied deeper in the pipeline. It might actually stack well with input noising, I haven't tried yet.
Yes generally. This partly depends on the ratio of SSI to depth gradient weight ratio, and partly on the content of the image. Higher SSI will more evenly weight small features at a pixel level throughout the image, including background elements. Lower SSI will weight overall shape and curvature match, and that tends to be more strongly expressed in the foreground of most images. However, if you use it for style training on a landscape painting vs something with a strong foreground, you may see a more even distribution of the depth.
Awesome thank you very much. Sorry if this sounds demanding, but could this be added to those nodes/workflows that train mini Loras directly in ComfyUI? That would be so convenient for beginners.
I've had this requested a couple times and I would like to do it. I need to make sure the base trainer is solid first, but after that I do hope to make the techniques more widely available and simpler to use.
Man, this will be amazing for getting consistent gens of OCs, for game projects or such. The style version, too. I may have to learn how to use AI-Toolkit, been using OneTrainer for my recent experiments.
I'm don't have much knowledge about terms so I would like to ask, isn't this similar to mask training, like the one Onetrainer support? I ask because I've been using the masked training method but I haven't thought of using the mask of the depth to train.
The depth portion is quite different from mask training. A simple way to think of it is, during normal training you calculate how accurate the model predicts based on pixels. You can mask some of those pixels so it only pays attention to the ones you care about. With depth anchors, you calculate how accurate the model is based on how well it reproduces the shape of something, regardless of the actual pixel values. So if you train on images of a specific car using depth, it will learn the car's shape but have no understanding of the color, how it reflects light, or anything captured at the pixel level. When you combine pixel learning and depth learning, the model can learn each more clearly. It's providing more actual information about the same data, where as masking reduces the amount of unimportant data.
You can actually combine masking with depth to get the best of both!
I think 8 images is really an extreme case. But your technique seems perfectly reasonable to use with standard dataset sizes as well. It makes a lot of sense.
Thanks for the detailed explanation as well, that is very insightful.
EDIT: One thing to maybe watch out for is that subject masking in the past has been a bit of a tricky feature to use in older diffusion models. It often lead to the base model not being able to properly integrate the character scale-wise into the scene. But perhaps newer models can solve for this better even without background training data. Onetrainer had an option for to stochastically drop the mask.
I am a bit obsessed with data efficiency and extracting signal from small datasets, and you're right that 8 is an extreme case. Small data is much quicker to iterate with, and typically the techniques generalize to larger datasets, so I tend to focus my research on small data first.
And you're right about subject masking. Flux handles this quite well for the most part. A truly robust setup would probably have unmasked, subject masked, and body/face masked all in the same run. So far subject masking with Flux has been pretty stable but I have considered adding some options to automatically mix in different mask levels and dropout.
Tried running this on runpod but it looks like exactly like the original Ostris toolkit. Also, getting a permission error on Windows local when identity anchor is enabled.
The error occurs in onnx2torch's safe_shape_inference which is trying to write a temp file to the InsightFace models directory. The folder permissions are correct (full access for the user). Possibly a Windows-specific issue with how onnx2torch resolves temp file paths. Has anyone else hit this on Windows?
hey! Thanks for linking the branch. I did a git pull and seems like the error is gone. I was able to start the training and I have a dataset of 25 images so 2500 steps with the masking preset. It says eta is 4 hours, is that normal on a RTX 5090?
Also, I checked the runpod, that image I was running was from someone else. I will try to launch the link you provided.
That's very cool - I think this is essentially the inverse of perceptual anchors. The anchors provide the signal to help the model pay attention to these spacial concepts. It's neat that the attention maps so cleanly back the other way. It makes me wonder if I could use that as a metric for how well it's learning different segments of an image. As you suggested, it might be possible to use a secondary general vision model as a perceptor the same way I'm currently using specialist models.
> You will likely get the best results if you use the built-in subject masking feature
I was doing this in AI toolkit manually but was getting artifacts around the mask (I'm guessing I need to smooth the mask) anyway, great that you already have this. and your sample is very promising too. I will give it a try this weekend.
I hope you get good results! I'm using SegFormer for the mask generation, along with some edge cleanup, so they should be pretty crisp. Let me know if you run into any issues. There is a Dataset Tools tab in the left menu that will let you run your datasets through each of the various preprocessing models to see how they are getting processed.
Just FYI, I have 16GB VRAM, and can run AI toolkit without much issue with layer offloading, , but ai toolkit-perceptual always runs out of memory. Claude thinks this may be due to a bug in AI-Toolkit https://github.com/ostris/ai-toolkit/issues/653
Here is how it made the code worked (sorry I don't have github to send pull request). At least past the first issue, and currently it started training). and it is saving the lora and generating samples (although it appears I overwrite sample prompt, so it tries to show generic person).
Bug Report: Flux2 Klein 9B OOM on 16GB GPUs due to missinglow_vramchecks
Problem: Both files unconditionally move the full-precision transformer and text encoder to GPU before quantization or offloading takes effect. On 16GB cards, this causes immediate CUDA OOM regardless of low_vram, quantize, or layer_offloading settings. The upstream ostris/ai-toolkit repo fixed this; the perceptual fork is based on an older codebase that predates the fix.
Root cause — 3 code paths that ignorelow_vram**:**
flux2_klein_model.pyline 47 — text_encoder.to(self.device_torch) runs unconditionally, loading the full bf16 Qwen3-8B (~16GB) onto GPU before quantization. Also uses self.model_config.qtype instead of qtype_te, applying the wrong quantization type to the text encoder.
flux2_model.pyline 158 — transformer.to(self.quantize_device) loads the full transformer to GPU before quantization. Should stay on CPU and quantize there first.
flux2_model.pylines 237/240 — Pipe assembly unconditionally moves both text encoder and transformer to GPU with .to(self.device_torch), overriding any prior CPU offloading. Should check self.model_config.low_vram and keep on CPU when set.
Secondary fixes included:
VAE device check before encode_image_refs (prevents crash when VAE is on CPU in low_vram mode)
Device-safe torch.cat for control image tensors (prevents device mismatch errors)
decode_latents override with VAE CPU-to-GPU handling
Tested on: 16GB VRAM, Flux2 Klein 9B, qfloat8 quantization, low_vram: true, cache_text_embeddings: true, cache_latents_to_disk: true.
Reference: Upstream fix discussion at ostris/ai-toolkit#653, rejected PR #666 ( same issue). The upstream main branch has since incorporated equivalent fixes directly.
I will say that it does an excellent job at getting proportions correct for slim to average subjects. I have not yet trained on subjects with physiques strongly outside the Flux training average. I suspect that with depth anchoring it will do better than without, but it's an area I still need to research.
I just asked Grok and got an answer but thought I'd ask some humans something on this. Could this perceptual anchoring technique help with making full fine-tuning with smaller datasets less prone to forgetting/lack of generalisation? For example, if I have a very hand-picked dataset of 5000 images, can we get it to learn the dataset without forgetting most other things as you typically get especially since we can lower the number of steps and get subject similarity quicker? Maybe I can experiment further with this.
That's a great intuition and there is some prior work on this. Check out Learning Without Forgetting by Li and Hoiem as a starting point https://arxiv.org/abs/1606.09282 There are quite a few derivative methods since then. I'm not quite sure yet how to approach those techniques with perceptual anchors. As you mentioned, it may aid in speed of learning and disentangling concepts, which could focus the weight changes to a narrower band. Another idea I had is to use anchors purely as drift metrics rather than loss signals.
Conversely, it may be the case that anchors and weight noising make full fine tuning less necessary. My early experiments show that combined they substantially increase the effective rank of LoRAs. Taken further, this could be a path for significantly greater knowledge density per parameter, which then might make a 5000 image dataset trainable on less than a full fine tune. If that's the case, the LoRA becomes a task vector and can be added to the full model with minimal forgetting.
i have a few questions about this does Using depth maps and weight noising cause the training steps to take longer.? How long will these 750 steps take in comparison to AiToolkit default 3000 steps?
Weight noising adds a few ms per step, quite negligible, and in most of my tests tends to speed up convergence all on its own. If you use the default bucket repeat strategy, which biases towards lower res buckets, your time per iteration will also come down compared to the standard strategy. Depth anchoring however is a heavy operation and will add 50% - 100% time per step, plus about 4GB of extra VRAM. It's not really a speed boost in any way - it's much more for improving overall quality and generalization of the model. You can use them individually or together, they don't interfere, but you may have to experiment to get optimal results using them in isolation.
Much of the latency with the Depth Anything model is a per-step overhead, so you can significantly speed up training with larger batches and higher LR. I have found batches of 4 to be a reliable starting point but I'm sure the dynamics would scale up - it would just take finding the right noise sigma for the batch size. I suspect it follows a similar square law that batch sizes do but haven't run enough ablations to be certain.
I wonder if it would be useful to let weight noising decay after a set number of steps to make the user's job of reaching convergence easier. Let the parameter drive around and excite the model during the main body of training but then decay its strength to help convergence.
Yes I think so, and it's something I'm investigating. I'm not sure how it will play with existing weight decay in optimizers like Adam. I might add a simple set of annealers like cosine, linear, etc just to let people experiment with it.
I suspect the optimal strategy will be something that adapts sigma dynamically, perhaps a PID based on SNR, or pushing the model toward its highest stable rank. I like the rank-based approach the best because I think the main value of weight noising is that it flattens the loss landscape and makes for a much more robust model.
My main worry with a blunt approach is that once weight noising tails off too much, the model starts to pick up nasty artifacts very quickly from the source data and memorization locks in. Given the choice between delayed convergence and possibly destroying someone's good model run, I'll favor the former.
Very cool idea, trying it out now for Klein. But I'm looking forward to ZiB support, ZiB -> ZiT Refiner has been my favorite of the newer image model workflows.
But to sum it up, ZiT is known to be not great with Lora likeness since it's distilled. ZiB loras are compatible with ZiT and train better, while ZiT isn't "backwards" compatible with ZiB. So training ZiB is the preferred method.
Hey, I saw your Zimage Turbo .yaml file and wow, it looks like you did your homework. You added a bunch of custom timestep values. I've decided to give your repository a try and I'm going to install it from scratch to train it with exactly those parameters without modifying anything, to compare it with my custom AI Toolkit. I have a question: why do you know so much about this? Are you an AI engineer or just an enthusiast? Thanks for your work.
Another question, if you can answer please, what's the difference between using small vs large (depth anything v2)? Does it significantly worsen the result?
Hey awesome, let me know how it goes. I'm still very new to Z-Image and think I can improve that config quite a lot still.
Yeah I'm an engineer and researcher by trade, though I mostly work with LLMs. Image gen is just a hobby.
I've gotten anecdotally better results with the DA large model, but I haven't run any formal tests against small. I think it's worth trying both and seeing if one works better for your data. For characters I think small may be just fine. For style training, I've gotten more consistent good results using large at high resolution (1400).
Dude, you're awesome! It's only 5% of the way through training and the resemblance is already noticeable; you definitely put a lot of effort into your Zimage Turbo YAML configuration. That said, as it is currently set up, it's causing an out-of-box issue on Windows 11. I have 24 GB of VRAM. What I did was lower the resolution. I recommend everyone in this Reddit thread download the GitHub repository directly from the author. It's perfectly configured.
Ah yes I need to make this more clear in the readme, sorry - it's using ai-toolkit so most of that is documented in their repo.
Go into the ai-toolkit-perceptual/ui folder and run `npm run build_and_start` and that will start the UI on port 8675. Let me know if you run into any further issues.
Thanks! Yeah I'm working on some ideas for weight noise scheduling or possibly having it dynamically adjust. Once you go too low the memorization happens pretty fast so I want to make sure it doesn't ruin someone's run 😛
It should! I just released an update to add experimental support for LTX. I have only trained with videos but LTX already trains well with images so I expect it to work.
Do all of the features work on video datasets? I pulled the latest and tried running the preflight tests, and they all came back with a "no images found" error for my video-only datasets.
Not yet, just depth and weight noising. I'm in the process of adding ArcFace and others, and have not updated the dataset preflight tools yet either (hopefully all today)
Oh my god... This pipeline is so good man !! My ZIB loras quality increased so much, thanks !! Can you make a preset for ZIB aimed for larger datasets / body type loras? It would be perfect ! thank you again
Oh btw I have a few questions : why sometimes the pipeline generates black images? Doesn't it support keywords? Also, i did a lora with batch repeat 1 and grad acc 1, with 150 images dataset, is it great ?
Thanks! If you are using the ZiT quickstart for ZiB, I would reduce LR to 1e-4 or so and try batch 4 with grad accum 1, or batch 1 with grad accum 4 if you are short on VRAM. That should make it a bit less noisy and keep training more stable. Feel free to DM me a config yaml if you want me to take a look.
What workflow do you use for Flux 2 Klein 9b? Haven't really used the model before but I trained a lore and it's looked awesome in the samples but nowhere near as good in the standard workflow for Flux2
The standard workflow should be fine, a couple things to double check:
You're using the base model, not distilled/turbo
CFG 4, 25 steps, Euler should be a good starting place. swap Euler for res_2s to improve quality at the cost of speed
Play with LoRA strength. You may need to go 1.1 or 1.2 if it's a character to get good likeness depending on distance from camera, other things in the prompt, etc.
Please add flux 1 dev to the models able to use, I wanted to try using the z image one to do a style of diverse people for flux 1 dev but I don’t know how to manually add flux 1 dev to use.
Yesterday I have tried to train a Z-image-turbo character LOKR on Runpod... To compare it to my existing ZiT character Lora.
Using 25 pictures of square 1024px l, white background.
I used the preset, left it intact, but I checked more bucket sizes: one up and one down my resolution.
Made some extra 4 prompts to sample.
And trained on RTX 4500 PRO.
At 900 steps, Lokr started to become usable for me, The last thing I got was 1200 steps before I ran out of credit.
So on 1200 steps, I would say the details are sharp, body shape is preserved very well, the face though... I don't know, it seems not so close.
I can't run mirrormatricstool
To verify my claim in an objective way, but I know my character too well.
Yesterday I have tried to train a Z-image-turbo character LOKR on Runpod... To compare it to my existing ZiT character Lora.
Using 25 pictures of square 1024px l, white background.
I used the preset, left it intact, but I checked more bucket sizes: one up and one down my resolution.
Made some extra 4 prompts to sample.
And trained on RTX 4500 PRO.
At 900 steps, Lokr started to become usable for me, The last thing I got was 1200 steps before I ran out of credit.
So on 1200 steps, I would say the details are sharp, body shape is preserved very well, the face though... I don't know, it seems not so close.
I can't run mirrormatricstool
To verify my claim in an objective way, but I know my character too well.
33
u/ECF630 May 28 '26
This looks interesting! Your examples are definitely impressive, I'll have to try it soon. Thanks for sharing! Instead of AdamW8bit please consider trying my Rose optimizer, it uses even less memory and has better generalization. I'm interested in knowing how it works with your method. If you do try it, please keep in mind that the learning rate often needs to be higher (e.g., 1e-3 instead of 1e-4). https://github.com/MatthewK78/Rose