StreamDiffusionとか

StableDiffusionが出たときぶりに画像生成AIをローカルに構築してみる

まずはcolabでお試し

from google.colab import drive
drive.mount('./gdrive')
!mkdir -p ./gdrive/MyDrive/pip/
import sys
sys.path.append('/content/gdrive/MyDrive/pip')
#!mkdir -p ./diffusers/

# PyTorchとxformersのインストール
import sys
sys.path.append('/content/gdrive/MyDrive/pip')
!pip install --target /content/gdrive/MyDrive/pip  torch==2.1.0 torchvision==0.16.0 xformers --index-url https://download.pytorch.org/whl/cu121
!pip install --target /content/gdrive/MyDrive/pip diffusers
!pip install --target /content/gdrive/MyDrive/pip accelerate
!pip install --target /content/gdrive/MyDrive/pip git+https://github.com/cumulo-autumn/StreamDiffusion.git@main#egg=streamdiffusion

import torch
from diffusers import AutoencoderTiny, StableDiffusionPipeline

from streamdiffusion import StreamDiffusion
from streamdiffusion.image_utils import postprocess_image
import time
import numpy as np
import matplotlib.pyplot as plt
from IPython.display import display  # 試した環境では省略しても動作しました
pipe = StableDiffusionPipeline.from_pretrained("KBlueLeaf/kohaku-v2.1",cache_dir = "/content/gdrive/MyDrive/diffusers/kohaku/").to(
    device=torch.device("cuda"),
    dtype=torch.float16,
)

# Wrap the pipeline in StreamDiffusion
# Requires more long steps (len(t_index_list)) in text2image
# You should use cfg_type="none" when text2image
stream = StreamDiffusion(
    pipe,
    t_index_list=[0, 16, 32, 45],
    torch_dtype=torch.float16,
    cfg_type="none",
)

# If the loaded model is not LCM, merge LCM
stream.load_lcm_lora()
stream.fuse_lora()
# Use Tiny VAE for further acceleration
stream.vae = AutoencoderTiny.from_pretrained("madebyollin/taesd").to(device=pipe.device, dtype=pipe.dtype)

# Enable acceleration
pipe.enable_xformers_memory_efficient_attention()
#prompt = "1girl with dog hair, thick frame glasses" # @param {type:"string"}
prompt = "boy" # @param {type:"string"}

# ファイルが存在しない場合は新しく作成されます
with open("prompt.txt", "w") as file:
    # 文字列をファイルに書き込む
    file.write(prompt)
    file.close()
# Prepare the stream
stream.prepare(prompt)

# Warmup >= len(t_index_list) x frame_buffer_size
for i in range(4):
    stream()

# Run the stream infinitely
fig = plt.figure(figsize=(5, 5))
display_handle = display(fig, display_id=True)
shinprompt=prompt
for i in range(30):
    with open('prompt.txt', 'r') as file:
      prompt2 = file.read()
    # ストリームの実行
    if shinprompt != prompt2:
      stream.update_prompt(prompt2)
    x_outputlist = stream.txt2img()
    pics = postprocess_image(x_outputlist, output_type="np")
    for count, pic in enumerate(pics):
      #arrPIL = np.asarray(pic)
      #output_image.save(f"/content/sdoutput/{i:02}.png")
      plt.imshow(pic)
      display_handle.update(fig)
      plt.clf()
    shinprompt=prompt2
    time.sleep(2.0)

サンプルを適当にいじって実行
画像が保存されるのがなんとなくいやだからmatlab使って表示するだけに

で、実行してみた結果だけど流石にはやい。どれが効いているのかは浦島すぎるけど、確か1つの画像に4秒から10秒くらいかかってたのが、一瞬で出る。すごい。

で、colab使ってると制限かかったからローカル環境でも試してみる。
ReedMeに従ってanacondaとvenv使った上で環境作成して後は同じ。VS Codeにjupytor Notebook入れてるからコードもそのまま
30秒で1枚だったのがsleep入れないと流れるのが速すぎるくらいだから速い

ただ、自分の弱点としてこれといって表現したいものは無いから、適当に遊ぶ以上のことが特にできない。
なんかないものか

そういえばローカルの環境書いてないな
メモリ16GBでグラボはGTX1080
VRやるために強いグラボ買ったおかげでまあ楽しめてる
次の買い替えがなあ