import torch
from diffusers import HunyuanDiTPipeline
import time

pipe = HunyuanDiTPipeline.from_pretrained("HunyuanDiT-Diffusers/", torch_dtype=torch.float16)
pipe.to("cuda")

base_count = 0
prompt = "一个宇航员在骑马"
for i in range(2):
    start = time.time()
    image = pipe(prompt).images[0]
    print(f"the {i} text-to-image use time {time.time()-start}")
    base_count += 1
    image.save(f"Hunyuan-{i:05}.png")