SimpleTuner/moss-image-sft-anyflow-stage1

MOSS-Image SFT AnyFlow stage1

Forward MeanFlow pretraining. Stage 2 initializes from the final EMA adapter.

Status: training. Observed updates: 19,640 / 30,000.

Training recipe

This SFT-only run follows the Anima AnyFlow experiments, using the stage1-v2 and dmd-v5 settings. Both MOSS stages use rank-128 BF16 LoRA, zero LoRA/caption dropout, eight H100 GPUs, microbatch 2, gradient accumulation 1, and global batch 16. Both stages use the same 50/50 sampling mix of CC12M with structured long captions and e621 with its original captions. Each source is capped at 65,536 images. CC12M uses long_caption explicitly and preserves caption paragraphs; it does not substitute short captions. Random square crops match the Anima recipe. Dataset files are not redistributed here. The SFT model revision is fc5de0d56fc910cbca6f6625a70f5f582b26deb1. Training uses MOSS's native unshifted noise grid.

Stage 1 uses learning rate 6e-5 for 30,000 updates with per-step EMA decay 0.995. Stage 2 starts from stage 1's final EMA adapter, with fresh optimizer and scheduler state, then runs 1,500 updates at learning rate 2e-6 with per-step EMA decay 0.99. The guidance-fusion scale is 3; stage 2 uses real-score guidance 3 and co-trained forward loss. Validation uses 4 sampling steps, CFG 1, seed 42, at 512px and 1024px, with an unshifted MOSS schedule and native grid snapping.

See training_config.json, dataloader.json, and prompts.json. An external publisher uploads adapter checkpoints and previews while the trainers run.

Checkpoints

Stage 1 uses 512px buckets through update 20,000. At that restart, 2048px training buckets are added with an even 512px/2048px sampling mix, keeping the source balance at 50/50 and global batch 16. Stage 2 uses the same mixed-resolution datasets. Existing 512px samplers resume their checkpoint state; new 2048px samplers start fresh.

Schedule correction: the stage-1 run resumed from update 1,000 with the flow-schedule fix. Subsequent previews use the corrected inference grid.

Pre-correction previews remain available.

1,000, 2,000, 3,000, 4,000, 5,000, 6,000, 7,000, 8,000, 9,000, 10,000, 11,000, 12,000, 13,000, 14,000, 15,000, 16,000, 17,000, 18,000, 19,000

All validation previews include live and EMA comparisons at both resolutions.

Latest live adapter: update 19,000.

Inference

Use SimpleTuner's MOSSImagePipeline with the original SFT checkpoint and the stage-2 live adapter:

import torch
from simpletuner.helpers.models.moss_image.pipeline import MOSSImagePipeline, MOSSImageScheduler
from simpletuner.helpers.distillation.anyflow.scheduler import AnyFlowValidationScheduler

pipe = MOSSImagePipeline.from_pretrained("OpenMOSS-Team/MOSS-Image-SFT", torch_dtype=torch.bfloat16)
pipe.load_lora_weights("SimpleTuner/moss-image-sft-anyflow-stage2", flowmap_deltatime_type="r", flowmap_gate_value=0.25)
pipe.scheduler = AnyFlowValidationScheduler(MOSSImageScheduler(shift=1.0), num_train_timesteps=1000)
pipe.scheduler.install_pipeline_hooks(pipe)
pipe.to("cuda")
image = pipe("A friendly red fox in a sunlit forest.", height=512, width=512, num_inference_steps=4, guidance_scale=1.0).images[0]
image.save("moss-anyflow.png")