Stable Diffusionの環境設定

by 水凪
2022-09-192022-09-19
ソフトウェア

タグ:Python

GPUを換装したのでStable Diffusionを実行してみました。

環境構築方法

GitHubのCompVis / stable-diffusionから自分で環境を構築する方法もありますが、動かすだけでよければHugging Faceから環境ごとダウンロードしてしまうのが楽です。

GitHub CompVis/stable-diffusion – https://github.com/CompVis/stable-diffusion
HuggingFace CompVis – https://huggingface.co/CompVis

$ conda create -n pyenv38_ldm python=3.8
$ conda activate pyenv38_ldm
$ conda install cudatoolkit==11.3
$ conda install pytorch==1.11.0
$ conda install torchvision==0.12.0
$ conda install numpy==1.19.2
$ pip install --upgrade diffusers transfomers scipi
$ huggingface-cli login

サンプルコード

import PIL.Image
import torch
import enum
from torch import autocast
from diffusers import StableDiffusionPipeline
from diffusers import StableDiffusionImg2ImgPipeline

MODEL_ID = "CompVis/stable-diffusion-v1-4"
DEVICE = "cuda"


class E_PIPE(enum.IntEnum):
    TXT2IMG = 0x01
    IMG2IMG = 0x02


def processing(e_type: int) -> any:

    if e_type == E_PIPE.TXT2IMG:
        o_pipe = StableDiffusionPipeline.from_pretrained(
            MODEL_ID, revision="fp16", torch_dtype=torch.float16, use_auth_token=True
        )
    elif e_type == E_PIPE.IMG2IMG:
        o_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
            MODEL_ID, revision="fp16", torch_dtype=torch.float16, use_auth_token=True
        )
    else:
        o_pipe = None

    assert o_pipe is not None
    return o_pipe.to(DEVICE)


def main():

    e_type = E_PIPE.IMG2IMG
    prompt = "anime"

    o_pipe = processing(e_type)

    if e_type == E_PIPE.TXT2IMG:
        with autocast("cuda"):
            im_dst = o_pipe(prompt, guidance_scale=7.5).images[0]
            im_dst.save("txt2img_export.png")

    elif e_type == E_PIPE.IMG2IMG:
        im_src = PIL.Image.open("import.png").convert("RGB")
        with autocast("cuda"):
            im_dst = o_pipe(
                prompt=prompt,
                init_image=im_src,
                strength=0.75,
                guidance_scale=7,
                num_inference_steps=50,
                generator=None,
            )["sample"][0]
            im_dst.save("img2img_export.png")


if __name__ == "__main__":
    main()

関連