SimpleTuner/moss-image-sft-anyflow-stage2

MOSS-Image SFT AnyFlow stage2

On-policy DMD with co-trained forward MeanFlow. Use the live adapter for few-step inference.

Status: completed. Observed updates: 1,500 / 1,500.

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 100,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. Existing 512px samplers resume their checkpoint state; new 2048px samplers start fresh.

At update 25,000, stage 1 resumes its full checkpoint and adds 1024px datasets. The remaining 75,000 updates sample 512px and 1024px equally, with 2048px disabled, retaining the 50/50 CC12M/e621 balance and global batch 16. The 512px sampler states and caches are reused. The 1024px datasets use crop=false with aspect-ratio buckets, require a source area of at least 1,048,576 pixels, and use fresh samplers and VAE caches. Stage 2 initializes from stage 1's 100,000-update EMA adapter and uses the same two-resolution mix.

Undersized aspect-ratio buckets in the uncropped 1024px datasets use automatic per-bucket padding to form complete distributed batches.

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

250, 500, 750, 1,000, 1,250, 1,500

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

Latest live adapter: update 1,500.

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")