From 9aa1791ec84c3b7fb116fe8d2e2dc8017cdd45c9 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 17 Oct 2026 22:00:51 +0000 Subject: [PATCH] LoRA training: encoded clips and captions are kept on disk and reused by the next run Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01MpnF7GwNGMbngU11wjym7P --- modules/training/lora/tests/test_train_h3.py | 22 ++++++++ modules/training/lora/train.py | 55 +++++++++++++++++++- 2 files changed, 75 insertions(+), 2 deletions(-) diff --git a/modules/training/lora/tests/test_train_h3.py b/modules/training/lora/tests/test_train_h3.py index e221e180..e412842b 100644 --- a/modules/training/lora/tests/test_train_h3.py +++ b/modules/training/lora/tests/test_train_h3.py @@ -403,6 +403,28 @@ def test_a_failed_run_is_restored_from_its_saved_state_in_comfyuis_inference_mod assert not loop.has_checkpoint(run) +def test_a_second_run_on_the_same_data_does_not_encode_it_again(run_env, monkeypatch): + from modules.training.lora import train + tmp_path, jobs = run_env + calls = {"video": 0, "text": 0} + real_item, real_text = train.encode_item, train._encode_text + + def item(*a, **k): + calls["video"] += 1 + return real_item(*a, **k) + + def text(*a, **k): + calls["text"] += 1 + return real_text(*a, **k) + monkeypatch.setattr(train, "encode_item", item) + monkeypatch.setattr(train, "_encode_text", text) + train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=2)) + first = dict(calls) + assert first["video"] > 0 and first["text"] > 0 and any((tmp_path / "training" / "cache").glob("latent-*.pt")) + train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=2)) + assert calls == first, "the clips and captions came from the cache" + + def test_pause_button_waits_at_the_next_step_and_stop_ends_a_paused_run(run_env, monkeypatch): from modules.training.lora import jobs as j, train tmp_path, jobs = run_env diff --git a/modules/training/lora/train.py b/modules/training/lora/train.py index 0e14984d..ccf76467 100644 --- a/modules/training/lora/train.py +++ b/modules/training/lora/train.py @@ -62,6 +62,50 @@ def plan_items(set_name: str, use: str): return found, problems +def _fingerprint(obj) -> str: + """Tells one set of VAE or text-encoder weights from another cheaply: class, size and a few values.""" + import hashlib + if obj is None: + return "" + inner = getattr(obj, "first_stage_model", None) or getattr(obj, "cond_stage_model", None) or getattr(obj, "model", None) or obj + h = hashlib.sha256(type(inner).__name__.encode()) + try: + params = list(inner.parameters()) + h.update(str(sum(p.numel() for p in params)).encode()) + for p in (params[0], params[-1]) if params else (): + h.update(p.detach().flatten()[:16].float().cpu().numpy().tobytes()) + except Exception: # noqa: BLE001 -- an object with no parameters: its class has to do + pass + return h.hexdigest()[:16] + + +def _pair(result): + return result[0].clone(), None if result[1] is None else result[1].clone() + + +def _cached(kind, parts, make, torch): + """What an item encodes to, kept on disk for the next run: same file, same settings, same weights give the same answer. Returns (value, was_cached).""" + import hashlib + import json + from ..._core import config + folder = config.TRAINING_DIR / "cache" + path = folder / f"{kind}-{hashlib.sha256(json.dumps(parts, default=str).encode()).hexdigest()[:24]}.pt" + if path.is_file(): + try: + return torch.load(path, weights_only=False), True + except Exception: # noqa: BLE001 -- a damaged or foreign-device file: encode again + pass + value = make() + try: + folder.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(".tmp") + torch.save(value, tmp) + tmp.replace(path) + except Exception as exc: # noqa: BLE001 -- the cache is only a speed-up + log.warning("FunPack LoRA training", f"could not keep {kind} in the cache: {exc}") + return value, False + + def _encode_text(clip, text: str): tokens = clip.tokenize(text) return clip.encode_from_tokens_scheduled(tokens) @@ -186,15 +230,22 @@ def _run(model, clip, vae, audio_vae, set_name: str, cfg: dict, prompt_id: str = if adapt and adapt.get("video"): cfg = _fit_frames(cfg, adapt, items) videos, sounds, conds = [], [], [] + hits = 0 with torch.inference_mode(): + vae_id, audio_id, clip_id = _fingerprint(vae), _fingerprint(audio_vae), _fingerprint(clip) 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) + shape = [str(Path(path).resolve()), Path(path).stat().st_mtime_ns, Path(path).stat().st_size, item["media_kind"], list(size), cfg["frames"], cfg["fps"], cfg["start"], + (adapt or {}).get("name"), vae_id, audio_id if audio_vae is not None else None] + (latent, sound), hit = _cached("latent", shape, lambda: _pair(encode_item(item, path, size, cfg, vae, adapt, audio_vae)), torch) + cond, hit_text = _cached("text", [text, clip_id], lambda: _encode_text(clip, text), torch) + hits += hit and hit_text videos.append(latent.clone()) sounds.append(None if sound is None else sound.clone()) - conds.append(_encode_text(clip, text)) + conds.append(cond) + log.info("FunPack LoRA training", f"prepared {len(items)} items, {hits} entirely from the cache") 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)