Persist ACE-Step XL-SFT quality defaults

experiments/acestep15-xl-sft/overrides/user_preferences.py

@@ -0,0 +1,258 @@
+"""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)),
+    )