"""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 = 'Paste a Hugging Face path above and press Enter to inspect it.
'
IDLE_UPLOAD = 'Drop a .safetensors LoRA above to inspect it.
'
# (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'{inner}
'
def _head(n, title, hint=""):
hint = f'{hint}' if hint else ""
return f'{n} {title} {hint}
'
def _error(msg):
return _card("err", f'❌ {html.escape(str(msg))}
')
def _busy(msg, note=""):
note = f'{html.escape(note)}
' if note else ""
return _card("busy", f'{html.escape(msg)}'
f'
{note}
{value}{html.escape(label)}
')
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", '%d LoRA files in this repo — '
'pick the one you want above.
' % 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'{html.escape(name)}
'
f'Nothing to do — no {label} layers in this LoRA '
f'({total} layers, all kept). It should already load in H3.
')
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'→ {H.human_size(summary["output_size"])}',
"size after")
)
body = f'{html.escape(name)}
{stats}
'
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'Nothing to do — no {html.escape(pattern)}'
f' layers in {html.escape(name)}.
')
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'→ {H.human_size(size)}', "size")
+ _stat(f"{len(meta)}" if meta else "—", "metadata entries kept")
)
body = (f'✅ Pruned {html.escape(os.path.basename(out_path))}
'
f'{stats}
'
f'Verified: the output re-opens cleanly and every kept layer is '
f'byte-identical to the source.
')
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(
'✂️ H3 LoRA Pruner
'
'
Strips the adaln_proj layers out of a LoRA so it loads in H3.