使用 Diffusers 中的 LattePipeline 进行文本生成视频从 torch.compile 加速到 8-bit 量化【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusersLatteLatent Diffusion Transformer for Video Generation论文见 Latte: Latent Diffusion Transformer for Video Generation 所引用的 2401.03048是一种以 Transformer 为骨干的潜空间扩散模型本仓库 Diffusers在 LattePipeline 中提供了面向文本生成视频T2V任务的官方实现。本文以 docs/source/en/api/pipelines/latte.md 为主线结合仓库源码与测试完整讲解如何加载LattePipeline、使用torch.compile降低推理延迟、通过 bitsandbytes 量化压缩显存并深入剖析其底层LatteTransformer3DModel的空间/时间双重注意力架构与全部可调参数帮助你快速上手并深入定制 Latte 视频生成。Latte 是什么为视频生成设计的潜空间扩散 TransformerLatte 由 Monash University、Shanghai AI Lab、Nanjing University 与 Nanyang Technological University 联合提出。其核心思路是先从输入视频中提取时空 token再利用一系列 Transformer 块在潜空间中建模视频分布。论文摘要指出为处理视频中数量庞大的 tokenLatte 从分解输入视频的空间与时间维度这一视角引入了四种高效变体同时通过严格的实验分析确定了最佳实践包括视频片段 patch embedding、模型变体、timestep-类别信息注入、时间位置编码与学习策略等。Latte 在四个标准视频生成基准FaceForensics、SkyTimelapse、UCF101、Taichi-HD上取得了当时的 SOTA 表现并将其能力扩展到了文本生成视频T2V任务。本仓库中的 LattePipeline 即为此 T2V 方向贡献的实现由 maxin-cn 贡献原始权重位于maxin-cn/Latte-1。在 Diffusers 的组件结构中LattePipeline的组成见 pipeline_latte.py 的__init__与文档字符串如下组件类型职责transformerLatteTransformer3DModel文本条件化的 3D Transformer对编码后的视频 latents 去噪vaeAutoencoderKL在视频与潜表示之间编码/解码text_encoderT5EncoderModel冻结的文本编码器使用 T5具体为t5-v1_1-xxl变体tokenizerT5TokenizerT5 分词器schedulerKarrasDiffusionSchedulers与transformer配合完成去噪其中tokenizer与text_encoder被声明为可选组件源码第 170 行的_optional_components [tokenizer, text_encoder]这意味着你可以在加载 pipeline 后移除这两个组件改为直接传入预先算好的prompt_embeds与negative_prompt_embeds来省去文本编码器的显存占用这一点在 test_latte.py 的test_save_load_optional_components中有完整验证。同时源码第 171 行定义了model_cpu_offload_seq text_encoder-transformer-vae因此调用enable_model_cpu_offload()时会按此顺序将组件逐个卸载到 CPU。快速上手加载 pipeline 并生成视频基础推理流程使用LattePipeline.from_pretrained加载官方权重并生成视频的最简代码如下来自文档示例与 EXAMPLE_DOC_STRINGimport torch from diffusers import LattePipeline from diffusers.utils import export_to_gif pipe LattePipeline.from_pretrained(maxin-cn/Latte-1, torch_dtypetorch.float16) # 启用模型 CPU 卸载以节省显存 pipe.enable_model_cpu_offload() prompt A small cactus with a happy face in the Sahara desert. videos pipe(prompt).frames[0] export_to_gif(videos, latte.gif)要点说明权重与精度maxin-cn/Latte-1是官方权重仓库加载时建议使用torch.float16以显著降低显存与带宽开销。输出结构pipe(prompt)返回LattePipelineOutputreturn_dictTrue时其唯一字段frames是形状为(batch, channels, frames, height, width)的张量列表.frames[0]取出第一个视频可直接交给export_to_gif导出为 GIF。若return_dictFalse则返回只含该张量的tuple。文本编码细节encode_prompt将 prompt 通过 T5 分词器编码为固定 120 token长度源码第 257 行max_length 120超长部分截断并给出警告随后通过mask_text_embeddings依据 attention mask 掩掉 padding 区域mask_featureTrue时这也是 T5 编码器输出由 120 个 token 被压缩到实际有效 token 数的原因。Latte 的 negative prompt 应使用空字符串详见negative_prompt_embeds的文档说明。使用 torch.compile 降低推理延迟LattePipeline的推理瓶颈在 Transformer 与 VAE 解码。文档推荐使用torch.compile减少推理延迟分三步进行第一步加载 pipeline支持cuda、mps、xpu、cpu等设备import torch from diffusers import LattePipeline pipeline LattePipeline.from_pretrained( maxin-cn/Latte-1, dtypetorch.float16 ).to(cuda) # or mps, xpu, cpu第二步将transformer与vae的内存布局切换为torch.channels_lastpipeline.transformer.to(memory_formattorch.channels_last) pipeline.vae.to(memory_formattorch.channels_last)第三步编译组件并运行推理pipeline.transformer torch.compile(pipeline.transformer) pipeline.vae.decode torch.compile(pipeline.vae.decode) video pipeline(promptA dog wearing sunglasses floating in space, surreal, nebulae in background).frames[0][!TIP]torch.compile首次调用会触发编译预热真正收益体现在连续多次推理上。另外pipeline 的decode_latents实现pipeline_latte.py 中decode_latents方法内部通过is_compiled_module(self.vae)判断 VAE 是否已被编译并据此调用self.vae._orig_mod.forward因此编译后的 VAE 也能正确拿到num_frames参数——编译与分块解码可以安全共存。文档给出了在 80GB A100 机器上的 benchmark 结果原始基准链接见文档正文Without torch.compile(): Average inference time: 16.246 seconds. With torch.compile(): Average inference time: 14.573 seconds.即在此硬件上torch.compile带来了约 10% 的推理加速。需要说明的是该数字取决于具体硬件、驱动与 PyTorch 版本实际收益请以本机实测为准。量化推理用 bitsandbytes 把 Latte 装进更小的显存量化通过以更低精度的数据类型存储模型权重来降低超大模型的内存需求但对不同视频模型量化对视频质量的影响不一。文档示例演示了如何使用 bitsandbytes 加载量化版LattePipeline分别对text_encoderT5用 transformers 的BitsAndBytesConfig与transformer用 diffusers 的BitsAndBytesConfig做 8-bit 量化。import torch from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig, LatteTransformer3DModel, LattePipeline from diffusers.utils import export_to_gif from transformers import BitsAndBytesConfig, T5EncoderModel quant_config BitsAndBytesConfig(load_in_8bitTrue) text_encoder_8bit T5EncoderModel.from_pretrained( maxin-cn/Latte-1, subfoldertext_encoder, quantization_configquant_config, dtypetorch.float16, ) quant_config DiffusersBitsAndBytesConfig(load_in_8bitTrue) transformer_8bit LatteTransformer3DModel.from_pretrained( maxin-cn/Latte-1, subfoldertransformer, quantization_configquant_config, dtypetorch.float16, ) pipeline LattePipeline.from_pretrained( maxin-cn/Latte-1, text_encodertext_encoder_8bit, transformertransformer_8bit, dtypetorch.float16, device_mapbalanced, ) prompt A small cactus with a happy face in the Sahara desert. video pipeline(prompt).frames[0] export_to_gif(video, latte.gif)关键点两个量化配置类不要混淆T5 文本编码器来自transformers使用transformers.BitsAndBytesConfigTransformer 是 diffusers 模型使用diffusers.BitsAndBytesConfig。按 subfolder 加载子组件maxin-cn/Latte-1仓库中模型按text_encoder、transformer等子目录组织量化加载时同样需要指定subfolder。device_mapbalanced在 pipeline 层面按显存/内存均衡地把各组件分配到可用设备。更全面的量化后端介绍与后端选择方法请参阅 量化总览。LattePipeline 完整参数说明LattePipeline.__call__的完整签名与默认值如下源码见 pipeline_latte.py 第 615 行起__call__( promptNone, # str 或 list[str]引导视频生成的提示词 negative_prompt, # 负向提示词Latte 中应为 不使用 guidance 时忽略 num_inference_steps50, # 去噪步数越多质量越高、速度越慢 timestepsNone, # 自定义 timesteps需降序覆盖调度器的等间隔策略 guidance_scale7.5, # CFG 引导强度1 时启用无分类器引导 num_images_per_prompt1, # 每个 prompt 生成的视频数量 video_length16, # 生成的视频帧数16 帧约等于 8fps 下的 2 秒 height512, width512, # 视频高宽必须能被 8 整除 eta0.0, # DDIM 论文中的 η仅对 DDIMScheduler 生效 generatorNone, # torch.Generator用于可复现生成 latentsNone, # 预生成噪声 latents可跨 prompt 复用 prompt_embedsNone, # 预计算的正向文本嵌入可做 prompt weighting negative_prompt_embedsNone, # 预计算的负向文本嵌入 output_typepil, # 输出格式PIL 或 np.array 等 return_dictTrue, # 返回 LattePipelineOutput 还是 tuple callback_on_step_endNone, # 每步去噪结束时的回调 callback_on_step_end_tensor_inputs[latents], clean_captionTrue, # 编码前是否清洗 caption需安装 beautifulsoup4 与 ftfy mask_featureTrue, # 是否掩掉文本嵌入的 padding 部分 enable_temporal_attentionsTrue,# 是否启用时间注意力 decode_chunk_size14, # VAE 每次解码的帧数越小越省显存 )主要参数的行为与限制guidance_scale与 CFGdo_classifier_free_guidance guidance_scale 1.0源码属性do_classifier_free_guidance与__call__第 750 行。启用时encode_prompt会额外编码空字符串得到负向嵌入去噪循环中将条件/无条件预测拼接为一批做两次前向再按noise_pred_uncond guidance_scale * (noise_pred_text - noise_pred_uncond)合成第 832–834 行。guidance_scale越高视频与文本越贴合但通常以降低视频质量为代价。height/width必须能被 8 整除check_inputs中会显式校验第 383–384 行否则抛出ValueError。未显式传入时默认取transformer.config.sample_size * vae_scale_factor第 722–723 行。video_length与临时位置编码帧数决定临时位置编码的长度。LatteTransformer3DModel在初始化时按video_length生成 1D sincos 位置编码temp_pos_embed并注册为 bufferlatte_transformer_3d.py 第 159–162 行仅在num_frame 1时加入第一层时间块之前。decode_chunk_sizeVAE 解码按此大小分块进行以避免 OOMdecode_latents方法第 887–914 行分块更大则时间一致性更好但更耗显存默认 14 是质量与内存的折中。timesteps可传入自定义降序 timesteps 覆盖调度器的等间隔策略若同时传入sigmas会报错。该逻辑由从StableDiffusionPipeline复制的retrieve_timesteps统一处理第 81–137 行需要调度器set_timesteps支持对应参数。clean_caption为True时通过ftfy与BeautifulSoup清洗 caption去 URL、HTML、CJK 字符、连续标点等逻辑见_clean_caption方法未安装依赖时会回退为原始 prompt 并打印警告。latents/prompt_embeds两者均可用于复现或精细控制。prepare_latents会按video_length、height//vae_scale_factor、width//vae_scale_factor构造(batch, latent_channels, frames, h, w)形状的噪声并用scheduler.init_noise_sigma缩放第 565–588 行。深入底层LatteTransformer3DModel 的空间/时间分解架构文档对应的底层去噪网络是LatteTransformer3DModel。它把视频视为形状(batch, channels, num_frames, height, width)的 5D 数据核心设计是将空间与时间注意力分解为两组 Transformer 块输入层PatchEmbed将每一帧 patch 化并加入空间位置编码支持interpolation_scale处理非 64 的sample_size。transformer_blocks空间块num_layers个BasicTransformerBlock带cross_attention_dim的交叉注意力接收经PixArtAlphaTextProjection投影后的 T5 文本嵌入。temporal_transformer_blocks时间块num_layers个BasicTransformerBlockcross_attention_dimNone纯自注意力负责建模帧间时序关系。条件注入AdaLayerNormSingle(use_additional_conditionsFalse)将 timestep 编码为调制参数配合scale_shift_table与norm_out完成 FiLM 式调制caption_projection把 T5 的 caption 通道映射到隐藏维度。输出层proj_out将每个 patch 投影回patch_size * patch_size * out_channels后 unpatchify 还原为视频张量。forward的数据流第 205–329 行会反复在空间布局与时间布局间重塑张量先将(B, C, F, H, W)展开为(B*F, C, H, W)逐帧过空间块enable_temporal_attentionsTrue时再重塑为(B*tokens, F, hidden)过时间块第一层时间块前叠加temp_pos_embed时间位置编码。整个流程中timestep 与文本嵌入通过repeat_interleave按帧或按 patch 广播保证空间与时间块都能拿到条件信号。这也是论文从空间与时间维度分解输入视频思想的具体落地。模型主要配置参数__init__默认值参数默认值说明num_attention_heads16多头注意力头数attention_head_dim88每个注意力头的通道数num_layers1空间/时间 Transformer 块层数in_channels/out_channelsNone输入/输出通道数out_channels缺省取in_channelscross_attention_dimNone文本嵌入维度交叉注意力caption_channelsNonecaption 投影的输入通道数sample_size64潜空间分辨率用于 patch embedding 位置编码patch_sizeNonepatch embedding 的 patch 尺寸num_embeds_ada_normNone训练时的扩散步数AdaLayerNorm 嵌入数量推理步数不可超过它norm_typelayer_norm归一化类型layer_norm或ada_layer_normLatte 权重实际使用ada_norm_single见测试配置activation_fngeglu前馈网络激活函数video_length16视频帧数决定临时位置编码长度测试与验证仓库如何保障 Latte 质量仓库为 Latte 提供了两套测试可作为你改造或复现时的行为契约Pipeline 测试tests/pipelines/latte/test_latte.pyLattePipelineTesterConfig中通过 dummy 组件LatteTransformer3DModelAutoencoderKLDDIMScheduler tiny T5验证 pipeline 的基本推理、批量一致性、可选组件tokenizer/text_encoder移除后仍可用预计算嵌入推理TestLattePipelineMemory覆盖 CPU offload、group offload 与 layerwise castingTestLattePipelinePyramidAttentionBroadcast与TestLattePipelineFasterCache分别验证 Pyramid Attention Broadcast 与 FasterCache 加速方案对空间/时间/交叉注意力的跳过策略TestLattePipelineIntegration标记slow则用真实权重maxin-cn/Latte-1在加速器上做端到端生成并比对余弦相似度。Transformer 模型测试tests/models/transformers/test_models_transformer_latte.pyLatteTransformerTesterConfig以hidden_states为主输入、(4, 1, 8, 8)为输入形状覆盖前向、内存优化MemoryTesterMixin、注意力AttentionTesterMixin与梯度检查点TrainingTesterMixin验证LatteTransformer3DModel支持gradient_checkpointing。进阶阅读探索调度器速度与质量之间的权衡Schedulers 指南学习如何高效地在多个 pipeline 间复用组件如共享 VAE 或文本编码器Loading 指南组件复用小节了解支持的量化后端与选型建议Quantization 总览【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考