zjwfufu / SMD

Simplifying and Improving On-Policy Distillation 
 for Few-Step Diffusion
*SMD is a simple and scalable method for few-step diffusion distillation across image, video, audio-video and 3D generation.*
## Models
See the model card for Hugging Face weights, standalone inference and example outputs.
The configurations below use four sampling steps unless noted otherwise. Model names link to their training configs.
| Task | Models |
| --- | --- |
| Text to image | FLUX.1-dev, Qwen-Image, SD3.5-Medium, FLUX.2 Klein 4B, Lens |
| Image editing | Qwen-Image-Edit-2511 |
| Text to video | Wan2.1 1.3B and 14B, causal Wan2.1 1.3B |
| Audio-video | MiniMax H3 T2VA (4-step, 8-step) and Ref2VA (4-step, 8-step) |
| Image to 3D | Hunyuan3D-2.1 Shape, TRELLIS.2 Shape and Texture |
Additional configs cover SD3.5 reward specialists and multiple teachers, plus FLUX.2 9B distillation into Klein 4B and Lens.
## Installation
Requirements:
- Linux x86_64 and Conda.
- NVIDIA GPU with a driver compatible with CUDA 12.6.
- CUDA toolkit and C++ compiler for building FlashAttention from source.
The setup script creates an environment with Python 3.12, PyTorch 2.7.1, CUDA 12.6 wheels and FlashAttention.
Run from the repository root:
 [code] 
For Hunyuan3D mesh previews and TRELLIS.2, install the additional environment:
 [code] 
## Inference
For released weights, follow the standalone examples in the Hugging Face model collection.
For your own training checkpoints, export the student weights first, then run the inference script.
For example, with Qwen-Image:
 [code] 
Pass your exported model directory to `--weights` and choose where to save samples with `--output`. The model directory must contain both student weights and sampling metadata.
See the inference guide for other models and options.
> [!NOTE]
> H3 uses a student LoRA and cached conditions; no fused base checkpoint is needed. Export the adapter with the sampling metadata from the same checkpoint.
>
> Launching with `python` defaults to one GPU; use `--n-gpus` to choose the GPU count. See H3 export and inference.
## Training
Choose a configuration from the training experiments and prepare your data using the data guide.
For example, train Qwen-Image on two GPUs with a Parquet file containing a `prompt` column:
 [code] 
Replace the example data file and output directory with your own paths.
The `--download-resources` option downloads required resources automatically. Parameters after the standalone `-` override the selected configuration.
The `SMD_OUTPUT_ROOT` environment variable sets the root directory for checkpoints and previews. W&B logging defaults to offline mode.
See the training guide for run directories, resume, logging and preview settings.
## Agent skills
To make SMD easier to use and extend, we provide agent skills for environment setup, training and inference, and model integration in the skills directory.
Each skill includes instructions and references that you or your coding assistant can follow.
## Ackno