from diffusers import DiffusionPipeline
import torch

device = "cuda"
generator = torch.Generator(device).manual_se2500)
prompt = "An astronaut riding a green horse"

# Hugging Faceから重みを引っ張ってきます。
pipe = DiffusionPipeline.from_pretrained("SimLuo/LCM_Dreamshaper_v7")

# GPUを指定し、精度はfloat32かVRAMが足りない合にはfloat16を指定してください。今回は高速性確保するため、float16を採用してみます。
pipe.to(torch_device=device, torch_dtype=torcloat16)

# 1から50の数字をステップ数に指定します。
num_inference_steps = 2 #4と8もベンチマークしみます。

# あとは画像生成を実行します。
images = pipe(prompt=prompt, num_inference_st=num_inference_steps, guidance_scale=8.0, lcmigin_steps=50, output_type="pil",generator=geator).images[0]
images
