Add latent sidecars to fast saver

2026-07-07 15:31:13 +02:00
parent aae0bdf746
commit 4b0b4b6c5a
5 changed files with 368 additions and 20 deletions
@@ -0,0 +1,136 @@
+# Fast Absolute Saver Latent Sidecars Implementation Plan
+
+> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
+
+**Goal:** Save unmodified optional latents next to `FastAbsoluteSaver` media outputs and load them back by absolute path.
+
+**Architecture:** Keep latent save/load behavior in `fast_saver.py` beside the existing saver node. Add small helpers for sidecar path derivation and direct `torch.save`/`torch.load` persistence so media naming and latent naming stay coupled.
+
+**Tech Stack:** Python 3.10+, PyTorch, pytest, existing ComfyUI node mapping conventions.
+
+---
+
+### Task 1: Latent Sidecar Tests
+
+**Files:**
+- Create: `tests/test_fast_saver_latent.py`
+- Modify: `fast_saver.py`
+
+- [ ] **Step 1: Write the failing tests**
+
+```python
+import torch
+
+from fast_saver import FastAbsoluteSaver
+
+
+def test_png_save_writes_matching_latent_sidecar(tmp_path):
+    saver = FastAbsoluteSaver()
+    images = torch.zeros((1, 2, 2, 3), dtype=torch.float32)
+    latent = {"samples": torch.arange(4, dtype=torch.float32).reshape(1, 1, 2, 2), "keep": {"value": 7}}
+
+    saver.save_images_fast(
+        images=images,
+        output_path=str(tmp_path),
+        filename_prefix="frame",
+        save_format="png",
+        use_timestamp=False,
+        auto_increment=False,
+        counter_digits=4,
+        max_threads=1,
+        filename_with_score=False,
+        metadata_key="sharpness_score",
+        save_workflow_metadata=False,
+        save_metadata_png=False,
+        webp_lossless=True,
+        webp_quality=100,
+        webp_method=4,
+        video_fps=24,
+        video_crf=18,
+        video_pixel_format="yuv420p",
+        video_bitrate=10,
+        prores_profile="hq",
+        gif_dither="sierra2_4a",
+        latent=latent,
+    )
+
+    latent_path = tmp_path / "frame_0000.latent"
+    loaded = torch.load(latent_path, map_location="cpu", weights_only=False)
+    assert torch.equal(loaded["samples"], latent["samples"])
+    assert loaded["keep"] == {"value": 7}
+
+
+def test_video_save_writes_latent_sidecar_next_to_video(tmp_path):
+    saver = FastAbsoluteSaver()
+    images = torch.zeros((2, 2, 2, 3), dtype=torch.float32)
+    latent = {"samples": torch.arange(8, dtype=torch.float32).reshape(2, 1, 2, 2)}
+    video_path = tmp_path / "clip_0001.mp4"
+
+    def fake_save_video(*args, **kwargs):
+        video_path.write_bytes(b"video")
+        return str(video_path)
+
+    saver.save_video = fake_save_video
+    saver.save_images_fast(
+        images=images,
+        output_path=str(tmp_path),
+        filename_prefix="clip",
+        save_format="mp4",
+        use_timestamp=False,
+        auto_increment=False,
+        counter_digits=4,
+        max_threads=1,
+        filename_with_score=False,
+        metadata_key="sharpness_score",
+        save_workflow_metadata=False,
+        save_metadata_png=False,
+        webp_lossless=True,
+        webp_quality=100,
+        webp_method=4,
+        video_fps=24,
+        video_crf=18,
+        video_pixel_format="yuv420p",
+        video_bitrate=10,
+        prores_profile="hq",
+        gif_dither="sierra2_4a",
+        latent=latent,
+    )
+
+    loaded = torch.load(tmp_path / "clip_0001.latent", map_location="cpu", weights_only=False)
+    assert torch.equal(loaded["samples"], latent["samples"])
+
+
+def test_load_latent_absolute_round_trips_saved_object(tmp_path):
+    from fast_saver import JDL_LoadLatentAbsolute
+
+    path = tmp_path / "sample.latent"
+    latent = {"samples": torch.ones((1, 4, 8, 8)), "noise_mask": torch.zeros((1, 1, 8, 8))}
+    torch.save(latent, path)
+
+    loaded, = JDL_LoadLatentAbsolute().load_latent(str(path))
+
+    assert torch.equal(loaded["samples"], latent["samples"])
+    assert torch.equal(loaded["noise_mask"], latent["noise_mask"])
+```
+
+- [ ] **Step 2: Run tests to verify they fail**
+
+Run: `python -m pytest tests/test_fast_saver_latent.py -q`
+
+Expected: tests fail because `latent` is not an accepted input and `JDL_LoadLatentAbsolute` does not exist.
+
+- [ ] **Step 3: Implement minimal production code**
+
+Add `latent` to `FastAbsoluteSaver.INPUT_TYPES()["optional"]`, accept it in `save_images_fast`, save `.latent` sidecars with `torch.save`, and register `JDL_LoadLatentAbsolute`.
+
+- [ ] **Step 4: Run focused tests**
+
+Run: `python -m pytest tests/test_fast_saver_latent.py -q`
+
+Expected: both tests pass.
+
+- [ ] **Step 5: Run broader verification**
+
+Run: `python -m pytest tests/test_fast_saver_latent.py -q && python -m compileall fast_saver.py image_preview.py string_utils.py json_loader_dynamic.py`
+
+Expected: pytest passes and compileall reports no syntax errors.
@@ -0,0 +1,24 @@
+# Fast Absolute Saver Latent Sidecars Design
+
+## Goal
+
+Add an optional `latent` input to `FastAbsoluteSaver` that saves the unmodified latent object next to the generated media, and add an absolute-path latent loader that can read those sidecars back into ComfyUI.
+
+## Behavior
+
+- `FastAbsoluteSaver` accepts an optional `latent` input of type `LATENT`.
+- When `latent` is connected, the saver writes a sidecar file with the exact media base name and `.latent` extension.
+- Video outputs produce one sidecar: `clip.mp4` writes `clip.latent`.
+- Image sequence outputs write a sidecar for each saved image. Each sidecar stores the full connected latent object unchanged rather than slicing per frame.
+- The latent is persisted with `torch.save` as provided. The saver must not prune keys, convert tensors, detach tensors, clone tensors, or strip metadata.
+- A new `Load Latent Absolute` node accepts an absolute path and returns a `LATENT`.
+- Missing or invalid latent files raise clear errors rather than silently blocking downstream nodes.
+
+## Files
+
+- `fast_saver.py` owns the saver and the new absolute latent load node.
+- `tests/test_fast_saver_latent.py` covers sidecar naming, exact object persistence, and absolute loading.
+
+## Performance Notes
+
+The implementation should keep the current fast batch image conversion path. Any performance improvement should be low risk and localized, such as reusing output path decisions for media and sidecars rather than recomputing or guessing names after save.