MiniCPM-V-Test / src /streamlit_app.py
syedkhalid076's picture
Update src/streamlit_app.py
d0e6f56 verified
# app.py
import streamlit as st
from PIL import Image
import io
import torch
from transformers import AutoModel, AutoTokenizer
@st.cache_resource(show_spinner=True)
def load_model():
model = AutoModel.from_pretrained("openbmb/MiniCPM-V", trust_remote_code=True, torch_dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained("openbmb/MiniCPM-V", trust_remote_code=True)
# send to GPU/CPU as per availability
if torch.cuda.is_available():
model = model.to(device='cuda', dtype=torch.bfloat16)
else:
model = model.to(device='cpu')
model.eval()
return model, tokenizer
model, tokenizer = load_model()
st.set_page_config(page_title="MiniCPM-V Chat", layout="wide")
st.title("📄 MiniCPM-V Chat — Image/Text → Markdown / Chat")
if "history" not in st.session_state:
st.session_state.history = []
# Sidebar: upload or text input
with st.sidebar:
uploaded_file = st.file_uploader("Upload image / pdf-page (image) or enter text:", type=["jpg","jpeg","png","pdf","txt"])
text_input = st.text_area("Or paste text here:")
# Main chat interface
for msg in st.session_state.history:
with st.chat_message(msg["role"]):
st.markdown(msg["content"])
def run_minicpm_v(input_image=None, input_text=None, history=None):
"""
input_image: PIL.Image or None
input_text: str or None
history: list of prior messages (role, content)
"""
msgs = []
if history:
msgs = history.copy()
# Compose new user message
user_content = ""
if input_image is not None:
user_content = "[Image Uploaded]\n" # or some marker + optional prompt
if input_text:
user_content += input_text
msgs.append({"role": "user", "content": user_content})
# Run the multimodal chat
res, context, _ = model.chat(
image=input_image,
msgs=msgs,
context=None,
tokenizer=tokenizer,
sampling=True,
temperature=0.7,
)
return res
if uploaded_file is not None or text_input:
with st.chat_message("user"):
if uploaded_file is not None:
st.image(uploaded_file, caption="Uploaded")
if text_input:
st.markdown(text_input)
# Process input
input_image = None
input_text = None
if uploaded_file is not None:
# try open as image
try:
input_image = Image.open(uploaded_file).convert("RGB")
except Exception as e:
st.error("Could not open uploaded file as image.")
if text_input:
input_text = text_input
with st.spinner("Thinking..."):
reply = run_minicpm_v(input_image=input_image, input_text=input_text, history=st.session_state.history)
st.session_state.history.append({"role": "assistant", "content": reply})
with st.chat_message("assistant"):
st.markdown(reply)
st.chat_input(placeholder="Send more text or upload another file…") # optional extra prompt