app.py

import os
import time
import shutil
from pathlib import Path
from typing import Optional

import gradio as gr
from huggingface_hub import snapshot_download
from PIL import Image

# Import your existing inference endpoint implementation
from handler import EndpointHandler


# ------------------------------------------------------------------------------
# Asset setup: download weights/tags/mapping so local filenames are unchanged
# ------------------------------------------------------------------------------

REPO_ID = os.environ.get("ASSETS_REPO_ID", "pixai-labs/pixai-tagger-v0.9")
REVISION = os.environ.get("ASSETS_REVISION")  # optional pin, e.g. "main" or a commit
MODEL_DIR = os.environ.get("MODEL_DIR", "./assets")  # where the handler will look

# Optional: Hugging Face token for private repos
HF_TOKEN = (
    os.environ.get("HUGGINGFACE_HUB_TOKEN")
    or os.environ.get("HF_TOKEN")
    or os.environ.get("HUGGINGFACE_TOKEN")
    or os.environ.get("HUGGINGFACEHUB_API_TOKEN")
)

REQUIRED_FILES = [
    "model_v0.9.pth",
    "tags_v0.9_13k.json",
    "char_ip_map.json",
]

def ensure_assets(repo_id: str, revision: Optional[str], target_dir: str):
    """
    1) snapshot_download the upstream repo (cached by HF Hub)
    2) copy the required files into `target_dir` with the exact filenames expected
    """
    target = Path(target_dir)
    target.mkdir(parents=True, exist_ok=True)

    # Only download if something is missing
    missing = [f for f in REQUIRED_FILES if not (target / f).exists()]
    if not missing:
        return

    # Download snapshot (optionally filtered to speed up)
    snapshot_path = snapshot_download(
        repo_id=repo_id,
        revision=revision,
        allow_patterns=REQUIRED_FILES,  # only pull what we need
        token=HF_TOKEN,  # authenticate if repo is private
    )

    # Copy files into target_dir with the required names
    for fname in REQUIRED_FILES:
        src = Path(snapshot_path) / fname
        dst = target / fname
        if not src.exists():
            raise FileNotFoundError(
                f"Expected '{fname}' not found in snapshot for {repo_id} @ {revision or 'default'}"
            )
        shutil.copyfile(src, dst)


# Fetch assets (no-op if they already exist)
ensure_assets(REPO_ID, REVISION, MODEL_DIR)


# ------------------------------------------------------------------------------
# Initialize the handler
# ------------------------------------------------------------------------------

handler = EndpointHandler(MODEL_DIR)
DEVICE_LABEL = f"Device: {handler.device.upper()}"


# ------------------------------------------------------------------------------
# Gradio wiring
# ------------------------------------------------------------------------------

def run_inference(
    source_choice: str,
    image: Optional[Image.Image],
    url: str,
    general_threshold: float,
    character_threshold: float,
    mode_val: str,
    topk_general_val: int,
    topk_character_val: int,
    include_scores_val: bool,
    underscore_mode_val: bool,
):
    # Determine which input to use based on which Run button invoked the function.
    # We'll pass a string flag via source_choice: either "url" or "image".
    if source_choice == "image":
        if image is None:
            raise gr.Error("Please upload an image.")
        inputs = image
    else:
        if not url or not url.strip():
            raise gr.Error("Please provide an image URL.")
        inputs = {"url": url.strip()}

    params = {
        "general_threshold": float(general_threshold),
        "character_threshold": float(character_threshold),
        "mode": mode_val,
        "topk_general": int(topk_general_val),
        "topk_character": int(topk_character_val),
        "include_scores": bool(include_scores_val),
    }
    data = {"inputs": inputs, "parameters": params}

    started = time.time()
    try:
        out = handler(data)
    except Exception as e:
        raise gr.Error(f"Inference error: {e}") from e
    latency = round(time.time() - started, 4)

    # Individual outputs
    if underscore_mode_val:
        characters = " ".join(out.get("character", [])) or "—"
        ips = " ".join(out.get("ip", [])) or "—"
        features = " ".join(out.get("feature", [])) or "—"
    elif include_scores_val:
        gen_scores = out.get("feature_scores", {})
        char_scores = out.get("character_scores", {})
        characters = ", ".join(
            f"{k.replace('_', ' ')} ({char_scores[k]:.2f})" for k in sorted(char_scores, key=char_scores.get, reverse=True)
        ) or "—"
        ips = ", ".join(tag.replace("_", " ") for tag in out.get("ip", [])) or "—"
        features = ", ".join(
            f"{k.replace('_', ' ')} ({gen_scores[k]:.2f})" for k in sorted(gen_scores, key=gen_scores.get, reverse=True)
        ) or "—"
    else:
        characters = ", ".join(sorted(t.replace("_", " ") for t in out.get("character", []))) or "—"
        ips = ", ".join(tag.replace("_", " ") for tag in out.get("ip", [])) or "—"
        features = ", ".join(sorted(t.replace("_", " ") for t in out.get("feature", []))) or "—"

    # Combined output: probability-descending if scores available; else character, IP, general
    if underscore_mode_val:
        combined = " ".join(out.get("character", []) + out.get("ip", []) + out.get("feature", [])) or "—"
    else:
        char_scores = out.get("character_scores") or {}
        gen_scores = out.get("feature_scores") or {}
        if include_scores_val and (char_scores or gen_scores):
            # Build (tag, score) pairs
            char_pairs = [(k, float(char_scores.get(k, 0.0))) for k in out.get("character", [])]
            ip_pairs = [(k, 1.0) for k in out.get("ip", [])]  # IP has no score; treat equally
            gen_pairs = [(k, float(gen_scores.get(k, 0.0))) for k in out.get("feature", [])]
            all_pairs = char_pairs + ip_pairs + gen_pairs
            all_pairs.sort(key=lambda t: t[1], reverse=True)
            combined = ", ".join(
                [f"{k.replace('_', ' ')} ({score:.2f})" if (k in char_scores or k in gen_scores) else k.replace('_', ' ') for k, score in all_pairs]
            ) or "—"
        else:
            combined = ", ".join(
                list(sorted(t.replace("_", " ") for t in out.get("character", []))) +
                [tag.replace("_", " ") for tag in out.get("ip", [])] +
                list(sorted(t.replace("_", " ") for t in out.get("feature", [])))
            ) or "—"

    meta = {
        "device": handler.device,
        "latency_s_total": latency,
        **out.get("_timings", {}),
        "params": out.get("_params", {}),
    }

    return features, characters, ips, combined, meta, out


theme = gr.themes.Soft(primary_hue="indigo", secondary_hue="violet", radius_size="lg",)

with gr.Blocks(title="PixAI Tagger v0.9 — Demo", fill_height=True, theme=theme, analytics_enabled=False) as demo:
    gr.Markdown(
        """
        # PixAI Tagger v0.9 — Gradio Demo
        """
    )
    with gr.Row():
        gr.Markdown(f"**{DEVICE_LABEL}** — adjust thresholds or switch to Top-K mode.")

    with gr.Accordion("Settings", open=False):
        mode = gr.Radio(
            choices=["threshold", "topk"], value="threshold", label="Mode"
        )
        with gr.Group(visible=True) as threshold_group:
            general_threshold = gr.Slider(
                minimum=0.0, maximum=1.0, step=0.01, value=0.30, label="General threshold"
            )
            character_threshold = gr.Slider(
                minimum=0.0, maximum=1.0, step=0.01, value=0.85, label="Character threshold"
            )
        with gr.Group(visible=False) as topk_group:
            topk_general = gr.Slider(
                minimum=0, maximum=100, step=1, value=25, label="Top-K general"
            )
            topk_character = gr.Slider(
                minimum=0, maximum=100, step=1, value=10, label="Top-K character"
            )
        include_scores = gr.Checkbox(value=False, label="Include scores in output")
        underscore_mode = gr.Checkbox(value=False, label="Underscore-separated output")

        def toggle_mode(selected):
            return (
                gr.update(visible=(selected == "threshold")),
                gr.update(visible=(selected == "topk")),
            )

        mode.change(toggle_mode, inputs=[mode], outputs=[threshold_group, topk_group])

    with gr.Row(variant="panel"):
        with gr.Column(scale=2):
            image = gr.Image(label="Upload image", type="pil", visible=True, height="420px")
            url = gr.Textbox(label="Image URL", placeholder="https://…", visible=True)

            def toggle_inputs(choice):
                return (
                    gr.update(visible=(choice == "Upload image")),
                    gr.update(visible=(choice == "From URL")),
                )



        with gr.Column(scale=3):
            # No source choice; show both inputs and two run buttons
            with gr.Row():
                run_image_btn = gr.Button("Run from image", variant="primary")
                run_url_btn = gr.Button("Run from URL")
                clear_btn = gr.Button("Clear")

            gr.Markdown("### Combined Output (character → IP → general)")
            combined_out = gr.Textbox(label="Combined tags", lines=10,)
            copy_combined = gr.Button("Copy combined")

    with gr.Row():
        with gr.Column():
            gr.Markdown("### Character / General / IP")
            with gr.Row():
                with gr.Column():
                    characters_out = gr.Textbox(label="Character tags", lines=5,)
                with gr.Column():
                    features_out = gr.Textbox(label="General tags", lines=5,)
                with gr.Column():
                    ip_out = gr.Textbox(label="IP tags", lines=5,)
            with gr.Row():
                copy_characters = gr.Button("Copy character")
                copy_features = gr.Button("Copy general")
                copy_ip = gr.Button("Copy IP")

    with gr.Accordion("Metadata & Raw Output", open=False):
        with gr.Row():
            with gr.Column():
                meta_out = gr.JSON(label="Timings/Device")
                raw_out = gr.JSON(label="Raw JSON")
                copy_raw = gr.Button("Copy raw JSON")

    examples = gr.Examples(
        label="Examples (URL mode)",
        examples=[
            [None, "https://cdn.donmai.us/sample/50/b7/__komeiji_koishi_touhou_drawn_by_cui_ying__sample-50b7006f16e0144d5b5db44cadc2d22f.jpg", 0.30, 0.85, "threshold", 25, 10, False, False],
        ],
        inputs=[image, url, general_threshold, character_threshold, mode, topk_general, topk_character, include_scores, underscore_mode],
        cache_examples=False,
    )

    def clear():
        return (None, "", 0.30, 0.85, "", "", "", "", {}, {})

    # Bind buttons separately with a flag for source
    run_url_btn.click(
        run_inference,
        inputs=[
            gr.State("url"), image, url,
            general_threshold, character_threshold,
            mode, topk_general, topk_character, include_scores, underscore_mode,
        ],
        outputs=[features_out, characters_out, ip_out, combined_out, meta_out, raw_out],
        api_name="predict_url",
    )

    run_image_btn.click(
        run_inference,
        inputs=[
            gr.State("image"), image, url,
            general_threshold, character_threshold,
            mode, topk_general, topk_character, include_scores, underscore_mode,
        ],
        outputs=[features_out, characters_out, ip_out, combined_out, meta_out, raw_out],
        api_name="predict_image",
    )

    # Copy buttons
    copy_combined.click(lambda x: x, inputs=[combined_out], outputs=[combined_out])
    copy_characters.click(lambda x: x, inputs=[characters_out], outputs=[characters_out])
    copy_features.click(lambda x: x, inputs=[features_out], outputs=[features_out])
    copy_ip.click(lambda x: x, inputs=[ip_out], outputs=[ip_out])
    copy_raw.click(lambda x: x, inputs=[raw_out], outputs=[raw_out])
    clear_btn.click(
        clear,
        inputs=None,
        outputs=[
            image, url, general_threshold, character_threshold,
            features_out, characters_out, ip_out, meta_out, raw_out
        ],
    )

if __name__ == "__main__":
    demo.queue(max_size=8).launch()