🔮混元文生视频官网 | 🌟Github代码仓库 | 🎬 Demo 体验 | 📝技术报告 | 😍Hugging Face
文章目录
- 论文详解
-
- 基础介绍
- [数据预处理 (Data Pre-processing)](#数据预处理 (Data Pre-processing))
- [模型结构设计(Model Architecture Design)](#模型结构设计(Model Architecture Design))
-
- 3DVAE (3D Variational Auto-encoder Design)
- [统一图像视频生成架构 (Unified Image and Video Generative Architecture)](#统一图像视频生成架构 (Unified Image and Video Generative Architecture))
- [文本编码器(Text Encoder)](#文本编码器(Text Encoder))
- [模型缩放(Model Scaling)](#模型缩放(Model Scaling))
- 模型预训练(Model-Pretraining)
-
- [训练目标(Training Objective)](#训练目标(Training Objective))
- [图像预训练(Image Pre-Training)](#图像预训练(Image Pre-Training))
- [视频-图像联合训练(Video-Image Joint Training)](#视频-图像联合训练(Video-Image Joint Training))
- [提示重写(Prompt Rewrite)](#提示重写(Prompt Rewrite))
- [高效模型微调(High-performance Model Fine-tuning)](#高效模型微调(High-performance Model Fine-tuning))
- [模型加速(Model Acceleration)](#模型加速(Model Acceleration))
-
- [推理时间步减少(Inference Step Reduction)](#推理时间步减少(Inference Step Reduction))
- [文本引导蒸馏(Text-guidance Distillation)](#文本引导蒸馏(Text-guidance Distillation))
- [高效和扩展训练(Efficient and Scalable Training)](#高效和扩展训练(Efficient and Scalable Training))
-
- [硬件架构(Hardware Infrastucture)](#硬件架构(Hardware Infrastucture))
- [并行策略(Parallel Strategy)](#并行策略(Parallel Strategy))
- 优化(Optimization)
- [自动容错(Automatic fault tolerance)](#自动容错(Automatic fault tolerance))
- 应用
-
- [基于视频的音频生成(Audio Generation based on Video)](#基于视频的音频生成(Audio Generation based on Video))
- [混元图生视频(Hunyuan Image-to-Video)](#混元图生视频(Hunyuan Image-to-Video))
- [化身动画(Avatar animation)](#化身动画(Avatar animation))
- 代码实操
-
- 环境配置
- Demo运行
- 代码讲解
-
- 主函数 (sample_video.py)
- 模型推理(hyvideo/inference.py)
- 推理管道(hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py)
- 模型结构(hyvideo/modules/models.py)
-
- [双流块 (MMDoubleStreamBlock)](#双流块 (MMDoubleStreamBlock))
- [单流块 (MMSingleStreamBlock)](#单流块 (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个训练数据集,即,256p,360p,540p和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 (T+1)×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}) (ctT+1)×C×(csH)×(csW)的潜在特征。在实现中,设定 c t = 4 c_t = 4 ct=4, c s = 8 c_s = 8 cs=8, C = 16 C = 16 C=16。
T + 1 T+1 T+1的原因是,对视频的第一帧进行单独的编码,作为后续帧的参考帧。
这种压缩方法显著减少了后续扩散变换器模型所需的令牌数量,使其能够以原始分辨率和帧率训练视频,同时保持较高的效率和质量。模型结构如下图所示。
(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} Loss=L1+0.1Llpips+0.05Ladv+10−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,在各种设置下均表现出比传统文本编码器更优的性能。
(4)CLIP 特征的全局指导作用
尽管 MLLM 是主要的文本编码器,CLIP 文本特征作为文本信息的摘要也很有价值:
- 全局指导:HunyuanVideo 采用 CLIP-Large 文本特征的最终非填充 token 作为全局指导信息,集成到双流和单流的 DiT(Diffusion 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} Nopt=a1Cb1,Dopt=a2Cb2- 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 a1=5.48×10−4,b1=0.5634
- a 2 = 0.324 , b 2 = 0.4325 a_2 = 0.324, b_2 = 0.4325 a2=0.324,b2=0.4325
- 表明模型规模和数据规模与计算资源之间遵循幂律关系。
4. 模型优化
- 在不同计算预算下,通过扩展规律计算出模型和数据集的最优配置。
- 下图显示了不同规模模型的损失曲线,证明 DiT-T2X(I)家族符合幂律关系。
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 a1=0.0189,b1=0.3618
- a 2 = 0.0108 , b 2 = 0.6289 a_2 = 0.0108, b_2 = 0.6289 a2=0.0108,b2=0.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} Nopt=a1Cb1,Dopt=a2Cb2
- 如上图(d,e,f)展示了 T2X(V) 模型的扩展规律,拟合得到以下参数:
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} ut=dtdxt,指导样本 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 Lgeneration=Et,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)
背景和动机
- 早期实验表明,预训练模型显著加速了视频训练的收敛速度,并提升了视频生成性能。
- 为此,提出了一种两阶段的渐进式图像预训练策略,用于视频训练的热启动。
阶段 1:256px 图像训练
- 目标:模型首先在低分辨率(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} t=1−Qq,其中 t t t 表示噪声初始化为 t = 1 t = 1 t=1,并在 t = 0 t = 0 t=0时停止生成。
- 为优化推理步骤,我们将 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 的一个主要缺点是其高计算成本和推理延迟:
- 每个推理步骤都需要为无条件输入生成额外的输出。
- 尤其是在处理大型视频模型 和高分辨率视频生成时,同时生成文本条件与无条件视频的推理负担极其昂贵。
- 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
提示:1)720p x 1280p 分辨率大约需要76G显存,544p x 960p分辨率大约需要43G显存,合理租卡 2)预训练模型比较大,至少扩容50GB
- 安装huggingface_hub
shell
python -m pip install "huggingface_hub[cli]"
- 下载预训练模型
在HunyuanVideo
下新建ckpts
文件夹
shell
# 学术资源加速
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
- 模型预处理
shell
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_encoder
Demo运行
具体的参数配置见hyvideo/config.py
720p x 1280p
分辨率
python
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 ./results
544p x 960p 分辨率
shell
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()
函数,用于加载视频生成模型、生成视频样本并保存生成的结果。
python
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(f"`models_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_ok=True)
# 加载指定路径下的预训练模型,并传入 args 参数
hunyuan_video_sampler = HunyuanVideoSampler.from_pretrained(models_root_path, args=args)
# 更新 args 为模型内部参数,确保之后使用的一致性
args = hunyuan_video_sampler.args
# 生成视频样本
outputs = hunyuan_video_sampler.predict(
prompt=args.prompt, # 文本提示词,用于引导生成内容。
height=args.video_size[0],# 视频帧的分辨率(高度和宽度)
width=args.video_size[1],
video_length=args.video_length, # 视频时长(帧数)。
seed=args.seed, # 随机种子,用于结果的可重复性
negative_prompt=args.neg_prompt, # 负向提示词,指定生成时需避免的特性
infer_steps=args.infer_steps, # 推理步数
guidance_scale=args.cfg_scale,# 引导系数,控制生成质量。
num_videos_per_prompt=args.num_videos, # 每个提示生成的视频数量。
flow_shift=args.flow_shift, # 时间帧间的流动控制参数
batch_size=args.batch_size, # 批处理大小
embedded_guidance_scale=args.embedded_cfg_scale # 内嵌引导系数,用于调节特定特征
)
samples = outputs['samples']
# 保存视频样本
# 检查环境变量 LOCAL_RANK 是否存在,用于分布式训练的本地进程控制:如果不存在,或者值为 0(即主进程),则继续保存样本。
if 'LOCAL_RANK' not in os.environ or int(os.environ['LOCAL_RANK']) == 0:
# 遍历生成的samples
for 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('/','')}.mp4"
save_videos_grid(sample, save_path, fps=24) # 保存视频
logger.info(f'Sample save to: {save_path}') # 日志记录
模型推理(hyvideo/inference.py)
load_diffusion_pipeline
函数
用于加载视频生成的扩散推理管道(pipeline),包括模型、调度器(scheduler)、设备配置等必要的组件。
python
def load_diffusion_pipeline(
self,
args,
vae,
text_encoder,
text_encoder_2,
model,
scheduler=None,
device=None,
progress_bar_config=None,
data_type="video",
):
"""Load the denoising scheduler for inference."""
# 去噪调度器的初始化
if scheduler is None:
if args.denoise_type == "flow":
# 流动匹配的去噪策略,离散去噪调度器,可能用于视频生成任务中时间帧之间的一致性建模。
# 负责指导扩散模型逐步还原噪声,生成清晰的视频帧。
scheduler = FlowMatchDiscreteScheduler(
shift=args.flow_shift, # 流动偏移值。
reverse=args.flow_reverse, # 是否反向计算。
solver=args.flow_solver, # 去噪求解器的类型
)
else:
raise ValueError(f"Invalid denoise type {args.denoise_type}")
# 构建推理pipeline
pipeline = HunyuanVideoPipeline(
vae=vae, # 负责特征编码和解码的模块。
text_encoder=text_encoder, # 用于处理文本提示,生成与视频生成相关的特征。
text_encoder_2=text_encoder_2,
transformer=model, # 主扩散模型,生成视频的核心模块。
scheduler=scheduler, # 去噪调度器,控制扩散生成的时间步长和过程
progress_bar_config=progress_bar_config, # 可选的进度条配置,用于显示推理进度。
args=args, # 配置参数的集合
)
# 配置计算资源
if self.use_cpu_offload:
# 将部分计算任务卸载到 CPU。这是显存不足时的优化策略,可以大幅降低 GPU 的显存占用。
pipeline.enable_sequential_cpu_offload()
else:
# 如果为 False,直接将管道加载到指定的 device(如 GPU)上运行
pipeline = pipeline.to(device)
return pipeline
get_rotary_pos_embed
函数
其作用是计算用于视频处理的多维旋转位置嵌入(Rotary Positional Embedding, RoPE)
。它主要根据输入的视频维度、网络配置以及位置嵌入参数,生成对应的正弦和余弦频率嵌入,用于在 Transformer 等模型的注意力机制中进行位置编码。
python
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), (
f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.model.patch_size}), "
f"but got {latents_size}."
)
# 如果整除,计算 RoPE 的输入尺寸 rope_sizes(将 latents_size 每一维除以 patch_size)
rope_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] == 0
for idx, s in enumerate(latents_size)
), (
f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.model.patch_size}), "
f"but 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_list
if 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,
theta=self.args.rope_theta, #控制位置嵌入频率。
use_real=True, # 表示使用真实数值而非复数形式。
theta_rescale_factor=1, # 无缩放因子。
)
#返回: freqs_cos: 余弦频率嵌入。freqs_sin: 正弦频率嵌入。
return freqs_cos, freqs_sin
predict
函数
一个用于从文本生成视频或图像的预测函数,名为 predict
。函数通过输入文本 prompt
,结合其他参数(如视频分辨率、帧数、推理步数等),生成指定数量的视频或图像。
参数说明:
-
prompt
(str
orList[str]
):- 输入的文本描述,生成的图像或视频将基于此描述。
- 可以是单个字符串或字符串列表。
-
height
(int
):- 生成视频的高度(分辨率)。默认值为
192
。
- 生成视频的高度(分辨率)。默认值为
-
width
(int
):- 生成视频的宽度(分辨率)。默认值为
336
。
- 生成视频的宽度(分辨率)。默认值为
-
video_length
(int
):- 生成视频的帧数。默认值为
129
。
- 生成视频的帧数。默认值为
-
seed
(int
orList[int]
):- 随机种子,用于控制生成结果的随机性。
- 如果在分布式环境下运行(即模型被并行化),必须指定种子,以保证结果一致性。
-
negative_prompt
(str
orList[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
:- 其他额外参数,可以在调用时传入,以灵活配置生成过程。
python
@torch.no_grad()
def predict(
self,
prompt,
height=192,
width=336,
video_length=129,
seed=None,
negative_prompt=None,
infer_steps=50,
guidance_scale=6,
flow_shift=5.0,
embedded_guidance_scale=None,
batch_size=1,
num_videos_per_prompt=1,
**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 + i
for _ 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]) + j
for 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(
f"Length of seed must be equal to number of prompt(batch_size) or "
f"batch_size * num_videos_per_prompt ({batch_size} * {num_videos_per_prompt}), got {seed}."
)
else:
raise ValueError(
f"Seed 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(
f"`height` 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(
f"`video_length-1` must be a multiple of 4, got {video_length}"
)
# 日志记录
logger.info(
f"Input (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(f"`prompt` 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_prompt
if not isinstance(negative_prompt, str):
raise TypeError(
f"`negative_prompt` must be a string, but got {type(negative_prompt)}"
)
negative_prompt = [negative_prompt.strip()]
# ========================================================================
# 设置调度器 (Scheduler)
# ========================================================================
scheduler = FlowMatchDiscreteScheduler( # 处理流(Flow)的调度
shift=flow_shift, # 控制流动调度器的偏移量。flow_shift 通常与时序或流动模型相关,例如调整时间步之间的关系。
reverse=self.args.flow_reverse, # 决定是否反向调度(可能是在推理过程中逆序生成帧)
solver=self.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 = f"""
height: {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(
prompt=prompt, # 文本提示,用于指导生成内容。
height=target_height, # 生成图像或视频帧的分辨率。
width=target_width, #
video_length=target_video_length, # 视频的帧数。如果 video_length > 1,表示生成视频;否则生成单张图像。
num_inference_steps=infer_steps, # 推理步数,决定生成过程的细粒度程度,步数越多,生成结果越精细。
guidance_scale=guidance_scale, # 指导比例,控制生成与 prompt 的一致性程度。
negative_prompt=negative_prompt, # 负面提示,用于约束生成内容,避免不期望的结果。
num_videos_per_prompt=num_videos_per_prompt, # 每条提示生成的视频数量。
generator=generator, # 随机生成器对象,用于控制生成过程中的随机性,通常与随机种子结合。
output_type="pil", # 指定输出格式为 PIL.Image 对象,便于后续处理
freqs_cis=(freqs_cos, freqs_sin), # 旋转位置嵌入 (RoPE) 的频率矩阵,增强时空位置感知能力。
n_tokens=n_tokens, # 输入序列的总标记数,用于指导生成过程。
embedded_guidance_scale=embedded_guidance_scale, # 嵌入式指导比例,用于进一步优化嵌入向量的生成。
data_type="video" if target_video_length > 1 else "image", # 指定生成目标为视频或图像,取决于帧数。
is_progress_bar=True, # 显示推理进度条,方便监控生成进度。
vae_ver=self.args.vae, # 使用指定版本的 VAE(变分自编码器),决定生成内容的潜在空间。
enable_tiling=self.args.vae_tiling, # 启用 VAE 分块处理,提高内存效率,特别适用于高分辨率生成。
)[0] # 返回生成的样本,通常是一个 PIL.Image 或视频帧序列
# 保存生成结果
out_dict["samples"] = samples
out_dict["prompts"] = prompt
# 计算并记录推理时间
gen_time = time.time() - start_time
logger.info(f"Success, time: {gen_time}")
return out_dict
推理管道(hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py)
__call__
方法接受用户输入的提示、生成图像或视频的尺寸,以及其他生成过程的参数,完成推理并返回生成的图像或视频。
参数详解
1. 提示相关参数
-
prompt
用于引导生成的文本提示,可以是字符串(如
"海边的晴天"
)或字符串列表。示例:
prompt="在森林中散步的机器人"
-
negative_prompt
用于指示生成时需要避免的内容(负面提示)。可以是字符串或字符串列表。
示例:
negative_prompt="没有树木"
2. 图像或视频生成的尺寸参数
-
height
和width
指定生成的图像或视频帧的高度和宽度(单位:像素)。
示例:
height=512, width=512
表示生成 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
指定生成输出的类型,默认为
"pil"
(Pillow 图像对象)。可选"np.array"
(NumPy 数组)。 -
return_dict
是否以字典形式返回结果。如果为
True
,则返回类似HunyuanVideoPipelineOutput
的对象,否则返回元组。
6. 其他高级参数
-
cross_attention_kwargs
传递给交叉注意力模块的额外参数,用于调整注意力机制。
-
callback_on_step_end
回调函数,在每次去噪步骤结束时调用,可用于自定义处理逻辑。
-
clip_skip
跳过 CLIP 模型的层数,用于调整文本嵌入的生成层。
-
guidance_rescale
用于调整引导时的比例因子,以修正某些生成过程中的过曝问题。
返回值
- 如果
return_dict=True
,返回一个包含生成结果的字典对象(如HunyuanVideoPipelineOutput
)。 - 如果
return_dict=False
,返回一个元组,第一个元素是生成的图像/视频列表,第二个元素是布尔列表,指示生成内容是否包含敏感内容。
python
@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,
):
r"""
The 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 the
expense of slower inference.
timesteps (`List[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed 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 in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will 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 text
`prompt` 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 to
pass `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 applies
to 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 make
generation deterministic.
latents (`torch.Tensor`, *optional*):
Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
tensor 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 not
provided, 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). If
not 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 a
plain 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 are
Flawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure when
using 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 that
the 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 of
each 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 a
list 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 list
will 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 the
second element is a list of `bool`s indicating whether the corresponding generated image contains
"not-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_ver=vae_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]
# 确定设备的device
device = torch.device(f"cuda:{dist.get_rank()}") if dist.is_initialized() else self._execution_device
# 3. Encode input prompt
# 处理 LoRA(Low-Rank Adaptation)缩放系数:通过 cross_attention_kwargs 提取或设置缩放比例 lora_scale
lora_scale = (
self.cross_attention_kwargs.get("scale", None)
if self.cross_attention_kwargs is not None
else 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_embeds=prompt_embeds,
attention_mask=attention_mask,
negative_prompt_embeds=negative_prompt_embeds,
negative_attention_mask=negative_attention_mask,
lora_scale=lora_scale,
clip_skip=self.clip_skip,
data_type=data_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_embeds=None,
attention_mask=None,
negative_prompt_embeds=None,
negative_attention_mask=None,
lora_scale=lora_scale,
clip_skip=self.clip_skip,
text_encoder=self.text_encoder_2,
data_type=data_type,
)
else:
prompt_embeds_2 = None
negative_prompt_embeds_2 = None
prompt_mask_2 = None
negative_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 + 1
elif "888" in vae_ver:
video_length = (video_length - 1) // 8 + 1
else:
video_length = video_length
# 5. Prepare latent variables
num_channels_latents = self.transformer.config.in_channels
latents = 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 pipeline
extra_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.order
self._num_timesteps = len(timesteps)
# if is_progress_bar:
# progress_bar 用于显示推理过程的进度,num_inference_steps 是总推理步数。
with self.progress_bar(total=num_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_guidance
else 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],
dtype=torch.float32,
device=device,
).to(target_dtype)
* 1000.0
if embedded_guidance_scale is not None
else None
)
# 使用 Transformer 模型预测噪声残差
with torch.autocast(
device_type="cuda", dtype=target_dtype, enabled=autocast_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_states=prompt_embeds, # 与文本提示相关的嵌入向量 [2, 256, 4096]
text_mask=prompt_mask, # [2, 256]
text_states_2=prompt_embeds_2, # [2, 768]
freqs_cos=freqs_cis[0], # 频率信息,用于特定的时间步缩放 [seqlen, head_dim]
freqs_sin=freqs_cis[1], # [seqlen, head_dim]
guidance=guidance_expand, # 指导参数,用于条件生成
return_dict=True,
)[
"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.pdf
noise_pred = rescale_noise_cfg(
noise_pred,
noise_pred_text,
guidance_rescale=self.guidance_rescale,
)
# 使用调度器更新潜变量
# compute the previous noisy sample x_t -> x_t-1
latents = self.scheduler.step(
noise_pred, t, latents, **extra_step_kwargs, return_dict=False
)[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_dim=True,标记后续需要移除该额外维度。
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:
pass
else:
raise ValueError(
f"Only 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_factor
with torch.autocast(
device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled
):
if enable_tiling:
# 调用 VAE 的 enable_tiling() 方法,可能用于解码较大的图像块。
self.vae.enable_tiling()
# 使用 VAE(变分自编码器)的 decode 方法将潜变量解码为图像。
image = self.vae.decode(
latents, return_dict=False, generator=generator
)[0]
else:
image = self.vae.decode(
latents, return_dict=False, generator=generator
)[0]
# 如果添加了时间维度(expand_temporal_dim=True),或者解码出的图像在时间维度上只有一个帧,则移除时间维度。
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_dict=False),则直接返回处理后的图像
if not return_dict:
return image
return HunyuanVideoPipelineOutput(videos=image)
模型结构(hyvideo/modules/models.py)
双流块 (MMDoubleStreamBlock)
python
class MMDoubleStreamBlock(nn.Module):
"""
A multimodal dit block with seperate modulation for
text and image/video, see more details (SD3): https://arxiv.org/abs/2403.03206
(Flux.1): https://github.com/black-forest-labs/flux
"""
def __init__(
self,
hidden_size: int, # 模型隐藏层维度。
heads_num: int, # 多头注意力的头数。
mlp_width_ratio: float, # MLP 中隐藏层宽度与 hidden_size 的比率。
mlp_act_type: str = "gelu_tanh", # 激活函数的类型(默认 gelu_tanh)
qk_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 = False
self.heads_num = heads_num
head_dim = hidden_size // heads_num
mlp_hidden_dim = int(hidden_size * mlp_width_ratio)
### 图像模态
# 模态调制模块:使用 ModulateDiT,为图像和文本生成 6 组参数(shift、scale、gate)。
self.img_mod = ModulateDiT(
hidden_size,
factor=6,
act_layer=get_activation_layer("silu"),
**factory_kwargs,
)
# 归一化
self.img_norm1 = nn.LayerNorm(
hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs
)
# QKV 投影层:通过全连接层计算 Query、Key 和 Value
self.img_attn_qkv = nn.Linear(
hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs
)
# 归一化模块
qk_norm_layer = get_norm_layer(qk_norm_type)
self.img_attn_q_norm = (
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
if qk_norm
else nn.Identity()
)
self.img_attn_k_norm = (
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
if qk_norm
else nn.Identity()
)
self.img_attn_proj = nn.Linear(
hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs
)
self.img_norm2 = nn.LayerNorm(
hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs
)
self.img_mlp = MLP(
hidden_size,
mlp_hidden_dim,
act_layer=get_activation_layer(mlp_act_type),
bias=True,
**factory_kwargs,
)
### 文本模态
self.txt_mod = ModulateDiT(
hidden_size,
factor=6,
act_layer=get_activation_layer("silu"),
**factory_kwargs,
)
self.txt_norm1 = nn.LayerNorm(
hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs
)
self.txt_attn_qkv = nn.Linear(
hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs
)
self.txt_attn_q_norm = (
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
if qk_norm
else nn.Identity()
)
self.txt_attn_k_norm = (
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
if qk_norm
else nn.Identity()
)
self.txt_attn_proj = nn.Linear(
hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs
)
self.txt_norm2 = nn.LayerNorm(
hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs
)
self.txt_mlp = MLP(
hidden_size,
mlp_hidden_dim,
act_layer=get_activation_layer(mlp_act_type),
bias=True,
**factory_kwargs,
)
self.hybrid_seq_parallel_attn = None
def enable_deterministic(self):
self.deterministic = True
def disable_deterministic(self):
self.deterministic = False
def 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, shift=img_mod1_shift, scale=img_mod1_scale
)
# 得到 Query、Key 和 Value
img_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", K=3, H=self.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_first=False)
assert (
img_qq.shape == img_q.shape and img_kk.shape == img_k.shape
), f"img_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, shift=txt_mod1_shift, scale=txt_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", K=3, H=self.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), dim=1)
k = torch.cat((img_k, txt_k), dim=1)
v = torch.cat((img_v, txt_v), dim=1)
assert (
cu_seqlens_q.shape[0] == 2 * img.shape[0] + 1
), f"cu_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_q=cu_seqlens_q,
cu_seqlens_kv=cu_seqlens_kv,
max_seqlen_q=max_seqlen_q,
max_seqlen_kv=max_seqlen_kv,
batch_size=img_k.shape[0],
)
else:
attn = parallel_attention(
self.hybrid_seq_parallel_attn,
q,
k,
v,
img_q_len=img_q.shape[1],
img_kv_len=img_k.shape[1],
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_kv=cu_seqlens_kv
)
# 最终将注意力结果拆分为图像部分 img_attn 和文本部分 txt_attn
img_attn, txt_attn = attn[:, : img.shape[1]], attn[:, img.shape[1] :]
'''图像模态的更新'''
# 将注意力结果通过残差连接更新图像特征,并通过 MLP 进一步增强
img = img + apply_gate(self.img_attn_proj(img_attn), gate=img_mod1_gate)
img = img + apply_gate(
self.img_mlp(
modulate(
self.img_norm2(img), shift=img_mod2_shift, scale=img_mod2_scale
)
),
gate=img_mod2_gate,
)
'''文本模态的更新'''
txt = txt + apply_gate(self.txt_attn_proj(txt_attn), gate=txt_mod1_gate)
txt = txt + apply_gate(
self.txt_mlp(
modulate(
self.txt_norm2(txt), shift=txt_mod2_shift, scale=txt_mod2_scale
)
),
gate=txt_mod2_gate,
)
# 返回更新后的图像特征和文本特征
return img, txt
单流块 (MMSingleStreamBlock)
python
class MMSingleStreamBlock(nn.Module):
"""
A DiT block with parallel linear layers as described in
https://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/flux
"""
def __init__(
self,
hidden_size: int, # 隐藏层的维度大小,用于表示特征的维度。
heads_num: int, # 注意力头的数量。
mlp_width_ratio: float = 4.0, # 用于确定多层感知机 (MLP) 的隐藏层宽度比例,默认值为 4.0
mlp_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 = False
self.hidden_size = hidden_size
self.heads_num = heads_num
head_dim = hidden_size // heads_num
mlp_hidden_dim = int(hidden_size * mlp_width_ratio)
self.mlp_hidden_dim = mlp_hidden_dim
self.scale = qk_scale or head_dim ** -0.5
# qkv and mlp_in
self.linear1 = nn.Linear(
hidden_size, hidden_size * 3 + mlp_hidden_dim, **factory_kwargs
)
# proj and mlp_out
self.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_affine=True, eps=1e-6, **factory_kwargs)
if qk_norm
else nn.Identity()
)
self.k_norm = (
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
if qk_norm
else nn.Identity()
)
self.pre_norm = nn.LayerNorm(
hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs
)
self.mlp_act = get_activation_layer(mlp_act_type)()
self.modulation = ModulateDiT(
hidden_size,
factor=3,
act_layer=get_activation_layer("silu"),
**factory_kwargs,
)
self.hybrid_seq_parallel_attn = None
def enable_deterministic(self):
self.deterministic = True
def disable_deterministic(self):
self.deterministic = False
def 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), shift=mod_shift, scale=mod_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", K=3, H=self.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_first=False)
assert (
img_qq.shape == img_q.shape and img_kk.shape == img_k.shape
), f"img_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), dim=1)
k = torch.cat((img_k, txt_k), dim=1)
# Compute attention.
assert (
cu_seqlens_q.shape[0] == 2 * x.shape[0] + 1
), f"cu_seqlens_q.shape:{cu_seqlens_q.shape}, x.shape[0]:{x.shape[0]}"
# attention computation start
if not self.hybrid_seq_parallel_attn:
# 如果没有启用并行注意力机制,调用标准注意力函数 attention
attn = attention(
q,
k,
v,
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_kv=cu_seqlens_kv,
max_seqlen_q=max_seqlen_q,
max_seqlen_kv=max_seqlen_kv,
batch_size=x.shape[0],
)
else:
# 否则,使用并行注意力机制 parallel_attention
attn = parallel_attention(
self.hybrid_seq_parallel_attn,
q,
k,
v,
img_q_len=img_q.shape[1],
img_kv_len=img_k.shape[1],
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_kv=cu_seqlens_kv
)
# attention computation end
# 将注意力结果和 MLP 激活结果拼接,通过线性层投影回输入维度。
output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
# 使用 mod_gate 进行门控融合,将残差连接后的结果返回。
return x + apply_gate(output, gate=mod_gate)
混元主干网络(HYVideoDiffusionTransformer)
python
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.03206
Parameters
----------
args: argparse.Namespace
传入的命令行参数,用于设置模型的配置。
patch_size: list
输入特征的分块尺寸。一般用于图像或视频的分块操作。
in_channels: int
输入数据的通道数(如 RGB 图像为 3 通道)。
out_channels: int
模型输出的通道数。
hidden_size: int
Transformer 模块中隐藏层的维度。
heads_num: int
多头注意力机制中的注意力头数量,通常用来分配不同的注意力特征。
mlp_width_ratio: float
MLP(多层感知机)中隐藏层维度相对于 hidden_size 的比例。
mlp_act_type: str
MLP 使用的激活函数类型,例如 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: str
QK 归一化的类型。
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_config
def __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_size
self.in_channels = in_channels
self.out_channels = in_channels if out_channels is None else out_channels
self.unpatchify_channels = self.out_channels # 用来重新拼接patch时的通道数。
self.guidance_embed = guidance_embed
self.rope_dim_list = rope_dim_list
# Text projection. Default to linear projection.
# Alternative: TokenRefiner. See more details (LI-DiT): http://arxiv.org/abs/2406.11831
self.use_attention_mask = use_attention_mask
self.text_projection = text_projection
self.text_states_dim = args.text_states_dim
self.text_states_dim_2 = args.text_states_dim_2
# 确保每个头的维度是整数。
if hidden_size % heads_num != 0:
raise ValueError(
f"Hidden 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(
f"Got {rope_dim_list} but expected positional dim {pe_dim}"
)
self.hidden_size = hidden_size
self.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, depth=2, **factory_kwargs
)
else:
raise NotImplementedError(
f"Unsupported 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_embed
else None
)
# MMDoubleStreamBlock:多模态双流块,融合了图像流和文本流信息。
self.double_blocks = nn.ModuleList(
[
MMDoubleStreamBlock(
self.hidden_size,
self.heads_num,
mlp_width_ratio=mlp_width_ratio,
mlp_act_type=mlp_act_type,
qk_norm=qk_norm,
qk_norm_type=qk_norm_type,
qkv_bias=qkv_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_ratio=mlp_width_ratio,
mlp_act_type=mlp_act_type,
qk_norm=qk_norm,
qk_norm_type=qk_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 = x
txt = text_states
_, _, ot, oh, ow = x.shape
# 得到划分patch后的t,h,w
tt, 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(
"Didn't 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(
f"Unsupported 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 attention
cu_seqlens_q = get_cu_seqlens(text_mask, img_seq_len)
cu_seqlens_kv = cu_seqlens_q
max_seqlen_q = img_seq_len + txt_seq_len
max_seqlen_kv = max_seqlen_q
freqs_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"] = img
return out
return img
def unpatchify(self, x, t, h, w):
# 是将被切分为小块(patches)的特征重新还原成原始的张量形状,通常用于图像处理任务中,
# 例如在 ViT(Vision Transformer)模型的输出阶段将 patch 还原为完整图像的形式。
"""
x: (N, T, patch_size**2 * C) (批量大小,时间帧数,每个patch中的通道数)
imgs: (N, H, W, C)
"""
c = self.unpatchify_channels
pt, ph, pw = self.patch_size
assert 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 imgs
def 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["attn+mlp"] = counts["double"] + counts["single"]
return counts