diffusion-pipe: pipeline parallel training for large image and video diffusion models
A pipeline parallel training script for diffusion models. Delete cache folder or use regenerate_cache or else you might get Tensor shape errors from the old cached files.
At a glance
- What is it?
- diffusion-pipe is a DeepSpeed-based training script that splits diffusion model layers across multiple GPUs, supporting a long list of architectures from SDXL to Flux 2. It trades native Windows support and a stable dependency set for flexibility and model coverage.
- Who is it for?
- Adopt diffusion-pipe if you train diffusion models on Linux or WSL 2, need pipeline parallelism to exceed single-GPU memory limits, and want to cover many models without writing new trainers. Do not use it if you need native Windows support, require pinned dependencies, or want to avoid GPL-3.0 obligations.
- Can I use it commercially?
- Yes, with conditions. GPL-3.0 is a copyleft licence: if you distribute software that includes it, you must release that software's source code under the same licence. Running it internally without distributing it does not trigger that obligation.
- Is it still maintained?
- Yes. The repository last received commits 1 day ago.
- What is it written in?
- Mainly Python, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 14, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
The problem it solves: training models that exceed one GPU
diffusion-pipe addresses a specific pain: training large diffusion models that do not fit into the memory of a single GPU. The README states its core feature is pipeline parallelism, which splits the model across multiple GPUs. This is different from data parallelism, where each GPU holds a full copy of the model. For a model like Flux 2 or HunyuanVideo, a single consumer GPU is not enough. The script is aimed at individual developers and small teams who want to fine-tune or train LoRAs on these models without renting a cluster of high-memory GPUs. It also unifies image and video training in one codebase, so you do not need separate tools for each modality.
How it works: DeepSpeed pipeline parallelism and a unified model interface
The entire training script is built around DeepSpeed pipeline parallelism, as the README states. DeepSpeed is a hard requirement, not an optional accelerator. The script organizes the model into pipeline stages, each assigned to a different GPU, and data flows through them sequentially. The README also mentions that new models can be added by implementing a single subclass, which suggests a common interface for all supported models. Latent and text embeddings are pre-cached in a multi-process, multi-GPU manner, which speeds up training by avoiding repeated forward passes through the text encoder. Metrics are logged to Tensorboard, and you can compute them on a held-out eval set. The training state can be checkpointed and resumed, which is essential for long training runs that may be interrupted.
Installation and configuration: what the README actually tells you
Installation starts with cloning the repository. The README gives this command: git clone --recurse-submodules https://github.com/tdrussell/diffusion-pipe. If you forgot the submodules, you run git submodule init and git submodule update. You need Miniconda, then create a Python 3.12 environment with conda create -n diffusion-pipe python=3.12. PyTorch is intentionally not in requirements.txt because different GPUs need different versions. The README notes that as of October 26, 2025, PyTorch 2.9.0 with CUDA 12.8 works on a 4090 with flash-attn 2.8.3. You install torch and torchvision, then nvcc, then pip install -r requirements.txt. Flash Attention is optional but required for some models. Training is launched with a deepspeed command, for example: NCCL_P2P_DISABLE="1" NCCL_IB_DISABLE="1" deepspeed --num_gpus=1 train.py --deepspeed --config examples/hunyuan_vi... The config files in the examples directory are meant to be read first, with the main example containing commented settings.
The cache regeneration trap: a concrete failure mode
A notable limitation is the cache invalidation problem. The README warns explicitly: delete the cache folder or use --regenerate_cache, or else you might get Tensor shape errors from the old cached files. This is not a hypothetical issue; it appears in the changelog for 2026-06-07 and again for 2026-08-06. When latent scaling logic changes, as it did for Z-Image, Flux2, and Ernie-Image, the cached files become stale. If you forget to regenerate, training crashes with shape mismatches. This is a real operational burden, especially for users who train multiple models or update frequently. The README does not provide an automated cache versioning mechanism, so the responsibility falls on the user. This is a design choice that prioritizes speed over safety, and it will catch you at least once.
Model coverage is broad but uneven: read the per-model docs
The README lists a huge number of supported models: SDXL, Flux, LTX-Video, HunyuanVideo, Cosmos, Lumina Image 2.0, Wan2.1, 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, Anima, Ernie-Image, LTX 2.3, Ideogram4, Krea 2, and MiniMax H3. But coverage is not uniform. The changelog notes that some models have partial support, for example LTX 2.3 supports only T2I and T2V, and no audio. HunyuanVideo-1.5 also only supports T2I and T2V. For ComfyUI-based models, you can train directly from quantized weights like fp8_scaled, which is useful for memory savings. The README directs you to docs/supported_models.md for per-model configuration, so you must check that file before assuming a model works fully.
Windows, dependencies, and the GPL-3.0 license: what to know before adopting
Native Windows training is effectively unsupported. DeepSpeed has only partial Windows support, and the README recommends WSL 2 instead. This is a hard constraint for anyone on Windows. Dependencies are mostly unpinned, which means updates can introduce breaking changes. The README says to run pip install -r requirements.txt -U to update, but without pins you have no reproducible environment. The license is GPL-3.0, which has implications for derivative works and distribution. This is not legal advice, but you should be aware that GPL-3.0 can be a concern for commercial or closed-source projects. The README also mentions that some models require additional dependencies, such as TransformerEngine for Cosmos, which is not in requirements.txt and requires special environment variables to install.
Alternatives: how diffusion-pipe compares to other training approaches
A common alternative is to use the native training scripts provided by each model's repository, such as the diffusers library for Stable Diffusion or the official Flux training code. The key difference is that those scripts typically use data parallelism or simple single-GPU training, not pipeline parallelism. They also do not unify image and video models. diffusers, for example, offers a Trainer class that handles mixed precision and distributed training via Accelerate, but it does not automatically split a model across GPUs in the pipeline sense. If your model fits on one GPU, diffusers is simpler and has a larger community. If you need to train a model larger than one GPU, diffusion-pipe is one of the few open-source options that explicitly implements pipeline parallelism for a wide range of models. Another alternative is to use DeepSpeed directly, writing your own pipeline stage partitioning, which gives you full control but requires significant engineering effort.
Maintenance and upgrade cost: what the changelog reveals
The project is actively maintained, with changes dated from 2025-12-20 through 2026-08-08. The changelog shows frequent model additions and bug fixes. Upgrading is straightforward: git pull and git submodule update. But the cost is that you must monitor the changelog for cache regeneration requirements and dependency updates. The README notes that DeepSpeed was updated in 2026-01-16 and recommends updating all requirements. Since most dependencies are unpinned, you should expect occasional breakage when you pull new commits. The lack of a release process, as indicated by no recent releases being retrieved, means you are always on the main branch, which can be a moving target. If you need stability, you may want to pin a specific commit, but the README does not document a stable release workflow.
Editorial conclusion
Adopt diffusion-pipe if you train diffusion models on Linux or WSL 2, need pipeline parallelism to exceed single-GPU memory limits, and want to cover many models without writing new trainers. Do not use it if you need native Windows support, require pinned dependencies, or want to avoid GPL-3.0 obligations. Before adopting, verify your GPU and CUDA combination works with the required PyTorch and flash-attn versions, and always run --regenerate_cache after updating the repository, as the README warns that stale cached files can cause Tensor shape errors.
Community notes