로컬에 streamlit 로 llama 3.2 Vision 채팅창 만들기

로컬에 streamlit 로 llama 3.2 Vision 채팅창 만들기

 

streamlit 를 이용해서 llama 3.2 vision 을 이용해 이미지를 분석하는 채팅창을 만들어 봤습니다.

 

코드

import streamlit as st
import requests
import torch
from PIL import Image
from transformers import MllamaForConditionalGeneration, AutoProcessor
from io import BytesIO

# Streamlit 페이지 설정
st.set_page_config(page_title="LLaMA Vision Chat", page_icon="🦙")

@st.cache_resource
def load_model():
    model_id = "meta-llama/Llama-3.2-11B-Vision-Instruct"
    model = MllamaForConditionalGeneration.from_pretrained(
        model_id,
        torch_dtype=torch.bfloat16,
        device_map={"": 0}
    )
    processor = AutoProcessor.from_pretrained(model_id)
    return model, processor

# 모델 및 프로세서 로드
model, processor = load_model()

# CUDA 사용 가능 여부 확인
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
st.sidebar.write(f"Using device: {device}")

# 사이드바에 이미지 업로드 위젯 추가
uploaded_file = st.sidebar.file_uploader("이미지를 업로드하세요", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.sidebar.image(image, caption="업로드된 이미지", use_column_width=True)

# 메인 페이지 제목
st.title("LLaMA Vision Chat")

# 채팅 기록을 저장할 세션 상태 초기화
if "messages" not in st.session_state:
    st.session_state.messages = []

# 채팅 기록 표시
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# 사용자 입력 받기
if prompt := st.chat_input("메시지를 입력하세요..."):
    # 이미지가 업로드되었는지 확인
    if uploaded_file is None:
        st.error("먼저 이미지를 업로드해주세요.")
    else:
        # 사용자 메시지를 채팅 기록에 추가
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)

        # 봇 응답 생성
        with st.chat_message("assistant"):
            message_placeholder = st.empty()
            full_response = ""

            messages = [
                {"role": "user", "content": [
                    {"type": "image"},
                    {"type": "text", "text": prompt}
                ]}
            ]
            input_text = processor.apply_chat_template(messages, add_generation_prompt=True)

            inputs = processor(
                image,
                input_text,
                add_special_tokens=False,
                return_tensors="pt",
            ).to(device)

            output = model.generate(**inputs, max_new_tokens=1000)
            full_response = processor.decode(output[0], skip_special_tokens=True)

            message_placeholder.markdown(full_response)

        # 봇 응답을 채팅 기록에 추가
        st.session_state.messages.append({"role": "assistant", "content": full_response})

# Streamlit 실행 방법 안내
st.sidebar.markdown("## 실행 방법")
st.sidebar.code("streamlit run your_script_name.py")

 

 

결과

왼쪽에 이미지 업로드 후 설명해 달라고 하였습니다.