lvladikov's picture
Upload README.md
68c5641 verified
|
Raw History Blame
17.1 kB
metadata
base_model: krea/Krea-2-Turbo
base_model_relation: adapter
license: other
license_name: krea-2-community-license
license_link: >-
  https://huggingface.co/lvladikov/Krea2-Turbo-Distill-2step-LoRA/blob/main/LICENSE.pdf
library_name: diffusers
tags:
  - lora
  - text-to-image
  - distillation
  - step-distillation
  - distribution-matching
  - krea-2
pipeline_tag: text-to-image

Krea 2 Turbo β€” 2-Step Distillation LoRA

🚧 Work in progress. This page is the live account of the project β€” what the adapter is meant to do, how it is being trained, and where it stands. It is updated as the recipe changes. There are no shipped files yet; the samples and side-by-side comparisons will be added when there is a candidate worth showing.

A LoRA for Krea 2 Turbo that takes the model from its usual 8 steps down to 2 β€” Turbo's own weights and its own two sigmas, guidance 0.0, a quarter of the denoising passes β€” with the aim of matching what the 4-step LoRA delivers today, at every resolution it delivers it at.

  • 🎯 The bar β€” the 4-step LoRA's output quality, across the same 12 resolutions, in two steps.
  • πŸ”Œ Drop-in, no exceptions β€” a plain LoRA sampled by stock Euler at sigmas [1.0, 0.5128] in diffusers, ComfyUI or MLX. No custom sampler, no policy head, no per-step tricks. If the quality needs a special sampler it is not this project.
  • 🧬 Same shape as the 4-step adapter β€” rank 64 on the same 228 modules; a second adapter exists during training only and never ships.
  • πŸ“Š Distribution matching, not imitation β€” the training objective that finally moved this (see Method).
  • 🎲 The same 13,750 recorded teacher trajectories the 4-step adapter trained on, reused without a single teacher re-run.
  • πŸ–₯️ One RTX 3090, and a recipe shaped by its 24 GB.

Where it stands

lineage 4-step LoRA β†’ 2-step trajectory distillation β†’ distribution matching β†’ a spectral match against the teacher's own images on top
current run training with the spectral match at raised strength (The spectral match below); the first change since this page was written that improved the renders without damaging anything else
best result so far the current run's latest checkpoint, measured on the full 12-bucket sweep: at every size up to 1024Γ—768 it is at parity with the teacher's spectrum β€” fine texture and both grid bands within a few percent β€” where the checkpoint that opened this page ran 20 to 50 percent over; at the large sizes the excess is down by a quarter to a third but remains, rising with pixel count to about 1.4Γ— at 1440Γ—1440 on the hardest subjects (spacecraft, clay, birds); distance to the teacher and saturation unchanged everywhere; seed-to-seed variety at the teacher's level; no pairwise losses to the teacher at 1440Γ—1440; graded 7.8/10 at 512Γ—512 and 8.1/10 at 768Γ—1024 by the judge, unchanged. The same coherent 512Γ—512 and 768Γ—1024 renders as before, with the grain gone from portraits
known gaps at 1280Γ—1280 and above a fine grain remains on the most textured subjects, well below where it was but still above the teacher; at those sizes skin now reads slightly smoother and less saturated than the teacher's, and fur (the snow leopard) slightly softer at every size β€” the two-sided term pulling toward smooth where the push was pulling toward grain; faces at 512Γ—512 still carry a faint doubled contour

How I got here

The 4-step adapter closed its page with a promise: a 2-step LoRA as the next project, and a guess at the lever it would need β€” matching the teacher's distribution rather than its trajectory. That guess turned out to be the whole story.

The project began where the 4-step one ended, from its final weights, and ran the same recipe at two steps: progressive distillation on the recorded teacher trajectories, each student call covering four teacher steps, with the LADD-style critic as the finisher. Well into that run, every number had stopped moving and the pictures had a signature the numbers could not see: doubled contours on faces and limbs, soft fine texture, crowds averaged into translucent overlaps. Several variations followed β€” the critic re-weighted, judged per token, a heavier hand on the final call, the student's own first-step output fed into its second β€” and each traded one of those faults for another without moving past them. A capacity probe ruled out adapter rank; a learning-rate shock ruled out the optimiser.

The reason is structural, and worth stating plainly because it decides the whole design. A regression loss asks the student to land on the teacher's specific image for each prompt. When a two-step jump is wide enough that several images are plausible, the answer that minimises the squared error is their average β€” and the average of two sharp images is a blurred one with doubled edges. Every earlier recipe rewarded that average. Tuning its weights could not change what it rewarded.

Distribution matching asks a different question: not "does your image match this one" but "would the teacher plausibly have produced your image". The first run of that objective, on top of the trajectory-distilled weights, produced in a fraction of the old recipe's training what all of it never had β€” and it did so while every latent distance to the teacher rose, which is exactly what a mode-seeking objective predicts and what a mean-seeking metric punishes. The distances are reported on this page; they are not optimised for, and they are not what decides a checkpoint. Pictures are, at fixed seeds, at every resolution, with faces viewed at 1:1.

The recipe adjustments so far, each made on the measurement of the one before:

  1. progressive distillation at two steps from the 4-step adapter's weights, with the LADD critic β€” the baseline
  2. the critic made to judge the first call's endpoint against the teacher's mid-states, which removed the gross ghosting and left the doubled contours
  3. the critic's weight, its per-token form and the student's own first-step output as the second call's input β€” tried one at a time; the objective flip that mattered was not among them
  4. distribution matching (the DMD2 family) as the primary objective, trajectory regression demoted to an anchor at half weight, the critic switched off to read the new term alone
  5. the running average of the weights restarted at the objective switch, after it was caught averaging two lineages into composites
  6. a spectral match of the student's image against the teacher's own, on the whole latent and on decoded pixel windows β€” the term that finally reached the grain at large resolutions, after a critic, per-resolution weights and a filtered push had each been tried against it and retired

Method

Distribution matching with a trajectory anchor, Krea 2 Turbo as its own teacher, on the recorded 8-step trajectories.

The student makes two calls, at Οƒ = 1.0 and Οƒ = 0.5128 β€” the first and fifth points of the teacher's 8-step grid at mu = 1.15 β€” and stock Euler carries it between them. That grid is what makes the objective a drop-in: Euler's first step from pure noise lands exactly on the flow-matching interpolant at Οƒ = 0.5128 with the same noise and the student's own clean-image prediction as the data point. So the student's first-call output is a legitimate image prediction that can be judged as an image, and the second call is fed from it during training the way it will be at inference.

The distribution term. For an image the student produces, two denoisers estimate how it should be cleaned up from a freshly noised copy: the frozen teacher, and a second small adapter on the same frozen base β€” the fake score β€” that is trained online to denoise whatever the student currently makes. Where the two disagree is the direction that makes the image more like the teacher's work and less like the student's habits, and the student is pushed that way (the DMD2 gradient, per-sample normalised). Averaging is never rewarded, so the student commits. The fake adapter is rank 32, starts as an exact copy of the teacher, updates twice per student step, and is discarded at the end.

The anchor. Plain trajectory regression on the teacher's recorded chords stays in at half weight. It keeps the student on the teacher's two-step grid so the distribution term cannot wander into a different sampler behaviour, and it is what the earlier trajectory distillation had already satisfied β€” which is why the first distribution-matching run moved so far so fast.

Per resolution. The distribution term's push grows with resolution: the fake adapter sees few large-bucket samples and under-fits fine structure there, and a per-pixel normaliser lands harder as pixel counts grow. A full 12-bucket sweep of the first distribution-matching checkpoint located the problem at 1 megapixel and above (fine-texture energy 1.4–1.8Γ— the teacher's at the five largest buckets); scaling the push per bucket from that measurement was tried and did not hold, and the spectral match replaced it. The same sweep of the current checkpoint, live weights, fixed seed, 15 prompts per bucket, every image measured against the teacher's render of the same prompt and seed (in brackets: the first distribution-matching checkpoint on the same prompts):

bucket fine texture vs the teacher 16-px grid band 8-px grid band distance to the teacher
512x512 1.01 (1.16) 1.02 (1.23) 1.04 (1.23) 0.44 (0.44)
512x768 0.97 (1.23) 1.06 (1.37) 1.01 (1.33) 0.40 (0.41)
768x512 1.02 (1.22) 1.10 (1.37) 1.07 (1.26) 0.47 (0.47)
768x768 1.10 (1.43) 1.09 (1.46) 1.13 (1.50) 0.41 (0.43)
768x1024 1.03 (1.36) 1.02 (1.38) 1.07 (1.42) 0.41 (0.43)
1024x768 1.03 (1.35) 1.06 (1.40) 1.04 (1.38) 0.45 (0.46)
1024x1024 1.10 (1.47) 1.13 (1.55) 1.15 (1.56) 0.43 (0.44)
1280x960 1.14 (1.50) 1.16 (1.59) 1.20 (1.56) 0.45 (0.42)
960x1280 1.27 (1.60) 1.25 (1.63) 1.32 (1.70) 0.44 (0.43)
1280x1280 1.24 (1.65) 1.25 (1.68) 1.26 (1.69) 0.44 (0.44)
1440x1280 1.31 (1.59) 1.32 (1.66) 1.39 (1.66) 0.42 (0.42)
1440x1440 1.46 (1.81) 1.39 (1.79) 1.51 (1.89) 0.42 (0.41)

The spectral match

Distribution matching delivers its push one latent token at a time, and a token is a 16Γ—16-pixel block. At large resolutions the push that tells the student to add detail also, unintentionally, adds detail in the shape of that block: a faint checkerboard at the 16-pixel and 8-pixel periods that reads as grain on birds, food and foliage. It is measurable in the images' Fourier spectra β€” the student carried about twice the teacher's energy at exactly those two periods β€” and several ways of reshaping or re-weighting the push (a critic, per-region weights, a filtered push) were tried and retired: none of them could reach a pattern that lives inside the token.

What can is a term that judges the picture, not the push. For every second-call sample, the student's estimate of the finished image and the teacher's own image for the same prompt and seed are compared through their radial power spectra β€” how much energy each has at every spatial scale β€” as an L1 on the log spectrum, in two places: on the whole latent every sample, and on one random 256-pixel window decoded through the VAE so colour and decoder behaviour are judged too. The comparison is two-sided, so too much fine energy and too little are both penalised; a blurred image does not satisfy it. Its gradient is added to the distribution push and capped per sample as a fraction of it, so it refines rather than takes over. At its first strength it brought every resolution closer to the teacher's spectrum without touching adherence, layout or variety; raised, it began closing the grain on the hardest subjects too, and that is the setting the current run trains with.

What the LoRA touches

Rank 64, alpha = rank, bf16, the same 228 modules as the 4-step adapter: the 8 attention and feed-forward linears of all 28 transformer blocks, plus the four global linears β€” time_embed.linear_1, time_embed.linear_2, time_mod_proj, final_layer.linear β€” that a step-count change needs most. Nothing about the base model changes.

Training data

The 13,750 recorded teacher trajectories of the 4-step project β€” Krea 2 Turbo's own 8-step run at mu = 1.15 and guidance 0.0, every latent and velocity stored β€” serve unchanged: a 2-step chord is two of the 4-step chords end to end. 203 held-out prompts measure the student–teacher gap on unseen prompts and never receive a gradient. The spectral match reads the teacher's finals for the training prompts; the 43,044 real-photo crops of the 4-step project are not used in the current run.

Resolutions

The same 12 buckets as the 4-step adapter, interleaved in proportion to their remaining samples:

512Γ—512 512Γ—768 768Γ—512
768Γ—768 768Γ—1024 1024Γ—768
1024Γ—1024 960Γ—1280 1280Γ—960
1280Γ—1280 1440Γ—1280 1440Γ—1440

Hardware

One RTX 3090 (24 GB). The frozen base is weight-only int8; the student's checkpointed block inputs stage to pinned host memory above 0.3 megapixels; the student, the fake adapter and the spectral term each build and free their own graph in turn, so their peaks never overlap; a hard memory ceiling sits below the driver's paging threshold so a step that does not fit fails loudly. A full step with every term live reserves about 21.5 GB at 1440Γ—1440. The price of the objective is throughput: about 160 training samples an hour against the 4-step recipe's 470 β€” three times the cost per sample, and so far a small fraction of the samples.

How it is judged

At regular intervals, both the live weights and their running average are pulled, merged and rendered at fixed seeds on 15 fixed prompts across four resolutions (512Γ—512, 768Γ—1024, 1280Γ—1280, 1440Γ—1440); milestone checkpoints get the same render at all 12 buckets, which is where the per-resolution table above comes from. Every image is measured against the teacher's render of the same prompt and seed: distance, fine-texture energy, the 16-pixel and 8-pixel grid bands, skin and flat-region grain, saturation, faces cut out at 1:1, a graded judge, a pairwise preference against the teacher, and seed-to-seed variety across four seeds per prompt. Latent distances β€” the held-out chord gap and the two-step rollout error β€” are recorded but never used to keep or stop a run: this project's clearest lesson is that they reward blur, and a run that improved them while its pictures collapsed was stopped by its pictures. The eye has the last word at every gate.

Notes and limitations

  • 🎯 Krea 2 Turbo only, at 2 steps, guidance 0.0 (cfg 1.0 in ComfyUI), mu = 1.15 β€” the two training sigmas are anchored to that grid.
  • 🚧 Not released. Everything on this page is the current state of a run in progress.

What's next

The current run trains with the spectral term at its raised strength. Its probes decide whether the remaining large-resolution grain closes; if the spectrum matches but structure still reads wrong, the next lever is a perceptual term on the same decoded windows, which judges what things look like rather than how much detail they carry. The shipping point, when there is one, is chosen by the sweep and the eye, not by distance run β€” the same discipline as the 4-step adapter.

License

The adapter is a derivative of Krea 2 Turbo and is covered by the Krea 2 Community License Agreement (LICENSE.pdf in this repository), as the 4-step adapter is. The attribution notice the license requires of a derivative ships with the files.