当前位置：首页 > news >正文

在亚马逊云科技上对Stable Diffusion模型提示词、输出图像内容进行安全审核

news 来源：原创 2024/9/20 19:42:09

项目简介：

小李哥将继续每天介绍一个基于亚马逊云科技AWS云计算平台的全球前沿AI技术解决方案，帮助大家快速了解国际上最热门的云计算平台亚马逊云科技AWS AI最佳实践，并应用到自己的日常工作里。

本次介绍的是如何在亚马逊云科技机器学习托管服务SageMaker上部署开源大模型Stable Diffusion，利用亚马逊云科技Comprehend对模型输入提示词进行有害性检测，并利用亚马逊云科技Rekognition服务对生成图像内容进行有害性检测，构建负责任的AI防止大模型被滥用。本架构设计全部采用了云原生Serverless架构，提供可扩展和安全的AI解决方案。本方案的解决方案架构图如下：

方案所需基础知识

什么是 Amazon SageMaker？

Amazon SageMaker 是亚马逊云科技提供的一站式机器学习服务，帮助开发者和数据科学家轻松构建、训练和部署机器学习模型。SageMaker 提供了全面的工具，从数据准备、模型训练到部署和监控，覆盖了机器学习项目的全生命周期。通过 SageMaker，用户可以加速机器学习模型的开发和上线，并确保模型在生产环境中的稳定性和性能。

什么是 Amazon Comprehend？

Amazon Comprehend 是亚马逊云科技提供的一项自然语言处理（NLP）服务，能够自动从文本中提取有价值的信息。通过机器学习技术，Comprehend 可以识别文本中的实体、情感、关键词、语言、主题等，帮助企业更好地理解和分析大量非结构化数据。它适用于客户反馈分析、内容分类、文档处理等场景，使得信息挖掘和数据洞察变得更加简单和高效。

什么是 Amazon Rekognition？

Amazon Rekognition 是亚马逊云科技提供的一项图像和视频分析服务。它使用深度学习技术来检测、识别和分析图像中的对象、场景、面部表情、文字等。Rekognition 可以应用于多种场景，如面部识别、内容审核、对象检测和人群统计等，帮助企业自动化处理图像和视频数据，提升效率并增强安全性。

什么是 Stable Diffusion？

Stable Diffusion 是一种先进的生成式 AI 模型，专门用于生成高质量的图像。通过扩散模型技术，Stable Diffusion 能够将简单的文本描述转化为逼真的图像。这个模型具有强大的生成能力，可以应用于艺术创作、广告设计、游戏开发等领域，为用户提供丰富的视觉内容生成工具。

为什么要对 Stable Diffusion 输入输出内容进行安全审核？

防止不当内容生成：

Stable Diffusion 可以根据输入的文本生成图像，但如果输入的文本内容不当或恶意，可能会生成带有敏感、违法或不道德内容的图像。对输入输出内容进行审核，能够有效防止此类内容的生成和传播，确保模型的使用符合道德和法律标准。

保护用户隐私：

在生成图像时，可能涉及到用户的私人信息或敏感数据。通过审核输入输出内容，可以确保这些信息不会被意外泄露或滥用，保护用户的隐私权。

遵守法律法规：

各国对生成和传播图像内容有不同的法律规定。通过对内容进行审核，企业可以确保生成的图像符合所在国家或地区的法律法规，避免法律风险。

维护品牌声誉：

对内容进行安全审核，有助于防止不符合公司价值观或可能损害品牌声誉的内容生成，从而维护品牌的形象和公众信任。

本方案包括的内容

1. 在SageMaker上部署开源大模型Stable Diffusion

2. 在SageMaker上调用Stable Diffusion模型API生成图片

3. 将Stable Diffusion模型API节点集成到云端应用上

4. 评估大模型输入问题的有害性

5. 对大模型输出图片进行安全审核

项目搭建具体步骤：

1. 打开亚马逊云科技控制台，进入Amazon SageMaker服务主页，点击Open Studio进入模型开发环境。

2. 创建一个新的Jupyte NoteBook文件，复制以下代码安装必要依赖并指明Stable Diffusion模型ID。

%pip install --upgrade sagemaker --quiet
model_id = "model-imagegeneration-stabilityai-stable-diffusion-xl-base-1-0"

3. 运行以下代码列举出JumpStart中，可以快速部署的用于生成图片的所有Stable Diffusion大模型

import IPython
from ipywidgets import Dropdown
from sagemaker.jumpstart.notebook_utils import list_jumpstart_models
from sagemaker.jumpstart.filters import Andfilter_value = And("task == imagegeneration")
ss_models = list_jumpstart_models(filter=filter_value)dropdown = Dropdown(value=model_id,options=ss_models,description="Sagemaker Pre-Trained Image Generation Models:",style={"description_width": "initial"},layout={"width": "max-content"},
)
display(IPython.display.Markdown("## Select a pre-trained model from the dropdown menu"))
display(dropdown)

4. 运行以下代码开始部署Stable Diffusion大模型。

# Deploy the model
from sagemaker.jumpstart.model import JumpStartModel
from sagemaker.serializers import JSONSerializer
import time# The model is deployed on an ml.g5.4xlarge instance. To see all the supported parameters by the JumpStartModel
# class use this link - https://sagemaker.readthedocs.io/en/stable/api/inference/model.html#sagemaker.jumpstart.model.JumpStartModel
my_model = JumpStartModel(model_id=dropdown.value)
predictor = my_model.deploy()
# Wait for a few seconds so model the is properly loaded.
time.sleep(60)

5. 运行以下代码，导入调用大模型的必要依赖，配置图片生成请求参数，这里我们的图片生成提示词为”生成一个亚马逊雨林中的美洲虎图片“。同时我们定一个图片解码函数”decode_and_show“用于显示生成的图片，最后调用图片生成API "Predictor.predict()"生成图片。

from PIL import Image
import io
import base64
import json
import boto3
from typing import Union, Tuple
import ospayload = {"text_prompts": [{"text": "jaguar in the Amazon rainforest"}],"width": 1024,"height": 1024,"sampler": "DPMPP2MSampler","cfg_scale": 7.0,"steps": 50,"seed": 133,"use_refiner": True,"refiner_steps": 40,"refiner_strength": 0.2,
}def decode_and_show(model_response) -> None:"""Decodes and displays an image from SDXL outputArgs:model_response (GenerationResponse): The response object from the deployed SDXL model.Returns:None"""image = Image.open(io.BytesIO(base64.b64decode(model_response)))display(image)image.close()response = predictor.predict(payload)
# If you get a time out error, check the endpoint logs in Amazon CloudWatch for the model loading status
# and invoke it again.
decode_and_show(response["generated_image"])

我们在开发环境里可以看到大模型生成的图片内容。

6. 接下来我们进入到无服务器计算服务Lambda中，创建一个函数”check_toxicity_function“，用于调用Amazon Comprehend服务的API，模型检测输入文字的有害性并返回到客户端。我们复制以下代码到Lambda函数中

import json
import boto3
import oscomprehend = boto3.client('comprehend')
THRESHOLD = float(os.environ['THRESHOLD'])def check_toxicity(text_prompts):detected_labels = []for prompt in text_prompts:response = comprehend.detect_toxic_content(TextSegments=[{"Text": prompt['text']}],LanguageCode='en')labels = response['ResultList'][0]['Labels']# DIY section# Replace l['Name'] with {l['Name']:l['Score']} so that detected# is an array of json objectsdetected = [l['Name']for l in labels if l['Score'] > THRESHOLD]if detected:detected_labels.extend(detected)return detected_labelsdef lambda_handler(event, context):print("event is ", json.dumps(event))try:text_prompts = [json.loads(event['body'].strip('"'))]detected_labels = check_toxicity(text_prompts)if detected_labels:return {'statusCode': 200,'headers': {'Content-Type': 'application/json','Access-Control-Allow-Headers': 'Content-Type','Access-Control-Allow-Origin': '*','Access-Control-Allow-Methods': 'OPTIONS,POST'},'body': json.dumps({'detected_labels': detected_labels})}else:return {'statusCode': 200,'headers': {'Content-Type': 'application/json','Access-Control-Allow-Headers': 'Content-Type','Access-Control-Allow-Origin': '*','Access-Control-Allow-Methods': 'OPTIONS,POST'},'body': json.dumps({'detected_labels': 'non-toxic content and safe to proceed'})}except Exception as e:print(f"Error: {e}")return {'statusCode': 500,'headers': {'Content-Type': 'application/json','Access-Control-Allow-Headers': 'Content-Type','Access-Control-Allow-Origin': '*','Access-Control-Allow-Methods': 'OPTIONS,POST'},'body': json.dumps({'error': 'An error occurred while processing the request'})}

7. 我们再建一个新的Lambda函数”classifier_lambda_function“，调用Amazon Rekognition服务API对Stable Diffusion生成的图片进行内容审核。复制以下代码到Lambda中。

import io
import base64
import json
import boto3
import os
import uuid
import astcomprehend = boto3.client('comprehend')
sagemaker_runtime = boto3.client("runtime.sagemaker")
rekognition = boto3.client('rekognition')
s3_client = boto3.client('s3')
s3 = boto3.resource('s3')ENDPOINT_NAME = os.environ["ENDPOINT_NAME"]
bucket_name = os.environ['BUCKET_NAME']
THRESHOLD = 0.2
s3_folder = 'generated_images/'def query_endpoint(prompt):response = sagemaker_runtime.invoke_endpoint(EndpointName=ENDPOINT_NAME, ContentType="application/json", Body=json.dumps(prompt,separators=(',', ':')).encode("utf-8"))print("response is ",response)result = json.loads(response["Body"].read().decode())return resultdef detect_moderation(img_bytes):confidence_data = [ ]response = rekognition.detect_moderation_labels(Image={'Bytes': base64.b64decode(img_bytes)})for label in response['ModerationLabels']:confidence = label['Name'] + ' : ' + str(label['Confidence'])print (label['Name'] + ' : ' + str(label['Confidence']))print("confidence is ", confidence)confidence_data.append(confidence + "\n")return confidence_datadef lambda_handler(event,context):print("event is ",json.dumps(event))pm_str=json.loads(event["body"].strip('"'))prompt = {"text_prompts":  [(pm_str)],}print(prompt)response = query_endpoint(prompt)if "generated_image" in response:image_data = response["generated_image"]confLevel = detect_moderation(image_data)print(confLevel, len(confLevel))if len(confLevel) > 0:return {'statusCode': 400,'headers': {'Access-Control-Allow-Headers': 'Content-Type','Access-Control-Allow-Origin': '*','Access-Control-Allow-Methods': 'OPTIONS,POST'},'body': json.dumps(confLevel)}else:imageBytes = io.BytesIO(base64.b64decode(image_data))file_name = f'generated-image-{uuid.uuid4()}.jpg's3_client.upload_fileobj(imageBytes,bucket_name,f'{s3_folder}{file_name}',ExtraArgs={'ContentType': 'image/jpeg'})return {'statusCode': 200,'headers': {'Content-Type': 'image/png','Access-Control-Allow-Headers': 'Content-Type','Access-Control-Allow-Origin': '*','Access-Control-Allow-Methods': 'OPTIONS,POST'},'body': json.dumps(file_name),'isBase64Encoded': True}else:return {'statusCode': 400,'headers': {'Access-Control-Allow-Headers': 'Content-Type','Access-Control-Allow-Origin': '*','Access-Control-Allow-Methods': 'OPTIONS,POST'},'body': json.dumps({'error': 'Response is not in the expected format'})}

8. 接下来我们为Lambda函数前面添加一个API Gateway，作为API管理服务并提供对外暴露的API端点，在该服务中我们定义不同的HTTP方法、路径，绑定不同的Lambda函数来管理API。

如使用POST方法调用路径/classifier时，我们触发Lambda函数：”classifier_lambda_function“。使用POST方法调用路径/classifier/checkToxicity时，我们出发函数：”check_toxicity_function“。

同时API Gateway服务提供了端点URL供用户访问。