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 os
import re
import logging
from typing import List, Any
from threading import Thread

import torch
import gradio as gr
from transformers import AutoModelForCausalLM, TextIteratorStreamer

# 모델 및 토크나이저 로딩
model_name = 'AIDC-AI/Ovis2-8B'
use_thread = False

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    multimodal_max_length=8192,
    trust_remote_code=True
).to(device='cuda')

text_tokenizer = model.get_text_tokenizer()
visual_tokenizer = model.get_visual_tokenizer()
streamer = TextIteratorStreamer(text_tokenizer, skip_prompt=True, skip_special_tokens=True)
image_placeholder = '<image>'
cur_dir = os.path.dirname(os.path.abspath(__file__))

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def initialize_gen_kwargs():
    return {
        "max_new_tokens": 1536,
        "do_sample": False,
        "top_p": None,
        "top_k": None,
        "temperature": None,
        "repetition_penalty": 1.05,
        "eos_token_id": model.generation_config.eos_token_id,
        "pad_token_id": text_tokenizer.pad_token_id,
        "use_cache": True
    }

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

@gradio.routes.no_temp_folder()
@gradio.gpu()
def ovis_chat(chatbot: List[List[str]], image_input: Any):
    conversations, model_inputs = prepare_inputs(chatbot, image_input)
    gen_kwargs = initialize_gen_kwargs()

    with torch.inference_mode():
        generate_func = lambda: model.generate(**model_inputs, **gen_kwargs, streamer=streamer)
        
        if use_thread:
            thread = Thread(target=generate_func)
            thread.start()
        else:
            generate_func()

        response = ""
        for new_text in streamer:
            response += new_text
            chatbot[-1][1] = response
            yield chatbot

        if use_thread:
            thread.join()

    log_conversation(chatbot)

def prepare_inputs(chatbot: List[List[str]], image_input: Any):
    conversations = []
    for query, response in chatbot[:-1]:
        conversations.extend([
            {"from": "human", "value": query},
            {"from": "gpt", "value": response}
        ])
    
    last_query = chatbot[-1][0].replace(image_placeholder, '')
    conversations.append({"from": "human", "value": last_query})

    if image_input is not None:
        # 이미지가 포함되면 첫 번째 human 메시지에 이미지 태그 추가
        for conv in conversations:
            if conv["from"] == "human":
                conv["value"] = f'{image_placeholder}\n{conv["value"]}'
                break

    logger.info(conversations)
    
    prompt, input_ids, pixel_values = model.preprocess_inputs(conversations, [image_input], max_partition=16)
    attention_mask = torch.ne(input_ids, text_tokenizer.pad_token_id)
    
    model_inputs = {
        "inputs": input_ids.unsqueeze(0).to(device=model.device),
        "attention_mask": attention_mask.unsqueeze(0).to(device=model.device),
        "pixel_values": [pixel_values.to(dtype=visual_tokenizer.dtype, device=visual_tokenizer.device)] if image_input is not None else [None]
    }
    
    return conversations, model_inputs

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

def clear_chat():
    return [], None, ""

# 로고 SVG 로드 및 스타일 수정
with open(f"{cur_dir}/resource/logo.svg", "r", encoding="utf-8") as svg_file:
    svg_content = svg_file.read()
font_size = "2.5em"
svg_content = re.sub(r'(<svg[^>]*)(>)', rf'\1 height="{font_size}" style="vertical-align: middle; display: inline-block;"\2', svg_content)
html = f"""
<p align="center" style="font-size: {font_size}; line-height: 1;">
    <span style="display: inline-block; vertical-align: middle;">{svg_content}</span>
    <span style="display: inline-block; vertical-align: middle;">{model_name.split('/')[-1]}</span>
</p>
<center>
    <font size=3>
        <b>Ovis</b> has been open-sourced on 
        <a href='https://huggingface.co/{model_name}'>😊 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>
"""

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
    }]

text_input = gr.Textbox(label="Prompt", placeholder="Enter your text here...", lines=1, container=False)

# 커스텀 CSS (배경 그라데이션, 반투명 컨테이너, 버튼 애니메이션 등)
custom_css = """
body {
    background: linear-gradient(135deg, #667eea, #764ba2);
    font-family: 'Helvetica Neue', Helvetica, Arial, sans-serif;
    color: #333;
    margin: 0;
    padding: 0;
}
.gradio-container {
    background: rgba(255, 255, 255, 0.95);
    border-radius: 15px;
    padding: 30px 40px;
    box-shadow: 0 8px 30px rgba(0, 0, 0, 0.3);
    margin: 40px auto;
    max-width: 1200px;
}
.gradio-container h1 {
    color: #333;
    text-shadow: 1px 1px 2px rgba(0, 0, 0, 0.2);
}
.fillable { 
    width: 95% !important; 
    max-width: unset !important;
}
#examples_container {
    margin: auto;
    width: 90%;
}
#examples_row {
    justify-content: center;
}
.sidebar {
    background: rgba(255, 255, 255, 0.98);
    border-radius: 10px;
    padding: 20px;
    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.2);
}
button, .btn {
    background: linear-gradient(90deg, #ff8a00, #e52e71);
    border: none;
    color: #fff;
    padding: 12px 24px;
    text-transform: uppercase;
    font-weight: bold;
    letter-spacing: 1px;
    border-radius: 5px;
    cursor: pointer;
    transition: transform 0.2s ease-in-out;
}
button:hover, .btn:hover {
    transform: scale(1.05);
}
"""

with gr.Blocks(css=custom_css, title=model_name.split('/')[-1]) as demo:
    gr.HTML(html)
    with gr.Row():
        with gr.Column(scale=3):
            image_input = gr.Image(label="Image", height=350, type="pil")
            gr.Examples(
                examples=[
                    [f"{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\nProvide a step-by-step solution to the problem, and conclude with 'the answer is' followed by the final solution."],
                    [f"{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\nProvide a step-by-step solution to the problem, and conclude with 'the answer is' followed by the final solution."],
                    [f"{cur_dir}/examples/ovis2_figure0.png", "Explain this model."],
                    [f"{cur_dir}/examples/ovis2_figure1.png", "Organize the notes about GRPO in the figure."],
                    [f"{cur_dir}/examples/ovis2_multi0.jpg", "Posso avere un frappuccino e un caffè americano di taglia M? Quanto costa in totale?"],
                ],
                inputs=[image_input, text_input]
            )
        with gr.Column(scale=7):
            chatbot = gr.Chatbot(label="Ovis", layout="panel", height=600, show_copy_button=True, latex_delimiters=latex_delimiters_set)
            text_input.render()
            with gr.Row():
                send_btn = gr.Button("Send")
                clear_btn = gr.Button("Clear")

    send_click_event = send_btn.click(
        submit_chat, 
        inputs=[chatbot, text_input], 
        outputs=[chatbot, text_input]
    ).then(
        ovis_chat, 
        inputs=[chatbot, image_input], 
        outputs=chatbot
    )
    submit_event = text_input.submit(
        submit_chat, 
        inputs=[chatbot, text_input], 
        outputs=[chatbot, text_input]
    ).then(
        ovis_chat, 
        inputs=[chatbot, image_input], 
        outputs=chatbot
    )
    clear_btn.click(clear_chat, outputs=[chatbot, image_input, text_input])

demo.launch()