Waifu Diffusionをローカル環境で動かした手順

Google Colabで実行するだけのほうが楽ではありますが
ローカル環境のほうがvram多いGPU積んだのでせっかくなので使いたい人とか向けの記事。

※Colabを使用したい方はこちらから"Open in Colab"でノートブックが開けますので
ひたすらそれを実行すると良いです。
https://huggingface.co/hakurei/waifu-diffusion

ローカル環境
-Ubuntu 20.04
-RTX 3090

1 Waifu-diffusionからリポジトリをローカル環境にクローンする。

https://github.com/harubaru/waifu-diffusion

git clone https://github.com/harubaru/waifu-diffusion
cd waifu-diffusion

2 anacondaでldmの環境を構築する。

conda env create -f environment.yaml
conda activate ldm

environment.yamlの中身は下記の通り

name: ldm
channels:
  - pytorch
  - defaults
dependencies:
  - git
  - python=3.8.5
  - pip=20.3
  - cudatoolkit=11.3
  - pytorch=1.11.0
  - torchvision=0.12.0
  - numpy=1.19.2
  - pip:
    - albumentations==0.4.3
    - opencv-python==4.1.2.30
    - pudb==2019.2
    - imageio==2.9.0
    - imageio-ffmpeg==0.4.2
    - pytorch-lightning==1.4.2
    - omegaconf==2.1.1
    - test-tube>=0.7.5
    - streamlit>=0.73.1
    - einops==0.3.0
    - torch-fidelity==0.3.0
    - transformers==4.19.2
    - torchmetrics==0.6.0
    - kornia==0.6
    - gradio==3.1.6
    - -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
    - -e git+https://github.com/openai/CLIP.git@main#egg=clip
    - -e git+https://github.com/hlky/k-diffusion-sd#egg=k_diffusion
    - -e .

3 diffusersが入ってないとエラーで起動しない & pytorchがcudaのバージョンと整合性がないと

実行出来ない様なので最新の環境に更新しておく

conda install pytorch torchvision -c pytorch
pip install transformers==4.19.2 diffusers invisible-watermark
pip install -e .

4 下記のコードをwaifu-diffusionのディレクトリに保存します。

sample.py

import gradio as gr
import torch
from torch import autocast
from diffusers import StableDiffusionPipeline

model_id = "hakurei/waifu-diffusion"
device = "cuda"


pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision='fp16')
pipe = pipe.to(device)


block = gr.Blocks(css=".container { max-width: 800px; margin: auto; }")

num_samples = 2

def infer(prompt):
    with autocast("cuda"):
        images = pipe([prompt] * num_samples, guidance_scale=7.5)["sample"]

    return images


with block as demo:
    gr.Markdown("<h1><center>Waifu Diffusion</center></h1>")
    gr.Markdown(
        "waifu-diffusion is a latent text-to-image diffusion model that has been conditioned on high-quality anime images through fine-tuning."
    )
    with gr.Group():
        with gr.Box():
            with gr.Row().style(mobile_collapse=False, equal_height=True):

                text = gr.Textbox(
                    label="Enter your prompt", show_label=False, max_lines=1
                ).style(
                    border=(True, False, True, True),
                    rounded=(True, False, False, True),
                    container=False,
                )
                btn = gr.Button("Run").style(
                    margin=False,
                    rounded=(False, True, True, False),
                )
               
        gallery = gr.Gallery(label="Generated images", show_label=False).style(
            grid=[2], height="auto"
        )
        text.submit(infer, inputs=[text], outputs=gallery)
        btn.click(infer, inputs=[text], outputs=gallery)

    gr.Markdown(
        """___
   <p style='text-align: center'>
   Created by https://huggingface.co/hakurei
   <br/>
   </p>"""
    )


demo.launch(debug=True)

5 sample.pyを実行するとこんな感じでローカルで提供されますのでブラウザでアクセス

~/git/waifu-diffusion$ python sample.py
Running on local URL:  http://127.0.0.1:7860/

6 プロンプトを入力すると10-20秒位で画像が生成されます。

あとは適当に画像を保存してなにか目的に合うものに使用します。

補足

vramの使用量は最大時11GB程でした。
optimized版でなければGPUはvramが少なくとも12GB以上ないと厳しそうです。