Training MOSS Serenity

How the model was trained, how to reproduce it, and how to continue it with more people. Back to the model card.

Trainer: maltrainer, train_moss_image_finetune.py. Its own guides go deeper: docs/finetune-tutorial.md (a walkthrough from install to the finished model) and docs/moss-image-finetuning.md (every setting).

What was trained


Base	MOSS-Image Base, bf16 (`moss_image_base_bf16.safetensors`)
Trained	the Linear weights of the 42 joint blocks: 4.76B of 5.08B parameters
Frozen	the text adapter (its projection and 2 blocks), norms, biases, embedders, modulation, output head
Steps	24,000 at batch size 1, about 25 epochs
Resolution	512 px area, aspect buckets on a 16 px grid
Optimizer	Automagic3, fused (each weight steps inside backward, so no full gradient set is held), stochastic rounding into bf16
Learning rate	1e-5, pinned (`learning_rate` = `automagic_min_lr` = `automagic_max_lr`)
EMA	0.999, updated every 4 steps, no warmup, in CPU RAM; the released file is the average
Timesteps	MOSS-Image's final-pretraining distribution: logit-normal, mean 1.09, std 0.95, truncated at sigma 0.95, 1% uniform, on the 1000-step grid
Caption dropout	5%
Prior preservation	weight 1.0, 8 recorded inputs per image, every step
Regularisation	real photos of women and men, 0.2× the step
Hardware	one RTX 5090 (32 GB): ~1.3 s/step with prior and EMA, 24.9 GB peak VRAM

Why these choices:

The text adapter stays frozen (OpenMOSS's advice, and Fizgig's): it reads every prompt, so training it on a few hundred captions would change how all prompts are read.
A pinned 1e-5: left to adapt, Automagic3's rate fell to ~2.5e-7 on MOSS-Image's flat loss, where a step moves a bf16 weight by a fraction of one rounding step and the change came out as noise. 1e-5 is Fizgig's finetune rate.
EMA 0.999: at batch size 1 each step's direction comes from one image at one noise level, so the raw weights wander (the run's drift grew like a random walk). The long average keeps the learning and cancels the wandering. Previews of the raw weights changed look every few hundred steps; the EMA row changed slowly and steadily.
Prior preservation: before step 1 the trainer records the original model's predictions for each image with its trigger replaced by the class word ("photo of feliciaday, a woman ..." → "photo of woman, a woman ..."). Each step then holds the model to them on that caption, so "a woman" without a trigger stays the base model's. Over the run the prior term drifted down (0.015 → 0.011): the people did not leak into the class words.

Dataset

One folder per person, named after their trigger, with a prefix.txt that puts the trigger in front of every caption:

datasets/finetunes/moss/        <- dataset.path
  people/
    prefix.txt                 "photo of [folder]"
    feliciaday/
      001.jpg
      001.txt                  "a woman with reddish-brown wavy hair smiling slightly, ..."
    keanureeves/
      ...
datasets/reg/                  <- dataset.reg_path
  women/  men/                 real photos of the class, captioned

Captions are one natural sentence: who, expression, framing, pose, clothing, background, light. [folder] in prefix.txt becomes the folder name, so the caption trains as photo of feliciaday, a woman with ....
Balancing with dataset.folder_repeats: people who learned slowly ran at 3× per epoch, two who learned quickly at 0.5× (a random half of their images each epoch).
Encodings are cached next to the images (.maltrainer_cache/), so restarts and resumes only encode what is new.

The config

The relevant parts (paths shortened; the full template is maltrainer's templates/moss_image_finetune_template.yaml):

model:
  transformer_path: ".../diffusion_models/moss_image_base_bf16.safetensors"
  vae_path: ".../vae/flux2_vae.safetensors"
  text_encoder_path: ".../text_encoders/qwen3.5_9b_int8_convrot.safetensors"

dataset:
  path: datasets/finetunes/moss
  resolution: 512
  aspect_mode: bucket
  trigger: ''                # every image has a caption; a missing one trains on nothing
  default_caption: ''
  folder_repeats: {feliciaday: 3, keanureeves: 0.5}   # per person, as needed
  cache: true
  reg_path: datasets/reg
  reg_default_caption: a photo of a woman

finetune:
  trainable_prefixes: [layers.]
  linears_only: true

prior:
  weight: 1.0
  classes: {feliciaday: woman, keanureeves: man}      # every trigger: its class word
  samples_per_image: 8
  every: 1

train:
  batch_size: 1
  max_steps: 24000
  save_every: 0              # save on the s / q keys and at the end
  save_state: true
  timestep_distribution: logit_normal
  timestep_logit_mean: 1.09
  timestep_logit_std: 0.95
  timestep_cutoff: 0.95
  timestep_uniform_mix: 0.01
  optimizer: automagic3
  learning_rate: 1.0e-05
  automagic_min_lr: 1.0e-05
  automagic_max_lr: 1.0e-05
  automagic_fused: true
  ema_decay: 0.999
  ema_every: 4
  ema_warmup: false
  caption_dropout: 0.05
  reg_lr_multiplier: 0.2

sample:
  every: 500
  rows: [ema]
  prompts:
    - "photo of feliciaday, a woman smiling slightly, close-up with face visible, soft light from the front"
  steps: 30
  cfg: 3.0
  width: 1024
  height: 1024

Running it

python train_moss_image_finetune.py --config my_run.yaml --suggest-steps     # max_steps for 22 epochs per person
python train_moss_image_finetune.py --config my_run.yaml --dry-run           # check the schedule
python train_moss_image_finetune.py --config my_run.yaml                     # train
python train_moss_image_finetune.py --config my_run.yaml --resume-from output/.../<name>_state

Keys while training: s save and continue, q save and stop, p pause, e preview menu, h help.

How this run went

Two stages. The first 8,505 steps trained a smaller set: 482 samples per epoch, 136 regularisation photos, prior recorded for 270 images. The dataset then grew to 2,465 samples per epoch, with more people and 432 regularisation photos, and the run resumed from its state: weights, optimizer, EMA and step count carried on, with a fresh epoch. The prior for the 444 new images was recorded from the original base model first (~5 minutes), for 714 in total.
Likeness kept improving to the end: people present from the start were recognisable by step 3,000-6,000 and refined after; people added at step 8,505 improved clearly between 12,000 and 24,000.
Health: no NaN or inf gradients in 24,000 steps; loss flat by noise band, as expected for flow matching; prior term slowly falling.
Cost: ~1.3 s/step at 512 px with prior preservation every step and EMA. Each save is a 10 GB model plus a ~35 GB resume state (raw weights, EMA, optimizer).

Continuing it

To add people, photos or regularisation images to Serenity, there are two ways:

From the resume state (exact): only the person who trained it has it. Add the new folders and their prior.classes entries, keep model.transformer_path on the original MOSS-Image Base (the prior for new images is recorded from it), raise max_steps (it counts from the start of the run), and resume.
From the released weights (anyone): point model.transformer_path at serenity_base_v1_24000_bf16.safetensors and train as above. The optimizer and EMA start fresh, and prior preservation then records from Serenity itself, so it protects what Serenity does now rather than what the base model did.

Converting to INT8 ConvRot

The INT8 file was made from the bf16 one with maltrainer's scripts/convert_int8_convrot.py, run with ComfyUI's Python because it calls comfy-kitchen's own quantize_int8_convrot_weight:

<ComfyUI>/venv/Scripts/python.exe scripts/convert_int8_convrot.py serenity_base_v1_24000_bf16.safetensors serenity_base_v1_24000_int8_convrot.safetensors --reference moss_image_sft_int8_convrot.safetensors

The reference is any INT8 ConvRot file of the same architecture (here Comfy-Org's MOSS-Image SFT release): it decides which Linears are quantized (each block's qkv, o, ffn_up_gate, ffn_down: 168 in all) and supplies the blobs an INT8 release carries (each block's comfy_attention.config). Everything else is copied from the bf16 file unchanged. The result has the reference's exact layout (757 tensors), and the quantized Linears are within 0.88% of the bf16 weights on average (0.98% at most). The conversion takes seconds on a GPU.

Finetune, continue and extract LoRAs from the bf16 file; the INT8 one is for generating.

Extracting the LoRAs

The lora-extract files came from the difference between Serenity and its base:

python scripts/extract_lora.py moss_image_base_bf16.safetensors serenity_base_v1_24000_bf16.safetensors --rank 64 128 256

One SVD pass writes all three ranks (~4 minutes on the RTX 5090). By size, they drop 96% (r64), 93% (r128) and 87% (r256) of each layer's change, mostly the batch-1 noise described above.