亚马逊AWS官方博客

云端生成式 AI – 基于 Amazon EKS 的 Stable Diffusion 图像生成方案

Stable Diffusion 是当下生成式 AI 领域最受欢迎的开源多模态语言-图像模型,由于其易用的接口和良好的使用体验,受到了开源社区和广大设计行业从业者的追捧。Stable Diffusion 模型版本正在快速迭代,并带动了各行各业的生产力变革。目前市场上有多种围绕 Stable Diffusion 模型的行业解决方案,比如亚马逊云科技推出的 Stable Diffusion 插件解决方案。本文中,我们将介绍一种更加适合 toC 场景的解决方案,它将帮助客户从容应对大并发量的请求,达到分钟级冷启动;使用更简单的控制平面,为客户提供更灵活的架构,轻松实现各种功能模块设计;同时,它还利用 Amazon Spot instances,帮助客户有效控制成本。

Stable Diffusion 文生图原理

Stable Diffusion 图生图原理

方案特点

我们根据实际客户需求和案例经验,提炼出该方案的特点如下:

  • 更快速的扩容,基于 KEDAKarpenterBottlerocket,提供可自定义的任务粒度的自动扩容能力,且最快可以达到 1 分钟内业务就绪,适合高并发且对响应时间有要求的情况。
  • 大幅降低成本,通过使用 Amazon Spot instancesAmazon EFS,使得计算成本相较于 On-demand instances 可节省最高约 70%,多节点场景下的存储成本相较于 EBS 更可节省数倍。
  • 运维非常方便,方案是云原生架构,基于 Amazon EKS、无服务器(Serverless)和事件驱动(Event-Driven),大部分组件无需自行管理,适合运维人手不足的情况。
  • 可视化监控,提供调用流程详细的可视化,历史统计与异常检测功能,有效缩短 MTTI(Mean Time to Identify)。
  • 自动化部署,方案使用 AWS CDK,借助编程语言的强大表达能力和高效执行效率,仅需 30-40 分钟就可以自动完成构建。

方案架构

方案流程如下:

  • 用户将提示词和参数通过 http 请求的方式,经过 Amazon CloudFront 加速,传递到应用服务器(可通过 AWS Fargate 部署)。
  • 应用服务器完成相关业务处理后(如:认证鉴权/task 跟踪/参数优化等),调用 Amazon API Gateway 提供的推理 API 接口,之后请求交由 AWS Lambda 校验并发布到 Amazon SNS 服务,并拿到同步响应。
  • Amazon SNS 借助可自定义的 filter 规则将请求以异步消息的方式发布到对应的队列服务 Amazon SQS 中,多队列的设计可灵活适应复杂的业务需求。
  • KEDA 会根据 Amazon SQS 中的消息数量,动态扩容 Amazon EKS 集群中的工作负载 Pod。
  • 扩容动作会借助 Karpenter 启动 Amazon GPU Spot instances 作为工作节点,并使用基于 Bottlerocket 构建的 snapshot 来快速启动 Stable Diffusion WebUI 工作负载。
  • Stable Diffusion WebUI 以指定配置参数启动,并从 Amazon EFS 中加载基础模型文件。
  • Queue Agent 从任务队列中拉取任务,完成相关处理(下载输入图片/切换基础模型等),调用 Stable Diffusion WebUI 的 API 完成推理并生成结果,将结果图片保存到 Amazon S3 中。
  • Queue Agent 将任务响应数据发布到 Amazon SNS 服务,以便应用服务器能够通过多种方式(Amazon SQS/Http Endpoint)获取异步响应数据并进行后续业务处理。
  • 方案使用 Amazon CloudWatchAWS Distro for OpenTelemetryAWS X-Ray 来采集、追踪和展示相关过程中产生的日志和指标。
  • 方案还使用了 AWS IAM 来控制资源的访问授权,使用 Amazon ECR 存储容器镜像,使用 AWS DataSync 同步模型文件,使用 AWS CDK 来完成自动化部署。

主要服务列表:

  • Amazon Elastic Kubernetes Service(Amazon EKS):是一项托管 Kubernetes 服务,用于在 AWS Cloud 上运行 Kubernetes。在云中,Amazon EKS 可自动管理负责安排容器、管理应用程序可用性、存储集群数据和其他关键任务的 Kubernetes 控制面板节点的可用性和可扩展性。
  • Amazon ECR:Amazon Elastic Container Registry是完全托管式容器注册表,提供高性能托管,让您能在任何地方可靠地部署应用程序映像和构件。
  • AWS Lambda:AWS Lambda 是一项事件驱动、随用随付的计算服务,让您无需预置或管理服务器即可运行代码。
  • Amazon SQS:Amazon Simple Queue Service(SQS)是一种消息队列服务,可让您分离和扩展微服务、分布式系统和无服务器应用程序。
  • Amazon SNS:Amazon Simple Notification Service(SNS)是一项用于应用与应用之间(A2A)以及应用与人之间(A2P)通信的完全托管式消息收发服务。
  • Amazon API Gateway:Amazon API Gateway 是一种完全托管式服务,可以轻松创建、发布任意规模的 API。
  • Amazon CloudWatch:Amazon CloudWatch 会在自动化控制面板中收集并可视化实时日志、指标和事件数据,从而简化您的基础设施和应用程序维护。
  • AWS X-Ray:AWS X-Ray 通过您的应用程序时会提供完整的请求视图,并通过无代码和低代码移动过滤负载、函数、轨迹、服务、API 等方面的可视化数据。
  • Amazon S3:Amazon Simple Storage Service(Amazon S3)是一项对象存储服务,旨在存储和保护任意量的数据。
  • Amazon EFS:适用于构建人员的无服务器、完全弹性的文件系统,可以轻松设置、扩展高可用性共享存储并对其进行成本优化。
  • AWS DataSync:安全的,可自动化并加速本地存储和 AWS 存储服务之间的数据移动服务。

方案实践

前置条件

本博客假设您已熟悉 Amazon EC2、AWS Cloud Development Kit(AWS CDK)、Docker、Amazon Elastic Container Registry(Amazon ECR)、Amazon Elastic Kubernetes Service(Amazon EKS)。

本方案后续步骤的前置条件如下:

  • Kubernetes cluster 1.22+
  • Helm CLI
  • Node.js 16+
  • AWS CDK CLI
  • An AWS Account
  • Administrator or equivalent privilege

本方案默认使用 G5 系列机型(如需使用其它系列机型可修改运行时 Helm Chart 的配置 karpenter.provisioner.instanceType),推荐的区域(Region)如下:

  • US East (N. Virginia) us-east-1
  • US West (Oregon) us-west-2

本方案建议准备一个可访问互联网的调试部署环境,该环境需要能运行 AWS cli、Git、Docker、NPM、AWS CDK、kubectl、eksctl 等命令。我们建议您使用 AWS Cloud9Amazon SageMaker Notebook 来进行部署。

环境部署

环境准备:

我们的方案已在 Github 上开源,地址为:https://github.com/aws-samples/stable-diffusion-on-eks

运行以下命令以获取源代码和部署脚本:

git clone --recursive https://github.com/aws-samples/stable-diffusion-on-ekscd stable-diffusion-on-eks

请参考我们的实施指南,会帮助您一步一步完成方案的环境部署:https://aws-samples.github.io/stable-diffusion-on-eks/zh/implementation-guide/deployment/

方案效果

方案部署完成以后,我们可以使用代码或者 Postman 等方式构造 http 请求并发送,将获取的 API Endpoint 做为请求 URL,Method 为 Post,Body 为 application/json,示例 Payload 如下(请将注释移除后使用):

文生图:

{
    "alwayson_scripts": {
        // 必要,任务类型
        "task": "text-to-image",
        // 必要,基础模型名称,关联队列分发或模型切换
        "sd_model_checkpoint": "v1-5-pruned-emaonly.safetensors",
        // 必要,任务ID,在上传结果图片和返回响应时会用到
        "id_task": "21123",
        "uid": "456",
        "save_dir": "outputs"
    },
    // 以下皆为官方参数,使用默认值或者直接传入即可
    "prompt": "A dog",
    "steps": 30,
    "width": 768,
    "height": 768
}

文生图结果:

图生图:

{
    "alwayson_scripts": {
        // 必要,任务类型
        "task": "image-to-image",
        // 必要,输入图片的url
        "image_link": "https://www.segmind.com/sd-img2img-input.jpeg",
        // 必要,任务ID,在上传结果图片和返回响应时会用到
        "id_task": "31311",
        // 必要,基础模型名称,关联队列分发或模型切换
        "sd_model_checkpoint": "v1-5-pruned-emaonly.safetensors",
        // 非必要,用户id
        "uid": "456"
    },
    // 以下皆为官方参数,使用默认值或者直接传入即可
    "prompt": "A fantasy landscape, trending on artstation, mystical sky",
    "steps": 16,
    "width": 512,
    "height": 512
}

图生图结果:

注:结果图片和响应数据,会存储到相关的 S3 output bucket 中以便后续访问。

方案细节

更快的启动时间:toC 业务往往具有高并发的特点,为了迅速响应请求,本方案使用 Karpenter 进行 node 的弹性伸缩,底层直接调用了 EC2 fleet API,结合基于 Bottlerocket 构建的 snapshot,使得服务冷启动时间最快可达到 1 分钟内。

# a sample config for karpenter nodeTemplate
modelsRuntime:
- name: "sample1"
  namespace: "sample1"
  modelFilename: "v1-5-pruned-emaonly.safetensors"
  extraValues:
    karpenter: 
      nodeTemplate:
        amiFamily: Bottlerocket
        dataVolume:
          volumeSize: 80Gi
          volumeType: gp3
          deleteOnTermination: true
          iops: 4000
          throughput: 1000
          snapshotID: snap-0c31434acc0a8a558

更低的成本:通过使用 Spot 实例和混合 GPU 实例大幅度降低推理成本。虽然 Spot 实例有被回收的风险,但由于 Stable Diffusion 单任务耗时短且无状态,而且在 Spot 实例回收前会处理相应事件并重新启动新的实例,故 Spot 实例完全可以胜任该场景推理需求。Karpenter 默认使用 pricing-capacity-optimized 的 Spot 实例选择策略,会优先使用回收风险最低,且最廉价的 Spot 实例类型。当遇到 Spot 实例无法启动时,将以 On-demand 方式启动,且支持异构和不同家族的实例。

# a sample config for karpenter provisioner
modelsRuntime:
- name: "sample1"
  namespace: "sample1"
  modelFilename: "v1-5-pruned-emaonly.safetensors"
  extraValues:
    karpenter: 
      provisioner:
        instanceType:
        - "g5.xlarge"
        - "g5.2xlarge"
        - "g4dn.xlarge"
        - "g4dn.2xlarge"

更简单的控制平面:EKS 可直接使用 KarpenterKEDA 等社区组件进行任务调度和弹性伸缩,使用 Amazon SQS 的消息数作为弹性指标更贴近实际业务需求。

# a sample config for scaling
modelsRuntime:
- name: "sample1"
  namespace: "sample1"
  modelFilename: "v1-5-pruned-emaonly.safetensors"
  extraValues:
    sdWebuiInferenceApi:
      scaling:
        queueLength: 10
        cooldownPeriod: 60
        pollingInterval: 1
        maxReplicaCount: 10
        minReplicaCount: 1
      inferenceApi:
        image:
          repository: 123456789012.dkr.ecr.us-west-2.amazonaws.com/sd-on-eks/inference-api
          tag: abcd1
      queueAgent:
        image:
          repository: 123456789012.dkr.ecr.us-west-2.amazonaws.com/sd-on-eks/queue-agent
          tag: latest
      persistence:
        enabled: true
        existingClaim: "efs-model-storage-pvc"

更灵活的设计:可以使用多个 Amazon SQS 队列并结合 Amazon SNS 的 filter 功能,来满足不同业务场景的架构设计需求。例如通过创建针对特定基础模型的专用队列,避免基础模型频繁切换以加快特定模型的生图速度,而对于需要切换主模型的请求可以将其路由到专门的动态队列处理。例如为付费用户的请求设置单独的专用队列,用更多的资源提供更好的服务等。

// lib/runtime/sdRuntime.ts
if (!this.options.dynamicModel) {
  this.options.inputSns!.addSubscription(new aws_sns_subscriptions.SqsSubscription(inputQueue, {
    filterPolicy: {
      sd_model_checkpoint:
        sns.SubscriptionFilter.stringFilter({
          allowlist: [this.options.sdModelCheckpoint]
        })
    }
  }))
  generatedValues.sdWebuiInferenceApi.queueAgent.dynamicModel = false
} else {
  this.options.inputSns!.addSubscription(new aws_sns_subscriptions.SqsSubscription(inputQueue, {
    filterPolicy: {
      sd_model_checkpoint:
        sns.SubscriptionFilter.stringFilter({
          denylist: this.options.allModels
        })
    }
  }))
  generatedValues.sdWebuiInferenceApi.queueAgent.dynamicModel = true
}

更精细的并发控制:将 Queue Agent 与 SD WebUI 部署在同一 Pod,以拉的模式来处理异步任务,能够对并发进行灵活精细的控制,缓解了 SD WebUI 负载过重无法响应或超时等问题,同时也方便调整超时时间和重试策略等参数来进行针对性优化。

# src/backend/queue_agent/main.py
queue = sqsRes.Queue(sqs_queue_url)
SQS_WAIT_TIME_SECONDS = 20
check_readiness()
while True:
    received_messages = receive_messages(queue, 1, SQS_WAIT_TIME_SECONDS)
    for message in received_messages:
        try:
            if taskType == 'text-to-image':
                r = invoke_txt2img()
            elif taskType == 'image-to-image':
                r = invoke_img2img()
        except Exception as e:
            traceback.print_exc()
        finally:
            handle_outputs()
            delete_message(message)

更好的兼容性:对 SD WebUI Runtime 无侵入,可兼容官方各版本的 Runtime 或客户定制的 Runtime,将定制化的部分放到 Queue Agent 以方便二次开发。支持异步响应以多种方式(Http/Https/SQS)发送到目标和 Amazon SageMaker 异步响应规范。

更直观的监控:集成了 AWS X-RayAmazon CloudWatchAWS Distro for OpenTelemetry,可以更直观地跟踪整个流程,查看日志和相关指标(例如:流程各环节调用状态和耗时/共享存储的 IO/工作节点的负载等)。

更方便易用:方案基于 AWS CDK 自动进行部署(根据不同环境估算约 30-40 分钟),并提供了模型文件自动同步(Amazon S3 → Amazon EFS)、模型文件批量下载上传、图片浏览应用等辅助工具。

结论

来源:Gartner(2023 年 8 月)

根据 Gartner 披露的“人工智能技术成熟度曲线”,生成式 AI 仍处于需求的快速增长期,其广阔的应用场景和巨大需求空间吸引着大量资本和技术的投入,预计将在 2-5 年内实现规模化应用。作为开源多模态生成式 AI 领域的领导者,Stability.ai 推出了更为强大的 Stable Diffusion XL 1.0,并发布在 Amazon BedRock 上。

本文介绍了一种扩容迅速、成本低、维护方便的 Stable Diffusion 开源方案,适合熟悉容器,无服务器并需要以快速扩展和低成本方式部署 Stable Diffusion 模型的学生、设计师和开发人员,应用在高速发展的 toC 业务中。AWS 致力于为客户提供丰富而功能强大的图像生成方案,也支持开源并持续推动优化,为业务创新提供更好的助力。

本篇作者

陆从尧

AWS 解决方案架构师,负责基于 AWS 云平台的解决方案咨询和设计. 从业十七年,担任过研发经理、架构师等多种角色,有多年的互联网软件研发、系统架构设计及开发经验。

孙华

亚马逊云科技资深无服务器产品专家。在过去的 20 多年一直从事软件架构、程序开发以及技术推广等领域的工作。他擅长 Web 领域应用、SaaS 系统和云呼叫中心开发,也从事过多个大型软件项目的设计、开发与项目管理。目前他专注与云计算以及互联网等技术领域,致力于帮助中国的 开发者构建基于云计算的新一代的互联网应用。

于昺蛟

亚马逊云科技现代化应用解决方案架构师,负责亚马逊云科技容器和无服务器产品的架构咨询和设计。在容器平台的建设和运维,应用现代化,DevOps 等领域有多年经验,致力于容器技术和现代化应用的推广。