chenfengxu714/StreamDiffusionV2

StreamDiffusionV2 – 基于扩散模型的实时视频生成

是什么 – StreamDiffusionV2 是一个开源库，可让你将基于扩散的视频生成作为 流式 过程运行。无需等待整个视频生成完成，系统会实时处理短视频片段，并通过滚动 KV 缓存，使每个新片段都能复用已有的工作。支持单 GPU、多 GPU 流水线并行，以及用于交互式使用的轻量级 Web UI。

为何重要 – 传统的文本到视频扩散模型是批处理导向且缓慢的（每秒视频需数分钟）。StreamDiffusionV2 通过以下方式实现高帧率（实时或近实时）：

• 将视频分割为重叠的片段，并通过环形缓冲 KV 缓存复用潜在状态。
• 允许每个片段灵活的去噪步数（速度与质量的权衡）。
• 提供可选的加速路径：快速的 TAEHV-VAE 解码器、Flash-Attention 和 TensorRT 支持。
• 从单 GPU 扩展到多 GPU 流水线并行，支持“schedule-block”选项以在大型设备上获得更高吞吐量。

核心组件

组件	作用
StreamDiffusionV2Pipeline	高级 Python API，负责分块、编码、去噪和解码。
run_v2v.sh	统一的 CLI 入口点，用于离线推理（单 GPU、单 GPU 无流批处理、流水线）。
demo/	用于交互式生成的最小 FastAPI/Gradio Web UI。
streamv2v/	展示如何以编程方式驱动管道的示例脚本。
KV 缓存实现	环形缓冲区，用于存储因果 DiT 主干的注意力键值，支持滑动窗口注意力。

安装 – 该包已发布在 PyPI，可安装并选择性启用 Flash-Attention 支持：

conda create -n streamdiffusionv2 python=3.10 -y
conda activate streamdiffusionv2
pip install streamdiffusionv2            # 核心
pip install "streamdiffusionv2[flash-attn]"  # 更快的注意力（需 CUDA）
# 适用于最新的 NVIDIA Blackwell GPU
pip install torch==2.11.0 torchvision==0.26.0

若偏好本地检出，pip install . 也适用。

获取模型 – 该库仅提供推理代码，需自行下载检查点：

# 1.3B 模型（默认用于演示）
huggingface-cli download --resume-download Wan-AI/Wan2.1-T2V-1.3B --local-dir wan_models/Wan2.1-T2V-1.3B
huggingface-cli download --resume-download jerryfeng/StreamDiffusionV2 --local-dir ./ckpts --include "wan_causal_dmd_v2v/*"

# 14B 模型（更高画质）
huggingface-cli download --resume-download Wan-AI/Wan2.1-T2V-14B --local-dir wan_models/Wan2.1-T2V-14B
huggingface-cli download --resume-download jerryfeng/StreamDiffusionV2 --local-dir ./ckpts --include "wan_causal_dmd_v2v_14b/*"

可选地拉取轻量级 TAEHV-VAE 解码器：

curl -L https://github.com/madebyollin/taehv/raw/main/taew2_1.pth -o ckpts/taew2_1.pth

快速 Python 示例（单 GPU）

import numpy as np
from streamdiffusionv2 import (
    StreamDiffusionV2Pipeline,
    load_video,
    export_video,
)

pipe = StreamDiffusionV2Pipeline(
    checkpoint_folder="ckpts/wan_causal_dmd_v2v",
    mode="single",
)
pipe.prepare("A dog walks on the grass, realistic")

video = load_video("examples/original.mp4", height=480, width=832)
chunks = []
noise = pipe.noise_scale

for chunk in pipe.chunk_video(video):
    enc = pipe.encode_chunk(video, chunk, previous_noise_scale=noise,
                           initial_noise_scale=pipe.noise_scale)
    noise = enc.noise_scale
    den = pipe.denoise_chunk(enc)
    if den is None:
        continue
    chunks.append(pipe.decode_chunk(den))

out = np.concatenate(chunks, axis=0)
export_video(out, "outputs/python_single.mp4", fps=16)

相同管道可在 pipeline 模式下用于多 GPU 推理，仅需一次调用：

from streamdiffusionv2 import run_video_to_video
run_video_to_video(
    mode="pipe",
    checkpoint_folder="ckpts/wan_causal_dmd_v2v",
    video_path="examples/original.mp4",
    prompt="A dog walks on the grass, realistic",
    output_path="outputs/python_pipe.mp4",
    gpu_ids=[0,1],
    num_gpus=2,
)

CLI 使用 – run_v2v.sh 脚本抽象了相同步骤：

# 单 GPU 流式
./run_v2v.sh single
# 单 GPU 无流批处理（内存模式略有不同）
./run_v2v.sh single-wo
# 多 GPU 流水线
./run_v2v.sh pipe
# 添加性能分析
./run_v2v.sh pipe --profile
# 启用快速 TAEHV 解码器
./run_v2v.sh single --use_taehv

常见标志（--config_path、--checkpoint_folder、--video_path、--prompt_file_path、--height、--width、--fps、--step、--use_taehv）也可通过环境变量提供。

Web 演示 – 安装 Node.js 18 后，使用以下命令启动 UI：

cd demo
USE_TAEHV=1 python -m streamdiffusionv2.web_demo   # 或使用提供的启动脚本

然后在浏览器中打开 http://localhost:7860。

训练 – 仓库包含一个独立的 user-study 分支，基于 CausVid 代码库构建。提供用于训练推理时使用的因果 DiT 主干的脚本。

路线图 – 作者列出了未来工作，如动态调度器、FP8 量化和更深入的 TensorRT 集成。

引用 – 若在研究中使用该库，请引用 MLSys 2026 论文（README 中提供 BibTeX）。

---

StreamDiffusionV2 是任何希望实验或部署实时扩散视频生成的用户的生产就绪工具包，提供简单 Python API 和可扩展的多 GPU 后端。