Attempt to remove flickering

app.py

@@ -315,7 +315,7 @@ with gr.Blocks(title="Search-TTA (Simplified)", theme=gr.themes.Base()) as demo:
         with gr.Column():
             gr.Markdown("### Live Heatmap Output")
             # gr.Markdown("### Live Heatmap (with TTA)")
-            display_img_tta = gr.Image(label="Heatmap (TTA)", type="filepath", height=400)  # 512
+            display_img_tta = gr.Image(label="Heatmap (TTA per 20 steps)", type="filepath", height=400)  # 512
             status_tta = gr.Markdown("")
             # gr.Markdown("### Live Heatmap (without TTA)")
             display_img_no_tta = gr.Image(label="Heatmap (no TTA)", type="filepath", height=400)  # 512

env.py

@@ -843,7 +843,10 @@ class Env():
         plt.tight_layout()
         os.makedirs(save_dir, exist_ok=True)
         out_path = os.path.join(save_dir, f"{step}.png")
-        fig.savefig(out_path, dpi=100)
+        # Save atomically: write to temp file then move into place so the poller never sees a partial file.
+        tmp_path = out_path + ".tmp"
+        fig.savefig(tmp_path, dpi=100, format='png')
+        os.replace(tmp_path, out_path)  # atomic on same filesystem
         plt.close(fig)
         return out_path