Spaces:
Sleeping
Sleeping
| """Gradio entry-point wiring the patch-view service layer.""" | |
| from __future__ import annotations | |
| import os | |
| import socket | |
| from typing import Dict, List | |
| import gradio as gr | |
| from .constraints import ConstraintDefinition | |
| from .services import PatchService | |
| def create_app(service: PatchService | None = None) -> gr.Blocks: | |
| """Create and configure the Gradio application.""" | |
| patch_service = service or PatchService.from_data_root() | |
| audio_choices = patch_service.audio_ids() | |
| with gr.Blocks(title="BLT Patching Visualization", theme=gr.themes.Default()) as app: | |
| gr.Markdown("# BLT Patching Visualization Explorer") | |
| gr.Markdown("Explore entropy-based patching strategies for Byte Latent Transformers") | |
| patches_state = gr.State([]) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("## Audio Selection") | |
| audio_dropdown = gr.Dropdown( | |
| choices=audio_choices, | |
| value=audio_choices[0] if audio_choices else None, | |
| label="Audio File", | |
| interactive=True, | |
| ) | |
| file_info = gr.Markdown() | |
| gr.Markdown("## Patching Parameters") | |
| metric_radio = gr.Radio( | |
| choices=["Entropy", "Perplexity"], | |
| value="Entropy", | |
| label="Metric", | |
| info="Underlying metric for patching decisions", | |
| ) | |
| gr.Markdown("**Constraints** (select one or more)") | |
| gr.Markdown("_Patch boundary created when ANY constraint is satisfied (OR logic)_") | |
| constraint_defs: List[ConstraintDefinition] = list( | |
| patch_service.constraint_definitions() | |
| ) | |
| constraint_choices = [(definition.label, definition.name) for definition in constraint_defs] | |
| constraint_selector = gr.CheckboxGroup( | |
| choices=constraint_choices, | |
| value=["global"], | |
| label="Active Constraints", | |
| info="Select multiple constraints to combine with OR logic", | |
| ) | |
| guard_config = patch_service.guard_slider_config("entropy") | |
| min_metric_slider = gr.Slider( | |
| value=guard_config.default, | |
| minimum=guard_config.minimum, | |
| maximum=guard_config.maximum, | |
| step=guard_config.step, | |
| label=guard_config.label, | |
| info=guard_config.description, | |
| ) | |
| gr.Markdown("_Guard constraint uses AND logic with all selected constraints_") | |
| threshold_inputs: Dict[str, gr.Slider] = {} | |
| for definition in constraint_defs: | |
| range_cfg = definition.range_for_metric("entropy") | |
| threshold_inputs[definition.name] = gr.Slider( | |
| value=range_cfg.default, | |
| minimum=range_cfg.minimum, | |
| maximum=range_cfg.maximum, | |
| step=range_cfg.step, | |
| label=f"{definition.label} Threshold", | |
| info=definition.description, | |
| visible=(definition.name == "global"), | |
| ) | |
| threshold_names = list(threshold_inputs.keys()) | |
| apply_btn = gr.Button("Apply Patching", variant="primary", size="lg") | |
| with gr.Column(scale=1): | |
| stats_output = gr.Markdown("Click 'Apply Patching' to see statistics") | |
| gr.Markdown("## Visualization") | |
| plot_output = gr.Plot(label="Patch Visualization") | |
| gr.Markdown("## Audio Playback & Comparison") | |
| with gr.Tab("Original vs Reconstructed"): | |
| gr.Markdown( | |
| "**Compare original audio with WavTokenizer reconstruction** " | |
| "(listen for compression artifacts)" | |
| ) | |
| with gr.Row(): | |
| original_audio = gr.Audio(label="Original Audio", type="numpy") | |
| reconstructed_audio = gr.Audio(label="WavTokenizer Reconstructed", type="numpy") | |
| with gr.Tab("Play Until Boundary"): | |
| gr.Markdown( | |
| "**Play audio from a token position until the next patch boundary** " | |
| "(hear what's in a single patch)" | |
| ) | |
| with gr.Row(): | |
| current_token_input = gr.Number( | |
| value=0, | |
| label="Token Position", | |
| precision=0, | |
| info="Enter token index (0 to num_tokens-1)", | |
| ) | |
| play_boundary_btn = gr.Button("▶ Play Until Next Boundary", variant="secondary") | |
| patch_info_output = gr.Markdown() | |
| boundary_audio = gr.Audio(label="Patch Audio", type="numpy", autoplay=True) | |
| # ------------------------------------------------------------------ | |
| # Callbacks | |
| # ------------------------------------------------------------------ | |
| def _update_file_info(file_id: str): | |
| return patch_service.file_info_markdown(file_id) | |
| def _update_threshold_visibility(selected: List[str]): | |
| return { | |
| threshold_inputs[name]: gr.update(visible=(name in selected)) | |
| for name in threshold_inputs | |
| } | |
| def _update_threshold_ranges(metric_name: str): | |
| metric_defaults = patch_service.threshold_defaults(metric_name) | |
| guard = patch_service.guard_slider_config(metric_name) | |
| updates = { | |
| min_metric_slider: gr.update( | |
| value=guard.default, | |
| minimum=guard.minimum, | |
| maximum=guard.maximum, | |
| step=guard.step, | |
| label=guard.label, | |
| info=guard.description, | |
| ) | |
| } | |
| for definition in constraint_defs: | |
| params = definition.range_for_metric(metric_name) | |
| updates[threshold_inputs[definition.name]] = gr.update( | |
| value=metric_defaults[definition.name], | |
| minimum=params.minimum, | |
| maximum=params.maximum, | |
| step=params.step, | |
| ) | |
| return updates | |
| def _handle_patching(file_id, metric_name, selected_constraints, min_threshold, *slider_values): | |
| if not file_id: | |
| return None, "No file selected", None, None, [] | |
| thresholds = {name: float(slider_values[idx]) for idx, name in enumerate(threshold_names)} | |
| try: | |
| return patch_service.run_patching( | |
| file_id=file_id, | |
| metric_name=metric_name, | |
| selected_constraints=selected_constraints, | |
| guard_threshold=min_threshold, | |
| thresholds=thresholds, | |
| ) | |
| except ValueError as exc: | |
| return None, f"⚠️ {exc}", None, None, [] | |
| def _handle_play(file_id, current_token, patches): | |
| if not file_id: | |
| return None, "No file selected", 0 | |
| if not patches: | |
| return None, "No patches computed yet - click 'Apply Patching' first", current_token | |
| try: | |
| return patch_service.play_until_boundary(file_id, int(current_token), patches) | |
| except ValueError as exc: | |
| return None, f"⚠️ {exc}", current_token | |
| audio_dropdown.change(fn=_update_file_info, inputs=[audio_dropdown], outputs=[file_info]) | |
| constraint_selector.change( | |
| fn=_update_threshold_visibility, | |
| inputs=[constraint_selector], | |
| outputs=list(threshold_inputs.values()), | |
| ) | |
| metric_radio.change( | |
| fn=_update_threshold_ranges, | |
| inputs=[metric_radio], | |
| outputs=[min_metric_slider, *threshold_inputs.values()], | |
| ) | |
| apply_btn.click( | |
| fn=_handle_patching, | |
| inputs=[ | |
| audio_dropdown, | |
| metric_radio, | |
| constraint_selector, | |
| min_metric_slider, | |
| *threshold_inputs.values(), | |
| ], | |
| outputs=[plot_output, stats_output, original_audio, reconstructed_audio, patches_state], | |
| ) | |
| play_boundary_btn.click( | |
| fn=_handle_play, | |
| inputs=[audio_dropdown, current_token_input, patches_state], | |
| outputs=[boundary_audio, patch_info_output, current_token_input], | |
| ) | |
| app.load(fn=_update_file_info, inputs=[audio_dropdown], outputs=[file_info]) | |
| return app | |
| def _find_open_port(preferred: int, attempts: int = 50) -> int: | |
| """Find an available TCP port starting from a preferred value.""" | |
| if preferred <= 0: | |
| with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: | |
| try: | |
| sock.bind(("127.0.0.1", 0)) | |
| return sock.getsockname()[1] | |
| except OSError: | |
| return 7860 | |
| port = preferred | |
| for _ in range(attempts): | |
| with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: | |
| try: | |
| sock.bind(("127.0.0.1", port)) | |
| except OSError: | |
| port += 1 | |
| continue | |
| return port | |
| raise OSError(f"Could not find open port starting at {preferred}") | |
| def main() -> None: | |
| """CLI entry point for running the Gradio app.""" | |
| app = create_app() | |
| port_env = os.getenv("GRADIO_SERVER_PORT") | |
| try: | |
| requested_port = int(port_env) if port_env is not None else 7860 | |
| except ValueError: | |
| requested_port = 7860 | |
| try: | |
| port = _find_open_port(requested_port) | |
| except OSError as exc: | |
| raise SystemExit(f"Failed to locate an open port near {requested_port}: {exc}") from exc | |
| launch_kwargs = dict(server_name="0.0.0.0", share=True, debug=True, server_port=port) | |
| print(f"Open in browser: http://localhost:{port}") | |
| app.launch(**launch_kwargs) | |
| # Gradio Spaces expects a variable named `demo` | |
| demo = create_app() | |
| if __name__ == "__main__": | |
| main() | |