Stable Diffusion（Diffusersライブラリ）をGoogle Colab上のflaskで動かす #JavaScript

やったこと

Google Colab上にたてたflaskサーバ上でStable diffusionを実行し、
外部のHTMLからJavascriptで画像を取得した。

参考
https://gammasoft.jp/blog/stable-diffusion-with-diffusers-library/

前提①

外部公開に、ngrokを使用するので、ngrokの認証用のTokenを取得する。
ngrok
https://ngrok.com/

前提②

Google ColabのRuntime はGPUに設定していること。

!nvidia-smi

Diffusersライブラリインストール

!pip install diffusers transformers ftfy accelerate

flaskインストール

!pip install flask-ngrok
!pip install flask_cors
!pip install pyngrok

ngrok認証設定

ngrokで取得した認証トークンを設定

!ngrok authtoken ＜auth Token＞

flask実行

promptを受け取り、base64でエンコードした文字列を返す。
v2にあげると、OutOfMemoryになったので、v1-4を指定

from flask import Flask,request
from flask_ngrok import run_with_ngrok
import torch
from diffusers import StableDiffusionPipeline
import base64
from io import BytesIO
from flask_cors import CORS

model_id = "CompVis/stable-diffusion-v1-4"
device = "cuda"

app = Flask(__name__)
CORS(app)
run_with_ngrok(app)

#prompt = "Mt. Fuji in the style of Gauguin"
pipe = StableDiffusionPipeline.from_pretrained(model_id, revision="fp16", torch_dtype=torch.float16)
pipe = pipe.to(device)

@app.route("/" ,methods=['POST'])
def home():
	prompt = request.form['prompt']
	print("prompt:" + prompt )
	generator = torch.Generator(device).manual_seed(0) # 
	image = pipe(prompt, guidance_scale=7.5, generator=generator).images[0]

	# 生成した画像の保存
	jpeg_data = BytesIO()
	image.save(jpeg_data, format='JPEG')
	image_bytes = jpeg_data.getvalue()
	base64_string = base64.b64encode(image_bytes).decode('utf-8')

	return "data:image/jpeg;base64," + base64_string
	
app.run()

公開用URL取得

実行後、公開用URLが出力されるので、メモする。

 * Serving Flask app '__main__'
 * Debug mode: off
INFO:werkzeug:WARNING: This is a development server. Do not use it in a production deployment. Use a production WSGI server instead.
 * Running on http://127.0.0.1:5000
INFO:werkzeug:Press CTRL+C to quit
 * Running on http://154d-34-134-107-25.ngrok-free.app ←　外部アクセス用URL
 * Traffic stats available on http://127.0.0.1:4040
prompt:town of Europe

確認用HTML

入力されたpromptを投げて、取得した画像をimgタグに設定するHTMLで確認。

<html>
<head>
<script src="https://ajax.googleapis.com/ajax/libs/jquery/3.6.0/jquery.min.js"></script>
</head>
<body>
	<input id="prompt" type="text" style="width:500px" value="" /><BR>
	<input type="hidden" id="base64Image" value="" />
	<button type="button" id="createImage"  onclick="sendPrompt()" />create image</button>
	<BR>
	<img id="createImg" src="" width="256px" height="256px">
	
<script>
function sendPrompt(){

	let prompt = document.getElementById('prompt').value ;
		
	const obj = {"prompt" : prompt};
	const method = "POST";
	const body = Object.keys(obj).map((key)=>key+"="+encodeURIComponent(obj[key])).join("&");
	const headers = {
		'Content-Type': 'application/x-www-form-urlencoded'
	};
	const url = "http://154d-34-134-107-25.ngrok-free.app" // URL指定
	
	fetch(url, 
		{method, headers, body})
		.then( (res) => res.text() )
		.then((text) => {
			document.getElementById('createImg').src = text;
			document.getElementById('base64Image').value = text;
		})
		.catch(error=>{
			console.log(error);
		});
}

</script>
</body>
</html>

つくったものはこちらでも利用