Perceptual LoRA Toolkit

An extension of AI Toolkit by Ostris that adds perceptual anchoring to LoRA training. The idea is straightforward: instead of training only on per-pixel match against your dataset, you also tell the LoRA to match specific properties of those images using pre-trained vision models. There's an anchor for depth (the geometric structure of a scene), one for facial identity, one for body proportions, and a face-suppression option for when you don't want faces baked in at all. The depth anchor is the most useful one in practice. It lets the LoRA pick up the shapes in your dataset without locking in the colors, textures, or lighting, so trained models stay sharper on small datasets and generalize better to new prompts.

Supported models: SDXL, FLUX.2 Klein 9B

Contents

• Perceptual Anchoring: depth, identity, body, face suppression
• Auto-Masking: body / clothing / subject masks for region-weighted loss
• Reg Dataset Semantics: how reg samples are treated in this extension
• Training Metrics: what gets logged each step
• Training Previews: what each anchor saves to disk
• Dataset-Tools UI: preflight passes for masks, depth, faces
• Examples
  • Sketchwave Style (single-image style LoRA)
  • Yoshitaka Amano Style (small-dataset style LoRA)
  • Handsome Squidward (single-image LoRA)
• Configuration Reference: every extension-specific config option
• Upstream: AI Toolkit by Ostris
• Installation

Perceptual Anchoring

The standard LoRA training loss is per-pixel MSE in latent space. It tells the model "match this exact image." On small datasets that turns into a strong instruction to memorize, which is why you often see washed-out colors, baked-in lighting, and "burn-in" (stippling, JPEG ghosts) showing up in every generation.

Perceptual anchors give the LoRA more targeted guidance. Each one is a frozen vision model that scores a single property of the generated image, like its depth or its facial identity, and the LoRA gets rewarded for matching the training images on that property alone. You pick which properties matter for what you're training.

flowchart TD
    LegendNote["∇ = gradients flow back<br/>along this edge during backprop"]
    GT([Training image])
    LegendNote ~~~ GT
    GT --> Encode[VAE encode]
    Encode --> Z0[Clean latent z₀]
    Z0 --> Noise[Add noise at step t]
    Noise --> Zt[Noisy latent z_t]
    Zt --> Model[/LoRA model/]
    Model <-->|∇| Zhat[Predicted z₀']

    Z0 -.-> Diff["Diffusion loss<br/>(MSE in latent space)"]
    Zhat <-.->|∇| Diff

    subgraph Perceptual["Perceptual anchor path (this extension)"]
        Decode[VAE decode]
        RGBp[Predicted RGB]
        Pp["Frozen perceptor<br/>(DA2 / ArcFace / ViTPose)"]
        Pg[Same frozen perceptor]
        Anchor["Perceptual anchor loss<br/>(compares predicted vs. clean ground truth perceptor outputs,<br/>not pixels)"]
    end

    Zhat <-->|∇| Decode
    Decode <-->|∇| RGBp
    RGBp <-->|∇| Pp
    GT --> Pg
    Pp <-.->|∇| Anchor
    Pg -.-> Anchor

    Diff --> Total((Total loss))
    Anchor --> Total

    classDef frozen fill:#e8eaf6,stroke:#3949ab,color:#1a237e
    classDef trainable fill:#fff8e1,stroke:#f57c00,color:#e65100
    classDef loss fill:#e8f5e9,stroke:#2e7d32,color:#1b5e20
    classDef anchor fill:#f3e5f5,stroke:#6a1b9a,color:#4a148c
    classDef legendNode fill:#fafafa,stroke:#bbb,color:#555,stroke-dasharray:3 3

    class Encode frozen
    class Model trainable
    class Diff,Total loss
    class Decode,RGBp,Pp,Pg,Anchor anchor
    class LegendNote legendNode
    style Perceptual fill:#faf5fc,stroke:#6a1b9a,stroke-dasharray:5 4,color:#4a148c

The anchor path (purple) is what this extension adds. Both the GT image and the LoRA's prediction go through the same frozen perceptor, and the loss is computed on its outputs (a depth map for DA2, a face embedding for ArcFace, a keypoint heatmap for ViTPose). Gradients flow back through the perceptor and VAE decoder, translating the perceptual loss into a latent-space update for the LoRA. The weights most strongly nudged are the ones whose latents most affected the property the perceptor measures (depth, identity, pose); others barely move. Loss splitting (described below) takes this further by running the diffusion-loss step and the anchor-loss step alternately rather than summing them every step.

Depth-Consistency Anchor

Tells the LoRA to keep the geometric structure of the training images while ignoring everything else. It separates "what's in the scene" (which is the LoRA's job) from "how it looks in this particular photo" (which can be left to the model's prior). Useful for:

• Subject LoRAs that generalize. The model learns the subject's shape and pose without baking in the outfit, lighting, or backdrop of each training photo.
• Style transfer that keeps scene composition but changes appearance.
• Reducing texture burn-in and stippling on small datasets. Depth doesn't reinforce per-pixel artifacts, so fine-detail memorization slows down a lot.

Powered by Depth-Anything-V2 (Small by default; Base or Large can be selected for stronger geometry).

Quick start:

depth_consistency:
  loss_weight: 0.1                       # default; 0 disables
  model_id: depth-anything/Depth-Anything-V2-Small-hf
  mask_source: subject                   # 'none' | 'subject' | 'body'
  loss_min_t: 0.0
  loss_max_t: 1.0
  preview_every: 100

The default of 0.1 is calibrated for DA2-Small (the default perceptor). If you switch to DA2-Large, drop the weight to around 0.001, since the larger model produces much higher-magnitude gradients and 0.1 will overpower the diffusion loss. DA2-Base sits between the two; start at 0.01 and tune from there. If outputs look washed-out, over-smoothed, or the LoRA seems to be ignoring color and texture, the depth weight is too high. Halve it and retry.

Per-dataset overrides (handy when different folders need different strengths):

datasets:
  - folder_path: /path/to/portraits
    depth_loss_weight: 0.2               # stronger structure on portraits
    depth_loss_min_t: 0.5                # only fire on noisy timesteps for this set

Ground-truth depth maps are cached automatically at job start, so the anchor adds no per-step preprocessing cost once training begins.

Loss splitting (on by default whenever depth anchoring is active). When the diffusion loss and depth anchor pull in different directions, having them fire on alternating optimizer steps instead of competing every step turns out to work better than running them together for almost every workflow we've tested. As of this version, the trainer turns this on automatically for every dataset whose effective depth-consistency weight is > 0, so you usually don't need to set anything. If you want to be explicit, you can flip it on or off globally:

train:
  loss_split: diffusion_depth   # force on for all datasets
  # loss_split: null            # force off everywhere
  # (omit the key entirely for autodetect, which is the default)

Or override per dataset, which always wins over the global setting:

datasets:
  - folder_path: /path/to/data
    loss_split: diffusion_depth

This separates structure-learning (depth) from appearance-learning (diffusion) into distinct optimizer steps. In practice it acts as a strong implicit regularizer against burn-in: fine-texture parameters update much more slowly than coarse-structure parameters, since the two losses only really agree on the latter. The autodetect default means turning on the depth anchor is enough to get the splitting behavior; you only need to touch this if you want the old summed-every-step behavior back.

Identity Anchor (ArcFace)

Keeps the trained subject's face recognizable across poses, expressions, and lighting. Useful when you're training on diverse appearances of the same person and the diffusion loss alone isn't enough to lock in identity. Recommended weight: 0.01 to 0.1.

Body Proportion Anchor (ViTPose)

Keeps body proportions (limb lengths, torso ratio) consistent with the training images. Useful for full-body subject LoRAs where the body shape should stay recognizable across generated poses. Recommended weight: 0.1 to 0.2.

Face Suppression

The inverse of the identity anchor: it tells the LoRA to ignore faces. The diffusion loss is downweighted (or zeroed) inside detected face regions, so the model doesn't learn to reproduce the faces in your dataset. Use this when training a style or clothing LoRA on a dataset that happens to contain people, and you want the style or outfit but not the faces.

Set face_id.face_suppression_weight between 0 (off) and 1 (full suppression). Per-dataset overrides are supported.

Quick-start config

depth_consistency:
  loss_weight: 0.1                       # primary anchor (DA2-Small default; use 0.001 for DA2-Large)
face_id:
  identity_loss_weight: 0.1              # secondary
  body_proportion_loss_weight: 0.1       # secondary
  face_suppression_weight: 0.5           # optional
  identity_metrics: true                 # log id_sim without applying loss

Per-dataset overrides let you tune each anchor for the dataset's content:

datasets:
  - folder_path: /path/to/portraits
    depth_loss_weight: 0.2               # stronger structure on portraits
    identity_loss_weight: 0.1            # stronger face on close-ups
  - folder_path: /path/to/fullbody
    body_proportion_loss_weight: 0.15    # preserve body shape
    face_suppression_weight: 1.0         # full suppression, don't learn faces

See config/examples/train_lora_flux_identity_24gb.yaml for a complete example.

Auto-Masking

Splits each training image into regions (body, clothing, and subject = body ∪ clothing) so different parts of the image can be weighted differently in the loss. Useful for:

• Subject LoRAs that should focus on the person, not the background. Set the background weight low and the body weight high.
• Clothing LoRAs that should learn outfit details while ignoring face and body.
• Letting the depth anchor focus on the subject instead of computing depth-consistency over the whole image, where most of the frame is usually background you don't care about.

Masks are generated per image at job start and cached.

subject_mask:
  enabled: true
  body_close_radius: 5

datasets:
  - folder_path: /path/to/data
    background_loss_weight: 0.3
    clothing_loss_weight: 0.7
    body_loss_weight: 1.0
    perceptual_restrict_to_body: true    # restrict perceptual anchors to body region

The depth anchor picks which mask it uses via depth_consistency.mask_source (subject, body, or none).

QC tiles for visual inspection are saved at job start and can be regenerated from the dataset-tools UI.

Reg Dataset Semantics

Reg datasets (is_reg: true) work the classic Dreambooth way: they're prior-preservation samples that train the model on generic non-subject images alongside your subject samples, so it doesn't forget how to make non-subject content while it's learning the subject. In this extension, reg semantics are tightened up:

• All perceptual anchors are turned off on reg samples. Only the diffusion loss fires, scaled by train.reg_weight.
• Subject conditioning is stripped. No clip-image or trigger-word injection.

The effect is that reg samples teach the model "produce sharp prior-distribution images" without contaminating any of the subject-specific anchors. The 50/50 reg/train alternation runs at the optimizer-step level, so the gradient stays clean under any accumulation setting. train.reg_weight (default 1.0) controls how strongly reg pulls vs. train.

Training Metrics

Every active loss is logged so you can see during a run whether each anchor is doing its job. The training UI shows live charts plus per-sample tooltips on each point (which images drove the loss this step).

Gradient-cosine diagnostic. When you suspect two anchors are pulling in opposite directions, this measures how aligned their gradients are. Cosine near +1 means they reinforce each other, near 0 means they're independent, negative means they're fighting. Off by default; enable with train.gradient_cosine_log_every: 50.

Training Previews

Visual previews are saved during training so you can see at a glance what each anchor is responding to.

Dataset-Tools UI

Before training, the web UI provides preflight passes that prepare the cached data each anchor needs:

• Depth preflight. Runs depth estimation across the dataset and shows visual QC tiles so you can spot bad masks or odd crops before they cost you a training run.
• Subject-mask preflight. Generates and caches the body, clothing, and subject masks with overlays for review.
• Face-detection preflight. Caches face bounding boxes and identity embeddings.

All three run as non-blocking background jobs. Start them and come back when they're done.

The scripts/sample_dataset.py utility builds a smaller dataset directory by sampling N random images (with their captions) from a larger source. Useful for building reg sets, running ablations, or making smoke-test datasets without copying everything.

Examples

Note on inference target. All example configs in this README and under config/examples/ are tuned for inference against the distilled model (Flux 2 Klein). If you plan to apply the trained LoRAs against the base (non-distilled) model instead, checkpoints in the 500–800 step range are usually closer to optimal than the 1000–1200 step range the configs save out.

Example: Sketchwave Style (single-image style LoRA)

Training a style LoRA from a single image. One training image, one caption, and the LoRA picks up an entire visual vocabulary.

Sketchwave is a specific look: sketchy graphite-style linework over warm cream paper, with restricted earthy palettes (olive-green, ochre, wine-red, sepia) and slightly painterly shading. There's one training image, a portrait. The goal is for the LoRA to apply the look to anything the base model can paint, including subjects with nothing in common with the portrait.

Dataset layout:

examples/sketchwave/dataset/
├── 1.webp     # single training image
└── 1.txt      # caption

The caption opens with the trigger phrase sketchwave style. and then describes the image in detail: figure, clothing, lighting, and an explicit enumeration of the palette ("warm cream-yellow background, tan-and-ochre skin, dark sepia-brown linework, dark brown-black hair..."). Calling out the palette is deliberate. In LoRA training, anything you describe in the caption stays controllable at inference, while anything you leave out becomes part of what the trigger word bakes in. Naming the colors here teaches the LoRA that those are content choices in this particular image, not the essence of sketchwave style itself, so the trigger applies later with whatever palette you prompt for.

The training image itself:

Full config is at examples/sketchwave/config.yaml. Key bits:

• LoKr, linear/alpha 32, conv/alpha 16, full-rank, factor 8.
• 1200 steps, batch size 1, gradient accumulation 2.
• Resolution 768.
• 1 image, num_repeats: 50.
• Depth anchor: weight 0.005, DA2-Large at input_size: 1400, mask_source: none.
• Loss splitting on the dataset (loss_split: diffusion_depth).

The interesting part is how the trained LoRA generalizes. None of these outputs share a subject with the training image. The LoRA carries the linework, palette, and paper-like shading onto identities and scenes with nothing in common with the training portrait, including an animal and an outdoor landscape:

To reproduce:

python run.py examples/sketchwave/config.yaml

Edit model.name_or_path in the config to point at your local Flux 2 Klein checkpoint first.

Example: Yoshitaka Amano Style (small-dataset style LoRA)

Training a style LoRA from a small dataset of 14 illustrations. With depth anchoring the LoRA learns enough of the artist's visual language to carry it onto subjects nowhere in the dataset.

Yoshitaka Amano is the illustrator behind the original Final Fantasy character art and a long-running body of solo watercolor portrait work. Flux 2 Klein 9B doesn't reproduce his look from a prompt alone; it defaults to generic anime or oil-paint stylings.

One illustration from the dataset is shown below to give a feel for what the LoRA is asked to learn: loose ink linework, watercolor washes, ornate costuming, hair drawn as long flowing tendrils.

Key bits (full config at output/amano/config.yaml after a run):

• LoKr, linear/alpha 32, conv/alpha 16, full-rank, factor 8.
• 4000 steps, batch size 1, gradient accumulation 2.
• Resolution 768.
• Depth anchor: weight 0.005, DA2-Large at input_size: 1400, mask_source: none.
• Loss splitting on the dataset (loss_split: diffusion_depth).

You can watch the depth anchor converge across training. Ground-truth pair (RGB | depth) first, then predicted pairs from an early step and a late step at a comparable noise level. Early on the predicted depth has heavy halo artifacts and doesn't track the figure cleanly; by the end it's a much closer match.

Ground truth (RGB | depth):

Early prediction (step 383, t=0.82), depth_consistency_loss: 17.17:

Late prediction (step 3941, t=0.81), depth_consistency_loss: 6.67:

None of these subjects appear in the training set. The LoRA carries Amano's linework, color treatment, and composition language onto subjects from very different IPs:

With a style dataset this small, the diffusion loss alone tends to overfit on the specific compositions of the training images; every output starts looking like a slight variation on the same handful of poses and figures. The depth anchor pushes the LoRA toward what's invariant across the artist's work (linework, paper texture, color treatment) and away from what's incidental (this exact figure, in this exact pose, against this exact background). Loss splitting reinforces the separation: the diffusion-step focuses on appearance, the depth-step on structure, and they only really agree on the high-level "this looks like Amano" signal.

Example: Handsome Squidward (single-image LoRA)

Training a subject LoRA from a single image. One illustration, one caption, and the LoRA learns a character the base model can't reliably reproduce.

Handsome Squidward is a side character from a single SpongeBob SquarePants episode. Flux 2 Klein 9B doesn't reliably reproduce him out of the box; prompts default to regular Squidward or a confused human-squid hybrid. We trained a LoRA on a single official illustration to teach the model what he looks like, then tested whether the trained LoRA could generalize to angles and contexts that don't exist anywhere in the source material.

Dataset layout:

examples/squidward/dataset/
├── 1.webp     # single training image
└── 1.txt      # caption

The caption: "a cartoon illustration of handsome squidward. he is standing confidently with his arms flared and his tentacle-hands on his hips. he is wearing a tight yellow shirt with a brown belt and gold buckle. he has four legs, two on each side close together. his legs are spread apart. there is a logo at the bottom of the frame."

Full config is at examples/squidward/config.yaml. Key bits:

• LoKr, linear/alpha 32, conv/alpha 16, full-rank, factor 8.
• 1200 steps, batch size 1, gradient accumulation 2.
• 1 image, num_repeats: 50.
• Depth anchor: weight 0.005, DA2-Large at input_size: 1400, mask_source: none.
• Loss splitting on the dataset (loss_split: diffusion_depth).

These settings were tuned for the single-image case and likely won't transfer to multi-image training as-is. We're still researching optimal settings for multi-image runs (depth weight, input_size, mask_source, and loss-split behavior all interact differently once the dataset has variety to draw on).

Each preview tile shows (GT RGB | GT depth | Pred RGB | Pred depth) side by side. At the start of training the predicted depth is unstructured noise; by the end it tracks the GT depth closely.

Early (step 21), depth_consistency_loss: 26.6:

Late (step 1199), depth_consistency_loss: 1.6:

The training image is a confident front-three-quarter pose. Generations from the trained LoRA hold the character identity in poses, framings, and contexts that don't exist in the source material:

The first output is a near-direct front view, and there is no front-view reference anywhere in the source material, let alone the dataset. Both outputs maintain the character's distinctive identity (chiseled face, squid morphology, the specific drawn-on aesthetic) while placing him in contexts the LoRA wasn't trained on.

From a single image, per-pixel diffusion MSE alone would just memorize the training photo. The depth anchor adds a structural objective that gets reinforced on the same one image, so the LoRA picks up a 3D-ish understanding of the character's shape that lets it interpolate to unseen angles. Loss splitting keeps the diffusion and depth gradients from interfering with each other, which dramatically reduces the texture burn-in that's the typical failure mode of one-shot LoRAs.

To reproduce:

python run.py examples/squidward/config.yaml

Edit model.name_or_path in the config to point at your local Flux 2 Klein checkpoint first.

Configuration Reference

Every extension-specific config option, grouped by the YAML block it lives in. Defaults shown match what you get if you omit the option entirely.

depth_consistency.*

The depth anchor.

face_id.*

The identity-related anchor losses.

Identity anchor (the loss that keeps face recognizable):

Body proportion anchor (ViTPose bone-length ratios):

Face suppression (the inverse anchor; tells the LoRA to not learn faces):

subject_mask.*

Auto-masking pipeline.

train.* (extension-specific additions)

apply_assistant_networks[] (stack frozen LoRA/LoKr)

Use this to apply one or more frozen adapters on top of the base model during training (stacked adapters), while still training your main network normally.

You can define this list either at the root of the YAML or under model.apply_assistant_networks. If both are present, the root-level value is used.

apply_assistant_networks:
  - path: "/path/to/lora.safetensors"
    strength: 1.0
  - path: "/path/to/lokr.safetensors"
    strength: 0.5
    type: lokr

Notes:

• Assistant networks are loaded before the trainable network, set to eval mode, and kept frozen (requires_grad = False).
• They stay active during training with their configured strength.
• They are not added to optimizer parameter groups and are not the target of checkpoint saves for your trainable adapter.

ComfyUI key compatibility

For Flux/Flux2 adapter loading, ComfyUI-style naming conventions are supported when loading assistant adapters, including:

• transformer.<module_path>.* (diffusers/simpletrainer style)
• diffusion_model.<module_path>.*
• unet.<module_path>.*
• lora_transformer_<module_path_with_underscores>.* (OneTrainer style)
• lycoris_<module_path_with_underscores>.* (LyCORIS/LoKr style)
• plain module paths without prefix (DiffSynth-style)
• text-encoder keys such as lora_te_*, lora_te1_*, lora_te2_*, lora_te3_*, and text_encoders.*

This is meant to match the Flux mapping behavior in ComfyUI's comfy/lora.py as closely as practical for training-time stacking.

Per-dataset overrides (datasets[].*)

Every entry in datasets: accepts these extension-specific overrides. null or omitted = inherit the global value.

---

Upstream: AI Toolkit by Ostris

This extension is based on AI Toolkit, an all-in-one training suite for diffusion models on consumer hardware.

Support the Original Author

Sponsor on GitHub | Support on Patreon | Donate on PayPal

---

Installation

Requirements:

• Python >3.10
• Nvidia GPU with enough VRAM for what you're training
• Python venv
• git

Linux:

git clone https://github.com/BuffaloBuffaloBuffaloBuffalo/ai-toolkit-perceptual.git
cd ai-toolkit-perceptual
python3 -m venv venv
source venv/bin/activate
# install torch first
pip3 install --no-cache-dir torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu126
pip3 install -r requirements.txt

Windows:

git clone https://github.com/BuffaloBuffaloBuffaloBuffalo/ai-toolkit-perceptual.git
cd ai-toolkit-perceptual
python -m venv venv
.\venv\Scripts\activate
pip install --no-cache-dir torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu126
pip install -r requirements.txt

For devices running DGX OS (including DGX Spark), follow these instructions.