stable-diffusion AI 绘画

kyaa111 8月前 ⋅ 233 阅读
  1. git clone https://github.com/CompVis/stable-diffusion.git
  2. 进入stable-diffusion目录
  3. 在这里注册一个账号: https://huggingface.co/ 并生成个token
  4. 安装CUDA https://blog.csdn.net/qq_35930739/article/details/128167167
  5. pip install torch -f https://download.pytorch.org/whl/torch_stable.html
  6. pip install -e git+ https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
  7. conda install torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia 选择合适的pytorch-cuda版本
  8. pip install transformers==4.25.1 diffusers invisible-watermark
  9. pip install -e .
  10. huggingface-cli login 输入生成的token
from torch import autocast
from diffusers import StableDiffusionPipeline
import torch

if __name__ == '__main__':
    # https://github.com/CompVis/stable-diffusion/issues/69#issuecomment-1260722801
    torch.backends.cudnn.benchmark = True
    torch.backends.cudnn.enabled = True

    pipe = StableDiffusionPipeline.from_pretrained(
        "CompVis/stable-diffusion-v1-4",
        torch_dtype=torch.float16,
        use_auth_token=True,
        safety_checker=None,
        requires_safety_checker=False
    ).to("cuda")


    prompt = "a photo of an astronaut riding a horse on mars"
    with autocast("cuda"):
        result = pipe(prompt, height=256, width=256)

    image = result[0][0]

    image.save("test.png")