databricks Model Servingで画像生成モデルを動かしてみました - CCCMKホールディングス TECH LABの Tech Blog

こんにちは、CCCMKホールディングス三浦です。

最近は色々な画像生成モデルが出てきています。その中にはHugging Faceで公開されているものもあります。そういったモデルを試すことが出来る環境が欲しいな、と考えていました。

このブログでも何度か紹介しているdatabricksのModel Servingでは、T4もしくはA100が搭載されたサーバを選択することが出来、それらを使うことで画像生成モデルをModel Servingで動かしてAPI経由で利用することも出来るのでは、と考え、今回実際に試してみようと思いました。

結果、狙い通りにModel Servingで画像生成モデル(stable-diffusion-xl-base-1.0)を動かすことが出来ました。ただ、結構詰まるポイントが個人的に多かったので、今回の記事ではそういった点に触れつつ手順をまとめていこうと思います。

Model Servingにデプロイするとdatabricks-claude-3-7-sonnetなどの基盤モデルと組み合わせてチャット形式で画像生成が出来るアプリケーションを作る、といったことも可能になります。アプリケーションはDatabricks Appsで動かしています。

手順の概要について

基本的には前回のPDFやPPTXからテキストを抽出するAPIを作った時と同様の手順になります。画像生成の機能をMLflowのPyFuncモデルに実装し、log_modelでExperimentに記録し、register_modelでUnity Catalogに登録したのちModel Servingにデプロイします。

※前回の記事はこちら

techblog.cccmkhd.co.jp

この手順の中で特に詰まることが多かったのがModel Servingにデプロイするステップです。今回画像生成モデルを利用するためにHugging Faceの"Diffusers"というライブラリを使ったのですが、このライブラリと関連するライブラリのバージョンが少しでも違うとデプロイに失敗することが頻繁に発生しました。ですので全てのステップで主要なライブラリのバージョンを統一させることが重要です。

またModel Servingのサーバのスペックをそれなりに高いものにしないとデプロイが制限時間内に完了せずに失敗することがあるのでこちらも重要な確認ポイントです。

Modelを登録するまで

PyFuncモデルの定義からUnity Catalogへの登録はNotebookで実行しました。Notebookを実行した環境は次の通りです。

前準備

Notebook上で、まず最初に元々インストールされているflash-attnのアンインストールを実行しました。もともと入っているflash-attnとなんらかの不整合があるのか、この手順を取らないとモデルのダウンロードに失敗する現象が発生しました。最適なソリューション・・・とは思えないのですが、応急処置で次のような処理を実行しました。

%pip uninstall -y flash-attn
!rm -rf /databricks/python3/lib/python3.12/site-packages/flash_attn*
!rm -rf /databricks/python3/lib/python3.12/site-packages/flash_attn-2*
dbutils.library.restartPython()

次に必要なライブラリをインストールしました。

%pip install "diffusers==0.35.1" "huggingface-hub==0.34.4" "transformers==4.51.3" "mlflow==3.2.0"
dbutils.library.restartPython()

PyFuncモデルの定義

PyFuncモデル本体の定義をします。モデルのメインの処理を書くpredict関数はプロンプトや画像生成パラメータを辞書型(の配列)で受け取り、生成画像はbase64文字列に変換して返すようにしました。

DiffusersのモデルはPillowのImageで返すため、APIで利用できるようbase64に変換する関数をutils.pyに以下のように実装しておきました。

import base64
from io import BytesIO
from PIL import Image

def pil_to_base64(pil_image, format="PNG"):
    buf = BytesIO()
    pil_image.save(buf, format=format)
    return base64.b64encode(buf.getvalue()).decode("utf-8")

以下がモデルの本体の実装です。

import os
import torch
import mlflow
import mlflow.pyfunc
from diffusers import DiffusionPipeline
from typing import Any, Dict, List

from utils import pil_to_base64

class DiffusersPyfunc(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        dtype = torch.float16 if torch.cuda.is_available() else torch.float32
        self.model = DiffusionPipeline.from_pretrained(
            "stabilityai/stable-diffusion-xl-base-1.0",
            torch_dtype=dtype, 
            use_safetensors=True
        )
        if torch.cuda.is_available():
            self.model = self.model.to("cuda")
        else:
            self.model = self.model.to("cpu")
        
    
    def predict(self, context, model_input: List[Dict[str, Any]])->List[Dict[str, str]]:

        # 画像生成時のパラメータは`model_input`で受け取る。
        model_input = model_input[0]
        prompt = model_input.get("prompt")
        negative_prompt = model_input.get("negative_prompt",None)
        num_inference_steps = model_input.get("num_inference_steps", 30)
        guidance_scale = model_input.get("guidance_scale", 7.0)
        height = model_input.get("height", 512)
        width = model_input.get("width", 512)
        seed = model_input.get("seed",None)

        # 生成結果の固定
        generator = None
        if seed is not None:
            generator = torch.Generator(
                device="cuda" if torch.cuda.is_available() else "cpu"
            ).manual_seed(seed)

        # 画像生成処理
        out = self.model(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            height=height,
            width=width,
            generator=generator
        )
        img = out.images[0]
        b64 = pil_to_base64(img)
        return [{"image_base64": b64}]

ローカルでテストをする

定義したモデルが正しく動作するか、ローカルでテストを行います。Model Servingのデプロイは結構時間がかかります。デプロイしたのにモデルの定義がおかしくて動作しない、といった事態を防ぐために、ローカルでのテストは必須だと感じました。

# Test Model
import base64
import matplotlib.pyplot as plt
from mlflow.pyfunc import PythonModelContext
from PIL import Image

model = DiffusersPyfunc()

# ローカル実行時は空のcontextを渡す。
test_context = PythonModelContext(artifacts={}, model_config={})
model.load_context(test_context)

model_input = [{
    "prompt": "a high quality photo of a cat"
}]

b64_image = model.predict(None, model_input)[0]
with open("sample.png","wb") as f:
    f.write(base64.b64decode(
        b64_image[0]["image_base64"])
    )
plt.imshow(Image.open("sample.png"))

画像が表示されたら確認OKです。

モデルの登録

Unity Catalogへのモデルの登録を実行します。MLflowにlog_modelで記録した後はGUIで登録する方法があるのですが、その方法で登録しようとすると1日たっても処理が完了しませんでした。おそらくモデルのサイズが大きいことも影響しているのでは、と考えています。

次のようなコードを実行すると、数分で登録が完了します。また、log_modelはpip_requirementsで明示的に使用するライブラリを指定することが出来ます。モデルの動作が確認できたローカルの環境とModel Servingの環境をなるべく同一にするため、必ず指定したほうが良いと思います。

import mlflow

# 主要ライブラリのバージョンを統一する
requirements = [
    'diffusers==0.35.1',
    'huggingface-hub==0.34.4',
    'transformers==4.51.3'
]

# 入力例
input_example = {
    "prompt": "a high quality photo of a cat",
    "negative_prompt": None,
    "num_inference_steps": 30,
    "guidance_scale": 7.0,
    "height": 512,
    "width": 512,
    "seed": 42
}

model_name = "image_generator"

with mlflow.start_run() as run:
    model = DiffusersPyfunc()
    # Experimentへのモデルの記録
    model_info = mlflow.pyfunc.log_model(
        name=model_name,
        code_paths=["utils.py"],
        python_model=model,
        input_example=[input_example],
        pip_requirements=requirements,
    )
    # Unity Catalogへの登録
    model_uri = model_info.model_uri
    reg = mlflow.register_model(model_uri, f"{db}.{schema}.{model_name}")

Model Servingにデプロイする

Unity Catalogで対象のモデルを開き、Model Servingにデプロイします。具体的な手順は前回の記事の内容と同様です。

前回のモデルと異なり、今回のモデルはデプロイにかなりの時間が必要になります。実は以下の記事にあるようにデプロイにはタイムアウトの時間が設定されていて、既定の時間内に処理が完了できないとエラーになってしまいます。

docs.databricks.com

CPU及びGPUのsmallだと処理能力も制限されている上にタイムアウトの時間も短いため、今回のモデルをデプロイすることが出来ませんでした。

色々試した結果、"Compute type"を"GPU Large"(A100), "Compute scale-out"を"Small"に設定してデプロイを成功させることが出来ました。

使ってみる！

Model Servingにリクエストを送るのに次のようなコードを実行しました。

# テスト

# エンドポイントURL
url = f"{workspace_host}/serving-endpoints/image-generator/invocations"

# アクセストークンを取得しておく
token = os.environ.get("DATABRICKS_TOKEN")

# header
headers = {
    "Authorization": f"Bearer {token}",
    "Content-Type": "application/json"
}

# request body
input = {
    "inputs": [
        {
        "prompt": "cute cat",
        "negative_prompt": None,
        "num_inference_steps": 30,
        "guidance_scale": 7,
        "height": 512,
        "width": 512,
        "seed": 42
        }
    ]
}

try:
    response = requests.post(url, headers=headers, data=json.dumps(input))
    response_json = response.json()
except Exception as e:
    print(f"Error: {e}")
    print(response.text)

# base64で受け取り、bytesに変換する
image_base64 = response_json["predictions"][0]["image_base64"]
image_bytes = base64.b64decode(image_base64)

# ファイルに書き込む
fname = "sample.png"
with open(fname, "wb") as f:
    f.write(image_bytes)