9fo912 / model_test.py
pengdaqian
fix more
91920f4
raw
history blame
509 Bytes
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
import torch
model_id = "stabilityai/stable-diffusion-2-1-base"
scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(
model_id,
scheduler=scheduler,
# safety_checker=None,
revision="fp16")
pipe = pipe.to("cpu")
prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]
image.save("astronaut_rides_horse.png")