stable-diffusion AI 绘画

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")