Merge 484081e00c into 51bf508a0b

comfy_extras/nodes_dataset.py

@@ -1,10 +1,14 @@
 import logging
 import os
 import json
+import pickle
+import struct
 
 import numpy as np
+import safetensors.torch
 import torch
 from PIL import Image
+from safetensors import safe_open
 from typing_extensions import override
 
 import folder_paths
@@ -1235,6 +1239,285 @@ class MergeTextListsNode(TextProcessingNode):
 # ========== Training Dataset Nodes ==========
 
 
+# Sentinel key used in the "skeleton" to mark where a tensor lived in the
+# original nested structure. The skeleton is pickled; tensors live in the
+# accompanying safetensors file under the referenced key.
+_TREF_KEY = "__tref__"
+
+
+def _split_tensors(obj, out_tensors, prefix):
+    """Walk obj recursively. Pull tensors out into out_tensors (keyed by f"{prefix}_{N}")
+    and return a "skeleton" with the same structure but each tensor replaced by
+    {"__tref__": key}. Everything that isn't a tensor / dict / list / tuple
+    (Hook objects, floats, strings, custom extension types, ...) passes through
+    untouched and will be handled by pickle.
+    """
+    if isinstance(obj, torch.Tensor):
+        key = f"{prefix}_{len(out_tensors)}"
+        out_tensors[key] = obj.detach().cpu().clone()
+        return {_TREF_KEY: key}
+    elif isinstance(obj, dict):
+        return {k: _split_tensors(v, out_tensors, prefix) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [_split_tensors(v, out_tensors, prefix) for v in obj]
+    elif isinstance(obj, tuple):
+        return tuple(_split_tensors(v, out_tensors, prefix) for v in obj)
+    return obj
+
+
+def _rejoin_tensors(obj, tensor_getter):
+    """Inverse of _split_tensors. Walk skeleton, fetch tensors via tensor_getter(key)
+    wherever a {"__tref__": ...} marker appears.
+    """
+    if isinstance(obj, dict):
+        if len(obj) == 1 and _TREF_KEY in obj:
+            return tensor_getter(obj[_TREF_KEY])
+        return {k: _rejoin_tensors(v, tensor_getter) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [_rejoin_tensors(v, tensor_getter) for v in obj]
+    if isinstance(obj, tuple):
+        return tuple(_rejoin_tensors(v, tensor_getter) for v in obj)
+    return obj
+
+
+# safetensors dtype strings -> torch dtype, used to read shapes/dtypes from the
+# header without loading any tensor bytes.
+_ST_STR_TO_DTYPE = {
+    "F64": torch.float64, "F32": torch.float32, "F16": torch.float16,
+    "BF16": torch.bfloat16, "I64": torch.int64, "I32": torch.int32,
+    "I16": torch.int16, "I8": torch.int8, "U8": torch.uint8, "BOOL": torch.bool,
+}
+
+
+def _read_safetensors_header(path):
+    """Read the safetensors header (dtype + shape per tensor key) without reading
+    any tensor data. The file starts with an 8-byte little-endian header length
+    followed by that many bytes of JSON."""
+    with open(path, "rb") as f:
+        n = struct.unpack("<Q", f.read(8))[0]
+        header = json.loads(f.read(n))
+    header.pop("__metadata__", None)
+    return header
+
+
+class RealizeRequired(RuntimeError):
+    """Raised when lazy on-disk dataset data is used where real tensors are
+    needed. Realize it first: .realize() in code, or the Realize Lazy Latents /
+    Realize Lazy Conditionings nodes in a workflow."""
+
+
+def _need_realize(self, *args, **kwargs):
+    raise RealizeRequired(
+        f"{type(self).__name__} is lazy on-disk data and does not support this "
+        f"operation. Realize it first (.realize() or a Realize node)."
+    )
+
+
+class LazyTensorInfo:
+    """Shape/dtype of one on-disk tensor, read from the safetensors header — no
+    tensor bytes. Anything beyond .shape/.dtype/.ndim raises RealizeRequired."""
+
+    def __init__(self, shape, dtype):
+        self.shape = torch.Size(shape)
+        self.dtype = dtype
+        self.ndim = len(self.shape)
+
+    def __repr__(self):
+        return f"LazyTensorInfo(shape={tuple(self.shape)}, dtype={self.dtype})"
+
+    __getattr__ = _need_realize
+
+
+class LazyLatent:
+    """One dataset sample's latent dict ({"samples": tensor, ...}) on disk.
+
+    Carries the sample's skeleton, so latent["samples"] serves shape/dtype from
+    the safetensors header with zero I/O. Tensor values require realization:
+    realize() -> real latent dict, realize_samples() -> real "samples" tensor.
+    Realization is never cached; a persistent list[LazyLatent] stays near-zero
+    RAM (the OS page cache handles re-read locality).
+    """
+
+    def __init__(self, reader, skeleton):
+        self._reader = reader
+        self._skel = skeleton
+
+    def __getitem__(self, name):
+        v = self._skel[name]
+        if isinstance(v, dict) and len(v) == 1 and _TREF_KEY in v:
+            key = v[_TREF_KEY]
+            return LazyTensorInfo(self._reader.shape(key), self._reader.dtype(key))
+        return v  # plain non-tensor value (e.g. batch_index)
+
+    def realize(self):
+        """Read this sample's tensors from disk; return the real latent dict."""
+        return _rejoin_tensors(self._skel, self._reader.get_tensor)
+
+    def realize_samples(self):
+        """Read and return just the real "samples" tensor."""
+        return self._reader.get_tensor(self._skel["samples"][_TREF_KEY])
+
+    def __repr__(self):
+        info = self["samples"]
+        return f"LazyLatent(samples={tuple(info.shape)}, dtype={info.dtype})"
+
+
+class LazyConditioning:
+    """One dataset sample's conditioning on disk. Content is an arbitrary pickled
+    structure, so the only access is realize() -> list of [tensor, dict] entries."""
+
+    def __init__(self, reader, skeleton):
+        self._reader = reader
+        self._skel = skeleton
+
+    def realize(self):
+        """Read the full conditioning for this sample from disk."""
+        return _rejoin_tensors(self._skel, self._reader.get_tensor)
+
+    realize_entries = realize  # a realized conditioning IS its entry list
+
+    def __repr__(self):
+        return "LazyConditioning(on-disk)"
+
+
+class LazyCondEntry:
+    """One entry of a LazyConditioning — emitted by ResolutionBucket so each
+    bucket row pairs with exactly one conditioning entry."""
+
+    def __init__(self, lazy_cond, index):
+        self._cond = lazy_cond
+        self._index = index
+
+    def realize(self):
+        return self._cond.realize()[self._index]
+
+    def realize_entries(self):
+        return [self.realize()]
+
+    def __repr__(self):
+        return f"LazyCondEntry(index={self._index})"
+
+
+class LazyBatchSamples:
+    """The "samples" batch of one resolution bucket: N equal-shape rows backed by
+    on-disk LazyLatents (stored (1, *row_shape)), or already-real row tensors
+    when eager and lazy inputs are mixed. .shape/.dtype come from metadata;
+    realize_rows(indices) reads only the selected rows — the per-training-step
+    read unit."""
+
+    def __init__(self, rows):
+        self.rows = list(rows)
+        first = self.rows[0]
+        if isinstance(first, LazyLatent):
+            info = first["samples"]
+            row_shape, self.dtype = tuple(info.shape[1:]), info.dtype
+        else:
+            row_shape, self.dtype = tuple(first.shape), first.dtype
+        self.shape = torch.Size((len(self.rows), *row_shape))
+        self.ndim = len(self.shape)
+
+    def _row(self, i):
+        r = self.rows[int(i)]
+        return r.realize_samples()[0] if isinstance(r, LazyLatent) else r
+
+    def realize_rows(self, indices):
+        """Read only the selected rows; return them stacked (len(indices), *row_shape)."""
+        return torch.stack([self._row(i) for i in indices], dim=0)
+
+    def realize(self):
+        """Read and stack all rows: (N, *row_shape)."""
+        return self.realize_rows(range(len(self.rows)))
+
+    def __repr__(self):
+        return f"LazyBatchSamples(shape={tuple(self.shape)}, dtype={self.dtype})"
+
+
+_LAZY_DATASET_TYPES = (LazyLatent, LazyConditioning, LazyCondEntry, LazyBatchSamples)
+
+# Any op a lazy class doesn't define itself (indexing, iteration, math,
+# truthiness, pickling) raises RealizeRequired instead of silently misbehaving.
+for _cls in (LazyTensorInfo, *_LAZY_DATASET_TYPES):
+    for _op in ("__getitem__", "__iter__", "__len__", "__bool__", "__reduce__",
+                "__add__", "__radd__", "__sub__", "__rsub__", "__mul__", "__rmul__",
+                "__truediv__", "__matmul__", "__neg__"):
+        if _op not in _cls.__dict__:
+            setattr(_cls, _op, _need_realize)
+
+
+def _realize_structure(obj):
+    """Recursively replace lazy dataset objects with their realized (in-RAM)
+    values. Real tensors and plain values pass through unchanged."""
+    if isinstance(obj, _LAZY_DATASET_TYPES):
+        return obj.realize()
+    if isinstance(obj, dict):
+        return {k: _realize_structure(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [_realize_structure(v) for v in obj]
+    if isinstance(obj, tuple):
+        return tuple(_realize_structure(v) for v in obj)
+    return obj
+
+
+class _ShardReader:
+    """Random-access reader for a single shard.
+
+    Loads the small skeleton pickle eagerly; opens the safetensors file lazily
+    and uses safe_open's per-tensor random access so read_sample(i) only pulls
+    the tensors belonging to sample i. read_sample_lazy(i) pulls nothing — it
+    returns (LazyLatent, LazyConditioning) handles that read on demand.
+    """
+
+    def __init__(self, shard_path, skeleton_path):
+        with open(skeleton_path, "rb") as f:
+            self.skeletons = pickle.load(f)
+        self.shard_path = shard_path
+        self._st = None
+        self._header = None
+
+    def _open(self):
+        if self._st is None:
+            self._st = safe_open(self.shard_path, framework="pt")
+        return self._st
+
+    @property
+    def header(self):
+        if self._header is None:
+            self._header = _read_safetensors_header(self.shard_path)
+        return self._header
+
+    def shape(self, key):
+        return tuple(self.header[key]["shape"])
+
+    def dtype(self, key):
+        return _ST_STR_TO_DTYPE[self.header[key]["dtype"]]
+
+    def get_tensor(self, key):
+        return self._open().get_tensor(key)
+
+    def get_slice(self, key):
+        return self._open().get_slice(key)
+
+    def __len__(self):
+        return len(self.skeletons)
+
+    def read_sample(self, local_idx):
+        """Return (latent_dict, conditioning_list) for one sample, reading its
+        tensors eagerly."""
+        latent_skel, cond_skel = self.skeletons[local_idx]
+        st = self._open()
+        latent = _rejoin_tensors(latent_skel, st.get_tensor)
+        cond = _rejoin_tensors(cond_skel, st.get_tensor)
+        return latent, cond
+
+    def read_sample_lazy(self, local_idx):
+        """Return (LazyLatent, LazyConditioning) handles for one sample — no
+        tensor bytes are read. The handles carry the sample's skeleton, so
+        latent["samples"].shape/.dtype come from the safetensors header and
+        realize() reads only this sample's tensors."""
+        latent_skel, cond_skel = self.skeletons[local_idx]
+        return LazyLatent(self, latent_skel), LazyConditioning(self, cond_skel)
+
+
 class ResolutionBucket(io.ComfyNode):
     """Bucket latents and conditions by resolution for efficient batch training."""
 
@@ -1274,8 +1557,9 @@ class ResolutionBucket(io.ComfyNode):
 
     @classmethod
     def execute(cls, latents, conditioning):
-        # latents: list[{"samples": tensor}] where tensor is (B, C, H, W), typically B=1
+        # latents: list of latent dicts {"samples": (B, C, H, W)} and/or LazyLatent
-        # conditioning: list[list[cond]]
+        # conditioning: list of conds (each a list of [tensor, dict] entries)
+        #               and/or LazyConditioning
 
         # Validate lengths match
         if len(latents) != len(conditioning):
@@ -1283,50 +1567,56 @@ class ResolutionBucket(io.ComfyNode):
                 f"Number of latents ({len(latents)}) does not match number of conditions ({len(conditioning)})."
             )
 
-        # Flatten latents and conditions to individual samples
+        # Group rows by (H, W). Lazy latents are grouped by header metadata only
-        flat_latents = []  # list of (C, H, W) tensors
+        # (no tensor bytes read); buckets with any lazy row become LazyBatchSamples.
-        flat_conditions = []  # list of condition lists
+        buckets = {}  # (h, w) -> {"rows": [...], "conds": [...]}
+        any_lazy = False
 
-        for latent_dict, cond in zip(latents, conditioning):
+        for latent, cond in zip(latents, conditioning):
-            samples = latent_dict["samples"]  # (B, C, H, W)
+            if isinstance(latent, LazyLatent):
-            batch_size = samples.shape[0]
+                info = latent["samples"]
-
+                if int(info.shape[0]) != 1:
-            # cond is a list of conditions with length == batch_size
+                    raise RealizeRequired(
-            for i in range(batch_size):
+                        "ResolutionBucket: lazy latents with stored batch size > 1 "
-                flat_latents.append(samples[i])  # (C, H, W)
+                        "are not supported; insert a Realize Lazy Latents node first."
-                flat_conditions.append(cond[i])  # single condition
+                    )
-
+                any_lazy = True
-        # Group by resolution (H, W)
+                h, w = int(info.shape[-2]), int(info.shape[-1])
-        buckets = {}  # (H, W) -> {"latents": list, "conditions": list}
+                bucket = buckets.setdefault((h, w), {"rows": [], "conds": []})
-
+                bucket["rows"].append(latent)
-        for latent, cond in zip(flat_latents, flat_conditions):
+                bucket["conds"].append(
-            # latent shape is (..., H, W) (B, C, H, W) or (B, T, C, H ,W)
+                    LazyCondEntry(cond, 0) if isinstance(cond, LazyConditioning) else cond[0]
-            h, w = latent.shape[-2], latent.shape[-1]
+                )
-            key = (h, w)
+            else:
-
+                samples = latent["samples"]  # (B, C, H, W) real tensor
-            if key not in buckets:
+                h, w = int(samples.shape[-2]), int(samples.shape[-1])
-                buckets[key] = {"latents": [], "conditions": []}
+                bucket = buckets.setdefault((h, w), {"rows": [], "conds": []})
-
+                # cond is a list of entries with length == batch size
-            buckets[key]["latents"].append(latent)
+                for i in range(samples.shape[0]):
-            buckets[key]["conditions"].append(cond)
+                    bucket["rows"].append(samples[i])
-
+                    bucket["conds"].append(
-        # Convert buckets to output format
+                        LazyCondEntry(cond, i) if isinstance(cond, LazyConditioning) else cond[i]
-        output_latents = []  # list[{"samples": tensor}] where tensor is (Bi, ..., H, W)
-        output_conditions = []  # list[list[cond]] where each inner list has Bi conditions
-
-        for (h, w), bucket_data in buckets.items():
-            # Stack latents into batch: list of (..., H, W) -> (Bi, ..., H, W)
-            stacked_latents = torch.stack(bucket_data["latents"], dim=0)
-            output_latents.append({"samples": stacked_latents})
-
-            # Conditions stay as list of condition lists
-            output_conditions.append(bucket_data["conditions"])
-
-            logging.info(
-                f"Resolution bucket ({h}x{w}): {len(bucket_data['latents'])} samples"
                     )
 
-        logging.info(f"Created {len(buckets)} resolution buckets from {len(flat_latents)} samples")
+        output_latents = []      # list[{"samples": (Bi, *row_shape)}]
+        output_conditions = []   # list[list[cond entry]] with Bi entries each
+
+        total = 0
+        for (h, w), bucket_data in buckets.items():
+            rows = bucket_data["rows"]
+            total += len(rows)
+            if any(isinstance(r, LazyLatent) for r in rows):
+                samples = LazyBatchSamples(rows)
+            else:
+                samples = torch.stack(rows, dim=0)
+            output_latents.append({"samples": samples})
+            output_conditions.append(bucket_data["conds"])
+            logging.info(f"Resolution bucket ({h}x{w}): {len(rows)} samples")
+
+        logging.info(
+            f"Created {len(buckets)} resolution buckets from {total} samples "
+            f"({'lazy' if any_lazy else 'eager'})"
+        )
         return io.NodeOutput(output_latents, output_conditions)
 
 
@@ -1464,7 +1754,8 @@ class SaveTrainingDataset(io.ComfyNode):
         shard_size = shard_size[0]
 
         # latents: list[{"samples": tensor}]
-        # conditioning: list[list[cond]]
+        # conditioning: list[list[[cond_tensor, dict]]]   (encode_from_tokens_scheduled output;
+        #   dicts may contain arbitrary extension types — Hook objects, floats, strings, etc.)
 
         # Validate lengths match
         if len(latents) != len(conditioning):
@@ -1473,45 +1764,55 @@ class SaveTrainingDataset(io.ComfyNode):
                 f"Something went wrong in dataset preparation."
             )
 
-        # Create output directory
+        # [TODO] can save to anywhere <- need to be resolve
         output_dir = os.path.join(folder_paths.get_output_directory(), folder_name)
         os.makedirs(output_dir, exist_ok=True)
 
-        # Prepare data pairs
         num_samples = len(latents)
-        num_shards = (num_samples + shard_size - 1) // shard_size  # Ceiling division
+        num_shards = (num_samples + shard_size - 1) // shard_size
 
         logging.info(
             f"Saving {num_samples} samples to {num_shards} shards in {output_dir}..."
         )
 
-        # Save data in shards
         for shard_idx in range(num_shards):
             start_idx = shard_idx * shard_size
             end_idx = min(start_idx + shard_size, num_samples)
 
-            # Get shard data (list of latent dicts and conditioning lists)
+            # Per shard: one safetensors holding every tensor (bulk bytes, partial-loadable)
-            shard_data = {
+            # plus one .skeleton.pkl holding the nested-structure shells with __tref__ markers.
-                "latents": latents[start_idx:end_idx],
+            shard_tensors = {}
-                "conditioning": conditioning[start_idx:end_idx],
+            shard_skeletons = []  # list of (latent_skeleton, cond_skeleton) per sample
-            }
 
-            # Save shard
+            for local_idx, i in enumerate(range(start_idx, end_idx)):
-            shard_filename = f"shard_{shard_idx:04d}.pkl"
+                # Lazy inputs are realized per sample; at most one shard is in RAM.
-            shard_path = os.path.join(output_dir, shard_filename)
+                latent_skel = _split_tensors(
+                    _realize_structure(latents[i]), shard_tensors, f"s{local_idx}_lat"
+                )
+                cond_skel = _split_tensors(
+                    _realize_structure(conditioning[i]), shard_tensors, f"s{local_idx}_cond"
+                )
+                shard_skeletons.append((latent_skel, cond_skel))
 
-            with open(shard_path, "wb") as f:
+            shard_path = os.path.join(output_dir, f"shard_{shard_idx:04d}.safetensors")
-                torch.save(shard_data, f)
+            skeleton_path = os.path.join(
-
+                output_dir, f"shard_{shard_idx:04d}.skeleton.pkl"
-            logging.info(
+            )
-                f"Saved shard {shard_idx + 1}/{num_shards}: {shard_filename} ({end_idx - start_idx} samples)"
+
+            safetensors.torch.save_file(shard_tensors, shard_path)
+            with open(skeleton_path, "wb") as f:
+                pickle.dump(shard_skeletons, f, protocol=pickle.HIGHEST_PROTOCOL)
+
+            logging.info(
+                f"Saved shard {shard_idx + 1}/{num_shards}: {end_idx - start_idx} samples, "
+                f"{len(shard_tensors)} tensors"
             )
 
-        # Save metadata
         metadata = {
             "num_samples": num_samples,
             "num_shards": num_shards,
             "shard_size": shard_size,
+            "format_version": 2,
         }
         metadata_path = os.path.join(output_dir, "metadata.json")
         with open(metadata_path, "w") as f:
@@ -1522,15 +1823,22 @@ class SaveTrainingDataset(io.ComfyNode):
 
 
 class LoadTrainingDataset(io.ComfyNode):
-    """Load encoded training dataset from disk."""
+    """Load encoded training dataset from disk as lazy references.
+
+    Outputs list[LazyLatent] and list[LazyConditioning] — one handle per sample,
+    near-zero RAM. Latent shapes/dtypes are readable from metadata (e.g. by
+    Resolution Bucket) without any I/O; tensor bytes are read per batch inside
+    the lazy-aware trainer. For any other consumer, insert the Realize Lazy
+    Latents / Realize Lazy Conditionings nodes to get standard in-RAM data.
+    """
     @classmethod
     def define_schema(cls):
         return io.Schema(
             node_id="LoadTrainingDataset",
-            search_aliases=["import dataset", "training data"],
+            search_aliases=["import dataset", "training data", "lazy", "streaming"],
             display_name="Load Training Dataset",
             category="model/training",
-            description="Load encoded training dataset (latents + conditioning) from disk for use in training.",
+            description="Load an encoded training dataset from disk as lazy references; tensors are read on demand during training instead of all at once.",
             is_experimental=True,
             inputs=[
                 io.String.Input(
@@ -1555,47 +1863,127 @@ class LoadTrainingDataset(io.ComfyNode):
 
     @classmethod
     def execute(cls, folder_name):
-        # Get dataset directory
         dataset_dir = os.path.join(folder_paths.get_output_directory(), folder_name)
 
         if not os.path.exists(dataset_dir):
             raise ValueError(f"Dataset directory not found: {dataset_dir}")
 
-        # Find all shard files
         shard_files = sorted(
-            [
             f
             for f in os.listdir(dataset_dir)
-                if f.startswith("shard_") and f.endswith(".pkl")
+            if f.startswith("shard_") and f.endswith(".safetensors")
-            ]
         )
 
         if not shard_files:
-            raise ValueError(f"No shard files found in {dataset_dir}")
+            raise ValueError(
+                f"No shard files found in {dataset_dir} "
+                f"(expected shard_*.safetensors + shard_*.skeleton.pkl)."
+            )
 
-        logging.info(f"Loading {len(shard_files)} shards from {dataset_dir}...")
+        logging.info(f"Lazy-loading {len(shard_files)} shards from {dataset_dir}...")
 
-        # Load all shards
+        all_latents = []        # list[LazyLatent]
-        all_latents = []  # list[{"samples": tensor}]
+        all_conditioning = []   # list[LazyConditioning]
-        all_conditioning = []  # list[list[cond]]
 
         for shard_file in shard_files:
             shard_path = os.path.join(dataset_dir, shard_file)
+            skeleton_path = os.path.join(
+                dataset_dir, shard_file[: -len(".safetensors")] + ".skeleton.pkl"
+            )
 
-            with open(shard_path, "rb") as f:
+            # Reads only the skeleton pickle + safetensors header, no tensor bytes.
-                shard_data = torch.load(f, weights_only=True)
+            reader = _ShardReader(shard_path, skeleton_path)
+            for local_idx in range(len(reader)):
+                latent, cond = reader.read_sample_lazy(local_idx)
+                all_latents.append(latent)
+                all_conditioning.append(cond)
 
-            all_latents.extend(shard_data["latents"])
+            logging.info(f"Indexed {shard_file}: {len(reader)} samples")
-            all_conditioning.extend(shard_data["conditioning"])
-
-            logging.info(f"Loaded {shard_file}: {len(shard_data['latents'])} samples")
 
         logging.info(
-            f"Successfully loaded {len(all_latents)} samples from {dataset_dir}."
+            f"Lazy-loaded {len(all_latents)} samples from {dataset_dir} "
+            f"(no tensor data read yet)."
         )
         return io.NodeOutput(all_latents, all_conditioning)
 
 
+class RealizeLazyLatents(io.ComfyNode):
+    """Read all lazy latent tensors from disk into RAM, producing standard latent
+    dicts.
+
+    Insert before any node that is not lazy-aware (one that stacks or does tensor
+    math on the latents). Real latents pass through unchanged, so it is safe to
+    apply unconditionally.
+    """
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="RealizeLazyLatents",
+            search_aliases=["realize", "materialize", "load to ram", "realize latents"],
+            display_name="Realize Lazy Latents",
+            category="model/training",
+            description="Read all lazy latent tensors from disk into memory, producing standard in-RAM latent dicts.",
+            is_experimental=True,
+            is_input_list=True,
+            inputs=[
+                io.Latent.Input("latents", tooltip="Lazy (or real) latent dicts."),
+            ],
+            outputs=[
+                io.Latent.Output(
+                    display_name="latents",
+                    is_output_list=True,
+                    tooltip="Realized (in-RAM) latent dicts",
+                ),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, latents):
+        real_latents = [_realize_structure(x) for x in latents]
+        logging.info(f"Realized {len(real_latents)} latents into RAM.")
+        return io.NodeOutput(real_latents)
+
+
+class RealizeLazyConditionings(io.ComfyNode):
+    """Read all lazy conditioning tensors from disk into RAM, producing standard
+    conditioning.
+
+    Insert before any node that is not lazy-aware. Real conditioning passes
+    through unchanged, so it is safe to apply unconditionally.
+    """
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="RealizeLazyConditionings",
+            search_aliases=["realize", "materialize", "load to ram", "realize conditioning"],
+            display_name="Realize Lazy Conditionings",
+            category="model/training",
+            description="Read all lazy conditioning tensors from disk into memory, producing standard in-RAM conditioning.",
+            is_experimental=True,
+            is_input_list=True,
+            inputs=[
+                io.Conditioning.Input(
+                    "conditioning", tooltip="Lazy (or real) conditioning."
+                ),
+            ],
+            outputs=[
+                io.Conditioning.Output(
+                    display_name="conditioning",
+                    is_output_list=True,
+                    tooltip="Realized (in-RAM) conditioning",
+                ),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, conditioning):
+        real_conditioning = [_realize_structure(x) for x in conditioning]
+        logging.info(f"Realized {len(real_conditioning)} conditionings into RAM.")
+        return io.NodeOutput(real_conditioning)
+
+
 # ========== Extension Setup ==========
 
 
@@ -1635,6 +2023,8 @@ class DatasetExtension(ComfyExtension):
             MakeTrainingDataset,
             SaveTrainingDataset,
             LoadTrainingDataset,
+            RealizeLazyLatents,
+            RealizeLazyConditionings,
             ResolutionBucket,
         ]
 

comfy_extras/nodes_train.py

@@ -1,5 +1,7 @@
 import logging
 import os
+import sys
+import importlib.util
 
 import numpy as np
 import safetensors
@@ -26,6 +28,33 @@ from comfy_api.latest import ComfyExtension, io, ui
 from comfy.utils import ProgressBar
 
 
+def _import_like_node_loader(filename):
+    path = os.path.join(os.path.dirname(__file__), filename)
+    key = os.path.splitext(path)[0]  # exactly the key load_custom_node uses
+    module = sys.modules.get(key)
+    if module is None:
+        spec = importlib.util.spec_from_file_location(key, path)
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[key] = module
+        spec.loader.exec_module(module)
+    return module
+
+
+_nodes_dataset = _import_like_node_loader("nodes_dataset.py")
+LazyLatent = _nodes_dataset.LazyLatent
+LazyBatchSamples = _nodes_dataset.LazyBatchSamples
+RealizeRequired = _nodes_dataset.RealizeRequired
+
+
+def _realize_latent(x):
+    """Return a real samples tensor from a LazyLatent (reads this one sample
+    from disk) or pass an already-real tensor through unchanged. This is the
+    per-sample realize point of the streaming training path."""
+    if isinstance(x, LazyLatent):
+        return x.realize_samples()
+    return x
+
+
 class TrainGuider(comfy_extras.nodes_custom_sampler.Guider_Basic):
     """
     CFGGuider with modifications for training specific logic
@@ -146,6 +175,7 @@ class TrainSampler(comfy.samplers.Sampler):
         real_dataset=None,
         bucket_latents=None,
         use_grad_scaler=False,
+        lazy_conds=None,
     ):
         self.loss_fn = loss_fn
         self.optimizer = optimizer
@@ -160,6 +190,9 @@ class TrainSampler(comfy.samplers.Sampler):
         self.bucket_latents: list[torch.Tensor] | None = (
             bucket_latents  # list of (Bi, C, Hi, Wi)
         )
+        # When set (one lazy cond per sample), conditioning is realized per batch
+        # in fwd_bwd instead of up front in the guider.
+        self.lazy_conds = lazy_conds
         # GradScaler for fp16 training
         self.grad_scaler = torch.amp.GradScaler() if use_grad_scaler else None
         # Precompute bucket offsets and weights for sampling
@@ -181,6 +214,26 @@ class TrainSampler(comfy.samplers.Sampler):
         # Weights for sampling buckets proportional to their size
         self.bucket_weights = torch.tensor(bucket_sizes, dtype=torch.float32)
 
+    def _build_batch_conds(self, model_wrap, indicies, batch_noise, batch_latent):
+        """Realize the sampled conditioning entries from disk and run the standard
+        convert_cond + process_conds pass on them, bounded to this batch."""
+        entries = []
+        for i in indicies:
+            c = self.lazy_conds[i]
+            if hasattr(c, "realize_entries"):
+                entries.extend(c.realize_entries())
+            else:
+                entries.append(c)  # already-real [tensor, dict] entry
+        converted = comfy.sampler_helpers.convert_cond(entries)
+        processed = comfy.samplers.process_conds(
+            model_wrap.inner_model,
+            batch_noise,
+            {"positive": converted},
+            batch_noise.device,
+            latent_image=batch_latent,
+        )
+        return processed["positive"]
+
     def fwd_bwd(
         self,
         model_wrap,
@@ -203,6 +256,11 @@ class TrainSampler(comfy.samplers.Sampler):
             False,
         )
 
+        if self.lazy_conds is not None:
+            model_wrap.conds["positive"] = self._build_batch_conds(
+                model_wrap, indicies, batch_noise, batch_latent
+            )
+        else:
             model_wrap.conds["positive"] = [cond[i] for i in indicies]
         batch_extra_args = make_batch_extra_option_dict(
             extra_args, indicies, full_size=dataset_size
@@ -247,6 +305,10 @@ class TrainSampler(comfy.samplers.Sampler):
         # Convert to absolute indices for fwd_bwd (cond is flattened, use absolute index)
         absolute_indices = [bucket_offset + idx for idx in relative_indices]
 
+        if isinstance(bucket_latent, LazyBatchSamples):
+            # Reads only this batch's rows from disk.
+            batch_latent = bucket_latent.realize_rows(relative_indices).to(latent_image)
+        else:
             batch_latent = bucket_latent[relative_indices].to(latent_image)  # (actual_batch_size, C, H, W)
         batch_noise = noisegen.generate_noise({"samples": batch_latent}).to(
             batch_latent.device
@@ -297,7 +359,8 @@ class TrainSampler(comfy.samplers.Sampler):
         indicies = torch.randperm(dataset_size)[: self.batch_size].tolist()
         total_loss = 0
         for index in indicies:
-            single_latent = self.real_dataset[index].to(latent_image)
+            # Realize one sample at a time (reads from disk for a lazy dataset).
+            single_latent = _realize_latent(self.real_dataset[index]).to(latent_image)
             batch_noise = noisegen.generate_noise(
                 {"samples": single_latent}
             ).to(single_latent.device)
@@ -540,13 +603,20 @@ def _process_latents_bucket_mode(latents):
     """Process latents for bucket mode training.
 
     Args:
-        latents: list[{"samples": tensor}] where each tensor is (Bi, C, Hi, Wi)
+        latents: list[{"samples": tensor | LazyBatchSamples}] per bucket,
+            each (Bi, C, Hi, Wi)
 
     Returns:
-        list of latent tensors
+        list of bucket batches (tensor or LazyBatchSamples)
     """
     bucket_latents = []
     for latent_dict in latents:
+        if isinstance(latent_dict, LazyLatent):
+            raise RealizeRequired(
+                "bucket_mode expects Resolution Bucket output, but got raw lazy "
+                "latents. Insert a Resolution Bucket node (it is lazy-aware) "
+                "between the dataset loader and the trainer."
+            )
         bucket_latents.append(latent_dict["samples"])  # (Bi, C, Hi, Wi)
     return bucket_latents
 
@@ -555,16 +625,28 @@ def _process_latents_standard_mode(latents):
     """Process latents for standard (non-bucket) mode training.
 
     Args:
-        latents: list of latent dicts or single latent dict
+        latents: list of latent dicts and/or LazyLatent handles
 
     Returns:
-        Processed latents (tensor or list of tensors)
+        Processed latents (tensor, or list of tensors / LazyLatent handles)
     """
     if len(latents) == 1:
-        return latents[0]["samples"]  # Single latent dict
+        only = latents[0]
+        if isinstance(only, LazyLatent):
+            return [only]
+        return only["samples"]  # Single latent dict
 
     latent_list = []
     for latent in latents:
+        if isinstance(latent, LazyLatent):
+            # Kept as a handle; realized one sample at a time in the train loop.
+            if int(latent["samples"].shape[0]) != 1:
+                raise RealizeRequired(
+                    "Lazy latents with stored batch size > 1 are not supported in "
+                    "the streaming path; insert a Realize Lazy Latents node first."
+                )
+            latent_list.append(latent)
+            continue
         latent = latent["samples"]
         bs = latent.shape[0]
         if bs != 1:
@@ -579,15 +661,18 @@ def _process_conditioning(positive):
     """Process conditioning - either single list or list of lists.
 
     Args:
-        positive: list of conditioning
+        positive: list of conditioning (cond entry lists and/or LazyConditioning)
 
     Returns:
         Flattened conditioning list
     """
     if len(positive) == 1:
-        return positive[0]  # Single conditioning list
+        only = positive[0]
+        if hasattr(only, "realize_entries"):
+            return [only]  # a lazy cond is one sample, not a list to unwrap
+        return only  # Single conditioning list
 
-    # Multiple conditioning lists - flatten
+    # Multiple conditioning lists - flatten (lazy handles stay whole)
     flat_positive = []
     for cond in positive:
         if isinstance(cond, list):
@@ -609,19 +694,34 @@ def _prepare_latents_and_count(latents, dtype, bucket_mode):
         tuple: (processed_latents, num_images, multi_res)
     """
     if bucket_mode:
-        # In bucket mode, latents is list of tensors (Bi, C, Hi, Wi)
+        # latents: list of bucket batches (Bi, C, Hi, Wi). LazyBatchSamples stay
-        latents = [t.to(dtype) for t in latents]
+        # lazy; their rows are read and cast per training step.
         num_buckets = len(latents)
-        num_images = sum(t.shape[0] for t in latents)
+        num_images = sum(int(t.shape[0]) for t in latents)
+        latents = [t if isinstance(t, LazyBatchSamples) else t.to(dtype) for t in latents]
         multi_res = False  # Not using multi_res path in bucket mode
 
         logging.debug(f"Bucket mode: {num_buckets} buckets, {num_images} total samples")
         for i, lat in enumerate(latents):
-            logging.debug(f"  Bucket {i}: shape {lat.shape}")
+            logging.debug(f"  Bucket {i}: shape {tuple(lat.shape)}")
         return latents, num_images, multi_res
 
-    # Non-bucket mode
+    # Non-bucket mode. A single lazy handle becomes a one-element per-sample list.
+    if isinstance(latents, LazyLatent):
+        latents = [latents]
+
     if isinstance(latents, list):
+        if any(isinstance(t, LazyLatent) for t in latents):
+            # Lazy: route to the per-sample (multi_res) path; samples are
+            # realized on demand in the train loop.
+            num_images = len(latents)
+            logging.info(
+                f"Lazy dataset: {num_images} samples will stream from disk one "
+                f"sample at a time. For batched streaming, insert a Resolution "
+                f"Bucket node and enable bucket_mode."
+            )
+            return latents, num_images, True
+
         all_shapes = set()
         latents = [t.to(dtype) for t in latents]
         for latent in latents:
@@ -905,7 +1005,7 @@ def _create_loss_function(loss_function_name):
 
 
 def _run_training_loop(
-    guider, train_sampler, latents, num_images, seed, bucket_mode, multi_res
+    guider, train_sampler, latents, num_images, seed, bucket_mode, multi_res, dtype=None
 ):
     """Execute the training loop.
 
@@ -917,13 +1017,18 @@ def _run_training_loop(
         seed: Random seed
         bucket_mode: Whether bucket mode is enabled
         multi_res: Whether multi-resolution mode is enabled
+        dtype: dtype for the dummy latent (lazy data is stored uncast on disk)
     """
     sigmas = torch.tensor(range(num_images))
     noise = comfy_extras.nodes_custom_sampler.Noise_RandomNoise(seed)
 
     if bucket_mode:
-        # Use first bucket's first latent as dummy for guider
+        # Use first bucket's first latent as dummy for guider (one disk read if lazy)
-        dummy_latent = latents[0][:1].repeat(num_images, 1, 1, 1)
+        first = latents[0]
+        row = first.realize_rows([0]) if isinstance(first, LazyBatchSamples) else first[:1]
+        if dtype is not None:
+            row = row.to(dtype)
+        dummy_latent = row.repeat(num_images, 1, 1, 1)
         guider.sample(
             noise.generate_noise({"samples": dummy_latent}),
             dummy_latent,
@@ -932,8 +1037,11 @@ def _run_training_loop(
             seed=noise.seed,
         )
     elif multi_res:
-        # use first latent as dummy latent if multi_res
+        # use first latent as dummy latent if multi_res (one disk read if lazy)
-        latents = latents[0].repeat(num_images, 1, 1, 1)
+        row = _realize_latent(latents[0])
+        if dtype is not None:
+            row = row.to(dtype)
+        latents = row.repeat(num_images, 1, 1, 1)
         guider.sample(
             noise.generate_noise({"samples": latents}),
             latents,
@@ -1233,6 +1341,17 @@ class TrainLoraNode(io.ComfyNode):
             def loss_callback(loss):
                 loss_map["loss"].append(loss)
 
+            # Lazy conds are realized per batch in the train loop; the guider
+            # only needs one realized template cond to initialize.
+            lazy_conds = None
+            guider_positive = positive
+            if any(hasattr(c, "realize_entries") for c in positive):
+                lazy_conds = positive
+                first = positive[0]
+                guider_positive = (
+                    first.realize_entries() if hasattr(first, "realize_entries") else [first]
+                )
+
             # Create sampler
             if bucket_mode:
                 train_sampler = TrainSampler(
@@ -1246,6 +1365,7 @@ class TrainLoraNode(io.ComfyNode):
                     training_dtype=dtype,
                     bucket_latents=latents,
                     use_grad_scaler=use_grad_scaler,
+                    lazy_conds=lazy_conds,
                 )
             else:
                 train_sampler = TrainSampler(
@@ -1259,11 +1379,12 @@ class TrainLoraNode(io.ComfyNode):
                     training_dtype=dtype,
                     real_dataset=latents if multi_res else None,
                     use_grad_scaler=use_grad_scaler,
+                    lazy_conds=lazy_conds,
                 )
 
             # Setup guider
             guider = TrainGuider(mp, offloading=offloading)
-            guider.set_conds(positive)
+            guider.set_conds(guider_positive)
 
             # Inject bypass hooks if bypass mode is enabled
             bypass_injections = None
@@ -1284,6 +1405,7 @@ class TrainLoraNode(io.ComfyNode):
                     seed,
                     bucket_mode,
                     multi_res,
+                    dtype=latents_dtype,
                 )
             finally:
                 comfy.model_management.in_training = False

nodes.py

@@ -2402,8 +2402,8 @@ async def init_builtin_extra_nodes():
         "nodes_ideogram4.py",
         "nodes_bounding_boxes.py",
         "nodes_json_prompt.py",
-        "nodes_train.py",
         "nodes_dataset.py",
+        "nodes_train.py",
         "nodes_sag.py",
         "nodes_perpneg.py",
         "nodes_stable3d.py",