使用 Python 从头开始​​构建 AI 文本生成视频模型

OpenAI 的 Sora、Stability AI 的 Stable Video Diffusion 以及许多其他已经问世或未来将出现的文本转视频模型,都是继大语言模型之后,2024 年最流行的 AI 趋势之一(LLMs)。在本博客中,我们将从头开始构建一个小型文本生成视频模型。我们将输入一个文本提示,我们训练的模型将根据该提示生成一个视频。该博客将涵盖从理解理论概念到编码整个架构并生成最终结果的所有内容。

由于我没有高级的 GPU,所以我编写了小型架构。以下是在不同处理器上训练模型所需时间的比较:

Training Videos Epochs CPU GPU A10 GPU T4
10K 30 more than 3 hr 1 hr 1 hr 42m
30K 30 more than 6 hr 1 hr 30 2 hr 30
100K 30 - 3-4 hr 5-6 hr

在 CPU 上运行显然需要更长的时间来训练模型。如果您需要快速测试代码中的更改并查看结果,CPU 并不是最佳选择。我建议使用 ColabKaggle 的 T4 GPU 来实现更高效、更快的训练。

为了避免复制和粘贴此博客中的代码,以下是包含笔记本文件以及所有代码和信息的 GitHub 存储库:
AI-text-to-video-model-from-scratch

以下是博客链接,指导您如何从头开始创建Stable Diffusion:
从头开始编码Stable Diffusion

目录

  1. 我们正在建设什么
  2. 先决条件
  3. 了解 GAN 架构
  4. 搭建舞台
  5. 对训练数据进行编码
  6. 预处理我们的训练数据
  7. 实现文本嵌入层
  8. 实现生成器层
  9. 实施鉴别器层
  10. 编码训练参数
  11. 训练循环编码
  12. 保存训练后的模型
  13. 生成人工智能视频
  14. 少了什么东西?

我们正在建设什么

我们将遵循与传统机器学习或深度学习模型类似的方法,在数据集上进行训练,然后在未见过的数据上进行测试。在文本到视频的背景下,假设我们有一个包含 10 万个狗捡球和猫追老鼠视频的训练数据集。我们将训练我们的模型来生成猫捡球或狗追老鼠的视频。 尽管此类训练数据集很容易在互联网上获得,但所需的计算能力非常高。因此,我们将使用由 Python 代码生成的移动对象的视频数据集。

我们将使用 GAN(生成对抗网络)架构来创建我们的模型,而不是 OpenAI Sora 使用的扩散模型。我尝试使用扩散模型,但由于内存要求而崩溃,这超出了我的能力。另一方面,GAN 的训练和测试更容易、更快捷。

先决条件

我们将使用 OOP(面向对象编程),因此您必须对它和神经网络有基本的了解。 GAN(生成对抗网络)的知识不是强制性的,因为我们将在这里介绍它们的架构。

主题 链接
面向对象编程 视频链接
神经网络理论 视频链接
生成式对抗网络架构 视频链接
Python 基础知识 视频链接

了解 GAN 架构

理解 GAN 很重要,因为我们的大部分架构都依赖于它。让我们来探讨一下它是什么、它的组件等等。

什么是GAN?

生成对抗网络 (GAN) 是一种深度学习模型,其中两个神经网络相互竞争:一个从给定的数据集中创建新数据(例如图像或音乐),另一个尝试判断数据是真实的还是虚假的。这个过程一直持续到生成的数据与原始数据无法区分为止。

实际应用

  1. 生成图像:GAN 根据文本提示创建逼真的图像或修改现有图像,例如增强分辨率或为黑白照片添加颜色。
  2. 数据增强:它们生成合成数据来训练其他机器学习模型,例如为欺诈检测系统创建欺诈交易数据。
  3. 补全缺失的信息:GAN 可以填充缺失的数据,例如从地形图生成地下图像以用于能源应用。
  4. 生成 3D 模型:它们将 2D 图像转换为 3D 模型,这在医疗保健等领域非常有用,可以为手术规划创建逼真的器官图像。

GAN 是如何工作的?

它由两个深度神经网络组成:生成器和鉴别器。这些网络在对抗性设置中一起训练,其中一个网络生成新数据,另一个网络评估数据是真实的还是虚假的。

以下是 GAN 工作原理的简单概述:

  1. 训练集分析:生成器分析训练集以识别数据属性,而判别器独立分析相同的数据以学习其属性。
  2. 数据修改:生成器向数据的某些属性添加噪声(随机变化)。
  3. 数据传递:修改后的数据然后被传递到鉴别器。
  4. 概率计算:判别器计算生成的数据来自原始数据集的概率。
  5. 反馈循环:鉴别器向生成器提供反馈,指导生成器减少下一个周期的随机噪声。
  6. 对抗性训练:生成器试图最大化判别器的错误,而判别器则试图最小化自己的错误。通过多次训练迭代,两个网络都得到改进和发展。
  7. 平衡状态:训练继续,直到判别器无法再区分真实数据和合成数据,表明生成器已成功学会生成真实数据。至此,训练过程就完成了。

GAN 训练示例

让我们以图像到图像转换的示例来解释 GAN 模型,重点是修改人脸。

  1. 输入图像:输入是真实的人脸图像。
  2. 属性修改:生成器修改脸部的属性,例如为眼睛添加墨镜。
  3. 生成的图像:生成器创建一组添加了太阳镜的图像。
  4. 鉴别器的任务:鉴别器接收真实图像(戴太阳镜的人)和生成图像(添加太阳镜的人脸)的混合。
  5. 评估:鉴别器试图区分真实图像和生成图像。
  6. 反馈循环:如果鉴别器正确识别出假图像,则生成器会调整其参数以产生更令人信服的图像。如果生成器成功欺骗了鉴别器,鉴别器就会更新其参数以改进其检测。

通过这个对抗过程,两个网络都在不断改进。生成器在创建真实图像方面变得更好,鉴别器在识别赝品方面也变得更好,直到达到平衡,鉴别器不再能够区分真实图像和生成图像之间的差异。至此,GAN 已成功学会产生现实的修改。

搭建舞台

我们将使用一系列 Python 库,让我们导入它们。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
# 操作系统模块,用于与操作系统交互
import os

# 生成随机数的模块
import random

# 数值运算模块
import numpy as np

# 用于图像处理的 OpenCV 库
import cv2

# 用于图像处理的 Python 图像库
from PIL import Image, ImageDraw, ImageFont

# 用于深度学习的 PyTorch 库
import torch

# 用于在 PyTorch 中创建自定义数据集的数据集类
from torch.utils.data import Dataset

# 图像转换模块
import torchvision.transforms as transforms

# PyTorch 中的神经网络模块
import torch.nn as nn

# PyTorch 中的优化算法
import torch.optim as optim

# PyTorch 中填充序列的函数
from torch.nn.utils.rnn import pad_sequence

# PyTorch 中保存图像的函数
from torchvision.utils import save_image

# 用于绘制图形和图像的模块
import matplotlib.pyplot as plt

# 用于在 IPython 环境中显示丰富内容的模块
from IPython.display import clear_output, display, HTML

# 用于将二进制数据编码和解码为文本的模块
import base64
现在我们已经导入了所有库,下一步是定义我们将用来训练 GAN 架构的训练数据。

对训练数据进行编码

我们需要至少 10,000 个视频作为训练数据。为什么?嗯,因为我用较小的数字进行了测试,结果很差,几乎没有什么可看的。下一个大问题是:这些视频是关于什么的?我们的训练视频数据集由一个以不同运动向不同方向移动的圆圈组成。那么,让我们对其进行编码并生成 10,000 个视频来看看它是什么样子。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# 创建一个名为“training_dataset”的目录
os.makedirs('training_dataset', exist_ok=True)

# 定义为数据集生成的视频数量
num_videos = 10000

# 定义每个视频的帧数(1 秒视频)
frames_per_video = 10

# 定义数据集中每个图像的大小
img_size = (64, 64)

# 定义形状的大小(圆形)
shape_size = 10
设置一些基本参数后,接下来我们需要定义训练数据集的文本提示,根据该文本提示将生成训练视频。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
# 定义圆圈的文本提示和相应的动作
prompts_and_movements = [
# 向下移动圆圈
("circle moving down", "circle", "down"),
# 向左移动圆圈
("circle moving left", "circle", "left"),
# 向右移动圆圈
("circle moving right", "circle", "right"),
# 向右对角移动圆圈
("circle moving diagonally up-right", "circle", "diagonal_up_right"),
# 向左下对角移动圆圈
("circle moving diagonally down-left", "circle", "diagonal_down_left"),
# 向左斜上方移动圆圈
("circle moving diagonally up-left", "circle", "diagonal_up_left"),
# 向右下对角移动圆圈
("circle moving diagonally down-right", "circle", "diagonal_down_right"),
# 顺时针旋转圆圈
("circle rotating clockwise", "circle", "rotate_clockwise"),
# 逆时针旋转圆圈
("circle rotating counter-clockwise", "circle", "rotate_counter_clockwise"),
# 缩小圆圈
("circle shrinking", "circle", "shrink"),
# 扩大圆圈
("circle expanding", "circle", "expand"),
# 垂直弹跳圆圈
("circle bouncing vertically", "circle", "bounce_vertical"),
# 水平弹跳圆圈
("circle bouncing horizontally", "circle", "bounce_horizontal"),
# 垂直锯齿形圆圈
("circle zigzagging vertically", "circle", "zigzag_vertical"),
# 水平锯齿形圆圈
("circle zigzagging horizontally", "circle", "zigzag_horizontal"),
# 向左上方移动圆圈
("circle moving up-left", "circle", "up_left"),
# 向右下移动圆圈
("circle moving down-right", "circle", "down_right"),
# 向左下移动圆圈
("circle moving down-left", "circle", "down_left"),
]
我们使用这些提示定义了圆圈的几种运动。现在,我们需要编写一些数学方程来根据提示移动该圆圈。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
# 定义带参数的函数
def create_image_with_moving_shape(size, frame_num, shape, direction):

# 创建指定大小、白色背景的新 RGB 图像
img = Image.new('RGB', size, color=(255, 255, 255))

# 为图像创建绘图上下文
draw = ImageDraw.Draw(img)

# 计算图像的中心坐标
center_x, center_y = size[0] // 2, size[1] // 2

# 所有动作均以中心为初始位置
position = (center_x, center_y)

# 定义字典将方向映射到各自的位置调整或图像转换
direction_map = {
# 根据帧数向下调整位置
"down": (0, frame_num * 5 % size[1]),
# 根据帧数向左调整位置
"left": (-frame_num * 5 % size[0], 0),
# 根据帧数向右调整位置
"right": (frame_num * 5 % size[0], 0),
# 调整位置对角向上和向右
"diagonal_up_right": (frame_num * 5 % size[0], -frame_num * 5 % size[1]),
# 调整位置对角向下和向左
"diagonal_down_left": (-frame_num * 5 % size[0], frame_num * 5 % size[1]),
# 调整位置对角向上和向左
"diagonal_up_left": (-frame_num * 5 % size[0], -frame_num * 5 % size[1]),
# 向右下方调整位置
"diagonal_down_right": (frame_num * 5 % size[0], frame_num * 5 % size[1]),
# 根据帧数顺时针旋转图像
"rotate_clockwise": img.rotate(frame_num * 10 % 360, center=(center_x, center_y), fillcolor=(255, 255, 255)),
# 根据帧数逆时针旋转图像
"rotate_counter_clockwise": img.rotate(-frame_num * 10 % 360, center=(center_x, center_y), fillcolor=(255, 255, 255)),
# 垂直调整弹跳效果的位置
"bounce_vertical": (0, center_y - abs(frame_num * 5 % size[1] - center_y)),
# 调整水平弹跳效果的位置
"bounce_horizontal": (center_x - abs(frame_num * 5 % size[0] - center_x), 0),
# 垂直调整锯齿形效果的位置
"zigzag_vertical": (0, center_y - frame_num * 5 % size[1]) if frame_num % 2 == 0 else (0, center_y + frame_num * 5 % size[1]),
# 水平调整锯齿效果的位置
"zigzag_horizontal": (center_x - frame_num * 5 % size[0], center_y) if frame_num % 2 == 0 else (center_x + frame_num * 5 % size[0], center_y),
# 根据帧数向上和向右调整位置
"up_right": (frame_num * 5 % size[0], -frame_num * 5 % size[1]),
# 根据帧数向上和向左调整位置
"up_left": (-frame_num * 5 % size[0], -frame_num * 5 % size[1]),
# 根据帧数向下和向右调整位置
"down_right": (frame_num * 5 % size[0], frame_num * 5 % size[1]),
# 根据帧数向下和向左调整位置
"down_left": (-frame_num * 5 % size[0], frame_num * 5 % size[1])
}

# 检查方向是否在方向图中
if direction in direction_map:
# 检查方向是否映射到位置调整
if isinstance(direction_map[direction], tuple):
# 根据调整更新位置
position = tuple(np.add(position, direction_map[direction]))
else: # 如果方向映射到图像变换
# 根据变换更新图像
img = direction_map[direction]

# 将图像作为 numpy 数组返回
return np.array(img)
上面的函数用于根据所选方向移动每一帧的圆。我们只需要在其上运行一个循环,直到达到视频数量的次数即可生成所有视频。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# 迭代要生成的视频数量
for i in range(num_videos):
# 从预定义列表中随机选择提示和动作
prompt, shape, direction = random.choice(prompts_and_movements)

# 为当前视频创建目录
video_dir = f'training_dataset/video_{i}'
os.makedirs(video_dir, exist_ok=True)

# 将所选提示写入视频目录中的文本文件
with open(f'{video_dir}/prompt.txt', 'w') as f:
f.write(prompt)

# 为当前视频生成帧
for frame_num in range(frames_per_video):
# 根据当前帧数、形状和方向创建具有移动形状的图像
img = create_image_with_moving_shape(img_size, frame_num, shape, direction)

# 将生成的图像保存为PNG文件在视频目录中
cv2.imwrite(f'{video_dir}/frame_{frame_num}.png', img)
运行上述代码后,它将生成我们的整个训练数据集。这是我们的训练数据集文件的结构。 每个训练视频文件夹都包含其帧及其文本提示。让我们看一下训练数据集的样本。 在我们的训练数据集中,我们没有包含圆圈向上移动然后向右移动的运动。我们将使用它作为我们的测试提示来评估我们在未见过的数据上训练的模型。

需要注意的更重要的一点是,我们的训练数据确实包含许多样本,其中物体远离场景或部分出现在相机前面,类似于我们在 OpenAI Sora 演示视频中观察到的情况。

在我们的训练数据中包含此类样本的原因是为了测试当圆圈从最角落进入场景而不破坏其形状时,我们的模型是否能够保持一致性。

现在我们的训练数据已经生成,我们需要将训练视频转换为张量,这是 PyTorch 等深度学习框架中使用的主要数据类型。此外,执行归一化等转换有助于通过将数据扩展到更小的范围来提高训练架构的收敛性和稳定性。

预处理我们的训练数据

我们必须为文本到视频任务编写一个数据集类,它可以从训练数据集目录中读取视频帧及其相应的文本提示,使其可在 PyTorch 中使用。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
# 定义一个继承自torch.utils.data.Dataset的数据集类
class TextToVideoDataset(Dataset):
def __init__(self, root_dir, transform=None):
# 使用根目录和可选转换初始化数据集
self.root_dir = root_dir
self.transform = transform
# 列出根目录下的所有子目录
self.video_dirs = [os.path.join(root_dir, d) for d in os.listdir(root_dir) if os.path.isdir(os.path.join(root_dir, d))]
# 初始化列表来存储帧路径和相应的提示
self.frame_paths = []
self.prompts = []

# 循环遍历每个视频目录
for video_dir in self.video_dirs:
# 列出视频目录中的所有PNG文件并存储它们的路径
frames = [os.path.join(video_dir, f) for f in os.listdir(video_dir) if f.endswith('.png')]
self.frame_paths.extend(frames)
# 读取视频目录下的提示文本文件并存储其内容
with open(os.path.join(video_dir, 'prompt.txt'), 'r') as f:
prompt = f.read().strip()
# 对视频中的每一帧重复提示并存储在提示列表中
self.prompts.extend([prompt] * len(frames))

# 返回数据集中的样本总数
def __len__(self):
return len(self.frame_paths)

# 从给定索引的数据集中检索样本
def __getitem__(self, idx):
# 获取给定索引对应的帧的路径
frame_path = self.frame_paths[idx]
# 使用PIL(Python图像库)打开图像
image = Image.open(frame_path)
# 获取给定索引对应的提示
prompt = self.prompts[idx]

# 如果指定则应用转换
if self.transform:
image = self.transform(image)

# 返回转换后的图像和提示
return image, prompt
在继续对架构进行编码之前,我们需要规范化我们的训练数据。我们将使用 16 的批量大小并对数据进行打乱以引入更多随机性。
1
2
3
4
5
6
7
8
9
10
11
12
# 定义一组要应用于数据的转换
transform = transforms.Compose([
# 将 PIL 图像或 numpy.ndarray 转换为张量
transforms.ToTensor(),
# 使用平均值和标准差标准化图像
transforms.Normalize((0.5,), (0.5,))
])

# 使用定义的转换加载数据集
dataset = TextToVideoDataset(root_dir='training_dataset', transform=transform)
# 创建一个数据加载器来迭代数据集
dataloader = torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=True)

实现文本嵌入层

您可能已经在 Transformer 架构中看到过,其中的起点是将我们的文本输入转换为嵌入,以便在多头注意力中进一步处理,与这里类似,我们必须编写一个文本嵌入层,基于该层,GAN 架构训练将在我们的嵌入数据上进行和图像张量。

1
2
3
4
5
6
7
8
9
10
11
12
13
# 定义文本嵌入类
class TextEmbedding(nn.Module):
# 带有 vocab_size 和 embed_size 参数的构造方法
def __init__(self, vocab_size, embed_size):
# 调用超类构造函数
super(TextEmbedding, self).__init__()
# 初始化嵌入层
self.embedding = nn.Embedding(vocab_size, embed_size)

# 定义前向传播方法
def forward(self, x):
# 返回输入的嵌入表示
return self.embedding(x)
词汇量大小将基于我们的训练数据,稍后我们将计算这些数据。嵌入大小将为 10。如果使用更大的数据集,您还可以使用自己选择的 Hugging Face 上提供的嵌入模型。

实现生成器层

现在我们已经知道生成器在 GAN 中的作用,让我们对这一层进行编码,然后了解其内容。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class Generator(nn.Module):
def __init__(self, text_embed_size):
super(Generator, self).__init__()

# 以噪声和文本嵌入作为输入的全连接层
self.fc1 = nn.Linear(100 + text_embed_size, 256 * 8 * 8)

# 转置卷积层对输入进行上采样
self.deconv1 = nn.ConvTranspose2d(256, 128, 4, 2, 1)
self.deconv2 = nn.ConvTranspose2d(128, 64, 4, 2, 1)
self.deconv3 = nn.ConvTranspose2d(64, 3, 4, 2, 1) # 输出有 3 个 RGB 图像通道

# 激活函数
self.relu = nn.ReLU(True) # ReLU激活函数
self.tanh = nn.Tanh() # 最终输出的 Tanh 激活函数

def forward(self, noise, text_embed):
# 沿通道维度连接噪声和文本嵌入
x = torch.cat((noise, text_embed), dim=1)

# 全连接层,然后重塑为 4D 张量
x = self.fc1(x).view(-1, 256, 8, 8)

# 使用 ReLU 激活通过转置卷积层进行上采样
x = self.relu(self.deconv1(x))
x = self.relu(self.deconv2(x))

# 最后一层使用 Tanh 激活来确保输出值在 -1 和 1 之间(对于图像)
x = self.tanh(self.deconv3(x))

return x
这个 Generator 类负责根据随机噪声和文本嵌入的组合创建视频帧。它的目的是根据给定的文本描述生成逼真的视频帧。该网络从全连接层 ( nn.Linear ) 开始,它将噪声向量和文本嵌入组合成单个特征向量。然后,该向量被重塑并通过一系列转置卷积层 ( nn.ConvTranspose2d ),这些层逐渐将特征图上采样到所需的视频帧大小。

这些层使用 ReLU 激活 ( nn.ReLU ) 实现非线性,最后一层使用 Tanh 激活 ( nn.Tanh ) 将输出缩放到范围 [-1, 1]。因此,生成器将抽象的高维输入转换为连贯的视频帧,直观地表示输入文本。

实施鉴别器层

对生成器层进行编码后,我们需要实现另一半,即鉴别器部分。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class Discriminator(nn.Module):
def __init__(self):
super(Discriminator, self).__init__()

# 用于处理输入图像的卷积层
self.conv1 = nn.Conv2d(3, 64, 4, 2, 1) # 3 个输入通道 (RGB)、64 个输出通道、内核大小 4x4、步长 2、填充 1
self.conv2 = nn.Conv2d(64, 128, 4, 2, 1) # 64 个输入通道,128 个输出通道,内核大小 4x4,步长 2,填充 1
self.conv3 = nn.Conv2d(128, 256, 4, 2, 1) # 128 个输入通道,256 个输出通道,内核大小 4x4,步长 2,填充 1

# 用于分类的全连接层
self.fc1 = nn.Linear(256 * 8 * 8, 1) # 输入大小256x8x8(最后一个卷积的输出大小),输出大小1(二元分类)

# 激活函数
self.leaky_relu = nn.LeakyReLU(0.2, inplace=True) # 负斜率 0.2 的 Leaky ReLU 激活
self.sigmoid = nn.Sigmoid() # 最终输出的 Sigmoid 激活(概率)

def forward(self, input):
# 使用 LeakyReLU 激活将输入传递到卷积层
x = self.leaky_relu(self.conv1(input))
x = self.leaky_relu(self.conv2(x))
x = self.leaky_relu(self.conv3(x))

# 展平卷积层的输出
x = x.view(-1, 256 * 8 * 8)

# 通过带有 Sigmoid 激活的全连接层进行二元分类
x = self.sigmoid(self.fc1(x))

return x
Discriminator 类充当二元分类器,区分真实视频帧和生成的视频帧。其目的是评估视频帧的真实性,从而指导生成器产生更真实的输出。该网络由卷积层 ( nn.Conv2d ) 组成,这些卷积层从输入视频帧中提取分层特征,并使用 Leaky ReLU 激活 ( nn.LeakyReLU ) 添加非线性,同时允许负梯度较小价值观。然后特征图被展平并通过全连接层 ( nn.Linear ),最终形成 sigmoid 激活 ( nn.Sigmoid ),输出一个概率分数,指示帧是真实的还是假的。

通过训练鉴别器对帧进行准确分类,同时训练生成器以创建更有说服力的视频帧,因为它的目的是欺骗鉴别器。

编码训练参数

我们必须设置训练 GAN 的基本组件,例如损失函数、优化器等。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# 检查 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 为文本提示创建简单的词汇表
all_prompts = [prompt for prompt, _, _ in prompts_and_movements] # 从提示和动作列表中提取所有提示
vocab = {word: idx for idx, word in enumerate(set(" ".join(all_prompts).split()))} # 创建词汇词典,其中每个唯一单词都分配有一个索引
vocab_size = len(vocab) # 词汇量的大小
embed_size = 10 # 文本嵌入向量的大小

def encode_text(prompt):
# 使用词汇表将给定提示编码为索引张量
return torch.tensor([vocab[word] for word in prompt.split()])

# 初始化模型、损失函数和优化器
text_embedding = TextEmbedding(vocab_size, embed_size).to(device) # 使用 vocab_size 和 embed_size 初始化 TextEmbedding 模型
netG = Generator(embed_size).to(device) # 使用 embed_size 初始化生成器模型
netD = Discriminator().to(device) # 初始化判别器模型
criterion = nn.BCELoss().to(device) # 二元交叉熵损失函数
optimizerD = optim.Adam(netD.parameters(), lr=0.0002, betas=(0.5, 0.999)) # 判别器的 Adam 优化器
optimizerG = optim.Adam(netG.parameters(), lr=0.0002, betas=(0.5, 0.999)) # Adam 生成器优化器
这是我们必须将代码转换为在 GPU 上运行(如果可用)的部分。我们已编写代码来查找 vocab_size,并且我们对生成器和判别器使用 ADAM 优化器。如果您愿意,可以选择自己的优化器。在这里,我们将学习率设置为较小的值 0.0002,嵌入大小为 10,这与其他可供公众使用的 Hugging Face 模型相比要小得多。

训练循环编码

就像所有其他神经网络一样,我们将以类似的方式编码 GAN 架构训练。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
# 纪元数
num_epochs = 13

# 迭代每个纪元
for epoch in range(num_epochs):
# 迭代每批数据
for i, (data, prompts) in enumerate(dataloader):
# 将真实数据移动到设备
real_data = data.to(device)

# 将提示转换为列表
prompts = [prompt for prompt in prompts]

# 更新鉴别器
netD.zero_grad() # 将鉴别器的梯度归零
batch_size = real_data.size(0) # 获取批量大小
labels = torch.ones(batch_size, 1).to(device) # 为真实数据创建标签 (一)
output = netD(real_data) # 通过判别器前向传递真实数据
lossD_real = criterion(output, labels) # 计算真实数据的损失
lossD_real.backward() # 向后传递计算梯度

# 生成虚假数据
noise = torch.randn(batch_size, 100).to(device) # Generate random noise
text_embeds = torch.stack([text_embedding(encode_text(prompt).to(device)).mean(dim=0) for prompt in prompts]) # 将提示编码为文本嵌入
fake_data = netG(noise, text_embeds) # 从噪声和文本嵌入生成虚假数据
labels = torch.zeros(batch_size, 1).to(device) # 为虚假数据创建标签(零)
output = netD(fake_data.detach()) # 通过鉴别器正向传递假数据(分离以避免梯度流回生成器)
lossD_fake = criterion(output, labels) # 计算虚假数据的损失
lossD_fake.backward() # 向后传递计算梯度
optimizerD.step() # 更新鉴别器参数

# Update Generator
netG.zero_grad() # 将生成器的梯度归零
labels = torch.ones(batch_size, 1).to(device) # 为虚假数据创建标签来欺骗鉴别器
output = netD(fake_data) # 通过鉴别器正向传递虚假数据(现已更新)
lossG = criterion(output, labels) # 根据判别器的响应计算生成器的损失
lossG.backward() # 向后传递计算梯度
optimizerG.step() # 更新发电机参数

# 打印纪元信息
print(f"Epoch [{epoch + 1}/{num_epochs}] Loss D: {lossD_real + lossD_fake}, Loss G: {lossG}")
通过反向传播,我们的损失将针对生成器和鉴别器进行调整。我们使用了 13 个 epoch 进行训练循环。我测试了不同的值,但如果纪元高于此值,结果不会显示出太大差异。而且,遇到过拟合的风险很高。如果我们有一个更多样化的数据集,有更多的运动和形状,我们可以考虑使用更高的纪元,但在这种情况下不行。

当我们运行此代码时,它会开始训练并在每个时期后打印生成器和鉴别器的损失。

保存训练后的模型

训练完成后,我们需要保存训练好的 GAN 架构的判别器和生成器,这只需两行代码即可实现。

1
2
3
4
5
# 将生成器模型的状态字典保存到名为“generator.pth”的文件中
torch.save(netG.state_dict(), 'generator.pth')

# 将鉴别器模型的状态字典保存到名为“discriminator.pth”的文件中
torch.save(netD.state_dict(), 'discriminator.pth')

生成人工智能视频

正如我们所讨论的,我们在未见过的数据上测试模型的方法与我们的训练数据涉及狗捡球和猫追老鼠的示例相当。因此,我们的测试提示可能涉及诸如猫取球或狗追老鼠之类的场景。

在我们的具体情况下,圆圈向上然后向右移动的运动不存在于我们的训练数据中,因此模型不熟悉这种特定运动。然而,它已经接受了其他动作的训练。我们可以使用这个动作作为测试我们训练的模型并观察其性能的提示。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# 根据给定的文本提示生成视频的推理功能
def generate_video(text_prompt, num_frames=10):
# 根据文字提示为生成的视频帧创建目录
os.makedirs(f'generated_video_{text_prompt.replace(" ", "_")}', exist_ok=True)

# 将文本提示编码为文本嵌入张量
text_embed = text_embedding(encode_text(text_prompt).to(device)).mean(dim=0).unsqueeze(0)

# 为视频生成帧
for frame_num in range(num_frames):
# 生成随机噪声
noise = torch.randn(1, 100).to(device)

# 使用生成器网络生成假帧
with torch.no_grad():
fake_frame = netG(noise, text_embed)

# 将生成的假帧保存为图像文件
save_image(fake_frame, f'generated_video_{text_prompt.replace(" ", "_")}/frame_{frame_num}.png')
# 使用generate_video函数和特定的文本提示
generate_video('circle moving up-right')
当我们运行上面的代码时,它将生成一个目录,其中包含我们生成的视频的所有帧。我们需要使用一些代码将所有这些帧合并成一个短视频。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
# 定义包含 PNG 帧的文件夹的路径
folder_path = 'generated_video_circle_moving_up-right'


# 获取文件夹中所有PNG文件的列表
image_files = [f for f in os.listdir(folder_path) if f.endswith('.png')]

# 按名称对图像进行排序(假设它们按顺序编号)
image_files.sort()

# 创建一个列表来存储帧
frames = []

# 读取每个图像并将其附加到帧列表中
for image_file in image_files:
image_path = os.path.join(folder_path, image_file)
frame = cv2.imread(image_path)
frames.append(frame)

# 将帧列表转换为 numpy 数组以便于处理
frames = np.array(frames)

# 定义帧速率(每秒帧数)
fps = 10

# 创建视频编写器对象
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter('generated_video.avi', fourcc, fps, (frames[0].shape[1], frames[0].shape[0]))

# 将每一帧写入视频
for frame in frames:
out.write(frame)

# 释放视频作者
out.release()
确保文件夹路径指向新生成的视频所在的位置。运行此代码后,您的AI视频将已成功创建。让我们看看它是什么样子的。 我以相同的时期数进行了多次训练。在这两种情况下,圆圈都是从底部出现的一半开始。好的部分是我们的模型尝试在这两种情况下执行直立运动。例如,在尝试 1 中,圆圈沿对角线向上移动,然后执行向上运动,而在尝试 2 中,圆圈沿对角线移动,同时缩小尺寸。在这两种情况下,圆圈都没有向左移动或完全消失,这是一个好兆头。

少了什么东西?

我测试了该架构的各个方面,发现训练数据是关键。通过在数据集中包含更多运动和形状,您可以增加可变性并提高模型的性能。由于数据是通过代码生成的,因此生成更多样的数据不会花费太多时间;相反,您可以专注于完善逻辑。

此外,本博客中讨论的 GAN 架构相对简单。您可以通过集成先进技术或使用语言模型嵌入 (LLM) 而不是基本的神经网络嵌入来使其变得更加复杂。此外,调整嵌入大小等参数可以显着影响模型的有效性。

文章转载于:
使用 Python 从头开始​​构建 AI 文本到视频模型