commit

custom_nodes/bernini_chunk/nodes.py

@@ -0,0 +1,135 @@
+"""
+Bernini-R chunk nodes — state caching for chunked ZeroGPU generation.
+Dropped into custom_nodes/bernini_chunk/ of the Bernini-R-Lightning space.
+
+Works with KSamplerAdvanced (Kijai) instead of SamplerCustomAdvanced+Rudra fork.
+No relay mask — Bernini-R doesn't need it. Pure save/load + sigma step control.
+"""
+from __future__ import annotations
+import os
+import torch
+
+_TMP = "/tmp"
+
+
+def _lat_path(session_id: str) -> str:
+    return os.path.join(_TMP, f"bernini_lat_{session_id}.pt")
+
+
+def _cond_path(session_id: str) -> str:
+    return os.path.join(_TMP, f"bernini_cond_{session_id}.pt")
+
+
+class BerniniChunkSaveLatent:
+    """Save KSamplerAdvanced output latent to /tmp for the next chunk."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {
+            "samples":    ("LATENT",),
+            "session_id": ("STRING", {"default": ""}),
+        }}
+
+    RETURN_TYPES  = ("LATENT",)
+    RETURN_NAMES  = ("samples",)
+    FUNCTION      = "save"
+    CATEGORY      = "Bernini/Chunk"
+    OUTPUT_NODE   = True
+
+    def save(self, samples, session_id):
+        if session_id:
+            path = _lat_path(session_id)
+            torch.save({"samples": samples["samples"].cpu()}, path)
+            print(f"[BerniniChunk] saved latent -> {path}", flush=True)
+        return (samples,)
+
+
+class BerniniChunkLoadLatent:
+    """Load previously saved latent. Falls back to a provided default if none exists."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {
+            "fallback":   ("LATENT",),
+            "session_id": ("STRING", {"default": ""}),
+        }}
+
+    RETURN_TYPES  = ("LATENT",)
+    RETURN_NAMES  = ("samples",)
+    FUNCTION      = "load"
+    CATEGORY      = "Bernini/Chunk"
+
+    def load(self, fallback, session_id):
+        if session_id:
+            path = _lat_path(session_id)
+            if os.path.exists(path):
+                data = torch.load(path, map_location="cpu", weights_only=False)
+                print(f"[BerniniChunk] loaded latent <- {path}", flush=True)
+                return ({"samples": data["samples"]},)
+        return (fallback,)
+
+
+class BerniniChunkSaveCond:
+    """Save positive/negative conditioning after step 0 (planning/text-encode)."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {
+            "positive":   ("CONDITIONING",),
+            "negative":   ("CONDITIONING",),
+            "session_id": ("STRING", {"default": ""}),
+        }}
+
+    RETURN_TYPES  = ("CONDITIONING", "CONDITIONING")
+    RETURN_NAMES  = ("positive", "negative")
+    FUNCTION      = "save"
+    CATEGORY      = "Bernini/Chunk"
+    OUTPUT_NODE   = True
+
+    def save(self, positive, negative, session_id):
+        if session_id:
+            path = _cond_path(session_id)
+            torch.save({"positive": positive, "negative": negative}, path)
+            print(f"[BerniniChunk] saved conditioning -> {path}", flush=True)
+        return (positive, negative)
+
+
+class BerniniChunkLoadCond:
+    """Load saved conditioning. Falls back to fresh if no saved file exists."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {
+            "fallback_pos": ("CONDITIONING",),
+            "fallback_neg": ("CONDITIONING",),
+            "session_id":   ("STRING", {"default": ""}),
+        }}
+
+    RETURN_TYPES  = ("CONDITIONING", "CONDITIONING")
+    RETURN_NAMES  = ("positive", "negative")
+    FUNCTION      = "load"
+    CATEGORY      = "Bernini/Chunk"
+
+    def load(self, fallback_pos, fallback_neg, session_id):
+        if session_id:
+            path = _cond_path(session_id)
+            if os.path.exists(path):
+                data = torch.load(path, map_location="cpu", weights_only=False)
+                print(f"[BerniniChunk] loaded conditioning <- {path}", flush=True)
+                return (data["positive"], data["negative"])
+        return (fallback_pos, fallback_neg)
+
+
+NODE_CLASS_MAPPINGS = {
+    "BerniniChunkSaveLatent":  BerniniChunkSaveLatent,
+    "BerniniChunkLoadLatent":  BerniniChunkLoadLatent,
+    "BerniniChunkSaveCond":    BerniniChunkSaveCond,
+    "BerniniChunkLoadCond":    BerniniChunkLoadCond,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "BerniniChunkSaveLatent":  "Bernini · Chunk Save Latent",
+    "BerniniChunkLoadLatent":  "Bernini · Chunk Load Latent",
+    "BerniniChunkSaveCond":    "Bernini · Chunk Save Conditioning",
+    "BerniniChunkLoadCond":    "Bernini · Chunk Load Conditioning",
+}