DOKを使ってFLUX.1でテキストから画像を生成する

DOKはコンテナー型のGPUサービスで、NVIDIA V100やH100を実行時間課金で利用できるサービスです。

コンテナー型GPUクラウドサービス 高火力 DOK（ドック） | さくらインターネット

今回はこのDOKを使って、FLUX.1による画像生成を試してみました。

結果

結果は以下のようになっています。テキストは「Anime style, smiling Japanese boys and Japanese girls enjoy party outside on the grass around cherry blossom.」としています。ちゃんとそれっぽい画像が生成されています。

2回目。

参考

Stable Diffusionのオリジナル開発陣が発表した画像生成AIモデルFLUX.1([dev]/[schnell])をMacBook(M2)で動かしてみた #Python - Qiitaを参考にさせてもらっています。

コンテナイメージの作成と登録

DOKで必要なコンテナイメージを作成する手順は以下の通りです。ベースとして、Stable Diffusionのオリジナル開発陣が発表した画像生成AIモデルFLUX.1([dev]/[schnell])をMacBook(M2)で動かしてみた #Python - Qiitaを参考にしつつ、DOK向けに改造しています。

完成版はgoofmint/dok-flux-1にありますので、実装時の参考にしてください。

アクセス権の取得

事前にblack-forest-labs/FLUX.1-dev · Hugging Faceへアクセスして、利用許諾を取得してください。利用許諾を取得した後、Hugging Faceのトークンを取得してください（Hagging Faceのアカウントが必須です）。設定はReadのみで大丈夫そうです。

requirements.txt の作成

DOKで実行する際に必要なライブラリを追加します。boto3は、今回は利用していませんが、生成した画像を任意のオブジェクトストレージへ保存する歳に利用します。

sentencepiece
torch
torchvision
torchaudio
accelerate
protobuf
transformers
git+https://github.com/huggingface/diffusers.git
boto3

Dockerfile の作成

Dockerfileの内容は以下の通りです。必要なライブラリをインストールし、ファイルをコピーしています。また、 requirements.txt を使ってライブラリをインストールしています。

FROM nvidia/cuda:12.5.1-runtime-ubuntu22.04
ENV DEBIAN_FRONTEND=noninteractive

# 必要なパッケージのインストール
# /appはアプリのディレクトリ、/opt/artifactはアウトプット先のディレクトリ
RUN apt-get update && \
    apt-get install -y \
        git \
        python3 \
        python3-pip \
      && \
    mkdir /app /opt/artifact && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

WORKDIR /app
COPY requirements.txt runner.py .
# 依存ライブラリのインストール
RUN pip install -r requirements.txt

# Dockerコンテナー起動時に実行するスクリプト（後で作成）
COPY docker-entrypoint.sh /
# 実行権限を付与
RUN chmod +x /docker-entrypoint.sh /

WORKDIR /
# Dockerコンテナー起動時に実行するスクリプトを指定して実行
CMD ["/bin/bash", "/docker-entrypoint.sh"]

docker-entrypoint.sh の作成

docker-entrypoint.sh はDockerコンテナー起動時に実行するスクリプトです。ここでは環境変数をチェックして、 runner.py を呼び出します。以下の内容で作成します。

#!/bin/bash
 
set -ue
shopt -s nullglob

export TZ=${TZ:-Asia/Tokyo}

# アウトプット先ディレクトリ（自動付与） /opt/artifact固定です
if [ -z "${SAKURA_ARTIFACT_DIR:-}" ]; then
  echo "Environment variable SAKURA_ARTIFACT_DIR is not set" >&2
  exit 1
fi

# DOKのタスクID（自動付与）
if [ -z "${SAKURA_TASK_ID:-}" ]; then
  echo "Environment variable SAKURA_TASK_ID is not set" >&2
  exit 1
fi

# 読み上げるテキスト（環境変数で指定）
if [ -z "${PROMPT:-}" ]; then
  echo "Environment variable PROMPT is not set" >&2
  exit 1
fi

# Hagging Faceのトークン（環境変数で指定）
if [ -z "${HF_TOKEN:-}" ]; then
  echo "Environment variable HF_TOKEN is not set" >&2
  exit 1
fi

# S3_はすべてboto3用の環境変数です
pushd /app
  python3 runner.py \
	  --id="${SAKURA_TASK_ID}" \
	  --output="${SAKURA_ARTIFACT_DIR}" \
	  --prompt="${PROMPT}" \
      --hf-token="${HF_TOKEN}" \
	  --s3-bucket="${S3_BUCKET:-}" \
	  --s3-endpoint="${S3_ENDPOINT:-}" \
	  --s3-secret="${S3_SECRET:-}" \
	  --s3-token="${S3_TOKEN:-}"
popd

runner.py の作成

runner.py は実際に処理を行うスクリプトです。まず必要なライブラリをインポートします。

import torch
from diffusers import FluxPipeline
import argparse

パラメータの取得

docker-entrypoint.sh から渡されたパラメータを取得します。

arg_parser = argparse.ArgumentParser()

arg_parser.add_argument(
    '--output',
    default='/opt/artifact',
    help='出力先ディレクトリを指定します。',
)
arg_parser.add_argument(
    '--prompt',
    default='',
    help='プロンプト',
)
arg_parser.add_argument(
    '--id',
    default='',
    help='タスクIDを指定します。',
)
arg_parser.add_argument(
    '--hf-token',
    default='',
    help='Hagging Faceのトークンを指定します。',
)
arg_parser.add_argument('--s3-bucket', help='S3のバケットを指定します。')
arg_parser.add_argument('--s3-endpoint', help='S3互換エンドポイントのURLを指定します。')
arg_parser.add_argument('--s3-secret', help='S3のシークレットアクセスキーを指定します。')
arg_parser.add_argument('--s3-token', help='S3のアクセスキーIDを指定します。')

args = arg_parser.parse_args()

FLUX.1の準備

FLUX.1のモデルを読み込みます。ここでHagging Faceのトークンが必要です。

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev",
                                    torch_dtype=torch.bfloat16,
                                    token=args.hf_token
)
pipe.to(torch.device("cuda"))  # deviceとしてCUDAを指定する

画像を生成する

実際の画像生成は、以下のコードだけで行えますので、とても簡単です。

image = pipe(
    args.prompt,
    height=1024,
    width=1024,
    guidance_scale=3.5,
    output_type="pil",
    num_inference_steps=50,
    max_sequence_length=512,
    generator=torch.Generator("cuda").manual_seed(0)
).images[0]
image.save(output_path)

全体の処理

runner.py の全体の処理は以下の通りです。

# 必要なライブラリをインポート
import torch
from diffusers import FluxPipeline  # DiffusersライブラリからFluxPipelineをインポート
import argparse  # コマンドライン引数を処理するためのモジュール

# コマンドライン引数の処理を行うためのArgumentParserを作成
arg_parser = argparse.ArgumentParser()

# 出力先ディレクトリを指定する引数を追加
arg_parser.add_argument(
    '--output',
    default='/opt/artifact',
    help='出力先ディレクトリを指定します。',
)

# 画像生成に使用するプロンプト（指示文や説明文）を指定する引数を追加
arg_parser.add_argument(
    '--prompt',
    default='',
    help='プロンプト',
)

# タスクIDを指定する引数を追加
arg_parser.add_argument(
    '--id',
    default='',
    help='タスクIDを指定します。',
)

# Hugging Faceの認証トークンを指定する引数を追加
arg_parser.add_argument(
    '--hf-token',
    default='',
    help='Hugging Faceのトークンを指定します。',
)

# オプションとして、S3（クラウドストレージ）関連の設定を行うための引数を追加
arg_parser.add_argument('--s3-bucket', help='S3のバケットを指定します。')
arg_parser.add_argument('--s3-endpoint', help='S3互換エンドポイントのURLを指定します。')
arg_parser.add_argument('--s3-secret', help='S3のシークレットアクセスキーを指定します。')
arg_parser.add_argument('--s3-token', help='S3のアクセスキーIDを指定します。')

# コマンドライン引数を解析して取得
args = arg_parser.parse_args()

# 出力画像の保存パスを指定（例: output-<タスクID>.png）
output_path = f'{args.output}/output-{args.id}.png'

# FluxPipeline（画像生成モデル）を事前学習済みのデータからロード
pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",  # 使用するモデルの名前
    torch_dtype=torch.bfloat16,  # 計算で使用するデータ型（省メモリ型）
    token=args.hf_token  # Hugging Faceの認証トークンを指定
)

# モデルをGPU（CUDAデバイス）に移動
pipe.to(torch.device("cuda"))

# 画像を生成
image = pipe(
    args.prompt,  # 指定されたプロンプトを使用
    height=1024,  # 出力画像の高さ（ピクセル）
    width=1024,  # 出力画像の幅（ピクセル）
    guidance_scale=3.5,  # ガイダンススケール（生成画像の方向性を制御）
    output_type="pil",  # 出力形式をPIL画像形式に設定
    num_inference_steps=50,  # 推論ステップ数（多いほど高品質だが時間がかかる）
    max_sequence_length=512,  # プロンプトの最大長
    generator=torch.Generator("cuda").manual_seed(0)  # 再現性のため乱数シードを設定
).images[0]  # 最初に生成された画像を取得

# 生成された画像を指定したパスに保存
image.save(output_path)

Dockerイメージのビルド

上記の内容で、Dockerイメージをビルドします。Linux環境などで行います。

コンテナレジストリの用意

Dockerイメージを登録するコンテナレジストリを作成します。さくらのクラウドではLAB機能で、コンテナレジストリを提供しています。さくらのクラウドにログインしたら さくらのクラウド を選択します。

左側のメニューの LAB の中にある コンテナレジストリ を選択します。

追加 を押して、コンテナレジストリを作成します。最低限、以下の入力が必要です。

項目	設定
名前	分かりやすい、任意の名前を入力してください
コンテナレジストリ名	ドメイン名に使われます。以下では、 EXAMPLE.sakuracr.jp として説明します
公開設定	Pullのみとします

ユーザーの作成

コンテナレジストリを作成したら、作成したコンテナレジストリを一覧でダブルクリックします。

詳細表示にて、ユーザータブをクリックします。

追加ボタンを押し、ユーザーを作成します。 YOUR_USER_NAME と PASSWORD は任意のものを指定してください。

項目	設定
ユーザー名	YOUR_USER_NAME
パスワード	YOUR_PASSWORD
ユーザ権限設定	All

Dockerイメージのビルド

DockerイメージはLinuxで行います。今回はUbuntu 24.04を使っています。Dockerが使える環境であれば、Windows + WSL2でも問題ありません。macOSの場合、アーキテクチャが異なるので動かせないかも知れません（未検証です）。

EXAMPLE.sakuracr.jp の部分は、作成したコンテナレジストリのドメイン名に置き換えてください。また、 flux-1 は任意の名前で大丈夫です（以下はその名称で読み替えてください）。

sudo docker build -t EXAMPLE.sakuracr.jp/flux-1:latest .

コンテナレジストリへのログイン

作成したコンテナレジストリにログインします。ログインIDとパスワードが求められるので、作成したものを入力してください。

sudo docker login EXAMPLE.sakuracr.jp

イメージのプッシュ

作成したイメージをコンテナレジストリにプッシュします。イメージサイズが大きいので、数十分かかります。

sudo docker push EXAMPLE.sakuracr.jp/flux-1:latest

タスクを作成する

後は最初と同じようにDOKでタスクを作成、実行します。

項目	設定
イメージ	EXAMPLE.sakuracr.jp/flux-1
環境変数	PROMPT = Anime style, smiling Japanese boys and Japanese girls enjoy party outside on the grass around cherry blossom. HF_TOKEN = (Hagging Faceのトークン)

なお、プランはh100-80gbを指定してください。デフォルトのv100-32gbではメモリが足りないため、エラーが発生します。

まとめ

今回はFLUX.1を使って、DOK上で画像を作成する手順を解説しました。かなり高品質な画像が生成できるので、ぜひ試してみてください。DOKでは、ｌこういった非同期タスクを多数立ち上げて、後は結果を待つのみと言った使い方ができます。ぜひAI・機械学習に活用してください。

コンテナー型GPUクラウドサービス 高火力 DOK（ドック） | さくらインターネット