機械学習初学者が Amazon SageMaker で Stable Diffusion を動かしてみた

こちらは AWS for Games Advent Calendar 2022 の12月7日分のエントリーです。

はじめに

機械学習の入門ということで、Stable DiffusionをAmazon SageMaker上で動かし、AIイラストを出力してみようと思います。

Amazon SageMakerには多くの機能があり初学者はどこから初めてよいか分かりづらいと思いますが、この記事がその取っ掛かりとなれば幸いです。

目標

• Amazon SageMaker上でStable Diffusionを動かし、AIイラストを作成する
• 機械学習の流れを理解する

TL;DR

• 機械学習とは
• Amazon SageMakerとは
• Hugging Faceアカウント作成
  • Stable Diffusion ライセンス同意
  • アクセストークン作成
• ノートブックインスタンス作成
  • ノートブックインスタンス設定
  • アクセス許可と暗号化
  • Jupyter Notebookログイン
• 推論モデル作成
• SageMakerエンドポイント作成
• 画像生成
• お片付け
• 終わりに

機械学習とは

そもそも機械学習とは何なのかですが、簡単に言うと膨大なデータをコンピュータに登録することでそのデータからルールやパターンを発見させ、その結果を基にさまざまな分析や予測に利用するというものです。
我々の身近なところで機械学習が利用されているケースとしては、ECサイトのおすすめ商品です。
これはユーザーの購入履歴から他の似たようなユーザーの購入情報を利用して次に購入しそうな商品を提示しています。
今回扱うStable Diffusionでも膨大な画像データを学習しており、「猫」というテキスト情報からこれまでに学習した画像データを基にして新たに猫の画像を生成します。

Amazon SageMakerとは

そんな機械学習をモデル作成から動作まで一貫して管理してくれるサービスです。
機械学習モデルの作成には多くの計算が必要なので、強力なCPU・GPU等のマシンリソースやそれを動かすためのライブラリといった準備が必要不可欠なのですが、それらもAmazon SageMakerが用意してくれているのでお金を払うだけでいつでも利用できます。
今回は既成モデルを利用するので、モデル作成の費用は不要です。
ただし、コマンド実行用のローエンドEC2と画像生成用のGPUを積んだミドルレンジEC2を借りる費用だけは必要です。

Hugging Faceアカウント作成

Hugging FaceのWebページ( https://huggingface.co/ )にアクセスし、右上部の「Sing Up」からアカウント登録を行います。

Stable Diffusion ライセンス同意

Stable DiffusionのWebページ( https://huggingface.co/CompVis/stable-diffusion-v1-4 )のライセンスに同意しておきます。
同意しておかなければ後々にモデルをダウンロードする際に403エラーになります。(1敗)

アクセストークン作成

アクセストークンを作成します。
画面右上部のアカウントメニューからSettingsを開きます。

左サイドメニューから「Access Tokens」を選ぶと、トークンを作成する画面があります。

「New token」をクリックしてわかりやすい名前とread Roleを設定して、作成したアクセストークンをメモしておきます。

モデルダウンロード時に必要になります。(1敗)

ノートブックインスタンス作成

ここからAWS上で作業していきます。
WebコンソールのAmazon SageMakerを開き、ノートブックインスタンスからノートブックインスタンスの新規作成を行います。

ノートブックインスタンスとはJupyter Notebookがインストール済なEC2インスタンスです。
WebブラウザからPythonのプログラムを実行できるので今後の作業が非常に楽になります。

ノートブックインスタンス設定

Stable Diffusionのモデルをダウンロードして圧縮する必要があるので、追加設定からボリュームサイズのみ30GBに拡張しておきます。(1敗)
その他はデフォルト設定で大丈夫です。

アクセス許可と暗号化

今回はデフォルトのIAMロールで行います。

ノートブックインスタンスを作成し、ステータスが「InService」になるまで待ちましょう。

Jupyter Notebookログイン

ノートブックインスタンス作成が完了したら、Jupyter Notebookにログインしましょう。
アクションからJupyterを開くをクリックします。

画面右上のNewからconda_pythorch_p38を選択します。

このような画面になれば準備完了です。

この後はこのテキストボックスにpythonプログラムを入力し、Shift + Enterで実行していきます。

model.tar.gz作成

Stable Diffusionのモデルをダウンロードし、SageMakerエンドポイントで使えるようにtar.gzで固めていきます。
上から順番にJupyter Notebook上で実行してください。

!pip install "sagemaker==2.116.0" "huggingface_hub==0.10.1" --upgrade --quiet

import sagemaker
import boto3
sess = sagemaker.Session()
# sagemaker session bucket -> used for uploading data, models and logs
# sagemaker will automatically create this bucket if it not exists
sagemaker_session_bucket=None
if sagemaker_session_bucket is None and sess is not None:
    # set to default bucket if a bucket name is not given
    sagemaker_session_bucket = sess.default_bucket()

try:
    role = sagemaker.get_execution_role()
except ValueError:
    iam = boto3.client('iam')
    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']

sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)

print(f"sagemaker role arn: {role}")
print(f"sagemaker bucket: {sess.default_bucket()}")
print(f"sagemaker session region: {sess.boto_region_name}")

!mkdir code

%%writefile code/requirements.txt
diffusers==0.6.0
transformers==4.23.1

%%writefile code/inference.py
import base64
import torch
from io import BytesIO
from diffusers import StableDiffusionPipeline


def model_fn(model_dir):
    # Load stable diffusion and move it to the GPU
    pipe = StableDiffusionPipeline.from_pretrained(model_dir, torch_dtype=torch.float16)
    pipe = pipe.to("cuda")

    return pipe


def predict_fn(data, pipe):

    # get prompt & parameters
    prompt = data.pop("inputs", data)
    # set valid HP for stable diffusion
    num_inference_steps = data.pop("num_inference_steps", 50)
    guidance_scale = data.pop("guidance_scale", 7.5)
    num_images_per_prompt = data.pop("num_images_per_prompt", 4)

    # run generation with parameters
    generated_images = pipe(
        prompt,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        num_images_per_prompt=num_images_per_prompt,
    )["images"]

    # create response
    encoded_images = []
    for image in generated_images:
        buffered = BytesIO()
        image.save(buffered, format="JPEG")
        encoded_images.append(base64.b64encode(buffered.getvalue()).decode())

    # create response
    return {"generated_images": encoded_images}

from distutils.dir_util import copy_tree
from pathlib import Path
from huggingface_hub import snapshot_download
import random

HF_MODEL_ID="CompVis/stable-diffusion-v1-4"
HF_TOKEN="" # Hugging Faceのアクセストークンを入力してください
assert len(HF_TOKEN) > 0, "Please set HF_TOKEN to your huggingface token. You can find it here: https://huggingface.co/settings/tokens"

# download snapshot
snapshot_dir = snapshot_download(repo_id=HF_MODEL_ID,revision="fp16",use_auth_token=HF_TOKEN)

# create model dir
model_tar = Path(f"model-{random.getrandbits(16)}")
model_tar.mkdir(exist_ok=True)

# copy snapshot to model dir
copy_tree(snapshot_dir, str(model_tar))

# copy code/ to model dir
copy_tree("code/", str(model_tar.joinpath("code")))

import tarfile
import os

# helper to create the model.tar.gz
def compress(tar_dir=None,output_file="model.tar.gz"):
    parent_dir=os.getcwd()
    os.chdir(tar_dir)
    with tarfile.open(os.path.join(parent_dir, output_file), "w:gz") as tar:
        for item in os.listdir('.'):
          print(item)
          tar.add(item, arcname=item)
    os.chdir(parent_dir)

compress(str(model_tar))

model.tar.gzのS3アップロード

固めたmodel.tar.gzをS3にアップロードします。
今回はSageMakerのデフォルトバケット内に保存します。

from sagemaker.s3 import S3Uploader

# upload model.tar.gz to s3
s3_model_uri=S3Uploader.upload(local_path="model.tar.gz", desired_s3_uri=f"s3://{sess.default_bucket()}/stable-diffusion-v1-4")

print(f"model uploaded to: {s3_model_uri}")

SageMaker エンドポイント作成

SageMakerエンドポイントを作成します。

from sagemaker.huggingface.model import HuggingFaceModel


# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
   model_data=s3_model_uri,      # path to your model and script
   role=role,                    # iam role with permissions to create an Endpoint
   transformers_version="4.17",  # transformers version used
   pytorch_version="1.10",       # pytorch version used
   py_version='py38',            # python version used
)

# deploy the endpoint endpoint
predictor = huggingface_model.deploy(
    initial_instance_count=1,
    instance_type="ml.g4dn.xlarge"
)

画像生成

実際にSageMakerエンドポイントにアクセスしてイラストを出力してみましょう。

from PIL import Image
from io import BytesIO
from IPython.display import display
import base64
import matplotlib.pyplot as plt

# helper decoder
def decode_base64_image(image_string):
  base64_image = base64.b64decode(image_string)
  buffer = BytesIO(base64_image)
  return Image.open(buffer)

# display PIL images as grid
def display_images(images=None,columns=3, width=100, height=100):
    plt.figure(figsize=(width, height))
    for i, image in enumerate(images):
        plt.subplot(int(len(images) / columns + 1), columns, i + 1)
        plt.axis('off')
        plt.imshow(image)

response = predictor.predict(data={
  "inputs": "1 girl",
  "num_images_per_prompt" : 1,
  "num_inference_steps" : 25,
  "guidance_scale" : 7.5
  }
)

# decode images
decoded_images = [decode_base64_image(image) for image in response["generated_images"]]

# visualize generation
display_images(decoded_images)

暫く待つとAIイラストが表示されます。
inputs部分を書き換えれば自由に好きなイラストを作成できます。

お片付け

SageMakerのエンドポイントを削除します。

predictor.delete_model()
predictor.delete_endpoint()

ノートブックインスタンスを削除します。
インスタンスを停止してから削除します。

S3上のmodel.tar.gzを削除します。

以上でお片付け完了です。
お疲れ様でした。

終わりに

Amazon SageMakerを利用したStable Diffusionの利用方法についてご紹介しました。
Stable Diffusion以外にもHugging Faceには様々なモデルがあるので是非動かしてみてください。

それでは、残りの年末も良い機械学習ライフをお過ごし下さい。

(免責) 本記事の内容はあくまでも個人の意見であり、所属する企業や団体とは関係がございません。