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()