Ouzhang/T-Stitch

# 🖇 T-Stitch: Accelerating Sampling in Pre-trained Diffusion Models with Trajectory Stitching This is the official PyTorch implementation of T-Stitch: Accelerating Sampling in Pre-trained Diffusion Models with Trajectory Stitching > [Zizheng Pan](https://zizhengpan.github.io/)¹, [Bohan Zhuang](https://bohanzhuang.github.io/)¹, [De-An Huang](https://ai.stanford.edu/~dahuang/)², [Weili Nie](https://weilinie.github.io/)², [Zhiding Yu](https://chrisding.github.io/)², [Chaowei Xiao](https://xiaocw11.github.io/)^2,3, [Jianfei Cai](https://jianfei-cai.github.io/)¹, [Anima Anandkumar](http://tensorlab.cms.caltech.edu/users/anima/) ⁴ > > Monash University¹, NVIDIA², University of Wisconsin, Madison³, Caltech⁴ > > [[Paper](https://arxiv.org/abs/2402.14167)] [[Project Page](https://t-stitch.github.io/)] ## 📰 A Gentle Introduction We introduce sampling Trajectory Stitching (**T-Stitch**), a simple yet efficient technique to improve the generation efficiency with little or no loss in the generation quality. Instead of solely using a large DPM for the entire sampling trajectory, T-Stitch first leverages a smaller DPM in the initial steps as a cheap drop-in replacement of the larger DPM and switches to the larger DPM at a later stage, thus achieving flexible speed and quality trade-offs.  One example of stitching more DiT-S steps to achieve faster sampling for DiT-XL, where the time cost is measured by generating 8 images on one RTX 3090 in seconds (s).  By directly adopting a small SD in the model zoo, T-Stitch naturally interpolates the speed, style, and image contents with a large styled SD, which also potentially improves the prompt alignment, e.g., “New York City” and “tropical beach” in the above examples.  T-Stitch is completely complementary to previous techniques that focus on reducing the sampling steps, e.g., directly reduce the number of steps, advanced samplers, distillation.  ## 🛠 Setup For basic usage with diffusers, you can create an environment following our provided `requirements.txt`. Create a conda environment and activate it ```bash conda create -n tstitch python=3.9 -y conda activate tstitch pip install -r requirements.txt ``` ### Docker For a containerized workflow, the repo now provides two Dockerfiles: - `Dockerfile`: primary image for `dit` and `SDXL` workflows - `Dockerfile.ldm`: separate image for the older `ldm` training stack Build the primary image: ```bash docker build -t tstitch:latest -f Dockerfile . ``` Run the primary image with GPU access: ```bash docker run --gpus all --rm -it \ -v /home/featurize/T-Stitch:/workspace/T-Stitch \ -v /path/to/data:/data \ -v /path/to/hf-cache:/workspace/.cache/huggingface \ tstitch:latest ``` Build the LDM image: ```bash docker build -t tstitch-ldm:latest -f Dockerfile.ldm . ``` Run the LDM image with GPU access: ```bash docker run --gpus all --rm -it \ -v /home/featurize/T-Stitch:/workspace/T-Stitch \ -v /path/to/data:/data \ -v /path/to/hf-cache:/workspace/.cache/huggingface \ tstitch-ldm:latest ``` ## 🪄 Gradio Demo  ```bash python sd/gradio_demo.py ``` ## ⚙️ DiT Experiments Please refer to the folder [dit](./dit) for detailed usage. ## ⚙️ U-Net Experiments Please refer to the folder [ldm](./ldm) for detailed usage. ## ⚙️ Stable Diffusion Experiments Using T-Stitch for stable diffusion models is easy. At the root of this repo, do ```python import torch from sd.tstitch_sd_utils import get_tstitch_pipepline import os large_sd = "Envvi/Inkpunk-Diffusion" small_sd = "nota-ai/bk-sdm-tiny" pipe_sd = get_tstitch_pipepline(large_sd, small_sd) prompt = 'a squirrel in the park, nvinkpunk style' latent = torch.randn(1, 4, 64, 64, device="cuda", dtype=torch.float16) save_dir = f'figures/inkpunk' if not os.path.exists(save_dir): os.makedirs(save_dir) ratios = [round(item, 1) for item in torch.arange(0, 1.1, 0.1).tolist()] for ratio in ratios: image = pipe_sd(prompt, unet_s_ratio=ratio, latents=latent, height=512, width=512).images[0] image.save(f"{save_dir}/sample-ratio-{ratio}.png") ``` The above script will create images by gradually increasing the fraction of small sd at the early sampling steps. Please feel free to try other stylized SD and other prompts. Also note that both models are required to process latents of the same shape.   ### Accelerating SDXL T-Stitch provides a smooth speed and quality trade-off between a compressed SSD-1B and the original SDXL. Try the following command for this demo, ```bash python sd/sdxl_demo.py ```  ### Training SDXL T-Stitch The repo now includes `sd/train_sdxl_tstitch.py` for training a T-Stitch-specific SDXL small UNet while freezing the large UNet, text encoders, and VAE. Fixed ratio training: ```bash accelerate launch sd/train_sdxl_tstitch.py \ --train_data_dir /path/to/images \ --metadata_file /path/to/metadata.jsonl \ --pretrained_model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \ --small_model_name_or_path segmind/SSD-1B \ --ratio 0.3 \ --ratio_schedule fixed \ --train_batch_size 1 \ --max_train_steps 10000 \ --learning_rate 1e-5 ``` Curriculum ratio training: ```bash accelerate launch sd/train_sdxl_tstitch.py \ --train_data_dir /path/to/images \ --metadata_file /path/to/metadata.jsonl \ --pretrained_model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \ --small_model_name_or_path segmind/SSD-1B \ --ratio 0.9 \ --ratio_schedule curriculum \ --ratio_start 0.1 \ --train_batch_size 1 \ --max_train_steps 10000 \ --learning_rate 1e-5 ``` Enable epsilon distillation from the frozen large UNet: ```bash accelerate launch sd/train_sdxl_tstitch.py \ --train_data_dir /path/to/images \ --metadata_file /path/to/metadata.jsonl \ --pretrained_model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \ --small_model_name_or_path segmind/SSD-1B \ --ratio 0.3 \ --ratio_schedule fixed \ --enable_distill \ --distill_weight 100.0 \ --train_batch_size 1 \ --max_train_steps 10000 \ --learning_rate 1e-5 ``` Resume training: ```bash accelerate launch sd/train_sdxl_tstitch.py \ --train_data_dir /path/to/images \ --metadata_file /path/to/metadata.jsonl \ --pretrained_model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \ --small_model_name_or_path segmind/SSD-1B \ --ratio 0.3 \ --resume_from_checkpoint latest ``` TensorBoard logs are written under `--output_dir/logs` and checkpoints under `--output_dir/checkpoints`. ### Accelerating SDXL + ControlNet T-Stitch is compatible with Controlnet, for example, To use canny edges with SDXL, run `python sd/sdxl_canny.py`  To use depth images with SDXL, run `python sd/sdxl_depth.py`  To use poses with SDXL, run `python sd/sdxl_pose.py`  ### Accelerating SDXL + LCM T-Stitch is compatible with step-distilled models such as LCM-SDXL to achieve further speedup. For example, by adopting a small LCM distilled SSD-1B, T-Stitch still obtains impressive speed and quality trade-offs. We provide a script to demonstrate this compatibility. ```bash python sd/sdxl_lcm_lora.py ```  ## Reproducing Paper Experiments SD-related code lives in [`sd/`](./sd), DiT code lives in [`dit/`](./dit), and LDM code lives in [`ldm/`](./ldm). See [EXPERIMENTS.md](./EXPERIMENTS.md) for a command index, and [REPRODUCIBILITY.md](./REPRODUCIBILITY.md) for additional reproduction notes and external baseline requirements. ## Acknowledgments Thanks to the open source codebases such as [DiT](https://github.com/facebookresearch/DiT), [ADM](https://github.com/openai/guided-diffusion), [Diffusers](https://github.com/huggingface/diffusers) and [LDM](https://github.com/CompVis/latent-diffusion). Our codebase is built on them. ## License T-Stitch is licensed under CC-BY-NC. See [LICENSE.txt](./LICENSE.txt) for details. Portions of the project are available under separate license terms: [LDM](https://github.com/CompVis/latent-diffusion) is licensed under the [MIT License](https://github.com/CompVis/latent-diffusion/blob/main/LICENSE).

搜集汇总

数据集介绍

构建方式

T-Stitch（Trajectory Stitching）是一种旨在加速预训练扩散模型采样的创新技术。其核心构建思路是通过引入一条由小模型主导的早期采样轨迹，替代常规采样中始终使用大模型的做法。具体而言，在扩散过程的初始阶段，该方法将一个小型扩散模型作为大模型的廉价替代品，执行前若干步的降噪过程；之后，在某一预定的切换点或根据可调的比例参数，将采样任务交还给原大模型，完成剩余步骤。这种在单次采样轨迹中拼接不同规模模型的前后段落的构建方式，实现了对采样速度与生成质量之间权衡的灵活调控。

使用方法

T-Stitch的使用方法高度集成且极为简便。对于Stable Diffusion模型，用户可通过调用官方提供的`get_tstitch_pipepline`函数，直接传入一个大模型和一个已训练的紧凑型小模型路径，即可构建出支持T-Stitch的管道。使用时，仅需在管道的`__call__`方法中传入一个介于0到1之间的`unet_s_ratio`参数，即可精确指定小型UNet在整个采样过程中所占的比例。该框架同时支持对SDXL、ControlNet以及LCM等进阶模型的加速，并提供了包括固定比例训练和课程学习比例训练在内的微调脚本，便于用户根据特定场景定制最优的采样策略。

背景与挑战

背景概述

在生成式人工智能迅速发展的浪潮中，扩散概率模型凭借其卓越的生成质量，在图像合成领域占据核心地位。然而，这些模型通常需要数十乃至上百步的迭代采样过程，导致高昂的计算延迟，成为其在实际应用中落地的关键瓶颈。在此背景下，2024年由澳大利亚莫纳什大学、美国英伟达研究院、威斯康星大学麦迪逊分校及加州理工学院联合团队提出了T-Stitch方法。该工作旨在解决预训练大模型采样加速这一核心研究问题，通过引入轨迹拼接思想，在初始阶段使用轻量小模型作为大模型的高效替代，于后期切换回大模型，从而实现采样速度与生成质量间的灵活权衡。T-Stitch以其简洁有效的方案问世即受广泛关注，为后续加速扩散模型采样提供了全新视角。

当前挑战

T-Stitch所面临的挑战首先源自扩散模型领域固有的生成效率难题：大模型尽管生成质量优异，但参数规模庞大、推理步数众多，难以胜任实时或资源受限场景下的部署需求，而简单减小模型规模又会显著损伤生成保真度与多样性。其次，构建过程中存在的技术挑战包括如何在采样轨迹中自适应地确定小模型与大模型的切换时机，以避免生成内容的剧烈跳变与风格不连贯；如何确保不同尺寸模型在潜在空间中的特征对齐，使得小模型输出的中间状态能够无缝衔接大模型的后续生成；以及如何扩展该方法以兼容现有各类加速技术，比如步数蒸馏与高级求解器，从而实现综合性能的进一步优化。

常用场景

经典使用场景

T-Stitch作为一种创新的采样加速技术，在预训练扩散模型的推理过程中展现出卓越的效能。其核心思想在于，不再单一依赖一个大型扩散模型走完整个采样轨迹，而是在初始阶段巧妙地引入一个小型扩散模型作为低成本替代，待到采样后期再切换至大型模型。这种‘轨迹拼接’策略，使得在生成质量几乎无损的前提下，灵活调节速度与质量的平衡，为高效生成开辟了新路径。

解决学术问题

该研究直面扩散模型采样速度迟缓这一关键瓶颈。传统方法常需数百步迭代才能生成高质量样本，计算成本高昂。T-Stitch通过模型间智能协作，在不牺牲生成质量的同时大幅缩短采样时间，有效缓解了扩散模型在实时或资源受限场景下的部署难题。其意义在于，从算法层面提出了一个与步数压缩、高级采样器、知识蒸馏等现有技术完全互补的加速新范式，丰富了扩散模型高效推理的理论工具箱。

实际应用

在实际应用中，T-Stitch展现出极强的适配性与实用价值。它能直接应用于主流扩散模型架构，如DiT、U-Net及Stable Diffusion系列。例如，在风格化图像生成中，通过混合不同规模的模型，T-Stitch不仅能加速生成，还能自然插值图像的风格与内容，甚至提升对提示词的遵循度。此外，它无缝兼容ControlNet和LCM等条件控制与步数蒸馏模型，在边缘图、深度图、姿态图引导的生成以及超快速采样（如四步生成）场景中均能稳定发挥加速作用。

数据集最近研究

最新研究方向

T-Stitch技术代表了扩散模型采样加速领域的前沿突破，通过巧妙的轨迹拼接策略，在不显著牺牲生成质量的前提下实现高效推理加速。该研究直击生成式AI落地的核心瓶颈——计算效率问题，利用大小模型在采样轨迹初期的互补特性，达成了速度与质量间的可调权衡。这一创新面向图像生成、视频合成及多模态内容创作等热点方向，为扩散模型的实时应用开辟了新路径。其与步长缩减、高级采样器及知识蒸馏等现有技术的正交兼容性，使其成为加速大模型推理的重要基础性工具，对推动AIGC在工业界的大规模部署具有里程碑意义。