From 92535deab2be6040fc7e6250de15440d289abcca Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Wed, 8 Apr 2026 23:39:30 +0200 Subject: [PATCH] fix(ti-scheduler): save comparison image after each completed experiment Previously the comparison PNG was only written at the very end of the sweep, so an interrupted run produced no image at all. Now _save_comparison() is called right after _write_summary() for every successful experiment, keeping loss_comparison.png current throughout the sweep. Co-Authored-By: Claude Sonnet 4.6 --- nodes/selva_ti_scheduler.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/nodes/selva_ti_scheduler.py b/nodes/selva_ti_scheduler.py index 65cdcb7..7f134ca 100644 --- a/nodes/selva_ti_scheduler.py +++ b/nodes/selva_ti_scheduler.py @@ -326,9 +326,18 @@ class SelvaTiScheduler: "experiments": [], } + comparison_img_path = output_root / "loss_comparison.png" + def _write_summary(): summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8") + def _save_comparison(): + try: + img = _draw_comparison_curves(all_curve_data) + img.save(str(comparison_img_path)) + except Exception as e: + print(f"[TI Scheduler] Comparison image failed: {e}", flush=True) + _write_summary() # ------------------------------------------------------------------ @@ -452,6 +461,7 @@ class SelvaTiScheduler: continue _write_summary() + _save_comparison() pbar_outer.update(1) # ------------------------------------------------------------------ @@ -462,10 +472,8 @@ class SelvaTiScheduler: print(f"\n[TI Scheduler] Sweep complete. Summary: {summary_path}", flush=True) # ------------------------------------------------------------------ - # 7. Comparison image + # 7. Comparison image (final update, then return to ComfyUI) # ------------------------------------------------------------------ + _save_comparison() comparison_img = _draw_comparison_curves(all_curve_data) - comparison_img.save(str(output_root / "loss_comparison.png")) - comparison_tensor = _pil_to_tensor(comparison_img) - - return (str(summary_path), comparison_tensor) + return (str(summary_path), _pil_to_tensor(comparison_img))