diff --git a/src/bridge/images.py b/src/bridge/images.py index 3deb343..0b36bf2 100644 --- a/src/bridge/images.py +++ b/src/bridge/images.py @@ -1,10 +1,27 @@ +import glob import os +import time import uuid import numpy as np import torch from PIL import Image +def _reap_old(out_dir: str, keep: str) -> None: + """Delete img_*.png in out_dir older than COMFY_BRIDGE_TMP_TTL (default + 3600s) so the temp dir doesn't grow forever. Never touches `keep`.""" + ttl = float(os.environ.get("COMFY_BRIDGE_TMP_TTL", "3600")) + cutoff = time.time() - ttl + for f in glob.glob(os.path.join(out_dir, "img_*.png")): + if f == keep: + continue + try: + if os.path.getmtime(f) < cutoff: + os.unlink(f) + except (FileNotFoundError, OSError): + pass # races / permission churn -> best-effort + + def save_tensor_png(tensor: torch.Tensor, out_dir: str) -> str: """ComfyUI IMAGE [B,H,W,C] float 0..1 -> PNG file, returns path.""" os.makedirs(out_dir, exist_ok=True) @@ -14,6 +31,7 @@ def save_tensor_png(tensor: torch.Tensor, out_dir: str) -> str: img = Image.fromarray(arr, mode="RGB") path = os.path.join(out_dir, f"img_{uuid.uuid4().hex}.png") img.save(path) + _reap_old(out_dir, keep=path) return path diff --git a/tests/bridge/test_images_gc.py b/tests/bridge/test_images_gc.py new file mode 100644 index 0000000..feb14a5 --- /dev/null +++ b/tests/bridge/test_images_gc.py @@ -0,0 +1,39 @@ +"""Age-based GC of temp PNGs: save_tensor_png reaps img_*.png older than the +TTL (env COMFY_BRIDGE_TMP_TTL, default 3600s) while keeping recent files and +the file it just wrote.""" +import os +import time + +import torch + +from src.bridge import images + + +def _tensor(): + return torch.zeros(1, 4, 4, 3) + + +def test_old_temp_image_is_reaped(tmp_path, monkeypatch): + monkeypatch.setenv("COMFY_BRIDGE_TMP_TTL", "3600") + old = tmp_path / "img_old.png" + old.write_bytes(b"stale") + backdated = time.time() - 7200 # 2h old, past the 1h TTL + os.utime(old, (backdated, backdated)) + + new_path = images.save_tensor_png(_tensor(), str(tmp_path)) + + assert not old.exists(), "stale temp image should have been reaped" + assert os.path.exists(new_path), "freshly written image must survive" + + +def test_recent_temp_image_is_kept(tmp_path, monkeypatch): + monkeypatch.setenv("COMFY_BRIDGE_TMP_TTL", "3600") + recent = tmp_path / "img_recent.png" + recent.write_bytes(b"fresh") + recent_mtime = time.time() - 60 # 1 min old, within TTL + os.utime(recent, (recent_mtime, recent_mtime)) + + new_path = images.save_tensor_png(_tensor(), str(tmp_path)) + + assert recent.exists(), "recent temp image within TTL must be kept" + assert os.path.exists(new_path)