app.py

"""
UI-UX Gradio Demo
Upload a mobile UI screenshot → select a diagnostic task → get UX defect diagnosis with reasoning.

Usage:
  # With local model
  python app.py --model /ix/share/models/UI_UX/checkpoint-700-merged

  # With HuggingFace model
  python app.py --model UI-UX/UI-UX

  # With vLLM backend (faster, recommended)
  python app.py --backend vllm --api-url http://localhost:8000/v1
"""

import os
import re
import argparse
import base64
from io import BytesIO

import gradio as gr

try:
    import spaces
    HAS_SPACES = True
except ImportError:
    HAS_SPACES = False

# ========================== Task Definitions ==========================

TASKS = {
    "PopupNoClose": {
        "name": "PopupNoClose",
        "description": "Check if a popup has an explicit close control",
        "dimension": "Usability",
        "prompt": """Core Task: Evaluate, based solely on the provided single UI screenshot, whether the 'pop-up' within it offers users an explicit control for closing it.
Options:
A. No modal pop-up is present
B. The modal pop-up lacks an explicit close control
C. The modal pop-up has an explicit close control
Output Format: Please analyze the provided UI screenshot, and output your evaluation result in the format $\\boxed{X}$ (where X is one of A-C).""",
    },
    "BubbleOcclT": {
        "name": "BubbleOcclT",
        "description": "Check if text overlay occludes page text",
        "dimension": "Usability",
        "prompt": """Core Task: Based solely on the provided single UI screenshot, determine whether the 'non-blocking text overlay that requires a close button' obscures the text on the current page.
Options:
A. The screenshot does not contain a 'non-blocking text overlay' that requires a close button.
B. Such an overlay exists but does not obscure the page text.
C. Such an overlay exists and obscures the page text.
Output Format: Please analyze the provided UI screenshot, and output your evaluation result in the format $\\boxed{X}$ (where X is one of A-C).""",
    },
    "BubbleOcclBtn": {
        "name": "BubbleOcclBtn",
        "description": "Check if text overlay occludes buttons",
        "dimension": "Usability",
        "prompt": """Core Task: Based solely on the provided single UI screenshot, determine whether the 'non-blocking text overlay that requires a close button' obscures the clickable hot-zone elements on the current page.
Options:
A. The screenshot does not contain a 'non-blocking text overlay' that requires a close button.
B. Such an overlay exists but does not obscure clickable elements.
C. Such an overlay exists and obscures clickable elements.
Output Format: Please analyze the screenshot, and output result in the format $\\boxed{X}$ (where X is one of A-C).""",
    },
    "PopupBlock": {
        "name": "PopupBlock",
        "description": "Check if popup blocks mini-program close button",
        "dimension": "Efficiency",
        "prompt": """Core Task: Based solely on the provided single UI screenshot, determine if the screenshot originates from a mini-program. If the screenshot is from a mini-program, determine whether the 'pop-up' window within its interface causes the mini-program framework's native 'Close button in the upper-right corner' to become unclickable.
Options:
A. Not a mini-program screenshot.
B. Clickable.
C. Not clickable.
Output Format: Please analyze the screenshot, and output result in the format $\\boxed{X}$ (where X is one of A-C).""",
    },
    "PopupStack": {
        "name": "PopupStack",
        "description": "Check if multiple modal dialogs are stacked",
        "dimension": "Efficiency",
        "prompt": """Core Task: Based solely on the provided single UI screenshot, determine whether the current interface contains two or more modal dialogs simultaneously.
Options:
A. Yes
B. No
Output Format: Please analyze the provided UI screenshot, and output your evaluation result in the format $\\boxed{X}$ (where X is one of A-B).""",
    },
    "MismatchBadge": {
        "name": "MismatchBadge",
        "description": "Check badge/functionality consistency",
        "dimension": "Trustworthiness",
        "content_label": "Badge & Functionality Description",
        "content_placeholder": "Describe the badge shown (e.g. '已认证企业') and the actual functionality of the entry point.",
        "prompt": """Core Task: Evaluate whether there is an "inconsistency between the badge and functionality" issue.
Badge & Functionality Description: {content}
Options:
A. Inconsistent
B. Consistent
Output Format: Please analyze the screenshot, and output result in the format $\\boxed{X}$ (where X is one of A-B).""",
    },
    "MismatchContent": {
        "name": "MismatchContent",
        "description": "Check service description consistency",
        "dimension": "Trustworthiness",
        "content_label": "Service Description",
        "content_placeholder": "Paste the service description text that should match the page content.",
        "prompt": """Core Task: The service description must match the overall functionality of the service page screenshot.
Service Description: {content}
Options:
A. Inconsistent
B. Consistent
Output Format: Please analyze the screenshot, and output result in the format $\\boxed{X}$ (where X is one of A-B).""",
    },
    "MismatchFunc": {
        "name": "MismatchFunc",
        "description": "Check ad/landing page consistency",
        "dimension": "Trustworthiness",
        "content_label": "Ad / Placement Information",
        "content_placeholder": "Describe the advertisement or placement creative (e.g. banner text, offer claimed).",
        "prompt": """Core Task: Based on the Landing Page Screenshot and Placement Information, determine whether there is an "inconsistency between the advertisement and the landing page" issue.
Placement Information: {content}
Options:
A. Inconsistency Between Advertisement and Landing Page
B. Consistency Between Advertisement and Landing Page
Output Format: Please analyze the screenshot, and output result in the format $\\boxed{X}$ (where X is one of A-B).""",
    },
}

ANSWER_LABELS = {
    "PopupNoClose": {"A": "No modal pop-up present", "B": "No close control", "C": "Has close control"},
    "BubbleOcclT": {"A": "No text overlay", "B": "Overlay exists, no occlusion", "C": "Overlay occludes text"},
    "BubbleOcclBtn": {"A": "No text overlay", "B": "Overlay exists, no occlusion", "C": "Overlay occludes buttons"},
    "PopupBlock": {"A": "Not a mini-program", "B": "Close button clickable", "C": "Close button blocked"},
    "PopupStack": {"A": "Yes, multiple dialogs", "B": "No stacking"},
    "MismatchBadge": {"A": "Inconsistent", "B": "Consistent"},
    "MismatchContent": {"A": "Inconsistent", "B": "Consistent"},
    "MismatchFunc": {"A": "Inconsistent", "B": "Consistent"},
}


# ========================== Inference Backends ==========================

class TransformersBackend:
    def __init__(self, model_path):
        import torch
        from transformers import AutoModelForImageTextToText, AutoProcessor

        hf_token = os.environ.get("HF_TOKEN", None)
        print(f"Loading model from: {model_path}")
        self.processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True, token=hf_token)
        self.model = AutoModelForImageTextToText.from_pretrained(
            model_path,
            dtype=torch.bfloat16,
            device_map="cuda",
            trust_remote_code=True,
            token=hf_token,
        )
        self.torch = torch  # removed — use import torch directly
        print("Model loaded.")

    def generate(self, image, prompt):
        import torch
        print(f"[Device] model device: {next(self.model.parameters()).device}")
        print(f"[Device] cuda available: {torch.cuda.is_available()}")
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image},
                    {"type": "text", "text": prompt},
                ],
            }
        ]

        inputs = self.processor.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_dict=True,
            return_tensors="pt",
        ).to(self.model.device)

        with torch.no_grad():
            output_ids = self.model.generate(
                **inputs,
                max_new_tokens=4096,
                temperature=0,
                do_sample=False,
            )

        generated_ids = output_ids[0][len(inputs.input_ids[0]):]
        raw_response = self.processor.decode(generated_ids, skip_special_tokens=False)
        clean_response = self.processor.decode(generated_ids, skip_special_tokens=True)
        return raw_response, clean_response


class VLLMBackend:
    def __init__(self, api_url, model_name="UI-UX/UI-UX-4B"):
        from openai import OpenAI
        api_key = os.environ.get("OPENAI_API_KEY", "dummy")
        self.client = OpenAI(base_url=api_url, api_key=api_key)
        self.model_name = model_name
        print(f"Using vLLM backend at: {api_url}")

    def generate(self, image, prompt):
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        img_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")

        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=[{
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_b64}"}},
                    {"type": "text", "text": prompt},
                ],
            }],
            max_tokens=4096,
            temperature=0,
        )

        content = response.choices[0].message.content
        reasoning = getattr(response.choices[0].message, "reasoning_content", None)
        if reasoning:
            raw_response = f"<think>{reasoning}</think>{content}"
        else:
            raw_response = content
        return raw_response, content


# ========================== Core Logic ==========================

def extract_answer(text):
    match = re.search(r'\\boxed\{([A-C])\}', text)
    if match:
        return match.group(1)
    match = re.search(r'boxed\{([A-C])\}', text)
    if match:
        return match.group(1)
    for letter in ['A', 'B', 'C']:
        if re.search(rf'(?:answer|Answer|ANSWER)\s*(?:is|:)\s*{letter}\b', text):
            return letter
    return None


def parse_response(raw_response, clean_response):
    thinking = None
    # Full <think>...</think> match
    think_match = re.search(r'<think>(.*?)</think>', raw_response, re.DOTALL)
    if think_match:
        thinking = think_match.group(1).strip()
    else:
        # <think> was stripped as special token but </think> is still present
        close_idx = raw_response.find('</think>')
        if close_idx != -1:
            thinking = raw_response[:close_idx].strip()

    # Strip thinking from full response display
    full_response = re.sub(r'<think>.*?</think>', '', clean_response, flags=re.DOTALL).strip()
    close_idx = full_response.find('</think>')
    if close_idx != -1:
        full_response = full_response[close_idx + len('</think>'):].strip()

    answer = extract_answer(full_response or clean_response)
    return answer, thinking, full_response


_backend = None  # set at startup

def predict(image, task_name, content=""):
    if image is None:
        return "", "", "", ""

    task = TASKS[task_name]
    prompt = task["prompt"]
    if "{content}" in prompt:
        prompt = prompt.replace("{content}", content or "")
    raw_response, clean_response = _backend.generate(image, prompt)
    answer, thinking, full_response = parse_response(raw_response, clean_response)

    if answer:
        label = ANSWER_LABELS.get(task_name, {}).get(answer, "")
        result_md = f"## Result: {answer}\n\n**{label}**"
    else:
        result_md = "## Result: Unable to parse\n\nModel response did not contain a clear answer."

    thinking_display = thinking if thinking else "(Model did not produce explicit reasoning)"

    return result_md, answer or "", thinking_display, full_response

if HAS_SPACES:
    predict = spaces.GPU(predict, duration=120)

# ========================== Gradio UI ==========================

def build_demo(backend):

    EXAMPLES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "images")

    with gr.Blocks(
        title="UI-UX: UX Defect Diagnosis",
        theme=gr.themes.Soft(),
        css="""
        .result-box { padding: 16px; border-radius: 8px; background: #f0fdf4; border: 1px solid #bbf7d0; }
        .result-box.defect { background: #fef2f2; border-color: #fecaca; }
        .task-badge { display: inline-block; padding: 2px 8px; border-radius: 4px; font-size: 0.8em; font-weight: 600; }
        .dim-usability { background: #dbeafe; color: #1d4ed8; }
        .dim-efficiency { background: #fef3c7; color: #92400e; }
        .dim-trustworthiness { background: #ede9fe; color: #5b21b6; }
        """
    ) as demo:

        gr.Markdown("""
# UI-UX: UX Defect Diagnosis with Multimodal LLMs

Upload a mobile UI screenshot and select a diagnostic task. The model will analyze the screenshot and identify potential UX defects with chain-of-thought reasoning.

<p>
<a href="https://huggingface.co/UI-UX/UI-UX-4B"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Model-UI--UX--4B-yellow"/></a>
<a href="https://huggingface.co/datasets/UI-UX/UXBench"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Dataset-UXBench-green"/></a>
<a href="https://cvpr.thecvf.com/virtual/2026/poster/41386"><img src="https://img.shields.io/badge/CVPR%202026-Findings-blue"/></a>
</p>
        """)

        with gr.Row(equal_height=False):
            # Left panel: Input
            with gr.Column(scale=1):
                image_input = gr.Image(
                    type="pil",
                    label="UI Screenshot",
                    height=400,
                )

                task_dropdown = gr.Dropdown(
                    choices=[(f"[{TASKS[k]['dimension']}] {k} — {TASKS[k]['description']}", k) for k in TASKS],
                    value="PopupNoClose",
                    label="Diagnostic Task",
                )

                CONTENT_TASKS = {"MismatchBadge", "MismatchContent", "MismatchFunc"}

                content_input = gr.Textbox(
                    label="Additional Content (required for Mismatch tasks)",
                    placeholder="For MismatchBadge: describe the badge and functionality. For MismatchContent: paste the service description. For MismatchFunc: describe the ad/placement creative.",
                    lines=3,
                    visible=True,
                )

                def on_task_change(task_name):
                    if task_name in CONTENT_TASKS:
                        task = TASKS[task_name]
                        return gr.update(
                            label=task["content_label"],
                            placeholder=task["content_placeholder"],
                        )
                    return gr.update(
                        label="Additional Content (required for Mismatch tasks)",
                        placeholder="For MismatchBadge: describe the badge and functionality. For MismatchContent: paste the service description. For MismatchFunc: describe the ad/placement creative.",
                    )

                task_dropdown.change(on_task_change, inputs=task_dropdown, outputs=content_input)

                with gr.Row():
                    clear_btn = gr.Button("Clear", variant="secondary")
                    submit_btn = gr.Button("Diagnose", variant="primary", size="lg")

            # Right panel: Output
            with gr.Column(scale=1):
                result_output = gr.Markdown(
                    value="*Upload an image and click Diagnose*",
                    label="Diagnosis Result",
                )
                answer_output = gr.Textbox(
                    label="Answer",
                    interactive=False,
                    max_lines=1,
                )

                with gr.Accordion("Reasoning (Chain-of-Thought)", open=False):
                    thinking_output = gr.Textbox(
                        label="Model Reasoning",
                        lines=12,
                        interactive=False,
                        show_copy_button=True,
                    )

                with gr.Accordion("Full Response", open=False):
                    full_output = gr.Textbox(
                        label="Complete Model Output",
                        lines=8,
                        interactive=False,
                        show_copy_button=True,
                    )

        # Events
        def on_predict(image, task_name, content):
            return predict(image, task_name, content)

        submit_btn.click(
            on_predict,
            inputs=[image_input, task_dropdown, content_input],
            outputs=[result_output, answer_output, thinking_output, full_output],
        )

        clear_btn.click(
            lambda: (None, "PopupNoClose", "", "*Upload an image and click Diagnose*", "", "", ""),
            outputs=[image_input, task_dropdown, content_input, result_output, answer_output, thinking_output, full_output],
        )

        # Examples
        EXAMPLES = [
            [os.path.join(EXAMPLES_DIR, "0004.png"), "PopupNoClose",    ""],
            [os.path.join(EXAMPLES_DIR, "0003.png"), "BubbleOcclT",     ""],
            [os.path.join(EXAMPLES_DIR, "0002.png"), "BubbleOcclBtn",   ""],
            [os.path.join(EXAMPLES_DIR, "0005.png"), "PopupBlock",      ""],
            [os.path.join(EXAMPLES_DIR, "0001.png"), "PopupStack",      ""],
            [os.path.join(EXAMPLES_DIR, "0004.png"), "MismatchBadge",
                "Badge: 'Free Claim'. Entry point: A recycling reward service where users exchange old items for credits — no free direct claim is available."],
            [os.path.join(EXAMPLES_DIR, "0003.png"), "MismatchContent",
                "Service Description: A daily check-in tool that lets users collect small cash rewards and withdraw once they reach the threshold."],
            [os.path.join(EXAMPLES_DIR, "0002.png"), "MismatchFunc",
                "Ad: 'Cainiao Guoguo — Free door-to-door pickup, ship anything for free with no minimum weight.' Landing page: Standard shipping service with weight-based pricing and a first-order ¥2 discount coupon."],
        ]

        gr.Examples(
            examples=EXAMPLES,
            inputs=[image_input, task_dropdown, content_input],
            label="Examples — click to load",
        )

        # Task info
        gr.Markdown("""
---

### Supported Tasks

| Dimension | Task | What it checks |
|-----------|------|----------------|
| **Usability** | PopupNoClose | Modal popup without close control |
| **Usability** | BubbleOcclT | Text overlay occluding page text |
| **Usability** | BubbleOcclBtn | Text overlay occluding buttons |
| **Efficiency** | PopupBlock | Popup blocking mini-program close |
| **Efficiency** | PopupStack | Multiple modal dialogs stacked |
| **Trustworthiness** | MismatchBadge | Badge/functionality inconsistency |
| **Trustworthiness** | MismatchContent | Service description mismatch |
| **Trustworthiness** | MismatchFunc | Ad/landing page inconsistency |

**Model**: UI-UX-4B (Qwen3.5-4B + GRPO) | **Benchmark**: 79.63% on UXBench | **Paper**: CVPR 2026 Findings
        """)

    return demo


# ========================== Entry Point ==========================

def parse_args():
    parser = argparse.ArgumentParser(description="UI-UX Gradio Demo")
    parser.add_argument("--model", type=str, default=None,
                        help="Model path (local or HuggingFace ID)")
    parser.add_argument("--backend", type=str, choices=["transformers", "vllm"], default="transformers",
                        help="Inference backend")
    parser.add_argument("--api-url", type=str, default="http://localhost:8000/v1",
                        help="vLLM API URL (only for --backend vllm)")
    parser.add_argument("--model-name", type=str, default="UI-UX/UI-UX-4B",
                        help="Model name for vLLM API")
    parser.add_argument("--port", type=int, default=7860, help="Server port")
    parser.add_argument("--share", action="store_true", help="Create public Gradio link")
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()

    model_path = args.model or os.environ.get("MODEL_ID", "afx-team/UI-UX")

    if args.backend == "vllm":
        _backend = VLLMBackend(args.api_url, args.model_name)
    else:
        _backend = TransformersBackend(model_path)

    demo = build_demo(_backend)
    demo.launch(
        server_name="0.0.0.0",
        server_port=args.port,
        share=args.share,
    )