使用Amazon SageMaker使用生成式AI构建个性化的头像
使用Amazon SageMaker构建个性化头像
生成式人工智能已经成为增强和加速创意过程的常用工具，应用于娱乐、广告和平面设计等各个行业。它可以为观众提供更个性化的体验，提高最终产品的整体质量。
生成式人工智能的一个重要优势是为用户创造独特和个性化的体验。例如，流媒体服务使用生成式人工智能生成个性化的电影标题和视觉效果，以增加观众的参与度，并基于用户的观看历史和偏好来生成标题的视觉效果。系统随后生成成千上万个标题艺术作品的变体，并测试它们以确定哪个版本最能吸引用户的注意力。在某些情况下，个性化的电视剧艺术作品相较于没有个性化艺术作品的节目，显著提高了点击率和观看率。
在本文中，我们演示了如何使用稳定扩散等生成式人工智能模型，在亚马逊SageMaker上构建个性化头像解决方案，并同时节省推断成本。该解决方案演示了通过上传自己的10-12张图像，您可以微调一个个性化模型，然后根据任何文本提示生成头像，如下面的屏幕截图所示。尽管此示例生成个性化头像，但您可以将该技术应用于基于特定对象或样式的任何创意艺术生成。
 |
 |
解决方案概述
以下架构图概述了我们头像生成器的端到端解决方案。
本文和我们提供的示例GitHub代码仅关注模型训练和推断编排（前述图中的绿色部分）。您可以参考完整的解决方案架构，并在其基础上进行扩展。
模型训练和推断可以分为以下四个步骤：
- 将图像上传到亚马逊简单存储服务（Amazon S3）。在此步骤中，我们要求您提供至少10张高分辨率的自己的图像。图像越多，结果越好，但训练时间也越长。
- 使用SageMaker异步推断微调稳定扩散2.1基础模型。我们将在本文后面解释使用推断端点进行训练的原理。微调过程从准备图像开始，包括人脸裁剪、背景变化和模型的调整大小。然后，我们使用低秩适应（LoRA）这种对大型语言模型（LLM）进行参数高效微调的技术，对模型进行微调。最后，在后处理中，我们将微调的LoRA权重与推断脚本和配置文件（tar.gz）打包，并将其上传到S3存储桶位置以供SageMaker MME使用。
- 使用GPU在SageMaker MME上托管经过微调的模型。SageMaker将根据每个模型的推断流量动态加载和缓存来自Amazon S3位置的模型。
- 使用经过微调的模型进行推断。在收到微调完成的亚马逊简单通知服务（Amazon SNS）通知后，您可以立即通过在调用MME时提供
target_model参数来使用该模型创建您的头像。
我们将在以下章节中详细解释每个步骤，并演示一些示例代码片段。
准备图像
为了从稳定扩散微调生成您的图像的最佳结果，通常需要为自己提供大量和多样化的照片，包括不同角度、不同表情和不同背景的照片。然而，通过我们的实现，您现在可以通过仅提供至少10张输入图像来获得高质量的结果。我们还添加了自动预处理功能，可以从每张照片中提取您的面部。您只需要以多个角度清晰地捕捉您的外貌特征。包括正面照片、侧面照片和角度之间的照片。您还应该包括带有不同面部表情的照片，如微笑、皱眉和中性表情。具有各种表情的照片将使模型更好地复制您独特的面部特征。输入图像决定了您可以生成的头像的质量。为确保正确完成此过程，我们建议提供一个直观的前端用户界面体验，以指导用户完成图像捕捉和上传过程。
以下是示例自拍图像，采用不同角度和面部表情。
调整稳定的扩散模型
在图像上传到Amazon S3后，我们可以调用SageMaker异步推理端点来启动我们的训练过程。异步端点适用于具有大负载（最高1 GB）和长处理时间（最高1小时）的推理用例。它还提供了一个内置的请求排队机制，并通过Amazon SNS提供任务完成通知机制，以及SageMaker托管的其他本机功能，如自动缩放。
尽管微调不是推理用例，但我们选择在这里使用它，而不是SageMaker训练作业，是因为它具有内置的排队和通知机制，以及托管的自动缩放功能，包括在服务不使用时能够缩减到0个实例的能力。这使我们能够轻松地将微调服务扩展到大量并发用户，并消除了实施和管理额外组件的需求。但是，它的缺点是1 GB的负载和1小时的最大处理时间。在我们的测试中，我们发现在ml.g5.2xlarge实例上，大约10个输入图像在20分钟内可以获得相当不错的结果。但是，对于更大规模的微调作业，建议使用SageMaker训练。
要托管异步端点，我们必须完成几个步骤。首先是定义我们的模型服务器。对于本文，我们使用大型模型推理容器（LMI）。LMI由DJL Serving驱动，它是一个高性能、编程语言无关的模型服务解决方案。我们选择这个选项是因为SageMaker托管的推理容器已经具备了我们所需的许多训练库，如Hugging Face Diffusers和Accelerate。这极大地减少了为我们的微调作业定制容器所需的工作量。
以下代码段显示了我们在示例中使用的LMI容器的版本：
inference_image_uri = (
f"763104351884.dkr.ecr.{region}.amazonaws.com/djl-inference:0.21.0-deepspeed0.8.3-cu117"
)
print(f"将要使用的镜像是 ---- > {inference_image_uri}")除此之外，我们还需要一个serving.properties文件来配置服务属性，包括要使用的推理引擎、模型工件的位置和动态批处理。最后，我们必须有一个model.py文件，它将模型加载到推理引擎中，并准备模型的数据输入和输出。在我们的示例中，我们使用model.py文件来启动微调作业，我们将在后面的章节中详细解释。serving.properties和model.py文件都提供在training_service文件夹中。
在定义我们的模型服务器之后，下一步是创建一个端点配置，定义我们的异步推理将如何提供服务。对于我们的示例，我们只定义了最大并发调用限制和输出S3位置。使用ml.g5.2xlarge实例，我们发现我们能够同时微调最多两个模型，而不会遇到内存不足（OOM）异常，因此我们将max_concurrent_invocations_per_instance设置为2。如果使用不同的调优参数或较小的实例类型，可能需要调整此数字。我们建议最初将其设置为1，并在Amazon CloudWatch中监视GPU内存使用情况。
# 创建异步端点配置
async_config = AsyncInferenceConfig(
output_path=f"s3://{bucket}/{s3_prefix}/async_inference/output" , # 结果的存储位置
max_concurrent_invocations_per_instance=2,
notification_config={
"SuccessTopic": "...",
"ErrorTopic": "...",
}, # 通知配置
)最后，我们创建一个SageMaker模型，将容器信息、模型文件和AWS Identity and Access Management（IAM）角色打包到一个对象中。使用我们之前定义的端点配置部署模型：
model = Model(
image_uri=image_uri,
model_data=model_data,
role=role,
env=env
)
model.deploy(
initial_instance_count=1,
instance_type=instance_type,
endpoint_name=endpoint_name,
async_inference_config=async_inference_config
)
predictor = sagemaker.Predictor(
endpoint_name=endpoint_name,
sagemaker_session=sagemaker_session
)当端点准备好后，我们使用以下示例代码调用异步端点并开始微调过程：
sm_runtime = boto3.client("sagemaker-runtime")
input_s3_loc = sess.upload_data("data/jw.tar.gz", bucket, s3_prefix)
response = sm_runtime.invoke_endpoint_async(
EndpointName=sd_tuning.endpoint_name,
InputLocation=input_s3_loc)有关SageMaker上的LMI的更多详细信息，请参阅使用DJLServing和DeepSpeed模型并行推理在Amazon SageMaker上部署大型模型。
调用后，异步端点开始排队我们的微调作业。每个作业都经过以下步骤：准备图像，执行Dreambooth和LoRA微调，并准备模型工件。让我们深入了解微调过程。
准备图像
正如我们之前提到的，输入图像的质量直接影响微调模型的质量。对于头像用例，我们希望模型专注于面部特征。我们不要求用户提供精心策划的确切大小和内容的图像，而是使用计算机视觉技术实现一个预处理步骤来减轻这个负担。在预处理步骤中，我们首先使用面部检测模型来隔离每个图像中最大的面部。然后我们将图像裁剪和填充到我们模型所需的512 x 512像素大小。最后，我们将面部从背景中分割出来，并添加随机背景变化。这有助于突出面部特征，使我们的模型能够从面部本身而不是背景中学习。以下图像说明了这个过程中的三个步骤。
 |
 |
 |
| 步骤1：使用计算机视觉进行面部检测 | 步骤2：将图像裁剪和填充到512 x 512像素 | 步骤3（可选）：分割并添加背景变化 |
Dreambooth和LoRA微调
对于微调，我们结合了Dreambooth和LoRA的技术。Dreambooth允许您个性化您的稳定扩散模型，使用唯一标识符将主题嵌入模型的输出领域，并扩展模型的语言视觉词典。它使用一种称为先验保持的方法来保持模型对主题类别（在本例中为人物）的语义知识，并使用类别中的其他对象来改善最终的图像输出。这就是Dreambooth可以通过仅使用少量主题输入图像实现高质量结果的原因。
以下代码片段显示了我们avatar解决方案中trainer.py类的输入。请注意，我们选择了<<TOK>>作为唯一标识符。这是有意为之的，以避免选择可能已存在于模型字典中的名称。如果名称已经存在，模型必须取消学习然后重新学习主题，这可能导致微调结果不佳。主题类别设置为“a photo of person”，这通过首先生成人们的照片作为微调过程中的附加输入来启用先验保持。这将有助于减少过拟合，因为模型试图使用先验保持方法保留对人员的先前知识。
status = trn.run(base_model="stabilityai/stable-diffusion-2-1-base",
resolution=512,
n_steps=1000,
concept_prompt="照片是<<TOK>>", # << 学科的唯一标识符
learning_rate=1e-4,
gradient_accumulation=1,
fp16=True,
use_8bit_adam=True,
gradient_checkpointing=True,
train_text_encoder=True,
with_prior_preservation=True,
prior_loss_weight=1.0,
class_prompt="一个人的照片", # << 学科类别
num_class_images=50,
class_data_dir=class_data_dir,
lora_r=128,
lora_alpha=1,
lora_bias="none",
lora_dropout=0.05,
lora_text_encoder_r=64,
lora_text_encoder_alpha=1,
lora_text_encoder_bias="none",
lora_text_encoder_dropout=0.05
)在配置中启用了多个节省内存的选项，包括 fp16、use_8bit_adam 和渐变累积。这将内存占用降低到不到 12 GB，允许在 ml.g5.2xlarge 实例上同时对两个模型进行微调。
LoRA 是一种高效的 LLMs 微调技术，它冻结大部分权重，并将一个小的适配器网络附加到预训练的 LLM 的特定层上，以实现更快的训练和优化的存储。对于 Stable Diffusion，适配器附加到推理流水线的文本编码器和 U-Net 组件上。文本编码器将输入提示转换为 U-Net 模型理解的潜在空间，而 U-Net 模型使用这个潜在含义在后续的扩散过程中生成图像。微调的输出仅是 text_encoder 和 U-Net 适配器的权重。在推理时，这些权重可以重新连接到基础 Stable Diffusion 模型上，以重现微调的结果。
下面的图是原始作者提供的 LoRA 微调的详细图解：Cheng-Han Chiang, Yung-Sung Chuang, Hung-yi Lee, “AACL_2022_tutorial_PLMs,” 2022
 |
 |
通过结合这两种方法，我们能够生成一个个性化模型，同时调整的参数数量减少了一个数量级。这导致训练时间大大缩短，GPU 利用率降低。此外，适配器权重的存储量仅为 70 MB，而完整的 Stable Diffusion 模型为 6 GB，大小减小了 99%。
准备模型文件
在微调完成后，后处理步骤将使用其余的模型服务文件一起将 LoRA 权重打包成 TAR 文件，用于 NVIDIA Triton。我们使用 Python 后端，这意味着 Triton 配置文件和用于推理的 Python 脚本是必需的。注意，Python 脚本的名称必须为 model.py。最终的模型 TAR 文件应具有以下文件结构：
|--sd_lora
|--config.pbtxt
|--1\
|--model.py
|--output #LoRA 权重
|--text_encoder\
|--unet\
|--train.sh使用具备 GPU 的 SageMaker MME 托管微调的模型
在完成模型微调后，我们使用 SageMaker MME 托管个性化 Stable Diffusion 模型。SageMaker MME 是一种强大的部署功能，允许在单个容器后面托管多个模型，并在单个端点上自动管理流量和路由，以优化资源利用率、节省成本，并减少管理成千上万个端点的操作负担。在我们的例子中，我们在 GPU 实例上运行，SageMaker MME 支持使用 Triton Server 的 GPU，这使您可以在单个 GPU 设备上运行多个模型，并利用加速计算的优势。有关如何在 SageMaker MME 上托管 Stable Diffusion 的详细信息，请参阅使用 Stable Diffusion 模型创建高质量图像，并在 Amazon SageMaker 上以成本效益方式部署。
对于我们的示例，我们在冷启动情况下进行了额外的优化，以更快地加载微调模型。这是由于LoRA的适配器设计所能实现的。由于基础模型权重和Conda环境对于所有微调模型来说是相同的，我们可以通过将这些共享资源预加载到托管容器中来共享它们。这样，只需要在第一次调用后从Amazon S3动态加载Triton配置文件、Python后端（model.py）和LoRA适配器权重。下图提供了一个并排比较的示意图。
这将大大减少模型TAR文件的大小，从大约6 GB减少到70 MB，因此加载和解压速度更快。在我们的示例中，我们创建了一个实用的Python后端模型models/model_setup来进行预加载。该脚本简单地将基础稳定扩散模型和Conda环境从Amazon S3复制到一个共享位置，以供所有微调模型共享。以下是执行此任务的代码片段：
def initialize(self, args):
#conda环境设置
self.conda_pack_path = Path(args['model_repository']) / "sd_env.tar.gz"
self.conda_target_path = Path("/tmp/conda")
self.conda_env_path = self.conda_target_path / "sd_env.tar.gz"
if not self.conda_env_path.exists():
self.conda_env_path.parent.mkdir(parents=True, exist_ok=True)
shutil.copy(self.conda_pack_path, self.conda_env_path)
#基础扩散模型设置
self.base_model_path = Path(args['model_repository']) / "stable_diff.tar.gz"
try:
with tarfile.open(self.base_model_path) as tar:
tar.extractall('/tmp')
self.response_message = "模型环境设置成功。"
except Exception as e:
# 打印异常消息
print(f"捕获到异常：{e}")
self.response_message = f"捕获到异常：{e}"然后，每个微调模型将指向容器上的共享位置。Conda环境在config.pbtxt中进行引用。
name: "pipeline_0"
backend: "python"
max_batch_size: 1
...
parameters: {
key: "EXECUTION_ENV_PATH",
value: {string_value: "/tmp/conda/sd_env.tar.gz"}
}稳定扩散基础模型从每个model.py文件的initialize()函数中加载。然后，我们将个性化的LoRA权重应用于unet和text_encoder模型，以重新生成每个微调模型：
...
class TritonPythonModel:
def initialize(self, args):
self.output_dtype = pb_utils.triton_string_to_numpy(
pb_utils.get_output_config_by_name(json.loads(args["model_config"]),
"generated_image")["data_type"])
self.model_dir = args['model_repository']
device='cuda'
self.pipe = StableDiffusionPipeline.from_pretrained('/tmp/stable_diff',
torch_dtype=torch.float16,
revision="fp16").to(device)
# 加载LoRA权重
self.pipe.unet = PeftModel.from_pretrained(self.pipe.unet, unet_sub_dir)
if os.path.exists(text_encoder_sub_dir):
self.pipe.text_encoder = PeftModel.from_pretrained(self.pipe.text_encoder, text_encoder_sub_dir)使用微调模型进行推断
现在，我们可以通过调用MME端点来尝试我们的微调模型。我们在示例中展示的输入参数包括prompt、negative_prompt和gen_args，如下面的代码片段所示。我们在字典中设置每个输入项的数据类型和形状，并将它们转换为JSON字符串。最后，将字符串载荷和TargetModel传入请求中生成您的头像图片。
import random
prompt = """<<TOK>> epic portrait, zoomed out, blurred background cityscape, bokeh,
perfect symmetry, by artgem, artstation ,concept art,cinematic lighting, highly
detailed, octane, concept art, sharp focus, rockstar games, post processing,
picture of the day, ambient lighting, epic composition"""
negative_prompt = """
beard, goatee, ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, blurry, bad anatomy, blurred,
watermark, grainy, signature, cut off, draft, amateur, multiple, gross, weird, uneven, furnishing, decorating, decoration, furniture, text, poor, low, basic, worst, juvenile,
unprofessional, failure, crayon, oil, label, thousand hands
"""
seed = random.randint(1, 1000000000)
gen_args = json.dumps(dict(num_inference_steps=50, guidance_scale=7, seed=seed))
inputs = dict(prompt = prompt,
negative_prompt = negative_prompt,
gen_args = gen_args)
payload = {
"inputs":
[{"name": name, "shape": [1,1], "datatype": "BYTES", "data": [data]} for name, data in inputs.items()]
}
response = sm_runtime.invoke_endpoint(
EndpointName=endpoint_name,
ContentType="application/octet-stream",
Body=json.dumps(payload),
TargetModel="sd_lora.tar.gz",
)
output = json.loads(response["Body"].read().decode("utf8"))["outputs"]
original_image = decode_image(output[0]["data"][0])
original_image清理
按照笔记本中清理部分的说明删除作为本文一部分提供的资源，以避免不必要的费用。有关推理实例费用的详细信息，请参考Amazon SageMaker定价。
结论
在本文中，我们演示了如何使用SageMaker上的稳定扩散创建个性化头像解决方案。通过使用仅几张图片对预训练模型进行微调，我们可以生成反映每个用户个性和个人特点的头像。这只是使用生成式人工智能为用户创建定制和独特体验的众多示例之一。可能性是无限的，我们鼓励您尝试这项技术并探索其增强创作过程的潜力。我们希望本文对您有所启发和指导。我们鼓励您尝试示例，并使用社交平台上的标签#sagemaker #mme #genai与我们分享您的作品。我们很想看到您的创作。
除了稳定扩散，Amazon SageMaker JumpStart还提供了许多其他生成式人工智能模型。请参考《入门Amazon SageMaker JumpStart》以了解它们的能力。
