fix: v4 regressions before the rental: taste judge trains under inference mode; LoRA keys unwrapped, misfits dropped, no-match said; unsupported dtype said; missed run finish read from history; prompt nudge no longer holds the model

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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;
     },
 

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");
+});

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)."))

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

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 = {}

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

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.<path>`
+# 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)

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})")

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")

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())

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()

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)

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)