app.py

import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
from PIL import Image
import io

@st.cache_resource
def load_model():
    model_name = "Qwen/Qwen2-VL-7B-Instruct"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
    return tokenizer, model

def generate_response(prompt, image, tokenizer, model):
    if image:
        image = Image.open(image).convert('RGB')
        inputs = tokenizer.from_pretrained(prompt, images=[image], return_tensors='pt').to(model.device)
    else:
        inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=100)
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

st.title("Чат с Qwen VL-7B-Instruct")

tokenizer, model = load_model()

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 "image" in message:
            st.image(message["image"])

prompt = st.chat_input("Введите ваше сообщение")
uploaded_file = st.file_uploader("Загрузите изображение (необязательно)", type=["png", "jpg", "jpeg"])

if prompt or uploaded_file:
    if uploaded_file:
        image = Image.open(uploaded_file)
        st.session_state.messages.append({"role": "user", "content": prompt or "Опишите это изображение", "image": uploaded_file})
        with st.chat_message("user"):
            if prompt:
                st.markdown(prompt)
            st.image(image)
    else:
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)
    
    with st.chat_message("assistant"):
        response = generate_response(prompt, uploaded_file, tokenizer, model)
        st.markdown(response)
    
    st.session_state.messages.append({"role": "assistant", "content": response})