*[Please keep in mind this is still WIP/Beta but latest versions seem to work well and I can train a Lokr with body training data in roughly 60 mins now using my 5090.]

Regular LoRA training is good at teaching a model your trigger word + your photos. What it’s bad at is making generated full-body shots look like you.

You can often get a decent face from captions and a small dataset. But body shape — height, build, proportions — often drifts. More steps and more images help a little; they don’t fix the core issue: normal training never checks “does this generated body match the real person?”

What I did

I forked ostris/ai-toolkit and added a second training pass that uses Meta’s SAM 3D Body to scan bodies in your reference photos and in images the model generates during training.

So the loss isn’t guessing from pixels alone — it’s comparing against an actual 3D body read of your subject (shape and proportions, not pose). Face matching uses a separate identity reward; body matching uses the scan.

Repo: github.com/CliffNodes/FedorAiToolkit

How it works (simple)

1. Stage 1 (~200 steps) — Normal LoRA/LoKr training on your photos (same as ai-toolkit).

2. Stage 2 (15–60 steps) — The model generates images, SAM 3D scans them, and training pushes the LoRA toward your scanned body (and face). You’re reinforcing what you actually look like, not what the captions imply.

Fewer reward steps = faster, rougher polish. More steps = stronger likeness, especially for body. The example configs use 60; you can dial it down to ~15 if you’re experimenting.

Works on Krea 2 today (LoRA and LoKr). Ideogram 4 support is currently WIP.

How long on an RTX 5090?

SetupTime (rough)Face only (no SAM)~30 minFace + SAM 3D body~60 min

(With fewer Stage 2 steps, body training can finish faster — at the cost of less refinement.)

Face-only is the fast path if you only care about portraits. Body scanning adds setup and time, but that’s what makes full-body likeness stick.

Try it

1. Clone FedorAiToolkit (not stock ai-toolkit).

2. Put your photos + captions in a folder (e.g. datasets/subject).

3. Start from config/examples/krea2_lokr_draft.yaml (or use the UI and turn on the DRaFT / reward stage).

4. Set your trigger word and point draft.reward.reference_images at the same folder.

5. Adjust draft.num_reward_steps (15–60) for Stage 2 length.

Face-only: set body_weight: 0 — no SAM install.

Face + body: you’ll need Hugging Face access to the SAM 3D Body model (gated). Steps are in the README.

python run.py config/examples/krea2_lokr_draft.yaml

Results on my runs: much more consistent full-body likeness than SFT-only LoKRs, without hand-tuning dozens of body captions.

Questions welcome — happy to help with setup.

What makes this implementation technically interesting?

Unlike a standard LoRA trainer that optimizes only against captioned training images, this fork adds a second optimization stage based on differentiable rewards. After a conventional supervised fine-tuning (SFT) pass, the LoRA is resumed and optimized using DRaFT-K, allowing gradients to flow through the final denoising steps of the diffusion process instead of treating image generation as a black box.

The reward itself is also more sophisticated than simple CLIP or pixel similarity. Generated images are evaluated using ArcFace embeddings to measure facial identity and, optionally, SAM 3D Body to estimate a canonical 3D body shape. Because the body reward compares neutralized body geometry rather than pixels or pose, the model is encouraged to preserve a person's proportions and physique instead of memorizing specific poses from the training set.

The implementation is also practical from a hardware standpoint. Rather than backpropagating through every diffusion step, it uses the DRaFT-K approach of differentiating only through the final denoising steps, making reward-based optimization feasible on enthusiast GPUs. It also resumes from a standard SFT-trained LoRA or LoKr checkpoint, so the reward stage acts as a refinement pass instead of replacing conventional training.

Overall, it's an interesting combination of supervised diffusion training, reinforcement-style reward optimization, identity embedding models, and differentiable 3D body estimation—all integrated into an existing AI Toolkit workflow rather than built as a standalone research prototype.