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: queued. Observed updates: 0 / 1,500.

Training recipe

This SFT-only run follows the Anima AnyFlow experiments, using the stage1-v2 and dmd-v5 settings. Both MOSS stages train with 512px equivalent area, 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 10,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

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

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