app.py

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)