app.py

import subprocess
subprocess.run('pip install flash-attn==2.7.0.post2 --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)

import spaces

import argparse
import os
import re
import logging
from typing import List, Optional, Tuple, Generator
from threading import Thread

import gradio as gr
import PIL.Image
import torch
import numpy as np
from moviepy.editor import VideoFileClip
from transformers import AutoModelForCausalLM, TextIteratorStreamer

logging.getLogger("httpx").setLevel(logging.WARNING)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# --- Global Model Variables ---
model = None
streamer = None
# This should point to the directory containing your SVG file.
CUR_DIR = os.path.dirname(os.path.abspath(__file__))


class MyTextIteratorStreamer(TextIteratorStreamer):
    def manual_end(self):
        """Flushes any remaining cache and prints a newline to stdout."""
        # Flush the cache, if it exists
        if len(self.token_cache) > 0:
            text = self.tokenizer.decode(self.token_cache, **self.decode_kwargs)
            printable_text = text[self.print_len :]
            self.token_cache = []
            self.print_len = 0
        else:
            printable_text = ""

        self.next_tokens_are_prompt = True
        self.on_finalized_text(printable_text, stream_end=True)
    
    def end(self):
        pass

def submit_chat(chatbot, text_input):
    response = ''
    chatbot.append([text_input, response])
    return chatbot, ''


# --- Helper Functions ---
latex_delimiters_set = [
    {
        "left": "\\(",
        "right": "\\)",
        "display": False
    },
    {
        "left": "\\begin{equation}",
        "right": "\\end{equation}",
        "display": True
    },
    {
        "left": "\\begin{align}",
        "right": "\\end{align}",
        "display": True
    },
    {
        "left": "\\begin{alignat}",
        "right": "\\end{alignat}",
        "display": True
    },
    {
        "left": "\\begin{gather}",
        "right": "\\end{gather}",
        "display": True
    },
    {
        "left": "\\begin{CD}",
        "right": "\\end{CD}",
        "display": True
    },
    {
        "left": "\\[",
        "right": "\\]",
        "display": True
    }
]

def load_video_frames(video_path: Optional[str], n_frames: int = 8) -> Optional[List[PIL.Image.Image]]:
    """Extracts a specified number of frames from a video file."""
    if not video_path:
        return None
    try:
        with VideoFileClip(video_path) as clip:
            total_frames = int(clip.fps * clip.duration)
            if total_frames <= 0: return None
            num_to_extract = min(n_frames, total_frames)
            indices = np.linspace(0, total_frames - 1, num_to_extract, dtype=int)
            frames = [PIL.Image.fromarray(clip.get_frame(index / clip.fps)) for index in indices]
        return frames
    except Exception as e:
        print(f"Error processing video {video_path}: {e}")
        return None

def parse_model_output(response_text: str, enable_thinking: bool) -> str:
    """Formats the model output, separating 'thinking' and 'response' parts if enabled."""
    if enable_thinking:
        # Use a more robust regex to handle nested content and variations
        think_match = re.search(r"<think>(.*?)</think>", response_text, re.DOTALL)
        if think_match:
            thinking_content = think_match.group(1).strip()
            # Remove the think block from the original text to get the response
            response_content = re.sub(r"<think>.*?</think>", "", response_text, flags=re.DOTALL).strip()
            return f"**Thinking:**\n```\n{thinking_content}\n```\n\n**Response:**\n{response_content}"
        else:
            return response_text # No think tag found, return as is
    else:
        # If thinking is disabled, strip the tags just in case the model still generates them
        return re.sub(r"<think>.*?</think>", "", response_text, flags=re.DOTALL).strip()


# --- MODIFIED Core Inference Logic (Now with Streaming) ---
@spaces.GPU
def run_inference(
    chatbot: List,
    image_input: Optional[PIL.Image.Image],
    video_input: Optional[str],
    do_sample: bool,
    max_new_tokens: int,
    enable_thinking: bool,
    enable_thinking_budget: bool, # NEWLY ADDED
    thinking_budget: int,         # NEWLY ADDED
):
    """
    Runs a single turn of inference and yields the output stream for a gr.Chatbot.
    This function is now a generator.
    """
    prompt = chatbot[-1][0]
    if (not image_input and not video_input and not prompt) or not prompt:
        gr.Warning("A text prompt is required for generation.")
        chatbot.pop(-1)
        # MODIFICATION: Yield the current state and return to avoid errors
        yield chatbot
        return

    content = []
    if image_input:
        content.append({"type": "image", "image": image_input})
    if video_input:
        frames = load_video_frames(video_input)
        if frames:
            content.append({"type": "video", "video": frames})
        else:
            gr.Warning("Failed to process the video file.")
            chatbot.pop(-1)
            yield chatbot
            return

    content.append({"type": "text", "text": prompt})

    messages = [{"role": "user", "content": content}]
    logger.info(messages)

    try:
        if video_input:
            input_ids, pixel_values, grid_thws = model.preprocess_inputs(messages=messages, add_generation_prompt=True, enable_thinking=enable_thinking, max_pixels=896*896)
        else:
            input_ids, pixel_values, grid_thws = model.preprocess_inputs(messages=messages, add_generation_prompt=True, enable_thinking=enable_thinking)
    except Exception as e:
        gr.Warning(f"Error during input preprocessing: {e}")
        chatbot.pop(-1)
        yield chatbot
        return

    input_ids = input_ids.to(model.device)
    if pixel_values is not None:
        pixel_values = pixel_values.to(model.device, dtype=torch.bfloat16)
    if grid_thws is not None:
        grid_thws = grid_thws.to(model.device)

    gen_kwargs = {
        "max_new_tokens": max_new_tokens,
        "do_sample": do_sample,
        "eos_token_id": model.text_tokenizer.eos_token_id,
        "pad_token_id": model.text_tokenizer.pad_token_id,
        "streamer": streamer,
        "use_cache": True,
        "enable_thinking": enable_thinking,
        "enable_thinking_budget": enable_thinking_budget,
        "thinking_budget": thinking_budget
    }

    with torch.inference_mode():
        thread = Thread(target=model.generate, kwargs={
            "inputs": input_ids,
            "pixel_values": pixel_values,
            "grid_thws": grid_thws,
            **gen_kwargs
        })
        thread.start()

        # MODIFICATION: Stream output token by token
        response_text = ""
        for new_text in streamer:
            response_text += new_text
            # Append only the new text chunk to the last response
            chatbot[-1][1] = response_text
            yield chatbot # Yield the updated history

        thread.join()

        # MODIFICATION: Format the final response once generation is complete
        formatted_response = parse_model_output(response_text, enable_thinking)
        chatbot[-1][1] = formatted_response
        yield chatbot # Yield the final, formatted response

        logger.info("\n[OVIS_CONV_START]")
        [print(f'Q{i}:\n {request}\nA{i}:\n {answer}\n') for i, (request, answer) in enumerate(chatbot, 1)]
        logger.info("[OVIS_CONV_END]")


# --- UI Helper Functions ---
def toggle_media_input(choice: str) -> Tuple:
    """Switches visibility between Image/Video inputs and their corresponding examples."""
    if choice == "Image":
        return gr.update(visible=True, value=None), gr.update(visible=False, value=None), gr.update(visible=True), gr.update(visible=False)
    else:  # Video
        return gr.update(visible=False, value=None), gr.update(visible=True, value=None), gr.update(visible=False), gr.update(visible=True)


# --- Build Gradio Application ---
def build_demo(model_path: str):
    """Builds the Gradio user interface for the model."""
    global model, streamer
    device = "cuda"
    print(f"Loading model {model_path} onto device {device}...")

    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True
    ).to(device).eval()

    text_tokenizer = model.text_tokenizer
    streamer = MyTextIteratorStreamer(text_tokenizer, skip_prompt=True, skip_special_tokens=True)

    print("Model loaded successfully.")

    model_name_display = model_path.split('/')[-1]

    logo_html = ""
    logo_svg_path = os.path.join(CUR_DIR, "resource", "logo.svg")
    if os.path.exists(logo_svg_path):
        with open(logo_svg_path, "r", encoding="utf-8") as svg_file:
            svg_content = svg_file.read()
        font_size = "2.5em"
        svg_content_styled = re.sub(r'(<svg[^>]*)(>)', rf'\1 height="{font_size}" style="vertical-align: middle; display: inline-block;"\2', svg_content)
        logo_html = f'<span style="display: inline-block; vertical-align: middle;">{svg_content_styled}</span>'
    else:
        logo_html = '<span style="font-weight: bold; font-size: 2.5em; display: inline-block; vertical-align: middle;">Ovis</span>'
        print(f"Warning: Logo file not found at {logo_svg_path}. Using text fallback.")

    html_header = f"""
    <p align="center" style="font-size: 2.5em; line-height: 1;">
        {logo_html}
        <span style="display: inline-block; vertical-align: middle;">{model_name_display}</span>
    </p>
    <center><font size=3><b>Ovis</b> has been open-sourced on <a href='https://huggingface.co/{model_path}'>😊 Huggingface</a> and <a href='https://github.com/AIDC-AI/Ovis'>🌟 GitHub</a>. If you find Ovis useful, a like❤️ or a star🌟 would be appreciated.</font></center>
    """

    # --- START: Slider synchronization logic functions ---
    def adjust_max_tokens(thinking_budget_val: int, max_new_tokens_val: int) -> gr.Slider:
        """Adjusts max_new_tokens to be at least thinking_budget + 128."""
        new_max_tokens = max(max_new_tokens_val, thinking_budget_val + 128)
        return gr.update(value=new_max_tokens)

    def adjust_thinking_budget(max_new_tokens_val: int, thinking_budget_val: int) -> gr.Slider:
        """Adjusts thinking_budget to be at most max_new_tokens - 128."""
        new_thinking_budget = min(thinking_budget_val, max_new_tokens_val - 128)
        return gr.update(value=new_thinking_budget)
    # --- END: Slider synchronization logic functions ---

    prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your text here and press ENTER", lines=1, container=False)
    with gr.Blocks(theme=gr.themes.Ocean()) as demo:
        gr.HTML(html_header)
        gr.Markdown("Note: The Thinking Budget mechanism is enabled only when `Deep Thinking` and `Thinking Budget` are both checked. Could tune down `Thinking Budget` for faster generation in `Deep Thinking` mode.")

        with gr.Row():
            with gr.Column(scale=4):
                input_type_radio = gr.Radio(choices=["Image", "Video"], value="Image", label="Select Input Type")
                image_input = gr.Image(label="Image Input", type="pil", visible=True)
                video_input = gr.Video(label="Video Input", visible=False)

                with gr.Accordion("Generation Settings", open=True):
                    do_sample = gr.Checkbox(label="Enable Sampling (Do Sample)", value=True)
                    enable_thinking = gr.Checkbox(label="Enable Deep Thinking", value=True)
                    enable_thinking_budget = gr.Checkbox(label="Enable Thinking Budget", value=True)
                    max_new_tokens = gr.Slider(minimum=256, maximum=4096, value=2048, step=32, label="Max New Tokens")
                    thinking_budget = gr.Slider(minimum=128, maximum=3968, value=1024, step=32, label="Thinking Budget")

                with gr.Column(visible=True) as image_examples_col:
                    gr.Examples(
                        examples=[
                            [os.path.join(CUR_DIR, "examples", "ovis2_math0.jpg"), "Each face of the polyhedron shown is either a triangle or a square. Each square borders 4 triangles, and each triangle borders 3 squares. The polyhedron has 6 squares. How many triangles does it have?\n\nEnd your response with 'Final answer: '."],
                            [os.path.join(CUR_DIR, "examples", "ovis2_math1.jpg"), "A large square touches another two squares, as shown in the picture. The numbers inside the smaller squares indicate their areas. What is the area of the largest square?\n\nEnd your response with 'Final answer: '."],
                            [os.path.join(CUR_DIR, "examples", "ovis2_figure0.png"), "Explain this model."],
                            # [os.path.join(CUR_DIR, "examples", "ovis2_figure1.png"), "Organize the notes about GRPO in the figure."],
                            [os.path.join(CUR_DIR, "examples", "ovis2_multi0.jpg"), "Posso avere un frappuccino e un caffè americano di taglia M? Quanto costa in totale?"],
                        ],
                        inputs=[image_input, prompt_input]
                    )
                with gr.Column(visible=False) as video_examples_col:
                     gr.Examples(examples=[[os.path.join(CUR_DIR, "examples", "video_demo.mp4"), "Describe the video."]],
                         inputs=[video_input, prompt_input])

            with gr.Column(scale=7):
                chatbot = gr.Chatbot(label="Ovis", height=600, show_copy_button=True, layout="panel", latex_delimiters=latex_delimiters_set)
                prompt_input.render()
                with gr.Row():
                    generate_btn = gr.Button("Send", variant="primary")
                    clear_btn = gr.Button("Clear", variant="secondary")

        # --- START: Event Handlers for UI Elements ---

        input_type_radio.change(
            fn=toggle_media_input,
            inputs=input_type_radio,
            outputs=[image_input, video_input, image_examples_col, video_examples_col]
        )

        # Event handlers for coupled sliders
        thinking_budget.release(
            fn=adjust_max_tokens,
            inputs=[thinking_budget, max_new_tokens],
            outputs=[max_new_tokens]
        )
        max_new_tokens.release(
            fn=adjust_thinking_budget,
            inputs=[max_new_tokens, thinking_budget],
            outputs=[thinking_budget]
        )

        # MODIFICATION: Update run_inputs to include new controls
        run_inputs = [chatbot, image_input, video_input, do_sample, max_new_tokens, enable_thinking, enable_thinking_budget, thinking_budget]

        generat_click_event = generate_btn.click(submit_chat, [chatbot, prompt_input], [chatbot, prompt_input]).then(run_inference, run_inputs, chatbot)
        submit_event = prompt_input.submit(submit_chat, [chatbot, prompt_input], [chatbot, prompt_input]).then(run_inference, run_inputs, chatbot)

        # MODIFICATION: Update clear button to reset new controls
        # clear_btn.click(
        #     fn=lambda: ([], None, None, "", "Image", True, 2048, True, True, 1024),
        #     outputs=[chatbot, image_input, video_input, prompt_input, input_type_radio, do_sample, max_new_tokens, enable_thinking, enable_thinking_budget, thinking_budget]
        # ).then(
        #      fn=toggle_media_input,
        #      inputs=input_type_radio,
        #      outputs=[image_input, video_input, image_examples_col, video_examples_col]
        # )
        clear_btn.click(
            fn=lambda: (list(), None, None, ""),
            outputs=[chatbot, image_input, video_input, prompt_input]
        )
        # --- END: Event Handlers for UI Elements ---

    return demo


# --- Main Execution Block ---
# def parse_args():
#     parser = argparse.ArgumentParser(description="Gradio interface for a single Multimodal Large Language Model.")
#     parser.add_argument("--model-path", type=str, default='AIDC-AI/Ovis2.5-9B', help="Path to the model checkpoint on Hugging Face Hub or local directory.")
#     parser.add_argument("--gpu", type=int, default=0, help="GPU index to run the model on.")
#     parser.add_argument("--port", type=int, default=7860, help="Port to run the Gradio server on.")
#     parser.add_argument("--server-name", type=str, default="0.0.0.0", help="Server name for the Gradio app.")
#     return parser.parse_args()


# if __name__ == "__main__":
#     args = parse_args()
model_path = 'AIDC-AI/Ovis2.5-2B'
demo = build_demo(model_path=model_path)
# demo = build_demo(model_path=args.model_path)
# demo.launch(server_name=args.server_name, server_port=args.port, share=False, ssl_verify=False, show_error=True)
demo.queue().launch()