Ethanfel
f9bdc75a2a
When SaveSnapshot receives an image tensor, extract a 200x150 JPEG thumbnail (base64) and include it in the snapshot record. Sidebar shows a small preview, tooltip displays the generated image instead of SVG, and the preview modal shows the image above the graph. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
71 lines
1.7 KiB
Python
71 lines
1.7 KiB
Python
import base64
|
|
import io
|
|
|
|
from server import PromptServer
|
|
|
|
|
|
class _AnyType(str):
|
|
def __ne__(self, other):
|
|
return False
|
|
|
|
|
|
ANY_TYPE = _AnyType("*")
|
|
|
|
|
|
def _make_thumbnail(value):
|
|
"""Convert an image tensor to a base64 JPEG thumbnail, or return None."""
|
|
try:
|
|
import torch
|
|
if not isinstance(value, torch.Tensor):
|
|
return None
|
|
if value.ndim != 4 or value.shape[3] not in (3, 4):
|
|
return None
|
|
|
|
from PIL import Image
|
|
|
|
frame = value[0] # first frame only
|
|
if frame.shape[2] == 4:
|
|
frame = frame[:, :, :3] # drop alpha
|
|
arr = frame.clamp(0, 1).mul(255).byte().cpu().numpy()
|
|
img = Image.fromarray(arr, mode="RGB")
|
|
img.thumbnail((200, 150), Image.LANCZOS)
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="JPEG", quality=75)
|
|
return base64.b64encode(buf.getvalue()).decode("ascii")
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
class SaveSnapshot:
|
|
CATEGORY = "Snapshot Manager"
|
|
FUNCTION = "execute"
|
|
RETURN_TYPES = (ANY_TYPE,)
|
|
OUTPUT_NODE = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"value": (ANY_TYPE, {}),
|
|
"label": ("STRING", {"default": "Node Trigger"}),
|
|
}
|
|
}
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, input_types):
|
|
return True
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, *args, **kwargs):
|
|
return float("NaN")
|
|
|
|
def execute(self, value, label):
|
|
payload = {"label": label}
|
|
thumbnail = _make_thumbnail(value)
|
|
if thumbnail is not None:
|
|
payload["thumbnail"] = thumbnail
|
|
PromptServer.instance.send_sync(
|
|
"snapshot-manager-capture", payload
|
|
)
|
|
return (value,)
|