dptlab-klein-lokr-subject5

Post-trained with dptlab using the lora recipe with a lokr adapter, on top of black-forest-labs/FLUX.2-klein-4B.

Trainable parameters: 0.9M.

Training config

{
  "model_key": "flux2-klein-4b",
  "recipe": "lora",
  "dataset_path": "/root/data/syncd/subject-5",
  "output_dir": "/root/outputs/klein-peft/lokr/subject-5",
  "resolution": 512,
  "learning_rate": 0.0001,
  "train_batch_size": 1,
  "gradient_accumulation_steps": 4,
  "max_train_steps": 500,
  "lora_rank": 16,
  "lora_alpha": 16,
  "peft_method": "lokr",
  "max_grad_norm": 1.0,
  "use_masks": true,
  "mixed_precision": "bf16",
  "seed": 42,
  "checkpointing_steps": 500,
  "validation_prompts": [],
  "validation_steps": 500,
  "extra": {
    "sampling_steps_for_shift": 4,
    "logit_mean": 0.0,
    "logit_std": 1.0,
    "lokr_decompose_factor": 16
  }
}

Benchmark

Split: heldout (6 prompts). Prompts describe settings no training image shows.

CLIP-T (prompt)	DINO (subject)	CLIP-I (subject)	Avg. latency (ms)
0.9502	0.2946	0.6665	2104

CLIP-T and DINO pull in opposite directions: an adapter that learned nothing scores well on the first, one that memorized its training shots scores well on the second. Read them together.

See RESULTS.md in the repo for all six methods and the confounds.

Usage

import torch
from diffusers import DiffusionPipeline
from huggingface_hub import snapshot_download

# `lokr` is not a format `pipe.load_lora_weights()` can read -- it would
# log "no LoRA keys found" and silently serve the base model. Inject via peft.
from dptlab.training.peft_methods import load_peft_checkpoint  # pip install dptlab

pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.2-klein-4B", torch_dtype=torch.bfloat16).to("cuda")
denoiser = pipe.unet if hasattr(pipe, "unet") else pipe.transformer
load_peft_checkpoint(pipe, denoiser, snapshot_download("OnePunchMonk101010/dptlab-klein-lokr-subject5"), "lokr", for_training=False)

image = pipe(prompt="your prompt here", num_inference_steps=4, guidance_scale=1.0).images[0]