"""CVRR — Causal Visual Recurrent Reasoning, on Qwen2.5-VL-7B.

Paper: "Reason Through the Latent! Making Latent Visual Reasoning Necessary"
       (arXiv:2609.06746, Park, Jung & Kang)
Code:  https://github.com/dmis-lab/CVRR
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  (must precede torch / CUDA-touching imports)

import time  # noqa: E402
import torch  # noqa: E402
import gradio as gr  # noqa: E402
import transformers  # noqa: E402
from transformers import AutoModelForImageTextToText, AutoProcessor  # noqa: E402

MODEL_ID = "dmis-lab/Qwen2.5-VL-7B-CVRR"

# --------------------------------------------------------------------------
# The released CVRR code re-instantiates the processor on every
# `prepare_inputs` call. Cache it so the cost is paid once, at import time,
# instead of inside every GPU window.
# --------------------------------------------------------------------------
_PROCESSORS: dict = {}
_orig_processor_from_pretrained = AutoProcessor.from_pretrained


def _cached_processor_from_pretrained(path, **kwargs):
    key = str(path)
    if key not in _PROCESSORS:
        _PROCESSORS[key] = _orig_processor_from_pretrained(path, **kwargs)
    return _PROCESSORS[key]


transformers.AutoProcessor.from_pretrained = _cached_processor_from_pretrained

print(f"Loading {MODEL_ID} ...", flush=True)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    dtype=torch.bfloat16,
    device_map="cpu",  # ZeroGPU: build on CPU, then hand the whole replica to .to("cuda")
).eval()
model = model.to("cuda")

RELEASE = model.config.release
NATIVE_DIR = model.release_path / "native_backbone"
processor = _cached_processor_from_pretrained(
    NATIVE_DIR, local_files_only=True, use_fast=False
)
tokenizer = processor.tokenizer
model.tokenizer = tokenizer
EOS_TOKEN_ID = tokenizer.eos_token_id
print("Model ready.", flush=True)


def _estimate_duration(
    image=None, question=None, max_new_tokens=64, max_visual_tokens=4096, *args, **kwargs
) -> int:
    """ZeroGPU reservation: cold-start weight streaming plus input-dependent decode."""
    try:
        tokens = int(max_new_tokens)
        visual = int(max_visual_tokens)
    except (TypeError, ValueError):
        tokens, visual = 64, 4096
    return int(min(75, 22 + tokens * 0.045 + visual * 0.0015))


@spaces.GPU(duration=_estimate_duration)
def answer_question(
    image,
    question: str,
    max_new_tokens: int = 64,
    max_visual_tokens: int = 4096,
    progress=gr.Progress(track_tqdm=True),
) -> tuple[str, str]:
    """Answer a question about an image with CVRR latent visual recurrent reasoning.

    Args:
        image: the input image (PIL image) to reason about.
        question: a question about the image, in natural language.
        max_new_tokens: maximum number of answer tokens to decode greedily.
        max_visual_tokens: ceiling on visual tokens fed to the recurrent transition.

    Returns:
        A tuple of (answer text, a short run-info line).
    """
    if image is None:
        raise gr.Error("Please provide an image.")
    question = (question or "").strip()
    if not question:
        raise gr.Error("Please type a question about the image.")

    image = image.convert("RGB")
    started = time.perf_counter()
    inputs = model.prepare_inputs(
        image, question, max_visual_tokens=int(max_visual_tokens)
    )
    n_visual = int(inputs["input_ids"].eq(model.runtime.config.image_token_id).sum())

    token_ids = model.generate(
        **inputs,
        do_sample=False,
        max_new_tokens=int(max_new_tokens),
        eos_token_id=EOS_TOKEN_ID,
    )
    answer = tokenizer.decode(token_ids[0], skip_special_tokens=True).strip()
    elapsed = time.perf_counter() - started

    info = (
        f"**{n_visual}** visual tokens · **{int(RELEASE['inference_T'])}** recurrent steps "
        f"(shared decoder layer {int(RELEASE['recurrent_layer'])}, boundary "
        f"ℓ\\*={int(RELEASE['ell_star'])}) · **{elapsed:.1f} s**"
    )
    return (answer or "(empty answer)"), info


DESCRIPTION = """
# 🔁 CVRR — Reason Through the Latent

**Causal Visual Recurrent Reasoning** on top of `Qwen2.5-VL-7B-Instruct`. The lower
decoder is frozen up to the causal visual-read boundary; a single decoder layer is then
reused as a **shared recurrent transition** that refines the question state while keeping
the native visual rows as fixed evidence. Only the final recurrent state reaches the
answer decoder — the original visual rows and multimodal prefix cache are cut off, which
makes the latent visual reasoning *causally necessary* rather than merely present.

Ask a question about an image — open-ended, or in the benchmark-style multiple-choice
format the model was tuned on. Decoding is deterministic (greedy) by design. The recurrent
transition is supervised on short VQA answers, so crisp, targeted questions work best;
long free-form descriptions drift.

[Paper](https://huggingface.co/papers/2609.06746) · [Code](https://github.com/dmis-lab/CVRR) · [Model](https://huggingface.co/dmis-lab/Qwen2.5-VL-7B-CVRR)
"""

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EXAMPLES = [
    ["examples/mmvp_11.jpg", "Has the peacock opened its tail? Answer briefly."],
    ["examples/mmvp_7.jpg", "Is the school bus driving towards or away from the camera?"],
    ["examples/mmvp_1.jpg", "Are the butterfly's wings closer to being open or closed?"],
    [
        "examples/mmvp_5.jpg",
        "Is the dog facing left or right from the camera's perspective?\n"
        "(a) Left (b) Right\nAnswer with only the letter of the correct option.",
    ],
]

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(DESCRIPTION)

        with gr.Row():
            with gr.Column():
                image_in = gr.Image(type="pil", label="Image", height=360)
            with gr.Column():
                question_in = gr.Textbox(
                    label="Question",
                    lines=3,
                    placeholder="Is the peacock's tail open or closed?",
                )
                run_btn = gr.Button("Answer", variant="primary")
                answer_out = gr.Textbox(label="Answer", lines=4)
                info_out = gr.Markdown()

        with gr.Accordion("Advanced settings", open=False):
            max_new_tokens_in = gr.Slider(
                minimum=1, maximum=256, value=64, step=1, label="Max new tokens"
            )
            max_visual_tokens_in = gr.Slider(
                minimum=256,
                maximum=8192,
                value=4096,
                step=256,
                label="Max visual tokens (image resolution budget)",
            )

        gr.Examples(
            examples=EXAMPLES,
            inputs=[image_in, question_in],
            outputs=[answer_out, info_out],
            fn=answer_question,
            cache_examples=True,
            cache_mode="lazy",
        )

    gr.on(
        triggers=[run_btn.click, question_in.submit],
        fn=answer_question,
        inputs=[image_in, question_in, max_new_tokens_in, max_visual_tokens_in],
        outputs=[answer_out, info_out],
        api_name="answer_question",
    )

if __name__ == "__main__":
    demo.queue().launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)