diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index 73fe75b7f..3b9a3c267 100644 --- a/comfy_extras/nodes_dataset.py +++ b/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(" 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 - # conditioning: list[list[cond]] + # latents: list of latent dicts {"samples": (B, C, H, W)} and/or LazyLatent + # 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 - flat_latents = [] # list of (C, H, W) tensors - flat_conditions = [] # list of condition lists + # Group rows by (H, W). Lazy latents are grouped by header metadata only + # (no tensor bytes read); buckets with any lazy row become LazyBatchSamples. + buckets = {} # (h, w) -> {"rows": [...], "conds": [...]} + any_lazy = False - for latent_dict, cond in zip(latents, conditioning): - samples = latent_dict["samples"] # (B, C, H, W) - batch_size = samples.shape[0] + for latent, cond in zip(latents, conditioning): + if isinstance(latent, LazyLatent): + info = latent["samples"] + if int(info.shape[0]) != 1: + raise RealizeRequired( + "ResolutionBucket: lazy latents with stored batch size > 1 " + "are not supported; insert a Realize Lazy Latents node first." + ) + any_lazy = True + h, w = int(info.shape[-2]), int(info.shape[-1]) + bucket = buckets.setdefault((h, w), {"rows": [], "conds": []}) + bucket["rows"].append(latent) + bucket["conds"].append( + LazyCondEntry(cond, 0) if isinstance(cond, LazyConditioning) else cond[0] + ) + else: + samples = latent["samples"] # (B, C, H, W) real tensor + h, w = int(samples.shape[-2]), int(samples.shape[-1]) + bucket = buckets.setdefault((h, w), {"rows": [], "conds": []}) + # cond is a list of entries with length == batch size + for i in range(samples.shape[0]): + bucket["rows"].append(samples[i]) + bucket["conds"].append( + LazyCondEntry(cond, i) if isinstance(cond, LazyConditioning) else cond[i] + ) - # cond is a list of conditions with length == batch_size - for i in range(batch_size): - flat_latents.append(samples[i]) # (C, H, W) - flat_conditions.append(cond[i]) # single condition - - # Group by resolution (H, W) - buckets = {} # (H, W) -> {"latents": list, "conditions": list} - - for latent, cond in zip(flat_latents, flat_conditions): - # latent shape is (..., H, W) (B, C, H, W) or (B, T, C, H ,W) - h, w = latent.shape[-2], latent.shape[-1] - key = (h, w) - - if key not in buckets: - buckets[key] = {"latents": [], "conditions": []} - - buckets[key]["latents"].append(latent) - buckets[key]["conditions"].append(cond) - - # Convert buckets to output format - output_latents = [] # list[{"samples": tensor}] where tensor is (Bi, ..., H, W) - output_conditions = [] # list[list[cond]] where each inner list has Bi conditions + 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(): - # 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}) + 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") - # 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") + 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) - shard_data = { - "latents": latents[start_idx:end_idx], - "conditioning": conditioning[start_idx:end_idx], - } + # Per shard: one safetensors holding every tensor (bulk bytes, partial-loadable) + # plus one .skeleton.pkl holding the nested-structure shells with __tref__ markers. + shard_tensors = {} + shard_skeletons = [] # list of (latent_skeleton, cond_skeleton) per sample - # Save shard - shard_filename = f"shard_{shard_idx:04d}.pkl" - shard_path = os.path.join(output_dir, shard_filename) + for local_idx, i in enumerate(range(start_idx, end_idx)): + # Lazy inputs are realized per sample; at most one shard is in RAM. + 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: - torch.save(shard_data, f) - - logging.info( - f"Saved shard {shard_idx + 1}/{num_shards}: {shard_filename} ({end_idx - start_idx} samples)" + shard_path = os.path.join(output_dir, f"shard_{shard_idx:04d}.safetensors") + skeleton_path = os.path.join( + output_dir, f"shard_{shard_idx:04d}.skeleton.pkl" + ) + + 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") - ] + f + for f in os.listdir(dataset_dir) + 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[{"samples": tensor}] - all_conditioning = [] # list[list[cond]] + all_latents = [] # list[LazyLatent] + all_conditioning = [] # list[LazyConditioning] 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: - shard_data = torch.load(f, weights_only=True) + # Reads only the skeleton pickle + safetensors header, no tensor bytes. + 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"]) - all_conditioning.extend(shard_data["conditioning"]) - - logging.info(f"Loaded {shard_file}: {len(shard_data['latents'])} samples") + logging.info(f"Indexed {shard_file}: {len(reader)} 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, ] diff --git a/comfy_extras/nodes_train.py b/comfy_extras/nodes_train.py index a27217b80..ef4416e24 100644 --- a/comfy_extras/nodes_train.py +++ b/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,7 +256,12 @@ class TrainSampler(comfy.samplers.Sampler): False, ) - model_wrap.conds["positive"] = [cond[i] for i in indicies] + 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,7 +305,11 @@ 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] - batch_latent = bucket_latent[relative_indices].to(latent_image) # (actual_batch_size, C, H, W) + 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 = [t.to(dtype) for t in latents] + # latents: list of bucket batches (Bi, C, Hi, Wi). LazyBatchSamples stay + # 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 - dummy_latent = latents[0][:1].repeat(num_images, 1, 1, 1) + # Use first bucket's first latent as dummy for guider (one disk read if lazy) + 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 - latents = latents[0].repeat(num_images, 1, 1, 1) + # use first latent as dummy latent if multi_res (one disk read if lazy) + 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 diff --git a/nodes.py b/nodes.py index 9043a8d0a..0b64b35ab 100644 --- a/nodes.py +++ b/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",