StreamDiffusionで高速にAI画像を生成することができます。
JupyterLab上でStreamDiffusionをインストールするには以下の手順を実行します。
1.pipenvのインストール
sudo apt-get -y install python3-pip python3-distutils python3-dev python3-testresources
python3 -m pip install --user pipenv
echo "export PIPENV_VENV_IN_PROJECT=true" >> ~/.profile
echo 'export PATH=$PATH:$HOME/.local/bin' >> ~/.profile
source ~/.profile
2. Jupyterのインストール
sudo mkdir -p /mnt/data/jupyter
※パスは適宜変更してください
sudo chown ubuntu:ubuntu /mnt/data/jupyter
※ユーザ名・グループ名は適宜変更してください
cd /mnt/data/jupyter
pipenv --python 3
pipenv install jupyterlab
pipenv install matplotlib
pipenv install pandas
3. Dream Diffusionのインストール
sudo apt-get -y install git
git clone https://github.com/cumulo-autumn/StreamDiffusion.git
cat << EOF >> Pipfile
[[source]]
url = "https://download.pytorch.org/whl/cu121"
verify_ssl = false
name = "pytorch"
EOF
pipenv install torch==2.1.0 torchvision==0.16.0 xformers --index=pytorch
pipenv install git+https://github.com/cumulo-autumn/StreamDiffusion.git@main#egg=streamdiffusion[tensorrt]
pipenv install accelerate
4. モデルのダウンロード
cd /mnt/data/jupyter/StreamDiffusion/models/Model
wget -4 https://huggingface.co/KBlueLeaf/kohaku-v2.1/resolve/main/kohaku-v2.1.safetensors
※必要に応じて他のモデルもダウンロードしてください。
5. Jupyter Labをサービスとして登録
cat << EOF | sudo tee /etc/systemd/system/jupyter.service
[Unit]
Description=Jupyter
[Service]
Type=simple
Environment=HF_HOME=/mnt/data/jupyter/StreamDiffusion/models/Model
ExecStart=/mnt/data/jupyter/.venv/bin/jupyter lab --port 8080 --ip=0.0.0.0 --allow-root --NotebookApp.token='jupyter'
User=ubuntu
Group=ubuntu
WorkingDirectory=/mnt/data/jupyter
Restart=always
RestartSec=10
[Install]
WantedBy=multi-user.target
EOF
※ディレクトリ、ユーザ、トークンは適宜変更してください
sudo systemctl enable jupyter
sudo systemctl start jupyter
6. ブラウザでJupyterにアクセスしてNotebookに以下のコードを貼り付けます。
http://{Jupyterをインストールしたホスト}:8080/lab?token=jupyter
Notebookに貼り付けるコードは以下です。デバイスにCPUを指定して、GPUなしのマシンでも画像を生成することができます。
import sys
sys.path.append("/mnt/data/jupyter/StreamDiffusion")
import torch
from utils.wrapper import StreamDiffusionWrapper
# モデル
model_id_or_path = "/mnt/data/jupyter/StreamDiffusion/models/Model/kohaku-v2.1.safetensors"
# 画像サイズ
width = 512
height = 512
# seed
seed = 2
# プロンプト
prompt="1girl and 1dog walking riverside under blue sky"
negative_prompt=""
stream = StreamDiffusionWrapper(
model_id_or_path=model_id_or_path,
t_index_list=[0, 16, 32, 45],
frame_buffer_size=1,
width=width,
height=height,
warmup=10,
acceleration="none",
mode="txt2img",
use_denoising_batch=False,
cfg_type="none",
seed=seed,
device="cpu",
dtype=torch.float32,
)
stream.prepare(
prompt=prompt,
negative_prompt=negative_prompt,
num_inference_steps=50,
)
for _ in range(stream.batch_size - 1):
stream()
output_image = stream()
output_image
実行結果