Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from schwarz_demo import DemoConfig, architecture_figure, preview, train, training_figure, operators_figure, solution_figure, summary | |
| def cfg(n, o, lr, cr, local, coarse, glob, assembly, steps): | |
| return DemoConfig( | |
| int(n), int(o), int(lr), int(cr), bool(local), bool(coarse), bool(glob), | |
| assembly == "RAS-style output weighting (experimental)", int(steps) | |
| ) | |
| def update(*args): | |
| try: | |
| c = cfg(*args) | |
| return preview(c), architecture_figure(c) | |
| except Exception as e: | |
| return f"⚠️ {e}", None | |
| def run(*args, progress=gr.Progress()): | |
| c = cfg(*args) | |
| r = train(c, lambda s, n: progress(s / n, desc=f"Training {s}/{n}")) | |
| return training_figure(r), operators_figure(r), solution_figure(r), summary(r) | |
| with gr.Blocks(title="Overlapping Schwarz Attention") as demo: | |
| gr.Markdown( | |
| "# 🧩 Overlapping Schwarz Attention Playground: [https://arxiv.org/abs/2606.18525](https://arxiv.org/abs/2606.18525)\n" | |
| "Explore **local-only**, **coarse-only**, and **two-level** attention, " | |
| "with an optional nearest-parameter global low-rank baseline." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("### Schwarz components") | |
| local = gr.Checkbox(True, label="Local overlapping attention blocks") | |
| coarse = gr.Checkbox(True, label="Coarse attention block") | |
| glob = gr.Checkbox(True, label="Train nearest-parameter global low-rank baseline") | |
| n = gr.Slider(2, 16, 8, step=1, label="Subdomains") | |
| o = gr.Slider(0, 8, 2, step=1, label="Overlap") | |
| lr = gr.Slider(1, 24, 4, step=1, label="Local attention rank") | |
| cr = gr.Slider(1, 16, 7, step=1, label="Coarse attention rank") | |
| steps = gr.Dropdown([500, 1000, 2000], value=500, label="Training steps") | |
| with gr.Accordion("Advanced options", open=False): | |
| assembly = gr.Radio( | |
| [ | |
| "Symmetric partition-of-unity weighting (paper default)", | |
| "RAS-style output weighting (experimental)", | |
| ], | |
| value="Symmetric partition-of-unity weighting (paper default)", | |
| label="Local assembly", | |
| ) | |
| gr.Markdown( | |
| "The RAS-style option is an exploratory comparison. " | |
| "The paper uses symmetric partition-of-unity weighting." | |
| ) | |
| button = gr.Button("Train and compare", variant="primary") | |
| with gr.Column(scale=2): | |
| info = gr.Markdown() | |
| arch = gr.Plot() | |
| gr.Markdown("### Training and operators") | |
| curve = gr.Plot() | |
| ops = gr.Plot() | |
| sol = gr.Plot() | |
| gr.Markdown("### Results and diagnostics") | |
| table = gr.Markdown() | |
| inputs = [n, o, lr, cr, local, coarse, glob, assembly, steps] | |
| for x in inputs: | |
| x.change(update, inputs, [info, arch]) | |
| button.click(run, inputs, [curve, ops, sol, table]) | |
| demo.load(update, inputs, [info, arch]) | |
| if __name__ == "__main__": | |
| demo.queue(default_concurrency_limit=1).launch() | |