Skip to content

Instantly share code, notes, and snippets.

@milhidaka
Last active August 24, 2022 10:50
Show Gist options
  • Save milhidaka/b28d074b130c5dad5d954e2c4bd4b454 to your computer and use it in GitHub Desktop.
Save milhidaka/b28d074b130c5dad5d954e2c4bd4b454 to your computer and use it in GitHub Desktop.
Stable Diffusionによる画像生成ツールをインストールしたjupyter notebookを動かせるdockerイメージの作成 (cpu only)
FROM python:3.9-buster
ENV PYTHONUNBUFFERED=1
RUN mkdir app
WORKDIR /app
RUN pip install --upgrade diffusers transformers scipy jupyter matplotlib
RUN curl http://localhost:8000/token > /dev/shm/token && python -c 'from diffusers import StableDiffusionPipeline; StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=open("/dev/shm/token").read());'
# Stable Diffusionによる画像生成ツールをインストールしたjupyter notebookを動かせるdockerイメージの作成
# CPU only。画像生成に5分程度かかる。
# ビルド方法
# huggingfaceのアクセストークン(read権限のみで良い)を予め準備。これがホストマシンの ~/.huggingface/token に書き込まれていると仮定する。
# モデルの利用規約への同意ボタンをクリックしたうえで、トークンをファイルに書き込む(エディタで直接 or huggingface_cli login)
# トークンの入手方法参考 https://gigazine.net/news/20220824-stable-diffusion-google-colaboratory/
# HTTP経由でこのトークンを渡すサーバを立てる
# docker run --name=secrets-server --rm --volume $PWD:/files busybox httpd -f -p 8000 -h /files
# この状態で別のシェルを開きビルド
# docker build . --network=container:secrets-server -t stable_diffusion_20220823
# 実行方法
# docker run --rm -it -p 8888:8888 stable_diffusion_20220823 jupyter notebook --allow-root --ip=0.0.0.0
# 表示されたURLにホスト側のWebブラウザでアクセス
# jupyter notebookを作成して以下のスクリプトで画像生成できる
# import torch
# from torch import autocast
# from diffusers import StableDiffusionPipeline
# model_id = "CompVis/stable-diffusion-v1-4"
# pipe = StableDiffusionPipeline.from_pretrained(model_id, local_files_only=True)
# # pipe.to('cuda') # nvidia-dockerでGPU使える場合
# prompt = "a photograph of an astronaut riding a horse"
# image = pipe(prompt)["sample"][0]
# plt.imshow(image)
# image.save("foo.png")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment