GPUコンテナの高火力DOKでJanus（画像生成）を試す

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

今回はこの高火力 DOKを使って、deepseek-ai/Janusを試してみました。

注意

JanusはDeepSeekが開発した画像生成AIです。人物やアニメ風画像生成などに特化したモデルとなっています。プロンプトベースで画像生成ができます。

参考

deepseek-ai/Janus: Janus-Series: Unified Multimodal Understanding and Generation Modelsにあるコードをベースに進めます。

とりあえず試す

コンテナレジストリにプッシュ済みのイメージを使って実行してみたい方は、高火力 DOKにて新しいタスクを作成し、以下の情報を入力してください。

項目	設定
イメージ	dok-handson.sakuracr.jp/janus
環境変数	PROMPT = A stunning boy from kabul in red, white traditional clothing, blue eyes, brown hair

PROMPTは生成する画像の指示内容です。この他、以下のプロパティが指定できます。

項目	説明	デフォルト
IMAGE_SIZE	画像のサイズを指定します	384
IMAGE_NUMBER	生成する画像の数を指定します	10
S3_BUCKET	オブジェクトストレージのバケット名を指定します
S3_ENDPOINT	オブジェクトストレージのエンドポイントを指定します
S3_SECRET	オブジェクトストレージのシークレットアクセスキーを指定します
S3_TOKEN	オブジェクトストレージのアクセスキーIDを指定します

S3_で始まる環境変数は、さくらのオブジェクトストレージに画像を保存するためのものです。保存しない場合には指定しなくても問題ありません（Amazon S3なども利用できます）。さくらのオブジェクトストレージを利用する場合には、オブジェクトストレージ サービス基本情報 | さくらのクラウド マニュアルを参照してください。

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

上記タスクで利用したDockerイメージを作成する手順は以下の通りです。完成版はgoofmint/dok-janusにありますので、実装時の参考にしてください。

Dockerfile の作成

deepseek-ai/Janus: Janus-Series: Unified Multimodal Understanding and Generation Modelsの内容に沿って、Dockerfileを作成します。

ベースイメージ

ベースは FROM pytorch/pytorch:2.2.1-cuda11.8-cudnn8-devel です。

FROM pytorch/pytorch:2.2.1-cuda11.8-cudnn8-devel

ライブラリのインストール

Pythonと、必要なライブラリをインストールします。

# 必要パッケージのインストール
RUN apt-get update && \
    apt-get install -y \
        libcupti-dev \
        git \
	wget \
        python3 \
        python3-pip && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

# 作業ディレクトリ
RUN mkdir /app /opt/artifact
WORKDIR /app

# Janus クローン
RUN git clone https://github.com/deepseek-ai/Janus.git .
RUN pip install -r requirements.txt --no-deps

そして、別途必要なライブラリをインストールします。PyTorchやFlashAttentionなどをインストールします。

RUN pip install --no-deps \
      boto3 \
      botocore \
      argparse \
      regex \
      python-dateutil \
      jmespath \
      safetensors>=0.4.3 \
      tokenizers \
      huggingface_hub
RUN pip install --index-url https://download.pytorch.org/whl/cu118 torch==2.2.1
RUN wget https://github.com/Dao-AILab/flash-attention/releases/download/v2.6.0.post1/flash_attn-2.6.0.post1+cu118torch2.2cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
RUN pip install flash_attn-2.6.0.post1+cu118torch2.2cxx11abiFALSE-cp310-cp310-linux_x86_64.whl

後は 後述する runner.py と、 docker-entrypoint.sh をコピーします。

# スクリプトをコピー
COPY runner.py /app/runner.py
COPY docker-entrypoint.sh /docker-entrypoint.sh
RUN chmod +x /docker-entrypoint.sh

# 起動
CMD ["/bin/bash", "/docker-entrypoint.sh"]

Dockerfile全体

Dockerfileの全体は以下の通りです。

FROM pytorch/pytorch:2.2.1-cuda11.8-cudnn8-devel

# 必要パッケージのインストール
RUN apt-get update && \
    apt-get install -y \
        libcupti-dev \
        git \
	wget \
        python3 \
        python3-pip && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

# 作業ディレクトリ
RUN mkdir /app /opt/artifact
WORKDIR /app

# Janus クローン
RUN git clone https://github.com/deepseek-ai/Janus.git .
RUN pip install -r requirements.txt --no-deps

RUN pip install --no-deps \
      boto3 \
      botocore \
      argparse \
      regex \
      python-dateutil \
      jmespath \
      safetensors>=0.4.3 \
      tokenizers \
      huggingface_hub
RUN pip install --index-url https://download.pytorch.org/whl/cu118 torch==2.2.1
RUN wget https://github.com/Dao-AILab/flash-attention/releases/download/v2.6.0.post1/flash_attn-2.6.0.post1+cu118torch2.2cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
RUN pip install flash_attn-2.6.0.post1+cu118torch2.2cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
# スクリプトをコピー
COPY runner.py /app/runner.py
COPY docker-entrypoint.sh /docker-entrypoint.sh
RUN chmod +x /docker-entrypoint.sh

# 起動
CMD ["/bin/bash", "/docker-entrypoint.sh"]

docker-entrypoint.sh の作成

docker-entrypoint.sh はDockerコンテナー起動時に実行するスクリプトです。ここでは環境変数をチェックして、 runner.py を呼び出します。以下の内容で作成します。PROMTは必須で、他は任意です。SAKURA_ ではじまる環境変数は、高火力 DOK実行時に自動的に渡される変数になります。

#!/bin/bash
export TZ=${TZ:-Asia/Tokyo}

# 必須環境変数のチェック
missing=0

check_required() {
  if [ -z "${!1}" ]; then
    echo "Error: $1 is required but not set."
    missing=1
  fi
}

check_required PROMPT
check_required SAKURA_ARTIFACT_DIR
check_required SAKURA_TASK_ID

if [ "$missing" -eq 1 ]; then
  echo "Please set all required environment variables."
  exit 1
fi

# runner.py を引数付きで実行
pushd /app
  python3 runner.py \
    --prompt "$PROMPT" \
    --output "$SAKURA_ARTIFACT_DIR" \
    --id "$SAKURA_TASK_ID" \
    --img_size "${IMAGE_SIZE:-384}" \
    --img_num "${IMAGE_NUMBER:-10}" \
    ${S3_BUCKET:+--s3-bucket "$S3_BUCKET"} \
    ${S3_ENDPOINT:+--s3-endpoint "$S3_ENDPOINT"} \
    ${S3_SECRET:+--s3-secret "$S3_SECRET"} \
    ${S3_TOKEN:+--s3-token "$S3_TOKEN"}
popd

runner.py の作成

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

import os
import argparse
import PIL.Image
import torch
import numpy as np
import boto3
from transformers import AutoModelForCausalLM
from janus.models import MultiModalityCausalLM, VLChatProcessor

パラメータの取得

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

# 引数のパース
parser = argparse.ArgumentParser()
parser.add_argument('--output', type=str, required=True, help='出力先ディレクトリ')
parser.add_argument('--id', type=str, required=True, help='出力ファイル名のID。プリフィックス')
parser.add_argument('--img_size', type=int, default=384, help='出力画像サイズ（1辺）')
parser.add_argument('--img_num', type=int, default=10, help='出力する画像枚数')
parser.add_argument('--prompt', type=str, required=True, help='画像生成プロンプト')
parser.add_argument('--s3-bucket', help='S3のバケットを指定します。')
parser.add_argument('--s3-endpoint', help='S3互換エンドポイントのURLを指定します。')
parser.add_argument('--s3-secret', help='S3のシークレットアクセスキーを指定します。')
parser.add_argument('--s3-token', help='S3のアクセスキーIDを指定します。')

args = parser.parse_args()

オブジェクトストレージ用のオブジェクトを準備

S3_ではじまる環境変数があれば、それを使ってS3オブジェクトを作成します。

s3 = None
if args.s3_token and args.s3_secret and args.s3_bucket:
    # S3クライアントの作成
    s3 = boto3.client(
        's3',
        endpoint_url=args.s3_endpoint if args.s3_endpoint else None,
        aws_access_key_id=args.s3_token,
        aws_secret_access_key=args.s3_secret)

モデルの準備

今回は Janus-1.3B を利用しています。この他、 Janus-Pro-1B や Janus-Pro-7B などが用意されています。

# モデルの準備
model_path = "deepseek-ai/Janus-1.3B"
vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
tokenizer = vl_chat_processor.tokenizer
vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
    model_path, trust_remote_code=True
)
vl_gpt = vl_gpt.to(torch.bfloat16).cuda().eval()

プロンプトの整形

入力されたプロンプトに基づいて、実行するためのプロンプトを作成します。

# プロンプト整形
conversation = [
    {"role": "User", "content": args.prompt},
    {"role": "Assistant", "content": ""},
]
sft_format = vl_chat_processor.apply_sft_template_for_multi_turn_prompts(
    conversations=conversation,
    sft_format=vl_chat_processor.sft_format,
    system_prompt="",
)
prompt = sft_format + vl_chat_processor.image_start_tag

画像生成

画像生成を行う generate 関数を作成します。この関数ではJanusを使って画像生成後、結果を指定したディレクトリへの保存と、必要があればオブジェクトストレージへ保存します。

コード自体はdeepseek-ai/Janus: Janus-Series: Unified Multimodal Understanding and Generation Modelsをカスタマイズしたものです。

@torch.inference_mode()
def generate(
    mmgpt: MultiModalityCausalLM,
    vl_chat_processor: VLChatProcessor,
    prompt: str,
    output_dir: str,
    prefix: str,
    img_size: int = 384,
    img_num: int = 10,
    temperature: float = 1,
    cfg_weight: float = 5,
    image_token_num_per_image: int = 576,
    patch_size: int = 16,
):
    input_ids = vl_chat_processor.tokenizer.encode(prompt)
    input_ids = torch.LongTensor(input_ids)

    tokens = torch.zeros((img_num * 2, len(input_ids)), dtype=torch.int).cuda()
    for i in range(img_num * 2):
        tokens[i, :] = input_ids
        if i % 2 != 0:
            tokens[i, 1:-1] = vl_chat_processor.pad_id

    inputs_embeds = mmgpt.language_model.get_input_embeddings()(tokens)

    generated_tokens = torch.zeros((img_num, image_token_num_per_image), dtype=torch.int).cuda()

    for i in range(image_token_num_per_image):
        outputs = mmgpt.language_model.model(
            inputs_embeds=inputs_embeds,
            use_cache=True,
            past_key_values=outputs.past_key_values if i != 0 else None
        )
        hidden_states = outputs.last_hidden_state

        logits = mmgpt.gen_head(hidden_states[:, -1, :])
        logit_cond = logits[0::2, :]
        logit_uncond = logits[1::2, :]

        logits = logit_uncond + cfg_weight * (logit_cond - logit_uncond)
        probs = torch.softmax(logits / temperature, dim=-1)

        next_token = torch.multinomial(probs, num_samples=1)
        generated_tokens[:, i] = next_token.squeeze(dim=-1)

        next_token = torch.cat([next_token.unsqueeze(dim=1)] * 2, dim=1).view(-1)
        img_embeds = mmgpt.prepare_gen_img_embeds(next_token)
        inputs_embeds = img_embeds.unsqueeze(dim=1)

    dec = mmgpt.gen_vision_model.decode_code(
        generated_tokens.to(dtype=torch.int),
        shape=[img_num, 8, img_size // patch_size, img_size // patch_size]
    )
    dec = dec.to(torch.float32).cpu().numpy().transpose(0, 2, 3, 1)
    dec = np.clip((dec + 1) / 2 * 255, 0, 255).astype(np.uint8)

    os.makedirs(output_dir, exist_ok=True)
    for i in range(img_num):
        save_path = os.path.join(output_dir, f"{prefix}_{i}.jpg")
        PIL.Image.fromarray(dec[i]).save(save_path)
        if s3 is not None:
            s3.upload_file(
                Filename=save_path,
                Bucket=args.s3_bucket,
                Key=os.path.basename(save_path))

関数を実行

最後に、 generate 関数を実行します。

generate(
    vl_gpt,
    vl_chat_processor,
    prompt,
    output_dir=args.output,
    prefix=args.id,
    img_size=args.img_size,
    img_num=args.img_num,
)

全体の処理

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

import os
import argparse
import PIL.Image
import torch
import numpy as np
import boto3
from transformers import AutoModelForCausalLM
from janus.models import MultiModalityCausalLM, VLChatProcessor

# 引数のパース
parser = argparse.ArgumentParser()
parser.add_argument('--output', type=str, required=True, help='出力先ディレクトリ')
parser.add_argument('--id', type=str, required=True, help='出力ファイル名のID。プリフィックス')
parser.add_argument('--img_size', type=int, default=384, help='出力画像サイズ（1辺）')
parser.add_argument('--img_num', type=int, default=10, help='出力する画像枚数')
parser.add_argument('--prompt', type=str, required=True, help='画像生成プロンプト')
parser.add_argument('--s3-bucket', help='S3のバケットを指定します。')
parser.add_argument('--s3-endpoint', help='S3互換エンドポイントのURLを指定します。')
parser.add_argument('--s3-secret', help='S3のシークレットアクセスキーを指定します。')
parser.add_argument('--s3-token', help='S3のアクセスキーIDを指定します。')

args = parser.parse_args()

s3 = None
if args.s3_token and args.s3_secret and args.s3_bucket:
    # S3クライアントの作成
    s3 = boto3.client(
        's3',
        endpoint_url=args.s3_endpoint if args.s3_endpoint else None,
        aws_access_key_id=args.s3_token,
        aws_secret_access_key=args.s3_secret)

# モデルの準備
model_path = "deepseek-ai/Janus-1.3B"
vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
tokenizer = vl_chat_processor.tokenizer
vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
    model_path, trust_remote_code=True
)
vl_gpt = vl_gpt.to(torch.bfloat16).cuda().eval()

# プロンプト整形
conversation = [
    {"role": "User", "content": args.prompt},
    {"role": "Assistant", "content": ""},
]
sft_format = vl_chat_processor.apply_sft_template_for_multi_turn_prompts(
    conversations=conversation,
    sft_format=vl_chat_processor.sft_format,
    system_prompt="",
)
prompt = sft_format + vl_chat_processor.image_start_tag

@torch.inference_mode()
def generate(
    mmgpt: MultiModalityCausalLM,
    vl_chat_processor: VLChatProcessor,
    prompt: str,
    output_dir: str,
    prefix: str,
    img_size: int = 384,
    img_num: int = 10,
    temperature: float = 1,
    cfg_weight: float = 5,
    image_token_num_per_image: int = 576,
    patch_size: int = 16,
):
    input_ids = vl_chat_processor.tokenizer.encode(prompt)
    input_ids = torch.LongTensor(input_ids)

    tokens = torch.zeros((img_num * 2, len(input_ids)), dtype=torch.int).cuda()
    for i in range(img_num * 2):
        tokens[i, :] = input_ids
        if i % 2 != 0:
            tokens[i, 1:-1] = vl_chat_processor.pad_id

    inputs_embeds = mmgpt.language_model.get_input_embeddings()(tokens)

    generated_tokens = torch.zeros((img_num, image_token_num_per_image), dtype=torch.int).cuda()

    for i in range(image_token_num_per_image):
        outputs = mmgpt.language_model.model(
            inputs_embeds=inputs_embeds,
            use_cache=True,
            past_key_values=outputs.past_key_values if i != 0 else None
        )
        hidden_states = outputs.last_hidden_state

        logits = mmgpt.gen_head(hidden_states[:, -1, :])
        logit_cond = logits[0::2, :]
        logit_uncond = logits[1::2, :]

        logits = logit_uncond + cfg_weight * (logit_cond - logit_uncond)
        probs = torch.softmax(logits / temperature, dim=-1)

        next_token = torch.multinomial(probs, num_samples=1)
        generated_tokens[:, i] = next_token.squeeze(dim=-1)

        next_token = torch.cat([next_token.unsqueeze(dim=1)] * 2, dim=1).view(-1)
        img_embeds = mmgpt.prepare_gen_img_embeds(next_token)
        inputs_embeds = img_embeds.unsqueeze(dim=1)

    dec = mmgpt.gen_vision_model.decode_code(
        generated_tokens.to(dtype=torch.int),
        shape=[img_num, 8, img_size // patch_size, img_size // patch_size]
    )
    dec = dec.to(torch.float32).cpu().numpy().transpose(0, 2, 3, 1)
    dec = np.clip((dec + 1) / 2 * 255, 0, 255).astype(np.uint8)

    os.makedirs(output_dir, exist_ok=True)
    for i in range(img_num):
        save_path = os.path.join(output_dir, f"{prefix}_{i}.jpg")
        PIL.Image.fromarray(dec[i]).save(save_path)
        if s3 is not None:
            s3.upload_file(
                Filename=save_path,
                Bucket=args.s3_bucket,
                Key=os.path.basename(save_path))

# 画像生成
generate(
    vl_gpt,
    vl_chat_processor,
    prompt,
    output_dir=args.output,
    prefix=args.id,
    img_size=args.img_size,
    img_num=args.img_num,
)

Dockerイメージのビルド

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

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

Dockerイメージを登録するコンテナレジストリを作成します。さくらのクラウドでは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 の部分は、作成したコンテナレジストリのドメイン名に置き換えてください。また、 janus は任意の名前で大丈夫です（以下はその名称で読み替えてください）。

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

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

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

sudo docker login EXAMPLE.sakuracr.jp

イメージのプッシュ

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

sudo docker push EXAMPLE.sakuracr.jp/janus:latest

タスクを作成する

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

項目	設定
イメージ	dok-handson.sakuracr.jp/janus
環境変数	PROMPT = A stunning boy from kabul in red, white traditional clothing, blue eyes, brown hair IMAGE_NUM = 20

実行が完了し、画像が生成できていれば成功です。

まとめ

今回はJanusを使って、高火力 DOK上で画像生成を行いました。まずは実行できるのみ、次にDockerイメージの作成と段階的に進められるようにしています。画像生成のように、処理に時間がかかるものを利用する際に高火力 DOKは便利です。

高火力 DOKはタスクを多数立ち上げて、後は結果を待つのみと言った使い方ができます。ぜひAI・機械学習に活用してください。

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