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,
):
    if source_choice == "Upload 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()}

    data = {
        "inputs": inputs,
        "parameters": {
            "general_threshold": float(general_threshold),
            "character_threshold": float(character_threshold),
        },
    }

    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)

    features = ", ".join(sorted(out.get("feature", []))) or "—"
    characters = ", ".join(sorted(out.get("character", []))) or "—"
    ips = ", ".join(out.get("ip", [])) or "—"

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

    return features, characters, ips, meta, out


with gr.Blocks(title="PixAI Tagger v0.9 — Demo", fill_height=True) as demo:
    gr.Markdown(
        """
        # PixAI Tagger v0.9 — Gradio Demo
        Downloads model assets from **pixai-labs/pixai-tagger-v0.9** on first run,
        then uses your imported `EndpointHandler` to predict **general**, **character**, and **IP** tags.

        Configure via env vars:
        - `ASSETS_REPO_ID` (default: `pixai-labs/pixai-tagger-v0.9`)
        - `ASSETS_REVISION` (optional)
        - `MODEL_DIR` (default: `./assets`)
        """
    )
    with gr.Row():
        gr.Markdown(f"**{DEVICE_LABEL}**")

    with gr.Row():
        source_choice = gr.Radio(
            choices=["Upload image", "From URL"],
            value="Upload image",
            label="Image source",
        )

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

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

            source_choice.change(toggle_inputs, [source_choice], [image, url])

        with gr.Column(scale=1):
            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"
            )
            run_btn = gr.Button("Run", variant="primary")
            clear_btn = gr.Button("Clear")

    with gr.Row():
        with gr.Column():
            gr.Markdown("### Predicted Tags")
            features_out = gr.Textbox(label="General tags", lines=4)
            characters_out = gr.Textbox(label="Character tags", lines=4)
            ip_out = gr.Textbox(label="IP tags", lines=2)

        with gr.Column():
            gr.Markdown("### Metadata & Raw Output")
            meta_out = gr.JSON(label="Timings/Device")
            raw_out = gr.JSON(label="Raw JSON")

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

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

    run_btn.click(
        run_inference,
        inputs=[source_choice, image, url, general_threshold, character_threshold],
        outputs=[features_out, characters_out, ip_out, meta_out, raw_out],
        api_name="predict",
    )
    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()