基础模型:

• THUDM/CogVideoX-2b 库名称: diffusers 许可证: apache-2.0 管道标签: 文本到视频 标签:
• lora

TDM: 通过轨迹分布匹配学习少步扩散模型

 

这是"通过轨迹分布匹配学习少步扩散模型"的官方仓库，作者为罗艺鸿、胡天阳、孙嘉诚、蔡宇君、唐静。

GitHub仓库: https://github.com/Luo-Yihong/TDM

用户研究时间！

您认为哪个更好？部分图像由Pixart-α（50 NFE）生成，部分图像由**TDM（4 NFE）**生成，后者以数据无关的方式从Pixart-α蒸馏而来，仅需500次训练迭代和2小时A800训练时间。

快速文本到视频生成

我们提出的TDM可以轻松扩展到文本到视频任务。

上方视频由CogVideoX-2B（100 NFE）生成。在相同时间内，TDM（4 NFE）可以生成25个视频，如下所示，实现了惊人的25倍加速且性能无损。（注：GIF中的噪点由压缩导致。）

使用方法

TDM-SD3-LoRA

import torch
from diffusers import StableDiffusion3Pipeline, AutoencoderTiny, DPMSolverMultistepScheduler
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from diffusers.utils import make_image_grid
pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16).to("cuda")
pipe.load_lora_weights('Luo-Yihong/TDM_sd3_lora', adapter_name = 'tdm') # 加载TDM-LoRA
pipe.set_adapters(["tdm"], [0.125])# 重要提示：请将LoRA缩放比例设为0.125
pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taesd3", torch_dtype=torch.float16) # 节省GPU内存
pipe.vae.config.shift_factor = 0.0
pipe = pipe.to("cuda")
pipe.scheduler = DPMSolverMultistepScheduler.from_pretrained("Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers", subfolder="scheduler")
pipe.scheduler.config['flow_shift'] = 6 # flow_shift可在1到6之间调整
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
generator = torch.manual_seed(8888)
image = pipe(
    prompt="一只可爱的熊猫举着写有'TDM SOTA!'的牌子",
    negative_prompt="",
    num_inference_steps=4,
    height=1024,
    width=1024,
    num_images_per_prompt = 1,
    guidance_scale=1.,
    generator = generator,
).images[0]

pipe.scheduler = DPMSolverMultistepScheduler.from_pretrained("Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers", subfolder="scheduler")
pipe.set_adapters(["tdm"], [0.]) # 卸载lora
generator = torch.manual_seed(8888)
teacher_image = pipe(
    prompt="一只可爱的熊猫举着写有'TDM SOTA!'的牌子",
    negative_prompt="",
    num_inference_steps=28,
    height=1024,
    width=1024,
    num_images_per_prompt = 1,
    guidance_scale=7.,
    generator = generator,
).images[0]
make_image_grid([image,teacher_image],1,2)

右侧为SD3使用56 NFE生成的样本，左侧为TDM使用4 NFE生成的样本。您觉得哪个效果更好？

TDM-Dreamshaper-v7-LoRA

import torch
from diffusers import DiffusionPipeline, UNet2DConditionModel, DPMSolverMultistepScheduler
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
repo_name = "Luo-Yihong/TDM_dreamshaper_LoRA"
ckpt_name = "tdm_dreamshaper.pt"
pipe = DiffusionPipeline.from_pretrained('lykon/dreamshaper-7', torch_dtype=torch.float16).to("cuda")
pipe.load_lora_weights(hf_hub_download(repo_name, ckpt_name))
pipe.scheduler = DPMSolverMultistepScheduler.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="scheduler")
generator = torch.manual_seed(317)
image = pipe(
    prompt="戴太阳镜的亚洲女性特写照片",
    negative_prompt="",
    num_inference_steps=4,
    num_images_per_prompt = 1,
    generator = generator,
    guidance_scale=1.,
).images[0]
image

TDM-CogVideoX-2B-LoRA

import torch
from diffusers import CogVideoXPipeline
from diffusers.utils import export_to_video
pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-2b", torch_dtype=torch.float16)
pipe.vae.enable_slicing() # 节省内存
pipe.vae.enable_tiling() # 节省内存
pipe.load_lora_weights("Luo-Yihong/TDM_CogVideoX-2B_LoRA")
pipe.to("cuda")
prompt = (
    "一只熊猫穿着红色小夹克，戴着迷你帽子，坐在竹林中的木凳上。"
    "熊猫用毛茸茸的爪子拨弄着一把微型原声吉他，奏出柔和的旋律。"
    "附近有几只其他熊猫好奇地围观，有的还跟着节奏拍掌。"
    "阳光透过高高的竹子洒落，为场景镀上温柔的光晕。"
    "熊猫表情生动，弹奏时展现出专注与喜悦。"
    "背景中有一条潺潺的小溪和翠绿的植被，为这场独特音乐表演增添了宁静而神奇的氛围"
)
# 我们在时间步[999, 856, 665, 399]上训练生成器
# CogVideo-X官方调度器使用均匀间隔，可能导致次优结果
# 但TDM-LoRA在4 NFE下仍表现良好
# 我们将很快更新TDM-CogVideoX-LoRA以获得更好性能！
generator = torch.manual_seed(8888)
frames = pipe(prompt, guidance_scale=1, 
              num_inference_steps=4, 
              num_frames=49,
              generator = generator).frames[0]
export_to_video(frames, "output-TDM.mp4", fps=8)

🔥 预训练模型

我们发布了一系列TDM-LoRA模型，欢迎使用！

• TDM-SD3-LoRA
• TDM-CogVideoX-2B-LoRA
• TDM-Dreamshaper-LoRA

🔥 待办事项

• 发布训练代码

联系方式

如有任何问题，请联系罗艺鸿 (yluocg@connect.ust.hk)。

引用

@misc{luo2025tdm,
      title={通过轨迹分布匹配学习少步扩散模型}, 
      author={罗艺鸿 and 胡天阳 and 孙嘉诚 and 蔡宇君 and 唐静},
      year={2025},
      eprint={2503.06674},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2503.06674}, 
}