blip-image-captioning-base01.py

# # 필요한 라이브러리 임포트
# from transformers import BlipProcessor, BlipForConditionalGeneration        # 이미지 → 텍스트 설명 생성용
# from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline    # 텍스트 생성 (동화 생성)용
# from PIL import Image                                                       # 이미지 처리용
# import gradio as gr                                                         # 웹 인터페이스용
# import torch                                                                # PyTorch (모델 구동)

# # 1. BLIP 모델: 이미지 설명 생성용 모델 및 전처리기 로딩
# caption_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
# caption_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

# # 2. KoT5 모델: 설명 → 동화 생성용 한국어 텍스트 생성 모델 로딩
# story_tokenizer = AutoTokenizer.from_pretrained("paust/pko-t5-base")
# story_model = AutoModelForSeq2SeqLM.from_pretrained("paust/pko-t5-base")

# # Hugging Face pipeline으로 간편한 생성 함수 구성
# story_generator = pipeline("text2text-generation", model=story_model, tokenizer=story_tokenizer)

# # 3. 핵심 함수: 이미지 입력 → 설명 생성 → 동화 생성
# def generate_fairytale(image):
#     # 이미지 → 설명 생성 (BLIP)
#     inputs = caption_processor(images=image, return_tensors="pt")              # 이미지 전처리
#     outputs = caption_model.generate(**inputs)                                 # 설명 생성
#     caption = caption_processor.decode(outputs[0], skip_special_tokens=True)   # 설명 텍스트 디코딩

#     # 프롬프트 구성: 설명을 바탕으로 동화를 써달라고 요청
#     prompt = f"다음 설명을 바탕으로 어린이 동화를 써줘:\n\"{caption}\""

#     # 설명 → 동화 생성 (KoT5)
#     story = story_generator(prompt, max_length=300, do_sample=True)[0]['generated_text']

#     # Gradio에 표시할 두 개 결과 반환: 설명, 동화
#     return caption, story

# # 4. Gradio 인터페이스 구성
# interface = gr.Interface(
#     fn=generate_fairytale,                                         # 호출할 함수
#     inputs=gr.Image(type="pil", label="🖼️ 그림을 업로드하세요"),   # 이미지 업로드 창
#     outputs=[
#         gr.Textbox(label="📌 이미지 설명"),                         # 텍스트 설명 출력
#         gr.Textbox(label="📖 생성된 동화")                          # 생성된 동화 출력
#     ],
#     title="🎨 AI 이미지 동화 생성기",                                # 앱 제목
#     description="이미지를 업로드하면 설명과 동화를 자동으로 생성해주는 서비스입니다.",
#     theme="soft"                                                   # 부드러운 Gradio 테마
# )

# # 5. 웹 앱 실행 (URL 콘솔 출력)
# interface.launch(share=True, debug=True, inbrowser=True)


from transformers import BlipProcessor, BlipForConditionalGeneration
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
from PIL import Image
import gradio as gr
import torch

# 1. BLIP 이미지 캡셔닝 모델
caption_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
caption_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

# 2. 영어 동화 생성 모델 (Flan-T5)
story_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
story_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
story_generator = pipeline("text2text-generation", model=story_model, tokenizer=story_tokenizer)

# 3. 메인 함수
def generate_fairytale(image):
    # (1) 이미지 → 캡션 생성
    inputs = caption_processor(images=image, return_tensors="pt").to(caption_model.device)
    outputs = caption_model.generate(**inputs)
    caption = caption_processor.decode(outputs[0], skip_special_tokens=True).strip()

    # (2) 캡션이 너무 짧으면 보정
    if len(caption.split()) < 5:
        caption += ". They appear to be magical characters from a fantasy world."

    # (3) 영어 프롬프트 구성
    prompt = (
        f"Write a short and magical fairytale for children based on the description below.\n"
        f"Description: \"{caption}\"\n"
        f"Story:"
    )

    # (4) 동화 생성
    story = story_generator(prompt, max_length=300, do_sample=True)[0]['generated_text']

    return caption, story

# 4. Gradio UI
interface = gr.Interface(
    fn=generate_fairytale,
    inputs=gr.Image(type="pil", label="🖼️ Upload an Image"),
    outputs=[
        gr.Textbox(label="📌 Image Description"),
        gr.Textbox(label="📖 Generated Fairytale")
    ],
    title="🧚 AI Fairytale Generator from Image",
    description="Upload an image, and this AI will create a magical story based on it!",
    theme="soft"
)

# 5. 앱 실행
interface.launch(share=True, debug=True, inbrowser=True)