diff --git a/app/ui/shell/run.js b/app/ui/shell/run.js index aa04ffe5..2d930bed 100644 --- a/app/ui/shell/run.js +++ b/app/ui/shell/run.js @@ -57,8 +57,33 @@ export function createRun({ for (const fn of listeners) { try { fn(state); } catch { /* one listener must not stop the others */ } } + watch(); }; + // The socket is the fast path, not the only one: a run whose finish arrived while the socket was down (a tunnel + // blip, a reload) would otherwise sit at "Generating" forever. ComfyUI's history is the record; asked on + // reconnect, on adopt, and every half minute while a run is going. + let poll = null; + function watch() { + const going = state.promptId && (state.phase === QUEUED || state.phase === RUNNING); + if (going && !poll) { poll = setInterval(reconcile, 30000); if (poll.unref) poll.unref(); } // unref: never what keeps a process alive + if (!going && poll) { clearInterval(poll); poll = null; } + } + async function reconcile() { + const id = state.promptId; + if (!id || !(state.phase === QUEUED || state.phase === RUNNING) || typeof doFetch !== "function") return; + let entry; + try { const res = await doFetch(`${base}/history/${encodeURIComponent(id)}`); entry = res.ok ? (await res.json())[id] : null; } catch { return; } + if (!entry || state.promptId !== id || !(state.phase === QUEUED || state.phase === RUNNING)) return; // still going, or moved on + const st = entry.status || {}, outs = Object.values(entry.outputs || {}); + if (st.status_str === "error") { + const m = ((st.messages || []).find((x) => x[0] === "execution_error") || [])[1] || {}; + emit({ phase: FAILED, progress: null, error: { node: m.node_type || m.node_id || null, message: m.exception_message || "the run failed", traceback: m.traceback || null } }); + } else if (st.completed) { + emit({ images: outs.flatMap((o) => o.images || []), audio: outs.flatMap((o) => o.audio || []), phase: DONE, progress: null, node: null }); + } + } + // Only messages carrying OUR prompt id are acted on. ComfyUI broadcasts some // of them (an interrupt is broadcast to every client), so a second tab // generating would otherwise drive this one's progress bar and, worse, mark @@ -234,7 +259,7 @@ export function createRun({ socket = connect(clientId); if (!socket) return null; - socket.addEventListener("open", () => { attempt = 0; emit({ connection: LIVE }); }); + socket.addEventListener("open", () => { attempt = 0; emit({ connection: LIVE }); reconcile(); }); socket.addEventListener("message", (event) => { // Binary frames are previews, which nothing here consumes yet. A parse // failure must not take the socket down with it. @@ -366,6 +391,7 @@ export function createRun({ emit({ phase: running ? RUNNING : QUEUED, images: [], audio: [], error: null, node: null }); settle(promptId); + reconcile(); // a run that finished before this page heard of it return true; }, diff --git a/app/ui/shell/tests/run.test.js b/app/ui/shell/tests/run.test.js index 89d85f61..1fe895f6 100644 --- a/app/ui/shell/tests/run.test.js +++ b/app/ui/shell/tests/run.test.js @@ -337,7 +337,7 @@ test("an adopted run can be cancelled by its own id", async () => { const { run, sent } = runner(ok({})); run.adopt("was-running"); await run.cancel(); - assert.equal(sent[0].url, "/api/jobs/was-running/cancel"); + assert.ok(sent.some((x) => x.url === "/api/jobs/was-running/cancel")); }); // --- a run that finishes while the page is loading --------------------------- @@ -470,3 +470,29 @@ test("a run adopted while it is still waiting its turn says Queued, not Working" run.handle(JSON.parse(message("execution_start", { prompt_id: "waiting" }).data)); assert.equal(run.state.phase, RUNNING); }); + +// --- a finish the socket never delivered ------------------------------------ + +test("a run whose finish was missed (socket down, or finished before the page heard of it) is read from ComfyUI's history", async () => { + const history = { p9: { status: { status_str: "success", completed: true }, outputs: { "7": { images: [{ filename: "late.mp4", type: "output" }] } } } }; + const { run } = runner(async (url) => ({ ok: true, status: 200, json: async () => (url === "/history/p9" ? history : {}) })); + run.adopt("p9"); + await new Promise((r) => setTimeout(r, 0)); + assert.equal(run.state.phase, DONE); + assert.deepEqual(run.state.images.map((i) => i.filename), ["late.mp4"]); +}); + +test("a missed failure is read from history with ComfyUI's own reason", async () => { + const history = { p9: { status: { status_str: "error", messages: [["execution_error", { node_type: "SaveVideo", exception_message: "disk full" }]] }, outputs: {} } }; + const { run } = runner(async (url) => ({ ok: true, status: 200, json: async () => (url === "/history/p9" ? history : {}) })); + run.adopt("p9"); + await new Promise((r) => setTimeout(r, 0)); + assert.deepEqual([run.state.phase, run.state.error.node, run.state.error.message], ["failed", "SaveVideo", "disk full"]); +}); + +test("a run still going (no history entry yet) stays as it is", async () => { + const { run } = runner(ok({})); + run.adopt("p9"); + await new Promise((r) => setTimeout(r, 0)); + assert.equal(run.state.phase, "running"); +}); diff --git a/core/readiness.py b/core/readiness.py index ac552abf..48ccb2fe 100644 --- a/core/readiness.py +++ b/core/readiness.py @@ -18,6 +18,14 @@ def row(level, text): return {"level": level, "text": text} +def xformers_on() -> bool: + try: + import comfy.model_management as mm + return bool(mm.xformers_enabled()) + except Exception: # noqa: BLE001 + return have("xformers") + + def have(module) -> bool: try: return importlib.util.find_spec(module) is not None @@ -64,7 +72,7 @@ def machine(disk_free_gb=None) -> List[dict]: out.append(row("ok", f"{props.name}, {props.total_memory / 1024 ** 3:.0f} GB, sm_{props.major}{props.minor}")) if props.major < 8: # bf16 arrived with Ampere (sm_80), on the card ComfyUI uses out.append(row("warn", "This GPU has no native bf16 (before sm_80): the H3 VAEs, which need it, run emulated and slowly.")) - if props.major >= 12 and have("xformers"): + if props.major >= 12 and xformers_on(): # installed is not active: --disable-xformers turns it off out.append(row("warn", "xformers has no masked-attention kernel for this GPU: start ComfyUI with --disable-xformers --use-sage-attention.")) if not have("sageattention"): # a CUDA-only library: nothing to install on a Mac out.append(row("warn", "sageattention is not installed: attention runs slower (and the H3 SLA setting has no fast path).")) diff --git a/core/tests/test_readiness.py b/core/tests/test_readiness.py index 3320fc84..7f56d5c0 100644 --- a/core/tests/test_readiness.py +++ b/core/tests/test_readiness.py @@ -54,3 +54,15 @@ def test_the_gpu_row_names_the_card_comfyui_uses(monkeypatch): lambda i: asked.append(i) or types.SimpleNamespace(name=f"GPU{i}", total_memory=32 * 1024 ** 3, major=12, minor=0)) rows = r.machine() assert asked == [1] and any("GPU1" in x["text"] for x in rows) + + +def test_the_blackwell_xformers_warning_follows_whether_xformers_is_on_not_installed(monkeypatch): + import types + import torch + monkeypatch.setattr(r, "device", lambda: "cuda") + monkeypatch.setattr(r, "gpu_index", lambda: 0) + monkeypatch.setattr(torch.cuda, "get_device_properties", + lambda i: types.SimpleNamespace(name="B", total_memory=96 * 1024 ** 3, major=12, minor=0)) + for on, warned in ((True, True), (False, False)): + monkeypatch.setattr(r, "xformers_on", lambda: on) + assert any("xformers" in x["text"] for x in r.machine()) is warned diff --git a/modules/loaders/common.py b/modules/loaders/common.py index 46462432..2c9ab3bd 100644 --- a/modules/loaders/common.py +++ b/modules/loaders/common.py @@ -39,6 +39,16 @@ def dtype_of(name): return _BY_NAME.get(str(name or "default")) +def unsupported_dtypes(model, weight_dtype, compute_dtype): + """A forced dtype the model does not run in, one line each. A model says which it runs in (H3: bf16/fp32; + v4 measured fp16 breaking it). Allowed, but said.""" + supported = getattr(getattr(getattr(model, "model", None), "model_config", None), "supported_inference_dtypes", None) or [] + runs_in = ", ".join(str(d).replace("torch.", "") for d in supported) + return [f"{what} dtype {picked} is not one this model runs in ({runs_in})" + for what, picked in (("weight", weight_dtype), ("compute", compute_dtype)) + if picked in ("fp16", "bf16", "fp32") and supported and dtype_of(picked) not in supported] + + def weight_model_options(weight_dtype): """comfy `model_options` for a weight dtype. The mapping core's UNETLoader uses.""" options = {} diff --git a/modules/loaders/diffusion_model/nodes.py b/modules/loaders/diffusion_model/nodes.py index 958a1576..41cc1cdb 100644 --- a/modules/loaders/diffusion_model/nodes.py +++ b/modules/loaders/diffusion_model/nodes.py @@ -14,7 +14,7 @@ from comfy_api.latest import io from ..._core import log from .. import gguf_support, sla_attention from .. import int8_convrot as int8_convrot_module -from ..common import (COMPUTE_DTYPES, WEIGHT_DTYPES, attention_choices, +from ..common import (COMPUTE_DTYPES, WEIGHT_DTYPES, attention_choices, unsupported_dtypes, attention_override, dtype_of, set_fp16_accumulation, weight_model_options) @@ -176,6 +176,9 @@ class FunPackDiffusionModelLoader(io.ComfyNode): f"Could not detect a diffusion model in {model_name}. Loading a text " f"encoder or a VAE file here is the usual cause.") + for line in unsupported_dtypes(model, weight_dtype, compute_dtype): + log.alert("FunPack Diffusion Model Loader", f"{line}; expect noise or broken sound. Set it to default.") + notes.append(line) dtype = dtype_of(compute_dtype) if dtype is not None and hasattr(model, "set_model_compute_dtype"): # Do NOT clear force_cast_weights afterwards. set_model_compute_dtype diff --git a/modules/loaders/lora/keys.py b/modules/loaders/lora/keys.py new file mode 100644 index 00000000..f3fc3c6d --- /dev/null +++ b/modules/loaders/lora/keys.py @@ -0,0 +1,62 @@ +"""Which of a LoRA's weights land in this model, and which cannot. + +ComfyUI converts the naming dialects (lora_up / lora_A / PEFT); it does not unwrap a file whose +keys sit under a training wrapper. v4 measured it: an H3 adapter keyed `base_model.model.dit.*` +matched 0 of 532 keys and did nothing, silently. Every known wrapper is tried and the one matching +the most weights wins. A pair that matches by name but not by shape is dropped: ComfyUI reshapes +any delta with the right element count, so a wrong-shaped one merges scrambled (all-NaN mid-render), +and a full-width LoRA on a curve-form H3 checkpoint churned ~25 GB failing per key. +""" + +import comfy.lora +import comfy.lora_convert + +# (prefix, replacement). Replacing, not only stripping: ComfyUI's key map has `diffusion_model.` +# and no bare form, so `base_model.model.dit.blocks.0...` must become `diffusion_model.blocks.0...`. +REKEY = ( + ("transformer.", ""), ("diffusion_model.", ""), ("model.diffusion_model.", ""), + ("base_model.model.", ""), ("base_model.model.dit.", ""), ("base_model.model.diffusion_model.", ""), + ("lora_model.", ""), ("net.", ""), + ("base_model.model.dit.", "diffusion_model."), ("base_model.model.", "diffusion_model."), + ("dit.", "diffusion_model."), ("model.dit.", "diffusion_model."), ("transformer.", "diffusion_model."), +) + + +def _rekey(lora, prefix, replacement): + if not any(k.startswith(prefix) for k in lora): + return None + return {(replacement + k[len(prefix):] if k.startswith(prefix) else k): v for k, v in lora.items()} + + +def _misfit(sd, key, adapter): + """True when a plain LoRA pair cannot make the weight it names (both dimensions, not the count).""" + w, target = getattr(adapter, "weights", None), sd.get(key) + if not w or len(w) < 2 or target is None or target.dim() < 2: + return False + if (len(w) > 3 and w[3] is not None) or (len(w) > 5 and w[5] is not None): # locon mid / padded target: not modelled + return False + try: + got = (int(w[0].flatten(start_dim=1).shape[0]), int(w[1].flatten(start_dim=1).shape[-1])) + except Exception: # noqa: BLE001 -- not a plain pair (LoKr, LoHa, diff): left to ComfyUI + return False + return got != (int(target.shape[0]), int(target.shape[1:].numel())) + + +def match(model, clip, lora): + """-> (patches, how, dropped): the best keying's patches, which keying won, and how many were dropped for shape.""" + key_map = comfy.lora.model_lora_keys_unet(model.model, {}) + if clip is not None: + key_map = comfy.lora.model_lora_keys_clip(clip.cond_stage_model, key_map) + converted = comfy.lora_convert.convert_lora(lora) + best, how = comfy.lora.load_lora(converted, key_map, log_missing=False), "as-is" + for prefix, replacement in REKEY: + variant = _rekey(converted, prefix, replacement) + if variant is not None: + got = comfy.lora.load_lora(variant, key_map, log_missing=False) + if len(got) > len(best): + best, how = got, f"{prefix} -> {replacement or '(none)'}" + sd = model.model.state_dict() + bad = [k for k, a in best.items() if _misfit(sd, k, a)] + for k in bad: + best.pop(k) + return best, how, len(bad) diff --git a/modules/loaders/lora/nodes.py b/modules/loaders/lora/nodes.py index 629b5199..920b0ab6 100644 --- a/modules/loaders/lora/nodes.py +++ b/modules/loaders/lora/nodes.py @@ -7,12 +7,12 @@ what ComfyUI's own LoraLoader does, every strength is a real socket, and the gra shows how many LoRAs are actually applied instead of hiding the count in a string. """ -import comfy.sd import comfy.utils import folder_paths from comfy_api.latest import io from ..._core import log +from .keys import match NONE = "None" @@ -61,11 +61,24 @@ class FunPackLoraLoader(io.ComfyNode): path = folder_paths.get_full_path_or_raise("loras", lora_name) lora = comfy.utils.load_torch_file(path, safe_load=True) - patched_model, patched_clip = comfy.sd.load_lora_for_models( - model, clip, lora, strength_model, strength_clip) + patches, how, dropped = match(model, clip, lora) + if not patches: + log.alert("FunPack LoRA Loader", f"{lora_name} matched NOTHING in this model, so it does nothing: " + "it is for another model, or names its layers in a way not known here") + return io.NodeOutput(model, clip, f"{lora_name}: matched nothing, not applied") + patched_model = model.clone() + patched_model.add_patches(patches, strength_model) + patched_clip = None + if clip is not None: + patched_clip = clip.clone() + patched_clip.add_patches(patches, strength_clip) + if dropped: + log.alert("FunPack LoRA Loader", f"{lora_name}: {dropped} of {len(patches) + dropped} weights do not fit this " + "model's shapes and were left out (trained on another variant of it)") applied = [f"model={strength_model}"] applied.append(f"clip={strength_clip}" if clip is not None else "clip=not wired") - log.info("FunPack LoRA Loader", f"{lora_name} applied ({', '.join(applied)})") + found = f"{len(patches)} weights" + (f", keys {how}" if how != "as-is" else "") + (f", {dropped} dropped" if dropped else "") + log.info("FunPack LoRA Loader", f"{lora_name} applied ({', '.join(applied)}; {found})") return io.NodeOutput(patched_model, patched_clip, - f"FunPack LoRA Loader | {lora_name} ({', '.join(applied)})") + f"FunPack LoRA Loader | {lora_name} ({', '.join(applied)}; {found})") diff --git a/modules/loaders/tests/test_loader_behaviour.py b/modules/loaders/tests/test_loader_behaviour.py index 055ccf8e..9fc309d7 100644 --- a/modules/loaders/tests/test_loader_behaviour.py +++ b/modules/loaders/tests/test_loader_behaviour.py @@ -100,22 +100,72 @@ def test_a_lora_at_zero_strength_passes_the_originals_through(monkeypatch): def test_a_model_only_lora_still_applies_at_zero_clip_strength(monkeypatch): """clip is unwired, so strength_clip is meaningless -- but strength_model is not, and skipping the whole thing would silently drop the LoRA.""" - import comfy.sd import comfy.utils import folder_paths from modules.loaders.lora import nodes - called = {} + added = {} + + class M: + def clone(self): + return self + + def add_patches(self, patches, strength): + added.update(n=len(patches), s=strength) + monkeypatch.setattr(folder_paths, "get_full_path_or_raise", lambda kind, name: f"/fake/{name}") monkeypatch.setattr(comfy.utils, "load_torch_file", lambda p, **kw: {"fake": 1}) - monkeypatch.setattr(comfy.sd, "load_lora_for_models", - lambda m, c, l, sm, sc: called.update(sm=sm, sc=sc) or ("patched", None)) + monkeypatch.setattr(nodes, "match", lambda m, c, l: ({"w": object()}, "as-is", 0)) + out = nodes.FunPackLoraLoader.execute(model=M(), lora_name="x.safetensors", strength_model=0.8, clip=None) + assert added == {"n": 1, "s": 0.8} + assert "1 weights" in out.result[2] - out = nodes.FunPackLoraLoader.execute( - model=object(), lora_name="x.safetensors", strength_model=0.8, clip=None) - assert called["sm"] == 0.8 - assert out.result[0] == "patched" +def _lora_world(monkeypatch, weight_shape=(4, 3)): + import torch + import comfy.lora + + class Inner: + def state_dict(self): + return {"diffusion_model.blocks.0.attn.weight": torch.zeros(weight_shape)} + + class M: + model = Inner() + monkeypatch.setattr(comfy.lora, "model_lora_keys_unet", lambda m, km: {**km, "diffusion_model.blocks.0.attn": "diffusion_model.blocks.0.attn.weight"}) + return M() + + +def test_a_lora_under_a_training_wrapper_is_unwrapped_and_lands(monkeypatch): + # v4: an H3 adapter keyed base_model.model.dit.* matched 0 of 532 keys and did nothing, silently. + import torch + from modules.loaders.lora.keys import match + model = _lora_world(monkeypatch) + lora = {"base_model.model.dit.blocks.0.attn.lora_up.weight": torch.zeros(4, 2), + "base_model.model.dit.blocks.0.attn.lora_down.weight": torch.zeros(2, 3)} + patches, how, dropped = match(model, None, lora) + assert list(patches) == ["diffusion_model.blocks.0.attn.weight"] and dropped == 0 and "dit." in how + + +def test_a_lora_pair_of_the_wrong_shape_is_dropped_not_merged_scrambled(monkeypatch): + import torch + from modules.loaders.lora.keys import match + model = _lora_world(monkeypatch, weight_shape=(6, 2)) # same element count (12), other shape + lora = {"diffusion_model.blocks.0.attn.lora_up.weight": torch.zeros(4, 2), + "diffusion_model.blocks.0.attn.lora_down.weight": torch.zeros(2, 3)} + patches, _how, dropped = match(model, None, lora) + assert patches == {} and dropped == 1 + + +def test_a_lora_that_matches_nothing_is_not_applied_and_says_so(monkeypatch): + import comfy.utils + import folder_paths + from modules.loaders.lora import nodes + monkeypatch.setattr(folder_paths, "get_full_path_or_raise", lambda kind, name: f"/fake/{name}") + monkeypatch.setattr(comfy.utils, "load_torch_file", lambda p, **kw: {}) + monkeypatch.setattr(nodes, "match", lambda m, c, l: ({}, "as-is", 0)) + model = object() + out = nodes.FunPackLoraLoader.execute(model=model, lora_name="x.safetensors", strength_model=1.0, clip=None) + assert out.result[0] is model and "matched nothing" in out.result[2] def test_the_attention_override_calls_the_unwrapped_backend(monkeypatch): @@ -336,3 +386,14 @@ def test_no_lora_picked_passes_the_originals_through_without_reading_a_file(): model, clip = object(), object() out = nodes.FunPackLoraLoader.execute(model=model, lora_name="None", strength_model=1.0, clip=clip) assert out.result[0] is model and out.result[1] is clip + + +def test_a_forced_dtype_the_model_does_not_run_in_is_said(): + # v4: fp16 broke MiniMax H3, which ComfyUI lists as bf16/fp32 only. + import torch + from types import SimpleNamespace as NS + from modules.loaders.common import unsupported_dtypes + h3 = NS(model=NS(model_config=NS(supported_inference_dtypes=[torch.bfloat16, torch.float32]))) + assert len(unsupported_dtypes(h3, "fp16", "default")) == 1 and "fp16" in unsupported_dtypes(h3, "fp16", "default")[0] + assert unsupported_dtypes(h3, "default", "bf16") == [] + assert unsupported_dtypes(h3, "fp8_e4m3fn", "fp16")[0].startswith("compute dtype fp16") diff --git a/modules/sampling/prompt_nudge/__init__.py b/modules/sampling/prompt_nudge/__init__.py index 946257aa..52b98915 100644 --- a/modules/sampling/prompt_nudge/__init__.py +++ b/modules/sampling/prompt_nudge/__init__.py @@ -60,11 +60,12 @@ def nudged(c, words, d, amount, width=None): return c + step * words.view(1, -1, 1).to(c.dtype) -def picture_width(patcher, c): - """How many leading channels of the words belong to the picture, or None when all do.""" +def picture_widths(patcher): + """-> c -> how many leading channels of the words belong to the picture, or None when all do. Reads the model + now: the wrapper must not hold the patcher it is installed on (a cycle that keeps the model alive).""" dm = getattr(getattr(patcher, "model", None), "diffusion_model", None) video, audio = getattr(dm, "cross_attention_dim", None), getattr(dm, "audio_cross_attention_dim", None) - return int(video) if video and audio and int(c.shape[-1]) == int(video) + int(audio) else None + return lambda c: int(video) if video and audio and int(c.shape[-1]) == int(video) + int(audio) else None def install(patcher, values, key): @@ -84,6 +85,7 @@ def install(patcher, values, key): captured = taste.collect(patcher, key, KIND, fresh=fresh) acted = {"yes": False} + picture_width = picture_widths(patcher) def apply_model(executor, x, t, *args, **kwargs): named = streams.model_args(args, kwargs) @@ -121,7 +123,7 @@ def install(patcher, values, key): _say("off this run: the learned direction was taught on a different model's text width; " "rate a few clips on this model") return executor(x, t, *args, **kwargs) - new_c = nudged(c, words, d, amount, picture_width(patcher, c)) + new_c = nudged(c, words, d, amount, picture_width(c)) if not acted.get("checked"): acted["checked"] = True moved = float((new_c - c).float().norm()) diff --git a/modules/sampling/prompt_nudge/tests/test_prompt_nudge.py b/modules/sampling/prompt_nudge/tests/test_prompt_nudge.py index ac4b7b48..c4d4cc9a 100644 --- a/modules/sampling/prompt_nudge/tests/test_prompt_nudge.py +++ b/modules/sampling/prompt_nudge/tests/test_prompt_nudge.py @@ -93,7 +93,7 @@ def test_off_installs_nothing(tiny_h3): def test_the_sounds_text_channels_are_left_alone(): - from modules.sampling.prompt_nudge import nudged, picture_width + from modules.sampling.prompt_nudge import nudged, picture_widths class Dm: cross_attention_dim, audio_cross_attention_dim = 6, 2 @@ -102,8 +102,8 @@ def test_the_sounds_text_channels_are_left_alone(): model = type("M", (), {"diffusion_model": Dm()})() c = torch.ones(1, 3, 8) - width = picture_width(P(), c) - assert width == 6 and picture_width(P(), torch.ones(1, 3, 5)) is None + width = picture_widths(P())(c) + assert width == 6 and picture_widths(P())(torch.ones(1, 3, 5)) is None out = nudged(c, torch.ones(3, dtype=torch.bool), torch.ones(8) / 8 ** 0.5, 1.0, width) assert (out[..., :6] > 1).all() and torch.equal(out[..., 6:], c[..., 6:]) @@ -176,3 +176,21 @@ def test_a_probe_does_not_teach_the_taste_key(tiny_h3): wrap(lambda x, t, *a, **k: x, torch.zeros(1, 1, 4), ts[3:4], None, c, None, {**_to(3), dit_hooks.PROBE: True}) pending = store._dir("fox") / "prompt_taste.pending.pt" assert not pending.exists() # nothing was captured from the probe + + +def test_the_installed_wrapper_does_not_hold_its_patcher(): + # v4 8bc7cb64: a wrapper closing over its own patcher keeps the model alive past the run ("Potential memory leak"). + import gc, weakref + from modules.sampling import prompt_nudge as pn + + class P: + model = None + p = P() + ref = weakref.ref(p) + fn = pn.picture_widths(p) + del p + gc.disable() + try: + assert ref() is None and fn is not None + finally: + gc.enable() diff --git a/modules/system/taste/tests/test_value.py b/modules/system/taste/tests/test_value.py index 341860da..ffdc1ae7 100644 --- a/modules/system/taste/tests/test_value.py +++ b/modules/system/taste/tests/test_value.py @@ -33,3 +33,13 @@ def test_its_push_raises_the_score_and_points_at_liked(): shift = value.describe(moved) - value.describe(video) assert torch.nn.functional.cosine_similarity(shift, liked_dir, dim=0) > 0.3 assert torch.allclose((moved - video).norm(), 0.05 * video.norm(), rtol=1e-3) + + +def test_the_judge_trains_and_steers_under_comfyuis_inference_mode(): + # ComfyUI runs every prompt under torch.inference_mode(): weights made there cannot be trained. + import torch + from modules.system.taste import value as v + with torch.inference_mode(): + judge = v.Judge.train_on([torch.randn(v.DIM) for _ in range(v.MIN_SAMPLES + 2)], [i % 2 for i in range(v.MIN_SAMPLES + 2)]) + assert judge is not None + assert judge.gradient(torch.randn(1, 4, 2, 8, 8)).shape == (1, 4, 2, 8, 8) diff --git a/modules/system/taste/value.py b/modules/system/taste/value.py index a7370e00..158c8898 100644 --- a/modules/system/taste/value.py +++ b/modules/system/taste/value.py @@ -57,14 +57,14 @@ class Judge(nn.Module): pairs = list(zip(descriptors, rewards))[-BUFFER:] if len(pairs) < MIN_SAMPLES or len({r for _d, r in pairs}) < 2: return None - judge = cls(seed) pick = random.Random(seed) device = "cuda" if torch.cuda.is_available() else "cpu" - xs = torch.stack([d.float().reshape(-1) for d, _r in pairs]).to(device) rs = [float(r) for _d, r in pairs] n = len(rs) + # ComfyUI runs the whole prompt under inference_mode: weights or inputs made there cannot be trained. with torch.inference_mode(False), torch.enable_grad(): - judge.to(device) + judge = cls(seed).to(device) + xs = torch.stack([d.float().reshape(-1).clone() for d, _r in pairs]).to(device) opt = torch.optim.Adam(judge.parameters(), lr=2e-3) for _ in range(STEPS): opt.zero_grad() @@ -92,9 +92,9 @@ class Judge(nn.Module): def gradient(self, video): """d(score)/d(video), same shape, scaled by the members' agreement.""" device = video.device - if next(self.parameters()).device != device: - self.to(device) with torch.inference_mode(False), torch.enable_grad(): + if next(self.parameters()).device != device: + self.to(device) # moved here, not under inference_mode: its copies must stay differentiable x = torch.empty(video.shape, dtype=torch.float32, device=device) x.copy_(video) x.requires_grad_(True)