1 控制随机种子
generator = torch.Generator("cuda").manual_seed(1024)
# 1 导入torch,pipline
import torch
from diffusers import StableDiffusionPipeline
from accelerate import Accelerator# 2 生成种子1024是宇航员种子
generator = torch.Generator("cuda").manual_seed(1024) # # 3 导入模型名
model_id = "runwayml/stable-diffusion-v1-5"# 4 生成器送入pipline
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, generator=generator)# 5 pipe送入cuda
pipe = pipe.to("cuda")# 6 生成提示信息
prompt = "a photo of an astronaut riding a horse on mars"
# 'An astronaut riding a horse on Mars'# 7 开始生成
image = pipe(prompt, generator=generator).images[0] # 8保存图
image.save("astronaut_rides_horse.png")
2 采用自己的unet或pipeline方法
from diffusers_inheritv2 import StableDiffusionPipeline
from diffusers_inheritv2.models.unet_2d_condition import UNet2DConditionModel as NewUNet# 生成pipline
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, generator=generator)# 生成unet
unet = NewUNet.from_pretrained(model_id, subfolder="unet").to(torch.float16) # -> yw# unet送入pipe
pipe.unet = unet.to(torch.float16)