SageMakerでのHuggingFace Diffusersデプロイガイド(1/3)：Stable Diffusion XLのデプロイ

背景

• SageMakerでHuggingFaceのStableDiffusion XLをデプロイします
• promptをinputとして渡して画像を生成するシンプルなものを作成します
• モデルはstable-diffusion-xl-base-1.0 を使います
• デプロイにはPythonLibraryの HuggingFaceModelを使います

実装

このドキュメントを参考にしました

モデルの設定とデプロイ

import sagemaker
from sagemaker.huggingface import HuggingFaceModel

sagemaker_session = sagemaker.Session()
role = sagemaker.get_execution_role()  # IAM role with necessary permissions

# Hugging Faceモデルの設定
hub = {
    'HF_MODEL_ID': 'diffusers/controlnet-depth-sdxl-1.0',
    'HF_TASK': 'text-to-image',
}

# HuggingFaceModelクラスのインスタンスを作成
huggingface_model = HuggingFaceModel(
    transformers_version='4.26.0',
    pytorch_version='1.13.1',
    py_version='py39',
    env=hub,
    role=role,
)

# モデルのデプロイ
predictor = huggingface_model.deploy(
    initial_instance_count=1,
    instance_type='ml.g5.xlarge',  # 適切なインスタンスタイプを選択
)

• ドキュメントの通りPython Libraryで実装
• HuggingFaceModel を使うと、シンプルに実装が可能できます

推論実行コード

import base64
from PIL import Image
import io

data = {
    'inputs': 'A beautiful sunset over the mountains'
}

# 推論の実行
response = predictor.predict(data)

# 推論結果をデコードして画像として保存
image_data = base64.b64decode(response['body'])  # 実際のキー名に合わせて調整
image = Image.open(io.BytesIO(image_data))
image.save('generated_image.png')

print("画像が 'generated_image.png' として保存されました。")

まとめ

HuggingFaceのモデルのpipe()の実行だけであれば、HuggingFaceModelを使えば非常にシンプルにデプロイすることが可能でした。次からは、カスタムの推論コードをデプロイするようなケースを紹介していきます。