Add latent sidecars to fast saver
This commit is contained in:
+32
-8
@@ -1,19 +1,43 @@
|
||||
from .json_loader_dynamic import (
|
||||
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 (
|
||||
)
|
||||
from .string_utils import (
|
||||
NODE_CLASS_MAPPINGS as _string_class_mappings,
|
||||
NODE_DISPLAY_NAME_MAPPINGS as _string_display_mappings,
|
||||
)
|
||||
from .image_preview import (
|
||||
)
|
||||
from .image_preview import (
|
||||
NODE_CLASS_MAPPINGS as _image_class_mappings,
|
||||
NODE_DISPLAY_NAME_MAPPINGS as _image_display_mappings,
|
||||
)
|
||||
from .fast_saver import (
|
||||
)
|
||||
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}
|
||||
|
||||
@@ -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.
|
||||
+80
-4
@@ -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",
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
Reference in New Issue
Block a user