KohakuMoss: MOSS-Image LoRA Training With ComfyUI Export

KohakuMoss is a finetuning framework for MOSS-Image-Base and MOSS-Image-SFT, published as version 0.1.0 on PyPI on October 18, 2026. It is built by KohakuBlueleaf, the author of LyCORIS, and it fills the part of the MOSS-Image toolchain that the model release and the community few-step adapters did not cover: training on your own images at native resolution, then exporting something ComfyUI can load without a conversion step.

The gap it fills

MOSS-Image arrived on October 16 with two checkpoints, layout control and same-day ComfyUI support, and community 4-step and 8-step LoRAs followed within a day. What neither shipped was a trainer aimed at this specific 5B diffusion transformer and its Qwen3.5-9B text encoder. Generic diffusion trainers assume a UNet or a Flux-shaped stack; MOSS-Image needs its own packed attention and its own key names. KohakuMoss supplies both: an independent MOSS-Image implementation, the LyCORIS adapter set, a data pipeline and the memory plans needed to run them on one card.

Adapters and full finetune

Adapter training covers LoRA and LoCon, LoHa, LoKr, DoRA, OFT, BOFT, GLoRA, IA3, DyLoRA and T-LoRA, all through LyCORIS, so its algorithms, save format and fused kernels come along unchanged. Dropout, LoRA+, bypass or merged forwards and targeting presets are settings rather than separate scripts.

Full finetune is supported as well, using Muon with AdamW and EMA after the MOSS technical report. It fits a single 24 GB card with the master weights in host memory, so the trade is system RAM rather than another GPU.

Running on 6 GB

Memory is handled as a set of layers rather than a single switch: block checkpointing, weight streaming, activation offload, gradient offload, and optimizer states on the host. Four named profiles span speed to minimal, kohakumoss-config picks one from the card it detects, and a catalog of 240 entries sits underneath for manual tuning. Adapter training reaches 8 GB cards, and the 6 GB layer fits when both latent and text caches are built offline. A 24 GB card still covers the full finetune case.

Precision follows the hardware: bf16 where it exists, fp16 on older cards with a residual fix, fp8 and MXFP8 compute on sm89 and newer for both the DiT and the text encoder in training and in inference, and int8 or fp8 weight-only quantization for a frozen base.

Kernels per architecture

The project ships its own kernels instead of depending on one stack: TileLang for sm70 and sm75, Triton for sm80 and newer as well as Intel Arc, with torch.compile and eager behind both. Varlen attention, the block operations and the fused linears each pick a backend per architecture, dtype and shape, measured once on the card and stored in a tuning table, so the same install is tuned differently on a V100 and an RTX 5090.

The model implementation runs packed variable-length batches: one flat token axis, with attention kept inside each sample. Against the official code the project reports an fp32 forward relative error of 4.7e-6 and a 50-step sampling LPIPS of 2.3e-6.

Data and resume

Input formats include kohya-style folders, jsonl, parquet, webdataset and native shards, with an incremental catalog that is cached after the first scan. Resize, bucket and crop planners handle native-resolution images, token batching keeps steps even, and latent and text caches can be built offline or filled during training. The sampling plan is a pure function of the data, the seed and the settings, so a resumed run replays the same order instead of approximating it.

Multi-GPU is Linux only: DDP under every memory profile including weight streaming, FSDP2, pipeline and a threaded pipeline.

Exporting to ComfyUI

The export step is what makes this more than another trainer. A run can produce a LyCORIS file, a merged model, a diffusers transformer, or a ComfyUI-format adapter written against ComfyUI's own MOSS-Image key layout, including the repack for its fused q/k/v projection. That file goes into models/loras and loads in a normal ComfyUI workflow, so a LoRA trained here does not need a third-party node or a conversion notebook. Export also works on any saved checkpoint, not only on the final step.

Measured speed

The release tables report 1 MP training on MOSS-Image-Base with a rank 32 LoRA in the balanced profile. On an RTX 5090 the run reaches 1.01 samples per second with no cache and 1.14 with both latent and text caches; an RTX 4090 goes from 0.49 to 0.73 and an RTX 3090 from 0.31 to 0.39. Step peaks sit near 6.6 GiB of torch-allocated memory in most of those rows, and the 5090 numbers include a resident 15.6 GiB text encoder, which is why its no-cache figure is the faster one.

The README example is closer to a real first project: a rank 8 fp16 style LoRA over an fp8 base, 100 steps at roughly 20k tokens per step, with the aggressive ckpt-ring4-actoff profile and the text encoder and VAE encoding online every step. On one RTX 4090 it holds 12 to 13 GB and produces about 0.85 images per second.

Configuration

Every script is a KohakuEngine script, so settings are module-level knobs and a config can layer presets on top of each other. kohakumoss-config writes a checked config for the detected GPU and the dataset, kohakumoss-config wizard asks for the GPU, algorithm, data and weights instead of taking flags, and fixed recipes such as lora_24g_bf16 can be extended with every knob listed.

Availability and license

KohakuMoss lives at github.com/KohakuBlueleaf/KohakuMossTrainer and installs from PyPI as kohakumoss. The install helpers detect the driver, pick the matching torch index and then run kohakumoss-doctor, which reports anything missing together with the command that fixes it. Python 3.10 or newer is required, with CUDA 12.6 or newer or an Intel Arc driver; Windows needs the Visual Studio Build Tools, and sm70 and sm75 cards need the CUDA toolkit for nvcc.

The license is the KohakuMoss License 1.0: Apache 2.0 plus one term, that a derivative work, package, product or service built on KohakuMoss must include "Kohaku" in its name. Models, adapters and images trained or generated with the framework are explicitly outside that term, so a LoRA or checkpoint produced with it can be named and licensed as its author wishes.

For the checkpoints themselves, see the MOSS-Image release coverage and the 4-step and 8-step adapters that made fast sampling practical.