File size: 2,325 Bytes
17c2c06 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
import streamlit as st
import torch
import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer
from PIL import Image
import warnings
# Disable warnings and progress bars
transformers.logging.set_verbosity_error()
transformers.logging.disable_progress_bar()
warnings.filterwarnings('ignore')
# Set device
device = 'cuda' if torch.cuda.is_available() else 'cpu'
torch.set_default_device(device)
@st.cache_resource
def load_model():
model_name = 'cognitivecomputations/dolphin-vision-72b'
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16,
device_map='auto',
trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(
model_name,
trust_remote_code=True
)
return model, tokenizer
def generate_response(model, tokenizer, prompt, image=None):
messages = [
{"role": "user", "content": f'<image>\n{prompt}' if image else prompt}
]
text = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)
text_chunks = [tokenizer(chunk).input_ids for chunk in text.split('<image>')]
input_ids = torch.tensor(text_chunks[0] + [-200] + text_chunks[1], dtype=torch.long).unsqueeze(0)
if image:
image_tensor = model.process_images([image], model.config).to(dtype=model.dtype)
else:
image_tensor = None
output_ids = model.generate(
input_ids,
images=image_tensor,
max_new_tokens=2048,
use_cache=True
)[0]
return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()
st.title("Chat with DolphinVision 🐬")
model, tokenizer = load_model()
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
image = None
if uploaded_file is not None:
image = Image.open(uploaded_file)
st.image(image, caption='Uploaded Image', use_column_width=True)
user_input = st.text_input("You:", "")
if st.button("Send"):
if user_input:
with st.spinner("Generating response..."):
response = generate_response(model, tokenizer, user_input, image)
st.text_area("DolphinVision:", value=response, height=200)
else:
st.warning("Please enter a message.") |