SMD README

The project's README from github.com/zjwfufu/SMD. As one markdown file: README.md.

Simplifying and Improving On-Policy Distillation
for Few-Step Diffusion

Project Page arXiv:2610.18472 Hugging Face Model Collection Datasets 中文文档

SMD examples: images on the left, 3D assets in the center, and video sequences on the right.

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:

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:

bash environments/setup_conda.sh
conda activate smd

For Hunyuan3D mesh previews and TRELLIS.2, install the additional environment:

bash environments/setup_3d_environment.sh
conda activate smd_3d_extend

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:

python tools/download_resources.py --model qwen_image
python infer.py qwen_image --weights /path/to/exported_student \
  --output /path/to/samples --prompt "A red ceramic mug on a wooden table."

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:

export SMD_OUTPUT_ROOT=/path/to/training_runs
torchrun --standalone --nproc_per_node=2 train.py \
  --config smd/configs/experiments/qwen_image.py --download-resources \
  - dataloader_train.dataset_path=/path/to/prompts.parquet \
  log_config.name=qwen-image-my-data

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.

Acknowledgments and license

SMD builds on NVIDIA FastGen's training infrastructure and uses model implementations from Hugging Face Diffusers.

We are grateful to the authors and maintainers of these projects for giving us a foundation to build on.

We thank the authors of FLUX.1, FLUX.2, Qwen-Image and Qwen-Image-Edit, Stable Diffusion 3.5, Wan2.1, MiniMax H3, Lens, Hunyuan3D-2.1 and TRELLIS.2 for sharing their models and code. These releases made it possible to explore SMD across image, video, audio-video and 3D generation.

We also thank the authors of Self Forcing and Causal Forcing for sharing their code and checkpoints. Our causal Wan implementation builds on Self Forcing and uses the chunkwise initialization and teacher checkpoints from Causal Forcing.

See LICENSE, NOTICE and third-party licenses for code licensing and attribution.

Model weights, teacher adapters and datasets retain their upstream licenses.

SMD (zjwfufu/SMD)

pytorch
Searching Database…

205 ·GitHub
SMD
What is this?SMD is a research method and codebase for distilling diffusion models into fast few-step generators, covering image, video, audio-video and 3D generation tasks.
Is this a ready-made app or a research codebase?It is a research codebase of training and inference scripts with configs, not a packaged application — you run Python scripts against configs and exported checkpoints.
How do I get it running?On a Linux x86_64 machine with Conda, run the provided setup script from the repository root to build a Python 3.12 / PyTorch 2.7.1 / CUDA 12.6 environment with FlashAttention compiled from source; a separate extra environment is needed for Hunyuan3D and TRELLIS.2.
Do I need an NVIDIA GPU and CUDA toolkit?Yes — an NVIDIA GPU with a CUDA 12.6-compatible driver is required, plus the CUDA toolkit and a C++ compiler to build FlashAttention from source.
Can I run this on Windows or Mac?No — the README requires Linux x86_64 with Conda, and there is no mention of Windows or macOS support.
Is everything here under Apache-2.0?The SMD code itself is Apache-2.0, but model weights, teacher adapters and datasets keep their own upstream licenses.
Which models can I use out of the box?Released student weights on Hugging Face cover 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), audio-video (MiniMax H3), and image-to-3D (Hunyuan3D-2.1, TRELLIS.2).
How do I train SMD on my own data?Pick a training configuration and prepare a Parquet file with a prompt column following the data guide, then launch training across your GPUs (e.g. Qwen-Image on two GPUs); SMD_OUTPUT_ROOT sets where checkpoints and previews land, and W&B logging defaults to offline.
Is inference the same for every model?Mostly: export student weights plus sampling metadata and run the inference script with --weights and --output. MiniMax H3 is an exception — it uses a student LoRA with cached conditions instead of a fused checkpoint.
What are the agent skills for?The skills directory contains agent skills for environment setup, training, inference and model integration, written so you or a coding assistant can follow them to use and extend SMD more easily.
What are similar projects?EasyWAM, DiffSynth-Studio, kohya_ss, CodeVideoModel, ai-toolkit, vdn-minimax-h3, generative-models, hyperflow
Which topics is it in?flash-attention, knowledge-distillation, 3d-generation, stable-diffusion, text-to-video, text-to-image, diffusion-models, lora, training, fine-tuning, multimodal, video-generation
Creator
zjwfufu
Size
27 MB
Stars
205
Updated
CommitOct 14
Age
13d
Commits · 90 Days