"""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()