Claude
45b60ac874
Tags from the tagger are added after a comma and the caption is never reworded. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MpnF7GwNGMbngU11wjym7P
393 lines
18 KiB
Python
393 lines
18 KiB
Python
"""The training run: the set's media become latents, its captions become conditioning, ComfyUI's own trainer does the learning.
|
|
|
|
Blocking, and run on ComfyUI's worker. Everything that needs torch or comfy is imported inside the
|
|
functions, so the rest of the module (sets, tags, the editor's routes) loads without them.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from ..._core import log
|
|
from . import captions, jobs, media_io, store
|
|
|
|
POSITIVE = ("own", "liked")
|
|
NEGATIVE = ("disliked", "bad_composition", "bad_image")
|
|
USES = {"normal": "liked and your own media", "avoid": "disliked results, to train an 'avoid' LoRA", "all": "everything in the set"}
|
|
ALGORITHMS = ("LoRA", "LoHa", "LoKr", "OFT")
|
|
|
|
|
|
class Stopped(Exception):
|
|
"""The person pressed Stop."""
|
|
|
|
|
|
def pick(doc: dict, use: str) -> list[dict]:
|
|
"""The items a run trains on. 'normal' leaves out what was disliked, 'avoid' keeps only that, 'all' keeps everything."""
|
|
keep = {"normal": POSITIVE, "avoid": NEGATIVE, "all": POSITIVE + NEGATIVE}.get(use)
|
|
if keep is None:
|
|
raise ValueError(f"use is one of {sorted(USES)}, not {use!r}")
|
|
return [i for i in doc["items"] if i["label"] in keep]
|
|
|
|
|
|
def roots() -> dict:
|
|
import folder_paths
|
|
return {"output": folder_paths.get_output_directory(), "temp": folder_paths.get_temp_directory(),
|
|
"input": folder_paths.get_input_directory()}
|
|
|
|
|
|
def bin_path(media_id: str):
|
|
from ..._core import media
|
|
try:
|
|
return media.path_for(media_id)
|
|
except Exception: # noqa: BLE001 -- a bin entry that is gone is a missing file, said by the caller
|
|
return None
|
|
|
|
|
|
def plan_items(set_name: str, use: str):
|
|
"""(items with their files, problems). Nothing here touches torch."""
|
|
doc = store.get(set_name)
|
|
if doc is None:
|
|
return [], [f"there is no training set called {set_name!r}"]
|
|
problems, found = [], []
|
|
chosen = pick(doc, use)
|
|
if not chosen:
|
|
problems.append(f"the set {set_name!r} has nothing to train on for 'use: {use}' ({USES[use]})")
|
|
for item in chosen:
|
|
path = store.resolve(item["src"], roots(), bin_path)
|
|
if path is None:
|
|
problems.append(f"{item['name'] or item['id']} is no longer there")
|
|
continue
|
|
found.append((item, path))
|
|
return found, problems
|
|
|
|
|
|
def _encode_text(clip, text: str):
|
|
tokens = clip.tokenize(text)
|
|
return clip.encode_from_tokens_scheduled(tokens)
|
|
|
|
|
|
def _encode_pixels(vae, pixels):
|
|
"""Pixels [N,H,W,3] -> latent tensor, in the VAE's own layout."""
|
|
return vae.encode(pixels[:, :, :, :3])
|
|
|
|
|
|
def encode_item(item, path, size, cfg, vae, adapt, audio_vae):
|
|
"""One item's latents. For video models a clip stays a clip (cfg['frames'] frames); for picture models a clip gives one frame."""
|
|
import torch
|
|
width, height = size
|
|
if adapt and adapt.get("video"):
|
|
frames = adapt["frames"](cfg["frames"])
|
|
if item["media_kind"] == "video":
|
|
pixels, rate = media_io.clip_frames(path, frames, cfg["fps"] or adapt.get("fps", 24), cfg["start"], width, height)
|
|
else:
|
|
from PIL import Image, ImageOps
|
|
with Image.open(path) as img:
|
|
still = media_io.cover(ImageOps.exif_transpose(img).convert("RGB"), width, height)
|
|
pixels, rate = media_io.to_tensor([still] * frames), float(adapt.get("fps", 24))
|
|
video = _encode_pixels(vae, pixels)
|
|
if video.ndim == 4: # a 2D VAE gave a stack of frames: [T,C,h,w] -> [1,C,T,h,w]
|
|
video = video.permute(1, 0, 2, 3)[None]
|
|
sound = None
|
|
if audio_vae is not None and adapt.get("silence"):
|
|
sound = _sound_latent(path, item["media_kind"], frames / rate, audio_vae, adapt)
|
|
return video, sound
|
|
img = media_io_picture(path, item["media_kind"], width, height)
|
|
return _encode_pixels(vae, media_io.to_tensor([img])), None
|
|
|
|
|
|
def media_io_picture(path, kind, width, height):
|
|
from . import frames
|
|
return media_io.cover(frames.picture(path), width, height)
|
|
|
|
|
|
def _sound_latent(path, kind, seconds, audio_vae, adapt):
|
|
import comfy.audio
|
|
wave = media_io.read_audio(path, seconds) if kind == "video" else None
|
|
if wave is None:
|
|
return adapt["silence"](audio_vae, seconds)
|
|
waveform, rate = wave
|
|
target = int(getattr(audio_vae, "audio_sample_rate", 44100))
|
|
if rate != target:
|
|
waveform = comfy.audio.resample(waveform, rate, target)
|
|
return audio_vae.encode(waveform.movedim(1, -1))
|
|
|
|
|
|
def run(model, clip, vae, audio_vae, set_name: str, cfg: dict, prompt_id: str = "") -> str:
|
|
"""Train and save. Returns the saved LoRA's file name inside the loras folder. Raises with a reason a person can act on."""
|
|
import torch
|
|
import comfy.model_management as mm
|
|
import comfy_extras.nodes_train as nt
|
|
import folder_paths
|
|
import safetensors.torch
|
|
|
|
from ..._core import registry
|
|
run_id = jobs.begin(set_name, cfg["name"], cfg["steps"], cfg, prompt_id) # first, so a refusal below is shown in the window and not only in ComfyUI's log
|
|
try:
|
|
if model is None:
|
|
model = load_model(cfg.get("model_name"), cfg.get("weight_dtype", "default"), cfg.get("compute_dtype", "default"))
|
|
adapt = registry.current().ask("training_adapt", model)
|
|
items, problems = plan_items(set_name, cfg["use"])
|
|
if problems and not items:
|
|
raise ValueError("; ".join(problems))
|
|
if problems:
|
|
log.warning("FunPack LoRA training", f"left out of the run: {'; '.join(problems[:5])}{' ...' if len(problems) > 5 else ''}")
|
|
if adapt is None and getattr(model, "cached_patcher_init", None) is None:
|
|
raise ValueError("this model was loaded in a way the trainer cannot copy; load it with the FunPack Diffusion Model Loader from a plain file")
|
|
batch = cfg["batch_size"]
|
|
if adapt and batch > adapt.get("batch_limit", batch):
|
|
log.warning("FunPack LoRA training", f"{adapt['name']} trains one sample at a time; batch size {batch} is lowered to {adapt['batch_limit']}")
|
|
batch = adapt["batch_limit"]
|
|
jobs.say(f"Reading {len(items)} items")
|
|
step = (adapt or {}).get("step") or 64
|
|
sizes = []
|
|
for item, path in items:
|
|
try:
|
|
sizes.append(media_io.size_of(path, item["media_kind"]))
|
|
except Exception as exc: # noqa: BLE001
|
|
raise ValueError(f"{item['name'] or item['id']} cannot be read: {exc}") from exc
|
|
single = bool(adapt) or not cfg["aspect_buckets"]
|
|
buckets = media_io.buckets_for([w / h for w, h in sizes], cfg["resolution"], step, single)
|
|
texts = [captions.with_parts(item["caption"], cfg["trigger"], cfg.get("postfix", "")) for item, _ in items]
|
|
|
|
if adapt and adapt.get("video"):
|
|
cfg = _fit_frames(cfg, adapt, items)
|
|
videos, sounds, conds = [], [], []
|
|
with torch.inference_mode():
|
|
for number, ((item, path), size, text) in enumerate(zip(items, buckets, texts), 1):
|
|
if jobs.stop_requested():
|
|
raise Stopped()
|
|
jobs.say(f"Preparing {number} of {len(items)}")
|
|
latent, sound = encode_item(item, path, size, cfg, vae, adapt, audio_vae)
|
|
videos.append(latent.clone())
|
|
sounds.append(None if sound is None else sound.clone())
|
|
conds.append(_encode_text(clip, text))
|
|
videos, sounds, conds = _normal(videos), _normal(sounds), _normal(conds) # made in inference mode: copies a backward pass can use
|
|
if adapt and adapt.get("pack") and audio_vae is not None and all(s is not None for s in sounds):
|
|
latents = adapt["pack"](videos, sounds)
|
|
payload = [latents]
|
|
elif adapt and adapt.get("pack"):
|
|
raise ValueError("this model trains picture and sound together and needs its audio VAE: add the audio VAE loader to the pipeline")
|
|
elif len({tuple(v.shape) for v in videos}) == 1:
|
|
payload = [{"samples": torch.cat(videos, 0)}]
|
|
else:
|
|
payload = [{"samples": v} for v in videos]
|
|
|
|
_sizes(payload, conds)
|
|
_card("before freeing", mm, torch)
|
|
del vae, clip, audio_vae # their work is done; the trainer copies the model next and needs the room
|
|
_free_the_card(mm, torch)
|
|
_card("after freeing", mm, torch)
|
|
_train_on_a_plain_copy(model)
|
|
losses = _loss_hook(nt, cfg["steps"] * cfg["grad_accumulation"])
|
|
jobs.say("Training", "training")
|
|
shim = adapt["shims"]() if adapt and adapt.get("shims") else _nothing()
|
|
with shim, losses, _watch(nt, mm, torch):
|
|
out = nt.TrainLoraNode.execute(
|
|
model=[model], latents=payload, positive=conds, batch_size=[batch], steps=[cfg["steps"]],
|
|
grad_accumulation_steps=[cfg["grad_accumulation"]], learning_rate=[cfg["learning_rate"]], rank=[cfg["rank"]],
|
|
optimizer=[cfg["optimizer"]], loss_function=[cfg["loss"]], seed=[cfg["seed"]], training_dtype=[cfg["training_dtype"]],
|
|
lora_dtype=[cfg["lora_dtype"]], quantized_backward=[False], algorithm=[cfg["algorithm"]],
|
|
gradient_checkpointing=[cfg["gradient_checkpointing"]], checkpoint_depth=[(adapt or {}).get("checkpoint_depth", 1)], offloading=[cfg["offloading"]],
|
|
existing_lora=["[None]"], bucket_mode=[False], bypass_mode=[cfg["bypass"]])
|
|
lora = (out.result if hasattr(out, "result") else out)[0]
|
|
file = _save(lora, cfg["name"], folder_paths, safetensors.torch, set_name, cfg)
|
|
jobs.finish("done", lora=file)
|
|
return file
|
|
except Stopped:
|
|
jobs.finish("stopped")
|
|
raise
|
|
except mm.InterruptProcessingException:
|
|
jobs.finish("stopped")
|
|
raise
|
|
except Exception as exc: # noqa: BLE001 -- recorded for the window, then raised for ComfyUI to show
|
|
jobs.finish("error", error=f"{type(exc).__name__}: {exc}")
|
|
raise
|
|
|
|
|
|
class _nothing:
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *exc):
|
|
return False
|
|
|
|
|
|
_WEIGHT = {"fp8_e4m3fn": "float8_e4m3fn", "fp8_e5m2": "float8_e5m2", "fp16": "float16", "bf16": "bfloat16", "fp32": "float32"}
|
|
|
|
|
|
def load_model(name, weight_dtype="default", compute_dtype="default"):
|
|
"""The model file, read once and never moved to the card here. The trainer's own copy shares it, so one model is in memory."""
|
|
import torch
|
|
import comfy.sd
|
|
import folder_paths
|
|
if not name:
|
|
raise ValueError("no model to train on: the pipeline's model loader names no file")
|
|
if str(name).lower().endswith(".gguf"):
|
|
raise ValueError("a GGUF model cannot be trained: pick a .safetensors file in the model loader")
|
|
path = folder_paths.get_full_path_or_raise("diffusion_models", name)
|
|
options = {}
|
|
if weight_dtype in _WEIGHT:
|
|
options["dtype"] = getattr(torch, _WEIGHT[weight_dtype])
|
|
with torch.inference_mode(False): # ComfyUI runs nodes in inference mode; weights made there cannot take part in a backward pass
|
|
model = comfy.sd.load_diffusion_model(path, model_options=options, disable_dynamic=True)
|
|
if compute_dtype in _WEIGHT and hasattr(model, "set_model_compute_dtype"):
|
|
model.set_model_compute_dtype(getattr(torch, _WEIGHT[compute_dtype]))
|
|
model.cached_patcher_init = (lambda *args, **kwargs: model, ()) # the trainer's "fresh copy" is this model: no second read of the file
|
|
return model
|
|
|
|
|
|
def _normal(value):
|
|
"""Tensors made in inference mode, copied so autograd may save them; lists, tuples and dicts are walked."""
|
|
import torch
|
|
if isinstance(value, torch.Tensor):
|
|
with torch.inference_mode(False):
|
|
return value.clone() if value.is_inference() else value
|
|
if isinstance(value, list):
|
|
return [_normal(v) for v in value]
|
|
if isinstance(value, tuple):
|
|
return tuple(_normal(v) for v in value)
|
|
if isinstance(value, dict):
|
|
return {k: _normal(v) for k, v in value.items()}
|
|
return value
|
|
|
|
|
|
def _fit_frames(cfg, adapt, items):
|
|
"""Every clip is read for the same number of frames, so the shortest clip sets the limit: fewer than asked is used, and said."""
|
|
rate = cfg["fps"] or adapt.get("fps", 24)
|
|
shortest = min((media_io.frames_at(path, rate) for item, path in items if item["media_kind"] == "video"), default=None)
|
|
wanted = adapt["frames"](cfg["frames"])
|
|
if shortest is None or shortest >= wanted:
|
|
return cfg
|
|
fit = adapt["fit"](shortest) if adapt.get("fit") else shortest
|
|
if fit < 1:
|
|
raise ValueError(f"the shortest clip has only {shortest} frames at {rate:g} fps: too short to train on")
|
|
log.warning("FunPack LoRA training", f"the shortest clip gives {shortest} frames at {rate:g} fps, so every clip is read for {fit} frames and not {wanted}")
|
|
return {**cfg, "frames": fit}
|
|
|
|
|
|
def _sizes(payload, conds):
|
|
"""One Log line: the shape of the first item's latent and of its text, to tell a run that is too big from one that is not."""
|
|
def shape(x):
|
|
tensors = getattr(x, "tensors", None)
|
|
if tensors is not None:
|
|
return " + ".join(str(list(t.shape)) for t in tensors)
|
|
return str(list(x.shape)) if hasattr(x, "shape") else type(x).__name__
|
|
try:
|
|
first = payload[0]["samples"] if isinstance(payload[0], dict) else payload[0]
|
|
text = conds[0][0] if isinstance(conds[0], (list, tuple)) else conds[0]
|
|
log.info("FunPack LoRA training", f"{len(conds)} text(s); first latent {shape(first)}, first text {shape(text)}")
|
|
except Exception as exc: # noqa: BLE001 -- a report must never stop a run
|
|
log.info("FunPack LoRA training", f"sizes: not readable ({type(exc).__name__})")
|
|
|
|
|
|
def _card(when, mm, torch):
|
|
"""One line in the Log: what the card holds and which models ComfyUI still has loaded."""
|
|
try:
|
|
device = mm.get_torch_device()
|
|
gib = 1024 ** 3
|
|
loaded = ", ".join(f"{type(getattr(m.model, 'model', m.model)).__name__} {m.model_loaded_memory() / gib:.1f} GiB" for m in mm.current_loaded_models) or "none"
|
|
log.info("FunPack LoRA training", f"card {when}: tensors {torch.cuda.memory_allocated(device) / gib:.1f} GiB, free {mm.get_free_memory(device) / gib:.1f} GiB, loaded models: {loaded}")
|
|
except Exception as exc: # noqa: BLE001 -- a report must never stop a run
|
|
log.info("FunPack LoRA training", f"card {when}: not readable ({type(exc).__name__})")
|
|
|
|
|
|
class _watch:
|
|
"""While open, a Log line after each stage of the trainer's setup, so an out-of-memory says which stage filled the card."""
|
|
STAGES = ("_prepare_latents_and_count", "_setup_lora_adapters", "_setup_lora_adapters_bypass", "_create_optimizer")
|
|
|
|
def __init__(self, nt, mm, torch):
|
|
self.nt, self.mm, self.torch, self.put = nt, mm, torch, []
|
|
|
|
def __enter__(self):
|
|
def wrap(owner, name, label):
|
|
original = getattr(owner, name, None)
|
|
if original is None:
|
|
return
|
|
|
|
def watched(*args, **kwargs):
|
|
out = original(*args, **kwargs)
|
|
_card(f"after {label}", self.mm, self.torch)
|
|
return out
|
|
setattr(owner, name, watched)
|
|
self.put.append((owner, name, original))
|
|
for name in self.STAGES:
|
|
wrap(self.nt, name, name.strip("_").replace("_", " "))
|
|
wrap(self.mm, "load_models_gpu", "loading the model")
|
|
return self
|
|
|
|
def __exit__(self, *exc):
|
|
for owner, name, original in reversed(self.put):
|
|
setattr(owner, name, original)
|
|
return False
|
|
|
|
|
|
def _free_the_card(mm, torch):
|
|
import gc
|
|
gc.collect()
|
|
mm.unload_all_models()
|
|
gc.collect()
|
|
mm.soft_empty_cache()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def _train_on_a_plain_copy(model):
|
|
"""The trainer copies the model with clone(force_deepcopy=True). On a dynamic-VRAM model that copy is dynamic too, and the
|
|
trainer then loads it in full beside the original: two models on the card. Ask for a plain copy, loaded once."""
|
|
is_dynamic = getattr(model, "is_dynamic", None)
|
|
if not (callable(is_dynamic) and is_dynamic()):
|
|
return
|
|
clone = model.clone
|
|
|
|
def plain(*args, **kwargs):
|
|
return clone(*args, **{**kwargs, "disable_dynamic": True})
|
|
model.clone = plain
|
|
|
|
|
|
class _loss_hook:
|
|
"""While open, every step's loss goes to the run record, and Stop (or ComfyUI's interrupt) ends the loop at the next step."""
|
|
|
|
def __init__(self, nt, total):
|
|
self.nt, self.total = nt, total
|
|
|
|
def __enter__(self):
|
|
import comfy.model_management as mm
|
|
nt = self.nt
|
|
self.original = nt.TrainSampler
|
|
|
|
def sampler(*args, **kwargs):
|
|
inner = kwargs.get("loss_callback")
|
|
|
|
def callback(loss):
|
|
if inner:
|
|
inner(loss)
|
|
jobs.step(float(loss))
|
|
mm.throw_exception_if_processing_interrupted()
|
|
if jobs.stop_requested():
|
|
raise Stopped()
|
|
kwargs["loss_callback"] = callback
|
|
return self.original(*args, **kwargs)
|
|
|
|
nt.TrainSampler = sampler
|
|
return self
|
|
|
|
def __exit__(self, *exc):
|
|
self.nt.TrainSampler = self.original
|
|
return False
|
|
|
|
|
|
def _save(lora, name, folder_paths, safetensors_torch, set_name, cfg) -> str:
|
|
"""Into the first loras folder, as <name>.safetensors, or <name>_2 and so on rather than over an earlier one."""
|
|
folder = Path(folder_paths.get_folder_paths("loras")[0])
|
|
folder.mkdir(parents=True, exist_ok=True)
|
|
stem = "".join(c for c in name if c.isalnum() or c in "-_ .") .strip() or "funpack_lora"
|
|
n, path = 1, folder / f"{stem}.safetensors"
|
|
while path.exists():
|
|
n += 1
|
|
path = folder / f"{stem}_{n}.safetensors"
|
|
meta = {"software": "ComfyUI-FunPack", "training_set": set_name, "algorithm": cfg["algorithm"], "rank": str(cfg["rank"]),
|
|
"steps": str(cfg["steps"]), "trigger": cfg["trigger"]}
|
|
safetensors_torch.save_file(lora, str(path), metadata=meta)
|
|
return path.name
|