1
2# Authenticate with Hugging Face
3from huggingface_hub import login
4
5# Log in to Hugging Face using the provided token
6hf_token = '----xxxxx----'
7login(hf_token)
8
9# Demo
10import gradio as gr
11import spaces
12from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, TextIteratorStreamer
13from qwen_vl_utils import process_vision_info
14import torch
15from PIL import Image
16import os
17import uuid
18import io
19from threading import Thread
20from reportlab.lib.pagesizes import A4
21from reportlab.lib.styles import getSampleStyleSheet
22from reportlab.lib import colors
23from reportlab.platypus import SimpleDocTemplate, Image as RLImage, Paragraph, Spacer
24from reportlab.pdfbase import pdfmetrics
25from reportlab.pdfbase.ttfonts import TTFont
26import docx
27from docx.enum.text import WD_ALIGN_PARAGRAPH
28
29# Define model options
30MODEL_OPTIONS = {
31 "ChemQwen": "prithivMLmods/ChemQwen2-vL",
32}
33
34# Preload models and processors into CUDA
35models = {}
36processors = {}
37for name, model_id in MODEL_OPTIONS.items():
38 print(f"Loading {name}...")
39 models[name] = Qwen2VLForConditionalGeneration.from_pretrained(
40 model_id,
41 trust_remote_code=True,
42 torch_dtype=torch.float16
43 ).to("cuda").eval()
44 processors[name] = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
45
46image_extensions = Image.registered_extensions()
47
48def identify_and_save_blob(blob_path):
49 """Identifies if the blob is an image and saves it."""
50 try:
51 with open(blob_path, 'rb') as file:
52 blob_content = file.read()
53 try:
54 Image.open(io.BytesIO(blob_content)).verify() # Check if it's a valid image
55 extension = ".png" # Default to PNG for saving
56 media_type = "image"
57 except (IOError, SyntaxError):
58 raise ValueError("Unsupported media type. Please upload a valid image.")
59
60 filename = f"temp_{uuid.uuid4()}_media{extension}"
61 with open(filename, "wb") as f:
62 f.write(blob_content)
63
64 return filename, media_type
65
66 except FileNotFoundError:
67 raise ValueError(f"The file {blob_path} was not found.")
68 except Exception as e:
69 raise ValueError(f"An error occurred while processing the file: {e}")
70
71@spaces.GPU
72def qwen_inference(model_name, media_input, text_input=None):
73 """Handles inference for the selected model."""
74 model = models[model_name]
75 processor = processors[model_name]
76
77 if isinstance(media_input, str):
78 media_path = media_input
79 if media_path.endswith(tuple([i for i in image_extensions.keys()])):
80 media_type = "image"
81 else:
82 try:
83 media_path, media_type = identify_and_save_blob(media_input)
84 except Exception as e:
85 raise ValueError("Unsupported media type. Please upload a valid image.")
86
87 messages = [
88 {
89 "role": "user",
90 "content": [
91 {
92 "type": media_type,
93 media_type: media_path
94 },
95 {"type": "text", "text": text_input},
96 ],
97 }
98 ]
99
100 text = processor.apply_chat_template(
101 messages, tokenize=False, add_generation_prompt=True
102 )
103 image_inputs, _ = process_vision_info(messages)
104 inputs = processor(
105 text=[text],
106 images=image_inputs,
107 padding=True,
108 return_tensors="pt",
109 ).to("cuda")
110
111 streamer = TextIteratorStreamer(
112 processor.tokenizer, skip_prompt=True, skip_special_tokens=True
113 )
114 generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
115
116 thread = Thread(target=model.generate, kwargs=generation_kwargs)
117 thread.start()
118
119 buffer = ""
120 for new_text in streamer:
121 buffer += new_text
122 # Remove <|im_end|> or similar tokens from the output
123 buffer = buffer.replace("<|im_end|>", "")
124 yield buffer
125
126def format_plain_text(output_text):
127 """Formats the output text as plain text without LaTeX delimiters."""
128 # Remove LaTeX delimiters and convert to plain text
129 plain_text = output_text.replace("\\(", "").replace("\\)", "").replace("\\[", "").replace("\\]", "")
130 return plain_text
131
132def generate_document(media_path, output_text, file_format, font_size, line_spacing, alignment, image_size):
133 """Generates a document with the input image and plain text output."""
134 plain_text = format_plain_text(output_text)
135 if file_format == "pdf":
136 return generate_pdf(media_path, plain_text, font_size, line_spacing, alignment, image_size)
137 elif file_format == "docx":
138 return generate_docx(media_path, plain_text, font_size, line_spacing, alignment, image_size)
139
140def generate_pdf(media_path, plain_text, font_size, line_spacing, alignment, image_size):
141 """Generates a PDF document."""
142 filename = f"output_{uuid.uuid4()}.pdf"
143 doc = SimpleDocTemplate(
144 filename,
145 pagesize=A4,
146 rightMargin=inch,
147 leftMargin=inch,
148 topMargin=inch,
149 bottomMargin=inch
150 )
151 styles = getSampleStyleSheet()
152 styles["Normal"].fontSize = int(font_size)
153 styles["Normal"].leading = int(font_size) * line_spacing
154 styles["Normal"].alignment = {
155 "Left": 0,
156 "Center": 1,
157 "Right": 2,
158 "Justified": 4
159 }[alignment]
160
161 story = []
162
163 # Add image with size adjustment
164 image_sizes = {
165 "Small": (200, 200),
166 "Medium": (400, 400),
167 "Large": (600, 600)
168 }
169 img = RLImage(media_path, width=image_sizes[image_size][0], height=image_sizes[image_size][1])
170 story.append(img)
171 story.append(Spacer(1, 12))
172
173 # Add plain text output
174 text = Paragraph(plain_text, styles["Normal"])
175 story.append(text)
176
177 doc.build(story)
178 return filename
179
180def generate_docx(media_path, plain_text, font_size, line_spacing, alignment, image_size):
181 """Generates a DOCX document."""
182 filename = f"output_{uuid.uuid4()}.docx"
183 doc = docx.Document()
184
185 # Add image with size adjustment
186 image_sizes = {
187 "Small": docx.shared.Inches(2),
188 "Medium": docx.shared.Inches(4),
189 "Large": docx.shared.Inches(6)
190 }
191 doc.add_picture(media_path, width=image_sizes[image_size])
192 doc.add_paragraph()
193
194 # Add plain text output
195 paragraph = doc.add_paragraph()
196 paragraph.paragraph_format.line_spacing = line_spacing
197 paragraph.paragraph_format.alignment = {
198 "Left": WD_ALIGN_PARAGRAPH.LEFT,
199 "Center": WD_ALIGN_PARAGRAPH.CENTER,
200 "Right": WD_ALIGN_PARAGRAPH.RIGHT,
201 "Justified": WD_ALIGN_PARAGRAPH.JUSTIFY
202 }[alignment]
203 run = paragraph.add_run(plain_text)
204 run.font.size = docx.shared.Pt(int(font_size))
205
206 doc.save(filename)
207 return filename
208
209# CSS for output styling
210css = """
211 #output {
212 height: 500px;
213 overflow: auto;
214 border: 1px solid #ccc;
215 }
216.submit-btn {
217 background-color: #cf3434 !important;
218 color: white !important;
219}
220.submit-btn:hover {
221 background-color: #ff2323 !important;
222}
223.download-btn {
224 background-color: #35a6d6 !important;
225 color: white !important;
226}
227.download-btn:hover {
228 background-color: #22bcff !important;
229}
230"""
231
232# Gradio app setup
233with gr.Blocks(css=css) as demo:
234 gr.Markdown("# ChemQwen Chemical Identifier")
235
236 with gr.Tab(label="Image Input"):
237
238 with gr.Row():
239 with gr.Column():
240 model_choice = gr.Dropdown(
241 label="Model Selection",
242 choices=list(MODEL_OPTIONS.keys()),
243 value="ChemQwen"
244 )
245 input_media = gr.File(
246 label="Upload Image", type="filepath"
247 )
248 text_input = gr.Textbox(label="Question", placeholder="Ask a question about the image...")
249 submit_btn = gr.Button(value="Submit", elem_classes="submit-btn")
250
251 with gr.Column():
252 output_text = gr.Textbox(label="Output Text", lines=10)
253 plain_text_output = gr.Textbox(label="Standardized Plain Text", lines=10)
254
255 submit_btn.click(
256 qwen_inference, [model_choice, input_media, text_input], [output_text]
257 ).then(
258 lambda output_text: format_plain_text(output_text), [output_text], [plain_text_output]
259 )
260
261 # Add examples directly usable by clicking
262 with gr.Row():
263 with gr.Column():
264 line_spacing = gr.Dropdown(
265 choices=[0.5, 1.0, 1.15, 1.5, 2.0, 2.5, 3.0],
266 value=1.5,
267 label="Line Spacing"
268 )
269 font_size = gr.Dropdown(
270 choices=["8", "10", "12", "14", "16", "18", "20", "22", "24"],
271 value="18",
272 label="Font Size"
273 )
274 alignment = gr.Dropdown(
275 choices=["Left", "Center", "Right", "Justified"],
276 value="Justified",
277 label="Text Alignment"
278 )
279 image_size = gr.Dropdown(
280 choices=["Small", "Medium", "Large"],
281 value="Small",
282 label="Image Size"
283 )
284 file_format = gr.Radio(["pdf", "docx"], label="File Format", value="pdf")
285 get_document_btn = gr.Button(value="Get Document", elem_classes="download-btn")
286
287 get_document_btn.click(
288 generate_document, [input_media, output_text, file_format, font_size, line_spacing, alignment, image_size], gr.File(label="Download Document")
289 )
290
291demo.launch(debug=True)