Initial commit: save/load latent nodes with absolute path support
Preserves full LATENT dict including tensors, non-tensor metadata, and original device placement. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
13
__init__.py
Normal file
13
__init__.py
Normal file
@@ -0,0 +1,13 @@
|
|||||||
|
from .nodes import SaveLatentAbsolute, LoadLatentAbsolute
|
||||||
|
|
||||||
|
NODE_CLASS_MAPPINGS = {
|
||||||
|
"SaveLatentAbsolute": SaveLatentAbsolute,
|
||||||
|
"LoadLatentAbsolute": LoadLatentAbsolute,
|
||||||
|
}
|
||||||
|
|
||||||
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||||
|
"SaveLatentAbsolute": "Save Latent (Absolute Path)",
|
||||||
|
"LoadLatentAbsolute": "Load Latent (Absolute Path)",
|
||||||
|
}
|
||||||
|
|
||||||
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||||
77
nodes.py
Normal file
77
nodes.py
Normal file
@@ -0,0 +1,77 @@
|
|||||||
|
import os
|
||||||
|
import json
|
||||||
|
import torch
|
||||||
|
import safetensors.torch
|
||||||
|
|
||||||
|
|
||||||
|
class SaveLatentAbsolute:
|
||||||
|
@classmethod
|
||||||
|
def INPUT_TYPES(cls):
|
||||||
|
return {
|
||||||
|
"required": {
|
||||||
|
"samples": ("LATENT",),
|
||||||
|
"path": ("STRING", {"default": "/path/to/latent.latent"}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
RETURN_TYPES = ("LATENT",)
|
||||||
|
FUNCTION = "save"
|
||||||
|
CATEGORY = "latent"
|
||||||
|
OUTPUT_NODE = True
|
||||||
|
|
||||||
|
def save(self, samples, path):
|
||||||
|
path = os.path.expanduser(path)
|
||||||
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||||
|
|
||||||
|
tensors = {}
|
||||||
|
non_tensors = {}
|
||||||
|
devices = {}
|
||||||
|
for key, value in samples.items():
|
||||||
|
if isinstance(value, torch.Tensor):
|
||||||
|
devices[key] = str(value.device)
|
||||||
|
tensors[key] = value.contiguous()
|
||||||
|
else:
|
||||||
|
non_tensors[key] = value
|
||||||
|
|
||||||
|
metadata = {"devices": json.dumps(devices)}
|
||||||
|
if non_tensors:
|
||||||
|
metadata["non_tensor_data"] = json.dumps(non_tensors)
|
||||||
|
|
||||||
|
safetensors.torch.save_file(tensors, path, metadata=metadata)
|
||||||
|
return (samples,)
|
||||||
|
|
||||||
|
|
||||||
|
class LoadLatentAbsolute:
|
||||||
|
@classmethod
|
||||||
|
def INPUT_TYPES(cls):
|
||||||
|
return {
|
||||||
|
"required": {
|
||||||
|
"path": ("STRING", {"default": "/path/to/latent.latent"}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
RETURN_TYPES = ("LATENT",)
|
||||||
|
FUNCTION = "load"
|
||||||
|
CATEGORY = "latent"
|
||||||
|
|
||||||
|
def load(self, path):
|
||||||
|
path = os.path.expanduser(path)
|
||||||
|
|
||||||
|
samples = safetensors.torch.load_file(path, device="cpu")
|
||||||
|
|
||||||
|
with safetensors.safe_open(path, framework="pt") as f:
|
||||||
|
meta = f.metadata()
|
||||||
|
|
||||||
|
# Restore original devices
|
||||||
|
if meta and "devices" in meta:
|
||||||
|
devices = json.loads(meta["devices"])
|
||||||
|
for key, device in devices.items():
|
||||||
|
if key in samples:
|
||||||
|
samples[key] = samples[key].to(device)
|
||||||
|
|
||||||
|
# Restore non-tensor data
|
||||||
|
if meta and "non_tensor_data" in meta:
|
||||||
|
non_tensors = json.loads(meta["non_tensor_data"])
|
||||||
|
samples.update(non_tensors)
|
||||||
|
|
||||||
|
return (samples,)
|
||||||
Reference in New Issue
Block a user