Spaces:
Runtime error
Runtime error
File size: 2,131 Bytes
979c542 |
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 |
import streamlit as st
import torch
from PIL import Image
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
# Inference steps taken from https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct
@st.cache_resource
def get_model(model_path):
try:
with st.spinner(f"Loading model {model_path}"):
device = "cuda" if torch.cuda.is_available() else "cpu"
# Load the model here
model_import = Qwen2VLForConditionalGeneration.from_pretrained(
model_path, torch_dtype="auto", device_map=device
)
processor_import = AutoProcessor.from_pretrained(model_path)
return model_import, processor_import
except Exception as e:
st.error(f"Error loading model: {e}")
return None, None
def run_inference(input_file, model_path, args):
model, processor = get_model(model_path)
if model is None or processor is None:
return "Error loading model."
image = Image.open(input_file)
conversation = [
{
"role": "user",
"content": [
{"type": "image"},
{"type": "text", "text": "Please generate TikZ code to draw the diagram of the given image."}
],
}
]
text_prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
inputs = processor(image, text_prompt, return_tensors="pt").to("cuda")
output_ids = model.generate(**inputs,
max_new_tokens=args.max_length,
do_sample=True,
top_p=args.top_p,
top_k=args.top_k,
num_return_sequences=1,
temperature=args.temperature
)
generated_ids = [
output_ids[len(input_ids):]
for input_ids, output_ids in zip(inputs.input_ids, output_ids)
]
output_text = processor.batch_decode(
generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
)
return output_text
|