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

项目简介：

小李哥将继续每天介绍一个基于亚马逊云科技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 And


filter_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 os

payload = {
    "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 output

    Args:
        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 os

comprehend = 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 objects
        detected = [l['Name']for l in labels if l['Score'] > THRESHOLD]
        if detected:
            detected_labels.extend(detected)
    return detected_labels

def 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 ast


comprehend = 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 result
    
    
def 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_data

def 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供用户访问。

9. 本架构中我们使用到了CloudFront对API和网页请求进行加速，我们进入CloudFront服务页面中，复制并打开URL。

10. 首先我们对提示词文字进行检测，我们输入问题得到了回复”提示词包含侮辱性词汇“。

11. 我们再在相同界面中输入”生成一个晴朗的一天“，该提示词通过了文字有害性检测，生成的图片也通过安全检查，成功显示在生成界面中。

以上就是在亚马逊云科技上利用亚马逊云科技上利用Amazon Sagemaker部署Stable Diffusion模型，并对输入提示词和输出图像内容进行安全审核,的全部步骤。欢迎大家未来与我一起，未来获取更多国际前沿的生成式AI开发方案。