app.py

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
from threading import Thread
from io import BytesIO
import uuid
import tempfile

import gradio as gr
import numpy as np
import torch
from PIL import Image
import supervision as sv


from transformers import (
    Qwen2_5_VLForConditionalGeneration,
    Qwen2VLForConditionalGeneration,
    AutoModelForCausalLM,
    AutoProcessor,
    TextIteratorStreamer,
)

# --- Constants and Model Setup ---
MAX_INPUT_TOKEN_LENGTH = 4096
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print("--- System Information ---")
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)
print("--------------------------")


# --- Model Loading ---

# Load Camel-Doc-OCR-062825
print("Loading Camel-Doc-OCR-062825...")
MODEL_ID_M = "prithivMLmods/Camel-Doc-OCR-062825"
processor_m = AutoProcessor.from_pretrained(MODEL_ID_M, trust_remote_code=True)
model_m = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID_M,
    trust_remote_code=True,
    torch_dtype=torch.float16
).to(device).eval()
print("Camel-Doc-OCR-062825 loaded.")

# MinerU2.5-2509
print("Loading MinerU2.5-2509...")
MODEL_ID_T = "opendatalab/MinerU2.5-2509-1.2B"
processor_t = AutoProcessor.from_pretrained(MODEL_ID_T, trust_remote_code=True)
model_t = Qwen2VLForConditionalGeneration.from_pretrained(
    MODEL_ID_T,
    trust_remote_code=True,
    torch_dtype=torch.float16
).to(device).eval()
print("MinerU2.5-2509 loaded.")

# Load moondream3
print("Loading moondream3-preview...")
MODEL_ID_MD3 = "moondream/moondream3-preview"
model_md3 = AutoModelForCausalLM.from_pretrained(
    MODEL_ID_MD3,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    device_map={"": "cuda"},
)
model_md3.compile()
print("moondream3-preview loaded and compiled.")


# --- Moondream3 Utility Functions ---

def create_annotated_image(image, detection_result, object_name="Object"):
    if not isinstance(detection_result, dict) or "objects" not in detection_result:
        return image
    
    original_width, original_height = image.size
    annotated_image = np.array(image.convert("RGB"))
  
    bboxes = []
    labels = []
    
    for i, obj in enumerate(detection_result["objects"]):
        x_min = int(obj["x_min"] * original_width)
        y_min = int(obj["y_min"] * original_height)
        x_max = int(obj["x_max"] * original_width)
        y_max = int(obj["y_max"] * original_height)
        
        x_min = max(0, min(x_min, original_width))
        y_min = max(0, min(y_min, original_height))
        x_max = max(0, min(x_max, original_width))
        y_max = max(0, min(y_max, original_height))
        
        if x_max > x_min and y_max > y_min:
            bboxes.append([x_min, y_min, x_max, y_max])
            labels.append(f"{object_name} {i+1}")
    
    if not bboxes:
        return image
        
    detections = sv.Detections(
        xyxy=np.array(bboxes, dtype=np.float32),
        class_id=np.arange(len(bboxes))
    )
    
    bounding_box_annotator = sv.BoxAnnotator(
        thickness=3,
        color_lookup=sv.ColorLookup.INDEX
    )
    label_annotator = sv.LabelAnnotator(
        text_thickness=2,
        text_scale=0.6,
        color_lookup=sv.ColorLookup.INDEX
    )
    
    annotated_image = bounding_box_annotator.annotate(
        scene=annotated_image, detections=detections
    )
    annotated_image = label_annotator.annotate(
        scene=annotated_image, detections=detections, labels=labels
    )
    
    return Image.fromarray(annotated_image)

def create_point_annotated_image(image, point_result):
    if not isinstance(point_result, dict) or "points" not in point_result:
        return image
    
    original_width, original_height = image.size
    annotated_image = np.array(image.convert("RGB"))
    
    points = []
    for point in point_result["points"]:
        x = int(point["x"] * original_width)
        y = int(point["y"] * original_height)
        points.append([x, y])
    
    if points:
        points_array = np.array(points).reshape(1, -1, 2)
        key_points = sv.KeyPoints(xy=points_array)
        vertex_annotator = sv.VertexAnnotator(radius=8, color=sv.Color.RED)
        annotated_image = vertex_annotator.annotate(
            scene=annotated_image, key_points=key_points
        )
    
    return Image.fromarray(annotated_image)

@spaces.GPU()
def detect_objects_md3(image, prompt, task_type, max_objects):
    STANDARD_SIZE = (1024, 1024)
    if image is None:
        raise gr.Error("Please upload an image.")
    image.thumbnail(STANDARD_SIZE)
    
    t0 = time.perf_counter()
    
    if task_type == "Object Detection":
        settings = {"max_objects": max_objects} if max_objects > 0 else {}
        result = model_md3.detect(image, prompt, settings=settings)
        annotated_image = create_annotated_image(image, result, prompt)
    elif task_type == "Point Detection":
        result = model_md3.point(image, prompt)
        annotated_image = create_point_annotated_image(image, result)
    elif task_type == "Caption":
        result = model_md3.caption(image, length="normal")
        annotated_image = image  
    else:  
        result = model_md3.query(image=image, question=prompt, reasoning=True)
        annotated_image = image  
          
    elapsed_ms = (time.perf_counter() - t0) * 1_000
    
    if isinstance(result, dict):
        if "objects" in result:
          output_text = f"Found {len(result['objects'])} objects:\n"
          for i, obj in enumerate(result['objects'], 1):
              output_text += f"\n{i}. Bounding box: ({obj['x_min']:.3f}, {obj['y_min']:.3f}, {obj['x_max']:.3f}, {obj['y_max']:.3f})"
        elif "points" in result:
            output_text = f"Found {len(result['points'])} points:\n"
            for i, point in enumerate(result['points'], 1):
                output_text += f"\n{i}. Point: ({point['x']:.3f}, {point['y']:.3f})"
        elif "caption" in result:
            output_text = result['caption']
        elif "answer" in result:
            output_text = f"Reasoning: {result.get('reasoning', 'N/A')}\n\nAnswer: {result['answer']}"
        else:
            output_text = json.dumps(result, indent=2)
    else:
        output_text = str(result)
    
    timing_text = f"Inference time: {elapsed_ms:.0f} ms"
    
    return annotated_image, output_text, timing_text


# --- Core Application Logic (for other models) ---
@spaces.GPU
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
):
    """
    Main generator function for models other than Moondream3.
    """
    if image is None:
        yield "Please upload an image."
        return
    if not prompt_input or not prompt_input.strip():
        yield "Please enter a prompt."
        return

    # Select processor and model based on dropdown choice
    if model_name == "Camel-Doc-OCR-062825 (OCR)":
        processor, model = processor_m, model_m
    elif model_name == "MinerU2.5-2509 (General)":
        processor, model = processor_t, model_t
    else:
        yield "Invalid model selected."
        return
            
    messages = [{"role": "user", "content": [{"type": "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, 
        "max_new_tokens": max_new_tokens,
        "temperature": temperature,
        "top_p": top_p,
        "top_k": top_k,
        "repetition_penalty": repetition_penalty,
        "do_sample": True if temperature > 0 else False
    }

    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    buffer = ""
    for new_text in streamer:
        buffer += new_text
        # Clean up potential model-specific tokens
        buffer = buffer.replace("<|im_end|>", "").replace("</s>", "")
        time.sleep(0.01)
        yield 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; }
    """
    with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
        gr.Markdown("# Multimodal VLM v1.0 ⚡")
        gr.Markdown("Explore the capabilities of various Vision Language Models for tasks like OCR, VQA, and Object Detection.")

        with gr.Tabs():
            # --- TAB 1: Document and General VLMs ---
            with gr.TabItem("📄 Document & General VLM"):
                with gr.Row():
                    with gr.Column(scale=1):
                        #gr.Markdown("### 1. Configure Inputs")
                        model_choice = gr.Dropdown(
                            choices=["Camel-Doc-OCR-062825 (OCR)", "MinerU2.5-2509 (General)"],
                            label="Select Model", value= "Camel-Doc-OCR-062825 (OCR)"
                        )
                        image_input_doc = gr.Image(label="Upload Image", type="pil", sources=['upload'], height=280)
                        prompt_input_doc = gr.Textbox(label="Query Input", placeholder="e.g., 'Transcribe the text in this document.'")
                        
                        with gr.Accordion("Advanced Generation Settings", open=False):
                            max_new_tokens = gr.Slider(minimum=256, maximum=4096, value=2048, step=128, label="Max New Tokens")
                            temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, step=0.1, value=0.7)
                            top_p = gr.Slider(label="Top-p", minimum=0.1, maximum=1.0, step=0.05, value=0.9)
                            top_k = gr.Slider(label="Top-k", minimum=1, maximum=100, step=1, value=40)
                            repetition_penalty = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.1)

                        process_btn = gr.Button("🚀 Process Image", variant="primary", elem_classes=["process-button"])
                        clear_btn = gr.Button("🗑️ Clear", variant="secondary")

                    with gr.Column(scale=2):
                        #gr.Markdown("### 2. View Output")
                        output_stream = gr.Textbox(label="Model Output", interactive=False, lines=20, show_copy_button=True)
                
                gr.Examples(
                    examples=[
                        ["examples/1.png", "Transcribe this receipt."],
                        ["examples/2.png", "Extract the table from this document as markdown."],
                        ["examples/3.png", "What information is presented in this infographic?"],
                    ],
                    inputs=[image_input_doc, prompt_input_doc]
                )
            
            # --- TAB 2: Moondream3 Lab ---
            with gr.TabItem("🌝 Moondream3"):
                with gr.Row():
                    with gr.Column(scale=1):
                        md3_image_input = gr.Image(label="Upload an image", type="pil", height=400)
                        md3_task_type = gr.Radio(
                            choices=["Object Detection", "Point Detection", "Caption", "Visual Question Answering"],
                            label="Task Type", value="Object Detection"
                        )
                        md3_prompt_input = gr.Textbox(
                            label="Prompt (object to detect/question to ask)",
                            placeholder="e.g., 'car', 'person', 'What's in this image?'", value="objects"
                        )
                        md3_max_objects = gr.Number(
                            label="Max Objects (for Object Detection only)",
                            value=10, minimum=1, maximum=50, step=1, visible=True
                        )
                        md3_generate_btn = gr.Button(value="✨ Generate", variant="primary")
                    with gr.Column(scale=1):
                        md3_output_image = gr.Image(type="pil", label="Result", height=400)
                        md3_output_textbox = gr.Textbox(label="Model Response", lines=10, show_copy_button=True)
                        md3_output_time = gr.Markdown()

                gr.Examples(
                    examples=[
                        ["https://huggingface.co/datasets/merve/vlm_test_images/resolve/main/candy.JPG", "Object Detection", "candy", 5],
                        ["https://huggingface.co/datasets/merve/vlm_test_images/resolve/main/candy.JPG", "Point Detection", "candy", 5],
                        ["https://moondream.ai/images/blog/moondream-3-preview/benchmarks.jpg", "Caption", "", 5],
                        ["https://moondream.ai/images/blog/moondream-3-preview/benchmarks.jpg", "Visual Question Answering", "how well does moondream 3 perform in chartvqa?", 5],
                    ],
                    inputs=[md3_image_input, md3_task_type, md3_prompt_input, md3_max_objects],
                    label="Click an example to populate inputs"
                )

        # --- Event Handlers ---
        
        # Document Tab
        process_btn.click(
            fn=process_document_stream,
            inputs=[model_choice, image_input_doc, prompt_input_doc, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
            outputs=[output_stream]
        )
        clear_btn.click(lambda: (None, "", ""), outputs=[image_input_doc, prompt_input_doc, output_stream])

        # Moondream3 Tab
        def update_max_objects_visibility(task):
            return gr.update(visible=(task == "Object Detection"))
        
        md3_task_type.change(fn=update_max_objects_visibility, inputs=[md3_task_type], outputs=[md3_max_objects])
        
        md3_generate_btn.click(
            fn=detect_objects_md3,
            inputs=[md3_image_input, md3_prompt_input, md3_task_type, md3_max_objects],
            outputs=[md3_output_image, md3_output_textbox, md3_output_time]
        )
        
    return demo

if __name__ == "__main__":
    # Create some example images if they don't exist
    if not os.path.exists("examples"):
        os.makedirs("examples")
    try:
        # Dummy image creation for examples to prevent errors if not present
        Image.new('RGB', (200, 100), color = 'red').save('examples/1.png')
        Image.new('RGB', (200, 100), color = 'green').save('examples/2.png')
        Image.new('RGB', (200, 100), color = 'blue').save('examples/3.png')
    except Exception as e:
        print(f"Could not create dummy example images: {e}")

    demo = create_gradio_interface()
    demo.queue(max_size=20).launch(share=True, show_error=True)