SD3-medium.py 623 Bytes
Newer Older
lijian6's avatar
lijian6 committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
import torch
import time
from diffusers import StableDiffusion3Pipeline
#from xformers.ops import MemoryEfficientAttentionFlashAttentionOp, MemoryEfficientAttentionTritonFwdFlashBwOp

pipe = StableDiffusion3Pipeline.from_pretrained("/public/home/lijian/model/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

prompt = "A cat holding a sign that says hello world"
for i in range(3):
    start = time.time()
    image = pipe(prompt, negative_prompt="", num_inference_steps=28, guidance_scale=7.0,).images[0]
    print(f"All time: {time.time()-start}")
    image.save(f"SD3-medium.png")