Update app.py

app.py

@@ -4,21 +4,33 @@ from transformers import pipeline
 # 1) 이미지 캡셔닝 파이프라인 초기화
 captioner = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
 
-# 2) 스타일 변환 함수
-def style_convert(raw_caption, style):
+def generate_caption(image, style):
+    # 1) 원본 캡션 얻기
+    output = captioner(image)
+    # output 예시: [{"generated_text": "..."}] 또는 [{"text": "..."}] 혹은 ["..."]
+    # 안전하게 추출
+    if isinstance(output, list) and output:
+        first = output[0]
+        if isinstance(first, dict):
+            raw_caption = first.get("generated_text") or first.get("text") or ""
+        else:
+            raw_caption = str(first)
+    else:
+        raw_caption = str(output)
+
+    raw_caption = raw_caption.strip()
+    # 2) 스타일 변환
     if style == "TOEIC Speaking Part 1":
         return f"Q: What do you see in the picture?\nA: {raw_caption.capitalize()}."
     elif style == "IELTS Describe a Photo":
-        return f"Describe the photo in two sentences:\n1. {raw_caption.capitalize()}.\n2. It also shows the context of daily life."
+        return (
+            "Describe the photo in two sentences:\n"
+            f"1. {raw_caption.capitalize()}.\n"
+            "2. It also shows the context of daily life."
+        )
     else:
         return raw_caption
 
-# 3) Gradio 인터페이스 함수
-def generate_caption(image, style):
-    result = captioner(image, max_length=30, num_beams=3)[0]["generated_text"]
-    return style_convert(result, style)
-
-# 4) Gradio Blocks 정의
 with gr.Blocks() as demo:
     gr.Markdown("## 📸 이미지 캡셔닝 → English Test 스타일 문장 생성")
     with gr.Row():