PaperBanana
Multi-Agent
Scientific Diagrams
Settings
') pipeline_mode = gr.Dropdown( choices=["demo_planner_critic", "demo_full"], value="demo_full", label="Pipeline Mode", info="Select which agent pipeline to use", ) pipeline_desc = gr.Textbox( label="Pipeline Description", value=PIPELINE_DESCRIPTIONS["demo_full"], interactive=False, lines=2, ) pipeline_mode.change( lambda m: PIPELINE_DESCRIPTIONS.get(m, ""), inputs=[pipeline_mode], outputs=[pipeline_desc], ) retrieval_setting = gr.Dropdown( choices=["auto", "manual", "random", "none"], value="auto", label="Retrieval Setting", info="How to retrieve reference diagrams", ) num_candidates = gr.Number( value=10, minimum=1, maximum=20, step=1, label="Number of Candidates", ) aspect_ratio = gr.Dropdown( choices=["16:9", "21:9", "3:2"], value="21:9", label="Aspect Ratio", ) figure_size = gr.Dropdown( choices=["1-3cm", "4-6cm", "7-9cm", "10-13cm", "14-17cm"], value="7-9cm", label="Figure Size", ) max_critic_rounds = gr.Slider( minimum=1, maximum=5, value=3, step=1, label="Max Critic Rounds", ) main_model_name = gr.Textbox( label="Model Name", info="Model name to use for reasoning", value=default_main_model, ) image_model_name = gr.Textbox( label="Image Generation Model", info="Model for generating diagram images", value=default_image_model, ) save_results = gr.Dropdown( choices=["Yes", "No"], value="Yes", label="Save Results", ) # ---------- RIGHT COLUMN: INPUT + OUTPUT ---------- with gr.Column(scale=3): gr.HTML('Input
') with gr.Row(): method_example = gr.Dropdown( choices=["None", "PaperBanana Framework"], value="PaperBanana Framework", label="Load Example (Method)", ) caption_example = gr.Dropdown( choices=["None", "PaperBanana Framework"], value="PaperBanana Framework", label="Load Example (Caption)", ) with gr.Row(): method_content = gr.Textbox( label="Method Content", value=EXAMPLE_METHOD, lines=12, max_lines=30, ) caption_input = gr.Textbox( label="Figure Caption", value=EXAMPLE_CAPTION, lines=12, max_lines=30, ) # Wire example selectors def load_method_example(choice): return EXAMPLE_METHOD if choice == "PaperBanana Framework" else "" def load_caption_example(choice): return EXAMPLE_CAPTION if choice == "PaperBanana Framework" else "" method_example.change(load_method_example, inputs=[method_example], outputs=[method_content]) caption_example.change(load_caption_example, inputs=[caption_example], outputs=[caption_input]) generate_btn = gr.Button( "✨ Generate Candidates", variant="primary", elem_classes=["orange-btn"], size="lg", ) # ---- Status ---- status_text = gr.Textbox(label="Status", interactive=False, lines=1) # ---- Results ---- gr.HTML('Generated Candidates
') results_gallery = gr.Gallery( label="Generated Candidates", columns=3, height="auto", object_fit="contain", ) with gr.Accordion("Evolution Timeline", open=False): evolution_html = gr.HTML("") with gr.Accordion("Download All (ZIP)", open=False): zip_file_output = gr.File(label="ZIP download") # ---- Generate handler ---- def run_generate( method_text, caption_text, pipe_mode, ret_setting, n_cands, ar, max_rounds, m_model, img_model, figure_size, save_results, progress=gr.Progress(track_tqdm=True), ): if not method_text or not caption_text: raise gr.Error("Please provide both method content and caption.") n_cands = int(n_cands) max_rounds = int(max_rounds) timestamp_str = datetime.now().strftime("%Y%m%d_%H%M%S") progress(0, desc="Preparing inputs...") input_data = create_sample_inputs( method_content=method_text, caption=caption_text, aspect_ratio=ar, num_copies=n_cands, max_critic_rounds=max_rounds, ) params = {"figure_size": figure_size} progress(0.1, desc=f"Generating {n_cands} candidates in parallel...") try: loop = asyncio.new_event_loop() results = loop.run_until_complete( process_parallel_candidates( input_data, exp_mode=pipe_mode, retrieval_setting=ret_setting, main_model_name=m_model, image_gen_model_name=img_model, ) ) loop.close() except Exception as e: raise gr.Error(f"Generation failed: {e}") progress(0.9, desc="Saving results...") # Save JSON results_dir = Path(__file__).parent / "results" / "demo" results_dir.mkdir(parents=True, exist_ok=True) json_filename = results_dir / f"demo_{timestamp_str}.json" try: with open(json_filename, "w", encoding="utf-8", errors="surrogateescape") as f: s = json.dumps(results, ensure_ascii=False, indent=4) s = s.encode("utf-8", "ignore").decode("utf-8") f.write(s) except Exception: json_filename = None # Build gallery images gallery_images = [] for idx, res in enumerate(results): img, _ = get_final_image(res, pipe_mode) if img: gallery_images.append((img, f"Candidate {idx}")) # Build evolution HTML evo_parts = [] for idx, res in enumerate(results): stages = get_evolution_stages(res, pipe_mode) if stages: evo_parts.append(f"No evolution data available.
" # Build ZIP zip_path = None if save_results != "No": try: zip_filename = results_dir / f"papervizagent_candidates_{timestamp_str}.zip" buf = BytesIO() with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf: for idx, res in enumerate(results): img, _ = get_final_image(res, pipe_mode) if img: ib = BytesIO() img.save(ib, format="PNG") zf.writestr(f"candidate_{idx}.png", ib.getvalue()) buf.seek(0) with open(zip_filename, "wb") as wf: wf.write(buf.getvalue()) zip_path = str(zip_filename) except Exception: pass status = f"Generated {len(results)} candidates at {datetime.now().strftime('%H:%M:%S')}." if json_filename and Path(str(json_filename)).exists(): status += f" JSON saved to {Path(str(json_filename)).name}." progress(1.0, desc="Done!") return ( gallery_images, # results_gallery evo_html, # evolution_html zip_path, # zip_file_output status, # status_text results, # gen_results_state pipe_mode, # gen_mode_state timestamp_str, # gen_timestamp_state ) generate_btn.click( fn=run_generate, inputs=[ method_content, caption_input, pipeline_mode, retrieval_setting, num_candidates, aspect_ratio, max_critic_rounds, main_model_name, image_model_name, figure_size, save_results, ], outputs=[ results_gallery, evolution_html, zip_file_output, status_text, gen_results_state, gen_mode_state, gen_timestamp_state, ], ) # ============================================================ # TAB 2 — Refine Image # ============================================================ with gr.TabItem("Refine Image"): gr.Markdown("### Refine and upscale your diagram to high resolution (2K/4K)") gr.Markdown("Upload an image, describe changes, and get a high-res version.") with gr.Row(): with gr.Column(): refine_upload = gr.Image(label="Upload Image", type="pil", height=400) with gr.Column(): refine_prompt = gr.Textbox( label="Edit Instructions", lines=6, placeholder="E.g., 'Change the color scheme to match academic paper style' or 'Keep everything the same but output in higher resolution'", ) with gr.Row(): refine_resolution = gr.Dropdown(choices=["2K", "4K"], value="2K", label="Resolution") refine_aspect = gr.Dropdown(choices=["21:9", "16:9", "3:2"], value="21:9", label="Aspect Ratio") refine_btn = gr.Button("Refine Image", variant="primary", elem_classes=["orange-btn"]) refine_status = gr.Textbox(label="Status", interactive=False) with gr.Row(): refine_before = gr.Image(label="Before", interactive=False, height=400) refine_after = gr.Image(label="After", interactive=False, height=400) refine_download = gr.File(label="Download refined image") def run_refine(pil_img, prompt, resolution, ar): if pil_img is None: raise gr.Error("Please upload an image first.") if not prompt: raise gr.Error("Please provide edit instructions.") buf = BytesIO() pil_img.save(buf, format="JPEG") image_bytes = buf.getvalue() loop = asyncio.new_event_loop() try: refined_bytes, msg = loop.run_until_complete( refine_image_with_nanoviz(image_bytes, prompt, aspect_ratio=ar, image_size=resolution) ) except Exception as e: raise gr.Error(f"Refinement error: {e}") finally: loop.close() if not refined_bytes: raise gr.Error(msg) refined_img = Image.open(BytesIO(refined_bytes)) # Save to temp file for download ts = datetime.now().strftime("%Y%m%d_%H%M%S") out_dir = Path(__file__).parent / "results" / "demo" out_dir.mkdir(parents=True, exist_ok=True) out_path = out_dir / f"refined_{resolution}_{ts}.png" refined_img.save(str(out_path), format="PNG") return pil_img, refined_img, str(out_path), msg refine_btn.click( fn=run_refine, inputs=[refine_upload, refine_prompt, refine_resolution, refine_aspect], outputs=[refine_before, refine_after, refine_download, refine_status], ) # ================================================================ # FOOTER # ================================================================ gr.HTML(""" """) return app # --------------------------------------------------------------------------- # Entry point # --------------------------------------------------------------------------- if __name__ == "__main__": app = build_app() app.launch( server_port=7860, share=False, css=CUSTOM_CSS, theme=gr.themes.Default( primary_hue=gr.themes.colors.amber, secondary_hue=gr.themes.colors.gray, neutral_hue=gr.themes.colors.gray, font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"], ), )