目录
介绍
如今，媒体和娱乐公司拥有大量在高分辨率技术出现之前以低分辨率或标准分辨率创建的传统媒体内容，例如电影、电视剧、音乐视频和体育精彩片段。鉴于当今显示技术的进步，观众普遍要求视频内容以高清格式（如高清或 4K）交付，以在较大屏幕上获得增强的观看体验。通过以更高分辨率（高清、4K 或更高）提供传统媒体资产，提高增加收入的机会。
深度学习和生成式 AI 模型的最新进展使得通过超分辨率等技术大幅提高图像质量成为可能。超分辨率可以增加图像的像素密度，使其显得更清晰，并包含更多细节。这些 AI 经过数百万张图像的训练，可以通过单次传递大规模提高图像质量，而无需额外步骤，如典型超分辨率工作流程中常见的去噪滤波器。与简单根据固定数学公式估计新像素的传统技术（如双线性或双三次插值）不同，生成式 AI 技术，如生成对抗网络（GAN）、视觉转换器和扩散模型（如 Stable Diffusion）学会智能地填充细节、纹理和边缘，同时保留细节并减少锯齿状边缘等伪影。
在本博客文章中，我们提出了一种端到端解决方案，结合使用 Real-ESRGANRReal-ESRGANeal-ESRGANReal-ESRGAN 和 SwinIRSwinIRSwinIR 以及 AWS 服务来编排工作流程，可以将低分辨率视频作为输入来生成 4 倍分辨率的视频。例如，我们可以在不到 60 分钟的时间内将一小时动画电影的分辨率提高 4 倍。该解决方案可用于自动化大规模超分辨率处理媒体内容的过程。
架构
该解决方案使用本地 AWS 服务构建，重点关注以下主要组件：
- 允许用户上传标准分辨率视频的用户界面
- 执行视频分析和帧提取的事件驱动流程
- 通过调用使用 SageMaker 推理工具包构建的 API 端点，并行应用 AI 图像超分辨率
- 使用上采样图像和从原始媒体提取的音频进行最终视频构建
考虑到可能从输入产生大量帧的可能性，视频超分辨率工作流程的设计以可扩展性和性能为首要目标。在本博客文章中，我们采用 HPC 方法来解决这些挑战。
首先，我们使用 AWS ParallelClusterAWSAWS ParallelClusterParallelClusterAWS ParallelCluster 来支持端到端视频超分辨率工作流程的计算基础架构。AWS ParallelCluster 是一项 AWS 服务，允许用户在云中快速构建 HPC 应用程序。对于存储层，我们使用 Amazon FSx For LustreAmazon FSx For LustreAmazon FSx For Lustre 作为跨所有计算节点挂载的共享文件系统。Amazon FSx for Lustre 提供亚毫秒级延迟、高达数百 GBps 的吞吐量和高达数百万 IOPS，优化了视频帧分析和处理的 I/O 吞吐量。
对于 AI 视频超分辨率任务，我们利用使用 Amazon SageMaker 推理工具包SageMaker SageMaker 推理工具包构建的自定义 docker 容器，这是一个在基于 GPU 的计算节点上的开源机器学习推理框架。所有 Amazon SageMaker 端点都与给定的 FSx for Lustre 文件系统集成，以实现高吞吐量和低延迟的并行图像生成。最后，我们提供了一个用户界面，允许用户上传视频，从而触发自动化端到端视频超分辨率过程的工作流程。
下图更详细地描述了我们的工作流程：
 |
图 1. 视频超分辨率架构图
解决方案演练
在下一节中，我们将在较高层次上演练部署视频超分辨率工作流程。
先决条件
- 安装了 docker 的开发环境（注意：此解决方案已在 SageMaker notebook上使用 ml.g5.2xlarge 实例类型进行了测试）
- 具有访问 Amazon Elastic Container Registry（Amazon ECR）的权限的 AWS Identity and Access Management（IAM），可创建存储库并将 docker 映像推送到注册表
- 至少有两个私有子网和一个公共子网的 Amazon Virtual Private Cloud（Amazon VPC）
- 用于存储媒体内容和工作流程所需相关脚本的 Amazon Simple Storage Service（Amazon S3）存储桶
- 通过 SSH 访问 ParallelCluster 头节点的密钥对：Create a key pair - AWS ParallelCluster
模型部署
为了在上采样图像中保持高保真度，我们测试了几种图像上采样模型并观察了结果。根据我们的评估，我们发现 Real-ESRGAN 和 SwinIR 在去除伪影和恢复纹理细节方面都取得了出色的质量。关于模型评估和技术，我们建议阅读与 Real-ESRGAN 论文和 SwinIR 论文相关的内容。
Real-ESRGAN 模型
通过基于图像结构质量和模型延迟的评估，我们发现 Real-ESRGAN 模型能够以低延迟实现图像的优质超分辨率，尤其是对于动画和动漫媒体内容。与其他生成式 AI 解决方案（如 Stable Diffusion）相比，我们发现该模型在保持高保真度的同时，能够实现更好的生成一致性，而不会引入伪影。我们的解决方案利用此模型对动画视频进行上采样。我们使用 SageMaker Pytorch 推理工具包构建了一个 docker 容器，它将图像帧作为输入。处理动画帧的推理代码片段如下所示：
def model_fn(model_dir):
realesr_gan_anime_model_path = os.path.join(model_dir, realesr_gan_anime_video_model_name)
#加载RealESRGan模型
realesr_gan_anime_video_model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=netscale, act_type='prelu')
real_esr_gan_anime_video_upsampler = RealESRGANer(
scale=netscale,
model_path=realesr_gan_anime_model_path,
dni_weight=dni_weight,
model=realesr_gan_anime_video_model,
tile=0,
tile_pad=10,
pre_pad=0,
half=True,
gpu_id=0)
model = {}
model['realesr_gan_anime'] = real_esr_gan_anime_video_upsampler
return model
def predict_fn(input_data, model):
# 前向传递
input_file_path = input_data['input_file_path']
output_file_path = input_data['output_file_path']
imgname, extension = os.path.splitext(os.path.basename(input_file_path))
img = cv2.imread(input_file_path, cv2.IMREAD_UNCHANGED)
try:
if ('is_anime' in input_data) and ((input_data['is_anime'].lower() == "yes") or (input_data['is_anime'].lower() == "true")):
upsampler = model['realesr_gan_anime']
else:
raise RuntimeError("不支持的模型,仅支持动画上采样。")
output, _ = upsampler.enhance(img, outscale=outscale)
except RuntimeError as error:
print('Error', error)
raise error
cv2.imwrite(output_file_path, output)
return { "status" : 200, "output_file_path" : output_file_path, "job_id" : job_id, "batch_id" : batch_id }
Swin2SR 模型
我们进行视频超分辨率的方法是为给定类型的视频提供最佳上采样模型。通过对图像质量和结构一致性的评估指标，我们采用了第二种基于视觉转换器的生成模型 Swin2SR，用于对真实图像进行上采样。该模型非常适合对包括音乐视频、电影、电视剧和体育剪辑在内的内容进行上采样。同样，我们使用 SageMaker Pytorch 推理工具包构建了一个 docker 容器来执行实时推理。处理真实图像帧的推理代码片段如下所示：
def model_fn(model_dir):
# 加载SwinIR模型
model_path = os.path.join(model_dir, model_name)
model = define_model(model_path, "real_sr", scale_factor)
model = model.to(device)
model.eval()
return model
def predict_fn(input_data, model):
# 前向传递
input_file_path = input_data['input_file_path']
output_file_path = input_data['output_file_path']
img_lq = cv2.imread(input_file_path, cv2.IMREAD_COLOR).astype(np.float32) / 255.
img_lq = np.transpose(img_lq if img_lq.shape[2] == 1 else img_lq[:, :, [2, 1, 0]], (2, 0, 1)) # HCW-BGR to CHW-RGB
img_lq = torch.from_numpy(img_lq).float().unsqueeze(0).to(device) # CHW-RGB to NCHW-RGB
with torch.no_grad():
# 将输入图像填充为窗口大小的倍数
_, _, h_old, w_old = img_lq.size()
h_pad = (h_old // window_size + 1) * window_size - h_old
w_pad = (w_old // window_size + 1) * window_size - w_old
img_lq = torch.cat([img_lq, torch.flip(img_lq, [2])], 2)[:, :, :h_old + h_pad, :]
img_lq = torch.cat([img_lq, torch.flip(img_lq, [3])], 3)[:, :, :, :w_old + w_pad]
output = model(img_lq)
output = output[..., :h_old * scale_factor, :w_old * scale_factor]
output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
if output.ndim == 3:
output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0)) # CHW-RGB to HCW-BGR
output = (output * 255.0).round().astype(np.uint8) # float32 to uint8
cv2.imwrite(f'{output_file_path}', output)
torch.cuda.empty_cache()
return { "status" : 200, "output_file_path" : output_file_path, "job_id" : job_id, "batch_id" : batch_id }
构建、发布这些模型到 ECR 注册表以及部署模型进行推理的完整代码示例可以在 GitHub 存储库中找到。
AWS ParallelCluster
您可以使用几种方法在 AWS 中部署和管理 ParallelCluster。这些包括 ParallelCluster CLI、ParallelCluster API、ParallelCluster UI、Python API 或 Cloudformation 自定义资源。在本博客文章中，我们利用 ParallelCluster CLI 来管理具有 Slurm 调度程序的 ParallelCluster 的生命周期，用于我们的工作流程。Slurm 是一个用于管理和调度 Linux 集群的系统。它是开源的、容错的和可扩展的，适用于各种规模的集群。从高层次来看，我们的 ParallelCluster 由以下组件组成：
- 一个编排集群扩展和将新节点附加到 Slurm 调度程序的头节点
- 一个Slurm调度程序，用于在集群中的计算节点上启动和监控作业。它还管理作业队列，并根据运行时的配置和容量将作业调度到适当的节点
- 执行视频分析、帧提取和视频编码任务的计算节点。这组计算节点是 CPU 密集型的
- 执行图像超分辨率任务的计算节点。这组计算节点是 GPU 内存密集型的
- 使用 FSx for Lustre 作为存储层的共享文件系统，用于读写图像帧和视频前/后处理
以下 yaml 模板文件显示了集群配置（cluster config）的示例：
Region: <AWS区域>
Image:
Os: alinux2
HeadNode:
InstanceType: t3.large
Networking:
SubnetId: {
{VPC_PUBLIC_SUBNET_ID}}
Ssh:
KeyName: {
{SSH_KEY_PAIR_NAME}}
CustomActions:
OnNodeConfigured:
Script: {
{ON_NODE_STARTED_SCRIPT_HEAD_NODE}}
Iam:
S3Access:
- BucketName: {
{S3_BUCKET}}
EnableWriteAccess: False
AdditionalIamPolicies:
- Policy: arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore
Scheduling:
Scheduler: slurm
SlurmSettings:
Dns:
UseEc2Hostnames: true
SlurmQueues:
- Name: q1
ComputeResources:
- Name: genai-api-service
Instances:
- InstanceType: g5.2xlarge
MinCount: 0
MaxCount: 1
Networking:
SubnetIds:
- {
{VPC_PRIVATE_SUBNET_ID}}
CustomActions:
OnNodeConfigured:
Script: {
{ON_NODE_STARTED_SCRIPT_GPU_COMPUTE_NODE}}
Iam:
S3Access:
- BucketName: {
{S3_BUCKET}}
EnableWriteAccess: False
AdditionalIamPolicies:
- Policy: arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore
- Policy: arn:aws:iam::aws:policy/EC2InstanceProfileForImageBuilderECRContainerBuilds
Image:
CustomAmi: {
{GPU_COMPUTE_NODE_CUSTOM_AMI}}
- Name: q2
ComputeResources:
- Name: video-processing
Instances:
- InstanceType: c6i.4xlarge
MinCount: 1
MaxCount: 1
Networking:
SubnetIds:
- {
{VPC_PRIVATE_SUBNET_ID}}
CustomActions:
OnNodeConfigured:
Script: {
{ON_NODE_STARTED_SCRIPT_CPU_COMPUTE_NODE}}
Iam:
S3Access:
- BucketName: {
{S3_BUCKET}}
EnableWriteAccess:True
AdditionalIamPolicies:
- Policy: arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore
- Policy: arn:aws:iam::aws:policy/EC2InstanceProfileForImageBuilderECRContainerBuilds
Image:
CustomAmi: {
{CPU_COMPUTE_NODE_CUSTOM_AMI}}
SharedStorage:
- MountDir: /fsx
Name: fsx
StorageType: FsxLustre
FsxLustreSettings:
StorageCapacity: 7200
DeploymentType: PERSISTENT_2
PerUnitStorageThroughput: 1000
DeletionPolicy: Delete
要使用上述配置创建 ParallelCluster，请运行以下 CLI 命令：
pcluster create-cluster --cluster-name <集群名称> --cluster-configuration <集群配置 (cluster config) yaml>集群设置过程大约需要 30 分钟。您可以通过运行以下 CLI 命令来检查集群的状态：
pcluster describe-cluster --cluster-name <集群名称>以下消息片段显示集群已成功创建。
...
"cloudFormationStackStatus": "CREATE_COMPLETE",
"clusterName": "test-cluster",
"computeFleetStatus": "STARTED",
"cloudformationStackArn": "arn:aws:cloudformation:us-east-1:123456789012:stack/test-cluster/9f2ea6a0-385e-11ee-a631-0aaaad7cc485",
"lastUpdatedTime": "2023-10-06T00:28:44.591Z",
"region": "us-east-1",
"clusterStatus": "CREATE_COMPLETE",
"scheduler": {
"type": "slurm"
}
}
自定义 Bootstrap 操作
在配置头节点和计算节点的过程中，ParallelCluster 可以在节点启动后或节点配置正确完成后执行任意代码。该代码以位于给定 S3 位置的 shell 脚本的形式提供，如上所示在集群配置 yaml 中进行配置。在这里，我们利用自定义 bootstrap 脚本来配置头节点和计算节点。对于头节点，我们安装了在 ParallelCluster 调度程序调度作业时将被调用的所需脚本。一个 bootstrap 脚本示例如下所示：
#!/bin/bash
echo "downloading bootstrap scripts from s3" >> /tmp/.node-status
aws s3 cp s3://{
{S3_BUCKET_NAME}}/bootstrap/bootstrap-headnode.tar.gz /tmp/genai-video-super-resolution/
echo "completed bootstrap scripts from s3" >> /tmp/.node-status
cd /tmp/genai-video-super-resolution
tar -xvzf bootstrap-headnode.tar.gz -C /home/ec2-user
与头节点类似，我们使用 shell 脚本来 bootstrap GPU 计算节点，以创建其他资源，例如自定义 Amazon CloudWatch 指标、通过 docker-compose 命令的本地 SageMaker 端点，并在作业可以通过 API 调用进行图像超分辨率之前，对这些端点进行就绪性测试。GPU 计算节点的 bootstrap 脚本示例如下所示：
#!/bin/bash
SECONDS=0
sudo usermod -aG docker $USER
sudo service docker restart
echo "downloading bootstrap scripts from s3" >> /tmp/.node-status
aws s3 cp s3://{
{S3_BUCKET_NAME}}/bootstrap/bootstrap.tar.gz /tmp/genai-video-super-resolution-pcluster/
echo "completed bootstrap scripts from s3" >> /tmp/.node-status
cd /tmp/genai-video-super-resolution-pcluster
tar -xvzf bootstrap.tar.gz
# 将 GPU 指标添加到 cloudwatch 代理
cp cloudwatch/cloudwatch-agent-gpu-config.json /opt/aws/amazon-cloudwatch-agent/bin/config.json
/opt/aws/amazon-cloudwatch-agent/bin/amazon-cloudwatch-agent-ctl -a fetch-config -s -m ec2 -c file:/opt/aws/amazon-cloudwatch-agent/bin/config.json &>> /tmp/.node-status
# 拉取最新镜像
aws ecr get-login-password —region us-east-1 | docker login —username AWS —password-stdin 123456789012.dkr.ecr.us-east-1.amazonaws.com
echo "login successful" >> /tmp/.node-status
echo "install docker-compose" >> /tmp/.node-status
sudo curl -L https://github.com/docker/compose/releases/latest/download/docker-compose-$(uname -s)-$(uname -m) -o /usr/bin/docker-compose
sudo chmod +x /usr/bin/docker-compose
echo "docker-compose installed successfully." >> /tmp/.node-status
docker-compose -f docker/docker-compose.yaml up -d &>> /tmp/.node-status
success=0
cd test
while [ $success -ne 1 ]
do
./test.sh realesrgan &>> /tmp/.node-status
if [ $? -eq 0 ]
then
echo "done initialize realesrgan" >> /tmp/.node-status
success=1
else
echo "failed to initialize realesrgan, keep trying.." >> /tmp/.node-status
sleep 15
fi
done
success=0
while [ $success -ne 1 ]
do
./test.sh swinir2 &>> /tmp/.node-status
if [ $? -eq 0 ]
then
echo "done initialize swinir2" >> /tmp/.node-status
success=1
else
echo "failed to initialize swinir, keep trying.." >> /tmp/.node-status
sleep 15
fi
done
cd ~
echo "bootstrap took $SECONDS seconds" >> /tmp/.node-status
除了图像超分辨率任务外，我们还为视频编码和解码任务配置了一组 CPU 计算节点。我们使用与用于 bootstrap GPU 计算节点的脚本类似的 bootstrap 方法安装依赖项。完整细节可以在 GitHub 存储库GitHub 存储库GitHub 存储库中找到。
工作流程编排
视频超分辨率工作流程实现了事件驱动架构。当用户将视频上传到给定的 S3 位置时，工作流程编排由 AWS Lambda 函数触发，通过头节点调度一系列批处理作业。下图显示了使用 ParallelCluster 中的 Slurm 调度程序执行超分辨率过程的顺序步骤：
 |
图 2. 视频超分辨率工作流程的顺序步骤
上述每个过程中涉及的主要脚本总结如下：
- create_extract_job_full.sh – 作为 SSM RunCommand从 Lambda 函数触发，以启动视频超分辨率工作流程。
- extract_frame_audio.sh – 从原始视频中提取图像和音频组件。输出写入到 FSx for Lustre 文件系统中指定的目录。提取过程完成后，将触发 Slurm 批处理数组作业，对每一帧执行 AI 超分辨率推理。此步骤在 ParallelCluster 中的 CPU 计算节点上运行。
- frame-super-resolution-array.sh – 在帧提取步骤完成后调用此脚本。它调用相应的 SageMaker 端点进行图像超分辨率任务。此作业安排在 GPU 计算节点上。根据可用的 GPU 资源，可以在同一计算节点上调度多个作业，以利用可用资源实现更好的吞吐量。此外，这些作业会自动跨计算节点调度，以实现水平扩展。集群扩展到的节点数量在 ParallelCluster 配置文件中配置。工作流程的吞吐量与可用于工作流程的计算节点数量成线性关系。
- encode-new-movie.sh – 在所有帧成功超分辨率后，该脚本对新电影进行编码。该脚本使用 ffmpeg 对图像和音频进行编码，安排在 CPU 计算节点上。新视频将上传到指定的 S3 位置。此作业完成视频超分辨率工作流程。
汇总
为了演示端到端视频超分辨率工作流程，我们构建了一个用户界面组件，使您更容易开始使用该解决方案。该 UI 是作为 Gradio 应用程序构建的，这是一个用于机器学习的开源 Python UI 框架。所有依赖项都打包在一个 Docker 镜像中。Docker 容器可以在本地运行，或者作为 Amazon ECS Fargate 无服务器解决方案进行调度。
这是视频超分辨率 UI 的屏幕截图。以下视频段展示了我们的视频超分辨率结果。标准分辨率视频（左侧）和使用我们的超分辨率工作流程进行超分辨率的视频（右侧）。
示例
 |