ShimaRinC/StepX-OPD-SD3.5

StepX-OPD (SD3.5-Medium)

Four-step text-to-image weights obtained by training Stable Diffusion 3.5 Medium with StepX-OPD, from the paper StepX-OPD: Rethinking Multi-Teacher On-Policy Distillation in Diffusion Models.

Files

Path	Purpose
pytorch_lora_weights.safetensors	Inference. A single LoRA that applies on top of stabilityai/stable-diffusion-3.5-medium and loads with diffusers load_lora_weights.
additive_init/	Training initialization. A PEFT adapter (adapter_config.json, adapter_model.safetensors) consumed by the StepX-OPD training code.

The root pytorch_lora_weights.safetensors folds the initialization LoRA and the trained adapter into one adapter, so its rank is the sum of the two (192 = 160 + 32). It already contains the additive initialization, so generating images needs nothing from additive_init/.

Training initialization

additive_init/ holds the SFT-augmented additive initialization: a PEFT LoRA for SD3Transformer2DModel with rank 160, lora_alpha 64, and 23 target modules, stored in float32 (1.32 GB). During training it is merged into the SD3.5-Medium transformer with PEFT merge_and_unload, and fresh rank-32 student adapters are then created on top of the merged weights. It is the starting point for training and is not meant to be used as a standalone four-step model.

The StepX-OPD training code takes this directory as init_lora_merge_path in configs/sd35m_wb.yaml, and its generation script takes it as --init_merge. To download only this directory:

huggingface-cli download ShimaRinC/StepX-OPD-SD3.5 --include "additive_init/*" --local-dir ./StepX-OPD-SD3.5

The adapter directory is then ./StepX-OPD-SD3.5/additive_init.

Usage

import numpy as np
import torch
from diffusers import StableDiffusion3Pipeline

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-medium", torch_dtype=torch.bfloat16
).to("cuda")
pipe.load_lora_weights("ShimaRinC/StepX-OPD-SD3.5")

image = pipe(
    "a corgi wearing sunglasses, studio photo",
    num_inference_steps=4,
    guidance_scale=1.0,
    sigmas=np.linspace(1.0, 1.0 / 4, 4),
    height=512,
    width=512,
    generator=torch.Generator("cuda").manual_seed(0),
).images[0]
image.save("out.png")

The model samples in four steps without classifier-free guidance, so guidance_scale stays at 1.0. The explicit sigmas reproduce the four-step timestep grid the model is evaluated with; omitting them falls back to the default grid of the base scheduler, which is a different set of timesteps.

Citation

@misc{wang2026stepxopdrethinkingmultiteacheronpolicy,
      title={StepX-OPD: Rethinking Multi-Teacher On-Policy Distillation in Diffusion Models}, 
      author={Chenru Wang and Xuanyang Zhang and Yucheng Han and Wenbo Nie and Wei Cheng and Xianfang Zeng and Gang YU and Chi Zhang},
      year={2026},
      eprint={2611.04386},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2611.04386}, 
}

License

This model is derived from Stable Diffusion 3.5 Medium and is subject to the Stability AI Community License Agreement. The full upstream license is provided in LICENSE-MODEL.

This model is fine-tuned by StepFun AI. Downstream users are permitted to use, modify, and distribute this model, subject to compliance with the Stability AI Community License Agreement.

NOTICE carries the required attribution and the statement of modifications, and additive_init/ carries its own copy of both files. The StepX-OPD source code is licensed separately under Apache-2.0 at https://github.com/WinterCCC/StepX-OPD.

Powered by Stability AI