Compute EngineにHuggingFaceのDiffusers 載せて StableDiffusionの画像生成API作ってみる

# 背景
- 大量に画像生成しまくれる環境が欲しい
- T4が乗ってるCompute EngineにDiffusers載せて画像を生成するAPIを作ってみます
- 今回は `stable diffusino v1.5` のpre-trainモデルと、 `epiCRealism` のcheckpointのAPIを用意します
- 77token以上のpromptを扱うためにCompelを利用します


https://huggingface.co/docs/diffusers/v0.27.2/en/index

https://github.com/damian0815/compel


# CompVis/stable-diffusion? Diffusers?

StableDiffusionのAPI作成方法を調べると `CompVis/stable-diffusion` を使う方法と, `Diffusers` の方法がよく出てきます.

https://github.com/CompVis/stable-diffusion

どうやら、ComVis/stable-diffusionの方が軽量で速いようです。内部的にはhaggingFace Diffusersが使われています。ただ、CompVisはドキュメントがあまりなかったので、Diffusersを使うことにしました。


# 環境
- Compute Engine
- machine type: n1-highmem-8
- 1 x NVIDIA T4
- 起動イメージ: Google, Deep Learning VM with CUDA 11.8, M116, Debian 11, Python 3.10. With CUDA 11.8 preinstalled

# 方法

### インストール

```py
# 80ポートで動かしたいのでsudoでやっていきます
sudo pip install --upgrade diffusers accelerate transformers flask gunicorn
```

### 実装

```py
from flask import Flask, request
from diffusers import DiffusionPipeline, StableDiffusionPipeline, DPMSolverMultistepScheduler
import base64
import torch
from compel import Compel


# epiCRealismの場合
# model_id = "emilianJR/epiCRealism"
# pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", use_safetensors=True)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
compel = Compel(tokenizer=pipe.tokenizer, text_encoder=pipe.text_encoder)
if torch.cuda.is_available():
    pipe = pipe.to("cuda")

app = Flask(__name__)

@app.route("/create", methods=['POST'])
def create():
    data_dict = request.get_json()
    prompt = data_dict['prompt']

    # ネガティブプロンプトの追加
    negative_prompt = """
    verybadimagenegative_v1.3, ng_deepnegative_v1_75t, (ugly face:0.8),cross-eyed,sketches, (worst quality:2), (low quality:2), (normal quality:2), lowres, normal quality, ((monochrome)), ((grayscale)), skin spots, acnes, skin blemishes, bad anatomy, DeepNegative, facing away, tilted head, {Multiple people}, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worstquality, low quality, normal quality, jpegartifacts, signature, watermark, username, blurry, bad feet, cropped, poorly drawn hands, poorly drawn face, mutation, deformed, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, extra fingers, fewer digits, extra limbs, extra arms,extra legs, malformed limbs, fused fingers, too many fingers, long neck, cross-eyed,mutated hands, polar lowres, bad body, bad proportions, gross proportions, text, error, missing fingers, missing arms, missing legs, extra digit, extra arms, extra leg, extra foot, ((repeating hair)), facial distortion, facial exaggerated features, unnatural facial proportions, unnatural human facial symmetry, any facial deformities
    """
    positive_scale = 7.5  # ポジティブなプロンプトの制御強度
    negative_scale = 5.0  # ネガティブプロンプトの制御強度

    # compel使わないとinputトークンは77tokenが上限になってしまうので、compelを使います
    # コンディショニングテンソルの作成
    conditioning = compel.build_conditioning_tensor(prompt)
    negative_conditioning = compel.build_conditioning_tensor(negative_prompt)

    pipe.to("cuda")
    # ガイダンススケールを使用して、プロンプトのポジティブおよびネガティブ効果を調整
    image = pipe(prompt_embeds=conditioning, negative_prompt_embeds=negative_conditioning, width=512, height=512,
                 num_inference_steps=30, guidance_scale=positive_scale, negative_guidance_scale=negative_scale).images[0]
    image.save('tmp.png')
    with open("tmp.png", "rb") as image_file:
        encoded_string = base64.b64encode(image_file.read()).decode('utf-8')

    # b64エンコードの文字列を返します
    return {"b64_json": encoded_string}

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=80)
```

### 起動

```sh
sudo gunicorn --bind :80 --workers 8 --timeout 0 main:app
```

### 実行

promptの画像が `hoge.png` に保存されます.

```py
import requests
import base64
import json

prompt = """"
Image of a cat in the park.
"""

r = requests.post(
    'http://IP_ADDRESS/create',
    headers = {'Content-Type': 'application/json'},
    data=json.dumps({'prompt': prompt})
)

image_bytes = base64.b64decode(r.json()['b64_json'])
with open('hoge.png', 'wb') as file:
    file.write(image_bytes)
```

# まとめ
HuggingFaceを使うと結構簡単にAPIが立ち上がりました。品質がイマイチなので今度はStableDiffusion XLを立ち上げてみようと思いまs。

# 参考

https://huggingface.co/emilianJR/epiCRealism

https://github.com/huggingface/diffusers/issues/2136