Ethanfel/ComfyUI-JSON-Dynamic
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add latent sidecars to fast saver

Browse Source
This commit is contained in:
Ethanfel
2026-07-07 15:31:13 +02:00
parent aae0bdf746
commit 4b0b4b6c5a
5 changed files with 368 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+40 -16
View File
@@ -1,19 +1,43 @@
from .json_loader_dynamic import (
NODE_CLASS_MAPPINGS as _json_class_mappings,
NODE_DISPLAY_NAME_MAPPINGS as _json_display_mappings,
)
from .string_utils import (
NODE_CLASS_MAPPINGS as _string_class_mappings,
NODE_DISPLAY_NAME_MAPPINGS as _string_display_mappings,
)
from .image_preview import (
NODE_CLASS_MAPPINGS as _image_class_mappings,
NODE_DISPLAY_NAME_MAPPINGS as _image_display_mappings,
)
from .fast_saver import (
NODE_CLASS_MAPPINGS as _saver_class_mappings,
NODE_DISPLAY_NAME_MAPPINGS as _saver_display_mappings,
)
if __package__:
from .json_loader_dynamic import (
NODE_CLASS_MAPPINGS as _json_class_mappings,
NODE_DISPLAY_NAME_MAPPINGS as _json_display_mappings,
)
from .string_utils import (
NODE_CLASS_MAPPINGS as _string_class_mappings,
NODE_DISPLAY_NAME_MAPPINGS as _string_display_mappings,
)
from .image_preview import (
NODE_CLASS_MAPPINGS as _image_class_mappings,
NODE_DISPLAY_NAME_MAPPINGS as _image_display_mappings,
)
from .fast_saver import (
NODE_CLASS_MAPPINGS as _saver_class_mappings,
NODE_DISPLAY_NAME_MAPPINGS as _saver_display_mappings,
)
else:
from json_loader_dynamic import (
NODE_CLASS_MAPPINGS as _json_class_mappings,
NODE_DISPLAY_NAME_MAPPINGS as _json_display_mappings,
)
from string_utils import (
NODE_CLASS_MAPPINGS as _string_class_mappings,
NODE_DISPLAY_NAME_MAPPINGS as _string_display_mappings,
)
try:
from image_preview import (
NODE_CLASS_MAPPINGS as _image_class_mappings,
NODE_DISPLAY_NAME_MAPPINGS as _image_display_mappings,
)
except ModuleNotFoundError as e:
if e.name != "folder_paths":
raise
_image_class_mappings = {}
_image_display_mappings = {}
from fast_saver import (
NODE_CLASS_MAPPINGS as _saver_class_mappings,
NODE_DISPLAY_NAME_MAPPINGS as _saver_display_mappings,
)
NODE_CLASS_MAPPINGS = {**_json_class_mappings, **_string_class_mappings, **_image_class_mappings, **_saver_class_mappings}
NODE_DISPLAY_NAME_MAPPINGS = {**_json_display_mappings, **_string_display_mappings, **_image_display_mappings, **_saver_display_mappings}
docs/superpowers/plans/2026-07-07-fast-absolute-saver-latent-sidecars.md
+136
View File
@@ -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.
docs/superpowers/specs/2026-07-07-fast-absolute-saver-latent-sidecars-design.md
+24
View File
@@ -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.
fast_saver.py
+80 -4
View File
@@ -142,6 +142,17 @@ VIDEO_FORMATS = {
}
def _latent_sidecar_path(media_path):
return os.path.splitext(media_path)[0] + ".latent"
def _load_latent_file(latent_path):
try:
return torch.load(latent_path, map_location="cpu", weights_only=False)
except TypeError:
return torch.load(latent_path, map_location="cpu")
class FastAbsoluteSaver:
@classmethod
def INPUT_TYPES(s):
@@ -185,6 +196,7 @@ class FastAbsoluteSaver:
"optional": {
"scores_info": ("STRING", {"forceInput": True}),
"audio": ("AUDIO", ),
"latent": ("LATENT", ),
},
# Hidden inputs used to capture the workflow graph
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
@@ -248,6 +260,12 @@ class FastAbsoluteSaver:
meta.add_text("workflow", workflow_json)
img.save(png_path, format="PNG", pnginfo=meta, compress_level=1)
def _save_latent_sidecar(self, latent, media_path):
latent_path = _latent_sidecar_path(media_path)
torch.save(latent, latent_path)
print(f"xx- FastSaver: Latent sidecar saved to {latent_path}")
return latent_path
def save_single_image(self, img_array, full_path, score, key_name, fmt, lossless, quality, method,
save_workflow, prompt_data, extra_data, force_png_metadata=False):
try:
@@ -489,7 +507,7 @@ class FastAbsoluteSaver:
webp_lossless, webp_quality, webp_method,
video_fps, video_crf, video_pixel_format,
video_bitrate, prores_profile, gif_dither,
scores_info=None, audio=None, prompt=None, extra_pnginfo=None):
scores_info=None, audio=None, latent=None, prompt=None, extra_pnginfo=None):
output_path = output_path.strip('"')
if not os.path.isabs(output_path):
@@ -519,6 +537,8 @@ class FastAbsoluteSaver:
extra_data=extra_pnginfo,
bitrate=video_bitrate, prores_profile=prores_profile,
gif_dither=gif_dither, audio=audio)
if latent is not None:
self._save_latent_sidecar(latent, out_file)
# Save metadata sidecar PNG next to the video file
if save_metadata_png:
png_path = os.path.splitext(out_file)[0] + ".png"
@@ -545,6 +565,7 @@ class FastAbsoluteSaver:
print(f"xx- FastSaver: Saving {batch_size} images to {output_path}...")
first_image_path = None
saved_image_paths = []
with concurrent.futures.ThreadPoolExecutor(max_workers=max_threads) as executor:
futures = []
@@ -575,15 +596,24 @@ class FastAbsoluteSaver:
# For PNG sequences: embed workflow metadata in the first file only
force_meta = (save_metadata_png and save_format == "png" and i == 0)
futures.append(executor.submit(
future = executor.submit(
self.save_single_image,
img_array, full_path, current_score, metadata_key,
save_format, webp_lossless, webp_quality, webp_method,
save_workflow_metadata, prompt, extra_pnginfo,
force_png_metadata=force_meta
))
)
futures.append((full_path, future))
concurrent.futures.wait(futures)
concurrent.futures.wait([future for _, future in futures])
for full_path, future in futures:
if future.result():
saved_image_paths.append(full_path)
if latent is not None:
for image_path in saved_image_paths:
self._save_latent_sidecar(latent, image_path)
# Save a single metadata sidecar PNG using the first image (skip for PNG sequences - handled above)
if save_metadata_png and save_format != "png" and first_image_path is not None:
@@ -596,10 +626,56 @@ class FastAbsoluteSaver:
return {"ui": {"images": []}}
class JDL_LoadLatentAbsolute:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"path": ("STRING", {"default": "", "multiline": False}),
},
}
RETURN_TYPES = ("LATENT",)
RETURN_NAMES = ("latent",)
FUNCTION = "load_latent"
CATEGORY = "JSON Dynamic/io"
def load_latent(self, path):
latent_path = path.strip().strip('"')
if not os.path.isabs(latent_path):
raise ValueError(f"Load Latent Absolute requires an absolute path: {path}")
if not os.path.isfile(latent_path):
raise FileNotFoundError(f"Latent file not found: {latent_path}")
latent = _load_latent_file(latent_path)
if not isinstance(latent, dict) or "samples" not in latent:
raise ValueError(f"Expected a latent dict with a 'samples' key: {latent_path}")
return (latent,)
@classmethod
def IS_CHANGED(s, path):
latent_path = path.strip().strip('"')
if not os.path.isfile(latent_path):
return f"missing:{latent_path}"
stat_result = os.stat(latent_path)
return f"{latent_path}:{stat_result.st_mtime_ns}:{stat_result.st_size}"
@classmethod
def VALIDATE_INPUTS(s, path):
latent_path = path.strip().strip('"')
if not os.path.isabs(latent_path):
return f"Path must be absolute: {path}"
if not os.path.isfile(latent_path):
return f"Latent file not found: {latent_path}"
return True
NODE_CLASS_MAPPINGS = {
"FastAbsoluteSaver": FastAbsoluteSaver,
"JDL_LoadLatentAbsolute": JDL_LoadLatentAbsolute,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"FastAbsoluteSaver": "Fast Absolute Saver (Metadata)",
"JDL_LoadLatentAbsolute": "Load Latent Absolute",
}
tests/test_fast_saver_latent.py
+88
View File
@@ -0,0 +1,88 @@
import pytest
import torch
from fast_saver import FastAbsoluteSaver
def _save_args(tmp_path, *, save_format="png", latent=None):
return {
"images": torch.zeros((1, 2, 2, 3), dtype=torch.float32),
"output_path": str(tmp_path),
"filename_prefix": "frame",
"save_format": save_format,
"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,
}
def test_png_save_writes_matching_latent_sidecar(tmp_path):
saver = FastAbsoluteSaver()
latent = {
"samples": torch.arange(4, dtype=torch.float32).reshape(1, 1, 2, 2),
"keep": {"value": 7},
}
saver.save_images_fast(**_save_args(tmp_path, latent=latent))
loaded = torch.load(tmp_path / "frame_0000.latent", 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()
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
args = _save_args(tmp_path, save_format="mp4", latent=latent)
args["images"] = torch.zeros((2, 2, 2, 3), dtype=torch.float32)
args["filename_prefix"] = "clip"
saver.save_images_fast(**args)
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"])
def test_load_latent_absolute_rejects_relative_paths():
from fast_saver import JDL_LoadLatentAbsolute
with pytest.raises(ValueError, match="absolute"):
JDL_LoadLatentAbsolute().load_latent("sample.latent")