# from diffusers import StableDiffusionPipeline
# import torch

# def generate_waifu(prompt: str):
#     model_id = "hakurei/waifu-diffusion"
#     pipe = StableDiffusionPipeline.from_pretrained(
#         model_id,
#         torch_dtype=torch.float16,
#         use_safetensors=False,
#         revision="fp16"
#     ).to("cuda")

#     image = pipe(prompt=prompt).images[0]
#     return image

# if __name__ == "__main__":
#     prompt = '그녀를 바라보는 한 남자의 야망'
#     image = generate_waifu(prompt)
#     image.save("output_waifu.png")
#     print("✅ 저장 완료: output_waifu.png")
#     image.show()

from diffusers import StableDiffusionPipeline
import torch

model_id = "hakurei/waifu-diffusion"
pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float32,
    use_safetensors=False
).to("cpu")  # CPU로 명시

def generate_waifu(prompt: str):
    image = pipe(prompt=prompt).images[0]
    return image


if __name__ == "__main__":
    prompt = '그녀를 바라보는 한 남자의 야망'
    image = generate_waifu(prompt)
    image.save("output_waifu_cpu.png")
    print("✅ 저장 완료: output_waifu_cpu.png")
    image.show()