Update app.py

app.py

@@ -1,28 +1,57 @@
+import spaces
+import json
+import math
 import os
+import traceback
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Tuple
+import re
 import time
-import threading
+from threading import Thread
+from io import BytesIO
+import uuid
+import tempfile
+
 import gradio as gr
-import spaces
+import requests
 import torch
-import numpy as np
 from PIL import Image
-import cv2
+import fitz
+import numpy as np
+
+
 from transformers import (
     Qwen2_5_VLForConditionalGeneration,
     Qwen2VLForConditionalGeneration,
-    #Glm4vForConditionalGeneration,
+    AutoModelForCausalLM,
     AutoProcessor,
     TextIteratorStreamer,
+    AutoTokenizer,
 )
-from qwen_vl_utils import process_vision_info
 
-# Constants for text generation
-MAX_MAX_NEW_TOKENS = 16384
-DEFAULT_MAX_NEW_TOKENS = 8192
-MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
+from transformers.image_utils import load_image as hf_load_image
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+from reportlab.lib.pagesizes import A4
+from reportlab.lib.styles import getSampleStyleSheet
+from reportlab.platypus import SimpleDocTemplate, Image as RLImage, Paragraph, Spacer
+from reportlab.lib.units import inch
 
+# --- Constants and Model Setup ---
+MAX_INPUT_TOKEN_LENGTH = 4096
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))
+print("torch.__version__ =", torch.__version__)
+print("torch.version.cuda =", torch.version.cuda)
+print("cuda available:", torch.cuda.is_available())
+print("cuda device count:", torch.cuda.device_count())
+if torch.cuda.is_available():
+    print("current device:", torch.cuda.current_device())
+    print("device name:", torch.cuda.get_device_name(torch.cuda.current_device()))
+
+print("Using device:", device)
+
+# --- Model Loading ---
 # Load Camel-Doc-OCR-062825
 MODEL_ID_M = "prithivMLmods/Camel-Doc-OCR-062825"
 processor_m = AutoProcessor.from_pretrained(MODEL_ID_M, trust_remote_code=True)
@@ -50,252 +79,249 @@ model_s = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     torch_dtype=torch.float16
 ).to(device).eval()
 
-# Load ViLaSR
-MODEL_ID_Y = "inclusionAI/ViLaSR"
-processor_y = AutoProcessor.from_pretrained(MODEL_ID_Y, trust_remote_code=True)
-model_y = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    MODEL_ID_Y,
+# moondream3
+MODEL_ID_MD3 = "moondream/moondream3-preview"
+model_md3 = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID_MD3,
     trust_remote_code=True,
-    torch_dtype=torch.float16
-).to(device).eval()
+    torch_dtype=torch.bfloat16,
+    device_map={"": "cuda"},
+)
+tokenizer_md3 = AutoTokenizer.from_pretrained(MODEL_ID_MD3)
 
-def downsample_video(video_path):
-    """
-    Downsample a video to evenly spaced frames, returning each as a PIL image with its timestamp.
-    """
-    vidcap = cv2.VideoCapture(video_path)
-    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
-    fps = vidcap.get(cv2.CAP_PROP_FPS)
-    frames = []
-    frame_indices = np.linspace(0, total_frames - 1, 10, dtype=int)
-    for i in frame_indices:
-        vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
-        success, image = vidcap.read()
-        if success:
-            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-            pil_image = Image.fromarray(image)
-            timestamp = round(i / fps, 2)
-            frames.append((pil_image, timestamp))
-    vidcap.release()
-    return frames
 
-@spaces.GPU
-def generate_image(model_name: str, text: str, image: Image.Image,
-                  max_new_tokens: int = 1024,
-                  temperature: float = 0.6,
-                  top_p: float = 0.9,
-                  top_k: int = 50,
-                  repetition_penalty: float = 1.2):
+# --- PDF Generation and Preview Utility Function ---
+def generate_and_preview_pdf(image: Image.Image, text_content: str, font_size: int, line_spacing: float, alignment: str, image_size: str):
     """
-    Generate responses using the selected model for image input.
+    Generates a PDF, saves it, and then creates image previews of its pages.
+    Returns the path to the PDF and a list of paths to the preview images.
     """
-    if model_name == "Camel-Doc-OCR-062825":
-        processor = processor_m
-        model = model_m
-    elif model_name == "MinerU2.5-2509":
-        processor = processor_t
-        model = model_t
-    elif model_name == "Video-MTR":
-        processor = processor_s
-        model = model_s
-    elif model_name == "ViLaSR-7B":
-        processor = processor_y
-        model = model_y
-    else:
-        yield "Invalid model selected.", "Invalid model selected."
-        return
+    if image is None or not text_content or not text_content.strip():
+        raise gr.Error("Cannot generate PDF. Image or text content is missing.")
 
-    if image is None:
-        yield "Please upload an image.", "Please upload an image."
-        return
+    # --- 1. Generate the PDF ---
+    temp_dir = tempfile.gettempdir()
+    pdf_filename = os.path.join(temp_dir, f"output_{uuid.uuid4()}.pdf")
+    doc = SimpleDocTemplate(
+        pdf_filename,
+        pagesize=A4,
+        rightMargin=inch, leftMargin=inch,
+        topMargin=inch, bottomMargin=inch
+    )
+    styles = getSampleStyleSheet()
+    style_normal = styles["Normal"]
+    style_normal.fontSize = int(font_size)
+    style_normal.leading = int(font_size) * line_spacing
+    style_normal.alignment = {"Left": 0, "Center": 1, "Right": 2, "Justified": 4}[alignment]
 
-    messages = [{
-        "role": "user",
-        "content": [
-            {"type": "image", "image": image},
-            {"type": "text", "text": text},
-        ]
-    }]
-    prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    inputs = processor(
-        text=[prompt_full],
-        images=[image],
-        return_tensors="pt",
-        padding=True,
-        truncation=False,
-        max_length=MAX_INPUT_TOKEN_LENGTH
-    ).to(device)
-    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
-    generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens}
-    thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-    buffer = ""
-    for new_text in streamer:
-        buffer += new_text
-        time.sleep(0.01)
-        yield buffer, buffer
+    story = []
 
+    img_buffer = BytesIO()
+    image.save(img_buffer, format='PNG')
+    img_buffer.seek(0)
+    
+    page_width, _ = A4
+    available_width = page_width - 2 * inch
+    image_widths = {
+        "Small": available_width * 0.3,
+        "Medium": available_width * 0.6,
+        "Large": available_width * 0.9,
+    }
+    img_width = image_widths[image_size]
+    img = RLImage(img_buffer, width=img_width, height=image.height * (img_width / image.width))
+    story.append(img)
+    story.append(Spacer(1, 12))
+
+    cleaned_text = re.sub(r'#+\s*', '', text_content).replace("*", "")
+    text_paragraphs = cleaned_text.split('\n')
+    
+    for para in text_paragraphs:
+        if para.strip():
+            story.append(Paragraph(para, style_normal))
+
+    doc.build(story)
+
+    # --- 2. Render PDF pages as images for preview ---
+    preview_images = []
+    try:
+        pdf_doc = fitz.open(pdf_filename)
+        for page_num in range(len(pdf_doc)):
+            page = pdf_doc.load_page(page_num)
+            pix = page.get_pixmap(dpi=150)
+            preview_img_path = os.path.join(temp_dir, f"preview_{uuid.uuid4()}_p{page_num}.png")
+            pix.save(preview_img_path)
+            preview_images.append(preview_img_path)
+        pdf_doc.close()
+    except Exception as e:
+        print(f"Error generating PDF preview: {e}")
+        
+    return pdf_filename, preview_images
+
+
+# --- Core Application Logic ---
 @spaces.GPU
-def generate_video(model_name: str, text: str, video_path: str,
-                  max_new_tokens: int = 1024,
-                  temperature: float = 0.6,
-                  top_p: float = 0.9,
-                  top_k: int = 50,
-                  repetition_penalty: float = 1.2):
+def process_document_stream(
+    model_name: str, 
+    image: Image.Image, 
+    prompt_input: str, 
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    repetition_penalty: float
+):
     """
-    Generate responses using the selected model for video input.
+    Main generator function that handles model inference tasks with advanced generation parameters.
     """
-    if model_name == "Camel-Doc-OCR-062825":
-        processor = processor_m
-        model = model_m
-    elif model_name == "MinerU2.5-2509":
-        processor = processor_t
-        model = model_t
-    elif model_name == "Video-MTR":
-        processor = processor_s
-        model = model_s
-    elif model_name == "ViLaSR-7B":
-        processor = processor_y
-        model = model_y
-    else:
-        yield "Invalid model selected.", "Invalid model selected."
+    if image is None:
+        yield "Please upload an image.", ""
         return
-
-    if video_path is None:
-        yield "Please upload a video.", "Please upload a video."
+    if not prompt_input or not prompt_input.strip():
+        yield "Please enter a prompt.", ""
         return
 
-    frames = downsample_video(video_path)
-    messages = [
-        {"role": "system", "content": [{"type": "text", "text": "You are a helpful assistant."}]},
-        {"role": "user", "content": [{"type": "text", "text": text}]}
-    ]
-    for frame in frames:
-        image, timestamp = frame
-        messages[1]["content"].append({"type": "text", "text": f"Frame {timestamp}:"})
-        messages[1]["content"].append({"type": "image", "image": image})
-    inputs = processor.apply_chat_template(
-        messages,
-        tokenize=True,
-        add_generation_prompt=True,
-        return_dict=True,
-        return_tensors="pt",
-        truncation=False,
-        max_length=MAX_INPUT_TOKEN_LENGTH
-    ).to(device)
+    # --- Special Handling for Moondream3 ---
+    if model_name == "Moondream3":
+        prompt_full = f"<image>\n{prompt_input}"
+        answer = model_md3.answer_question(
+            model_md3.encode_image(image),
+            prompt_full,
+            tokenizer=tokenizer_md3
+        )
+        yield answer, answer
+        return
+    
+    processor = None
+    model = None
+    
+    # --- Generic Handling for all other models ---
+    if model_name == "Camel-Doc-OCR-062825": processor, model = processor_m, model_m
+    elif model_name == "MinerU2.5-2509-1.2B": processor, model = processor_t, model_t
+    elif model_name == "Video-MTR": processor, model = processor_s, model_s
+    else:
+        yield "Invalid model selected.", ""
+        return
+            
+    messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": prompt_input}]}]
+    prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = processor(text=[prompt_full], images=[image], return_tensors="pt", padding=True, truncation=True, max_length=MAX_INPUT_TOKEN_LENGTH).to(device)
+    
     streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
+    
     generation_kwargs = {
-        **inputs,
-        "streamer": streamer,
+        **inputs, 
+        "streamer": streamer, 
         "max_new_tokens": max_new_tokens,
-        "do_sample": True,
         "temperature": temperature,
         "top_p": top_p,
         "top_k": top_k,
         "repetition_penalty": repetition_penalty,
+        "do_sample": True
     }
-    thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
+
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
+
     buffer = ""
     for new_text in streamer:
         buffer += new_text
         buffer = buffer.replace("<|im_end|>", "")
         time.sleep(0.01)
-        yield buffer, buffer
-
-# Define examples for image and video inference
-image_examples = [
-    ["convert this page to doc [text] precisely for markdown.", "images/1.png"],
-    ["explain the movie shot in detail.", "images/5.jpg"],
-    ["convert this page to doc [table] precisely for markdown.", "images/2.png"],
-    ["explain the movie shot in detail.", "images/3.png"],
-    ["fill the correct numbers.", "images/4.png"]
-]
-
-video_examples = [
-    ["explain the video in detail.", "videos/b.mp4"],
-    ["explain the ad video in detail.", "videos/a.mp4"]
-]
-
-# Updated CSS with model choice highlighting
-css = """
-.submit-btn {
-    background-color: #2980b9 !important;
-    color: white !important;
-}
-.submit-btn:hover {
-    background-color: #3498db !important;
-}
-.canvas-output {
-    border: 2px solid #4682B4;
-    border-radius: 10px;
-    padding: 20px;
-}
-"""
-
-# Create the Gradio Interface
-with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
-    gr.Markdown("# **[Multimodal VLM v1.0](https://huggingface.co/collections/prithivMLmods/multimodal-implementations-67c9982ea04b39f0608badb0)**")
-    with gr.Row():
-        with gr.Column():
-            with gr.Tabs():
-                with gr.TabItem("Image Inference"):
-                    image_query = gr.Textbox(label="Query Input", placeholder="✦︎ Enter your query here...")
-                    image_upload = gr.Image(type="pil", label="Image", height=290)
-                    image_submit = gr.Button("Submit", elem_classes="submit-btn")
-                    gr.Examples(
-                        examples=image_examples,
-                        inputs=[image_query, image_upload]
-                    )
-                with gr.TabItem("Video Inference"):
-                    video_query = gr.Textbox(label="Query Input", placeholder="✦︎ Enter your query here...")
-                    video_upload = gr.Video(label="Video", height=290)
-                    video_submit = gr.Button("Submit", elem_classes="submit-btn")
-                    gr.Examples(
-                        examples=video_examples,
-                        inputs=[video_query, video_upload]
-                    )
-            with gr.Accordion("Advanced options", open=False):
-                max_new_tokens = gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS)
-                temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.6)
-                top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9)
-                top_k = gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50)
-                repetition_penalty = gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2)
-        with gr.Column():
-            with gr.Column(elem_classes="canvas-output"):
-                gr.Markdown("## Output")
-                output = gr.Textbox(label="Raw Output Stream", interactive=False, lines=5)
-                with gr.Accordion("(Result.md)", open=False):
-                    markdown_output = gr.Markdown(label="(Result.md)")
-            model_choice = gr.Radio(
-                choices=["Camel-Doc-OCR-062825", "MinerU2.5-2509", "Video-MTR", "ViLaSR-7B"],
-                label="Select Model",
-                value="Camel-Doc-OCR-062825"
-            )
-            gr.Markdown("**Model Info 💻** | [Report Bug](https://huggingface.co/spaces/prithivMLmods/Multimodal-VLM-v1.0/discussions)")
-
-            gr.Markdown("> [Camel-Doc-OCR-062825](https://huggingface.co/prithivMLmods/Camel-Doc-OCR-062825) is a Qwen2.5-VL-7B-Instruct finetune, highly optimized for document retrieval, structured extraction, analysis, and direct Markdown generation from images and PDFs.")
-            gr.Markdown("> [MinerU2.5-2509](https://huggingface.co/opendatalab/MinerU2.5-2509-1.2B) is a 1.2B-parameter vision-language model for document parsing that achieves state-of-the-art accuracy with high computational efficiency by adopting a two-stage parsing strategy.")
-            gr.Markdown("> [ViLaSR-7B](https://huggingface.co/inclusionAI/ViLaSR) focuses on reinforcing spatial reasoning in visual-language tasks by combining interwoven thinking with visual drawing, making it especially suited for spatial reasoning and complex tip-based queries.")
-            gr.Markdown("> [Video-MTR](https://huggingface.co/Phoebe13/Video-MTR) introduces reinforced multi-turn reasoning for long-form video understanding, enabling iterative key segment selection and deeper question comprehension.")
-       
-            gr.Markdown("> ✋ ViLaSR-7B - demo only supports text-only reasoning, which doesn't reflect the full behavior of the model and may underrepresent its capabilities.")            
-            gr.Markdown("> ⚠️ Note: Models in this space may not perform well on video inference tasks.")
-    # Define the submit button actions
-    image_submit.click(fn=generate_image,
-                       inputs=[
-                           model_choice, image_query, image_upload,
-                           max_new_tokens, temperature, top_p, top_k,
-                           repetition_penalty
-                       ],
-                       outputs=[output, markdown_output])
-    video_submit.click(fn=generate_video,
-                       inputs=[
-                           model_choice, video_query, video_upload,
-                           max_new_tokens, temperature, top_p, top_k,
-                           repetition_penalty
-                       ],
-                       outputs=[output, markdown_output])
+        yield buffer , buffer
+
+    yield buffer, buffer
+
+
+# --- Gradio UI Definition ---
+def create_gradio_interface():
+    """Builds and returns the Gradio web interface."""
+    css = """
+    .main-container { max-width: 1400px; margin: 0 auto; }
+    .process-button { border: none !important; color: white !important; font-weight: bold !important; background-color: blue !important;}
+    .process-button:hover { background-color: darkblue !important; transform: translateY(-2px) !important; box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important; }
+    #gallery { min-height: 400px; }
+    """
+    with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
+        gr.HTML("""
+        <div class="title" style="text-align: center">
+            <h1>Multimodal VLM v1.0</h1>
+            <p style="font-size: 1.1em; color: #6b7280; margin-bottom: 0.6em;">
+                Multimodal VLM for Image Content Extraction and Understanding
+            </p>
+        </div>
+        """)
+
+        with gr.Row():
+            # Left Column (Inputs)
+            with gr.Column(scale=1):
+                model_choice = gr.Dropdown(
+                    choices=["Moondream3", "Camel-Doc-OCR-062825", "MinerU2.5-2509-1.2B", "Video-MTR"],
+                    label="Select Model", value= "Moondream3"
+                )
+                
+                prompt_input = gr.Textbox(label="Query Input", placeholder="✦︎ Enter the prompt")
+                image_input = gr.Image(label="Upload Image", type="pil", sources=['upload'])
+                
+                with gr.Accordion("Advanced Settings (PDF)", open=False):
+                    max_new_tokens = gr.Slider(minimum=512, maximum=8192, value=2048, step=256, label="Max New Tokens")
+                    temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.6)
+                    top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9)
+                    top_k = gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50)
+                    repetition_penalty = gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2)
+                    
+                    gr.Markdown("### PDF Export Settings")
+                    font_size = gr.Dropdown(choices=["8", "10", "12", "14", "16", "18"], value="12", label="Font Size")
+                    line_spacing = gr.Dropdown(choices=[1.0, 1.15, 1.5, 2.0], value=1.15, label="Line Spacing")
+                    alignment = gr.Dropdown(choices=["Left", "Center", "Right", "Justified"], value="Justified", label="Text Alignment")
+                    image_size = gr.Dropdown(choices=["Small", "Medium", "Large"], value="Medium", label="Image Size in PDF")
+
+                process_btn = gr.Button("🚀 Process Image", variant="primary", elem_classes=["process-button"], size="lg")
+                clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
+
+            # Right Column (Outputs)
+            with gr.Column(scale=2):
+                with gr.Tabs() as tabs:
+                    with gr.Tab("📝 Extracted Content"):
+                        raw_output_stream = gr.Textbox(label="Raw Model Output Stream", interactive=False, lines=15, show_copy_button=True)
+                        with gr.Row():
+                            examples = gr.Examples(
+                                examples=["examples/1.png", "examples/2.png", "examples/3.png",
+                                          "examples/4.png", "examples/5.png"],
+                                inputs=image_input, label="Examples"
+                            )
+                        gr.Markdown("[Report-Bug💻](https://huggingface.co/spaces/prithivMLmods/Tiny-VLMs-Lab/discussions) | [prithivMLmods🤗](https://huggingface.co/prithivMLmods)")
+                    
+                    with gr.Tab("📰 README.md"):
+                        with gr.Accordion("(Result.md)", open=True): 
+                            markdown_output = gr.Markdown()
+
+                    with gr.Tab("📋 PDF Preview"):
+                        generate_pdf_btn = gr.Button("📄 Generate PDF & Render", variant="primary")
+                        pdf_output_file = gr.File(label="Download Generated PDF", interactive=False)
+                        pdf_preview_gallery = gr.Gallery(label="PDF Page Preview", show_label=True, elem_id="gallery", columns=2, object_fit="contain", height="auto")
+
+        # Event Handlers
+        def clear_all_outputs():
+            return None, "", "Raw output will appear here.", "", None, None
+
+        process_btn.click(
+            fn=process_document_stream,
+            inputs=[model_choice, image_input, prompt_input, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
+            outputs=[raw_output_stream, markdown_output]
+        )
+        
+        generate_pdf_btn.click(
+            fn=generate_and_preview_pdf,
+            inputs=[image_input, raw_output_stream, font_size, line_spacing, alignment, image_size],
+            outputs=[pdf_output_file, pdf_preview_gallery]
+        )
+
+        clear_btn.click(
+            clear_all_outputs,
+            outputs=[image_input, prompt_input, raw_output_stream, markdown_output, pdf_output_file, pdf_preview_gallery]
+        )
+    return demo
 
 if __name__ == "__main__":
-    demo.queue(max_size=40).launch(share=True, mcp_server=True, ssr_mode=False, show_error=True)
+    demo = create_gradio_interface()
+    demo.queue(max_size=50).launch(share=True, mcp_server=True, ssr_mode=False, show_error=True)