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 <noreply@anthropic.com>
This commit is contained in:
2026-04-08 23:39:30 +02:00
parent 0b24207ca5
commit 92535deab2
+13 -5
View File
@@ -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))