Skip to content
BranchesTags

Repository files navigation

ChituDiffusion

中文  ·  English

支持 混合并行、弹性序列并行 与 SOTA特征缓存策略的 DiT 推理运行时

Python CUDA License

ChituDiffusion 基于 Diffusers 生命周期提供模型生成、上下文并行、缓存加速和常驻服务。

项目处于开发者预览阶段。服务不保证任意 rank 故障后的恢复。

核心能力

Fast CP：高性能序列并行

提供 Fast AGKV 与 Fast Ulysses，并保留 NCCL fallback。面向具备 GPU P2P、 NVSHMEM 和目标架构扩展的兼容单机环境优化通信路径。

EPE：弹性并行服务

根据实测代价和端到端 SLO，在 pulse 边界调整 CP lane，并通过统一 executor 处理 队列、状态迁移和 worker 生命周期。

FlexCache：Cache 加速 API 和评测

提供 MagCache、MeanCache、TeaCache、TaylorSeer 与 PAB 的统一 API，可叠加单卡或 静态 CP；另提供单卡 FreeCache Preview，支持任意整数 Fresh 预算。 查看 速度与质量评测。

Diffusers Plug and Play

保留 tokenizer、文本编码器、scheduler、VAE 和输出处理流程，只适配模型相关的 DiT 执行与 tensor 布局。

支持模型

模型 生成能力 Hybrid Parallelism EPE Serve FlexCache
FLUX.1 文生图 ✅ CP（NCCL / Fast CP） ✅ ✅
FLUX.2-klein 文生图 ✅ CP（固定 full-world 静态并行） ❌ ❌
Hunyuan Image 3 固定尺寸文生图与图生图基础服务链路 ✅ TP × CFG × CP × EP、独立 VAEP
CP：NCCL / Fast AGKV / Fast Ulysses		 ✅ 仅 static_cp ❌
LLaDA-Image 文生图、VQ 条件生成、单图编辑 ✅ CFG × CP（AGKV / Ulysses） ✅ 文生图 ❌
MiniMax-H3 T2VA、first/last-frame FL2VA 基础服务链路 ✅ TP × CP、独立 VAEP
CP：NCCL / Fast AGKV / Fast Ulysses		 ✅ ❌
Qwen-Image 文生图 ✅ CFG × CP（NCCL / Fast CP） ✅ ✅
Wan 2.1 T2V 文生视频 ✅ TP × CFG × CP（NCCL / Fast CP） ✅ ✅
Z-Image 文生图 ✅ TP × CFG × CP（NCCL / Fast CP） ✅ ✅

安装

要求 Linux、Python 3.12 或 3.13、NVIDIA GPU 和可用的 CUDA 环境。按机器 CUDA 版本选择 pyproject.toml 中的 PyTorch index，然后运行：

uv sync --group dev
source .venv/bin/activate
chitu --help

pyproject.toml 是唯一构建入口。

生成

chitu generate \
  --model zimage \
  --model-path /path/to/Z-Image \
  --prompt "a red cube on a white table" \
  --output outputs/zimage.png

静态多卡 NCCL CP：

torchrun --standalone --nproc-per-node=4 -m chitu_diffusion.cli \
  generate \
  --model flux1 \
  --model-path /path/to/FLUX.1-dev \
  --output outputs/flux1.png

单机 TP×CFG×CP×EP MoE，例如 Hunyuan Image 3 的默认 8 卡拓扑：

torchrun --standalone --nproc-per-node=8 -m chitu_diffusion.cli \
  generate \
  --model hunyuan-image3 \
  --model-path /path/to/HunyuanImage-3 \
  --steps 50 \
  --output outputs/hunyuan_image3.png

Fast CP 需要兼容的单机 GPU、GPU P2P、NVSHMEM 和针对目标环境编译的扩展。参数和安装 步骤见 Fast CP README。

FlexCache

chitu generate \
  --model wan \
  --model-path /path/to/Wan2.1-T2V-1.3B \
  --steps 50 \
  --cache-strategy magcache \
  --output outputs/wan-magcache.mp4

缓存支持与模型、步数和 profile 绑定。不支持的组合会在运行前报错。原理、支持范围和 测试结果见 FlexCache 文档。

FreeCache Preview 支持 Z-Image、Qwen-Image 和 FLUX.1-dev 的单卡 50-step 生成， Fresh 预算可选 1–50，建议从 25 开始。见 使用说明。

八个已支持模型的图像/视频解码共用静态逐层 VAEP：卷积 halo、全局归一化、 AGKV，以及模型特定的位置编码和时序适配。见 VAEP 架构 与 验证范围。

EPE 服务

chitu serve --stage-config examples/stage-zimage.yaml

所有 rank 进入服务生命周期，leader 提供 HTTP，其他 rank 执行 worker loop。配置和 限制见运行说明与 EPE 文档。

文档与开发

python -m pytest -q
python -m ruff check chitu_diffusion tests examples
python -m build
mkdocs build --strict

许可证

MIT License

ChituDiffusion

中文  ·  English

DiT inference with hybrid parallelism, elastic sequence parallelism, and SOTA feature caching strategies

Python CUDA License

ChituDiffusion provides model generation, context parallelism, cache acceleration, and persistent serving while preserving the Diffusers pipeline lifecycle.

The project is a developer preview. The service does not recover from an arbitrary rank failure.

Core capabilities

Fast CP: High-performance sequence parallelism

Fast AGKV and Fast Ulysses with an NCCL fallback. The fast paths target compatible single-node systems with GPU P2P, NVSHMEM, and extensions built for the target architecture.

EPE: Elastic parallel serving

Adjusts CP lanes at pulse boundaries using measured costs and end-to-end SLOs, with one executor lifecycle for queues, state migration, and workers.

FlexCache: Cache APIs and evaluation

One API for MagCache, MeanCache, TeaCache, TaylorSeer, and PAB, composable with single-GPU or static CP generation. FreeCache Preview adds single-GPU generation with any integer Fresh budget. See the speed and quality evaluations.

Diffusers Plug and Play

Preserves tokenizers, text encoders, schedulers, VAEs, and output processing. Adapters only supply model-specific DiT execution and tensor layouts.

Supported models

Model Generation Hybrid Parallelism EPE Serve FlexCache
FLUX.1 Text to image ✅ CP (NCCL / Fast CP) ✅ ✅
FLUX.2-klein Text to image ✅ CP (fixed full-world static parallelism) ❌ ❌
Hunyuan Image 3 Foundational fixed-size text-to-image and image-to-image service path ✅ TP × CFG × CP × EP, independent VAEP
CP: NCCL / Fast AGKV / Fast Ulysses		 ✅ static_cp only ❌
LLaDA-Image Text to image, VQ-conditioned generation, and single-image editing ✅ CFG × CP (AGKV / Ulysses) ✅ Text to image ❌
MiniMax-H3 Foundational T2VA and first/last-frame FL2VA service path ✅ TP × CP, independent VAEP
CP: NCCL / Fast AGKV / Fast Ulysses		 ✅ ❌
Qwen-Image Text to image ✅ CFG × CP (NCCL / Fast CP) ✅ ✅
Wan 2.1 T2V Text to video ✅ TP × CFG × CP (NCCL / Fast CP) ✅ ✅
Z-Image Text to image ✅ TP × CFG × CP (NCCL / Fast CP) ✅ ✅

Installation

ChituDiffusion requires Linux, Python 3.12 or 3.13, an NVIDIA GPU, and a working CUDA environment. Select the matching PyTorch index in pyproject.toml, then run:

uv sync --group dev
source .venv/bin/activate
chitu --help

pyproject.toml is the only build entry point.

Generation

chitu generate \
  --model zimage \
  --model-path /path/to/Z-Image \
  --prompt "a red cube on a white table" \
  --output outputs/zimage.png

Static multi-GPU NCCL CP:

torchrun --standalone --nproc-per-node=4 -m chitu_diffusion.cli \
  generate \
  --model flux1 \
  --model-path /path/to/FLUX.1-dev \
  --output outputs/flux1.png

A single-node TP×CFG×CP×EP MoE stage, for example Hunyuan Image 3 on its default eight-GPU topology:

torchrun --standalone --nproc-per-node=8 -m chitu_diffusion.cli \
  generate \
  --model hunyuan-image3 \
  --model-path /path/to/HunyuanImage-3 \
  --steps 50 \
  --output outputs/hunyuan_image3.png

Fast CP requires compatible single-node GPUs, GPU P2P, NVSHMEM, and extensions built for the target environment. See the Fast CP README.

FlexCache

chitu generate \
  --model wan \
  --model-path /path/to/Wan2.1-T2V-1.3B \
  --steps 50 \
  --cache-strategy magcache \
  --output outputs/wan-magcache.mp4

Cache support is tied to the model, step count, and profile. Unsupported combinations fail before execution. See the FlexCache documentation for principles, support, and measured results.

FreeCache Preview supports single-GPU, 50-step generation for Z-Image, Qwen-Image, and FLUX.1-dev with Fresh budgets from 1 to 50. Start with 25; see the usage guide.

All eight supported models share static layer-wise image/video VAEP with convolution halos, global normalization, AGKV, and model-specific position/time adapters. See the architecture and validation coverage.

EPE serving

chitu serve --stage-config examples/stage-zimage.yaml

All ranks enter the service lifecycle. The leader hosts HTTP while the remaining ranks execute worker loops. See the running guide and EPE documentation.

Documentation and development

python -m pytest -q
python -m ruff check chitu_diffusion tests examples
python -m build
mkdocs build --strict

License

MIT License

About

thu-pacman.github.io/ChituDiffusion/

Resources

Readme
MIT license
Activity
Custom properties

Stars

27 stars

Watchers

0 watching

Forks

3 forks
Report repository

Releases

Packages

Contributors

Languages