LakonLab: A Training Codebase for AsymFlow, pi-Flow and GMFlow

Three papers, one repository, and the audience that implies

LakonLab exists to publish the reference implementations of three specific methods: Asymmetric Flow Models (AsymFlow), pi-Flow, and Gaussian Mixture Flow Matching Models (GMFlow). The README lists each with its venue or arXiv identifier, its author list, and a project page or demo link. That framing answers the first question an engineer asks about a research repository: who is this for. It is for people who want to run the code that produced a paper's results, or to build on the scheduler and training machinery those papers introduced. It is not pitched at teams looking for a general text-to-image training stack.

The three methods sit at different points in the same problem space. GMFlow appeared at ICML 2025 and concerns flow matching with Gaussian mixture source distributions. pi-Flow is an ICLR 2026 paper on few-step generation via imitation distillation. AsymFlow is the newest, dated May 2026 in the news section, and covers asymmetric flow models. The shared thread is flow-based generative modelling and the step-reduction problem: getting from a many-step diffusion or flow sampler down to something that runs in a handful of steps.

That shared thread is why the repository is one codebase rather than three. The schedulers, the training loop, the storage layer and the evaluation harness are common infrastructure, and each paper's directory supplies the model-specific pieces. If your interest is only in one of the three, you still install the whole package, because the abstractions are not split into separate distributions.

FlowSDEScheduler and the h parameter as the central abstraction

The most concrete piece of engineering described in the README is the solver family. FlowSDEScheduler is presented as a generic flow SDE solver with an adjustable diffusion coefficient. The README gives three settings for the parameter it calls h, and they are worth reading carefully because they define the method's range. With h=0 the solver reduces to a flow ODE. With h=1 it becomes a standard flow SDE. With h='inf' it corresponds to the re-noising sampler from the original consistency models work.

That single parameter spanning ODE, SDE and re-noising is the design idea. A researcher comparing deterministic and stochastic sampling does not switch codebases or reimplement a sampler; they change a value. The README states that this scheduler powers the GM-SDE solver in GMFlow, which is the clearest evidence that the abstraction is load-bearing rather than decorative.

FlowMapSDEScheduler is the companion for few-step flow map models, described as similar to the above. The README does not spell out where the two diverge, and that is a gap. A reader can infer that flow map models need different treatment than plain flow models, but the documentation does not say what changes. If you are choosing between them, the source files under lakonlab/models/diffusions/schedulers/ are the place to look, not the README.

The file paths themselves are informative. flow_sde.py and flow_map_sde.py sit under a schedulers directory inside a models/diffusions tree, which is the layout you would expect from a codebase that treats sampling schedules as pluggable components rather than hardcoded loops.

Weight tying, parallelism modes and why the training code is the underrated part

The README's feature list for the codebase itself is where the practical value sits, and it is easy to skim past. Two items stand out. The first is switching between DDP, FSDP and FSDP2, all supporting gradient accumulation and mixed precision. The second is weight tying for LoRA fine-tuning: the base weights of the teacher, student and EMA models are tied and share the same underlying memory, and this is stated to be compatible with both DDP and FSDP.

Weight tying matters here because distillation setups hold several copies of a model at once. A teacher, a student and an exponential moving average of one of them, each with its own full set of base weights, is a large memory bill before any optimizer state is counted. Sharing the base weights across all three removes that duplication. The README does not quantify the saving, and I will not invent a number, but the mechanism is stated plainly and it is the kind of detail that separates a codebase written to run at scale from one written to demonstrate an idea.

The parallelism story is similarly practical. DDP, FSDP and FSDP2 are three different answers to the same question, and the README claims you can switch between them. Gradient accumulation and mixed precision working across all three is the harder claim, since accumulation interacts awkwardly with sharded optimizers. The README asserts compatibility without documenting the configuration keys that control it, which is a real limitation for anyone who needs to reproduce a specific multi-node setup.

Installation: uv, PyTorch 2.10 and the gated-model login step

The README gives an explicit environment. Linux, tested on Ubuntu 20 and above, with PyTorch 2.6 or newer. The install command from the repository root is pip install -e . --no-build-isolation. The worked example uses uv rather than plain pip, and the sequence is short enough to quote in full.

cd <PATH_TO_YOUR_LOCAL_REPO>, then uv venv --python 3.10, then source .venv/bin/activate, then uv pip install torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cu128, then uv pip install -e . --no-build-isolation.

Note the version gap between the two statements. The prose prerequisite says PyTorch 2.6 or newer; the example pins 2.10.0 with a CUDA 12.8 index. The pinned pair is what the maintainers are actually running, and the 2.6+ line is the floor. If you are on an older CUDA toolkit, the index URL will need changing, and the README points to the PyTorch previous-versions page for that.

The --no-build-isolation flag is not incidental. It means the build runs against the environment you already created, so PyTorch must be installed before the editable install. Reversing the order will fail.

There is a second, easy-to-miss step. Accessing FLUX models requires accepting the FLUX.2 klein Base 9B conditions and the FLUX.1 dev conditions on HuggingFace, then running hf auth login. Without that, the model download paths will not resolve. The README also documents a checkpoint loading scheme that accepts huggingface://<HF_REPO_NAME>/<PATH_TO_MODEL> alongside local paths and HTTP or HTTPS URLs, which is a nice touch: a config can point at a hosted checkpoint without a manual download step.

Where LakonLab stops being the right tool

The README never claims to be a general framework, and the omissions are consistent with that. There is no list of supported base architectures beyond the models the three papers use. There is no configuration reference: the README mentions that parallelism modes and gradient accumulation are supported but does not show the keys that select them. There is no inference API documentation, no CLI reference, and no stated policy on backward compatibility between minor versions.

The release cadence reinforces the point. v0.2.0 and v0.2.1 landed on the same day in June 2026, and v0.2.2 followed in July. A repository moving that fast at the 0.2.x level is telling you the interfaces are still settling. If you build a production pipeline on internal LakonLab APIs, expect to track changes.

There is also a licensing-adjacent constraint that is not about the Apache-2.0 licence at all. The code is permissively licensed, but the FLUX weights it can load are gated and carry their own conditions that you accept separately on HuggingFace. Apache-2.0 on the repository does not grant you anything with respect to those weights. That distinction is worth being explicit about, though it is not legal advice and the terms themselves are the authority.

The realistic wrong-tool case is a team that wants a stable, documented trainer for a model architecture of their own choosing. LakonLab's value is concentrated in the flow solver abstractions and the distillation training setup. Outside that, you are adopting a research codebase's conventions without the documentation that would make them cheap to learn.

How it differs from the ComfyUI extension and from a general trainer

The README points to a separate repository, ComfyUI-piFlow, as the way to use these methods in ComfyUI. The news entries describe it as supporting 4-step sampling of Qwen-Image and Flux.1 dev using 8-bit models on a single consumer-grade GPU, and note that AsymFlow support was added in May 2026. That is the inference-side, end-user-facing path, and it lives elsewhere.

The division is clean. LakonLab is the training and experimentation codebase: schedulers, LoRA weight tying, parallel training modes, checkpoint I/O across local filesystems and S3, and an evaluation harness. ComfyUI-piFlow is the sampling front end for people who want to generate images with the distilled models rather than train them. If your goal is producing images, the extension is the shorter route. If your goal is reproducing a paper, training a distillation run, or reusing the solvers, LakonLab is the one you install.

Against a general-purpose diffusion trainer, the difference is the solver layer. A typical trainer exposes a fixed sampler and a noise schedule; LakonLab exposes a diffusion coefficient you can move between ODE, SDE and re-noising regimes, plus a separate scheduler for flow map models. That is a narrower and more opinionated surface, aimed at the specific research question of how the sampling process itself should be parameterized. It is not a broader tool. It is a deeper one in a smaller area.

Evaluation, storage and the cost of keeping up

The evaluation harness is more complete than the rest of the documentation. The README lists online metrics including FID, KID, IS, Precision, Recall, CLIP similarity, VQAScore, HPSv2 and HPSv3, all reachable through lakonlab/evaluation/metrics.py. It also supports exporting results to offline evaluators: the HPSv3 Benchmark, DPG-Bench and GenEval. For a research codebase, having both the online metric path and the export path in the same repository saves a meaningful amount of glue code, since these benchmarks are otherwise assembled by hand per project.

The storage layer is the other piece with operational consequences. Dataloaders and checkpoint I/O both support local filesystems and AWS S3, according to the README. Checkpoints can additionally be loaded from HuggingFace via the huggingface:// scheme or from HTTP and HTTPS URLs. The practical effect is that a training config can reference remote artifacts directly, which matters when the same run is reproduced on a different machine.

Maintenance cost is where a prospective adopter should think hardest. There is no stated support window, no deprecation policy and no changelog beyond the release tags. The news section functions as the de facto changelog, and it is written for people following the research rather than for people upgrading a dependency. Upgrading across a minor version means reading release notes and checking whether the scheduler or training entry points you rely on have moved. The Apache-2.0 licence means you can fork and pin, and for a research codebase with this release cadence, pinning a tag and reading the diff before moving is the realistic posture.

Editorial conclusion

Adopt LakonLab if you are reproducing or extending AsymFlow, pi-Flow or GMFlow, or if you need the FlowSDEScheduler and FlowMapSDEScheduler abstractions under a permissive licence. Do not adopt it as a general-purpose diffusion training framework: the README documents no supported model list beyond the paper implementations, no configuration reference and no API stability guarantee, and the repository moved through three releases in roughly a month. Before committing, verify that your PyTorch version satisfies the 2.6+ prerequisite, that you can accept the FLUX.1 dev and FLUX.2 klein gated-model conditions on HuggingFace, and that the specific paper directory you need has a README under docs/ rather than only an arXiv link.

Official sources

1. Lakonik/LakonLab on GitHub
2. License: Apache-2.0
3. Project website
4. README
5. Releases