Yuanzhi/sd3.5-meanflow

SD3.5-Medium MeanFlow Checkpoint

Stage 2 MeanFlow distillation checkpoint for Stable Diffusion 3.5 Medium.

Details

• Base model: stabilityai/stable-diffusion-3.5-medium
• Training: MeanFlow distillation (continuous JVP), 10,000 iterations
• Config: lr=3e-5, batch_size=16, grad_accum=1, guidance_scale=3.5
• Architecture: MeanFlowSD3Transformer2DModel — extends SD3Transformer2DModel with r_embedder submodule for two-time velocity v(x_t, t, r)

Usage

This checkpoint is used as the Stage 2 base for MeanFlow RL (Stage 3) fine-tuning. Load via FastGen or Flow-Factory:

import torch

ckpt = torch.load("0010000.pth", map_location="cpu")
# Navigate checkpoint structure:
# ckpt['model']['ema_1'] contains the transformer weights
# Keys are prefixed with 'transformer.' — strip for loading into MeanFlowSD3Transformer2DModel

Citation

MeanFlow: https://arxiv.org/abs/2504.12419