app.py

import os
import argparse
import traceback
import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
logging.getLogger("http").setLevel(logging.WARNING)
logging.getLogger("httpx").setLevel(logging.WARNING)

import spaces
import gradio as gr
from conversation_public import default_conversation

auth_token = os.environ.get("TOKEN_FROM_SECRET")

##########################################
# LLM part
##########################################
import torch
from transformers import AutoProcessor, AutoTokenizer, AutoModelForCausalLM
from transformers import Qwen2_5_VLForConditionalGeneration, TextIteratorStreamer
from qwen_vl_utils import process_vision_info
from threading import Thread

# === Prompts ===
SYSTEM_PROMPT_LLM = "You are a helpful assistant."
SYSTEM_PROMPT_CAP = "You are given an image and a relevant question. Based on the query, please describe the image in details. Do not try to answer the question."

CAPTION_PROMPT = "Question: {}\nPlease describe the image. DO NOT try to answer the question!"
LLM_PROMPT = """In the following text, you will receive a detailed caption of an image and a relevant question. In addition, you will be provided with a tentative model response. You goal is to answer the question using these information.\n\n### The detailed caption of the provided image: {}\n\n### Note that the caption might contain incorrect solutions, do not be misguided by them.\n\n### A problem to be solved: {}\n\n### A tentative model response: {}\n\n### Note that the above tentative response might be inaccurate (due to calculation errors, incorrect logic/reasoning and so on), under such a case, please ignore it and give your own solutions. However, if you do not have enough evidence to show it is wrong, please output the tentative response."""

# === Initialize Models ===
MLLM_MODEL_PATH = "KaiChen1998/RACRO-7B-CRO-GRPO"
LLM_MODEL_PATH = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"

processor = AutoProcessor.from_pretrained(MLLM_MODEL_PATH)
tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_PATH)

mllm = Qwen2_5_VLForConditionalGeneration.from_pretrained(MLLM_MODEL_PATH, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2", device_map="auto")
llm = AutoModelForCausalLM.from_pretrained(LLM_MODEL_PATH, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2", device_map="auto")

mllm_sampling = dict(do_sample=False, temperature=0, max_new_tokens=8192)
llm_sampling = dict(temperature=0.6, top_p=0.95, max_new_tokens=8192)

# === Build Prompts ===
def build_messages(image_path, question):
    cap_msgs = [
        {"role": "system", "content": SYSTEM_PROMPT_CAP},
        {"role": "user", "content": [{"type": "image", "image": image_path}, {"type": "text", "text": CAPTION_PROMPT.format(question)}]}
    ]
    qa_msgs = [
        {"role": "user", "content": [{"type": "image", "image": image_path}, {"type": "text", "text": question + " Please think step by step. The final answer MUST BE put in \\boxed{}."}]}
    ]
    return cap_msgs, qa_msgs

##########################################
# Streaming
##########################################
mllm_streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15)
llm_streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15)

def stream_response(model, inputs, streamer, prompt, gen_kwargs):
    thread = Thread(target=model.generate, kwargs=dict(
        streamer=streamer,
        **inputs,
        **gen_kwargs
        )
    )
    thread.start()

    generated_text = prompt
    for new_text in streamer:
        generated_text += new_text
        yield generated_text

##########################################
# Gradio part
##########################################
no_change_btn = gr.Button()
enable_btn = gr.Button(interactive=True)
disable_btn = gr.Button(interactive=False)
server_error_msg = "**NETWORK ERROR DUE TO HIGH TRAFFIC. PLEASE REGENERATE OR REFRESH THIS PAGE.**"
server_oom_msg = "**OUT OF GPU MEMORY DETECTED. PLEASE DECREASE THE MAX OUTPUT TOKENS AND REGENERATE.**"

def load_demo_refresh_model_list():
    logging.info(f"load_demo.")
    state = default_conversation.copy()
    return state

def regenerate(state, image_process_mode):
    logging.info(f"regenerate.")
    state.messages = state.messages[:-3]
    prev_human_msg = state.messages[-1]
    if type(prev_human_msg[1]) in (tuple, list):
        prev_human_msg[1] = (*prev_human_msg[1][:2], image_process_mode, *prev_human_msg[1][3:])
    state.skip_next = False
    return (state, state.to_gradio_chatbot_public(), "", None) + (disable_btn,) * 2

def clear_history():
    logging.info(f"clear_history.")
    state = default_conversation.copy()
    return (state, state.to_gradio_chatbot_public(), "", None) + (disable_btn,) * 2

############
# Show prompt in the chatbot
# Input: [state, textbox, imagebox, image_process_mode]
# Return: [state, chatbot, textbox, imagebox] + btn_list
############
def add_text(state, text, image, image_process_mode):
    # Input legality checking
    logging.info(f"add_text. len: {len(text)}")
    if len(text) <= 0 or image is None:
        state.skip_next = True
        return (state, state.to_gradio_chatbot_public(), "", None) + (no_change_btn,) * 2
    
    # Deal with image inputs
    if image is not None:
        text = (text, image, image_process_mode, None)
    
    # Single round only
    state = default_conversation.copy()
    state.append_message(state.roles[0], text)
    state.skip_next = False
    logging.info(str(state.messages))
    return (state, state.to_gradio_chatbot_public(), "", None) + (disable_btn,) * 2

############
# Get response
# Input: [state]
# Return: [state, chatbot] + btn_list
############
@spaces.GPU
def http_bot(state):
    logging.info(f"http_bot.")

    if state.skip_next:
        yield (state, state.to_gradio_chatbot_public()) + (no_change_btn,) * 2
        return

    # Retrive prompt
    prompt = state.messages[-1][-1][0]
    all_images = state.get_images(return_pil=True)[0]
    pload = {"prompt": prompt, "images": f'List of {len(state.get_images())} images: {all_images}'}
    logging.info(f"==== request ====\n{pload}")
    
    # Construct prompt
    cap_msgs, qa_msgs = build_messages(all_images, prompt)
    cap_prompt = processor.apply_chat_template(cap_msgs, tokenize=False, add_generation_prompt=True)
    qa_prompt = processor.apply_chat_template(qa_msgs, tokenize=False, add_generation_prompt=True)
    image_tensor, _ = process_vision_info(cap_msgs)
    cap_inputs = processor(text=cap_prompt, images=image_tensor, return_tensors="pt").to(mllm.device)
    qa_inputs = processor(text=qa_prompt, images=image_tensor, return_tensors="pt").to(mllm.device)    
    
    # Step 1: Tentative Response
    state.append_message(state.roles[1], "# Tentative Response\n\n▌")
    try:
        for generated_text in stream_response(mllm, qa_inputs, mllm_streamer, qa_prompt, mllm_sampling):
            output = generated_text[len(qa_prompt):].strip()
            state.messages[-1][-1] = "# Tentative Response\n\n" + output + "▌"
            yield (state, state.to_gradio_chatbot_public()) + (disable_btn,) * 2
    except Exception as e:
        os.system("nvidia-smi")
        logging.info(traceback.print_exc())
        state.messages[-1][-1] = server_error_msg
        yield (state, state.to_gradio_chatbot_public()) + (enable_btn,) * 2
        return
    tentative_answer = output
    logging.info(f"Tentative Response: {tentative_answer}")
    state.messages[-1][-1] = state.messages[-1][-1][:-1]
    yield (state, state.to_gradio_chatbot_public()) + (disable_btn,) * 2
    
    # Step 2: Query-conditioned Caption
    state.append_message(state.roles[1], "# Query-conditioned Caption\n\n▌")
    try:
        for generated_text in stream_response(mllm, cap_inputs, mllm_streamer, cap_prompt, mllm_sampling):
            output = generated_text[len(cap_prompt):].strip()
            state.messages[-1][-1] = "# Query-conditioned Caption\n\n" + output + "▌"
            yield (state, state.to_gradio_chatbot_public()) + (disable_btn,) * 2
    except Exception as e:
        os.system("nvidia-smi")
        logging.info(traceback.print_exc())
        state.messages[-1][-1] = server_error_msg
        yield (state, state.to_gradio_chatbot_public()) + (enable_btn,) * 2
        return
    caption_text = output
    logging.info(f"Query-conditioned Caption: {caption_text}")
    state.messages[-1][-1] = state.messages[-1][-1][:-1]
    yield (state, state.to_gradio_chatbot_public()) + (disable_btn,) * 2
    
    # Step 3: Text-only Reasoning
    reason_msgs = [
        {"role": "system", "content": SYSTEM_PROMPT_LLM},
        {"role": "user", "content": LLM_PROMPT.format(caption_text, prompt, tentative_answer)}
    ]
    reason_prompt = tokenizer.apply_chat_template(reason_msgs, tokenize=False, add_generation_prompt=True)
    reason_inputs = tokenizer(reason_prompt, return_tensors="pt").to(llm.device)
    
    state.append_message(state.roles[1], "# Text-only Reasoning\n\n▌")
    try:
        for generated_text in stream_response(llm, reason_inputs, llm_streamer, reason_prompt, llm_sampling):
            output = generated_text[len(reason_prompt):].strip()
            state.messages[-1][-1] = "# Text-only Reasoning\n\n" + output + "▌"
            yield (state, state.to_gradio_chatbot_public()) + (disable_btn,) * 2
    except Exception as e:
        os.system("nvidia-smi")
        logging.info(traceback.print_exc())
        state.messages[-1][-1] = server_error_msg
        yield (state, state.to_gradio_chatbot_public()) + (enable_btn,) * 2
        return
    final_response = output
    logging.info(f"Text-only Reasoning: {final_response}")
    state.messages[-1][-1] = state.messages[-1][-1][:-1]
    yield (state, state.to_gradio_chatbot_public()) + (enable_btn,) * 2

############
# Layout Markdown
############
title_markdown = ("""
<div style="display: flex; align-items: center; padding: 20px; border-radius: 10px; background-color: #f0f0f0;">
  <div>
    <h1 style="margin: 0;">RACRO: Perceptual Decoupling for Scalable Multi-modal Reasoning via Reward-Optimized Captioning</h1>
    <h2 style="margin: 10px 0;">📃 <a href="https://www.arxiv.org/abs/2506.04559" style="font-weight: 400;">Paper</a> | 💻 <a href="https://github.com/gyhdog99/RACRO2" style="font-weight: 400;">Code</a> | 🤗 <a href="https://huggingface.co/collections/KaiChen1998/racro-6848ec8c65b3a0bf33d0fbdb" style="font-weight: 400;">HuggingFace</a></h2>
    <p  style="margin: 20px 0;">
      <strong>1. RACRO is designed for multi-modal reasoning, and thus, image inputs are <mark>ALWAYS</mark> necessary!</strong>
    </p>
  </div>
</div>
""")

learn_more_markdown = ("""
## Citation
<pre><code>@article{gou2025perceptual,
  author    = {Gou, Yunhao and Chen, Kai and Liu, Zhili and Hong, Lanqing and Jin, Xin and Li, Zhenguo and Kwok, James T. and Zhang, Yu}, 
  title     = {Perceptual Decoupling for Scalable Multi-modal Reasoning via Reward-Optimized Captioning},
  journal   = {arXiv preprint arXiv:2506.04559},
  year      = {2025},
}</code></pre>
""")

block_css = """
#buttons button {
    min-width: min(120px,100%);
}
.message-row img {
    margin: 0px !important;
}
.avatar-container img {
    padding: 0px !important;
}
"""

############
# Layout Demo
############
def build_demo(embed_mode):
    textbox = gr.Textbox(label="Text", show_label=False, placeholder="Enter text and then click 💬 Chat to talk with me ^v^", container=False)
    with gr.Blocks(title="RACRO", theme=gr.themes.Default(), css=block_css) as demo:
        state = gr.State()
        if not embed_mode:
            gr.HTML(title_markdown)

        ##############
        # Chatbot
        ##############
        with gr.Row(equal_height=True):
            with gr.Column(scale=1):
                imagebox = gr.Image(type="pil", label="Image")
                image_process_mode = gr.Radio(
                    ["Crop", "Resize", "Pad", "Default"],
                    value="Default",
                    label="Preprocess for non-square image", visible=False)

                gr.Examples(examples=[
                    ["./examples/image-text/demo_example.jpg", "When the canister is momentarily stopped by the spring, by what distance $d$ is the spring compressed?"],
                ], inputs=[imagebox, textbox], label='Examples')

            with gr.Column(scale=8):
                chatbot = gr.Chatbot(
                    type="messages",
                    elem_id="chatbot",
                    label="RACRO Chatbot",
                    layout="bubble",
                    avatar_images=["examples/user_avator.png", "examples/icon_256.png"]
                )
                textbox.render()
                with gr.Row(elem_id="buttons") as button_row:
                    submit_btn = gr.Button(value="💬  Chat", variant="primary")
                    # stop_btn = gr.Button(value="⏹️  Stop Generation", interactive=False)
                    regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
                    clear_btn = gr.Button(value="🗑️  Clear", interactive=False)
        
        if not embed_mode:
            gr.Markdown(learn_more_markdown)

        # Register listeners
        btn_list = [regenerate_btn, clear_btn]
        regenerate_btn.click(
            regenerate,
            [state, image_process_mode],
            [state, chatbot, textbox, imagebox] + btn_list
        ).then(
            http_bot,
            [state],
            [state, chatbot] + btn_list,
        )

        clear_btn.click(
            clear_history,
            None,
            [state, chatbot, textbox, imagebox] + btn_list,
            queue=False
        )

        # probably mean press enter
        textbox.submit(
            add_text,
            [state, textbox, imagebox, image_process_mode],
            [state, chatbot, textbox, imagebox] + btn_list,
            queue=False
        ).then(
            http_bot,
            [state],
            [state, chatbot] + btn_list,
        )

        submit_btn.click(
            add_text,
            [state, textbox, imagebox, image_process_mode],
            [state, chatbot, textbox, imagebox] + btn_list
        ).then(
            http_bot,
            [state],
            [state, chatbot] + btn_list,
        )

        ##############
        # Demo loading
        ##############
        demo.load(
            load_demo_refresh_model_list,
            None,
            [state],
            queue=False
        )
    return demo


parser = argparse.ArgumentParser()
parser.add_argument("--share", action="store_true")
parser.add_argument("--embed", action="store_true")
args = parser.parse_args()

demo = build_demo(args.embed)
demo.queue(
    max_size=10,
    api_open=False
).launch(
    favicon_path="./examples/icon_256.png",
    allowed_paths=["/"],
    share=args.share
)