259 lines
10 KiB
Python
259 lines
10 KiB
Python
"""Frontend user-preference persistence helpers for the Gradio UI.
|
|
|
|
Save side: A ``<script>`` injected via ``Blocks(head=…)`` listens for DOM
|
|
changes and writes the current preference values to ``localStorage``.
|
|
|
|
Restore side: ``wire_preference_restore`` attaches a ``demo.load()`` handler
|
|
whose *js* parameter reads ``localStorage`` on page load and feeds the saved
|
|
values straight into the Gradio component outputs. Because Gradio itself
|
|
applies the updates through its own Svelte reactivity, every component type
|
|
(dropdown, slider, checkbox, number) is updated correctly—no fragile DOM
|
|
hacking required.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from functools import partial
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
_ASSET_FILENAME = "user_preferences.js"
|
|
_STORAGE_KEY = "acestep.ui.user_preferences"
|
|
_SCHEMA_VERSION = 2
|
|
|
|
# Ordered list of preference keys. The order here MUST match the order of
|
|
# *outputs* passed to ``demo.load()`` in ``wire_preference_restore``.
|
|
PREF_KEYS: list[str] = [
|
|
"audio_format",
|
|
"mp3_bitrate",
|
|
"mp3_sample_rate",
|
|
"score_scale",
|
|
"enable_normalization",
|
|
"normalization_db",
|
|
"fade_in_duration",
|
|
"fade_out_duration",
|
|
"latent_shift",
|
|
"latent_rescale",
|
|
"lm_batch_chunk_size",
|
|
]
|
|
|
|
# Default values used when localStorage is empty or the schema version has
|
|
# changed. Keys must match ``PREF_KEYS``.
|
|
_DEFAULTS: dict[str, Any] = {
|
|
"audio_format": "flac",
|
|
"mp3_bitrate": "320k",
|
|
"mp3_sample_rate": 48000,
|
|
"score_scale": 0.5,
|
|
"enable_normalization": True,
|
|
"normalization_db": -1.0,
|
|
"fade_in_duration": 0.0,
|
|
"fade_out_duration": 0.0,
|
|
"latent_shift": 0.0,
|
|
"latent_rescale": 1.0,
|
|
"lm_batch_chunk_size": 8,
|
|
}
|
|
|
|
|
|
# ── Save-side: head script injection ────────────────────────────────────
|
|
|
|
|
|
def _load_preferences_script() -> str:
|
|
"""Load the external save-preferences JavaScript asset."""
|
|
asset_path = Path(__file__).with_name(_ASSET_FILENAME)
|
|
return asset_path.read_text(encoding="utf-8").strip()
|
|
|
|
|
|
def get_user_preferences_head() -> str:
|
|
"""Return Gradio head HTML that injects save-side preference persistence."""
|
|
script_source = _load_preferences_script()
|
|
return f"<script>\n{script_source}\n</script>"
|
|
|
|
|
|
# ── Restore-side: Gradio .load() wiring ─────────────────────────────────
|
|
|
|
|
|
def _build_restore_js(num_outputs: int) -> str:
|
|
"""Build the client-side JS that reads localStorage and returns values.
|
|
|
|
The returned function is passed as the ``js`` parameter to
|
|
``demo.load()``. It returns an array whose element order matches
|
|
``PREF_KEYS`` (and therefore the *outputs* list).
|
|
|
|
When localStorage has no saved preferences (first visit, cleared
|
|
storage, private browsing), the function returns an array of ``null``
|
|
sentinels so the Python side can skip the update and preserve whatever
|
|
values were already rendered from ``init_params``.
|
|
|
|
Args:
|
|
num_outputs: Total number of output components (preference keys
|
|
plus any extra outputs like ``mp3_controls_row``).
|
|
"""
|
|
keys_json = json.dumps(PREF_KEYS)
|
|
# Build a type map so the restore JS can validate each value.
|
|
type_map: dict[str, str] = {}
|
|
for k in PREF_KEYS:
|
|
v = _DEFAULTS[k]
|
|
if isinstance(v, bool):
|
|
type_map[k] = "boolean"
|
|
elif isinstance(v, (int, float)):
|
|
type_map[k] = "number"
|
|
else:
|
|
type_map[k] = "string"
|
|
type_map_json = json.dumps(type_map, ensure_ascii=False)
|
|
# Keys whose Gradio Dropdown choices are integers stored as strings in
|
|
# localStorage. Only actual dropdown keys with numeric defaults need
|
|
# coercion; sliders/numbers are already stored as numbers.
|
|
numeric_dropdown_keys_json = json.dumps(["mp3_sample_rate"])
|
|
# Sentinel array returned when there is nothing to restore. Using null
|
|
# lets the Python fn detect "no stored prefs" and return gr.update()
|
|
# for every output, preserving the values already rendered on the page.
|
|
skip_sentinel = f"new Array({num_outputs}).fill(null)"
|
|
return f"""() => {{
|
|
const STORAGE_KEY = {json.dumps(_STORAGE_KEY)};
|
|
const SCHEMA_VERSION = {_SCHEMA_VERSION};
|
|
const KEYS = {keys_json};
|
|
const TYPE_MAP = {type_map_json};
|
|
const NUMERIC_COERCE_KEYS = new Set({numeric_dropdown_keys_json});
|
|
const SKIP = {skip_sentinel};
|
|
try {{
|
|
const raw = window.localStorage.getItem(STORAGE_KEY);
|
|
if (!raw) return SKIP;
|
|
const prefs = JSON.parse(raw);
|
|
// Only reset on downgrade; forward-compatible additions of new
|
|
// keys are handled by skipping (preserving init_params).
|
|
if (prefs._version !== SCHEMA_VERSION) {{
|
|
return SKIP;
|
|
}}
|
|
const result = KEYS.map(k => {{
|
|
if (!(k in prefs)) return null;
|
|
let v = prefs[k];
|
|
// Type-check: fall back to null (skip) if the stored type
|
|
// does not match what the Gradio component expects.
|
|
const expected = TYPE_MAP[k];
|
|
if (expected && typeof v !== expected) {{
|
|
// Allow stringified numbers for dropdown coercion below.
|
|
if (!(NUMERIC_COERCE_KEYS.has(k) && typeof v === "string")) {{
|
|
return null;
|
|
}}
|
|
}}
|
|
// Coerce stringified numbers back for Dropdown choices that
|
|
// expect integers (e.g. mp3_sample_rate: 48000 not "48000").
|
|
if (NUMERIC_COERCE_KEYS.has(k) && typeof v === "string") {{
|
|
const n = Number(v);
|
|
if (Number.isFinite(n)) v = n;
|
|
else return null;
|
|
}}
|
|
return v;
|
|
}});
|
|
// If none of the keys had stored values, skip entirely.
|
|
if (result.every(v => v === null)) return SKIP;
|
|
// Compute mp3 control visibility from audio_format (index 0).
|
|
// Push 3 extra values: mp3_controls_row, mp3_bitrate, mp3_sample_rate
|
|
// matching the outputs of _update_mp3_control_visibility().
|
|
// When audioFormat is null (no stored value), push nulls so Python
|
|
// emits gr.update() and preserves whatever init_params set.
|
|
const audioFormat = result[0];
|
|
const mp3 = audioFormat === null ? null : audioFormat === "mp3";
|
|
result.push(mp3, mp3, mp3);
|
|
return result;
|
|
}} catch (_e) {{
|
|
return SKIP;
|
|
}}
|
|
}}"""
|
|
|
|
|
|
def restore_preferences(
|
|
*values: Any, _num_outputs: int = 0
|
|
) -> tuple[Any, ...]:
|
|
"""Map JS restore results into Gradio output values.
|
|
|
|
The JS function reads localStorage and produces an array:
|
|
- First ``len(PREF_KEYS)`` elements are preference values (or null).
|
|
- Next 3 elements are mp3 visibility booleans (or null):
|
|
[mp3_controls_row, mp3_bitrate, mp3_sample_rate].
|
|
|
|
``None`` (JSON ``null``) → ``gr.update()`` (no-op, preserves current).
|
|
Booleans beyond PREF_KEYS → visibility/interactivity updates matching
|
|
``_update_mp3_control_visibility()`` from the output controls module.
|
|
|
|
When the JS side returns no values (e.g. certain Gradio versions do not
|
|
forward the JS return value to the Python ``fn`` when ``inputs=None``),
|
|
``_num_outputs`` is used to produce the correct number of no-op updates
|
|
so Gradio does not raise a ``ValueError`` about mismatched output count.
|
|
"""
|
|
import gradio as gr
|
|
|
|
if not values:
|
|
return tuple(gr.update() for _ in range(_num_outputs))
|
|
|
|
n_prefs = len(PREF_KEYS)
|
|
results: list[Any] = []
|
|
for i, v in enumerate(values):
|
|
if v is None:
|
|
results.append(gr.update())
|
|
elif i == n_prefs and isinstance(v, bool):
|
|
# mp3_controls_row: visibility only.
|
|
results.append(gr.update(visible=v))
|
|
elif i > n_prefs and isinstance(v, bool):
|
|
# mp3_bitrate, mp3_sample_rate: visibility + interactivity.
|
|
results.append(gr.update(visible=v, interactive=v))
|
|
else:
|
|
results.append(v)
|
|
return tuple(results)
|
|
|
|
|
|
def wire_preference_restore(
|
|
demo: Any,
|
|
generation_section: dict[str, Any],
|
|
*,
|
|
service_mode: bool = False,
|
|
) -> None:
|
|
"""Attach a ``demo.load()`` handler that restores saved preferences.
|
|
|
|
Must be called **inside** the ``with gr.Blocks() as demo:`` context,
|
|
after all generation components have been created.
|
|
|
|
In service mode the function is a no-op: service-mode sessions use
|
|
server-side ``init_params`` and controls are locked
|
|
(``interactive=False``), so localStorage values must not override them.
|
|
|
|
Args:
|
|
demo: The ``gr.Blocks`` instance.
|
|
generation_section: Merged component dict that includes the output
|
|
control components (``audio_format``, ``mp3_bitrate``, etc.).
|
|
service_mode: When ``True``, skip wiring entirely so that
|
|
localStorage cannot override server-configured values.
|
|
"""
|
|
if service_mode:
|
|
return
|
|
|
|
outputs = []
|
|
for key in PREF_KEYS:
|
|
component = generation_section.get(key)
|
|
if component is None:
|
|
raise KeyError(
|
|
f"wire_preference_restore: missing component {key!r} in "
|
|
f"generation_section (available: {sorted(generation_section)})"
|
|
)
|
|
outputs.append(component)
|
|
|
|
# Also update mp3 control visibility so it stays in sync when the
|
|
# restored audio_format differs from the server-rendered default.
|
|
# Gradio does not fire .change() for load-time value assignments, so
|
|
# without this the MP3 row and its children could be visible/hidden
|
|
# incorrectly. The three extra outputs mirror the return of
|
|
# _update_mp3_control_visibility(): [row, bitrate, sample_rate].
|
|
for mp3_key in ("mp3_controls_row", "mp3_bitrate", "mp3_sample_rate"):
|
|
comp = generation_section.get(mp3_key)
|
|
if comp is not None:
|
|
outputs.append(comp)
|
|
|
|
demo.load(
|
|
fn=partial(restore_preferences, _num_outputs=len(outputs)),
|
|
inputs=None,
|
|
outputs=outputs,
|
|
js=_build_restore_js(num_outputs=len(outputs)),
|
|
)
|