展示用网站,经营范围网站建设,分销系统软件,企业信用公示网上查询平台#x1f52e;混元文生视频官网 | #x1f31f;Github代码仓库 | #x1f3ac; Demo 体验 | #x1f4dd;技术报告 | #x1f60d;Hugging Face 文章目录 论文详解基础介绍数据预处理 #xff08;Data Pre-processing#xff09;数据过滤 (Data Filtering)数据标注 (Data… 混元文生视频官网 | Github代码仓库 | Demo 体验 | 技术报告 | Hugging Face 文章目录 论文详解基础介绍数据预处理 Data Pre-processing数据过滤 (Data Filtering)数据标注 (Data Annotation) 模型结构设计Model Architecture Design3DVAE (3D Variational Auto-encoder Design)统一图像视频生成架构 Unified Image and Video Generative Architecture文本编码器Text Encoder模型缩放Model Scaling 模型预训练Model-Pretraining训练目标Training Objective图像预训练Image Pre-Training视频-图像联合训练Video-Image Joint Training 提示重写Prompt Rewrite高效模型微调High-performance Model Fine-tuning模型加速Model Acceleration推理时间步减少Inference Step Reduction 文本引导蒸馏Text-guidance Distillation高效和扩展训练Efficient and Scalable Training硬件架构Hardware Infrastucture并行策略Parallel Strategy优化Optimization自动容错Automatic fault tolerance 应用基于视频的音频生成Audio Generation based on Video混元图生视频Hunyuan Image-to-Video化身动画Avatar animation 代码实操环境配置Demo运行代码讲解主函数 (sample_video.py)模型推理hyvideo/inference.py推理管道hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py模型结构hyvideo/modules/models.py双流块 MMDoubleStreamBlock单流块 MMSingleStreamBlock混元主干网络HYVideoDiffusionTransformer 论文详解
基础介绍
HunyuanVideo 是一款全新的开源视频生成基础模型其生成性能可与业内顶尖的闭源模型媲美。拥有超过 130 亿个参数是当前规模最大的开源视频生成模型。 该模型集成了数据精选data curation、高级架构设计(architecture design)、渐进式模型扩展与训练progressive model scaling and training以及高效的基础设施 (efficient infrastructure)以支持大规模模型训练与推理。 HunyuanVideo 在视频生成的四个关键方面表现出色视觉质量、运动动态、视频-文本对齐和语义场景切换。 专业评测显示HunyuanVideo 的表现优于 Runway Gen-3、Luma 1.6 以及其他三款顶尖的中文视频生成模型。 通过开源模型代码和应用HunyuanVideo 致力于缩小开源与闭源社区之间的性能差距推动更具活力和创新的视频生成生态体系。
数据预处理 Data Pre-processing
HunyuanVideo 采用了图像-视频联合训练策略将训练数据精细划分为不同类别以满足各自的训练需求。视频数据被划分为五个不同组别图像数据则被划分为两个组别以确保在训练过程中充分发挥数据的特性和优势。本节将重点介绍视频数据的精细化筛选和准备过程。
原始数据池涵盖多个领域包括人物、动物、植物、风景、交通工具、物体、建筑及动画等多种类别的视频。所有视频采集均设定了基本的阈值要求如视频的最小时长等。此外还针对部分数据设定了更高的筛选标准包括空间分辨率、特定宽高比、构图、色彩和曝光等专业要求确保数据在技术质量和美学品质上均达标。
数据过滤 (Data Filtering)
HunyuanVideo分层数据筛选与高质量训练数据构建
(1) 数据采集与预处理 针对质量参差不齐的原始数据采用一系列预处理技术
使用 PySceneDetect 将视频分割为单镜头片段。利用 OpenCV 的拉普拉斯算子提取清晰帧作为视频片段的起始帧。通过内部 VideoCLIP 模型计算视频嵌入向量用于去重基于余弦距离和 k-means 聚类生成约 1 万个概念中心。
(2) 分层数据筛选管道 构建了一个分层数据筛选管道通过多维度的筛选技术来提升数据质量包括
使用 Dover 评估视频片段的美学和技术质量。通过自研模型检测清晰度剔除模糊视频。预测视频的运动速度过滤掉静态或低速运动视频基于光流估计。使用 PySceneDetect 和 TransNet v2 获取场景边界信息。利用内部 OCR 模型 移除带有过多文本或字幕的片段。采用类似 YOLOX 的视觉模型检测并去除水印、边框和标志等遮挡或敏感信息。
通过小规模模型实验验证筛选器的有效性并据此逐步优化数据筛选管道。
(3) 视频与图像训练数据构建 通过分层数据筛选为不同训练阶段构建了五个视频训练数据集视频分辨率从 256 × 256 × 65 逐步提升到 720 × 1280 × 129并根据训练阶段动态调整筛选阈值。
在最终微调阶段构建了一个包含约 100 万个样本的高质量微调数据集。该数据集通过人工标注确保视频在以下两个维度表现优异
视觉美学色彩和谐、光照、主体强调、空间布局动态内容运动速度、动作完整性、运动模糊程度
此外为图像数据建立了类似的分层筛选管道排除了与运动相关的筛选器构建了两个图像训练数据集
初始训练数据集数十亿样本用于第一阶段文本到图像预训练。优化训练数据集数亿样本用于第二阶段文本到图像预训练。
通过分层数据筛选管道与精细化微调数据集的构建HunyuanVideo 确保了在视觉质量、动态细节、文本对齐和语义场景切割等方面的卓越表现为开源视频生成模型树立了新的标杆。 上图展示了分层数据过滤管道。采用各种过滤器进行数据过滤并逐步增加其阈值以构建4个训练数据集即256p360p540p和720p而最终的SFT数据集是通过手动注释构建的。此图突出显示了在每个阶段使用的一些最重要的过滤器。每个阶段都会删除很大一部分数据从上一阶段数据的一半到五分之一不等。这里灰色条表示每个过滤器过滤出的数据量而彩色条表示每个阶段的剩余数据量。
数据标注 (Data Annotation)
(1) 结构化标注技术Structured Captioning 为了提升生成模型的提示响应能力和输出质量开发了内部视觉语言模型VLM为所有图像和视频生成结构化标注。这些标注采用 JSON 格式从多维度提供全面的描述信息包括
短描述Short Description概述场景的主要内容。密集描述Dense Description详细描述场景内容包括与视觉内容结合的场景转换和相机运动。背景Background描述主体所处的环境。风格Style如纪录片、电影感、现实主义、科幻等。镜头类型Shot Type如航拍、特写、中景、长景等。光照Lighting描述视频的光照条件。氛围Atmosphere如温馨、紧张、神秘等。 此外在 JSON 结构中加入了元数据标签如来源标签、质量标签等以增强信息的全面性。
为生成多样化的标注引入了随机丢弃机制及排列组合策略合成不同长度和模式的描述提升生成模型的泛化能力并防止过拟合。
(2) 相机运动类型标注Camera Movement Types 训练了一个相机运动分类器能够预测 14 种相机运动类型包括
变焦放大、缩小 平移上移、下移、左移、右移 俯仰上仰、下俯、左倾、右倾 绕拍绕左、绕右 静态镜头 手持镜头 高置信度的相机运动预测结果被集成到 JSON 格式的结构化标注中从而赋予生成模型对相机运动的控制能力。
模型结构设计Model Architecture Design
3DVAE (3D Variational Auto-encoder Design)
3DVAE压缩与视频图像编码
训练了一个3DVAE模型将像素空间的视频和图像压缩到紧凑的潜在空间。为了同时处理视频和图像采用的事CausalConv3D。
对于一个形状为 ( T 1 ) × 3 × H × W (T 1) \times 3 \times H \times W (T1)×3×H×W的视频 3DVAE 将其压缩为形状为 ( T c t 1 ) × C × ( H c s ) × ( W c s ) (\frac{T}{c_t} 1) \times C \times (\frac{H}{c_s}) \times (\frac{W}{c_s}) (ctT1)×C×(csH)×(csW)的潜在特征。在实现中设定 c t 4 c_t 4 ct4, c s 8 c_s 8 cs8, C 16 C 16 C16。 T 1 T1 T1的原因是对视频的第一帧进行单独的编码作为后续帧的参考帧。 这种压缩方法显著减少了后续扩散变换器模型所需的令牌数量使其能够以原始分辨率和帧率训练视频同时保持较高的效率和质量。模型结构如下图所示。 (1) 训练阶段
训练与损失函数 与以往的大多数研究 不同并未依赖预训练的图像VAE进行参数初始化而是从零开始训练模型。为了平衡视频和图像的重建质量我们按照4:1的比例混合了视频数据和图像数据。 在使用常规的 L 1 L_1 L1重建损失和 KL 损失 L k l L_{kl} Lkl 的基础上我们还引入了感知损失 L l p i p s L_{lpips} Llpips和对抗损失 L a d v L_{adv} Ladv以增强重建质量。完整的损失函数如下所示 Loss L 1 0.1 L l p i p s 0.05 L a d v 1 0 − 6 L k l \text{Loss} L_1 0.1L_{lpips} 0.05L_{adv} 10^{-6}L_{kl} LossL10.1Llpips0.05Ladv10−6Lkl
训练策略 在训练过程中我们采用了课程学习策略逐步从低分辨率短视频训练到高分辨率长视频。为了改善高运动视频的重建效果我们在采样帧时随机选择了 1 至 8 范围内的采样间隔确保从视频剪辑中均匀地抽取帧。
(2) 推理阶段 高分辨率长视频的编码与解码 在单块GPU上编码和解码高分辨率长视频可能会导致显存不足 (OOM) 错误。为了解决这一问题采用了一种时空切片策略将输入视频在空间和时间维度上划分为重叠的切片。每个切片单独编码/解码最终再将输出拼接在一起。对于重叠区域我们使用线性组合进行平滑融合。这一切片策略使我们能够在单块GPU上处理任意分辨率和时长的视频。
训练与推理的一致性 直接在推理时使用切片策略可能会由于训练和推理的不一致性而导致明显的伪影问题。为了解决这一问题我们引入了额外的微调阶段在此阶段随机启用/禁用切片策略进行训练。这种方法确保模型在使用或不使用切片策略时都能够保持训练和推理的一致性。
性能对比 在下表中将所提出的VAE与开源的最先进VAE进行了比较。在视频数据上Hunyuan的VAE在PSNR指标上显著优于其他视频VAE。在图像数据上其性能也超越了视频VAE和图像VAE。
下图展示了多个在256 × 256分辨率下的例子显示Hunyuan提出的VAE在处理文本、人脸和复杂纹理方面具有显著优势。
统一图像视频生成架构 Unified Image and Video Generative Architecture
HunyuanVideo 中的 Transformer 设计其采用了统一的全注意力机制 (Full Attention)并基于以下三大理由
性能优越全注意力机制相较于分离的时空注意力方法, 展现了更卓越的性能。统一生成该机制支持图像和视频的统一生成简化了训练过程并提升了模型的可扩展性。高效加速全注意力机制能够更有效地利用现有的大型语言模型LLM相关的加速能力进一步增强训练和推理效率。
模型的具体结构如下
(1)输入设计
对于给定的视频-文本对模型在3D潜在空间中运行。具体地 视频分支 对应上图的橙色小方块 输入首先被压缩为形状为 T × C × H × W T \times C \times H \times W T×C×H×W 的潜在表示。为统一输入处理模型将图像视为单帧视频。这些潜在表示通过一个 3D 卷积核大小为 k t × k h × k w k_t \times k_h \times k_w kt×kh×kw被划分为补丁并展开为 1D 序列令序列长度为 T k t ⋅ H k h ⋅ W k w \frac{T}{k_t} \cdot \frac{H}{k_h} \cdot \frac{W}{k_w} ktT⋅khH⋅kwW。 文本分支 (对应上图的绿色小方块) 首先使用一个先进的LLM模型将文本编码为嵌入序列以捕获细粒度的语义信息。同时利用 CLIP 模型提取一个包含全局信息的池化文本表示。
随后这个文本表示会通过扩展维度后与时间步嵌入相加再输入模型。
(2)模型设计
为了有效融合文本和视觉信息模型采用“双流到单流”的混合设计策略专注于视频生成任务 双流阶段 视频和文本的 token 通过多个 Transformer 块分别独立处理。这种独立处理方式使得每种模态能够学习其专属的调制机制避免互相干扰。 单流阶段 将视频和文本的 token 进行拼接并输入后续的 Transformer 块。通过这种方式实现对多模态信息的高效融合捕获视觉信息和语义信息之间的复杂交互。
这一设计显著增强了模型对多模态信息的理解能力从而提升了整体性能。
(3) 位置嵌入
为了支持多分辨率、多宽高比以及不同时长的生成任务模型在每个 Transformer 块中使用了 旋转位置嵌入Rotary Position Embedding, RoPE。以下是关键设计点 RoPE的基本优势 RoPE通过对嵌入应用旋转频率矩阵增强了模型捕获绝对位置和相对位置关系的能力。在大型语言模型LLM中RoPE表现出一定的外推能力。 扩展至视频的三维嵌入 由于视频数据中增加了时间维度的复杂性RoPE被扩展到三维空间时间 (T)、高度 (H)、宽度 (W)。对于每个维度分别计算旋转频率矩阵进一步提升时空建模能力。 具体实现 将查询 (query) 和键 (key) 的特征通道划分为三个段 d t , d h , d w d_t, d_h, d_w dt,dh,dw。每个段分别与对应维度时间、高度、宽度的旋转频率矩阵相乘。最后将这些处理后的段进行拼接生成包含位置信息的查询和键嵌入。 效果 该位置感知的查询和键嵌入在注意力计算中发挥作用使模型能更好地捕获三维空间中的复杂时空关系。
此设计显著提升了模型对视频生成中不同分辨率、宽高比和时长变化的适应能力。
模型的具体超参数设置如下表
文本编码器Text Encoder
在文本生成图像和视频的任务中文本编码器通过提供潜在空间中的指导信息起到关键作用。以下是 HunyuanVideo 模型在文本编码方面的创新设计
(1 采用多模态大型语言模型MLLM作为文本编码器 与传统方法相比HunyuanVideo 采用了一种基于 Decoder-Only 架构的预训练多模态大型语言模型 (MLLM) 作为文本编码器。这种选择具有以下优势 与 T5 的对比 T5 使用的是 Encoder-Decoder 结构而经过视觉指令微调的 MLLM 在特征空间中的图文对齐能力更强能够降低扩散模型中指令跟随的难度。 与 CLIP 的对比 CLIP 使用 Transformer Encoder 架构MLLM 在图像细节描述和复杂推理方面表现更优。 零样本学习能力 MLLM 可通过在用户提示语前添加系统指令作为一个零样本学习者帮助文本特征更关注关键信息。 2双向注意力增强 虽然 MLLM 使用基于因果注意力的 Decoder-Only 结构而 T5-XXL 使用的是双向注意力这使 T5 能为扩散模型提供更好的文本指导。因此
HunyuanVideo 引入了双向 Token 精炼器 (bidirectional token refiner)用于增强文本特征结合了因果注意力和双向注意力的优点。
3灵活配置的多种 MLLMs 根据不同任务需求HunyuanVideo 配置了一系列 MLLMs在各种设置下均表现出比传统文本编码器更优的性能。
4CLIP 特征的全局指导作用 尽管 MLLM 是主要的文本编码器CLIP 文本特征作为文本信息的摘要也很有价值
全局指导HunyuanVideo 采用 CLIP-Large 文本特征的最终非填充 token 作为全局指导信息集成到双流和单流的 DiTDiffusion Transformer块中。
综上HunyuanVideo 通过结合 MLLM 和 CLIP 的特性在文本指导和生成任务中实现了更高的对齐能力、更细致的图像生成和更强的多模态推理能力。
模型缩放Model Scaling
Neural Scaling Laws为理解和优化机器学习模型的性能提供了一种强有力的工具。通过揭示模型大小 ( N N N)、数据集大小 ( D D D) 和计算资源 ( C C C) 之间的关系这些规律推动了更高效和更强大模型的开发从而提升了大规模模型训练的成功率。以下是其在 HunyuanVideo 的应用
(1) 传统扩展规律的局限性 尽管此前已有关于大语言模型和图像生成模型的扩展规律研究视频生成模型通常依赖于预训练的图像模型。这种依赖性为视频生成任务引入了独特的挑战因为视频数据比静态图像具有更高的维度和复杂性。
(2) 基础扩展规律的建立 为了克服上述挑战研究的第一步是针对文本到图像生成模型建立基础扩展规律
确定文本到图像任务中的模型大小、数据集大小和计算需求之间的关系。这一基础规律为视频生成模型的扩展提供了参考框架。
(3) 从文本到图像到文本到视频的扩展 在文本到图像扩展规律的基础上进一步推导了适用于文本到视频模型的扩展规律
考虑到时间维度的复杂性结合了视频生成所需的额外资源和模型调整。推导结果为视频生成任务提供了数据集规模、模型容量和计算资源的最优配置方案。
(4) 两种扩展规律的集成 通过将文本到图像和文本到视频的扩展规律相结合研究团队能够系统性地确定
模型的设计与规模视频生成任务的最佳模型架构和大小。数据和资源的配置数据集规模与计算资源的平衡最大限度提高生成性能。
总结 HunyuanVideo 的研究以神经扩展规律为指导系统地设计了文本到视频生成模型的配置方案。通过整合图像生成和视频生成的扩展规律不仅优化了模型的训练效率还推动了视频生成任务性能的显著提升。 扩展规律在视觉生成模型中的应用 1. 研究背景 Kaplan 等人和 Hoffmann 等人 研究了语言模型的经验扩展规律尤其是交叉熵损失与模型规模的关系。在视觉生成领域
Li 等人研究了基于扩散的生成模型如 UNet的扩展性质。基于 Transformer 的方法如 DiT, U-ViT, Lumina-T2X, SD3则关注采样质量与网络复杂性的关系但未探索计算资源与均方误差MSE损失的幂律关系。
2. DiT-T2X 模型家族 为填补这一研究空白作者开发了一个 DiT-like 模型家族命名为 DiT-T2X
模型特点 T5-XXL作为文本编码器。前述 3D VAE 作为图像编码器。文本信息通过交叉注意力层注入模型。 模型规模 包括 7 种规模从 92M 到 6.6B 参数不等。 训练设置 使用 DDPM 和 v-prediction训练。数据集为分辨率 256px 的一致数据集保持一致的超参数设置。
3. 实验方法与扩展规律 作者借鉴 Hoffmann 等人的实验方法构建了神经扩展规律
公式定义 N opt a 1 C b 1 , D opt a 2 C b 2 N_{\text{opt}} a_1 C^{b_1}, \quad D_{\text{opt}} a_2 C^{b_2} Nopta1Cb1,Dopta2Cb2 N opt N_{\text{opt}} Nopt最优模型规模。 D opt D_{\text{opt}} Dopt最优数据集大小。 C C C计算资源以 Peta FLOPs 为单位。参数 a 1 , a 2 , b 1 , b 2 a_1, a_2, b_1, b_2 a1,a2,b1,b2 从实验中拟合得到。 实验结果 a 1 5.48 × 1 0 − 4 , b 1 0.5634 a_1 5.48 \times 10^{-4}, b_1 0.5634 a15.48×10−4,b10.5634 a 2 0.324 , b 2 0.4325 a_2 0.324, b_2 0.4325 a20.324,b20.4325表明模型规模和数据规模与计算资源之间遵循幂律关系。
4. 模型优化
在不同计算预算下通过扩展规律计算出模型和数据集的最优配置。下图显示了不同规模模型的损失曲线证明 DiT-T2XI家族符合幂律关系。
DiT-T2X 家族的研究填补了扩散模型在计算资源与 MSE 损失扩展规律方面的空白。通过建立神经扩展规律作者系统地优化了文本到图像/视频生成任务的模型和数据配置推动了视觉生成领域的进一步发展。 视频模型扩展规律Video Model Scaling Law
1. 基于图像扩展规律的初始化
本研究以 T2X(I) 模型的扩展规律为基础选择每个规模模型对应的最优图像检查点即损失曲线包络线上的模型作为视频扩展规律实验的初始化模型。通过这种方式T2X(V) 模型的训练能够充分利用图像模型的性能优势。
2. 视频扩展规律结果
实验结果 如上图d,e,f展示了 T2X(V) 模型的扩展规律拟合得到以下参数 a 1 0.0189 , b 1 0.3618 a_1 0.0189, b_1 0.3618 a10.0189,b10.3618 a 2 0.0108 , b 2 0.6289 a_2 0.0108, b_2 0.6289 a20.0108,b20.6289 模型规模 N opt N_{\text{opt}} Nopt和数据规模 D opt D_{\text{opt}} Dopt 与计算资源 C C C的幂律关系 N opt a 1 C b 1 , D opt a 2 C b 2 N_{\text{opt}} a_1 C^{b_1}, \quad D_{\text{opt}} a_2 C^{b_2} Nopta1Cb1,Dopta2Cb2
3. 模型规模选择
结合上图b,e的结果并综合考虑训练开销和推理成本最终将模型规模设定为 13B。根据图像和视频扩展规律计算出分别用于图像训练和视频训练的 token 数量详见上图(c,f)。 4. 分阶段训练的扩展规律需要注意的是 基于图像和视频扩展规律计算的训练 token 数量仅适用于图像和视频训练的 第一阶段。从低分辨率到高分辨率的渐进式训练的扩展特性将在未来的工作中进一步探索。 总结 视频模型扩展规律通过整合图像扩展规律结果构建了 T2X(V) 模型的高效训练流程。结合计算资源与训练消耗模型规模优化为 13B并确定了训练 token 的合理分配方案。
模型预训练Model-Pretraining
使用Flow Matching进行模型训练并将训练过程分为多个阶段。 首先在256px和512px的图像上预训练我们的模型然后在256px到960px的图像和视频上进行联合训练。
训练目标Training Objective
1. 框架概述
本研究使用 Flow Matching 框架 来训练图像和视频生成模型。Flow Matching 的核心思想是 目标将复杂的概率分布通过一系列变量变换转换为简单的概率分布。生成通过逆变换从简单分布生成新的数据样本。
2. 训练过程
输入表示使用训练集中图像或视频的潜在表示 x 1 \mathbf{x}_1 x1。时间采样 从 Logit-Normal 分布采样时间 t ∈ [ 0 , 1 ] t \in [0, 1] t∈[0,1]。初始化噪声样本 x 0 ∼ N ( 0 , I ) \mathbf{x}_0 \sim \mathcal{N}(0, \mathbf{I}) x0∼N(0,I)遵循高斯分布。 插值方法 通过线性插值方法构建训练样本 x t \mathbf{x}_t xt。 模型训练 目标是预测速度场 u t d x t d t \mathbf{u}_t \frac{d \mathbf{x}_t}{dt} utdtdxt指导样本 x t \mathbf{x}_t xt向样本 x 1 \mathbf{x}_1 x1移动。优化参数通过最小化预测速度 v t \mathbf{v}_t vt和真实速度 u t \mathbf{u}_t ut的均方误差 (MSE) 来优化模型参数 L generation E t , x 0 , x 1 ∥ v t − u t ∥ 2 \mathcal{L}_{\text{generation}} \mathbb{E}_{t, \mathbf{x}_0, \mathbf{x}_1} \| \mathbf{v}_t - \mathbf{u}_t \|^2 LgenerationEt,x0,x1∥vt−ut∥2
3. 推理过程
初始噪声样本从高斯分布 N ( 0 , I ) \mathcal{N}(0, \mathbf{I}) N(0,I)中抽取初始样本 x 0 \mathbf{x}_0 x0。解算器 使用一阶 Euler 常微分方程 (ODE) 求解器结合模型预测的 d x t / d t d\mathbf{x}_t / dt dxt/dt估计值逐步计算 x 1 \mathbf{x}_1 x1。生成样本最终生成 x 1 \mathbf{x}_1 x1即目标图像或视频。 总结 Flow Matching 框架通过建立速度场预测模型在训练过程中有效拟合复杂分布的动态变化并通过解算器实现高效的图像和视频生成。这一方法为高质量生成提供了理论和实践基础。
图像预训练Image Pre-Training
背景和动机
早期实验表明预训练模型显著加速了视频训练的收敛速度并提升了视频生成性能。为此提出了一种两阶段的渐进式图像预训练策略用于视频训练的热启动。 阶段 1256px 图像训练
目标模型首先在低分辨率256px图像上进行预训练。策略 多尺度训练在 256px 图像上启用多长宽比训练。优势 帮助模型学习生成宽广长宽比范围内的图像。避免了图像预处理阶段因裁剪操作导致的文本图像对齐问题。使用低分辨率样本可让模型从更大规模样本中学习低频概念。 作用为后续训练奠定基础使模型能够有效处理低分辨率样本。 阶段 2混合尺度训练
目标增强模型在高分辨率如 512px上的能力。问题 直接微调问题如果直接在 512px 图像上微调模型可能导致模型在 256px 图像生成上的性能严重退化从而影响后续视频预训练。 解决方案 提出 混合尺度训练 方法 在每次训练的全局批次中引入两个或多个尺度的多长宽比 buckets。每个尺度设定一个锚定大小anchor size。基于锚定大小构建多长宽比 buckets。 训练数据集 包括两个尺度数据集256px 和 512px 图像。同时学习高分辨率和低分辨率图像的生成能力。 其他优化 引入动态批次大小dynamic batch sizes用于不同图像尺度。效果最大化 GPU 内存和计算效率的利用率。 总结 通过两阶段渐进式预训练策略模型既能适应低分辨率图像的大规模训练又能平衡高分辨率图像的生成能力为后续视频训练提供了强大的支持和灵活性。
视频-图像联合训练Video-Image Joint Training 多长宽比和时长分桶策略 (Multiple aspect ratios and durations bucketization)
背景
数据过滤后视频数据具有不同的长宽比和时长。为了高效利用这些数据 创建了基于时长和长宽比的分桶策略 构建 B T B_T BT时长桶和 B A R B_{AR} BAR长宽比桶。产生总计 B T × B A R B_T \times B_{AR} BT×BAR 的桶数量。 优化分桶训练 每个桶分配一个最大批次大小以避免 GPU 内存溢出OOM。训练前将所有数据分配到最近的桶。训练中随机从每个桶中预取批次数据。 优势使模型在每一步都能处理不同尺寸的数据避免局限于单一尺寸训练增强模型的泛化能力。 渐进式视频-图像联合训练 (Progressive Video-Image Joint Training.)
背景
直接从文本生成高质量长时视频 常导致模型难以收敛、结果次优。 为此设计了一种渐进式课程学习策略 以 T2I 模型的参数初始化。逐步提升视频的时长和分辨率。
三阶段策略 低分辨率、短视频阶段 (Low-resolution, short video stage) 目标建立文本到视觉内容的基础映射。确保短期动作的一致性和连贯性。 低分辨率、长视频阶段 (Low-resolution, long video stage) 目标学习更复杂的时间动态和场景变化。确保更长时间范围内的时空一致性。 高分辨率、长视频阶段 (High-resolution, long video stage) 目标提升视频的分辨率和细节质量。同时维持时间连贯性并管理复杂的时间动态。 附加策略
每个阶段都按比例引入图像联合训练 提高高质量视频数据的稀缺性。扩展模型的多样性和知识覆盖面。优点 避免因视频与图像数据分布差异导致的灾难性遗忘。提升视频和图像语义空间的一致性。 总结 通过多阶段课程学习和灵活的分桶策略实现了模型从文本到长时高质量视频的高效生成具备更强的泛化性和动态管理能力。
提示重写Prompt Rewrite 背景
为应对用户提供的提示prompt在语言风格和长度上的多样性引入了 Hunyuan-Large 模型作为 prompt 重写工具。该模块能够在无训练框架下运行通过结合详细的指令和上下文学习示例提升模型的性能。 模块核心功能 多语言输入适配Multilingual Input Adaptation 设计用于处理和理解各种语言的用户提示。确保语义和上下文得以保留。 提示结构标准化Standardization of Prompt Structure 将提示重新表述为符合标准化信息结构的形式。类似于模型训练中使用的描述性标签。 复杂术语简化Simplification of Complex Terminology 将复杂的用户措辞转化为更直接、易懂的表达。同时保持用户意图不变。 附加策略
自我修订技术Self-Revision Technique 对最终的提示进行细化和优化。通过对原始提示和重写提示的比较分析确保输出内容既精准又与模型能力一致。 进一步优化
采用 LoRA 微调Low-Rank Adaptation进一步简化模块应用流程 微调训练使用来自无训练方法生成的高质量 prompt 重写对。提升了模块的应用精度同时保持实现过程高效。 总结 该模块通过多语言适配、标准化结构以及简化术语的策略显著优化了用户提示的处理效果使得生成结果更符合用户期望并提升模型适配性。
高效模型微调High-performance Model Fine-tuning 背景 在预训练阶段模型使用了大规模数据集进行训练。虽然数据集信息丰富但其数据质量存在较大差异性。 微调策略 为构建能够生成高质量动态视频并在连续动作控制和角色动画方面表现卓越的生成模型我们采取了以下关键措施 精选数据子集 从完整数据集中精心挑选了四个特定子集用于微调。数据筛选分为两个阶段 自动化数据过滤利用算法对数据进行初步筛查。人工审查确保最终数据的高质量和适用性。 模型优化策略 实施多种优化技术最大限度提高生成性能。优化重点包括视频生成质量、动态控制的精准性及角色动画的自然流畅性。 总结 通过筛选优质子集并采用先进的优化技术该微调流程显著提升了模型生成高质量、动态视频的能力同时强化了其在连续动作和角色动画方面的表现。
模型加速Model Acceleration
推理时间步减少Inference Step Reduction
相比图像生成视频生成在减少推理步骤的同时维持空间质量和时间质量更加具有挑战性。 策略减少推理步骤 为了解决这一问题我们重点研究如何减少视频生成所需的推理步骤数量。 时间步调整Time-Step Shifting 受到早期时间步在生成过程中贡献最大这一观察的启发我们提出了一种时间步调整策略。具体过程 假设推理步 q ∈ { 1 , 2 , … , Q } q \in \{1, 2, \dots, Q\} q∈{1,2,…,Q}输入的时间条件为 t 1 − q Q t 1 - \frac{q}{Q} t1−Qq其中 t t t 表示噪声初始化为 t 1 t 1 t1并在 t 0 t 0 t0时停止生成。为优化推理步骤我们将 t t t 转换为 t ′ t t′即 t ′ s t 1 ( s − 1 ) t t \frac{s t}{1 (s - 1) t} t′1(s−1)tst 其中 t ′ t t′为调整后的时间条件 s s s是一个调整因子。 关键观察 当 ( s 1 ) 时模型的推理更加集中于早期时间步。对于推理步骤较少的情况需要较大的 ( s ) 值 当推理步骤为 50 时经验上 ( s ) 设置为 7。当推理步骤小于 20 时( s ) 应增至 17。 优势 时间步调整策略能够让生成模型在减少推理步骤的同时达到更高推理步数的质量效果。 对比与结果
MovieGen 采用线性-二次调度器来实现类似目标但其性能在极低推理步如 10 步下不如我们的时间步调整方法。结论 时间步调整方法在极低推理步情况下效果更优而线性-二次调度器则在视觉质量上表现更差。
文本引导蒸馏Text-guidance Distillation 背景
Classifier-free Guidance (CFG) 是一种能够显著提升文本条件扩散模型样本质量和运动稳定性的技术。问题 CFG 的一个主要缺点是其高计算成本和推理延迟 每个推理步骤都需要为无条件输入生成额外的输出。尤其是在处理大型视频模型和高分辨率视频生成时同时生成文本条件与无条件视频的推理负担极其昂贵。 解决方案引导蒸馏 为了解决上述问题我们提出了一种蒸馏方法将无条件和有条件输入的组合输出压缩到一个学生模型中。 具体方法 学生模型 以引导比例guidance scale为条件输入。与教师模型共享相同的结构和超参数。 初始化 使用与教师模型相同的参数对学生模型进行初始化。 训练 引导比例随机从 1 到 8 中采样进行训练。 实验结果 通过文本引导蒸馏我们实现了约 1.9 倍的推理加速。 优势
蒸馏方法显著降低了推理成本和延迟同时保持了文本条件生成的高质量和运动稳定性。在大型视频模型场景中尤其适用于高分辨率视频生成任务。
高效和扩展训练Efficient and Scalable Training
背景 在大规模分布式训练中高效的通信是关键设计了一个专用的分布式训练框架。 解决方案 Tencent XingMai 网络 专门为实现高效的服务器间通信而构建的分布式训练框架。旨在优化大规模分布式系统中的通信性能。 GPU 任务调度 所有训练任务的 GPU 调度通过 Tencent Angel 机器学习平台完成。该平台提供强大的资源管理和任务调度能力有效提升训练效率。 优势
高效通信专用的分布式框架确保了服务器间的数据传输性能。资源优化借助 Angel 平台的调度功能最大化利用了硬件资源减少了资源浪费和调度延迟。
硬件架构Hardware Infrastucture 训练框架 我们使用腾讯 Angel 机器学习团队开发的大规模预训练框架 AngelPTM 对 HunyuanVideo 模型进行训练。该框架专为大规模模型的扩展性和高效训练而设计。 硬件与基础设施
硬件支持训练基于高性能计算资源构建提供强大的硬件支持。基础设施集成了分布式计算与存储系统确保高效的数据传输与处理。 并行训练与优化 模型并行方法 利用模型并行策略将模型参数分布到多个设备或节点上以支持更大规模模型的训练。通过优化并行通信机制降低训练中的通信开销。 优化方法 针对计算负载与通信延迟调整并行粒度和任务分配策略。引入动态负载均衡确保资源利用最大化。 容错机制
自动容错 实现了自动化故障容错机制可在节点或设备故障时自动切换备选资源。降低训练中断的风险提高训练过程的稳定性和可靠性。 优势
可扩展性强AngelPTM 提供了灵活的并行计算能力适配超大规模模型的训练需求。高效性优化的模型并行方法与容错机制显著提升了训练速度和资源利用效率。鲁棒性自动容错机制确保了训练过程的稳定性即使在复杂的大规模分布式环境中依然可靠。
并行策略Parallel Strategy
以下是对 HunyuanVideo 训练中使用的 5D 并行策略的介绍 5D 并行策略 HunyuanVideo 训练采用了五种并行策略包括
张量并行 (Tensor Parallelism, TP)序列并行 (Sequence Parallelism, SP)上下文并行 (Context Parallelism, CP)数据并行 (Data Parallelism, DP)ZeroCache 优化 1. 张量并行 (Tensor Parallelism, TP)
基本原理基于块矩阵计算的原则将模型参数 (张量) 分割并分配到不同的 GPU 上。目的 减少单个 GPU 的内存使用。加速矩阵运算。 GPU 任务分配每个 GPU 负责计算模型层中不同部分张量的操作。 2. 序列并行 (Sequence Parallelism, SP)
原理在 TP 的基础上对输入序列维度进行切片。优化目标 减少 LayerNorm 和 Dropout 等操作的重复计算。减少激活值的存储避免计算资源和 GPU 内存的浪费。 扩展能力 对不满足 SP 要求的输入数据引入等效 SP Padding 工程能力增强通用性。 3. 上下文并行 (Context Parallelism, CP)
特性在序列维度切片以支持长序列训练。任务分配 每个 GPU 负责计算不同序列切片的 Attention。采用 Ring Attention 机制提高 Attention 的计算效率。 4. 数据并行 ZeroCache 优化
作用 利用数据并行实现横向扩展以满足日益增长的训练数据需求。ZeroCache 优化策略进一步减少模型状态 (包括参数、梯度和优化器状态) 的冗余。 结果 显著降低 GPU 内存开销。提高数据并行的整体效率。 主要优势
高效长序列训练通过多 GPU 分工计算突破单 GPU 内存限制提升长序列处理效率。资源节约利用序列切片和 ZeroCache 优化减少内存与计算资源浪费。可扩展性5D 并行策略为大规模模型的高效训练提供了强大的支持。
优化Optimization Attention 优化
背景 随着序列长度的增加注意力计算成为训练的主要瓶颈。加速方法 采用 FusedAttention 技术显著加速注意力计算优化长序列的训练效率。 重计算与激活卸载优化 (Recomputation and Activations Offload Optimization) 重计算技术 (Recomputation) 核心思想通过增加计算量以换取存储资源。技术细节 指定重计算层或块选择特定层或块进行重计算。释放前向计算的激活值减少前向计算中激活值的 GPU 存储需求。反向计算中再计算激活值通过重新计算获取依赖的激活值从而显著降低训练中 GPU 内存的占用。 激活卸载 (Activation Offload) 策略基于层级的激活卸载方法。实施方式 在不降低训练性能的前提下将 GPU 内存中的激活值卸载到主机内存 (Host Memory)。 优势 进一步节省 GPU 内存使用。考虑了 PCIe 带宽和主机内存大小限制平衡性能与资源消耗。 主要优势
加速训练通过优化注意力计算加速长序列训练过程。节省内存采用重计算和激活卸载技术有效降低 GPU 内存占用。提升扩展性优化了高序列长度下的大规模分布式训练性能。
自动容错Automatic fault tolerance
HunyuanVideo 的大规模训练稳定性通过引入自动容错机制得到了保障。该机制能够快速应对常见的硬件故障避免频繁的人工任务恢复操作。通过自动检测错误并迅速替换健康节点以恢复训练任务系统实现了 99.5% 的训练稳定性。
应用
基于视频的音频生成Audio Generation based on Video 混元图生视频Hunyuan Image-to-Video 化身动画Avatar animation 代码实操
环境配置 AutoDL镜像 参考【文生视频-腾讯混元-HunyuanVideo】AutoDL配好环境开箱即用 在Autodl租用服务器后选择【社区镜像】并在其中搜索HunyuanVideo 提示1720p x 1280p 分辨率大约需要76G显存544p x 960p分辨率大约需要43G显存合理租卡 2预训练模型比较大至少扩容50GB 安装huggingface_hub
python -m pip install huggingface_hub[cli]下载预训练模型 在HunyuanVideo下新建ckpts文件夹
# 学术资源加速
source /etc/network_turbo
cd HunyuanVideo
# 下载HunyuanVideo 预训练模型
huggingface-cli download tencent/HunyuanVideo --local-dir ./ckpts
# 下载MLLM的Text Encoder模型
huggingface-cli download xtuner/llava-llama-3-8b-v1_1-transformers --local-dir ./ckpts/llava-llama-3-8b-v1_1-transformers
# 下载CLIP预训练模型
huggingface-cli download openai/clip-vit-large-patch14 --local-dir ./ckpts/text_encoder_2模型预处理
cd /root/HunyuanVideo/
python hyvideo/utils/preprocess_text_encoder_tokenizer_utils.py --input_dir ckpts/llava-llama-3-8b-v1_1-transformers --output_dir ckpts/text_encoderDemo运行
具体的参数配置见hyvideo/config.py
720p x 1280p 分辨率
python3 sample_video.py \--video-size 720 1280 \--video-length 129 \--infer-steps 50 \--prompt A cat walks on the grass, realistic style. \--flow-reverse \--use-cpu-offload \--save-path ./results544p x 960p 分辨率
python3 sample_video.py \--video-size 544 960 \--video-length 129 \--infer-steps 30 \--prompt A sexy beauty lying on the beach, realistic style. \--flow-reverse \--use-cpu-offload \--save-path ./results在单卡L20上生成544p x 960p 129frames 30steps的视频耗时约22分钟。 代码讲解
主函数 (sample_video.py)
sample_video.py中main() 函数用于加载视频生成模型、生成视频样本并保存生成的结果。
def main():# 调用 parse_args() 函数来解析命令行参数并将结果存储在 args 中。args parse_args()print(args)#使用 args.model_base 指定模型的基础路径。检查路径是否存在如果路径不存在抛出异常models_root_path Path(args.model_base)if not models_root_path.exists():raise ValueError(fmodels_root not exists: {models_root_path})# 创建保存目录save_path args.save_path if args.save_path_suffix else f{args.save_path}_{args.save_path_suffix}if not os.path.exists(args.save_path):os.makedirs(save_path, exist_okTrue)# 加载指定路径下的预训练模型并传入 args 参数hunyuan_video_sampler HunyuanVideoSampler.from_pretrained(models_root_path, argsargs)# 更新 args 为模型内部参数确保之后使用的一致性args hunyuan_video_sampler.args# 生成视频样本outputs hunyuan_video_sampler.predict(promptargs.prompt, # 文本提示词用于引导生成内容。heightargs.video_size[0],# 视频帧的分辨率高度和宽度widthargs.video_size[1],video_lengthargs.video_length, # 视频时长帧数。seedargs.seed, # 随机种子用于结果的可重复性negative_promptargs.neg_prompt, # 负向提示词指定生成时需避免的特性infer_stepsargs.infer_steps, # 推理步数guidance_scaleargs.cfg_scale,# 引导系数控制生成质量。num_videos_per_promptargs.num_videos, # 每个提示生成的视频数量。flow_shiftargs.flow_shift, # 时间帧间的流动控制参数batch_sizeargs.batch_size, # 批处理大小embedded_guidance_scaleargs.embedded_cfg_scale # 内嵌引导系数用于调节特定特征)samples outputs[samples]# 保存视频样本# 检查环境变量 LOCAL_RANK 是否存在用于分布式训练的本地进程控制如果不存在或者值为 0即主进程则继续保存样本。if LOCAL_RANK not in os.environ or int(os.environ[LOCAL_RANK]) 0:# 遍历生成的samplesfor i, sample in enumerate(samples):sample samples[i].unsqueeze(0)# 添加时间戳time_flag datetime.fromtimestamp(time.time()).strftime(%Y-%m-%d-%H:%M:%S)save_path f{save_path}/{time_flag}_seed{outputs[seeds][i]}_{outputs[prompts][i][:100].replace(/,)}.mp4save_videos_grid(sample, save_path, fps24) # 保存视频logger.info(fSample save to: {save_path}) # 日志记录模型推理hyvideo/inference.py load_diffusion_pipeline函数 用于加载视频生成的扩散推理管道pipeline包括模型、调度器scheduler、设备配置等必要的组件。 def load_diffusion_pipeline(self,args,vae,text_encoder,text_encoder_2,model,schedulerNone,deviceNone,progress_bar_configNone,data_typevideo,):Load the denoising scheduler for inference.# 去噪调度器的初始化if scheduler is None:if args.denoise_type flow:# 流动匹配的去噪策略离散去噪调度器可能用于视频生成任务中时间帧之间的一致性建模。# 负责指导扩散模型逐步还原噪声生成清晰的视频帧。scheduler FlowMatchDiscreteScheduler(shiftargs.flow_shift, # 流动偏移值。reverseargs.flow_reverse, # 是否反向计算。solverargs.flow_solver, # 去噪求解器的类型)else:raise ValueError(fInvalid denoise type {args.denoise_type})# 构建推理pipelinepipeline HunyuanVideoPipeline(vaevae, # 负责特征编码和解码的模块。text_encodertext_encoder, # 用于处理文本提示生成与视频生成相关的特征。text_encoder_2text_encoder_2,transformermodel, # 主扩散模型生成视频的核心模块。schedulerscheduler, # 去噪调度器控制扩散生成的时间步长和过程progress_bar_configprogress_bar_config, # 可选的进度条配置用于显示推理进度。argsargs, # 配置参数的集合)# 配置计算资源if self.use_cpu_offload:# 将部分计算任务卸载到 CPU。这是显存不足时的优化策略可以大幅降低 GPU 的显存占用。pipeline.enable_sequential_cpu_offload()else:# 如果为 False直接将管道加载到指定的 device如 GPU上运行pipeline pipeline.to(device)return pipelineget_rotary_pos_embed函数 其作用是计算用于视频处理的多维旋转位置嵌入Rotary Positional Embedding, RoPE。它主要根据输入的视频维度、网络配置以及位置嵌入参数生成对应的正弦和余弦频率嵌入用于在 Transformer 等模型的注意力机制中进行位置编码。 def get_rotary_pos_embed(self, video_length, height, width):# video_length: 视频的帧长度。# height, width: 视频的帧高和帧宽。 目标是根据这些维度计算位置嵌入。# 表示生成的 RoPE 的目标维度3D: 时间维度 空间高度和宽度target_ndim 3# 推导潜在特征latent feature所需维度的辅助变量ndim 5 - 2# 根据 self.args.vae 中的配置例如 VAE 模型类型 884 或 888确定潜在特征的空间尺寸 latents_size# 884: 时间维度下采样 4 倍1/4空间高宽下采样 8 倍1/8。if 884 in self.args.vae:latents_size [(video_length - 1) // 4 1, height // 8, width // 8]# 888: 时间维度下采样 8 倍1/8空间高宽下采样 8 倍。elif 888 in self.args.vae:latents_size [(video_length - 1) // 8 1, height // 8, width // 8]# 默认情况下不对时间维度下采样但高宽依然下采样 8 倍。else:latents_size [video_length, height // 8, width // 8]# 检查潜在空间尺寸是否与 Patch 尺寸兼容# 如果 self.model.patch_size 是单个整数检查潜在特征维度的每一维是否能被 patch_size 整除。if isinstance(self.model.patch_size, int):assert all(s % self.model.patch_size 0 for s in latents_size), (fLatent size(last {ndim} dimensions) should be divisible by patch size({self.model.patch_size}), fbut got {latents_size}.)# 如果整除计算 RoPE 的输入尺寸 rope_sizes将 latents_size 每一维除以 patch_sizerope_sizes [s // self.model.patch_size for s in latents_size]# 如果 self.model.patch_size 是一个列表分别对每一维进行整除检查和计算。elif isinstance(self.model.patch_size, list):assert all(s % self.model.patch_size[idx] 0for idx, s in enumerate(latents_size)), (fLatent size(last {ndim} dimensions) should be divisible by patch size({self.model.patch_size}), fbut got {latents_size}.)rope_sizes [s // self.model.patch_size[idx] for idx, s in enumerate(latents_size)]# 如果 rope_sizes 的维度数不足 target_ndim在开头补充时间维度值为 1。if len(rope_sizes) ! target_ndim:rope_sizes [1] * (target_ndim - len(rope_sizes)) rope_sizes # time axis# head_dim 是单个注意力头的维度大小由模型的 hidden_size 和 heads_num 计算得出。head_dim self.model.hidden_size // self.model.heads_num# rope_dim_list 是用于位置嵌入的维度分配列表# 如果未定义默认将 head_dim 平均分配到 target_ndim时间、高度、宽度。rope_dim_list self.model.rope_dim_listif rope_dim_list is None:rope_dim_list [head_dim // target_ndim for _ in range(target_ndim)]assert (sum(rope_dim_list) head_dim), sum(rope_dim_list) should equal to head_dim of attention layer# 调用 get_nd_rotary_pos_embed 函数计算基于目标尺寸 rope_sizes 和维度分配 rope_dim_list 的多维旋转位置嵌入。freqs_cos, freqs_sin get_nd_rotary_pos_embed(rope_dim_list,rope_sizes,thetaself.args.rope_theta, #控制位置嵌入频率。use_realTrue, # 表示使用真实数值而非复数形式。theta_rescale_factor1, # 无缩放因子。)#返回 freqs_cos: 余弦频率嵌入。freqs_sin: 正弦频率嵌入。return freqs_cos, freqs_sin predict 函数 一个用于从文本生成视频或图像的预测函数名为 predict。函数通过输入文本 prompt结合其他参数如视频分辨率、帧数、推理步数等生成指定数量的视频或图像。
参数说明 prompt (str or List[str]): 输入的文本描述生成的图像或视频将基于此描述。可以是单个字符串或字符串列表。 height (int): 生成视频的高度分辨率。默认值为 192。 width (int): 生成视频的宽度分辨率。默认值为 336。 video_length (int): 生成视频的帧数。默认值为 129。 seed (int or List[int]): 随机种子用于控制生成结果的随机性。如果在分布式环境下运行即模型被并行化必须指定种子以保证结果一致性。 negative_prompt (str or List[str]): 用于指定负面提示词影响生成结果。可以帮助模型避免生成某些不需要的特征或内容。 infer_steps (int): 推理过程中执行的步数影响生成质量。默认值为 50。 guidance_scale (float): 控制生成的引导强度指导生成更加贴近输入提示。值越高生成结果越受 prompt 的影响。默认值为 6。 flow_shift (float): 调整视频生成过程中帧之间的运动流变化。默认值为 5.0。 embedded_guidance_scale (float): 可选的额外控制参数用于嵌入特定的引导强度默认为 None。 batch_size (int): 每次生成的批量大小。默认值为 1。 num_videos_per_prompt (int): 每个 prompt 生成的视频数量。默认值为 1。 kwargs: 其他额外参数可以在调用时传入以灵活配置生成过程。 torch.no_grad()def predict(self,prompt,height192,width336,video_length129,seedNone,negative_promptNone,infer_steps50,guidance_scale6,flow_shift5.0,embedded_guidance_scaleNone,batch_size1,num_videos_per_prompt1,**kwargs,):Predict the image/video from the given text.Args:prompt (str or List[str]): The input text.kwargs:height (int): The height of the output video. Default is 192.width (int): The width of the output video. Default is 336.video_length (int): The frame number of the output video. Default is 129.seed (int or List[str]): The random seed for the generation. Default is a random integer.negative_prompt (str or List[str]): The negative text prompt. Default is an empty string.guidance_scale (float): The guidance scale for the generation. Default is 6.0.num_images_per_prompt (int): The number of images per prompt. Default is 1.infer_steps (int): The number of inference steps. Default is 100.# 分布式环境检查if self.parallel_args[ulysses_degree] 1 or self.parallel_args[ring_degree] 1:assert seed is not None, \You have to set a seed in the distributed environment, please rerun with --seed your-seed.# 满足分布式环境的条件调用 parallelize_transformer 函数并行化模型parallelize_transformer(self.pipeline)# 初始化一个空字典 out_dict用于存储最终的生成结果。out_dict dict()# # 根据传入的 seed 参数生成一组随机种子并将这些种子用于初始化随机数生成器 (torch.Generator) 来控制生成过程的随机性。# # 根据 seed 参数的类型None、int、list、tuple 或 torch.Tensor执行不同的逻辑生成用于控制随机数生成器的 seeds 列表if isinstance(seed, torch.Tensor):seed seed.tolist()if seed is None:seeds [random.randint(0, 1_000_000)for _ in range(batch_size * num_videos_per_prompt)]elif isinstance(seed, int):seeds [seed ifor _ in range(batch_size)for i in range(num_videos_per_prompt)]elif isinstance(seed, (list, tuple)):if len(seed) batch_size:seeds [int(seed[i]) jfor i in range(batch_size)for j in range(num_videos_per_prompt)]elif len(seed) batch_size * num_videos_per_prompt:seeds [int(s) for s in seed]else:raise ValueError(fLength of seed must be equal to number of prompt(batch_size) or fbatch_size * num_videos_per_prompt ({batch_size} * {num_videos_per_prompt}), got {seed}.)else:raise ValueError(fSeed must be an integer, a list of integers, or None, got {seed}.)# 对每个种子在指定设备self.device上创建一个 PyTorch 的随机数生成器 torch.Generator并使用对应的种子进行手动初始化# manual_seed(seed)。将这些生成器存储在列表 generator 中。generator [torch.Generator(self.device).manual_seed(seed) for seed in seeds]# 将生成的 seeds 列表存储在 out_dict 中供后续使用可能用于复现生成结果或记录生成过程的随机性。out_dict[seeds] seeds# # 检查和调整视频生成的输入参数height、width 和 video_length的合法性与对齐要求并计算出调整后的目标尺寸。# # 检查输入的 height、width 和 video_length 是否为正整数if width 0 or height 0 or video_length 0:raise ValueError(fheight and width and video_length must be positive integers, got height{height}, width{width}, video_length{video_length})# 检查 video_length - 1 是否为 4 的倍数if (video_length - 1) % 4 ! 0:raise ValueError(fvideo_length-1 must be a multiple of 4, got {video_length})# 日志记录logger.info(fInput (height, width, video_length) ({height}, {width}, {video_length}))# 目标高度和宽度对齐到 16 的倍数target_height align_to(height, 16)target_width align_to(width, 16)target_video_length video_length# 存储目标尺寸out_dict[size] (target_height, target_width, target_video_length)# # 检查和处理文本生成任务中的 prompt 和 negative_prompt 参数# # 确保输入的 prompt 是字符串类型if not isinstance(prompt, str):raise TypeError(fprompt must be a string, but got {type(prompt)})prompt [prompt.strip()] # 对 prompt 去除首尾多余的空格使用 .strip()然后包装成一个单元素列表# 处理 negative_prompt 参数if negative_prompt is None or negative_prompt :negative_prompt self.default_negative_promptif not isinstance(negative_prompt, str):raise TypeError(fnegative_prompt must be a string, but got {type(negative_prompt)})negative_prompt [negative_prompt.strip()]# # 设置调度器 (Scheduler)# scheduler FlowMatchDiscreteScheduler( # 处理流Flow的调度shiftflow_shift, # 控制流动调度器的偏移量。flow_shift 通常与时序或流动模型相关例如调整时间步之间的关系。reverseself.args.flow_reverse, # 决定是否反向调度可能是在推理过程中逆序生成帧solverself.args.flow_solver # 指定用于调度的解算器类型solver例如选择数值方法来优化时间步间的计算。)self.pipeline.scheduler scheduler# # 构建旋转位置嵌入 (Rotary Positional Embedding)# # 根据目标视频长度、高度和宽度生成正弦 (freqs_sin) 和余弦 (freqs_cos) 的频率嵌入。freqs_cos, freqs_sin self.get_rotary_pos_embed(target_video_length, target_height, target_width)# 表示视频中总的编码标记数tokens通常等于时间步数帧数与空间分辨率像素数相乘。n_tokens freqs_cos.shape[0]# # 打印推理参数# debug_str fheight: {target_height}width: {target_width}video_length: {target_video_length}prompt: {prompt}neg_prompt: {negative_prompt}seed: {seed}infer_steps: {infer_steps}num_videos_per_prompt: {num_videos_per_prompt}guidance_scale: {guidance_scale}n_tokens: {n_tokens}flow_shift: {flow_shift}embedded_guidance_scale: {embedded_guidance_scale}logger.debug(debug_str)# # Pipeline inference# start_time time.time()samples self.pipeline(promptprompt, # 文本提示用于指导生成内容。heighttarget_height, # 生成图像或视频帧的分辨率。widthtarget_width, #video_lengthtarget_video_length, # 视频的帧数。如果 video_length 1表示生成视频否则生成单张图像。num_inference_stepsinfer_steps, # 推理步数决定生成过程的细粒度程度步数越多生成结果越精细。guidance_scaleguidance_scale, # 指导比例控制生成与 prompt 的一致性程度。negative_promptnegative_prompt, # 负面提示用于约束生成内容避免不期望的结果。num_videos_per_promptnum_videos_per_prompt, # 每条提示生成的视频数量。generatorgenerator, # 随机生成器对象用于控制生成过程中的随机性通常与随机种子结合。output_typepil, # 指定输出格式为 PIL.Image 对象便于后续处理freqs_cis(freqs_cos, freqs_sin), # 旋转位置嵌入 (RoPE) 的频率矩阵增强时空位置感知能力。n_tokensn_tokens, # 输入序列的总标记数用于指导生成过程。embedded_guidance_scaleembedded_guidance_scale, # 嵌入式指导比例用于进一步优化嵌入向量的生成。data_typevideo if target_video_length 1 else image, # 指定生成目标为视频或图像取决于帧数。is_progress_barTrue, # 显示推理进度条方便监控生成进度。vae_verself.args.vae, # 使用指定版本的 VAE变分自编码器决定生成内容的潜在空间。enable_tilingself.args.vae_tiling, # 启用 VAE 分块处理提高内存效率特别适用于高分辨率生成。)[0] # 返回生成的样本通常是一个 PIL.Image 或视频帧序列# 保存生成结果out_dict[samples] samplesout_dict[prompts] prompt# 计算并记录推理时间gen_time time.time() - start_timelogger.info(fSuccess, time: {gen_time})return out_dict推理管道hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py
__call__ 方法接受用户输入的提示、生成图像或视频的尺寸以及其他生成过程的参数完成推理并返回生成的图像或视频。 参数详解
1. 提示相关参数 prompt 用于引导生成的文本提示可以是字符串如 海边的晴天或字符串列表。 示例prompt在森林中散步的机器人 negative_prompt 用于指示生成时需要避免的内容负面提示。可以是字符串或字符串列表。 示例negative_prompt没有树木
2. 图像或视频生成的尺寸参数 height 和 width 指定生成的图像或视频帧的高度和宽度单位像素。 示例height512, width512 表示生成 512×512 的图像。 video_length 生成视频的帧数。如果设为 1则生成单张图像而不是视频。 data_type 指定生成的数据类型可以是 video 或 image。
3. 推理过程相关参数 num_inference_steps 指定去噪的迭代步数。步数越多生成质量越高但推理时间会更长。默认为 50。 timesteps 用户可以自定义时间步降噪步骤。如果未定义则根据 num_inference_steps 自动生成。 sigmas 自定义噪声尺度σ用于降噪调度器。如果未定义则使用默认值。 guidance_scale 引导比例默认为 7.5。较高的值会让生成结果更贴合 prompt 提示但可能会牺牲生成质量。 eta DDIM 调度器的噪声参数控制生成的多样性。默认值为 0.0。
4. 嵌入和生成设置 prompt_embeds 和 negative_prompt_embeds 预生成的文本提示和负面提示的嵌入用于直接控制生成过程。如果未提供嵌入系统会根据 prompt 和 negative_prompt 自动生成嵌入。 latents 预生成的噪声张量。如果未提供系统会随机生成初始噪声。 generator 用于设置随机数生成器以便实现确定性固定随机种子时生成结果一致。
5. 输出相关参数 output_type 指定生成输出的类型默认为 pilPillow 图像对象。可选 np.arrayNumPy 数组。 return_dict 是否以字典形式返回结果。如果为 True则返回类似 HunyuanVideoPipelineOutput 的对象否则返回元组。
6. 其他高级参数 cross_attention_kwargs 传递给交叉注意力模块的额外参数用于调整注意力机制。 callback_on_step_end 回调函数在每次去噪步骤结束时调用可用于自定义处理逻辑。 clip_skip 跳过 CLIP 模型的层数用于调整文本嵌入的生成层。 guidance_rescale 用于调整引导时的比例因子以修正某些生成过程中的过曝问题。
返回值
如果 return_dictTrue返回一个包含生成结果的字典对象如 HunyuanVideoPipelineOutput。如果 return_dictFalse返回一个元组第一个元素是生成的图像/视频列表第二个元素是布尔列表指示生成内容是否包含敏感内容。 torch.no_grad()replace_example_docstring(EXAMPLE_DOC_STRING)def __call__(self,prompt: Union[str, List[str]],height: int,width: int,video_length: int,data_type: str video,num_inference_steps: int 50,timesteps: List[int] None,sigmas: List[float] None,guidance_scale: float 7.5,negative_prompt: Optional[Union[str, List[str]]] None,num_videos_per_prompt: Optional[int] 1,eta: float 0.0,generator: Optional[Union[torch.Generator, List[torch.Generator]]] None,latents: Optional[torch.Tensor] None,prompt_embeds: Optional[torch.Tensor] None,attention_mask: Optional[torch.Tensor] None,negative_prompt_embeds: Optional[torch.Tensor] None,negative_attention_mask: Optional[torch.Tensor] None,output_type: Optional[str] pil,return_dict: bool True,cross_attention_kwargs: Optional[Dict[str, Any]] None,guidance_rescale: float 0.0,clip_skip: Optional[int] None,callback_on_step_end: Optional[Union[Callable[[int, int, Dict], None],PipelineCallback,MultiPipelineCallbacks,]] None,callback_on_step_end_tensor_inputs: List[str] [latents],freqs_cis: Tuple[torch.Tensor, torch.Tensor] None,vae_ver: str 88-4c-sd,enable_tiling: bool False,n_tokens: Optional[int] None,embedded_guidance_scale: Optional[float] None,**kwargs,):rThe call function to the pipeline for generation.Args:prompt (str or List[str]):The prompt or prompts to guide image generation. If not defined, you need to pass prompt_embeds.height (int):The height in pixels of the generated image.width (int):The width in pixels of the generated image.video_length (int):The number of frames in the generated video.num_inference_steps (int, *optional*, defaults to 50):The number of denoising steps. More denoising steps usually lead to a higher quality image at theexpense of slower inference.timesteps (List[int], *optional*):Custom timesteps to use for the denoising process with schedulers which support a timesteps argumentin their set_timesteps method. If not defined, the default behavior when num_inference_steps ispassed will be used. Must be in descending order.sigmas (List[float], *optional*):Custom sigmas to use for the denoising process with schedulers which support a sigmas argument intheir set_timesteps method. If not defined, the default behavior when num_inference_steps is passedwill be used.guidance_scale (float, *optional*, defaults to 7.5):A higher guidance scale value encourages the model to generate images closely linked to the textprompt at the expense of lower image quality. Guidance scale is enabled when guidance_scale 1.negative_prompt (str or List[str], *optional*):The prompt or prompts to guide what to not include in image generation. If not defined, you need topass negative_prompt_embeds instead. Ignored when not using guidance (guidance_scale 1).num_videos_per_prompt (int, *optional*, defaults to 1):The number of images to generate per prompt.eta (float, *optional*, defaults to 0.0):Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only appliesto the [~schedulers.DDIMScheduler], and is ignored in other schedulers.generator (torch.Generator or List[torch.Generator], *optional*):A [torch.Generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to makegeneration deterministic.latents (torch.Tensor, *optional*):Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for imagegeneration. Can be used to tweak the same generation with different prompts. If not provided, a latentstensor is generated by sampling using the supplied random generator.prompt_embeds (torch.Tensor, *optional*):Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If notprovided, text embeddings are generated from the prompt input argument.negative_prompt_embeds (torch.Tensor, *optional*):Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). Ifnot provided, negative_prompt_embeds are generated from the negative_prompt input argument.output_type (str, *optional*, defaults to pil):The output format of the generated image. Choose between PIL.Image or np.array.return_dict (bool, *optional*, defaults to True):Whether or not to return a [HunyuanVideoPipelineOutput] instead of aplain tuple.cross_attention_kwargs (dict, *optional*):A kwargs dictionary that if specified is passed along to the [AttentionProcessor] as defined in[self.processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).guidance_rescale (float, *optional*, defaults to 0.0):Guidance rescale factor from [Common Diffusion Noise Schedules and Sample Steps areFlawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure whenusing zero terminal SNR.clip_skip (int, *optional*):Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means thatthe output of the pre-final layer will be used for computing the prompt embeddings.callback_on_step_end (Callable, PipelineCallback, MultiPipelineCallbacks, *optional*):A function or a subclass of PipelineCallback or MultiPipelineCallbacks that is called at the end ofeach denoising step during the inference. with the following arguments: callback_on_step_end(self:DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict). callback_kwargs will include alist of all tensors as specified by callback_on_step_end_tensor_inputs.callback_on_step_end_tensor_inputs (List, *optional*):The list of tensor inputs for the callback_on_step_end function. The tensors specified in the listwill be passed as callback_kwargs argument. You will only be able to include variables listed in the._callback_tensor_inputs attribute of your pipeline class.Examples:Returns:[~HunyuanVideoPipelineOutput] or tuple:If return_dict is True, [HunyuanVideoPipelineOutput] is returned,otherwise a tuple is returned where the first element is a list with the generated images and thesecond element is a list of bools indicating whether the corresponding generated image containsnot-safe-for-work (nsfw) content.# 处理与回调函数相关的参数同时对已弃用的参数发出警告deprecation warnings。# 它还检查了新的回调函数机制 callback_on_step_end 是否符合预期类型。callback kwargs.pop(callback, None)callback_steps kwargs.pop(callback_steps, None)if callback is not None:deprecate(callback,1.0.0,Passing callback as an input argument to __call__ is deprecated, consider using callback_on_step_end,)if callback_steps is not None:deprecate(callback_steps,1.0.0,Passing callback_steps as an input argument to __call__ is deprecated, consider using callback_on_step_end,)if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):callback_on_step_end_tensor_inputs callback_on_step_end.tensor_inputs# 0. Default height and width to unet# height height or self.transformer.config.sample_size * self.vae_scale_factor# width width or self.transformer.config.sample_size * self.vae_scale_factor# to deal with lora scaling and other possible forward hooks# 1. 验证输入参数是否合法。self.check_inputs(prompt,height,width,video_length,callback_steps, # 回调频率指定在生成过程中每隔多少步执行一次回调。negative_prompt,prompt_embeds, # 预嵌入的提示词和反向提示词。如果已经对文本进行了嵌入处理可以直接传递这些值而不是原始文本。negative_prompt_embeds,callback_on_step_end_tensor_inputs, # 与回调机制相关的数据张量。vae_vervae_ver, # 可选参数可能指定生成内容时使用的 VAE变分自动编码器的版本。)# 控制生成内容的引导强度。一般用于调整模型对 prompt提示词的依赖程度。较大的值会让生成内容更接近提示词但可能导致丢失多样性。self._guidance_scale guidance_scale# 用于重新调整指导比例可能是对 guidance_scale 的一种动态调整。用于平衡模型在特定生成任务中的表现。self._guidance_rescale guidance_rescale# 控制是否在 CLIP 模型中跳过某些层的计算。在某些生成任务中跳过部分层可以改善生成质量。self._clip_skip clip_skip# 与交叉注意力Cross Attention相关的参数。可能包括对注意力权重的控制比如调整注意力机制如何在提示词和生成内容之间分配权重。self._cross_attention_kwargs cross_attention_kwargs# 标志可能在生成过程的某些阶段被动态修改。# 如果 _interrupt 被设置为 True生成过程可能会被中止。这种设计通常用于在用户希望终止长时间生成任务时使用。self._interrupt False# 2. 根据输入的提示词 prompt 或嵌入 prompt_embeds确定生成任务的批量大小batch_size。if prompt is not None and isinstance(prompt, str):batch_size 1 # 如果 prompt 是单个字符串说明只有一个提示词。批量大小设置为 1。elif prompt is not None and isinstance(prompt, list):# 如果 prompt 是一个列表说明有多个提示词。# 此时批量大小等于提示词的数量即 len(prompt)。batch_size len(prompt)else:#如果 prompt 是 None说明提示词未提供可能直接使用预先计算的嵌入 prompt_embeds。# 此时批量大小由 prompt_embeds 的第一维通常是样本数量决定。batch_size prompt_embeds.shape[0]# 确定设备的devicedevice torch.device(fcuda:{dist.get_rank()}) if dist.is_initialized() else self._execution_device# 3. Encode input prompt# 处理 LoRALow-Rank Adaptation缩放系数通过 cross_attention_kwargs 提取或设置缩放比例 lora_scalelora_scale (self.cross_attention_kwargs.get(scale, None)if self.cross_attention_kwargs is not Noneelse None)# 对提示词进行编码将文本提示词 prompt 和负向提示词 negative_prompt 编码为嵌入向量并生成对应的注意力掩码。(prompt_embeds, # 正向提示词的嵌入向量。negative_prompt_embeds, # 负向提示词的嵌入向量。prompt_mask, # 正向提示词的注意力掩码。negative_prompt_mask, # 负向提示词的注意力掩码。) self.encode_prompt(prompt,device,num_videos_per_prompt,self.do_classifier_free_guidance,negative_prompt,prompt_embedsprompt_embeds,attention_maskattention_mask,negative_prompt_embedsnegative_prompt_embeds,negative_attention_masknegative_attention_mask,lora_scalelora_scale,clip_skipself.clip_skip,data_typedata_type,)# 处理多文本编码器若存在额外的文本编码器 text_encoder_2使用该编码器再次处理提示词。if self.text_encoder_2 is not None:(prompt_embeds_2,negative_prompt_embeds_2,prompt_mask_2,negative_prompt_mask_2,) self.encode_prompt(prompt,device,num_videos_per_prompt,self.do_classifier_free_guidance,negative_prompt,prompt_embedsNone,attention_maskNone,negative_prompt_embedsNone,negative_attention_maskNone,lora_scalelora_scale,clip_skipself.clip_skip,text_encoderself.text_encoder_2,data_typedata_type,)else:prompt_embeds_2 Nonenegative_prompt_embeds_2 Noneprompt_mask_2 Nonenegative_prompt_mask_2 None# 处理自由分类指导Classifier-Free Guidance为实现该技术合并正向和负向提示词嵌入避免多次前向传递。if self.do_classifier_free_guidance:# 功能如果启用了自由分类指导则将正向和负向提示词的嵌入和掩码合并为一个批次。# 原因自由分类指导需要两次前向传递一次处理负向提示词指导无条件生成一次处理正向提示词指导条件生成。# 为了提高效率将两组嵌入拼接在一起作为一个批次传递给模型避免两次单独的前向传递。prompt_embeds torch.cat([negative_prompt_embeds, prompt_embeds])if prompt_mask is not None:prompt_mask torch.cat([negative_prompt_mask, prompt_mask])if prompt_embeds_2 is not None:prompt_embeds_2 torch.cat([negative_prompt_embeds_2, prompt_embeds_2])if prompt_mask_2 is not None:prompt_mask_2 torch.cat([negative_prompt_mask_2, prompt_mask_2])# 4. Prepare timesteps# 准备调度器的额外参数extra_set_timesteps_kwargs self.prepare_extra_func_kwargs(self.scheduler.set_timesteps, {n_tokens: n_tokens})# 获取推理过程中需要用到的时间步 (timesteps) 和推理步数 (num_inference_steps)。timesteps, num_inference_steps retrieve_timesteps(self.scheduler,num_inference_steps,device,timesteps,sigmas,**extra_set_timesteps_kwargs,)# 根据 vae_ver 调整视频长度if 884 in vae_ver:video_length (video_length - 1) // 4 1elif 888 in vae_ver:video_length (video_length - 1) // 8 1else:video_length video_length# 5. Prepare latent variablesnum_channels_latents self.transformer.config.in_channelslatents self.prepare_latents(batch_size * num_videos_per_prompt,num_channels_latents,height,width,video_length,prompt_embeds.dtype,device,generator,latents,)# 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipelineextra_step_kwargs self.prepare_extra_func_kwargs(self.scheduler.step, # 扩散模型的调度器中的 step 方法负责更新噪声预测结果。{generator: generator, eta: eta}, # 一个字典包含生成器 generator 和步长相关参数 eta。)# 确定目标数据类型及自动混合精度的设置target_dtype PRECISION_TO_TYPE[self.args.precision]autocast_enabled (target_dtype ! torch.float32) and not self.args.disable_autocast# 确定 VAE 的数据类型及自动混合精度设置vae_dtype PRECISION_TO_TYPE[self.args.vae_precision]vae_autocast_enabled (vae_dtype ! torch.float32) and not self.args.disable_autocast# 7. 初始化去噪循环的预处理步骤# timesteps调度器生成的时间步序列。# num_inference_steps推理过程中真正的去噪步数。# self.scheduler.order调度器的阶数通常与预测算法的高阶插值相关。num_warmup_steps len(timesteps) - num_inference_steps * self.scheduler.orderself._num_timesteps len(timesteps)# if is_progress_bar:# progress_bar 用于显示推理过程的进度num_inference_steps 是总推理步数。with self.progress_bar(totalnum_inference_steps) as progress_bar:for i, t in enumerate(timesteps):if self.interrupt:continue# 如果启用了 分类器自由指导do_classifier_free_guidance则将 latents 复制两份用于同时计算 条件预测 和 无条件预测。# 否则仅使用原始 latents。latent_model_input (torch.cat([latents] * 2)if self.do_classifier_free_guidanceelse latents)# 调用 scheduler 的 scale_model_input 方法对 latent_model_input 在当前时间步 t 上进行预处理。# 这个缩放操作可能根据调度器的实现涉及到归一化或其他调整。latent_model_input self.scheduler.scale_model_input(latent_model_input, t)# t_expand 将时间步 t 扩展到与 latent_model_input 的批量维度一致。# 如果 embedded_guidance_scale 存在则创建扩展的指导参数 guidance_expand用于对模型预测进行额外控制。t_expand t.repeat(latent_model_input.shape[0])guidance_expand (torch.tensor([embedded_guidance_scale] * latent_model_input.shape[0],dtypetorch.float32,devicedevice,).to(target_dtype)* 1000.0if embedded_guidance_scale is not Noneelse None)# 使用 Transformer 模型预测噪声残差with torch.autocast(device_typecuda, dtypetarget_dtype, enabledautocast_enabled):noise_pred self.transformer( # For an input image (129, 192, 336) (1, 256, 256)latent_model_input, # 当前的潜变量输入 [2, 16, 33, 24, 42]t_expand, # 时间步信息 [2]text_statesprompt_embeds, # 与文本提示相关的嵌入向量 [2, 256, 4096]text_maskprompt_mask, # [2, 256]text_states_2prompt_embeds_2, # [2, 768]freqs_cosfreqs_cis[0], # 频率信息用于特定的时间步缩放 [seqlen, head_dim]freqs_sinfreqs_cis[1], # [seqlen, head_dim]guidanceguidance_expand, # 指导参数用于条件生成return_dictTrue,)[x]# 分类器自由指导的噪声调整if self.do_classifier_free_guidance:noise_pred_uncond, noise_pred_text noise_pred.chunk(2) # 无条件预测的噪声条件预测的噪声基于文本提示noise_pred noise_pred_uncond self.guidance_scale * (noise_pred_text - noise_pred_uncond)# 噪声重缩放if self.do_classifier_free_guidance and self.guidance_rescale 0.0:# Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdfnoise_pred rescale_noise_cfg(noise_pred,noise_pred_text,guidance_rescaleself.guidance_rescale,)# 使用调度器更新潜变量# compute the previous noisy sample x_t - x_t-1latents self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dictFalse)[0]# callback_on_step_end 函数则在每步结束时调用用于自定义操作如日志记录、结果保存。# 更新潜变量和提示嵌入向量。if callback_on_step_end is not None:callback_kwargs {}for k in callback_on_step_end_tensor_inputs:callback_kwargs[k] locals()[k]callback_outputs callback_on_step_end(self, i, t, callback_kwargs)latents callback_outputs.pop(latents, latents)prompt_embeds callback_outputs.pop(prompt_embeds, prompt_embeds)negative_prompt_embeds callback_outputs.pop(negative_prompt_embeds, negative_prompt_embeds)# 进度条更新与其他回调if i len(timesteps) - 1 or ((i 1) num_warmup_steps and (i 1) % self.scheduler.order 0):if progress_bar is not None:progress_bar.update()if callback is not None and i % callback_steps 0:step_idx i // getattr(self.scheduler, order, 1)callback(step_idx, t, latents)# 从潜变量latent space解码生成图像if not output_type latent:# 潜变量维度的扩展检查expand_temporal_dim False# 如果形状为 4D ([batch_size, channels, height, width])# 如果 VAE 是 3D 自回归模型AutoencoderKLCausal3D则对潜变量增加一个时间维度 (unsqueeze(2))。# 设置 expand_temporal_dimTrue标记后续需要移除该额外维度。if len(latents.shape) 4:if isinstance(self.vae, AutoencoderKLCausal3D):latents latents.unsqueeze(2)expand_temporal_dim True# 如果形状为 5D ([batch_size, channels, frames, height, width])则不需要操作。elif len(latents.shape) 5:passelse:raise ValueError(fOnly support latents with shape (b, c, h, w) or (b, c, f, h, w), but got {latents.shape}.)# 潜变量的缩放与偏移# 检查 VAE 配置中是否定义了 shift_factor偏移因子if (hasattr(self.vae.config, shift_factor)and self.vae.config.shift_factor): # 如果存在则对潜变量执行缩放和偏移操作latents (latents / self.vae.config.scaling_factor self.vae.config.shift_factor)else: # 如果 shift_factor 不存在仅进行缩放操作latents latents / self.vae.config.scaling_factorwith torch.autocast(device_typecuda, dtypevae_dtype, enabledvae_autocast_enabled):if enable_tiling:# 调用 VAE 的 enable_tiling() 方法可能用于解码较大的图像块。self.vae.enable_tiling()# 使用 VAE变分自编码器的 decode 方法将潜变量解码为图像。image self.vae.decode(latents, return_dictFalse, generatorgenerator)[0]else:image self.vae.decode(latents, return_dictFalse, generatorgenerator)[0]# 如果添加了时间维度expand_temporal_dimTrue或者解码出的图像在时间维度上只有一个帧则移除时间维度。if expand_temporal_dim or image.shape[2] 1:image image.squeeze(2)else:image latents# 图像归一化image (image / 2 0.5).clamp(0, 1)# 将图像移动到 CPU并转换为 float32 类型。这是为了确保图像兼容性无论之前是否使用了混合精度image image.cpu().float()# 调用 maybe_free_model_hooks() 方法可能会释放模型占用的内存资源尤其是在内存有限的 GPU 上有用。self.maybe_free_model_hooks()# 如果不需要返回字典return_dictFalse则直接返回处理后的图像if not return_dict:return imagereturn HunyuanVideoPipelineOutput(videosimage)
模型结构hyvideo/modules/models.py
双流块 MMDoubleStreamBlock
class MMDoubleStreamBlock(nn.Module):A multimodal dit block with seperate modulation fortext and image/video, see more details (SD3): https://arxiv.org/abs/2403.03206(Flux.1): https://github.com/black-forest-labs/fluxdef __init__(self,hidden_size: int, # 模型隐藏层维度。heads_num: int, # 多头注意力的头数。mlp_width_ratio: float, # MLP 中隐藏层宽度与 hidden_size 的比率。mlp_act_type: str gelu_tanh, # 激活函数的类型默认 gelu_tanhqk_norm: bool True, # 是否对 Query 和 Key 启用归一化。qk_norm_type: str rms, # Query 和 Key 归一化的方法默认 rms。qkv_bias: bool False, # QKV 投影中是否启用偏置项。dtype: Optional[torch.dtype] None, # 张量的数据类型和设备。device: Optional[torch.device] None,):factory_kwargs {device: device, dtype: dtype}super().__init__()self.deterministic Falseself.heads_num heads_numhead_dim hidden_size // heads_nummlp_hidden_dim int(hidden_size * mlp_width_ratio)### 图像模态# 模态调制模块使用 ModulateDiT为图像和文本生成 6 组参数shift、scale、gate。self.img_mod ModulateDiT(hidden_size,factor6,act_layerget_activation_layer(silu),**factory_kwargs,)# 归一化self.img_norm1 nn.LayerNorm(hidden_size, elementwise_affineFalse, eps1e-6, **factory_kwargs)# QKV 投影层通过全连接层计算 Query、Key 和 Valueself.img_attn_qkv nn.Linear(hidden_size, hidden_size * 3, biasqkv_bias, **factory_kwargs)# 归一化模块qk_norm_layer get_norm_layer(qk_norm_type)self.img_attn_q_norm (qk_norm_layer(head_dim, elementwise_affineTrue, eps1e-6, **factory_kwargs)if qk_normelse nn.Identity())self.img_attn_k_norm (qk_norm_layer(head_dim, elementwise_affineTrue, eps1e-6, **factory_kwargs)if qk_normelse nn.Identity())self.img_attn_proj nn.Linear(hidden_size, hidden_size, biasqkv_bias, **factory_kwargs)self.img_norm2 nn.LayerNorm(hidden_size, elementwise_affineFalse, eps1e-6, **factory_kwargs)self.img_mlp MLP(hidden_size,mlp_hidden_dim,act_layerget_activation_layer(mlp_act_type),biasTrue,**factory_kwargs,)### 文本模态self.txt_mod ModulateDiT(hidden_size,factor6,act_layerget_activation_layer(silu),**factory_kwargs,)self.txt_norm1 nn.LayerNorm(hidden_size, elementwise_affineFalse, eps1e-6, **factory_kwargs)self.txt_attn_qkv nn.Linear(hidden_size, hidden_size * 3, biasqkv_bias, **factory_kwargs)self.txt_attn_q_norm (qk_norm_layer(head_dim, elementwise_affineTrue, eps1e-6, **factory_kwargs)if qk_normelse nn.Identity())self.txt_attn_k_norm (qk_norm_layer(head_dim, elementwise_affineTrue, eps1e-6, **factory_kwargs)if qk_normelse nn.Identity())self.txt_attn_proj nn.Linear(hidden_size, hidden_size, biasqkv_bias, **factory_kwargs)self.txt_norm2 nn.LayerNorm(hidden_size, elementwise_affineFalse, eps1e-6, **factory_kwargs)self.txt_mlp MLP(hidden_size,mlp_hidden_dim,act_layerget_activation_layer(mlp_act_type),biasTrue,**factory_kwargs,)self.hybrid_seq_parallel_attn Nonedef enable_deterministic(self):self.deterministic Truedef disable_deterministic(self):self.deterministic Falsedef forward(self,img: torch.Tensor, # 图像张量 (B, L_img, hidden_size)txt: torch.Tensor, # 文本张量 (B, L_txt, hidden_size)vec: torch.Tensor, # 特征向量用于调制cu_seqlens_q: Optional[torch.Tensor] None, # Query 的累积序列长度cu_seqlens_kv: Optional[torch.Tensor] None, # Key/Value 的累积序列长度max_seqlen_q: Optional[int] None, # Query 最大序列长度max_seqlen_kv: Optional[int] None, # Key/Value 最大序列长度freqs_cis: tuple None, # 可选的旋转位置编码参数) - Tuple[torch.Tensor, torch.Tensor]:# vec 特征向量通过 ModulateDiT 模块分别为图像和文本模态生成 6 组调制参数(img_mod1_shift,img_mod1_scale,img_mod1_gate,img_mod2_shift,img_mod2_scale,img_mod2_gate,) self.img_mod(vec).chunk(6, dim-1)(txt_mod1_shift,txt_mod1_scale,txt_mod1_gate,txt_mod2_shift,txt_mod2_scale,txt_mod2_gate,) self.txt_mod(vec).chunk(6, dim-1)图像模态的前向处理# Layernorm 归一化img_modulated self.img_norm1(img)# 调制函数 modulate 进行标准化和缩放img_modulated modulate(img_modulated, shiftimg_mod1_shift, scaleimg_mod1_scale)# 得到 Query、Key 和 Valueimg_qkv self.img_attn_qkv(img_modulated)img_q, img_k, img_v rearrange(img_qkv, B L (K H D) - K B L H D, K3, Hself.heads_num)# 对 Query 和 Key 进行归一化。img_q self.img_attn_q_norm(img_q).to(img_v)img_k self.img_attn_k_norm(img_k).to(img_v)# 对 Query 和 Key 应用旋转位置编码。if freqs_cis is not None:img_qq, img_kk apply_rotary_emb(img_q, img_k, freqs_cis, head_firstFalse)assert (img_qq.shape img_q.shape and img_kk.shape img_k.shape), fimg_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}img_q, img_k img_qq, img_kk文本模态的前向处理txt_modulated self.txt_norm1(txt)txt_modulated modulate(txt_modulated, shifttxt_mod1_shift, scaletxt_mod1_scale)txt_qkv self.txt_attn_qkv(txt_modulated)txt_q, txt_k, txt_v rearrange(txt_qkv, B L (K H D) - K B L H D, K3, Hself.heads_num)# Apply QK-Norm if needed.txt_q self.txt_attn_q_norm(txt_q).to(txt_v)txt_k self.txt_attn_k_norm(txt_k).to(txt_v)# 将图像和文本的 Query、Key、Value 拼接q torch.cat((img_q, txt_q), dim1)k torch.cat((img_k, txt_k), dim1)v torch.cat((img_v, txt_v), dim1)assert (cu_seqlens_q.shape[0] 2 * img.shape[0] 1), fcu_seqlens_q.shape:{cu_seqlens_q.shape}, img.shape[0]:{img.shape[0]}# 多模态融合注意力计算if not self.hybrid_seq_parallel_attn:attn attention(q,k,v,cu_seqlens_qcu_seqlens_q,cu_seqlens_kvcu_seqlens_kv,max_seqlen_qmax_seqlen_q,max_seqlen_kvmax_seqlen_kv,batch_sizeimg_k.shape[0],)else:attn parallel_attention(self.hybrid_seq_parallel_attn,q,k,v,img_q_lenimg_q.shape[1],img_kv_lenimg_k.shape[1],cu_seqlens_qcu_seqlens_q,cu_seqlens_kvcu_seqlens_kv)# 最终将注意力结果拆分为图像部分 img_attn 和文本部分 txt_attnimg_attn, txt_attn attn[:, : img.shape[1]], attn[:, img.shape[1] :]图像模态的更新# 将注意力结果通过残差连接更新图像特征并通过 MLP 进一步增强img img apply_gate(self.img_attn_proj(img_attn), gateimg_mod1_gate)img img apply_gate(self.img_mlp(modulate(self.img_norm2(img), shiftimg_mod2_shift, scaleimg_mod2_scale)),gateimg_mod2_gate,)文本模态的更新txt txt apply_gate(self.txt_attn_proj(txt_attn), gatetxt_mod1_gate)txt txt apply_gate(self.txt_mlp(modulate(self.txt_norm2(txt), shifttxt_mod2_shift, scaletxt_mod2_scale)),gatetxt_mod2_gate,)# 返回更新后的图像特征和文本特征return img, txt
单流块 MMSingleStreamBlock
class MMSingleStreamBlock(nn.Module):A DiT block with parallel linear layers as described inhttps://arxiv.org/abs/2302.05442 and adapted modulation interface.Also refer to (SD3): https://arxiv.org/abs/2403.03206(Flux.1): https://github.com/black-forest-labs/fluxdef __init__(self,hidden_size: int, # 隐藏层的维度大小用于表示特征的维度。heads_num: int, # 注意力头的数量。mlp_width_ratio: float 4.0, # 用于确定多层感知机 (MLP) 的隐藏层宽度比例默认值为 4.0mlp_act_type: str gelu_tanh, # 激活函数类型qk_norm: bool True, # 决定是否对 Query 和 Key 应用归一化qk_norm_type: str rms, # 指定 Query 和 Key 的归一化方式例如 rms均方根归一化qk_scale: float None, # 自定义缩放因子用于注意力分数计算中的缩放dtype: Optional[torch.dtype] None, # 控制数据类型device: Optional[torch.device] None, # 控制缩放因子):factory_kwargs {device: device, dtype: dtype}super().__init__()self.deterministic Falseself.hidden_size hidden_sizeself.heads_num heads_numhead_dim hidden_size // heads_nummlp_hidden_dim int(hidden_size * mlp_width_ratio)self.mlp_hidden_dim mlp_hidden_dimself.scale qk_scale or head_dim ** -0.5# qkv and mlp_inself.linear1 nn.Linear(hidden_size, hidden_size * 3 mlp_hidden_dim, **factory_kwargs)# proj and mlp_outself.linear2 nn.Linear(hidden_size mlp_hidden_dim, hidden_size, **factory_kwargs)qk_norm_layer get_norm_layer(qk_norm_type)self.q_norm (qk_norm_layer(head_dim, elementwise_affineTrue, eps1e-6, **factory_kwargs)if qk_normelse nn.Identity())self.k_norm (qk_norm_layer(head_dim, elementwise_affineTrue, eps1e-6, **factory_kwargs)if qk_normelse nn.Identity())self.pre_norm nn.LayerNorm(hidden_size, elementwise_affineFalse, eps1e-6, **factory_kwargs)self.mlp_act get_activation_layer(mlp_act_type)()self.modulation ModulateDiT(hidden_size,factor3,act_layerget_activation_layer(silu),**factory_kwargs,)self.hybrid_seq_parallel_attn Nonedef enable_deterministic(self):self.deterministic Truedef disable_deterministic(self):self.deterministic Falsedef forward(self,x: torch.Tensor, # x: 输入特征张量形状为 (batch_size, seq_len, hidden_size)vec: torch.Tensor, # vec: 辅助特征向量通常来自调制器txt_len: int, # txt_len: 文本序列长度用于区分图像和文本部分。cu_seqlens_q: Optional[torch.Tensor] None, # 累积序列长度用于高效的分段注意力计算。cu_seqlens_kv: Optional[torch.Tensor] None,max_seqlen_q: Optional[int] None,# Query 和 Key/Value 的最大序列长度。max_seqlen_kv: Optional[int] None,freqs_cis: Tuple[torch.Tensor, torch.Tensor] None, # 可选的旋转位置编码RoPE) - torch.Tensor:# 调用 modulation 获取调制参数 mod_shift、mod_scale 和 mod_gate。mod_shift, mod_scale, mod_gate self.modulation(vec).chunk(3, dim-1)# 对输入 x 应用 LayerNorm并进行调制即元素级缩放和偏移x_mod modulate(self.pre_norm(x), shiftmod_shift, scalemod_scale)# 将 x_mod 映射到 qkv 和 mlp 两个部分。qkv, mlp torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim-1)# qkv 被分为 Query (q)、Key (k)、Value (v) 三个张量形状为 (batch_size, seq_len, heads_num, head_dim)。q, k, v rearrange(qkv, B L (K H D) - K B L H D, K3, Hself.heads_num)# 对 Query 和 Key 应用归一化。q self.q_norm(q).to(v)k self.k_norm(k).to(v)# 旋转位置编码 (RoPE)if freqs_cis is not None:img_q, txt_q q[:, :-txt_len, :, :], q[:, -txt_len:, :, :]img_k, txt_k k[:, :-txt_len, :, :], k[:, -txt_len:, :, :]# 分别对图像和文本部分应用旋转位置编码img_qq, img_kk apply_rotary_emb(img_q, img_k, freqs_cis, head_firstFalse)assert (img_qq.shape img_q.shape and img_kk.shape img_k.shape), fimg_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}img_q, img_k img_qq, img_kk# 图像部分和文本部分的 Query/Key 在编码后重新拼接。q torch.cat((img_q, txt_q), dim1)k torch.cat((img_k, txt_k), dim1)# Compute attention.assert (cu_seqlens_q.shape[0] 2 * x.shape[0] 1), fcu_seqlens_q.shape:{cu_seqlens_q.shape}, x.shape[0]:{x.shape[0]}# attention computation startif not self.hybrid_seq_parallel_attn:# 如果没有启用并行注意力机制调用标准注意力函数 attentionattn attention(q,k,v,cu_seqlens_qcu_seqlens_q,cu_seqlens_kvcu_seqlens_kv,max_seqlen_qmax_seqlen_q,max_seqlen_kvmax_seqlen_kv,batch_sizex.shape[0],)else:# 否则使用并行注意力机制 parallel_attentionattn parallel_attention(self.hybrid_seq_parallel_attn,q,k,v,img_q_lenimg_q.shape[1],img_kv_lenimg_k.shape[1],cu_seqlens_qcu_seqlens_q,cu_seqlens_kvcu_seqlens_kv)# attention computation end# 将注意力结果和 MLP 激活结果拼接通过线性层投影回输入维度。output self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))# 使用 mod_gate 进行门控融合将残差连接后的结果返回。return x apply_gate(output, gatemod_gate)
混元主干网络HYVideoDiffusionTransformer
class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):HunyuanVideo Transformer backbone该类继承了 ModelMixin 和 ConfigMixin使其与 diffusers 库的采样器例如 StableDiffusionPipeline兼容ModelMixin: 来自 diffusers 的模块提供了模型的保存和加载功能。ConfigMixin: 使模型能够以字典形式保存和加载配置信息。Reference:[1] Flux.1: https://github.com/black-forest-labs/flux[2] MMDiT: http://arxiv.org/abs/2403.03206Parameters----------args: argparse.Namespace传入的命令行参数用于设置模型的配置。patch_size: list输入特征的分块尺寸。一般用于图像或视频的分块操作。in_channels: int输入数据的通道数如 RGB 图像为 3 通道。out_channels: int模型输出的通道数。hidden_size: intTransformer 模块中隐藏层的维度。heads_num: int多头注意力机制中的注意力头数量通常用来分配不同的注意力特征。mlp_width_ratio: floatMLP多层感知机中隐藏层维度相对于 hidden_size 的比例。mlp_act_type: strMLP 使用的激活函数类型例如 ReLU、GELU 等。depth_double_blocks: int双 Transformer 块的数量。双块可能是指包含多层结构的单元。depth_single_blocks: int单 Transformer 块的数量。rope_dim_list: list为时空维度t, h, w设计的旋转位置编码ROPE的维度。qkv_bias: bool是否在 QKV查询、键、值线性层中使用偏置项。qk_norm: bool是否对 Q 和 K 应用归一化。qk_norm_type: strQK 归一化的类型。guidance_embed: bool是否使用指导嵌入guidance embedding来支持蒸馏训练。text_projection: str文本投影类型默认为 single_refiner可能用于文本引导的视频生成。use_attention_mask: bool是否在文本编码器中使用注意力掩码。dtype: torch.dtype模型参数的数据类型例如 torch.float32 或 torch.float16。device: torch.device模型的运行设备如 CPU 或 GPU。register_to_configdef __init__(self,args: Any, #patch_size: list [1, 2, 2],in_channels: int 4, # Should be VAE.config.latent_channels.out_channels: int None,hidden_size: int 3072,heads_num: int 24,mlp_width_ratio: float 4.0,mlp_act_type: str gelu_tanh,mm_double_blocks_depth: int 20,mm_single_blocks_depth: int 40,rope_dim_list: List[int] [16, 56, 56],qkv_bias: bool True,qk_norm: bool True,qk_norm_type: str rms,guidance_embed: bool False, # For modulation.text_projection: str single_refiner,use_attention_mask: bool True,dtype: Optional[torch.dtype] None,device: Optional[torch.device] None,):# 用来传递设备和数据类型如torch.float32的参数方便后续模块的初始化。factory_kwargs {device: device, dtype: dtype}super().__init__()self.patch_size patch_sizeself.in_channels in_channelsself.out_channels in_channels if out_channels is None else out_channelsself.unpatchify_channels self.out_channels # 用来重新拼接patch时的通道数。self.guidance_embed guidance_embedself.rope_dim_list rope_dim_list# Text projection. Default to linear projection.# Alternative: TokenRefiner. See more details (LI-DiT): http://arxiv.org/abs/2406.11831self.use_attention_mask use_attention_maskself.text_projection text_projectionself.text_states_dim args.text_states_dimself.text_states_dim_2 args.text_states_dim_2# 确保每个头的维度是整数。if hidden_size % heads_num ! 0:raise ValueError(fHidden size {hidden_size} must be divisible by heads_num {heads_num})pe_dim hidden_size // heads_num# 确保位置嵌入的维度与Transformer头的维度一致。if sum(rope_dim_list) ! pe_dim:raise ValueError(fGot {rope_dim_list} but expected positional dim {pe_dim})self.hidden_size hidden_sizeself.heads_num heads_num# 将输入图像分割为小块patch并映射到Transformer的隐藏空间hidden_size。# 每个patch相当于一个Transformer的输入token。self.img_in PatchEmbed(self.patch_size, self.in_channels, self.hidden_size, **factory_kwargs)# 根据text_projection参数选择不同的文本投影方式# TextProjection线性投影直接将文本特征映射到模型隐藏空间。if self.text_projection linear:self.txt_in TextProjection(self.text_states_dim,self.hidden_size,get_activation_layer(silu),**factory_kwargs,)# SingleTokenRefiner使用小型Transformer深度为2对文本特征进行细化处理。elif self.text_projection single_refiner:self.txt_in SingleTokenRefiner(self.text_states_dim, hidden_size, heads_num, depth2, **factory_kwargs)else:raise NotImplementedError(fUnsupported text_projection: {self.text_projection})# TimestepEmbedder时间步嵌入模块输入时间信息例如视频帧的索引并嵌入到Transformer隐藏空间。self.time_in TimestepEmbedder(self.hidden_size, get_activation_layer(silu), **factory_kwargs)# text modulation# MLPEmbedder用于处理来自文本或其他辅助信息的特征并投影到隐藏空间。self.vector_in MLPEmbedder(self.text_states_dim_2, self.hidden_size, **factory_kwargs)# guidance_in引导嵌入模块用于处理额外的控制信号如扩散模型中的引导提示。self.guidance_in (TimestepEmbedder(self.hidden_size, get_activation_layer(silu), **factory_kwargs)if guidance_embedelse None)# MMDoubleStreamBlock多模态双流块融合了图像流和文本流信息。self.double_blocks nn.ModuleList([MMDoubleStreamBlock(self.hidden_size,self.heads_num,mlp_width_ratiomlp_width_ratio,mlp_act_typemlp_act_type,qk_normqk_norm,qk_norm_typeqk_norm_type,qkv_biasqkv_bias,**factory_kwargs,)for _ in range(mm_double_blocks_depth)])# MMSingleStreamBlock单流块用于进一步处理多模态融合后的单一流特征。self.single_blocks nn.ModuleList([MMSingleStreamBlock(self.hidden_size,self.heads_num,mlp_width_ratiomlp_width_ratio,mlp_act_typemlp_act_type,qk_normqk_norm,qk_norm_typeqk_norm_type,**factory_kwargs,)for _ in range(mm_single_blocks_depth)])# FinalLayer将Transformer隐藏空间中的token重新解码为图像patch并还原到完整图像的分辨率。self.final_layer FinalLayer(self.hidden_size,self.patch_size,self.out_channels,get_activation_layer(silu),**factory_kwargs,)# 分别在模型中的 双流模块double_blocks 和 单流模块single_blocks 中启用或禁用确定性行为。def enable_deterministic(self):# 在深度学习中启用确定性行为意味着模型在同样的输入和参数初始化条件下无论多少次运行都能产生相同的输出结果。for block in self.double_blocks:block.enable_deterministic()for block in self.single_blocks:block.enable_deterministic()def disable_deterministic(self):# 禁用确定性行为可能会允许使用非确定性的操作如某些高效的并行实现从而提升计算效率。for block in self.double_blocks:block.disable_deterministic()for block in self.single_blocks:block.disable_deterministic()def forward(self,x: torch.Tensor, # 输入图像张量形状为 (N, C, T, H, W)。批量大小为 N通道数为 C时间步为 T高度和宽度为 H 和 W。t: torch.Tensor, # 时间步张量用于时间嵌入。范围应为 [0, 1000]可能对应扩散模型或时间相关的特征。text_states: torch.Tensor None, # 文本嵌入表示与图像配对的文本特征。text_mask: torch.Tensor None, # 文本掩码张量可选。当前未使用可能用于控制哪些文本特征参与计算。text_states_2: Optional[torch.Tensor] None, # 额外的文本嵌入用于进一步调制modulation。在模型中可能是辅助的文本特征表示freqs_cos: Optional[torch.Tensor] None, # 正弦和余弦频率用于位置编码或调制。freqs_sin: Optional[torch.Tensor] None,guidance: torch.Tensor None, # 引导调制强度形状可能是 cfg_scale x 1000。通常用于引导生成如扩散模型的分类引导。return_dict: bool True, # 是否返回一个字典结果。默认为 True。) - Union[torch.Tensor, Dict[str, torch.Tensor]]:out {}img xtxt text_states_, _, ot, oh, ow x.shape# 得到划分patch后的t,h,wtt, th, tw (ot // self.patch_size[0],oh // self.patch_size[1],ow // self.patch_size[2],)# Prepare modulation vectors.# 时间嵌入通过 self.time_in(t) 提取特征。vec self.time_in(t)# text modulation# 如果有额外文本嵌入 text_states_2则通过 self.vector_in 模块对 vec 进行调制。vec vec self.vector_in(text_states_2)# 启用了引导调制self.guidance_embed通过 self.guidance_in 引入引导特征。if self.guidance_embed:if guidance is None:raise ValueError(Didnt get guidance strength for guidance distilled model.)# our timestep_embedding is merged into guidance_in(TimestepEmbedder)vec vec self.guidance_in(guidance)# Embed image and text.# 图像嵌入img self.img_in(img)# 文本嵌入if self.text_projection linear: # 线性投影txt self.txt_in(txt)elif self.text_projection single_refiner: # 结合时间步 t 和文本掩码进行更复杂的处理。txt self.txt_in(txt, t, text_mask if self.use_attention_mask else None)else:raise NotImplementedError(fUnsupported text_projection: {self.text_projection})txt_seq_len txt.shape[1]img_seq_len img.shape[1]# 计算序列长度和累积序列索引# 用于 Flash Attention 的高效计算cu_seqlens_* 和 max_seqlen_* 控制序列长度和最大长度。# Compute cu_squlens and max_seqlen for flash attentioncu_seqlens_q get_cu_seqlens(text_mask, img_seq_len)cu_seqlens_kv cu_seqlens_qmax_seqlen_q img_seq_len txt_seq_lenmax_seqlen_kv max_seqlen_qfreqs_cis (freqs_cos, freqs_sin) if freqs_cos is not None else None# --------------------- Pass through DiT blocks ------------------------for _, block in enumerate(self.double_blocks):double_block_args [img,txt,vec,cu_seqlens_q,cu_seqlens_kv,max_seqlen_q,max_seqlen_kv,freqs_cis,]# 并行处理图像和文本信息使用输入参数包括嵌入和序列长度等逐步更新 img 和 txt。img, txt block(*double_block_args)# 合并图像和文本并通过单流模块x torch.cat((img, txt), 1)if len(self.single_blocks) 0:for _, block in enumerate(self.single_blocks):single_block_args [x,vec,txt_seq_len,cu_seqlens_q,cu_seqlens_kv,max_seqlen_q,max_seqlen_kv,(freqs_cos, freqs_sin),]x block(*single_block_args)# 分离图像特征img x[:, :img_seq_len, ...]# ---------------------------- Final layer ------------------------------# 图像特征通过 final_layer 提取最终结果img self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels)# 通过 unpatchify 恢复到原始分辨率。img self.unpatchify(img, tt, th, tw)if return_dict:out[x] imgreturn outreturn imgdef unpatchify(self, x, t, h, w):# 是将被切分为小块patches的特征重新还原成原始的张量形状通常用于图像处理任务中# 例如在 ViTVision Transformer模型的输出阶段将 patch 还原为完整图像的形式。x: (N, T, patch_size**2 * C) 批量大小时间帧数每个patch中的通道数imgs: (N, H, W, C)c self.unpatchify_channelspt, ph, pw self.patch_sizeassert t * h * w x.shape[1]x x.reshape(shape(x.shape[0], t, h, w, c, pt, ph, pw))x torch.einsum(nthwcopq-nctohpwq, x)imgs x.reshape(shape(x.shape[0], c, t * pt, h * ph, w * pw))return imgsdef params_count(self):# 计算模型的参数数量并将其按类别进行统计。它返回一个包含不同类别参数数量的字典通常用于分析模型的规模或复杂度。counts {double: sum( # double_blocks 模块的所有参数数量[sum(p.numel() for p in block.img_attn_qkv.parameters()) sum(p.numel() for p in block.img_attn_proj.parameters()) sum(p.numel() for p in block.img_mlp.parameters()) sum(p.numel() for p in block.txt_attn_qkv.parameters()) sum(p.numel() for p in block.txt_attn_proj.parameters()) sum(p.numel() for p in block.txt_mlp.parameters())for block in self.double_blocks]),single: sum( # single_blocks 模块的所有参数数量[sum(p.numel() for p in block.linear1.parameters()) sum(p.numel() for p in block.linear2.parameters())for block in self.single_blocks]),total: sum(p.numel() for p in self.parameters()),}# double 和 single 参数的总和主要聚焦于注意力和 MLP 层。counts[attnmlp] counts[double] counts[single]return counts
