mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-10 00:17:16 +08:00
Merge 484081e00c into 51bf508a0b
This commit is contained in:
Kohaku-Blueleaf
GitHub
commit
7bbdfcd33a
@ -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
|
||||
# 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,
|
||||
]
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user