🐕

カスタムDocker-Imageを作ってLCM-LoRAを使った画像生成APIの高速化

2023/12/03に公開

はじめに

前にLCM-LoRAを使った画像生成APIを作成する記事を書きました。こちらによって任意の画像生成モデル + LCM-LoRAを使って高速?で画像がAPI経由でできるようになりました。
しかし、画像の生成速度がAPIごとに30秒~60秒ほどかかっていました。
そこで高速化を試していきたいと思います

https://zenn.dev/midra_lab/articles/521a753373b3e3

処理の計測

だいたい計測してみると以下のような感じになります。

処理項目 処理時間(秒)
モデルのダウンロード 7.53
一時ファイルの処理時間 (キャッシュ移行含む) 7.022
モデルのロード 19.123
LoRAウェイトのロード 1.210
画像生成 2.138
画像のエンコード 0.031
合計時間 37.057

(ただ使用しているModalというサービスがコールドスタートの設定があるため、コールドスタート内であればもっと早いです)
https://modal.com/docs/guide/cold-start

最適化のアプローチ(カスタムDockerImageの作成)

今回は、いろいろ調べてより簡単にできそうな?モデルのダウンロード部分を最適化していきます

API上で使用しているDockerImageが以下のようになっています。そのため、ライブラリのインストールやモデルのダウンロードをdocker image側に事前に入れてしまえば速くなるのではないかと考えました

modal.Image.debian_slim().pip_install("omegaconf", "diffusers==0.24.0", "torch", "transformers"),

そこで必要なライブラリやモデルをimage内に入れてしまって、DockerHubにアップロードする方法で進めていきます

カスタムImageの作成

今回必要なのは、以下になります

  • pipでインストールしているライブラリ
  • 画像生成モデル(今回は Counterfeit-V3.0)

上記を含めたImageを作っていきます

Dockerfile

# ベースイメージ
FROM python:3.10.12-slim

# 作業ディレクトリの設定
WORKDIR /usr/src/app

# 必要なライブラリのインストール
COPY requirements.txt ./
RUN pip install --no-cache-dir -r requirements.txt

# Hugging Faceからモデルをダウンロード
RUN apt-get update && apt-get install -y wget
RUN wget https://huggingface.co/gsdf/Counterfeit-V3.0/resolve/main/Counterfeit-V3.0_fix_fp16.safetensors -O Counterfeit-V3.0_fix_fp16.safetensors

requirements.txt

omegaconf
diffusers==0.24.0
torch
transformers
accelerate

作成したものは、DockerHubにpushしておきます

https://hub.docker.com/repository/docker/midralab/generate-image-api-for-lcm/general

DockerHubへのpush及び自動化

Imageを作った後にほかからpullするためにDockerHubにアップロードします。しかし、毎回手動でするのは大変なので、GitHub Actionsを使った自動化を行います。

GitHub Actionsでは、以下のように処理をしてDockerHubにpushしてくれます

https://github.com/MidraLab/lcm-docker-image/actions/runs/7070811478/job/19247861174

name: Docker Build and Publish

on:
  push:
    branches:
      - main

jobs:
  build-and-push:
    runs-on: ubuntu-latest
    steps:
    - name: Checkout Repository
      uses: actions/checkout@v4.1.1

    - name: Login to Docker Hub
      uses: docker/login-action@v3.0.0
      with:
        username: ${{ secrets.DOCKER_HUB_USERNAME }}
        password: ${{ secrets.DOCKER_HUB_ACCESS_TOKEN }}

    - name: Build and Push Docker Image
      uses: docker/build-push-action@v5.1.0
      with:
        context: .
        file: ./Dockerfile
        push: true
        tags: midralab/generate-image-api-for-lcm:latest

カスタムImageを作る話は、以下の記事に書きましたのでこちらをご確認ください

https://ayousanz.hatenadiary.jp/entry/2023/12/02/184936

MidraLab(ミドラボ)

Discussion