project / manual_4cut_comic.py
seokhyun119's picture
4์ปท ์ด๋ฏธ์ง€ ์ƒ์„ฑ ๋ชจ๋ธ
cff38af
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)