😸

LCM-LoRA×DiffusersでリアルタイムAIお絵描きを試してみる

抹茶もなか

2023/11/17に公開

最近話題のリアルタイムAIお絵描きを実際にやってみようという回です。

使用技術

今回はタイトル通りDiffusersとLCM-Loraを使います。
AIによって画像生成をリアルタイムで行いながら、お絵描き補助をするというツールは実は画像生成AIが出てきた初期の方にあったりするのですが、今回はLCMという技術を使うことで画像生成速度が大幅に向上したことで、ほぼリアルタイムに書いたものがAIによって高品質化されるという点が相違点になります。

LCM及びLCM-LoRAとは

今回肝となるLCMですが、正式名称はlatent-consistency-modelといいます。
LCMが出てきたのは2023/10/6に公表された画像生成の高速化技術で、ものすごく簡潔に説明すると、従来の画像生成AIが複数のステップを踏みながらノイズを徐々に除去して画像生成するところを、一発でノイズ画像から生成結果を出す技術です（実際には、1ステップで生成すると品質の低下が著しいので、少ないステップ数で生成可能、くらいの効果）

ただし、このLCMは専用の学習が必要で、既存のStableDiffusionモデルと組み合わせて使うことができませんでした。
この問題を解決したのが、LCM-LoRAとなります。
生成AI界隈ではおなじみとなったLora（low rank adaptation）ですが、LCMがつい先日11月9日にLoRA形式で公開されました。

これにより、既存StableDiffusionモデルでも、他のLoRAを試すのと同じ感覚で画像生成の超高速化が試せるようになりました。

リアルタイムAIお絵描きの仕組み

本題に入ります。
上記で説明したLCM-LoRAを使いリアルタイムお絵描きを試します。
といっても、仕組みは非常に簡単で、PC画面に表示されたウィンドウをリアルタイムで取り込み、それをDiffusers経由でLCM-LoRAに渡すだけで基本的な実装が可能です。
（サービスとして公開されているリアルタイムお絵描きAIも、キャンバス側のメニューバーとか拾って生成結果にのってるので多分同じ仕組み）

結果だけ先に知りたい人用に動くコードはこちら

import cv2
import numpy as np
import pygetwindow as gw
import pyautogui

from diffusers import AutoPipelineForImage2Image, UNet2DConditionModel, LCMScheduler
from diffusers.utils import load_image
import torch

lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5"

torch_dtype=torch.float16).to("cuda")
pipe = AutoPipelineForImage2Image.from_pretrained(
    "852wa/SDHK", torch_dtype=torch.float16, use_safetensors=True
).to("cuda")


pipe.load_lora_weights(lcm_lora_id)
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)

if pipe.safety_checker is not None:
    pipe.safety_checker = lambda images, **kwargs: (images, [False])


prompt = "1girl, black hair, school uniform, 8k"


# 画像を表示するウィンドウの作成
cv2.namedWindow("Active Window Capture", cv2.WINDOW_NORMAL)

while True:
    # アクティブなウィンドウを取得
    active_window = gw.getActiveWindow()
    if active_window:
        # アクティブなウィンドウの位置とサイズを取得
        x, y, width, height = active_window.left, active_window.top, active_window.width, active_window.height

        # 指定された領域のスクリーンショットを取得
        screenshot = pyautogui.screenshot(region=(x, y, width, height))
        img_np = np.array(screenshot)

        

        generator = torch.Generator("cuda").manual_seed(2500)
        img = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)
        

        img = pipe(
            prompt=prompt,
            image=img,
            num_inference_steps=8,
            guidance_scale=1,
            generator=generator
        ).images[0]

        # PILイメージをnumpy配列に変換
        img = np.array(img)

        # OpenCVでは色の順番がBGRなので、RGBからBGRに変換
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

        # キャプチャした画像を同じウィンドウで更新して表示
        cv2.imshow("Active Window Capture", img)

        # 'q'を押したらループを抜ける
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

cv2.destroyAllWindows()

コード説明

まずはDiffusersの設定から

lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5"

torch_dtype=torch.float16).to("cuda")
pipe = AutoPipelineForImage2Image.from_pretrained(
    "852wa/SDHK", torch_dtype=torch.float16, use_safetensors=True
).to("cuda")


pipe.load_lora_weights(lcm_lora_id)
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)

if pipe.safety_checker is not None:
    pipe.safety_checker = lambda images, **kwargs: (images, [False])
prompt = "1girl, black hair, school uniform, 8k"

今回はSD1.5ベースのモデルである箱庭さんの852wa/SDHKモデルを使わせていただいています。
ベースが1.5なのでLCM-LoRAもSD1.5対応のものを使用（SDXL用のLoRAも公開されています）

ちなみにDiffusers公式にも、LCMの使い方サンプルがのっていますが、あちらはtxt2imageのサンプルのため、今回のリアルタイムAIお絵描きには使えないので注意。

i2iで利用する必要があるため、AutoPipelineForImage2Imageを利用します。

今回はお試しコードのため、テキストプロンプトもコード内でべた書きで書いてしまっています（普段使いしたい人はここは別ウィンドウで受け取れるようにした方が良いかも）
プロンプト部分は描きたい絵に沿って適宜変えてください。

生成モデルの設定は以上です。

# 画像を表示するウィンドウの作成
cv2.namedWindow("Active Window Capture", cv2.WINDOW_NORMAL)

while True:
    # アクティブなウィンドウを取得
    active_window = gw.getActiveWindow()
    if active_window:
        # アクティブなウィンドウの位置とサイズを取得
        x, y, width, height = active_window.left, active_window.top, active_window.width, active_window.height

        # 指定された領域のスクリーンショットを取得
        screenshot = pyautogui.screenshot(region=(x, y, width, height))
        img_np = np.array(screenshot)

        

        generator = torch.Generator("cuda").manual_seed(2500)
        img = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)
        

        img = pipe(
            prompt=prompt,
            image=img,
            num_inference_steps=8,
            guidance_scale=1,
            generator=generator
        ).images[0]

        # PILイメージをnumpy配列に変換
        img = np.array(img)

        # OpenCVでは色の順番がBGRなので、RGBからBGRに変換
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

        # PILの画像をOpenCVの形式に変換
        #img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)


        # キャプチャした画像を同じウィンドウで更新して表示
        cv2.imshow("Active Window Capture", img)

        # 'q'を押したらループを抜ける
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

cv2.destroyAllWindows()

アクティブな画面をリアルタイムで拾って、画像生成AIに流し別ウィンドウで再表示するスクリプトです。

処理手順は以下

cv2.namedWindowで生成結果を表示する画面を設定
gw.getActiveWindow()で現在アクティブ（直近クリックした画面）な画面を取得
pyautogui.screenshotでアクティブな画面の範囲のみスクリーンショットで撮影
pipe()でとったスクリーンショットを先ほど設定した画像生成モデルに渡しi2i結果を取得
「1」で設定した画面に生成結果を渡して完了

非常に簡単ですね。
生成モデルの設定さえうまくいってしまえば、特に問題なく動くかと思います。
注意点としては、WSLで実行した場合、アクティブ画面が出力できないので、WSL環境で普段スクリプトを動かしている人は、PowerShell等Windows側の環境から実行する必要があります。

使い方

動作確認をしたのはWindowsのみなのでご了承いただければと思います。

PowerShellで作成したプログラムを実行すると、生成結果表示画面が出てくるので、リアルタイムi2iをしたい画面をクリック。
リアルタイムAIお絵描きを簡単に試す場合は、Windows備え付けのペイントソフトとかを開いてクリックしておけば試すことができます。
無事アクティブウィンドウを読み込めるとひたすら高速i2i変換をし始めるので、ペイント側でお絵描きをする事で（ほぼ）リアルタイムに画像生成AIによる生成結果が表示されるようになります。

実際に使ってみた結果がこちら

最後に

という訳で話題のリアルタイムAIお絵描きの簡易版を作ってみました。
基本的な構造はこの記事に書いた通りだと思うので、あとはリッチなUIとか用意してあげるともう少し使いやすくなるかなと思います。

コードはgithubにも上げているので、そちらも参照していただければと思います。

AIものづくり研究会Publication

日本最大級のGenerative AI関連のDiscordサーバーコミュニティ

使用技術

LCM及びLCM-LoRAとは

リアルタイムAIお絵描きの仕組み

コード説明

使い方

最後に

Discussion