Skip to content

Instantly share code, notes, and snippets.

@c-bata
Last active October 28, 2023 01:04
Show Gist options
  • Save c-bata/0eed0dfb416a6994fa30fb23bb38d3ad to your computer and use it in GitHub Desktop.
Save c-bata/0eed0dfb416a6994fa30fb23bb38d3ad to your computer and use it in GitHub Desktop.
W&B 東京Meetup #3 - Optunaを使ったHuman-in-the-loop最適化の紹介
# Requirements:
# $ pip install optuna optuna-dashboard[preferential] diffusers transformers accelerate scipy safetensors xformers
#
# Process A: Launch a process that suggest new params and generate images.
# $ python main.py
#
# Process B: Launch an Optuna Dashboard process.
# $ optuna-dashboard sqlite:///db.sqlite3 --artifact-dir ./artifact
import time
import optuna
from optuna.trial import TrialState
import os
from optuna_dashboard import set_objective_names, register_objective_form_widgets, ChoiceWidget, save_note
from optuna_dashboard import save_note
from optuna_dashboard.artifact import upload_artifact, get_artifact_path
from optuna_dashboard.artifact.file_system import FileSystemBackend
from diffusers import StableDiffusionImg2ImgPipeline
import torch
from PIL import Image
device = "cuda:0"
torch_dtype = torch.float16
init_img = Image.open("./input.png")
init_img = init_img.resize((768, 768))
init_img = init_img.convert("RGB")
# 画像(Artifact)のアップロード先の設定 - 今回はFileSystemBackendを用いて、 "artifact" ディレクトリ以下に保存。
base_path = os.path.join(os.path.dirname(__file__), "artifact")
artifact_backend = FileSystemBackend(base_path=base_path)
def suggest_and_generate_image(study: optuna.Study, pipe: StableDiffusionImg2ImgPipeline):
# OptunaのTrialを生成し、パラメーターおよびプロンプトに含めるキーワードをサンプル
trial = study.ask()
guidance_scale = trial.suggest_float("guidance_scale", 1, 50)
strength = trial.suggest_float("strength", 0.70, 1.0)
num_inference_steps = trial.suggest_int("num_inference_steps", 5, 100)
prompts = ["a mascot character with two eyes and a mouth"]
prompts.append(trial.suggest_categorical("adjectives", ["cute", "funny", "memorable", "charming", "entertaining"]))
prompts.append(trial.suggest_categorical("style", ["anime", "photo", "painting", ""]))
prompts.append(trial.suggest_categorical("facial-expression", ["smiling", "frowning", "grinning", ""]))
negative_prompt = []
negative_prompt.append(trial.suggest_categorical("negative-quality", ["unnatural", "low-quality", ""]))
negative_prompt.append(trial.suggest_categorical("negative-adjectives", ["dull", "boring", "unfriendly", ""]))
# img2imgの実行
images = pipe(
", ".join(prompts),
negative_prompt=", ".join(negative_prompt),
generator=torch.Generator(device).manual_seed(0),
strength=strength,
image=init_img,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
num_images_per_prompt=3,
).images
# 画像アップロードとURLパスの取得
url_paths = []
for i, image in enumerate(images):
image_path = f"tmp/sample-{trial.number}-{i}.png"
image.save(image_path)
artifact_id = upload_artifact(artifact_backend, trial, image_path)
url_path = get_artifact_path(trial, artifact_id)
url_paths.append(url_path)
# ノート (Markdown Format) の保存
# Optuna DashboardのTrial Listタブから、各Trialを選択した際に表示されます。
note = f"""## Prompt
* Prompt: {", ".join(prompts)}
* Negative Prompt: {negative_prompt}
## Images
| Seed 0 | Seed 1 | Seed 2 |
| ------ | ------ | ------ |
| ![image0]({url_paths[0]}) | ![image1]({url_paths[1]}) | ![image2]({url_paths[2]}) |
"""
save_note(trial, note)
def main():
study = optuna.create_study(
storage="sqlite:///db.sqlite3",
sampler=optuna.samplers.TPESampler(n_startup_trials=5, constant_liar=True, multivariate=True),
direction="minimize",
study_name="Generate Optuna-kun (img2img)",
load_if_exists=True
)
# Optuna v3.1.1以前のバージョンでは、Optunaの _CachedStorage と呼ばれるキャッシュのためのコンポーネントの問題により、
# 本チュートリアルのコードが正しく動きません。v3.2で修正予定ですが、v3.1.1以前のバージョンをご使用中の方は次のコードを実行してください。
if isinstance(study._storage, optuna.storages._CachedStorage):
study._storage = study._storage._backend
set_objective_names(study, ["Score"])
register_objective_form_widgets(study, widgets=[
ChoiceWidget(
choices=["Good 👍", "So-so👌", "Bad 👎"],
values=[-1, 0, 1],
description="Please input your score!",
),
])
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-1",
torch_dtype=torch_dtype,
)
pipe = pipe.to(device)
pipe.enable_xformers_memory_efficient_attention()
# Start Human-in-the-loop Optimization
n_batch = 8
while True:
# 実行中のTrialが n_batch (8個) を下回ったら新しくパラメーターをサンプルして、画像を生成
running_trials = study.get_trials(deepcopy=False, states=(TrialState.RUNNING,))
if len(running_trials) >= n_batch:
time.sleep(2)
continue
suggest_and_generate_image(study, pipe)
if __name__ == '__main__':
main()
@c-bata
Copy link
Author

c-bata commented Apr 21, 2023

動作の様子

optuna-dashboard-img2img

実行方法

推奨実行環境

本Exampleを動かす際は Linux / NVIDIA GPUの利用を推奨 しています。もしGPUがなく動かすのが大変という方は、Optuna Dashboard公式ドキュメントの下記チュートリアルをお試しください。こちらはCPUのみで簡単に動かすことができ、Human-in-the-loop最適化のためのコードの理解には十分です。

https://optuna-dashboard.readthedocs.io/en/latest/tutorials/hitl.html

  • M1 Macbook Proでも動作しましたが、依存関係のインストールがやや大変な点やPyTorch MPS (Metal Performance Shaders) バックエンドではあまり速度が出なかったことから推奨しておりません。
  • Google Colaboratoryで動かしてみたいという方は、 https://optuna-dashboard.readthedocs.io/en/latest/getting-started.html#google-colaboratory を参考にOptuna Dashboardを起動してみてください。

依存関係のインストール

$ pip install optuna optuna-dashboard diffusers transformers accelerate scipy safetensors xformers

入力画像の用意

下記画像はOptunaのロゴに手足のようなものを足したものです。 input.png というファイル名で main.py 実行時のlocationと同じディレクトリ内に保存してください。 (※ 本画像の利用はデモの動作確認のみにとどめてください)

input

実行

次の2つのコマンドをそれぞれ実行してください。

  • python main.py
  • optuna-dashboard sqlite:///db.sqlite3 --artifact-dir ./artifact

macOSをお使いの方

xformersのインストール

本プログラムの実行に必要な xformers はmacOS向けのwheelバイナリを配布していないため、macOSではsdist (ソース配布パッケージ) からビルドされます。その際、Appleが提供しているClangにおいてOpenMPがデフォルトでは利用出来ないことから、次のようなエラーにぶつかることがあります。

$ pip install xformers
 ...
      clang -Wno-unused-result -Wsign-compare -Wunreachable-code -DNDEBUG -g -fwrapv -O3 -Wall -I...  -O3 -fopenmp -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE="_clang" -DPYBIND11_STDLIB="_libcpp" -DPYBIND11_BUILD_ABI="_cxxabi1002" -DTORCH_EXTENSION_NAME=_C -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++17
      clang: error: unsupported option '-fopenmp'
      error: command '/usr/bin/clang' failed with exit code 1
      [end of output]

この場合は libomp を別途インストールしたり、別のCコンパイラに切り替えることでインストールが可能です。例えば後者は次のようにできます。

$ brew install gcc
$ ls /usr/local/bin/gcc-*
/usr/local/bin/gcc-12 ...
$ CC=gcc-12 pip install xformers
...
Successfully installed xformers-0.0.16

main.py の変更

macOSで動かしたい方は、 main.py に次のpatchを適用してください。

-device = "cuda:0"
-torch_dtype = torch.float16
+device = "mps"
+torch_dtype = torch.float32
 pipe = pipe.to(device)
-pipe.enable_xformers_memory_efficient_attention()

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment