project / model /sd_turbo.py
seokhyun119's picture
๋ชจ๋ธ ์—…๋กœ๋“œ
4944c37
raw
history blame contribute delete
594 Bytes
from diffusers import StableDiffusionPipeline
import torch
def generate_sd_turbo(prompt: str):
model_id = "stabilityai/sd-turbo"
pipe = StableDiffusionPipeline.from_pretrained(
model_id,
torch_dtype=torch.float16,
use_safetensors=True
).to("cuda")
image = pipe(prompt=prompt, guidance_scale=0.0).images[0]
image.save("output_sd_turbo.png")
return image
if __name__ == "__main__":
prompt = "๊ทธ๋…€๋ฅผ ๋ฐ”๋ผ๋ณด๋Š” ํ•œ ๋‚จ์ž์˜ ์•ผ๋ง"
img = generate_sd_turbo(prompt)
print("โœ… ์ €์žฅ ์™„๋ฃŒ: output_sd_turbo.png")
img.show()