Jano: Adaptive Diffusion Generation with Early-stage Convergence Awareness

Jano is an inference acceleration framework for diffusion-based video/image generation models. It profiles the spatio-temporal dynamics of the latent space during a short warm-up phase and then selectively skips computations for low-dynamic regions, achieving significant speedup with minimal quality degradation.

Showcase / 效果展示

The following examples are from ./assets, with runtime (generate_e2e) and quality metrics from the corresponding JSON files.

以下展示来自 ./assets，时间（generate_e2e）与质量指标均读取自对应 JSON 文件。

FLUX.1-dev (Text-to-Image)

Prompt: A photorealistic cute cat, wearing a simple blue shirt, standing against a clear sky background.

ORI	Jano	PAB	TeaCache	ToCA
36.14 s 1.00× PSNR ----- / SSIM ----- / LPIPS -----	19.33 s 1.87× PSNR 28.19 / SSIM 0.940 / LPIPS 0.088	23.81 s 1.52× PSNR 23.87 / SSIM 0.898 / LPIPS 0.137	23.36 s 1.55× PSNR 23.60 / SSIM 0.908 / LPIPS 0.107	19.38 s 1.86× PSNR 14.79 / SSIM 0.761 / LPIPS 0.379

Wan2.1-1.3B (Text-to-Video)

Prompt: Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.

ORI ori_two_anthropomorphic_cats_comfy_t2v-1.3B.mp4 165.60 s | 1.00× | N/A	Jano jano_offload0_two_anthropomorphic_cats_comfy_t2v-1.3B.mp4 83.94 s | 1.97× | PSNR 17.96 / SSIM 0.782 / LPIPS 0.184
PAB s5c8_i4_two_anthropomorphic_cats_comfy_t2v-1.3B.mp4 139.94 s | 1.18× | PSNR 16.08 / SSIM 0.734 / LPIPS 0.243	TeaCache thresh0.2_two_anthropomorphic_cats_comfy_t2v-1.3B.mp4 106.10 s | 1.56× | PSNR 21.28 / SSIM 0.854 / LPIPS 0.100
ToCA toca_two_anthropomorphic_cats_comfy_t2v-1.3B.mp4 71.39 s | 2.32× | PSNR 13.97 / SSIM 0.670 / LPIPS 0.347

Jano supports the following workloads:

Model	Task
Wan2.1-1.3B-T2V	Text-to-Video
Wan2.1-14B-T2V	Text-to-Video
FLUX.1-dev	Text-to-Image

Baseline implementations included in this repo:

• TeaCache
• PAB
• TokenCache (ToCA)

---

Jano 是一个面向扩散模型视频/图像生成推理加速的框架。它在短暂的预热阶段对潜空间的时空动态性进行分析，然后对低动态区域选择性地跳过计算，在几乎不损失质量的情况下显著提升推理速度。

Project Structure / 项目结构

Jano/
├── jano/                   # Core Jano library
│   ├── __init__.py         # init_jano() entry point
│   ├── block_manager.py    # Latent-space block partitioning
│   ├── dynamic_analyzer.py # Spatio-temporal dynamics analysis
│   ├── stuff.py            # Shared utilities (timestep tracking, etc.)
│   ├── mask_manager/       # Per-model cache-mask managers
│   ├── modules/            # Modified model forward passes (Wan, Flux, SD3, CogVideoX)
│   └── dist/               # Distributed (CFG-parallel) utilities
├── wan/                    # Wan2.1 model code + baseline implementations
├── flux/                   # FLUX.1 model code + baseline implementations
├── utils/                  # Shared utilities (timer, logger, quality metrics, envs)
├── run_wan/                # Inference scripts for Wan2.1
├── run_flux/               # Inference scripts for FLUX.1
├── run_cvx/                # Inference scripts for CogVideoX
├── ras_exp/                # Experimental RAS baseline
├── requirements.txt
└── LICENSE

Installation / 安装

1. Set up the environment / 配置环境

conda create -n jano python=3.10
conda activate jano
pip install -r requirements.txt

2. Download models / 下载模型

pip install "huggingface_hub[cli]"

# (Optional) Use a mirror in China / 中国用户可使用镜像
export HF_ENDPOINT=https://hf-mirror.com

# Wan2.1 (choose one or both)
huggingface-cli download Wan-AI/Wan2.1-T2V-1.3B --local-dir ./Wan2.1-T2V-1.3B
huggingface-cli download Wan-AI/Wan2.1-T2V-14B  --local-dir ./Wan2.1-T2V-14B

# FLUX.1-dev
huggingface-cli download black-forest-labs/FLUX.1-dev --local-dir ./Flux-1

Usage / 运行

Please use the unified launcher launch.py. First set MODEL_PATH, then run with --model and --method:

# Wan2.1
export MODEL_PATH=<your model path>

python launch.py --model <wan/flux>  --method <ori/jano/teacache/pab/toca> --gpus-per-node <1/2> --partition <if specific>

参数修改说明 / Parameter configuration: Please edit generation parameters directly in the corresponding *_generate.py files under run_wan/, run_flux/, and run_cvx/.

请在 run_wan/、run_flux/、run_cvx/ 下对应的 *_generate.py 文件中修改具体参数。

Note: If you get ModuleNotFoundError, add the project root to your Python path first:

export PYTHONPATH=$PYTHONPATH:$(pwd)

Memory Optimization for Wan-14B / Wan-14B 内存优化

Method	Technique
Jano	Set KV_OFFLOAD=1 + 2-GPU parallel
PAB	Set LAYER_INTERVAL=2 (memory ÷ n) + 2-GPU parallel

License / 许可证

This project is licensed under the Apache License 2.0.

The Wan2.1 model weights are subject to their own license.
The FLUX.1-dev model weights are subject to the FLUX.1-dev Non-Commercial License.

Citation / 引用

If you use this code, please cite our paper:

@misc{chen2026janoadaptivediffusiongeneration,
      title={Jano: Adaptive Diffusion Generation with Early-stage Convergence Awareness}, 
      author={Yuyang Chen and Linqian Zeng and Yijin ZHou and Hengjie Li and Jidong Zhai},
      year={2026},
      eprint={2603.00519},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2603.00519}, 
}