Spaces:
No application file
No application file
File size: 6,035 Bytes
7e7196e | 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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | import streamlit as st
import torch
from PIL import Image
import io
import os
from transformers import AutoProcessor, Qwen2VLForConditionalGeneration, BitsAndBytesConfig
from peft import PeftModel
# =========================
# PAGE CONFIG
# =========================
st.set_page_config(
page_title="Markdown Generator - Qwen2-VL",
layout="wide",
initial_sidebar_state="expanded"
)
# =========================
# CUSTOM STYLING
# =========================
st.markdown("""
<style>
.main {
padding-top: 2rem;
}
.stButton button {
width: 100%;
padding: 0.5rem;
border-radius: 0.5rem;
}
.output-box {
background-color: #f0f2f6;
padding: 1rem;
border-radius: 0.5rem;
margin-top: 1rem;
}
</style>
""", unsafe_allow_html=True)
MODEL_NAME = "Qwen/Qwen2-VL-2B-Instruct"
CHECKPOINT_PATH = "./qlora-vlm"
IMAGE_SIZE = 512
@st.cache_resource
def load_model():
try:
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4"
)
processor = AutoProcessor.from_pretrained(MODEL_NAME)
base_model = Qwen2VLForConditionalGeneration.from_pretrained(
MODEL_NAME,
quantization_config=bnb_config,
device_map="auto",
torch_dtype=torch.float16,
trust_remote_code=True
)
# load lora adapters
if os.path.exists(CHECKPOINT_PATH):
model = PeftModel.from_pretrained(base_model, CHECKPOINT_PATH)
else:
model = base_model
model.eval()
return processor, model
except Exception as e:
st.error(f"Error loading model: {str(e)}")
return None, None
def generate_markdown(processor, model, image_path_or_pil):
try:
if isinstance(image_path_or_pil, str):
image = Image.open(image_path_or_pil).convert("RGB")
else:
image = image_path_or_pil.convert("RGB")
image = image.resize((IMAGE_SIZE, IMAGE_SIZE))
messages = [
{
"role": "user",
"content": [
{"type": "image", "image": image},
{"type": "text", "text": "Please describe this image in markdown format, including any text, equations, or structured content you see."}
],
}
]
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = processor(text=text, images=[image], return_tensors="pt")
inputs = {k: v.to(model.device) for k, v in inputs.items()}
with torch.no_grad():
output_ids = model.generate(**inputs, max_new_tokens=1024)
generated_text = processor.decode(output_ids[0], skip_special_tokens=True)
if "assistant" in generated_text:
result = generated_text.split("assistant")[-1].strip()
else:
result = generated_text
return result
except Exception as e:
return f"Error generating markdown: {str(e)}"
def main():
# Header
st.title("Markdown Generator")
st.subheader("Generate Markdown from Images using Qwen2-VL")
# Load model
with st.spinner("Loading model... This may take a moment on first run"):
processor, model = load_model()
if processor is None or model is None:
st.error("Failed to load model. Please check the model path and configuration.")
return
st.success("Model loaded successfully!")
st.divider()
tab1, tab2 = st.tabs(["Upload Image", "ℹAbout"])
with tab1:
col1, col2 = st.columns([2, 1])
with col1:
st.subheader("Upload or Paste Image")
uploaded_file = st.file_uploader(
"Choose an image",
type=["jpg", "jpeg", "png", "webp", "gif"],
help="Upload an image to generate markdown description"
)
with col2:
st.subheader("Quick Actions")
if st.button("Use Sample Image", use_container_width=True):
st.session_state.use_sample = True
st.divider()
# Process image
if uploaded_file is not None or st.session_state.get('use_sample', False):
if uploaded_file is not None:
image = Image.open(uploaded_file)
st.session_state.use_sample = False
else:
# Create a simple sample image if no file uploaded
image = Image.new('RGB', (512, 512), color='lightblue')
col1, col2 = st.columns(2)
with col1:
st.subheader("Input Image")
st.image(image, use_column_width=True)
with col2:
st.subheader("Generated Markdown")
if st.button("Generate Markdown", use_container_width=True, type="primary"):
with st.spinner("Generating markdown... This may take some time"):
markdown_result = generate_markdown(processor, model, image)
st.session_state.markdown_result = markdown_result
if st.session_state.get('markdown_result'):
result_text = st.session_state.markdown_result
st.markdown(result_text)
with st.expander("View Raw Text"):
st.code(result_text, language="markdown")
if __name__ == "__main__":
main()
|