app.py

import gradio as gr
import shutil
import torch
from transformers import (
    Blip2Processor,
    Blip2ForConditionalGeneration,
    AutoTokenizer,
    AutoModelForCausalLM,
)

# ✅ 모델 로드
device = "cpu"
processor = Blip2Processor.from_pretrained("Salesforce/blip2-flan-t5-xl")
blip_model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-flan-t5-xl").to(device)
llm_tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b-instruct")
llm_model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-7b-instruct").to(device)

# ✅ 기능 정의
def extract_caption(image):
    inputs = processor(images=image, return_tensors="pt").to(device)
    outputs = blip_model.generate(**inputs, max_new_tokens=50)
    return processor.tokenizer.decode(outputs[0], skip_special_tokens=True)

def build_prompt(caption):
    return (
        f"Based on the image description: \"{caption}\", write a children's fairytale.\n"
        "The story must:\n"
        "- Start with 'Once upon a time'\n"
        "- Be at least 10 full sentences long\n"
        "- Include named characters, a clear setting, emotions, a challenge, and a resolution\n"
        "- Avoid mentions of babies or unrelated royalty unless relevant\n"
        "Here is the story:\nOnce upon a time"
    )

def generate_fairytale(image):
    caption = extract_caption(image)
    prompt = build_prompt(caption)
    inputs = llm_tokenizer(prompt, return_tensors="pt").to(device)
    output = llm_model.generate(
        **inputs,
        max_new_tokens=500,
        do_sample=True,
        temperature=0.9,
        top_p=0.95,
        pad_token_id=llm_tokenizer.eos_token_id
    )
    result = llm_tokenizer.decode(output[0], skip_special_tokens=True)
    if "Once upon a time" in result:
        story = "Once upon a time" + result.split("Once upon a time", 1)[-1].strip()
    else:
        story = result
    return caption, story

# ✅ Gradio UI 구성 + 감성 디자인
with gr.Blocks(css="""
body {
    background-color: white;
    font-family: 'Comic Sans MS', cursive;
    color: #333;
}

/* 전체 그라디오 컨테이너 테두리 제거 */
.gradio-container {
    background-color: transparent !important;
    padding: 20px;
    border-radius: 15px;
}

/* 민트 테두리 + 흰 배경 */
.gr-box,
textarea, input, .input-image, .wrap.svelte-1ipelgc {
    background-color: white !important;
    border: 2px solid #a3e4e0 !important;
    border-radius: 12px !important;
    color: #000 !important;
}

/* 버튼 민트색 */
button {
    background-color: #a3e4e0 !important;
    color: #000 !important;
    font-size: 16px;
    border-radius: 10px;
    padding: 10px 20px;
}

/* 헤더 타이틀 */
h1, h3 {
    text-align: center;
    color: #333;
}        

/* 타원형 버튼 공통 스타일 */
.wrap.svelte-1ipelgc button {
    width: 50px !important;
    height: 50px !important;
    border-radius: 50% !important;
    background-size: 60% 60% !important;
    background-repeat: no-repeat !important;
    background-position: center !important;
    background-color: #a3e4e0 !important;
    border: none !important;
}

/* 개별 버튼 아이콘 삽입 */
.wrap.svelte-1ipelgc button:nth-child(1) {
    background-image: url("/file=스크린샷 2025-07-25 101331.png"); /* 업로드 */
}
.wrap.svelte-1ipelgc button:nth-child(2) {
    background-image: url("/file=스크린샷 2025-07-25 101348.png"); /* 카메라 */
}
.wrap.svelte-1ipelgc button:nth-child(3) {
    background-image: url("/file=스크린샷 2025-07-25 101356.png"); /* 갤러리 or 샘플 */
}
}
""") as demo:

    gr.HTML("""
        <div style="padding: 20px;">
            <h1>📖 AI Fairytale Creator</h1>
            <h3>Upload a character image and generate a magical children's story ✨</h3>
        </div>
    """)

    with gr.Column():
        image_input = gr.Image(type="pil", label="🖼️ Upload your character image")
        generate_button = gr.Button("✨ Create Fairytale")
        caption_output = gr.Textbox(label="📌 Image Description")
        story_output = gr.Textbox(label="📚 Generated Fairytale", lines=20)

    generate_button.click(fn=generate_fairytale, inputs=image_input, outputs=[caption_output, story_output])

# ✅ 실행 (필수: share=True)
if __name__ == "__main__":
    demo.launch(share=True)