최종 앱

app.py

@@ -1,57 +1,135 @@
 import gradio as gr
-from transformers import Blip2Processor, Blip2ForConditionalGeneration
-from PIL import Image
+import shutil
 import torch
+from transformers import (
+    Blip2Processor,
+    Blip2ForConditionalGeneration,
+    AutoTokenizer,
+    AutoModelForCausalLM,
+)
 
-# ✅ 모델 로딩
+# ✅ 모델 로드
+device = "cpu"
 processor = Blip2Processor.from_pretrained("Salesforce/blip2-flan-t5-xl")
-model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-flan-t5-xl").to("cuda" if torch.cuda.is_available() else "cpu")
+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)
 
-# 📌 Step 1: 이미지 설명 추출 (이미지 → 캡션)
+# ✅ 기능 정의
 def extract_caption(image):
-    inputs = processor(images=image, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=50)
-    caption = processor.tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return caption
+    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)
 
-# 📌 Step 2: 캡션 기반 동화 프롬프트 생성
-def build_prompt_from_caption(caption):
+def build_prompt(caption):
     return (
-        f"Write a magical and fun children's fairytale based on this description: \"{caption}\". "
-        "Start with 'Once upon a time' and continue for at least 7 sentences. "
-        "Include characters, emotions, the setting, and a twist. Make it feel like a real story."
+        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"
     )
 
-# 📌 Step 3: 동화 생성 (캡션 + 프롬프트 → 텍스트)
 def generate_fairytale(image):
     caption = extract_caption(image)
-    prompt = build_prompt_from_caption(caption)
-    inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(
+    prompt = build_prompt(caption)
+    inputs = llm_tokenizer(prompt, return_tensors="pt").to(device)
+    output = llm_model.generate(
         **inputs,
-        max_new_tokens=400,
+        max_new_tokens=500,
         do_sample=True,
-        temperature=0.95,
-        top_p=0.9,
+        temperature=0.9,
+        top_p=0.95,
+        pad_token_id=llm_tokenizer.eos_token_id
     )
-    story = processor.tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return story
+    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() as demo:
-    gr.Markdown("## 🧚‍♂️ 이미지 기반 AI 동화 생성기\n사진을 업로드하면 동화로 바꿔드립니다!")
+# ✅ Gradio UI 구성 + 감성 디자인
+with gr.Blocks(css="""
+body {
+    background-color: white;
+    font-family: 'Comic Sans MS', cursive;
+    color: #333;
+}
 
-    with gr.Row():
-        image_input = gr.Image(type="pil", label="📸 이미지 업로드")
+/* 전체 그라디오 컨테이너 테두리 제거 */
+.gradio-container {
+    background-color: transparent !important;
+    padding: 20px;
+    border-radius: 15px;
+}
 
-    with gr.Row():
-        generate_button = gr.Button("✨ 동화 만들기")
+/* 민트 테두리 + 흰 배경 */
+.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;
+}
 
-    with gr.Row():
-        output_text = gr.Textbox(label="📖 생성된 동화", lines=10)
+/* 버튼 민트색 */
+button {
+    background-color: #a3e4e0 !important;
+    color: #000 !important;
+    font-size: 16px;
+    border-radius: 10px;
+    padding: 10px 20px;
+}
 
-    generate_button.click(fn=generate_fairytale, inputs=[image_input], outputs=[output_text])
+/* 헤더 타이틀 */
+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)