Top AI Repos — open-source AI, indexed and scored

Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.

tdrussell/diffusion-pipe

A pipeline parallel training script for diffusion models.

54GPL-3.0Python

View on GitHub

Stars: 2,002
Forks: 279
Watchers: 18
Open issues: 260
Contributors: 18

Star history, last 90 days+1 over 3 snapshots

2026-07-31 · 2,0012026-08-06 · 2,002

Star history, last 90 days
Date	Stars
2026-07-31	2001
2026-08-04	2002
2026-08-06	2002

Today

— stars today

This week

— stars this week

This month

— stars this month

Momentum

0.0

growth rate 0.00%/day

From the README

# diffusion-pipe
A pipeline parallel training script for diffusion models.

Models supported: SDXL, Flux, LTX-Video, HunyuanVideo (t2v), Cosmos, Lumina Image 2.0, Wan2.1 (t2v and i2v), Chroma, HiDream, Stable Diffusion 3, Cosmos-Predict2, OmniGen2, Flux Kontext, Wan2.2, Qwen-Image, Qwen-Image-Edit, HunyuanImage-2.1, AuraFlow, Z-Image, HunyuanVideo-1.5, Flux 2 (Dev and Klein), Anima, Ernie-Image, LTX 2.3, Ideogram4, Krea 2.

## Features
- Pipeline parallelism, for training models larger than can fit on a single GPU
- Useful metrics logged to Tensorboard
- Compute metrics on a held-out eval set, for measuring generalization
- Training state checkpointing and resuming from checkpoint
- Efficient multi-process, multi-GPU pre-caching of latents and text embeddings
- Seemlessly supports both image and video models in a unified way
- Easily add new models by implementing a single subclass

## Recent changes
- 2026-06-24
- Support Krea 2.
- 2026-06-07
- Remove attention masking from Flux2. The model is supposed to take the full 512 length text embeds even with padding.
- This gives slightly lower starting loss, but seems to make almost no difference to the final trained result. Still, it should theoretically be better now.
- Delete cache folder or use `--regenerate_cache` or else you might get Tensor shape errors from the old cached files.
- For ComfyUI-based models (Z-Image and later), support training directly from Comfy quantized weights, e.g. fp8_scaled.
- Support Ideogram4.
- 2026-05-15
- Initial LTX 2.3 support. Only T2I and T2V training for now, and no audio.
- 2026-04-23
- Support Ernie-Image.
- 2026-02-04
- Support Anima.
- 2026-01-16
- Support Flux 2, both Dev and Klein.
- Updated DeepSpeed version. Should probably update all requirements: `pip install -r requirements.txt -U`.
- 2025-12-20
- Support HunyuanVideo-1.5. Currently only T2I and T2V training is supported.
- Add grad norm logging when using GenericOptim.
- 2025-11-29
- Change license to GPL-3 so I can use ComfyUI code. Going forward, model implementations will use ComfyUI backend code whenever possible.
- ComfyUI submodule has been added. Make sure to run ```git submodule update``` after pulling.
- Support Z-Image.

## Windows support
It will be difficult or impossible to make training work on native Windows. This is because Deepspeed only has [partial Windows support](https://github.com/microsoft/DeepSpeed/blob/master/blogs/windows/08-2024/README.md). Deepspeed is a hard requirement because the entire training script is built around Deepspeed pipeline parallelism. However, it will work on Windows Subsystem for Linux, specifically WSL 2. If you must use Windows I recommend trying WSL 2.

## Installing
Clone the repository:
```
git clone --recurse-submodules https://github.com/tdrussell/diffusion-pipe
```

If you alread cloned it and forgot to do --recurse-submodules:
```
git submodule init
git submodule update
```

Install Miniconda: https://docs.anaconda.com/miniconda/

Create the environment:
```
conda create -n diffusion-pipe python=3.12
conda activate diffusion-pipe
```

Install PyTorch first. It is not listed in the requirements file, because certain GPUs sometimes need different versions of PyTorch or CUDA, and you might have to find a combination that works for your hardware. As of this writing (October 26, 2025), PyTorch 2.9.0 with CUDA 12.8 works on my 4090, and is compatible with the current latest flash-attn 2.8.3:
```
pip install torch torchvision
```

Install nvcc: https://anaconda.org/nvidia/cuda-nvcc. Probably try to make it match the CUDA version of PyTorch.

Install the rest of the dependencies:
```
pip install -r requirements.txt
```

(Optional) Install Flash Attention. It's not in the requirements file. Some models need it.
```
pip install flash-attn
```

### Cosmos requirements
NVIDIA Cosmos (the original Cosmos video model, not Cosmos-Predict2) additionally requires TransformerEngine.

This dependency isn't

Excerpt of 10,626 characters

Read on GitHub