土鍋で雑多煮

UnityでXR・ゲーム開発をしています。学んだことや備忘録、趣味の記録などを書いていきます。

UnityからStable Diffusion web UI APIを叩いてテクスチャを生成する

Unity AI

どうも、土鍋です。
UnityからプロンプトをStableDiffusion web UI APIに送信して生成した画像をUnityで表示するということを一連でやっている記事を見かけなかったので書きました。

実装したもの

プロンプトを打ってボタンを押すと、画像を生成し、テクスチャに反映するものを実装しました。

Stable Diffusion web UIの起動

github.com

stable-diffusion-webuiディレクトリで.\webui-user.batを実行することで起動できます。
初回起動時はそれなりに時間がかかります。

※ここではStable Diffusion web UIの詳しい解説は書きません。

APIの起動

webui-user.batのset COMMANDLINE_ARGS= に --apiを追加することでweb UIとAPIを起動できます。

Text To Image

非同期処理部分はUniTaskを使用して実装してあります。

書いたコード

とりあえず、今回のコード全体を見たい方はこちら

今回のコード

Jsonを受け取ってAPI叩いて生成されたテクスチャを返すコード

using System;
using System.Text.RegularExpressions;
using System.Threading;
using Cysharp.Threading.Tasks;
using UnityEngine;
using UnityEngine.Networking;

public class TextToImage
{
    private string url = "http://127.0.0.1:7860";

    /// <summary>
    /// Text to ImageのAPIリクエスト
    /// </summary>
    /// <param name="json"></param>
    /// <param name="cancellationToken"></param>
    /// <returns></returns>
    public async UniTask<Texture2D> PostT2I(string json, CancellationToken cancellationToken = default)
    {
        cancellationToken.ThrowIfCancellationRequested();

        byte[] postData = System.Text.Encoding.UTF8.GetBytes(json);
        var request = new UnityWebRequest(url + "/sdapi/v1/txt2img", "POST");
        request.uploadHandler = (UploadHandler)new UploadHandlerRaw(postData);
        request.downloadHandler = (DownloadHandler)new DownloadHandlerBuffer();
        request.SetRequestHeader("Content-Type", "application/json");

        // APIリクエストを送信
        Debug.Log("Send Prompt");
        await request.SendWebRequest().WithCancellation(cancellationToken);

        if (request.result == UnityWebRequest.Result.Success)
        {
            Debug.Log("Request Success");
                
            // レスポンスのJsonを取得
            var response = request.downloadHandler.text;

            // 「"」で囲まれた文字列を抽出
            var matches = new Regex("\"(.+?)\"").Matches(response);

            // 画像データを取得
            var imageData = matches[1].ToString().Trim('"');

            // Base64をbyte型配列に変換
            byte[] data = Convert.FromBase64String(imageData);

            // byte型配列をテクスチャに変換
            Texture2D texture = new Texture2D(1, 1);
            texture.LoadImage(data);

            return texture;
        }
        else
        {
            Debug.Log("Error:" + request.result);
            
            Texture2D texture = new Texture2D(1, 1);
            return texture;
        }
    }
}

InputFieldにプロンプトを入力され、ボタンが押されると、画像を生成するコード

using UnityEngine;
using UnityEngine.UI;

public class Test2ImageTest : MonoBehaviour
{
    [SerializeField] private Image image;
    [SerializeField] private InputField inputField;
    [SerializeField] private Button sendButton;
    
    private TextToImage _t2I = new TextToImage();

    [System.Serializable]
    public class RequestData{ 
        public string prompt;
    }

    void Start()
    {
        sendButton.onClick.AddListener(()=> 
        {
            SendPrompt(inputField.text);
        });
    }

    async void SendPrompt(string prompt)
    {
        // Jsonに変換
        RequestData requestData = new RequestData();
        requestData.prompt = prompt;
        var json = JsonUtility.ToJson(requestData);
        
        // リクエスト
        var result = await _t2I.PostT2I(json);
        
        // Texture2DからSpriteに変換
        image.sprite = Sprite.Create(result, new Rect(0, 0, result.width, result.height), Vector2.zero);
    }
}

解説

http://127.0.0.1:7860/docs にアクセスすることでFastAPIのドキュメントページを確認できます。

今回はTextToImageなので/sdapi/v1/txt2imgを確認します。

リクエストJSON

何も記述しなければデフォルトのものが使われるので、とりあえずPromptだけのJsonにします。 UnityにはオブジェクトをJsonに変換するJsonUtilityがあるので、まずはJsonに対応したクラスを書きます。

[System.Serializable]
public class RequestData{ 
    public string prompt;
}

JsonUtility.ToJson()でオブジェクトをJsonに変換します。

RequestData requestData = new RequestData();
requestData.prompt = prompt;
var json = JsonUtility.ToJson(requestData);

POSTリクエスト

Jsonをbyte配列に変換し、UnityWebRequestをPOST形式で生成し、もろもろ設定。

byte[] postData = System.Text.Encoding.UTF8.GetBytes(json);
var request = new UnityWebRequest(url + "/sdapi/v1/txt2img", "POST");
request.uploadHandler = (UploadHandler)new UploadHandlerRaw(postData);
request.downloadHandler = (DownloadHandler)new DownloadHandlerBuffer();
request.SetRequestHeader("Content-Type", "application/json");

リクエストを投げる。

await request.SendWebRequest().WithCancellation(cancellationToken);

レスポンスJSON

上記のようなJsonレスポンスが返ってくるので、imageデータを抽出。

// レスポンスのJsonを取得
var response = request.downloadHandler.text;

// 「"」で囲まれた文字列を抽出
 var matches = new Regex("\"(.+?)\"").Matches(response);

 // 画像データを取得
var imageData = matches[1].ToString().Trim('"');

imageはBase64で返ってくるのでbyte配列に変換してあげる。

byte[] data = Convert.FromBase64String(imageData);

byte配列をテクスチャに変換

Texture2D texture = new Texture2D(1, 1);
texture.LoadImage(data);

最後にテクスチャからSpriteを生成して、画面に表示する。

image.sprite = Sprite.Create(result, new Rect(0, 0, result.width, result.height), Vector2.zero);

まとめ

今回はプロンプトのみのリクエストでしたが、テクスチャサイズによって生成サイズを変わるようにしたり、ChatGPTと組み合わせれば、面白いことができそうです。

参考