from diffusers import StableDiffusionXLPipeline, AutoencoderKL
import torch

def generate_animesai(prompt: str):
    model_id = "enhanceaiteam/AnimeSAI"
    vae = AutoencoderKL.from_pretrained(
        "madebyollin/sdxl-vae-fp16-fix",
        torch_dtype=torch.float16
    ).to("cuda")
    pipe = StableDiffusionXLPipeline.from_pretrained(
        model_id,
        vae=vae,
        torch_dtype=torch.float16,
        use_safetensors=True
    ).to("cuda")
    image = pipe(prompt=prompt, width=1024, height=1024, guidance_scale=7).images[0]
    image.save("output_animesai.png")
    print("✅ 저장 완료: output_animesai.png")
    return image

if __name__ == "__main__":
    prompt = "그녀를 바라보는 한 남자의 야망"
    generate_animesai(prompt)