app.py

import spaces
import torch
import gradio as gr
import re
import numpy as np
from PIL import Image

from transformers import AutoTokenizer, AutoImageProcessor, Qwen2_5_VLForConditionalGeneration

MODEL_ID = "openbmb/EVisRAG-7B"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True, padding_side="left")
image_processor = AutoImageProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda").eval()

# Token IDs for Qwen2.5-VL image placeholders
IMAGE_PAD_ID = tokenizer.convert_tokens_to_ids("<|image_pad|>")
MERGE_SIZE = image_processor.merge_size  # 2
MERGE_SQ = MERGE_SIZE ** 2  # 4


def evidence_prompt(query: str) -> str:
    """Build the evidence-guided multi-image reasoning prompt used by EVisRAG."""
    return f"""You are an AI Visual QA assistant. I will provide you with a question and several images. Please follow the four steps below:

Step 1: Observe the Images
First, analyze the question and consider what types of images may contain relevant information. Then, examine each image one by one, paying special attention to aspects related to the question. Identify whether each image contains any potentially relevant information.
Wrap your observations within <observe></observe> tags.

Step 2: Record Evidences from Images
After reviewing all images, record the evidence you find for each image within <evidence></evidence> tags.
If you are certain that an image contains no relevant information, record it as: [i]: no relevant information(where i denotes the index of the image).
If an image contains relevant evidence, record it as: [j]: [the evidence you find for the question](where j is the index of the image).

Step 3: Reason Based on the Question and Evidences
Based on the recorded evidences, reason about the answer to the question.
Include your step-by-step reasoning within <reasoning></reasoning> tags.

Step 4: Answer the Question
Provide your final answer based only on the evidences you found in the images.
Wrap your answer within <answer></answer> tags.
Avoid adding unnecessary contents in your final answer, like if the question is a yes/no question, simply answer "yes" or "no".
If none of the images contain sufficient information to answer the question, respond with <answer>insufficient to answer</answer>.

Formatting Requirements:
Use the exact tags <observe>, <evidence>, <reasoning>, and <answer> for structured output.
It is possible that none, one, or several images contain relevant evidence.
If you find no evidence or few evidences, and insufficient to help you answer the question, follow the instruction above for insufficient information.

Question and images are provided below. Please follow the steps as instructed.
Question: {query}
"""


def _parse_sections(text: str):
    """Extract <observe>, <evidence>, <reasoning>, and <answer> sections from the raw output."""
    sections = {"observe": "", "evidence": "", "reasoning": "", "answer": ""}
    for tag in sections:
        m = re.search(rf"<{tag}>(.*?)</{tag}>", text, re.DOTALL)
        if m:
            sections[tag] = m.group(1).strip()
    if not sections["answer"]:
        ans = re.search(r"<answer>(.*)", text, re.DOTALL)
        if ans:
            sections["answer"] = ans.group(1).strip()
    return sections


def _build_chat_text(question: str, num_images: int) -> str:
    """Apply the chat template manually for the Qwen2.5-VL format."""
    prompt = evidence_prompt(question)
    content = [{"type": "text", "text": prompt}]
    for i in range(num_images):
        content.append({"type": "image"})

    messages = [{"role": "user", "content": content}]

    # Use the tokenizer's chat template
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    return text


@spaces.GPU(duration=180)
def answer_question(
    question: str,
    image1=None,
    image2=None,
    image3=None,
    image4=None,
    image5=None,
    max_new_tokens: int = 2048,
    temperature: float = 0.1,
    top_k: int = 1,
    repetition_penalty: float = 1.05,
):
    """Answer a question using multiple evidence images with EVisRAG's evidence-guided reasoning.

    Upload 1-5 retrieved images and ask a question. The model observes each image,
    extracts per-image evidence, reasons over the evidence, and produces a final answer.
    """
    def _is_valid_image(img):
        # Only pass real image objects to the image processor. Guard against
        # non-image values (e.g. ints from gallery selection state / arg-order
        # mismatches) leaking into the images list.
        if img is None:
            return False
        if isinstance(img, Image.Image):
            return True
        if isinstance(img, str):
            return True  # file path
        if isinstance(img, np.ndarray):
            return True
        return False

    images = [img for img in [image1, image2, image3, image4, image5] if _is_valid_image(img)]
    # Normalize any file paths to PIL images so the image processor only ever
    # receives real image objects.
    images = [Image.open(img).convert("RGB") if isinstance(img, str) else img for img in images]
    if not images:
        return "Please upload at least one image.", "", "", ""
    if not question.strip():
        return "Please enter a question.", "", "", ""

    # Build chat text with image placeholders
    chat_text = _build_chat_text(question, len(images))

    # Process images
    image_inputs = image_processor(images=images, return_tensors="pt")
    pixel_values = image_inputs["pixel_values"].to("cuda", dtype=torch.bfloat16)
    image_grid_thw = image_inputs["image_grid_thw"].to("cuda")

    # Tokenize text
    inputs = tokenizer(
        text=[chat_text],
        return_tensors="pt",
        padding=True,
    )
    input_ids = inputs["input_ids"].to("cuda")
    attention_mask = inputs["attention_mask"].to("cuda")

    # Expand image placeholder tokens: the chat template produces one <|image_pad|> per image,
    # but the model expects (image_grid_thw[i].prod() // merge_size**2) copies per image.
    # We need to expand the single token into N copies in the right order.
    new_input_ids = []
    img_idx = 0
    for tok_id in input_ids[0].tolist():
        if tok_id == IMAGE_PAD_ID:
            # Number of image tokens for this image
            grid = image_grid_thw[img_idx]
            num_tokens = int(grid[0] * grid[1] * grid[2]) // MERGE_SQ
            new_input_ids.extend([IMAGE_PAD_ID] * num_tokens)
            img_idx += 1
        else:
            new_input_ids.append(tok_id)

    input_ids = torch.tensor([new_input_ids], dtype=input_ids.dtype, device="cuda")
    attention_mask = torch.ones_like(input_ids)

    with torch.no_grad():
        output_ids = model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            pixel_values=pixel_values,
            image_grid_thw=image_grid_thw,
            max_new_tokens=max_new_tokens,
            do_sample=temperature > 0,
            temperature=max(temperature, 1e-5) if temperature > 0 else 1.0,
            top_k=top_k if top_k > 0 else 1,
            repetition_penalty=repetition_penalty,
        )

    # Decode only the generated part
    input_len = input_ids.shape[1]
    generated_ids = output_ids[:, input_len:]
    raw_output = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

    sections = _parse_sections(raw_output)

    return sections["answer"], sections["observe"], sections["evidence"], sections["reasoning"]


# ---- UI ----

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
/* Restore normal page scrolling */
html, body { height: auto !important; overflow-y: auto !important; }
.gradio-container { height: auto !important; min-height: 0 !important; overflow: visible !important; }
"""

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # EVisRAG-7B: Evidence-Guided Multi-Image VQA
        Ask a question about multiple retrieved images. The model observes each image,
        extracts per-image evidence, reasons over the evidence, and produces a grounded answer.
        [Paper](https://huggingface.co/papers/2510.09733) · [Model](https://huggingface.co/openbmb/EVisRAG-7B) · [Code](https://github.com/OpenBMB/VisRAG)
        """
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            question = gr.Textbox(
                label="Question",
                placeholder="e.g. How many stores did Saint Laurent operate in Western Europe in 2020?",
                scale=4,
            )
            run_btn = gr.Button("Run", variant="primary", scale=1)

        gr.Markdown("### Evidence Images (upload 1-5 images)")
        with gr.Row():
            img1 = gr.Image(type="pil", label="Image 1")
            img2 = gr.Image(type="pil", label="Image 2")
            img3 = gr.Image(type="pil", label="Image 3")
            img4 = gr.Image(type="pil", label="Image 4")
            img5 = gr.Image(type="pil", label="Image 5")

        with gr.Row():
            answer_out = gr.Textbox(label="Answer", interactive=False, scale=1)

        with gr.Accordion("Reasoning details", open=False):
            observe_out = gr.Textbox(label="Observations", interactive=False, lines=6)
            evidence_out = gr.Textbox(label="Evidence", interactive=False, lines=6)
            reasoning_out = gr.Textbox(label="Reasoning", interactive=False, lines=6)

        with gr.Accordion("Advanced settings", open=False):
            max_tokens = gr.Slider(64, 4096, value=2048, step=64, label="Max new tokens")
            temp = gr.Slider(0.0, 2.0, value=0.1, step=0.05, label="Temperature")
            topk = gr.Slider(1, 100, value=1, step=1, label="Top-k (1 = greedy)")
            rep_penalty = gr.Slider(1.0, 2.0, value=1.05, step=0.01, label="Repetition penalty")

        gr.Examples(
            examples=[
                [
                    "How many stores did Saint Laurent operate in Western Europe in 2020?",
                    "example_images/example_0_img0.png",
                    "example_images/example_0_img1.png",
                    "example_images/example_0_img2.png",
                ],
                [
                    "In what year did online sales make up 6.8 percent of retail sales of jewelry, watches and accessories in Germany?",
                    "example_images/example_1_img0.png",
                    "example_images/example_1_img1.png",
                    "example_images/example_1_img2.png",
                ],
                [
                    "Which social media site was invented in 2002?",
                    "example_images/example_infovqa_0_img0.png",
                    "example_images/example_infovqa_0_img1.png",
                    "example_images/example_infovqa_0_img2.png",
                ],
            ],
            inputs=[question, img1, img2, img3],
            outputs=[answer_out, observe_out, evidence_out, reasoning_out],
            fn=answer_question,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        fn=answer_question,
        inputs=[question, img1, img2, img3, img4, img5, max_tokens, temp, topk, rep_penalty],
        outputs=[answer_out, observe_out, evidence_out, reasoning_out],
        api_name="answer_question",
    )

demo.launch(mcp_server=True, ssr_mode=False, theme=gr.themes.Citrus(), css=CSS)