import gradio as gr import spaces from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, TextIteratorStreamer from qwen_vl_utils import process_vision_info import torch from PIL import Image import os import uuid import io from threading import Thread from reportlab.lib.pagesizes import A4 from reportlab.lib.styles import getSampleStyleSheet from reportlab.lib import colors from reportlab.platypus import SimpleDocTemplate, Image as RLImage, Paragraph, Spacer from reportlab.lib.units import inch from reportlab.pdfbase import pdfmetrics from reportlab.pdfbase.ttfonts import TTFont import docx from docx.enum.text import WD_ALIGN_PARAGRAPH # Define model options MODEL_OPTIONS = { "Qwen2VL Base": "Qwen/Qwen2-VL-2B-Instruct", "Latex OCR": "prithivMLmods/Qwen2-VL-OCR-2B-Instruct", "Math Prase": "prithivMLmods/Qwen2-VL-Math-Prase-2B-Instruct", "Text Analogy Ocrtest": "prithivMLmods/Qwen2-VL-Ocrtest-2B-Instruct" } # Preload models and processors into CUDA models = {} processors = {} for name, model_id in MODEL_OPTIONS.items(): print(f"Loading {name}...") models[name] = Qwen2VLForConditionalGeneration.from_pretrained( model_id, trust_remote_code=True, torch_dtype=torch.float16 ).to("cuda").eval() processors[name] = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) image_extensions = Image.registered_extensions() def identify_and_save_blob(blob_path): try: with open(blob_path, 'rb') as file: blob_content = file.read() try: Image.open(io.BytesIO(blob_content)).verify() # Validate image extension = ".png" # Default extension media_type = "image" except (IOError, SyntaxError): raise ValueError("Unsupported media type. Please upload a valid image.") filename = f"temp_{uuid.uuid4()}_media{extension}" with open(filename, "wb") as f: f.write(blob_content) return filename, media_type except FileNotFoundError: raise ValueError(f"The file {blob_path} was not found.") except Exception as e: raise ValueError(f"Error processing file: {e}") @spaces.GPU def qwen_inference(model_name, media_input, text_input=None): model = models[model_name] processor = processors[model_name] if isinstance(media_input, str): media_path = media_input if media_path.endswith(tuple(image_extensions.keys())): media_type = "image" else: try: media_path, media_type = identify_and_save_blob(media_input) except Exception as e: raise ValueError("Unsupported media type. Please upload a valid image.") messages = [{ "role": "user", "content": [ {"type": media_type, media_type: media_path}, {"type": "text", "text": text_input}, ], }] text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) image_inputs, _ = process_vision_info(messages) inputs = processor(text=[text], images=image_inputs, padding=True, return_tensors="pt").to("cuda") streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True) generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024) thread = Thread(target=model.generate, kwargs=generation_kwargs) thread.start() buffer = "" for new_text in streamer: buffer += new_text buffer = buffer.replace("<|im_end|>", "") yield buffer def format_plain_text(output_text): return output_text.replace("\\(", "").replace("\\)", "").replace("\\[", "").replace("\\]", "") def generate_document(media_path, output_text, file_format, font_choice, font_size, line_spacing, alignment, image_size): plain_text = format_plain_text(output_text) if file_format == "pdf": return generate_pdf(media_path, plain_text, font_choice, font_size, line_spacing, alignment, image_size) elif file_format == "docx": return generate_docx(media_path, plain_text, font_choice, font_size, line_spacing, alignment, image_size) def generate_pdf(media_path, plain_text, font_choice, font_size, line_spacing, alignment, image_size): filename = f"output_{uuid.uuid4()}.pdf" doc = SimpleDocTemplate(filename, pagesize=A4, rightMargin=inch, leftMargin=inch, topMargin=inch, bottomMargin=inch) styles = getSampleStyleSheet() styles["Normal"].fontName = font_choice styles["Normal"].fontSize = int(font_size) styles["Normal"].leading = int(font_size) * line_spacing styles["Normal"].alignment = {"Left": 0, "Center": 1, "Right": 2, "Justified": 4}[alignment] font_path = f"font/{font_choice}" pdfmetrics.registerFont(TTFont(font_choice, font_path)) story = [] image_sizes = {"Small": (200, 200), "Medium": (400, 400), "Large": (600, 600)} img = RLImage(media_path, width=image_sizes[image_size][0], height=image_sizes[image_size][1]) story.append(img) story.append(Spacer(1, 12)) story.append(Paragraph(plain_text, styles["Normal"])) doc.build(story) return filename def generate_docx(media_path, plain_text, font_choice, font_size, line_spacing, alignment, image_size): filename = f"output_{uuid.uuid4()}.docx" doc = docx.Document() image_sizes = {"Small": docx.shared.Inches(2), "Medium": docx.shared.Inches(4), "Large": docx.shared.Inches(6)} doc.add_picture(media_path, width=image_sizes[image_size]) doc.add_paragraph() paragraph = doc.add_paragraph() paragraph.paragraph_format.line_spacing = line_spacing paragraph.paragraph_format.alignment = { "Left": WD_ALIGN_PARAGRAPH.LEFT, "Center": WD_ALIGN_PARAGRAPH.CENTER, "Right": WD_ALIGN_PARAGRAPH.RIGHT, "Justified": WD_ALIGN_PARAGRAPH.JUSTIFY }[alignment] run = paragraph.add_run(format_plain_text(output_text)) run.font.name = font_choice run.font.size = docx.shared.Pt(int(font_size)) doc.save(filename) return filename # CSS for compact styling css = """ #output { height: 300px; overflow: auto; border: 1px solid #ccc; } .submit-btn { background-color: #cf3434 !important; color: white !important; } .submit-btn:hover { background-color: #ff2323 !important; } .download-btn { background-color: #35a6d6 !important; color: white !important; } .download-btn:hover { background-color: #22bcff !important; } .compact { margin: 5px 0; } """ with gr.Blocks(css=css) as demo: gr.Markdown("# Qwen2VL: Compact Vision & Language Processing") with gr.Row(): with gr.Column(scale=1): model_choice = gr.Dropdown(label="Model", choices=list(MODEL_OPTIONS.keys()), value="Latex OCR", elem_classes="compact") input_media = gr.File(label="Upload Image", type="filepath", elem_classes="compact") text_input = gr.Textbox(label="Question", placeholder="Ask about the image...", elem_classes="compact") submit_btn = gr.Button("Submit", elem_classes="submit-btn compact") with gr.Column(scale=1): output_text = gr.Textbox(label="Output", lines=8, elem_classes="compact") plain_text_output = gr.Textbox(label="Plain Text", lines=8, elem_classes="compact") submit_btn.click(qwen_inference, [model_choice, input_media, text_input], [output_text] ).then(lambda txt: format_plain_text(txt), [output_text], [plain_text_output]) # Examples section remains compact gr.Examples( examples=[ ["examples/1.png", "summarize the letter", "Text Analogy Ocrtest"], ["examples/2.jpg", "Summarize the full image in detail", "Latex OCR"], ["examples/3.png", "Describe the photo", "Qwen2VL Base"], ["examples/4.png", "summarize and solve the problem", "Math Prase"], ], inputs=[input_media, text_input, model_choice], outputs=[output_text, plain_text_output], fn=lambda img, question, model: qwen_inference(model, img, question), cache_examples=False ) # Advanced options tucked into an accordion with gr.Accordion("Advanced Document Options", open=False): with gr.Row(): line_spacing = gr.Dropdown(choices=[0.5, 1.0, 1.15, 1.5, 2.0, 2.5, 3.0], value=1.5, label="Line Spacing", elem_classes="compact") font_size = gr.Dropdown(choices=["8", "10", "12", "14", "16", "18", "20", "22", "24"], value="18", label="Font Size", elem_classes="compact") with gr.Row(): font_choice = gr.Dropdown( choices=["DejaVuMathTeXGyre.ttf", "FiraCode-Medium.ttf", "InputMono-Light.ttf", "JetBrainsMono-Thin.ttf", "ProggyCrossed Regular Mac.ttf", "SourceCodePro-Black.ttf", "arial.ttf", "calibri.ttf", "mukta-malar-extralight.ttf", "noto-sans-arabic-medium.ttf", "times new roman.ttf", "ANGSA.ttf", "Book-Antiqua.ttf", "CONSOLA.TTF", "COOPBL.TTF", "Rockwell-Bold.ttf", "Candara Light.TTF", "Carlito-Regular.ttf", "Castellar.ttf", "Courier New.ttf", "LSANS.TTF", "Lucida Bright Regular.ttf", "TRTempusSansITC.ttf", "Verdana.ttf", "bell-mt.ttf", "eras-itc-light.ttf", "fonnts.com-aptos-light.ttf", "georgia.ttf", "segoeuithis.ttf", "youyuan.TTF", "TfPonetoneExpanded-7BJZA.ttf"], value="youyuan.TTF", label="Font Choice", elem_classes="compact") alignment = gr.Dropdown(choices=["Left", "Center", "Right", "Justified"], value="Justified", label="Alignment", elem_classes="compact") with gr.Row(): image_size = gr.Dropdown(choices=["Small", "Medium", "Large"], value="Small", label="Image Size", elem_classes="compact") file_format = gr.Radio(["pdf", "docx"], label="Format", value="pdf", elem_classes="compact") get_document_btn = gr.Button("Get Document", elem_classes="download-btn compact") get_document_btn.click( generate_document, [input_media, output_text, file_format, font_choice, font_size, line_spacing, alignment, image_size], gr.File(label="Download Document") ) demo.launch(debug=True)