《PYRAMIDAL FLOW MATCHING FOR EFFICIENT VIDEO GENERATIVE MODELING》
论文地址:2410.05954https://arxiv.org/pdf/2410.05954
项目地址:
jy0205/Pyramid-Flow: 用于高效视频生成建模的金字塔流匹配代码https://github.com/jy0205/Pyramid-Flow
论文提出了一种新的视频生成模型,通过金字塔流匹配算法(Pyramidal Flow Matching),有效降低了视频生成的计算复杂度。该方法通过在不同分辨率的金字塔阶段之间进行流匹配,实现了从噪声到数据的生成过程,并通过单一的Diffusion Transformer(DiT)进行端到端优化。
摘要详述
论文提出了一种高效的视频生成建模方法,称为金字塔流匹配,旨在通过降低计算复杂度来优化视频生成过程。该方法避免了直接在全分辨率下进行训练,而是将视频生成过程分解为多个在不同分辨率下运行的金字塔阶段,仅在最终阶段达到全分辨率。这种方法的主要优势包括:
-
连续性:不同金字塔阶段的生成轨迹相互链接,后续阶段继续从前一阶段生成,避免了每个阶段从纯噪声重新生成的需要。
-
统一模型:与为每个图像金字塔使用独立模型不同,金字塔流匹配算法将它们集成到一个统一的模型中,通过端到端优化实现更优雅的实现,并大幅加快训练速度。
在全分辨率下,在非常嘈杂的潜在值上花费大量计算。(b) 我们的方法利用流动匹配的 f 灵活性在不同分辨率的潜在变量之间进行插值。这允许以更好的计算效率同时生成和解压缩视觉内容。请注意,黑色箭头表示降噪轨迹,蓝色箭头表示时间条件。
方法详述
空间金字塔图示。(a) 金字塔流分为多个阶段,每个阶段都从像素化和嘈杂的起点到无像素化和更清晰的结果。(b) 在推理过程中,我们在跨阶段的跳跃点添加校正噪声,以确保概率的连续性
1. 金字塔流匹配 (Pyramidal Flow Matching)
论文提出了一个新颖的视频生成框架,称为金字塔流匹配,它通过将视频生成轨迹重新解释为不同尺度的压缩表示的金字塔阶段来解决视频生成中的高时空复杂性问题。具体来说,该方法只在最终阶段以全分辨率运行,而在早期阶段则在更低分辨率下运行,从而减少冗余计算。
-
流的构建:在金字塔流中,每个阶段都从带有噪声的像素化(压缩)潜在表示开始,到无像素化(解压缩)且更清晰的潜在表示结束。通过这种方式,只有最后一个阶段在全分辨率下执行,而大多数阶段在更低分辨率下执行,减少了计算量。
-
统一训练:为了统一不同阶段的建模,论文通过在不同噪声水平和分辨率之间进行插值来构建概率路径。这允许从低分辨率的噪声潜在表示生成更清晰、细节更丰富的高分辨率结果。
-
推理中的重噪声:在推理过程中,需要在不同分辨率的金字塔阶段之间的跳跃点仔细处理,以确保概率路径的连续性。为此,论文提出了一种添加校正高斯噪声的方法,以匹配不同阶段之间的分布。
2. 空间金字塔 (Spatial Pyramid)
-
流的分段:空间金字塔流被分为多个阶段,每个阶段从像素化且带噪声的起点到无像素化且更清晰的结果。每个阶段的流遵循类似的公式,插值在像素化(压缩)和更带噪声的潜在表示与无像素化(解压缩)且更清晰的潜在表示之间。
-
训练和推理:在训练阶段,通过插值不同分辨率的潜在表示来构建金字塔流。在推理阶段,每个阶段的输出通过添加校正高斯噪声重新噪声化,以维持连续性。
3. 时间金字塔 (Temporal Pyramid)
-
视频的时间复杂性:视频因其时间长度而呈现显著的挑战。现有的全序列扩散方法同时生成所有视频帧,限制了固定长度的生成。与之相对,自回归视频生成范式支持在推理期间灵活长度的生成。
-
压缩历史条件:考虑到全分辨率历史条件中的高冗余,论文提出使用压缩的、低分辨率的历史进行自回归视频生成。这显著减少了视频生成预训练的计算和内存开销。
图 3:时间金字塔图示。(a) 在每个金字塔阶段,生成都以压缩的、低分辨率的历史记录为条件,以提高自回归模型的训练效率,如行所示。(b) 设计了一种兼容的位置编码方案,该方案在空间金字塔中外推,但在时间金字塔中插值,以允许条件的空间对齐
实验
复现
复现了两种模式,一种为web ui 一种为推理,
1、下载及环境安装
bash
git clone https://github.com/jy0205/Pyramid-Flow
cd Pyramid-Flow
# create env using conda
conda create -n pyramid python==3.8.10
conda activate pyramid
pip install -r requirements.txt
其实环境不一定一模一样,我用的之前的环境,但是diffusion 和transformer最好和requirement一样,如果出现找不到pyramid模块之类的报错,检查版本。
2、下载权重
新建一个py文件
python
from huggingface_hub import snapshot_download
model_path = 'PATH' # The local directory to save downloaded checkpoint
snapshot_download("rain1011/pyramid-flow-miniflux", local_dir=model_path, local_dir_use_symlinks=False, repo_type='model')
修改model
3、web UI
调整app 内 model_path 变量,为上一步的model_path。注意目前使用的pyramid_flux而不是pyramid_mmdit
bash
python app.py
根据弹出的页面,简单移动,大家可以自行尝试调参。
4、本地推理
新建test.py ,复制下面代码,修改model_path,本地gpu 内存24g,如果内存小,跑再下面那个代码,卸载到cpu 的版本.两个都是384p版本,差不多本地要两分钟左右生成5s。A800需要50s左右。3090生成768p极慢,不推荐。
bash
import torch
from PIL import Image
from pyramid_dit import PyramidDiTForVideoGeneration
from diffusers.utils import load_image, export_to_video
torch.cuda.set_device(0)
model_dtype, torch_dtype = 'bf16', torch.bfloat16 # Use bf16 (not support fp16 yet)
model_path = ***
model = PyramidDiTForVideoGeneration(
model_path, # Pass the base model path
model_name="pyramid_flux" , # set to pyramid_flux or pyramid_mmdit
model_dtype=model_dtype, # Use bf16
model_variant='diffusion_transformer_384p', # Pass the variant directory name
cpu_offloading=True, # Pass the CPU offloading flag
)
model.vae.enable_tiling()
model.vae.to("cuda")
model.dit.to("cuda")
model.text_encoder.to("cuda")
from tqdm import tqdm
# if you're not using sequential offloading bellow uncomment the lines above ^
# model.enable_sequential_cpu_offload()
import json
prompts = []
with open(r"D:\T2V\KandinskyVideo-main\Qwen-Audio-main\prompts_dict_new.json", 'r', encoding='utf-8') as f:
datas = json.load(f)
for timestamp, data in datas.items():
prompts.append(data)
print(prompts)
for i, prompt in tqdm(enumerate(prompts)):
width = 640
height = 384
with torch.no_grad(), torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
frames = model.generate(
prompt=prompt,
num_inference_steps=[20, 20, 20],
video_num_inference_steps=[10, 10, 10],
height=height,
width=width,
temp=16, # temp=16: 5s, temp=31: 10s
guidance_scale=7.0, # The guidance for the first frame, set it to 7 for 384p variant
video_guidance_scale=5.0, # The guidance for the other video latent
output_type="pil",
save_memory=True, # If you have enough GPU memory, set it to `False` to improve vae decoding speed
)
export_to_video(frames, f"./demo/2/{i}.mp4", fps=24)
# prompt = "A wide shot of the sunflower field at sunset. The sky is now a deep orange and pink, with the sun setting behind the horizon. The sunflower petals are still swaying in the breeze, but the children have disappeared. A single butterfly lands on a sunflower, its wings shimmering in the warm light. The air is filled with the sound of crickets chirping."
# # used for 384p model variant
# width = 640
# height = 384
# # used for 768p model variant
# # width = 1280
# # height = 768
cpu版本
python
import torch
from PIL import Image
from pyramid_dit import PyramidDiTForVideoGeneration
from diffusers.utils import load_image, export_to_video
torch.cuda.set_device(0)
model_dtype, torch_dtype = 'bf16', torch.bfloat16 # Use bf16 (not support fp16 yet)
model = PyramidDiTForVideoGeneration(
'PATH', # The downloaded checkpoint dir
model_name="pyramid_flux",
model_dtype=model_dtype,
model_variant='diffusion_transformer_384p',
)
model.vae.enable_tiling()
# model.vae.to("cuda")
# model.dit.to("cuda")
# model.text_encoder.to("cuda")
# if you're not using sequential offloading bellow uncomment the lines above ^
model.enable_sequential_cpu_offload()
prompt = "A movie trailer featuring the adventures of the 30 year old space man wearing a red wool knitted motorcycle helmet, blue sky, salt desert, cinematic style, shot on 35mm film, vivid colors"
# used for 384p model variant
width = 640
height = 384
# used for 768p model variant
# width = 1280
# height = 768
with torch.no_grad(), torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
frames = model.generate(
prompt=prompt,
num_inference_steps=[20, 20, 20],
video_num_inference_steps=[10, 10, 10],
height=height,
width=width,
temp=16, # temp=16: 5s, temp=31: 10s
guidance_scale=7.0, # The guidance for the first frame, set it to 7 for 384p variant
video_guidance_scale=5.0, # The guidance for the other video latent
output_type="pil",
save_memory=True, # If you have enough GPU memory, set it to `False` to improve vae decoding speed
)
export_to_video(frames, "./text_to_video_sample.mp4", fps=24)