app.py

"""Gradio Space: prune adaln_proj layers out of H3 LoRAs.

Flow: point it at a LoRA, the app reads just the file header (a couple of range
requests, even for multi-GB files) and shows exactly what it would remove, then
one button prunes it and hands back a download.

See h3_prune.py for the actual surgery.
"""

import hashlib
import html
import os
import shutil
import tempfile
import time

import gradio as gr

import h3_prune as H

WORK_ROOT = os.path.join(tempfile.gettempdir(), "h3_pruner")
MAX_OUTPUT_AGE = 60 * 60  # keep finished results around for an hour

CSS = """
/* flat, square, 1px borders — everything lives in a bordered panel */
:root { --h3-lit: #fb923c; --h3-mid: #f97316; --h3-deep: #ea580c; }

/* Spaces sizes its iframe from the height the app reports, inside a shell that
   does not itself scroll. So the document must stay pinned to the frame and
   scroll internally — if it is allowed to grow to its content height the frame
   grows with it and everything past the fold becomes unreachable. */
html, body { height: 100% !important; max-height: 100% !important; }
body { overflow-y: auto !important; }
.gradio-container { overflow-y: visible !important; max-height: none !important; }

/* without an explicit width the container shrink-wraps its content, so a short
   status line collapses the whole page and a long filename snaps it back out */
.gradio-container { width: 100% !important; max-width: 820px !important; margin: 0 auto !important; }
.gradio-container .h3panel, .gradio-container .h3body, .gradio-container .gr-accordion,
.gradio-container .tabs, .gradio-container .tabitem { width: 100% !important; }
.card, .stats { width: 100%; box-sizing: border-box; }
.gradio-container button, .gradio-container input, .gradio-container textarea,
.gradio-container .block, .gradio-container .form, .gradio-container select { border-radius: 0 !important; }

/* --- panels ---------------------------------------------------------- */
.h3panel { border: 1px solid var(--border-color-primary, #e4e4e7) !important;
         background: var(--background-fill-primary, #fff) !important;
         padding: 0 !important; gap: 0 !important; margin-bottom: 16px !important; }
.h3head { padding: 13px 16px !important; margin: 0 !important;
              border-bottom: 1px solid var(--border-color-primary, #e4e4e7);
              background: var(--background-fill-secondary, #fafafa); }
.h3head .h { display: flex; align-items: baseline; gap: 10px;
                 font-weight: 600; font-size: .95rem; letter-spacing: -0.01em;
                 color: var(--body-text-color, #09090b); }
.h3head .n { flex: none; width: 20px; height: 20px; font-size: .72rem; font-weight: 600;
                 display: inline-flex; align-items: center; justify-content: center;
                 color: #fff; background: var(--h3-mid); align-self: center; }
.h3head .hint { font-weight: 400; font-size: .8rem;
                    color: var(--body-text-color-subdued, #71717a); }
.h3body { padding: 16px !important; gap: 12px !important; }
/* gradio stamps elem_classes on both the block and an inner div, so the header
   would otherwise get its padding and bottom border twice */
.h3head .h3head { padding: 0 !important; border: none !important; background: none !important; }

/* --- hero ------------------------------------------------------------ */
#hero { padding: 18px 16px; text-align: center;
        border: 1px solid var(--border-color-primary, #e4e4e7); margin-bottom: 16px; }
#hero .title { font-size: 1.6rem; font-weight: 650; letter-spacing: -0.03em; margin: 0;
               color: var(--body-text-color, #09090b); }
#hero .grad { color: var(--h3-deep); }
#hero .sub { color: var(--body-text-color-subdued, #71717a); margin: 7px 0 0; font-size: .88rem; }
#hero code, .card code { background: var(--background-fill-secondary, #f4f4f5); padding: 1px 5px;
                         border: 1px solid var(--border-color-primary, #e4e4e7); font-size: .85em; }

/* --- state cards ----------------------------------------------------- */
.card { border: 1px solid var(--border-color-primary, #e4e4e7); border-left-width: 3px;
        padding: 13px 15px; background: var(--background-fill-secondary, #fafafa); }
.card.idle { border: 1px dashed var(--border-color-primary, #e4e4e7); background: transparent;
             color: var(--body-text-color-subdued, #71717a); text-align: center;
             padding: 20px 15px; font-size: .9rem; }
/* amber would read as the brand colour now, so "nothing to do" goes neutral and
   only real failures get a hot colour */
.card.ok   { border-left-color: var(--h3-mid); }
.card.warn { border-left-color: #a1a1aa; }
.card.err  { border-left-color: #dc2626; }
.card.busy { border-left-color: var(--h3-lit); display: flex; align-items: center; gap: 12px; }
.card .name { font-weight: 600; word-break: break-all; color: var(--body-text-color, #09090b);
              font-size: .92rem; }
.card .msg  { color: var(--body-text-color, #09090b); font-size: .92rem; }
.card .msg b { font-weight: 600; }
.card .note { margin-top: 5px; font-size: .83rem; color: var(--body-text-color-subdued, #71717a); }

.spinner { flex: none; width: 16px; height: 16px; border-radius: 50% !important;
           border: 2px solid var(--border-color-primary, #e4e4e7); border-top-color: var(--h3-mid);
           animation: h3spin .7s linear infinite; }
@keyframes h3spin { to { transform: rotate(360deg); } }
@media (prefers-reduced-motion: reduce) { .spinner { animation-duration: 2.2s; } }

/* --- stat row -------------------------------------------------------- */
.stats { display: flex; flex-wrap: wrap; gap: 10px 24px; margin-top: 11px; }
.stat b { display: block; font-size: 1.05rem; font-weight: 600;
          color: var(--body-text-color, #09090b); font-variant-numeric: tabular-nums; }
.stat span { font-size: .69rem; letter-spacing: .06em; text-transform: uppercase;
             color: var(--body-text-color-subdued, #71717a); }
.stat.hl b { color: var(--h3-deep); }
.stat .arrow { color: var(--body-text-color-subdued, #71717a); font-weight: 400; margin: 0 3px; }

/* --- buttons --------------------------------------------------------- */
#go { font-weight: 550 !important; }
/* inverted ink so the payoff reads differently from the action above it:
   dark-on-light in light mode, light-on-dark in dark */
#dl { background: var(--body-text-color, #09090b) !important;
      color: var(--body-background-fill, #fff) !important;
      border: 1px solid var(--body-text-color, #09090b) !important;
      font-weight: 550 !important; margin-top: 12px; }
#dl:hover { opacity: .88; }

/* --- accordions as boxes --------------------------------------------- */
.gradio-container .gr-accordion {
    border: 1px solid var(--border-color-primary, #e4e4e7) !important;
    background: var(--background-fill-primary, #fff) !important;
    margin-bottom: 16px !important; }
.gradio-container .gr-accordion .label-wrap { font-weight: 500; }
/* gradio fills form wrappers with the secondary tone, which reads as a grey
   slab sitting inside the panel */
.gradio-container .form { background: transparent !important; border: none !important; }
/* flattening the blocks took the inputs' boxes with it; put them back */
.gradio-container textarea,
.gradio-container input[type="text"],
.gradio-container input[type="password"] {
    border: 1px solid var(--border-color-primary, #e4e4e7) !important;
    background: var(--background-fill-primary, #fff) !important;
    padding: 8px 10px !important; }
.gradio-container textarea:focus,
.gradio-container input[type="text"]:focus,
.gradio-container input[type="password"]:focus {
    border-color: var(--h3-mid) !important; outline: none !important; }
/* the dropdown's own input is .border-none — its box comes from the wrapper */
.h3drop .wrap { border: 1px solid var(--border-color-primary, #e4e4e7) !important;
                background: var(--background-fill-primary, #fff) !important; }
.h3drop input { border: none !important; padding: 8px 10px !important; }
.h3drop .wrap:focus-within { border-color: var(--h3-mid) !important; }

footer { display: none !important; }
"""

IDLE = '<div class="card idle">Paste a Hugging Face path above and press <b>Enter</b> to inspect it.</div>'
IDLE_UPLOAD = '<div class="card idle">Drop a <code>.safetensors</code> LoRA above to inspect it.</div>'

# (repo, revision, filename, token-hash) -> (header, size), so retyping a pattern
# or a stray blur event never re-fetches anything.
_peek_cache = {}
_list_cache = {}


def _tokey(token):
    return hashlib.sha256((token or "").encode()).hexdigest()[:16] if token else ""


def _bounded(cache):
    if len(cache) > 32:
        cache.clear()
    return cache


def _card(kind, inner):
    return f'<div class="card {kind}">{inner}</div>'


def _head(n, title, hint=""):
    hint = f'<span class="hint">{hint}</span>' if hint else ""
    return f'<div class="h"><span class="n">{n}</span> {title} {hint}</div>'


def _error(msg):
    return _card("err", f'<div class="msg">❌ {html.escape(str(msg))}</div>')


def _busy(msg, note=""):
    note = f'<div class="note">{html.escape(note)}</div>' if note else ""
    return _card("busy", f'<span class="spinner"></span><div><div class="msg">{html.escape(msg)}'
                         f'</div>{note}</div>')


def _stat(value, label, hl=False):
    return (f'<div class="stat{" hl" if hl else ""}"><b>{value}</b><span>{html.escape(label)}</span></div>')


def _workdir():
    """A fresh output directory, sweeping stale ones on the way in."""
    os.makedirs(WORK_ROOT, exist_ok=True)
    cutoff = time.time() - MAX_OUTPUT_AGE
    for name in os.listdir(WORK_ROOT):
        path = os.path.join(WORK_ROOT, name)
        try:
            if os.path.getmtime(path) < cutoff:
                shutil.rmtree(path, ignore_errors=True)
        except OSError:
            pass
    return tempfile.mkdtemp(dir=WORK_ROOT)


def _list_files(repo_id, revision, token):
    key = (repo_id, revision, _tokey(token))
    if key not in _list_cache:
        _bounded(_list_cache)[key] = H.list_safetensors(repo_id, revision, token)
    return _list_cache[key]


def _peek(repo_id, revision, filename, token):
    key = (repo_id, revision, filename, _tokey(token))
    if key not in _peek_cache:
        _bounded(_peek_cache)[key] = H.remote_header(repo_id, filename, revision, token)
    return _peek_cache[key]


def inspect(mode, ref, choice, upload, pattern, use_regex, token):
    """Look at the chosen LoRA and describe what pruning it would do.

    Runs on every change to the source or the options. Returns updates for the
    file dropdown, the preview card, and the hidden result widgets (a new
    inspection invalidates any previous run).
    """
    pattern = (pattern or "").strip() or H.DEFAULT_PATTERN
    token = (token or "").strip() or os.environ.get("HF_TOKEN") or None
    hide = (gr.update(visible=False), gr.update(visible=False), gr.update(visible=False))
    keep_dropdown = gr.update()

    try:
        if mode == "upload":
            if not upload:
                return keep_dropdown, IDLE_UPLOAD, *hide
            path = upload if isinstance(upload, str) else upload.name
            name = os.path.basename(path)
            header, _ = H.read_header(path)
            summary = H.summarize(header, os.path.getsize(path), pattern, use_regex)
        else:
            if not (ref or "").strip():
                return gr.update(choices=[], value=None, visible=False), IDLE, *hide

            repo_id, revision, filename = H.parse_hf_ref(ref)
            files = _list_files(repo_id, revision, token)
            if not files:
                return (gr.update(choices=[], value=None, visible=False),
                        _error(f"no .safetensors files in {repo_id}"), *hide)

            # A selection left over from another repo is meaningless here.
            if choice not in files:
                choice = filename if filename in files else (files[0] if len(files) == 1 else None)
            multi = len(files) > 1
            keep_dropdown = gr.update(choices=files, value=choice, visible=multi)

            if choice is None:
                return (keep_dropdown,
                        _card("warn", '<div class="msg"><b>%d</b> LoRA files in this repo — '
                                      'pick the one you want above.</div>' % len(files)), *hide)

            name = os.path.basename(choice)
            header, size = _peek(repo_id, revision, choice, token)
            summary = H.summarize(header, size, pattern, use_regex)
    except H.PruneError as exc:
        return keep_dropdown, _error(exc), *hide
    except Exception as exc:
        return keep_dropdown, _error(f"{type(exc).__name__}: {exc}"), *hide

    total = len(summary["drop"]) + len(summary["keep"])
    label = html.escape(pattern)
    if not summary["drop"]:
        body = (f'<div class="name">{html.escape(name)}</div>'
                f'<div class="note">Nothing to do — no <code>{label}</code> layers in this LoRA '
                f'({total} layers, all kept). It should already load in H3.</div>')
        return keep_dropdown, _card("warn", body), *hide

    stats = (
        _stat(f"{len(summary['drop'])}", f"{pattern} layers to remove", hl=True)
        + _stat(f"{len(summary['keep'])} / {total}", "layers kept")
        + _stat(f"{H.human_size(summary['total_size'])} "
                f'<span class="arrow">→</span> {H.human_size(summary["output_size"])}',
                "size after")
    )
    body = f'<div class="name">{html.escape(name)}</div><div class="stats">{stats}</div>'
    return keep_dropdown, _card("ok", body), *hide


def inspect_source(mode, ref, choice, upload, pattern, use_regex, token):
    """inspect(), but shows a spinner first.

    Wired to the triggers that can hit the network (a new path, a different
    file, a token). Option changes call inspect() directly instead — those are
    answered from the cached header, and a spinner that flashes on every
    keystroke just looks broken.
    """
    empty = (not upload) if mode == "upload" else not (ref or "").strip()
    if not empty:
        hide = (gr.update(visible=False),) * 3
        yield gr.update(), _busy("Reading the file header…"), *hide
    yield inspect(mode, ref, choice, upload, pattern, use_regex, token)


def _resolve_source(mode, ref, choice, upload, token, progress):
    """Return (local_path, display_name) for whichever tab the user is on."""
    if mode == "upload":
        if not upload:
            raise H.PruneError("drop a .safetensors file in first")
        path = upload if isinstance(upload, str) else upload.name
        if not path.endswith(".safetensors"):
            raise H.PruneError("that is not a .safetensors file")
        return path, os.path.basename(path)

    repo_id, revision, filename = H.parse_hf_ref(ref)
    filename = choice or filename

    if not filename:
        progress(0.05, desc=f"listing {repo_id}")
        files = _list_files(repo_id, revision, token)
        if not files:
            raise H.PruneError(f"no .safetensors files in {repo_id}")
        if len(files) > 1:
            raise H.PruneError(
                f"this repo has {len(files)} LoRA files — pick the one you want from the dropdown"
            )
        filename = files[0]

    progress(0.1, desc=f"downloading {os.path.basename(filename)}")
    local = H.download(repo_id, filename, revision, token)
    return local, os.path.basename(filename)


def prune(mode, ref, choice, upload, pattern, use_regex, token, progress=gr.Progress()):
    """Resolve the source, strip the matching keys, hand back a download."""
    pattern = (pattern or "").strip() or H.DEFAULT_PATTERN
    token = (token or "").strip() or os.environ.get("HF_TOKEN") or None
    # download button, removed-keys box, its accordion
    hidden = (gr.update(visible=False), gr.update(value=""), gr.update(visible=False))

    yield gr.update(value=_busy("Working on it…", "large LoRAs take a moment to download"),
                    visible=True), *hidden

    try:
        src, name = _resolve_source(mode, ref, choice, upload, token, progress)

        progress(0.35, desc="reading header")
        info = H.plan(src, pattern, use_regex)
        if not info["drop"]:
            card = _card("warn", f'<div class="msg">Nothing to do — no <code>{html.escape(pattern)}</code>'
                                 f' layers in <b>{html.escape(name)}</b>.</div>')
            yield gr.update(value=card, visible=True), *hidden
            return

        out_path = os.path.join(_workdir(), H.default_output_name(name, pattern))
        H.write_stripped(src, out_path, info, lambda f: progress(0.4 + 0.55 * f, desc="writing"))

        progress(0.97, desc="verifying")
        H.verify(out_path, info["keep"])
        size = os.path.getsize(out_path)
    except H.PruneError as exc:
        yield gr.update(value=_error(exc), visible=True), *hidden
        return
    except Exception as exc:
        yield gr.update(value=_error(f"pruning failed — {type(exc).__name__}: {exc}"),
                        visible=True), *hidden
        return

    meta = info["metadata"]
    stats = (
        _stat(f"{len(info['drop'])}", f"{pattern} layers removed", hl=True)
        + _stat(f"{len(info['keep'])}", "layers kept")
        + _stat(f"{H.human_size(info['input_size'])} "
                f'<span class="arrow">→</span> {H.human_size(size)}', "size")
        + _stat(f"{len(meta)}" if meta else "—", "metadata entries kept")
    )
    body = (f'<div class="msg">✅ Pruned <b>{html.escape(os.path.basename(out_path))}</b></div>'
            f'<div class="stats">{stats}</div>'
            f'<div class="note">Verified: the output re-opens cleanly and every kept layer is '
            f'byte-identical to the source.</div>')

    yield (
        gr.update(value=_card("ok", body), visible=True),
        gr.update(value=out_path, visible=True),
        gr.update(value="\n".join(info["drop"])),
        gr.update(visible=True),
    )


with gr.Blocks(title="H3 LoRA Pruner") as demo:
    gr.HTML(
        '<div id="hero"><p class="title">✂️ <span class="grad">H3 LoRA Pruner</span></p>'
        '<p class="sub">Strips the <code>adaln_proj</code> layers out of a LoRA so it loads in H3.</p></div>'
    )
    mode = gr.State("hub")

    with gr.Column(elem_classes="h3panel"):
        gr.HTML(_head(1, "Pick your LoRA"), elem_classes="h3head")
        with gr.Column(elem_classes="h3body"):
            with gr.Tabs():
                with gr.Tab("From Hugging Face") as tab_hub:
                    ref = gr.Textbox(
                        label="Model path or link",
                        placeholder="owner/repo   or   "
                                    "https://huggingface.co/owner/repo/blob/main/lora.safetensors",
                        info="Press Enter, or click away, to inspect it.",
                        autofocus=True,
                    )
                    choice = gr.Dropdown(label="Which file?", choices=[], visible=False,
                                         interactive=True, elem_classes="h3drop")
                with gr.Tab("Upload a file") as tab_upload:
                    upload = gr.File(label="LoRA (.safetensors)", file_types=[".safetensors"],
                                     type="filepath")

    with gr.Column(elem_classes="h3panel"):
        gr.HTML(_head(2, "Review & prune"), elem_classes="h3head")
        with gr.Column(elem_classes="h3body"):
            preview = gr.HTML(IDLE)
            go = gr.Button("Prune LoRA", variant="primary", size="lg", elem_id="go")
            result = gr.HTML(visible=False)
            download_btn = gr.DownloadButton("Download pruned LoRA", visible=False,
                                             size="lg", elem_id="dl")
            with gr.Accordion("Removed layers", open=False, visible=False) as removed_acc:
                removed = gr.Textbox(show_label=False, lines=10, max_lines=14,
                                     buttons=["copy"], interactive=False)

    with gr.Accordion("Options", open=False):
        with gr.Row():
            pattern = gr.Textbox(
                label="Remove layers whose name contains",
                value=H.DEFAULT_PATTERN,
                info="Leave this alone for H3.",
            )
            use_regex = gr.Checkbox(label="Treat as a regex", value=False)
        token = gr.Textbox(
            label="Hugging Face token",
            type="password",
            placeholder="hf_… — only for private or gated repos",
            info="Used for this request only, never stored.",
        )

    with gr.Accordion("Help & notes", open=False):
        gr.Markdown(
            "**What goes in the path box** — anything you can copy off the hub:\n"
            "`owner/repo`, `owner/repo@branch`, a full `.../blob/main/lora.safetensors` link, "
            "or a path to a file in a subfolder. If the repo holds more than one LoRA a picker "
            "appears so you can choose.\n\n"
            "**Private or gated repos** — paste a read token from "
            "[your HF settings](https://huggingface.co/settings/tokens) under **Options**. It is "
            "used for that one request and never stored.\n\n"
            "**What actually happens** — the safetensors header is rewritten and the surviving "
            "tensors are copied across as raw bytes. Weights are never decoded, so dtypes "
            "(bf16, fp8, …) come out exactly as they went in, and `__metadata__` is preserved. "
            "Every output is re-opened and checked before you get it.\n\n"
            "**Privacy** — pruned files live in this Space's temporary storage and are deleted "
            "after an hour."
        )

    # --- wiring -----------------------------------------------------------
    src_inputs = [mode, ref, choice, upload, pattern, use_regex, token]
    src_outputs = [choice, preview, result, download_btn, removed_acc]

    # Anything that changes what would be pruned re-runs the preview. Header
    # reads are cached, so repeated/duplicate triggers cost nothing. Deliberately
    # not on ref.change — that would hit the network on every keystroke.
    # show_progress must stay "hidden": the dropdown is one of these outputs, and
    # gradio's built-in loading state blanks an output's contents while the call
    # runs — so picking a file would wipe its own label. The busy card covers it.
    gr.on(
        [ref.submit, ref.blur, choice.change, upload.change, token.blur],
        inspect_source, src_inputs, src_outputs, show_progress="hidden",
    )
    # Answered from the cached header, so these repaint instantly and want no spinner.
    gr.on([pattern.change, use_regex.change], inspect, src_inputs, src_outputs,
          show_progress="hidden")

    tab_hub.select(lambda: "hub", None, mode).then(inspect_source, src_inputs, src_outputs)
    tab_upload.select(lambda: "upload", None, mode).then(inspect_source, src_inputs, src_outputs)

    go.click(prune, src_inputs, [result, download_btn, removed, removed_acc])

# Flat, square, neutral surfaces with one orange accent. Base rather than Soft:
# Soft ships gradients, shadows and tinted labels, all of which fight this.
THEME = gr.themes.Base(
    primary_hue="orange",
    neutral_hue="zinc",
    radius_size=gr.themes.sizes.radius_none,
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
).set(
    body_background_fill="#ffffff",
    body_background_fill_dark="#09090b",
    background_fill_primary="#ffffff",
    background_fill_primary_dark="#09090b",
    background_fill_secondary="#fafafa",
    background_fill_secondary_dark="#131316",
    border_color_primary="#e4e4e7",
    border_color_primary_dark="#27272a",
    # panels supply the boxes, so gradio's own blocks stay flat inside them
    block_background_fill="transparent",
    block_background_fill_dark="transparent",
    block_border_width="0px",
    block_shadow="none",
    block_label_background_fill="transparent",
    block_label_background_fill_dark="transparent",
    block_label_text_color="#71717a",
    block_label_text_color_dark="#a1a1aa",
    block_label_text_weight="500",
    block_title_text_color="#71717a",
    block_title_text_color_dark="#a1a1aa",
    input_background_fill="#ffffff",
    input_background_fill_dark="#09090b",
    input_border_color="#e4e4e7",
    input_border_color_dark="#27272a",
    input_shadow="none",
    button_border_width="1px",
    button_primary_background_fill="#f97316",
    button_primary_background_fill_dark="#f97316",
    button_primary_background_fill_hover="#ea580c",
    button_primary_background_fill_hover_dark="#ea580c",
    button_primary_border_color="#f97316",
    button_primary_border_color_dark="#f97316",
    button_primary_text_color="#ffffff",
    button_primary_text_color_dark="#ffffff",
)

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=2, max_size=20).launch(theme=THEME, css=CSS)