demo/app.py

"""
app.py
------
Gradio chatbot demo for the trained CXR-VLM.

Flow
----
  1. Upload a chest X-ray image.
  2. The model auto-generates a report — Findings + Impression (cascade,
     matching report_mode=split_cascade) — as the assistant's first message.
  3. Ask any follow-up question about the same image; it's answered VQA-style,
     the way the model was trained.

The image stays "loaded" for the whole conversation; uploading a new image
resets the chat and produces a fresh report.

Usage (Colab / cloud GPU — auto-download the checkpoint from the HF runs repo)
-----------------------------------------------------------------------------
  python -m demo.app \
      --hf_repo  hieu3636/cxr-vlm-runs \
      --run_id   MIMIC-CXR_resized_run_3 \
      --share                       # public Gradio link (needed on Colab)

Usage (local — checkpoint already on disk)
------------------------------------------
  python -m demo.app \
      --checkpoint checkpoints/IU-Xray_run_1/stage2_instruct   # dir OR .pt

Notes
-----
  * On Colab pass --share so you get a public URL (the in-VM port isn't
    reachable otherwise).
  * --hf_token is optional if the runs repo is public or $HF_TOKEN is set.
"""

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from demo.cxr_demo import CXRChatbot, download_checkpoint_from_hf


WELCOME = (
    "Upload a chest X-ray to begin. CXR-HIEU writes a radiology report for it, "
    "then you can ask any follow-up question about the same image."
)


def parse_args():
    p = argparse.ArgumentParser(description="CXR-VLM chatbot demo (Gradio)")
    # Checkpoint source — either a local path OR an HF runs-repo coordinate.
    p.add_argument("--checkpoint", type=str, default=None,
                   help="Local checkpoint dir (contains checkpoint_projection.pt "
                        "+ checkpoint_lora/) or a stage2_final.pt path. "
                        "If omitted, --hf_repo + --run_id are used to download.")
    p.add_argument("--hf_repo", type=str, default="hieu3636/cxr-vlm-runs",
                   help="HF runs repo to pull the checkpoint from.")
    p.add_argument("--run_id", type=str, default="MIMIC-CXR_resized_run_3",
                   help="Run id on the HF runs repo, e.g. MIMIC-CXR_resized_run_3.")
    p.add_argument("--stage", type=str, default="stage2",
                   choices=["stage2", "stage1"],
                   help="Which stage's checkpoint to load (default stage2).")
    p.add_argument("--which", type=str, default="best",
                   choices=["best", "last"],
                   help="best (lowest eval_loss) or last (default best).")
    p.add_argument("--hf_token", type=str, default=None,
                   help="HF token (defaults to $HF_TOKEN). Needed for private repos.")
    p.add_argument("--model_config", type=str, default="configs/model_config.yaml")
    p.add_argument("--device", type=str, default="auto",
                   choices=["auto", "cuda", "cpu"])
    p.add_argument("--max_new_tokens", type=int, default=300)
    p.add_argument("--share", action="store_true",
                   help="Create a public Gradio link (use this on Colab).")
    p.add_argument("--server_name", type=str, default="0.0.0.0")
    p.add_argument("--server_port", type=int, default=7860)
    return p.parse_args()


def resolve_checkpoint(args) -> str:
    if args.checkpoint:
        return args.checkpoint
    if not args.run_id:
        raise SystemExit(
            "Provide --checkpoint <path>, OR --run_id <id> (with --hf_repo) "
            "to auto-download from the HF runs repo."
        )
    print(f"[app] downloading {args.run_id}/{args.stage}/{args.which} "
          f"from {args.hf_repo} …")
    return download_checkpoint_from_hf(
        repo_id = args.hf_repo,
        run_id  = args.run_id,
        stage   = args.stage,
        which   = args.which,
        token   = args.hf_token,
    )


def build_ui(bot: CXRChatbot, max_new_tokens: int):
    import demo._gradio_compat  # noqa: F401  (fix gradio_client schema crash)
    import gradio as gr

    with gr.Blocks(title="CXR-HIEU Demo", theme=gr.themes.Soft()) as demo:
        gr.Markdown(
            "# 🩻 CXR-HIEU — Chest X-ray Report & Q&A\n"
            "Upload a chest X-ray. CXR-HIEU generates a radiology report and "
            "answers questions about the image."
        )

        # Holds the active conversation's session dict {img, pnu}. The image
        # (and its predicted PNU) persist across follow-up questions; a new
        # upload replaces it → the bot answers about the new image.
        sess_state = gr.State(value=None)

        with gr.Row():
            with gr.Column(scale=1):
                image_in = gr.Image(type="pil", label="Chest X-ray image",
                                    height=380)
                analyze_btn = gr.Button("🔍 Analyze image", variant="primary")
                gr.Markdown(
                    "_Upload an image and click **Analyze image** (it also "
                    "analyzes automatically on upload). Then ask questions on "
                    "the right._"
                )
            with gr.Column(scale=2):
                chatbot = gr.Chatbot(label="Conversation", height=460,
                                     type="messages")
                with gr.Row():
                    msg = gr.Textbox(
                        placeholder="Ask about this image, e.g. Is there pleural effusion?",
                        scale=5, show_label=False, container=False,
                    )
                    send_btn = gr.Button("Send", scale=1, variant="primary")
                clear_btn = gr.Button("🗑️ Clear conversation")

        # ── handlers ───────────────────────────────────────────────────────

        def on_analyze(image):
            """New image → new session, reset chat, generate the report."""
            if image is None:
                return None, [{
                    "role": "assistant",
                    "content": "⚠️ Please upload a chest X-ray first.",
                }]
            session = bot.prepare(image)
            report = bot.generate_report(session, max_new_tokens=max_new_tokens)
            history = [{"role": "assistant", "content": report}]
            return session, history

        def on_send(message, history, session):
            history = history or []
            if not message or not message.strip():
                return "", history
            if session is None:
                history.append({"role": "user", "content": message})
                history.append({
                    "role": "assistant",
                    "content": "⚠️ Please upload an image and click "
                               "**Analyze image** before asking questions.",
                })
                return "", history
            answer = bot.answer_question(session, message.strip(),
                                         max_new_tokens=max_new_tokens)
            history.append({"role": "user", "content": message})
            history.append({"role": "assistant", "content": answer})
            return "", history

        def on_clear():
            return None, None, ""

        # Auto-analyze on upload AND via the button.
        image_in.upload(on_analyze, inputs=image_in,
                        outputs=[sess_state, chatbot])
        analyze_btn.click(on_analyze, inputs=image_in,
                          outputs=[sess_state, chatbot])
        send_btn.click(on_send, inputs=[msg, chatbot, sess_state],
                       outputs=[msg, chatbot])
        msg.submit(on_send, inputs=[msg, chatbot, sess_state],
                   outputs=[msg, chatbot])
        clear_btn.click(on_clear, outputs=[sess_state, chatbot, msg])

        gr.Markdown(f"\n\n_{WELCOME}_")

    return demo


def main():
    args = parse_args()
    ckpt = resolve_checkpoint(args)
    bot = CXRChatbot(
        checkpoint_dir = ckpt,
        model_config   = args.model_config,
        device         = args.device,
    )
    demo = build_ui(bot, args.max_new_tokens)
    demo.queue().launch(
        share       = args.share,
        server_name = args.server_name,
        server_port = args.server_port,
    )


if __name__ == "__main__":
    main()