from diffusers import StableDiffusionPipeline
import torch
from PIL import Image
import os

# 이미지 생성 파이프라인
image_gen = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    torch_dtype=torch.float16,
).to("cuda" if torch.cuda.is_available() else "cpu")

def generate_comic_manual(cuts, output_dir="comic_outputs"):
    """
    cuts: 4개의 텍스트(컷 설명)가 담긴 리스트
    """
    if len(cuts) != 4:
        raise ValueError("정확히 4개의 컷 설명을 입력해주세요.")

    os.makedirs(output_dir, exist_ok=True)

    image_paths = []
    for i, cut in enumerate(cuts):
        print(f"컷 {i+1}: {cut}")
        image = image_gen(cut).images[0]
        path = os.path.join(output_dir, f"cut_{i+1}.png")
        image.save(path)
        image_paths.append(path)

    print(f"\n✅ 저장 완료! 총 {len(image_paths)}장의 컷이 생성되었습니다.")
    return image_paths

# 사용 예시
if __name__ == "__main__":
    cuts = [
    "A cat launches in a rocket into space",
    "A cat arrives on an alien planet",
    "The cat meets a robot friend",
    "The cat returns to Earth and writes in a diary"
]

    generate_comic_manual(cuts)