HuggingFaceとDatabricksApps、Mosaic AIを用いた生成AIアプリの構築

こちらのブログで紹介されている内容をウォークスルーします。猫の画像を見て手が止まりませんでした。

https://www.databricks.com/jp/blog/zero-genai-hero-building-your-genai-app-huggingface-and-databricks

ブログの内容そのままでは動かないので、適宜修正、補足しています。

ここで紹介されているのは、ベースとなる画像に対してプロンプトを与えることで、それに応じて画像を加工できるアプリを作るというものです。

# HuggingFaceモデルのモデルサービングエンドへのデプロイ

この章では、HuggingFaceモデルをDatabricksの[Mosaic AIモデルサービング](https://docs.databricks.com/aws/ja/machine-learning/model-serving/)にデプロイすることが主な作業となります。これによって、アプリからモデルにアクセスできるようになります。

ノートブックを作成して、サーバレスコンピュートにアタッチします。

```py
%pip install diffusers 
%pip install -U mlflow databricks-sdk
%restart_python
```

## サンプル画像をボリュームにロード

```py
import diffusers
from diffusers.utils import load_image, make_image_grid
from databricks.sdk import WorkspaceClient
from io import BytesIO

init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png")

# init_imageがPIL Imageオブジェクトであると仮定
w = WorkspaceClient()

# PIL Imageをバイトに変換
img_byte_arr = BytesIO()
init_image.save(img_byte_arr, format='PNG')  # または 'JPEG' を使用
img_byte_arr = img_byte_arr.getvalue()

# ボリューム内の保存先パスを定義
destination_path = "/Volumes/users/takaaki_yayoi/data/image_submitted.PNG"

# 画像をボリュームにアップロード
w.files.upload(destination_path, img_byte_arr, overwrite=True)
```

アップロードされました。
![Screenshot 2025-04-03 at 13.44.18.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/fb733011-d75b-4dc6-a46d-76a60d4bef1e.png)

## HuggingFaceからのモデルをUnity CatalogにMLflowで登録

モデルの前処理、後処理を含むメソッドを持つクラスで生成AIモデルをラッピングします。

使っているモデルはこちらです。

https://huggingface.co/kandinsky-community/kandinsky-2-2-decoder

```py
import mlflow

class StableDiffusionImgToImg(mlflow.pyfunc.PythonModel):
   def __init__(self):
       self.pipe = None

   def load_context(self, context):
       import torch
       from diffusers import AutoPipelineForImage2Image
       # モデルを初期化する。正しいデバイスで行うことを確認
       self.pipe = AutoPipelineForImage2Image.from_pretrained(
           "kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16, use_safetensors=True
       )
       self.pipe = self.pipe.to("cuda")

   def image_to_base64(self, image):
       from io import BytesIO
       import base64
       buffered = BytesIO()
       image.save(buffered, format="JPEG")
       return base64.b64encode(buffered.getvalue()).decode('utf-8')

   def base64_to_image(self, base64_string):
       from io import BytesIO
       import base64
       import PIL
       from PIL import Image
       # base64文字列をデコード
       img_data = base64.b64decode(base64_string)
  
       # デコードされたデータからBytesIOオブジェクトを作成
       buffer = BytesIO(img_data)
  
       # PILを使用して画像を開く
       image = Image.open(buffer)
  
       return image

   def predict(self, context, model_input):
       prompt = model_input["prompt"][0]
       init_image = self.base64_to_image(model_input["init_image"][0])
       negative_prompt = model_input["negative_prompt"][0]
       strength = model_input["strength"][0]
       guidance_scale = model_input["guidance_scale"][0]

       # 画像生成
       image = self.pipe(prompt, negative_prompt=negative_prompt, image=init_image, strength=strength, guidance_scale=guidance_scale).images[0]

       return self.image_to_base64(image)
```

MLflowでモデルを記録します。

```py
import base64
from mlflow.models import ModelSignature, infer_signature
from mlflow.types.schema import Schema, ColSpec, DataType
import pandas as pd

mlflow.set_registry_uri('databricks-uc')

def load_image_from_volume(volume_path):
     import PIL
     from PIL import Image
     with Image.open(volume_path) as img:
       return img.convert("RGB")

def image_to_base64(image):
   buffered = BytesIO()
   image.save(buffered, format="JPEG")
   return base64.b64encode(buffered.getvalue()).decode('utf-8')

# 入力スキーマ
input_schema = Schema([ColSpec(DataType.string, "prompt"),
                      ColSpec(DataType.string, "negative_prompt"),
                      ColSpec(DataType.string, "init_image"),
                      ColSpec(DataType.float, "strength"),
                      ColSpec(DataType.float, "guidance_scale")])

# 出力スキーマ
output_schema = Schema([ColSpec(DataType.string, "image")])

# ボリュームから画像を読み込み、base64に変換
image = image_to_base64(load_image_from_volume("/Volumes/users/takaaki_yayoi/data/image_submitted.PNG"))

# 入力例を定義
input_example=pd.DataFrame({"prompt":["cat grey bearded wizard, detailed, fantasy, cute, adorable, animated, cgi, entertainment, 8K"],
                           "negative_prompt" : ["deformed, ugly, blurry, low-quality, low-resolution"],
                           "init_image" : [image],
                           "strength" : [0.1],
                           "guidance_scale" : [18.0]})

# シグネチャを推定
signature = infer_signature(input_example, image)

# モデルの詳細（アーティファクト、pip要件、入力例など）と共にモデルをログ
with mlflow.start_run() as run: 
   mlflow.pyfunc.log_model(
       "model",
       python_model=StableDiffusionImgToImg(),
       input_example=input_example,
       signature=signature,
       registered_model_name="users.takaaki_yayoi.kandinsky-model",
       pip_requirements=["transformers", "torch", "accelerate", "diffusers", "huggingface_hub==0.23.0", "invisible-watermark>=0.2.0", "transparent-background", "opencv-python-headless"]
   )
```

入出力のスキーマと共に記録されました。

![Screenshot 2025-04-03 at 13.47.36.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/ea109265-e4cc-4009-b02a-c8f662971dfd.png)

## モデルサービングエンドポイントへのデプロイ

画像系のモデルなのでGPUを指定します。

```py
from mlflow.deployments import get_deploy_client

client = get_deploy_client("databricks")

endpoint = client.create_endpoint(
   name="taka-kandinsky",
   config={
       "served_entities": [{
           "entity_name": "users.takaaki_yayoi.kandinsky-model",
           "entity_version": "1",
           "workload_type": "GPU_SMALL",
           "workload_size": "Small",
           "scale_to_zero_enabled": True
       }],
       "traffic_config": {
           "routes": [{
               "served_model_name": "kandinsky-model-1",
               "traffic_percentage": 100
           }]
       }
   }
)
```

しばらく待つとエンドポイントが起動します。これでモデル側の準備は整い、REST API経由でモデルを利用できるようになりました。

![Screenshot 2025-04-03 at 13.49.38.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/8869e40f-c1ac-4d76-b87b-a41186600ce7.png)

# Databricks Appsによるフロントエンドの構築

[Databricks Apps]()を利用してフロントエンドとDatabricksリソースへのアクセスを構成します。今度は別のファイルとして`main.py`を作成します。また、タイトル画像を格納するための`static`フォルダも作成します。

![Screenshot 2025-04-03 at 13.51.37.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/05e915f7-64a5-4c13-a311-9f6d68a6adae.png)

:::note info
**注意**
執筆時点ではDatabricks Appsはまだ日本リージョンに来ていません。
:::


## アプリの実装

`main.py`ファイルの中身は以下のようになります。前のノートブックで実装した`image_to_base64`などの前処理、後処理はこちらでも必要となります。

```py:main.py
import io
import os
import PIL
from PIL import Image
from io import BytesIO
from databricks.sdk import WorkspaceClient
import base64
import pandas as pd
import json, requests

VOLUME_URI = os.getenv("VOLUME_URI", "/Volumes/users/takaaki_yayoi/data")

w = WorkspaceClient()

def image_to_base64(image):
   buffered = BytesIO()
   image.save(buffered, format="JPEG")
   return base64.b64encode(buffered.getvalue()).decode('utf-8')

def base64_to_image(base64_string):
    # Decode the base64 string
    img_data = base64.b64decode(base64_string)
    # Create a BytesIO object from the decoded data
    buffer = BytesIO(img_data)
    # Open the image using PIL
    image = Image.open(buffer)
  
    return image

def query_image_from_volume():
  
   image_name="image_submitted.PNG"

   image_path = os.path.join(VOLUME_URI, image_name)

   # シークレットから読み込んだパーソナルアクセストークンを環境変数に設定
   os.environ["TOKEN"] = w.dbutils.secrets.get(scope="demo-token-takaaki.yayoi", key="pat")

   response = w.files.download(image_path)
   image_data = io.BytesIO(response.contents.read())
   
   with Image.open(image_data) as img:
     return img.convert("RGB")

def query_model_endpoint(prompt):
   
    img_base64 = image_to_base64(query_image_from_volume())

    input_example=pd.DataFrame({"prompt":[prompt],
                                "negative_prompt" : [""],
                                "init_image" : [img_base64],
                                "strength" : [0.1],
                                "guidance_scale" : [18]})
  
    # モデルサービングエンドポイントへのアクセス
    url = 'https://e2-demo-west.cloud.databricks.com/serving-endpoints/taka-kandinsky/invocations'
    headers = {'Authorization': f'Bearer {os.environ.get("TOKEN")}', 'Content-Type': 'application/json'}
    ds_dict = {'dataframe_split': input_example.to_dict(orient='split')}
    data_json = json.dumps(ds_dict, allow_nan=True)
    response = requests.request(method='POST', headers=headers, url=url, data=data_json)
    final_img = base64_to_image(response.json()['predictions'])

    return final_img.convert("RGB")


import gradio as gr

with gr.Blocks() as demo:
   with gr.Row():
       gr.set_static_paths(paths=["static/databricks-logo.png"])
       gr.HTML(value='<img src="/file=static/databricks-logo.png", width="450", class="Databricks_logo"">', elem_id="Databricks Logo")
       gr.Markdown(value="# Picture Yourself in Generative AI")

   with gr.Row():
       inp = gr.Textbox(label="Character", placeholder="Enter the Character:")
  
   with gr.Row():
       image_ouput1 = gr.Image(value=query_image_from_volume, interactive=False, label="Before")
       image_output2 = gr.Image(label="After", show_download_button=True)
   with gr.Row():
       btn = gr.Button("Run")
       clear_button = gr.ClearButton([image_ouput1, image_output2, inp])
  
   btn.click(fn=query_model_endpoint, inputs=[inp], outputs=image_output2)


if __name__ == "__main__":
   demo.launch()
```

## リソースに対するアクセス権設定

ここで注意しなくてはならないのは、**DatabricksアプリからDatabricks管理リソースへのアクセス**です。

Databricksアプリは、アプリの作成とともに作成される**サービスプリンシパルの権限**で動作し、その権限でDatabricks管理のリソース(ボリュームなど)にアクセスします。このため、リソースに対するサービスプリンシパルの権限を適切に設定する必要があります。

![Screenshot 2025-04-03 at 13.58.27.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/ad727df8-44c4-49c6-b85b-41d40963fd58.png)

今回の例では、ボリュームに格納されている画像をアプリから読み込んでいるので、当該ボリュームではサービスプリンシパルに対して**READ_VOLUME**の権限を付与します。(実際にはアプリを一度作成してそのサービスプリンシパルの権限を後で付与します)

そして、モデルサービングエンドポイントへのアクセスには[パーソナルアクセストークン](https://docs.databricks.com/aws/ja/dev-tools/auth/pat)が必要です。これを平文でアプリに埋め込むのではなく、[シークレット](https://docs.databricks.com/aws/ja/security/secrets)で安全に管理されているパーソナルアクセストークンを参照するようにします。そのためには事前にパーソナルアクセストークンをシークレットに格納しておきます。私の場合はスコープ`demo-token-takaaki.yayoi`のキー`pat`にパーソナルアクセストークンを格納しています。

これらの準備ができたら、以下のように使用するリソースを指定してアプリを作成し、上記`main.py`が格納されているパスを指定してアプリをデプロイします。

![Screenshot 2025-04-03 at 14.06.28.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/ed58cc97-a8f1-453e-83b9-e8883974dee7.png)

稼働したら、上述のアクセス権を確認の上、アプリにアクセスします。

![Screenshot 2025-04-03 at 13.22.32.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/38359133-b3d1-45c6-a274-30225914765f.png)

`blue colored realistic wizard cat`のようなプロンプトを与えます。

![Screenshot 2025-04-03 at 14.11.38.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/97a9d58c-01f9-4691-beaa-b716e2c8a250.png)

動きました！

`red angry cat`

![Screenshot 2025-04-03 at 19.06.41.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/741ecf8f-f04c-4e7e-ab91-370e817dd4c9.png)

`purple sad cat`

![Screenshot 2025-04-03 at 19.07.25.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/ef6faf13-e45c-49a5-a1c9-51fe42e10fc1.png)

なにこれ楽しい

### はじめてのDatabricks

[はじめてのDatabricks](https://qiita.com/taka_yayoi/items/8dc72d083edb879a5e5d)

### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)