biegen / app.py
Pixelkram's picture
Deploy app
2fd30a7 verified
Raw History Blame Contribute Delete
22.5 kB
"""Gradio app "Biegen": Neural_Bending's tools for Stable Diffusion 1.5, on a map of the network.
The map (bendmap.js) is the control surface: pick tools, choose where they act by clicking the
network, drag their time windows. Every run renders the same seeds without tools too (cached), shows
both images with a wipe, and draws on the map how much each block of the UNet changed.
"""
from __future__ import annotations
import html
import json
import math
import os
import shutil
import tempfile
import threading
from collections import OrderedDict
from dataclasses import replace
os.environ.setdefault("MODEL_FAMILY", "sd15")
import gradio as gr # noqa: E402
import engine as E # noqa: E402
import measure as M # noqa: E402
import safety as S # noqa: E402
import sites as W # noqa: E402
import tools_sd15 as T # noqa: E402
from family import ACTIVE as FAM # noqa: E402
if FAM.key != "sd15":
raise SystemExit("Biegen runs on Stable Diffusion 1.5 only (MODEL_FAMILY=sd15).")
HERE = os.path.dirname(os.path.abspath(__file__))
WORKSHOP_PASSWORD = os.environ.get("WORKSHOP_PASSWORD")
OUT_DIR = os.environ.get("BIEGEN_OUT") or tempfile.mkdtemp(prefix="biegen-")
KEEP_RUNS = 40 # result folders kept on disk, older ones are deleted
PLAIN_CACHE = 12 # renders without tools kept in memory
PROMPT_A = "a photograph of a horse standing in a field"
PROMPT_B = "carved from cracked white marble, veined stone, museum lighting"
GRID_CHOICES = [(f"{k} × {k}", k * k) for k in range(1, 5)]
GRID_DEFAULT = 4
N_SLOTS = 3
SLOT_COLORS = ("#e8590c", "#1c7ed6", "#2f9e44") # orange, blue, green: tell apart in light and dark mode
ICONS = {"excise": "✂", "layered": "◫", "positions": "⇄", "blind": "◐", "erosion": "≋", "temperature": "✱",
"skips": "⤳", "mirror": "⧉", "freeze": "❄", "boil": "♨", "norm": "◩", "vae": "▦", "aphasia": "⌀",
"guidance": "↺", "bypass": "◌", "shakes": "⌇", "time": "⧗", "ghost": "▞"}
TIME_KNOB = ("blind", "freeze") # tools whose knob is a moment in the schedule
# Tools without a window act on the whole run. Why, in the words of the timeline:
WHOLE_RUN = {"aphasia": "ändert den Kontext, der in jedem Schritt derselbe ist"}
ENVELOPE = {"shakes": "sine"} # strength inside the window: half a sine, 0 at both ends
# Typed numbers may go past every slider. These limits stay.
HARD_LIMITS = {"seed": ("Seed", -2 ** 62, 2 ** 62), "steps": ("Schritte", 1, 250)}
# One job: the largest grid at the default steps, with and without tools. Nobody holds the queue longer.
MAX_WORK = 2 * max(n for _, n in GRID_CHOICES)
THEME = gr.themes.Soft(
primary_hue="indigo", neutral_hue="slate",
font=[gr.themes.GoogleFont("Atkinson Hyperlegible"), "system-ui", "sans-serif"],
font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"])
# Gradio's slider snaps a typed value back into its range when the number box loses focus.
FREE_SLIDER_JS = """<script>
window.addEventListener("blur", (e) => {
const box = e.target;
if (box instanceof HTMLInputElement && box.type === "number" && box.closest(".free-slider")) {
e.stopImmediatePropagation();
}
}, true);
</script>"""
class FreeSlider(gr.Slider):
"""A slider whose number box takes values beyond the slider's range (with FREE_SLIDER_JS).
No __init__ of its own: Gradio reads the component's settings from the __init__ signature."""
is_template = True
def preprocess(self, payload: float) -> float:
return gr.Number.round_to_precision(payload, self.precision)
def free_slider(minimum: float, maximum: float, **kwargs) -> FreeSlider:
return FreeSlider(minimum, maximum, elem_classes=["free-slider"], **kwargs)
def read(name: str) -> str:
with open(os.path.join(HERE, name), encoding="utf-8") as f:
return f.read()
# ---------------------------------------------------------------- engine
_engine = None
_engine_lock = threading.Lock()
_sec_per_image = None
_tokenizer = None
def get_engine():
global _engine, _sec_per_image
with _engine_lock:
if _engine is None:
_engine = E.load_engine()
print(f"[app] warm-up {_engine.warmup():.1f} s", flush=True)
seconds = _engine.warmup(n=GRID_DEFAULT, steps=E.DEFAULT_STEPS)
_sec_per_image = seconds / GRID_DEFAULT
print(f"[app] {GRID_DEFAULT} images at the defaults in {seconds:.1f} s", flush=True)
return _engine
def tokenizer():
"""CLIP's tokenizer on its own, so the map can show tokens before the model has loaded."""
global _tokenizer
if _tokenizer is None:
from transformers import CLIPTokenizer
_tokenizer = CLIPTokenizer.from_pretrained(E.MODEL_ID, subfolder="tokenizer")
return _tokenizer
def sec_per_image() -> float:
return _sec_per_image or FAM.sec_per_image
def duration(seconds: float) -> str:
return f"etwa {round(seconds / 60)} Minuten" if seconds >= 90 else f"etwa {round(seconds)} Sekunden"
# ---------------------------------------------------------------- the map's catalogue and state
def catalog() -> dict:
"""Everything bendmap.js needs to know about the tools and the network."""
tools = []
for t in T.TOOLS.values():
places = []
for label, value in (t.places or (("", ""),)):
v, start, end = t.defaults_for(value)
places.append({"value": value, "label": label, "sites": list(W.sites(t.key, value)),
"default": v, "start": round(start * 100), "end": round(end * 100)})
tools.append({"key": t.key, "label": t.label, "group": t.group, "what": t.what, "knob": t.knob,
"lo": t.lo, "hi": t.hi, "step": t.step, "neutral": t.neutral, "preset": t.preset,
"window": t.window is not None, "time_knob": t.key in TIME_KNOB,
"whole": WHOLE_RUN.get(t.key, "ändert Gewichte"), "envelope": ENVELOPE.get(t.key, ""),
"after_loop": t.key == "vae", "icon": ICONS.get(t.key, "•"), "places": places})
return {"tools": tools, "groups": list(dict.fromkeys(t.group for t in T.TOOLS.values())),
"colors": list(SLOT_COLORS), "max_slots": N_SLOTS}
def default_slot(key: str, place: str = "", word: str = "") -> dict:
t = T.TOOLS[key]
place = place or W.places(key)[0]
value, start, end = t.defaults_for(place)
return {"tool": key, "place": place, "value": value, "start": round(start * 100), "end": round(end * 100),
"word": word}
DEFAULT_STATE = {"slots": [default_slot("layered", "up")], "rev": 0}
def tokens(prompt: str) -> list[str]:
"""CLIP's tokens of a prompt as the map shows them, without start, end and padding."""
tok = tokenizer()
ids = tok(prompt or "", add_special_tokens=False, truncation=True,
max_length=tok.model_max_length - 2).input_ids
return [tok.convert_ids_to_tokens(i).replace("</w>", "") for i in ids]
def server_part(prompt_a: str, prompt_b: str) -> dict:
return {"prompt": prompt_a or "", "prompt_b": prompt_b or "", "tokens": tokens(prompt_a),
"tokens_b": tokens(prompt_b), "max_len": tokenizer().model_max_length}
def map_value(server: dict, state: dict | None = None) -> str:
out = {"server": server}
if state is not None:
out["state"] = state
return json.dumps(out, ensure_ascii=False)
def parse_state(value: str | None) -> dict:
"""The map's state from its JSON value. Anything malformed counts as no tools."""
try:
data = json.loads(value) if value else {}
except (TypeError, ValueError):
data = {}
state = data.get("state") if isinstance(data, dict) else None
if not isinstance(state, dict):
return {"slots": [], "rev": 0}
slots = [s for s in state.get("slots") or [] if isinstance(s, dict)][:N_SLOTS]
rev = state.get("rev", 0)
return {"slots": slots, "rev": rev if isinstance(rev, int) else 0}
def number(x, default: float) -> float:
try:
x = float(x)
except (TypeError, ValueError):
return default
return x if math.isfinite(x) else default
def tool_uses(slots: list[dict]) -> tuple[tuple, list[str]]:
"""ToolUse entries from the map's slots. Windows are clamped to 0..100 % with a note."""
uses, notes = [], []
for k, v in enumerate(slots):
key = v.get("tool")
if key not in T.TOOLS:
continue
t = T.TOOLS[key]
start, end = number(v.get("start"), 0.0) / 100, number(v.get("end"), 100.0) / 100
if t.window is not None and not (0 <= start <= 1 and 0 <= end <= 1):
notes.append(f"Eingriff {k + 1}: das Zeitfenster geht nur von 0 % bis 100 %, "
"gerechnet wird mit dem Teil darin.")
start, end = min(max(start, 0.0), 1.0), min(max(end, 0.0), 1.0)
place = v.get("place") if v.get("place") in W.places(key) else W.places(key)[0]
uses.append(T.ToolUse(key, number(v.get("value"), t.neutral), place, start, end,
str(v.get("word") or "")[:80]))
return tuple(uses), notes
def recipe(u: T.ToolUse) -> str:
"""One line per tool, to copy and to compare runs."""
t = T.TOOLS[u.key]
parts = [f'"{t.label}"']
if T.place_label(u):
parts.append(f'ort="{T.place_label(u)}"')
if u.key == "excise":
parts.append(f'wort="{u.word.strip()}"')
parts.append(f"{t.knob.split(' (')[0].lower()}={E.fmt_num(u.value, 2).rstrip('0').rstrip(',')}")
if t.window is not None:
parts.append(f"wirkt={round(u.start * 100)}–{round(u.end * 100)} %")
return f"biegen({', '.join(parts)})"
# ---------------------------------------------------------------- settings and budget
def settings(prompt_a, prompt_b, seed, steps, guidance, grid, uses) -> E.Settings:
return E.Settings(prompt_a or "", prompt_b or "", "seed_grid", seed=int(number(seed, 42)),
steps=int(round(number(steps, E.DEFAULT_STEPS))), guidance=number(guidance, FAM.g_default),
grid=int(grid or GRID_DEFAULT), te_layers=(12,), tools=uses)
def within_limits(s: E.Settings) -> tuple[E.Settings, list[str]]:
changes, notes = {}, []
for name, (label, lo, hi) in HARD_LIMITS.items():
value = getattr(s, name)
if not lo <= value <= hi:
changes[name] = min(max(value, lo), hi)
notes.append(f"{label}: möglich sind {lo} bis {hi}, gerechnet wird mit {changes[name]}.")
return replace(s, **changes), notes
def work(s: E.Settings) -> float:
"""GPU work of one render in images at the default steps."""
return (s.grid or GRID_DEFAULT) * s.steps / E.DEFAULT_STEPS
def too_much_work(s: E.Settings, renders: int) -> str | None:
w = work(s) * renders
if w <= MAX_WORK:
return None
return (f"Das ist zu viel auf einmal: Es würde {duration(w * sec_per_image())} dauern, und alle anderen "
"müssten warten. Nehmt ein kleineres Raster oder weniger Schritte.")
def plain_key(s: E.Settings) -> E.Settings:
"""Renders without tools do not depend on prompt B."""
return replace(s, tools=(), prompt_b="")
# ---------------------------------------------------------------- rendering
class Shot:
"""One render: image files, captions and the recorded activations."""
def __init__(self, paths: list[str], captions: list[str], images: list, acts: dict, seconds: float,
warnings: list[str]):
self.paths, self.captions, self.images, self.acts = paths, captions, images, acts
self.seconds, self.warnings = seconds, warnings
_plain: OrderedDict = OrderedDict()
_run_counter = 0
_run_lock = threading.Lock()
def run_dir() -> str:
"""A fresh folder for one run's images. Old folders go."""
global _run_counter
with _run_lock:
_run_counter += 1
path = os.path.join(OUT_DIR, f"run-{_run_counter:05d}")
os.makedirs(path, exist_ok=True)
stale = sorted(d for d in os.listdir(OUT_DIR) if d.startswith("run-"))[:-KEEP_RUNS]
keep = {os.path.dirname(p) for shot in _plain.values() for p in shot.paths[:1]}
for d in stale:
if os.path.join(OUT_DIR, d) not in keep:
shutil.rmtree(os.path.join(OUT_DIR, d), ignore_errors=True)
return path
def shoot(s: E.Settings, folder: str, tag: str, report) -> Shot:
eng = get_engine()
rec = M.Recorder(s.steps)
try:
with M.recording(eng, rec):
res = eng.render(s, E.sample_specs(s), max_seq=77, progress=report)
except gr.Error:
raise
except ValueError as exc:
raise gr.Error(f"Diese Einstellung geht nicht: {exc}")
items, hidden = S.filter_items(res.items)
if hidden:
res.warnings.append(f"Der Filter hat {hidden} von {len(items)} Bildern ausgeblendet.")
paths = []
for k, (im, _) in enumerate(items):
p = os.path.join(folder, f"{tag}-{k:02d}.jpg")
im.convert("RGB").save(p, quality=90)
paths.append(p)
return Shot(paths, [c for _, c in items], [im for im, _ in items], rec.tensors(), res.seconds, res.warnings)
def image_change(a, b) -> float:
"""Mean absolute pixel difference, 0 (same) to 1 (black against white)."""
import numpy as np
x = np.asarray(a.convert("RGB"), dtype=np.float32)
y = np.asarray(b.convert("RGB"), dtype=np.float32)
return float(np.abs(x - y).mean() / 255.0) if x.shape == y.shape else 1.0
def file_url(path: str) -> str:
return "/gradio_api/file=" + path
def compare_html(plain: Shot, bent: Shot | None, changes: list[float]) -> str:
cols = math.ceil(math.sqrt(len(plain.paths)))
cards = []
for k, p in enumerate(plain.paths):
cap = html.escape(plain.captions[k])
if bent is None:
cards.append(f'<figure class="bc-card"><img src="{file_url(p)}" alt="{cap}" loading="lazy">'
f"<figcaption>{cap}</figcaption></figure>")
continue
d = f"Δ {round(changes[k] * 100)} %"
cards.append(
f'<figure class="bc-card" data-both="1"><div class="bc-frame">'
f'<img class="bc-bent" src="{file_url(bent.paths[k])}" alt="{cap}, gebogen">'
f'<img class="bc-plain" src="{file_url(p)}" alt="{cap}, ohne Eingriff">'
f'<div class="bc-line"></div><span class="bc-tag bc-tag-plain">ohne</span>'
f'<span class="bc-tag bc-tag-bent">gebogen</span></div>'
f'<figcaption>{cap} <b title="mittlere Pixeländerung gegenüber dem Bild ohne Eingriff">{d}</b>'
f"</figcaption></figure>")
hint = ("Fahrt mit der Maus über ein Bild: links vom Strich seht ihr dasselbe Bild ohne Eingriff."
if bent is not None else "Kein Eingriff aktiv: das sind die Bilder ohne Eingriff.")
toggle = '<button type="button" class="bc-toggle">Alle ohne Eingriff zeigen</button>' if bent else ""
return (f'<div class="bc"><div class="bc-bar"><span>{hint}</span>{toggle}</div>'
f'<div class="bc-grid" style="grid-template-columns: repeat({cols}, minmax(0, 1fr))">'
f'{"".join(cards)}</div></div>')
EMPTY_RESULT = ('<div class="bc bc-empty">Hier erscheinen die Bilder: jedes einmal gebogen und einmal ohne '
"Eingriff, zum Vergleichen.</div>")
def run(prompt_a, prompt_b, seed, steps, guidance, grid, map_json, progress=gr.Progress()):
state = parse_state(map_json)
print(f"[run] rev {state['rev']}, tools {[v.get('tool') for v in state['slots']]}", flush=True)
uses, notes = tool_uses(state["slots"])
s, limit_notes = within_limits(settings(prompt_a, prompt_b, seed, steps, guidance, grid, uses))
notes = limit_notes + notes
live = [u for u in uses if T.active(u)]
key = plain_key(s)
cached = _plain.get(key)
renders = (cached is None) + bool(live)
refusal = too_much_work(s, renders)
if refusal:
gr.Warning(refusal)
return gr.update(), gr.update(), refusal
total = renders * s.steps * math.ceil(len(E.sample_specs(s)) / max(1, E.MAX_BATCH))
done = [0]
def report(i: int, n: int) -> None:
done[0] += 1
progress(min(done[0] / max(1, total), 1.0), desc="ohne Eingriff" if cached is None and done[0] <= total / renders
else "gebogen")
folder = run_dir()
if cached is None:
cached = shoot(key, folder, "plain", report)
_plain[key] = cached
while len(_plain) > PLAIN_CACHE:
_plain.popitem(last=False)
else:
_plain.move_to_end(key)
bent = shoot(s, folder, "bent", report) if live else None
server = server_part(prompt_a, prompt_b)
if bent is not None:
changes = [image_change(a, b) for a, b in zip(cached.images, bent.images)]
eff = M.effect(cached.acts, bent.acts)
eff.update(image=changes, thumb=file_url(bent.paths[0]), plain_thumb=file_url(cached.paths[0]),
rev=state["rev"])
server["effect"] = eff
seconds = bent.seconds + (cached.seconds if renders == 2 else 0.0)
lines = [f"{len(bent.paths)} Bilder gebogen in {E.fmt_num(seconds, 1)} s"
+ ("." if renders == 2 else ", die Bilder ohne Eingriff kamen aus dem Speicher.")]
lines += notes + bent.warnings
lines.append("```\n" + "\n".join(recipe(u) for u in live) + "\n```")
else:
changes = []
server["effect"] = {"rev": state["rev"], "steps": 0, "blocks": {}, "latent": [], "impact": {},
"first": None, "image": [], "thumb": file_url(cached.paths[0]), "none": True}
lines = [f"{len(cached.paths)} Bilder ohne Eingriff"
+ (f" in {E.fmt_num(cached.seconds, 1)} s." if renders else " aus dem Speicher.")]
lines += notes + cached.warnings
if uses:
lines.append("Die Werkzeuge stehen auf dem Wert, bei dem sie nichts ändern.")
return map_value(server), compare_html(cached, bent, changes), "\n\n".join(lines)
def show_tokens(prompt_a, prompt_b):
return map_value(server_part(prompt_a, prompt_b))
# ---------------------------------------------------------------- the app
INTRO = """# Biegen
Stable Diffusion 1.5 von innen, von oben nach unten: Text, UNet, VAE-Decoder. Wählt rechts ein Werkzeug: Auf der
Karte leuchtet auf, wo es im Netz eingreift. Klickt auf die Karte, um den Ort zu wechseln, und zieht auf der
Zeitleiste, wann es wirkt. Nach jedem Lauf zeigt die Karte, wie stark sich jeder Block gegenüber dem Lauf ohne
Eingriff verändert hat."""
LEGEND = """**So liest man die Karte.** Oben wird der Text zu 77 Vektoren (Kontext). Darunter das UNet in seiner
U-Form: der Abwärtsweg verkleinert das Bild von 64 × 64 auf 8 × 8 Stellen, der Aufwärtsweg baut es wieder auf, die
Skip-Verbindungen tragen Einzelheiten quer hinüber. In jedem Block rechnen **Res** (Faltungen, hier hört das Netz
auch den Zeitpunkt t), **Selbst** (jede Bildstelle schaut auf alle anderen) und **Quer** (jede Bildstelle schaut auf
den Kontext; nur hier kommt der Text ins Bild). Um das UNet läuft die Schleife: das UNet sagt das Rauschen voraus,
einmal mit und einmal ohne Prompt, die **Guidance** verstärkt den Unterschied, der Sampler zieht das Rauschen ab, und
das Latent geht mit dem nächsten **Zeitpunkt t** in den nächsten Schritt. Nach dem letzten Schritt übersetzt unten
der VAE-Decoder das Latent in Pixel. Ganz unten die Zeitleiste von Rauschen bis Bild.
**Was gemessen wird.** Jeder Lauf rechnet dieselben Seeds auch ohne Eingriff. Die Füllung eines Blocks zeigt, wie
stark sich sein Ausgang im ersten veränderten Schritt unterscheidet: Dort ist das Latent noch gleich, also ändern
sich nur der Ort des Eingriffs und alles danach. Die Linie im Block zeigt dieselbe Änderung über alle Schritte, die
Linie über der Zeitleiste die Änderung des Latents. Gemessen wird am ersten Bild, als relative Änderung."""
def biegen_app() -> gr.Blocks:
js = read("bendmap.js").replace("__CATALOG__", json.dumps(catalog(), ensure_ascii=False))
with gr.Blocks(title="Biegen") as demo:
gr.Markdown(INTRO, elem_id="intro")
with gr.Row():
prompt_a = gr.Textbox(PROMPT_A, label="Prompt A", max_lines=2, elem_id="prompt-a",
info="CLIP liest Englisch.")
prompt_b = gr.Textbox(PROMPT_B, label="Prompt B", max_lines=2, elem_id="prompt-b",
info="Nur für „Zweiter Prompt für die Oberfläche“.")
bendmap = gr.HTML(map_value(server_part(PROMPT_A, PROMPT_B), DEFAULT_STATE),
html_template='<div class="bm-host"></div>', js_on_load=js, apply_default_css=False,
elem_id="bendmap")
with gr.Row():
seed = gr.Number(value=42, precision=0, label="Seed", scale=1, min_width=120)
steps = free_slider(4, 50, value=E.DEFAULT_STEPS, step=1, label="Schritte", scale=2)
guidance = free_slider(1.0, 15.0, value=FAM.g_default, step=0.1, label="Guidance", scale=2)
grid = gr.Radio(GRID_CHOICES, value=GRID_DEFAULT, label="Raster", scale=3)
with gr.Column(scale=1, min_width=160):
go = gr.Button("Biegen", variant="primary", size="lg", elem_id="go")
status = gr.Markdown(elem_id="status")
result = gr.HTML(EMPTY_RESULT, js_on_load=read("compare.js"), apply_default_css=False, elem_id="result")
with gr.Accordion("So liest man die Karte", open=False):
gr.Markdown(LEGEND)
for box in (prompt_a, prompt_b):
box.change(show_tokens, [prompt_a, prompt_b], bendmap, queue=False, show_progress="hidden",
trigger_mode="always_last", api_visibility="private")
go.click(run, [prompt_a, prompt_b, seed, steps, guidance, grid, bendmap], [bendmap, result, status],
api_name="biegen", show_progress_on=[result])
return demo
def launch(**kwargs) -> gr.Blocks:
get_engine()
if S.SAFETY_FILTER:
S.classifier()
demo = biegen_app()
demo.queue(max_size=64)
demo.launch(theme=THEME, css=read("bendmap.css"), head=FREE_SLIDER_JS, allowed_paths=[OUT_DIR],
auth=("workshop", WORKSHOP_PASSWORD) if WORKSHOP_PASSWORD else None,
auth_message="Meldet euch mit dem Benutzernamen workshop an. Das Passwort steht an der Tafel.",
**kwargs)
return demo
if __name__ == "__main__":
launch()