Batch Training UX: real window + dynamic label + no auto-queue; remove sampler input
Per feedback: Batch Training is a complete UI thing driven by Studio (the hub), not a prompt, and it must NOT start generation itself (ComfyUI's Generate does). - Removed batch_iterations from the Chain Sampler entirely (input + param). The sampler now enters batch mode ONLY when Studio packs variant markers into positive; the batch count lives in Studio. _run_batch_training simplified to variant-entries-only. - The node's 'Batch Training' button now opens a dedicated window (openBatchPanel) that shows EITHER a config view (set batch size, Arm/Disarm) when no batch is pending, OR the rating list (preview + per-entry rating + Submit/Forget) when a batch has finished. It never queues — arming just sets refiner.batch_variants; the user presses Generate. - Dynamic button label reflects live state: '(off)', '(armed ×N)', '(Processing…)', '(Finished — rate)'. Driven by execution_start/success/error listeners, which also one-shot disarm after a batch runs so the next Generate is normal. The older in-Refiner-tab batch panel remains functional (secondary path); the node-button window is now the primary entry point. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
parent
a1ba067d9a
commit
a74108de1a
+20
-33
@@ -1621,10 +1621,6 @@ class FunPackLTXAVSceneChainSampler:
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"batch_iterations": ("INT", {
|
||||
"default": 1, "min": 1, "max": 32, "step": 1,
|
||||
"tooltip": "Batch Training: run the entire scene chain this many times with everything frozen except the seed (seed + i). Each result is saved under <refinement_key>_<timestamp> with a preview for rating in Studio. 1 = normal single generation. Requires refinement_key_input. Controlled batch = clean RLHF signal (same conditioning, only noise varies).",
|
||||
}),
|
||||
"decode_tile_size": ("INT", {
|
||||
"default": 0, "min": 0, "max": 4096, "step": 64,
|
||||
"tooltip": "Tile size for VAE decode (0 = no tiling). Set to e.g. 512 if decode OOMs.",
|
||||
@@ -2239,24 +2235,23 @@ class FunPackLTXAVSceneChainSampler:
|
||||
carry_i2v_guides=False,
|
||||
mid_scene_guide=False, mid_scene_guide_strength=0.4,
|
||||
embed_guidance=False, embed_guidance_strength=0.02,
|
||||
transition_duration=16, batch_iterations=1, decode_tile_size=0,
|
||||
transition_duration=16, decode_tile_size=0,
|
||||
refinement_key_input="", unique_id=None, prompt=None):
|
||||
if not isinstance(positive, list) or not positive:
|
||||
raise ValueError("positive conditioning must contain at least one scene entry.")
|
||||
if negative is None:
|
||||
negative = []
|
||||
|
||||
# Batch Training: run the whole chain N times, persisting each result for rating in
|
||||
# Studio. Triggered by batch_iterations > 1 (seed-only) OR by Studio packing variant
|
||||
# conditionings into positive (the 'funpack_batch_variant' markers — variant count wins).
|
||||
batch_n = max(1, int(batch_iterations or 1))
|
||||
if batch_n > 1 or self._split_batch_variants(positive) is not None:
|
||||
# Batch Training: Studio (the hub) packs N conditionings into positive, each scene entry
|
||||
# tagged 'funpack_batch_variant'. That marker is the only trigger — the sampler has no
|
||||
# batch-count input. Sample one chain per packed entry, persist each for rating in Studio.
|
||||
if self._split_batch_variants(positive) is not None:
|
||||
return self._run_batch_training(
|
||||
model, vae, positive, negative, sampler, sigmas, seed, latent_template,
|
||||
num_frames_per_scene, frame_overlap, cfg, max_scenes, use_same_seed,
|
||||
carry_i2v_guides, mid_scene_guide, mid_scene_guide_strength,
|
||||
embed_guidance, embed_guidance_strength, transition_duration,
|
||||
decode_tile_size, refinement_key_input, batch_n,
|
||||
decode_tile_size, refinement_key_input,
|
||||
)
|
||||
|
||||
max_scene_count = max(1, int(max_scenes))
|
||||
@@ -2486,11 +2481,10 @@ class FunPackLTXAVSceneChainSampler:
|
||||
latent_template, num_frames_per_scene, frame_overlap, cfg, max_scenes,
|
||||
use_same_seed, carry_i2v_guides, mid_scene_guide, mid_scene_guide_strength,
|
||||
embed_guidance, embed_guidance_strength, transition_duration,
|
||||
decode_tile_size, refinement_key_input, batch_n):
|
||||
"""Run the full scene chain batch_n times — frozen except the seed (seed + i) — and
|
||||
persist each result (latent + preview + manifest) under refinements/batches/<key>/<stamp>
|
||||
for controlled-batch rating in Studio. Conditioning is provably identical across the
|
||||
batch because each run reuses sample() with only the seed changed."""
|
||||
decode_tile_size, refinement_key_input):
|
||||
"""Sample one chain per Studio-packed variant entry (seed + index), persisting each result
|
||||
(latent + preview + per-entry cond + manifest) under ComfyUI temp for rating in Studio.
|
||||
Reuses sample() per entry with only the seed changed, so each entry is a clean generation."""
|
||||
import json as _json, time as _time
|
||||
key = str(refinement_key_input or "").strip()
|
||||
if not key:
|
||||
@@ -2501,38 +2495,31 @@ class FunPackLTXAVSceneChainSampler:
|
||||
carry_i2v_guides=carry_i2v_guides, mid_scene_guide=mid_scene_guide,
|
||||
mid_scene_guide_strength=mid_scene_guide_strength, embed_guidance=embed_guidance,
|
||||
embed_guidance_strength=embed_guidance_strength, transition_duration=transition_duration,
|
||||
batch_iterations=1, decode_tile_size=decode_tile_size, refinement_key_input="",
|
||||
decode_tile_size=decode_tile_size, refinement_key_input="",
|
||||
)
|
||||
stamp = _time.strftime("%Y-%m-%d_%H-%M-%S")
|
||||
batch_dir, safe, rel = self._batch_dir(key, stamp)
|
||||
scene_prompts = [self._scene_text(c, i) for i, c in enumerate(positive[:max(1, int(max_scenes))])]
|
||||
manifest = {"key": key, "created": stamp,
|
||||
"subfolder": rel.replace(os.sep, "/"), "scene_prompts": scene_prompts, "items": []}
|
||||
# Studio (the hub) owns conditioning. If it packed N shortcut VARIANTS into positive
|
||||
# (each scene entry tagged 'funpack_batch_variant'), sample one chain per variant and the
|
||||
# variant count wins. Otherwise this is a seed-only batch: one frozen conditioning sampled
|
||||
# batch_n times. Either way each entry's conditioning is saved so Submit trains per entry.
|
||||
variants = self._split_batch_variants(positive)
|
||||
if variants is not None:
|
||||
entries = list(enumerate(variants)) # [(idx, variant_positive_list), ...]
|
||||
mode = "variant"
|
||||
else:
|
||||
entries = [(i, positive) for i in range(batch_n)]
|
||||
mode = "seed-only"
|
||||
# Studio packed N entries into positive (each scene entry tagged 'funpack_batch_variant').
|
||||
# Sample one chain per entry. Entries may be genuinely different (shortcut variants) or
|
||||
# identical (seed-only batch — Studio packed N copies); either way each saves its own cond.
|
||||
variants = self._split_batch_variants(positive) or [positive]
|
||||
entries = list(enumerate(variants)) # [(idx, variant_positive_list), ...]
|
||||
manifest["iterations"] = len(entries)
|
||||
manifest["mode"] = mode
|
||||
last = None
|
||||
base_seed = int(seed)
|
||||
for idx, pos_i in entries:
|
||||
iter_seed = base_seed + idx
|
||||
print(f"[FunPackSceneChain] Batch Training {idx + 1}/{len(entries)} ({mode}, seed={iter_seed})")
|
||||
print(f"[FunPackSceneChain] Batch Training {idx + 1}/{len(entries)} (seed={iter_seed})")
|
||||
out = self.sample(
|
||||
model, vae, pos_i, negative, sampler, sigmas, iter_seed, latent_template,
|
||||
num_frames_per_scene, frame_overlap, cfg, max_scenes, use_same_seed=use_same_seed,
|
||||
carry_i2v_guides=carry_i2v_guides, mid_scene_guide=mid_scene_guide,
|
||||
mid_scene_guide_strength=mid_scene_guide_strength, embed_guidance=embed_guidance,
|
||||
embed_guidance_strength=embed_guidance_strength, transition_duration=transition_duration,
|
||||
batch_iterations=1, decode_tile_size=decode_tile_size, refinement_key_input=key,
|
||||
decode_tile_size=decode_tile_size, refinement_key_input=key,
|
||||
unique_id=None, prompt=None,
|
||||
)
|
||||
last = out
|
||||
@@ -2562,7 +2549,7 @@ class FunPackLTXAVSceneChainSampler:
|
||||
print(f"[FunPackSceneChain] batch decode failed: {e}")
|
||||
manifest["items"].append({
|
||||
"index": idx, "id": iid, "seed": iter_seed,
|
||||
"variant": idx if mode == "variant" else None,
|
||||
"variant": idx,
|
||||
"prompt": self._scene_text(pos_i[0], 0) if pos_i else None,
|
||||
"latent": os.path.basename(latent_path) if latent_path else None,
|
||||
"preview": os.path.basename(preview_path) if has_preview else None,
|
||||
@@ -2574,7 +2561,7 @@ class FunPackLTXAVSceneChainSampler:
|
||||
_json.dump(manifest, f, indent=2)
|
||||
except Exception as e:
|
||||
print(f"[FunPackSceneChain] batch manifest save failed: {e}")
|
||||
status = (f"Batch Training complete ({mode}): {manifest['iterations']} generations in "
|
||||
status = (f"Batch Training complete: {manifest['iterations']} generation(s) in "
|
||||
f"temp/{rel.replace(os.sep, '/')} (cleared on restart) — rate them in Studio.")
|
||||
print(f"[FunPackSceneChain] {status}")
|
||||
if last is None:
|
||||
|
||||
+196
-24
@@ -474,6 +474,169 @@ function showError(panel, msg) {
|
||||
if (tgt) tgt.textContent = msg;
|
||||
}
|
||||
|
||||
// ─── Batch Training ─────────────────────────────────────────────────────────
|
||||
const _fpBatchProcessing = new Set(); // node ids whose armed batch is currently running
|
||||
|
||||
function batchKey(node) {
|
||||
try { return (getSettings(node)?.refinement_key) || linkedRefinementKey(node) || ""; }
|
||||
catch { return ""; }
|
||||
}
|
||||
|
||||
async function refreshBatchButton(node) {
|
||||
const w = (node.widgets || []).find((x) => x._fpBatchBtn);
|
||||
if (!w) return;
|
||||
let label = "Batch Training (off)";
|
||||
try {
|
||||
const n = parseInt(getSettings(node)?.refiner?.batch_variants || 1, 10) || 1;
|
||||
if (_fpBatchProcessing.has(node.id)) {
|
||||
label = "Batch Training (Processing…)";
|
||||
} else {
|
||||
let pending = false;
|
||||
const key = batchKey(node);
|
||||
if (key) {
|
||||
try {
|
||||
const res = await api.fetchApi(`/funpack/batch/list?key=${encodeURIComponent(key)}`, { cache: "no-store" });
|
||||
pending = !!(await res.json()).found;
|
||||
} catch {}
|
||||
}
|
||||
if (pending) label = "Batch Training (Finished — rate)";
|
||||
else if (n > 1) label = `Batch Training (armed ×${n})`;
|
||||
else label = "Batch Training (off)";
|
||||
}
|
||||
} catch {}
|
||||
w.name = label;
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
}
|
||||
|
||||
function refreshAllBatchButtons() {
|
||||
try {
|
||||
for (const node of (app.graph?._nodes || [])) {
|
||||
if (node?.comfyClass === NODE_NAME || node?.type === NODE_NAME) refreshBatchButton(node);
|
||||
}
|
||||
} catch {}
|
||||
}
|
||||
|
||||
function openBatchPanel(node) {
|
||||
closePanel();
|
||||
injectStyles();
|
||||
const root = el("div", "funpack-studio-panel");
|
||||
root.style.cssText = "position:fixed;left:50%;top:50%;transform:translate(-50%,-50%);z-index:10000;width:min(460px,calc(100vw - 24px));";
|
||||
document.body.append(root);
|
||||
activePanel = root;
|
||||
|
||||
const header = el("div", "funpack-studio-header");
|
||||
header.append(el("div", "funpack-studio-title", "Batch Training"));
|
||||
const closeBtn = btn("Close");
|
||||
closeBtn.addEventListener("click", () => { closePanel(); refreshBatchButton(node); });
|
||||
header.append(closeBtn);
|
||||
root.append(header);
|
||||
|
||||
const errorEl = el("div", "funpack-studio-error");
|
||||
root.append(errorEl);
|
||||
const body = el("div", "funpack-studio-body");
|
||||
root.append(body);
|
||||
|
||||
const key = batchKey(node);
|
||||
|
||||
async function render() {
|
||||
body.replaceChildren();
|
||||
errorEl.textContent = "";
|
||||
if (!key) {
|
||||
body.append(el("div", "funpack-studio-hint", "Set a refinement key (Studio → Session) or wire refinement_key_input to use Batch Training."));
|
||||
return;
|
||||
}
|
||||
let data = { found: false, labels: [] };
|
||||
try {
|
||||
const res = await api.fetchApi(`/funpack/batch/list?key=${encodeURIComponent(key)}`, { cache: "no-store" });
|
||||
data = await res.json();
|
||||
} catch (e) {
|
||||
body.append(el("div", "funpack-studio-hint", `Batch list failed: ${e.message}`));
|
||||
return;
|
||||
}
|
||||
if (data.found) renderRatingList(data);
|
||||
else renderConfig();
|
||||
}
|
||||
|
||||
function renderConfig() {
|
||||
const settings = getSettings(node);
|
||||
const n = parseInt(settings?.refiner?.batch_variants || 1, 10) || 1;
|
||||
body.append(el("div", "funpack-studio-hint",
|
||||
"A batch is a set of generations made in one run with everything frozen except the seed (and shortcut options, if your prompt has any). Set the size and Arm it, then press ComfyUI's Generate. When it finishes, reopen this window to rate each one."));
|
||||
const sizeInput = numInput(n > 1 ? n : 10, 2, 64, 1);
|
||||
body.append(row("Batch size", sizeInput));
|
||||
body.append(el("div", "funpack-studio-hint", n > 1 ? `Armed: the next Generate produces ${n}.` : "Not armed — Generate runs a single normal generation."));
|
||||
const actions = el("div", "");
|
||||
actions.style.cssText = "display:flex;gap:6px;margin-top:8px;";
|
||||
const armBtn = btn(n > 1 ? "Re-arm" : "Arm batch");
|
||||
armBtn.addEventListener("click", () => {
|
||||
const v = Math.max(2, Math.min(64, parseInt(sizeInput.value, 10) || 0));
|
||||
const s = getSettings(node); s.refiner = s.refiner || {}; s.refiner.batch_variants = v; saveSettings(node, s);
|
||||
refreshBatchButton(node); render();
|
||||
});
|
||||
const offBtn = btn("Disarm", "secondary");
|
||||
offBtn.addEventListener("click", () => {
|
||||
const s = getSettings(node); s.refiner = s.refiner || {}; s.refiner.batch_variants = 1; saveSettings(node, s);
|
||||
refreshBatchButton(node); render();
|
||||
});
|
||||
actions.append(armBtn, offBtn);
|
||||
body.append(actions);
|
||||
}
|
||||
|
||||
function renderRatingList(data) {
|
||||
const labels = data.labels || [];
|
||||
const ratings = {};
|
||||
body.append(el("div", "funpack-studio-hint",
|
||||
`Finished batch ${data.created} — rate each, then Submit to teach the value function + repair memory. Forget discards the whole batch.`));
|
||||
data.items.forEach((it) => {
|
||||
const r = el("div", "");
|
||||
r.style.cssText = "display:flex;align-items:center;gap:6px;margin:2px 0;";
|
||||
const vtag = (it.variant != null) ? `[v${it.variant}] ` : "";
|
||||
const lab = it.prompt ? it.prompt : (it.preview || it.id);
|
||||
const name = el("span", "", `#${it.index ?? 0} ${vtag}${lab} (seed ${it.seed})`);
|
||||
name.title = it.preview || it.id;
|
||||
name.style.cssText = "flex:1;font-size:12px;opacity:0.85;white-space:nowrap;overflow:hidden;text-overflow:ellipsis;";
|
||||
const prev = btn("Preview", "secondary");
|
||||
let imgEl = null;
|
||||
prev.addEventListener("click", () => {
|
||||
if (imgEl) { imgEl.remove(); imgEl = null; return; }
|
||||
const url = `/view?filename=${encodeURIComponent(it.preview)}&type=temp&subfolder=${encodeURIComponent(data.subfolder)}`;
|
||||
imgEl = el("img", ""); imgEl.src = url; imgEl.style.cssText = "max-width:100%;border-radius:6px;margin:4px 0;"; r.after(imgEl);
|
||||
});
|
||||
if (!it.preview) prev.disabled = true;
|
||||
const sel = selectEl(["— rate —", ...labels], "— rate —");
|
||||
sel.addEventListener("change", () => { if (sel.value && sel.value !== "— rate —") ratings[it.id] = sel.value; else delete ratings[it.id]; });
|
||||
r.append(name, prev, sel);
|
||||
body.append(r);
|
||||
});
|
||||
const actions = el("div", "");
|
||||
actions.style.cssText = "display:flex;gap:6px;margin-top:8px;";
|
||||
const submitBtn = btn("Submit");
|
||||
const forgetBtn = btn("Forget batch", "secondary");
|
||||
submitBtn.addEventListener("click", async () => {
|
||||
if (Object.keys(ratings).length === 0) { errorEl.textContent = "Rate at least one generation first."; return; }
|
||||
submitBtn.disabled = true;
|
||||
try {
|
||||
const res = await api.fetchApi("/funpack/batch/submit", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ key, stamp: data.stamp, ratings }) });
|
||||
const out = await res.json();
|
||||
if (out.error) throw new Error(out.error);
|
||||
errorEl.textContent = `Submitted — trained on ${out.trained}.${out.info ? " " + out.info : ""}`;
|
||||
refreshBatchButton(node); render();
|
||||
} catch (e) { errorEl.textContent = `Submit failed: ${e.message}`; submitBtn.disabled = false; }
|
||||
});
|
||||
forgetBtn.addEventListener("click", async () => {
|
||||
try {
|
||||
const res = await api.fetchApi("/funpack/batch/forget", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ key, stamp: data.stamp }) });
|
||||
const out = await res.json(); if (out.error) throw new Error(out.error);
|
||||
refreshBatchButton(node); render();
|
||||
} catch (e) { errorEl.textContent = `Forget failed: ${e.message}`; }
|
||||
});
|
||||
actions.append(submitBtn, forgetBtn);
|
||||
body.append(actions);
|
||||
}
|
||||
|
||||
render();
|
||||
}
|
||||
|
||||
function openPanel(node) {
|
||||
closePanel();
|
||||
injectStyles();
|
||||
@@ -1661,30 +1824,13 @@ function setupNode(node) {
|
||||
openPanel(node);
|
||||
}, { serialize: false });
|
||||
|
||||
// Batch Training: one-shot — produce N generations (variants if the prompt has shortcut
|
||||
// options, else seed-only), then auto-revert so the next queue is a normal single run.
|
||||
node.widgets = (node.widgets || []).filter((w) => w.name !== "funpack_studio_batch");
|
||||
node.addWidget("button", "Batch Training", "funpack_studio_batch", async () => {
|
||||
const raw = window.prompt(
|
||||
"Batch Training — how many generations?\n(Different shortcut variants if your prompt has options; otherwise the same prompt with different seeds. Rate them in Studio afterwards.)",
|
||||
"10",
|
||||
);
|
||||
if (raw == null) return;
|
||||
const n = Math.max(2, Math.min(64, parseInt(raw, 10) || 0));
|
||||
if (n < 2) return;
|
||||
const settings = getSettings(node);
|
||||
settings.refiner = settings.refiner || {};
|
||||
settings.refiner.batch_variants = n;
|
||||
saveSettings(node, settings);
|
||||
try {
|
||||
await app.queuePrompt(0, 1);
|
||||
} finally {
|
||||
const s2 = getSettings(node);
|
||||
s2.refiner = s2.refiner || {};
|
||||
s2.refiner.batch_variants = 1;
|
||||
saveSettings(node, s2);
|
||||
}
|
||||
}, { serialize: false });
|
||||
// Batch Training: opens its own window to arm a batch or rate a finished one. It never
|
||||
// queues — ComfyUI's Generate runs an armed batch. The label reflects live state.
|
||||
node.widgets = (node.widgets || []).filter((w) => !w._fpBatchBtn && w.name !== "funpack_studio_batch");
|
||||
const batchBtn = node.addWidget("button", "Batch Training (off)", "funpack_studio_batch",
|
||||
() => openBatchPanel(node), { serialize: false });
|
||||
batchBtn._fpBatchBtn = true;
|
||||
refreshBatchButton(node);
|
||||
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
}
|
||||
@@ -1928,6 +2074,32 @@ function injectStyles() {
|
||||
|
||||
app.registerExtension({
|
||||
name: "funpack.studio",
|
||||
setup() {
|
||||
// Drive the Batch Training button label: Processing while an armed batch runs, then
|
||||
// Finished (and one-shot disarm so the next Generate is normal).
|
||||
const armedStudioNodes = () => (app.graph?._nodes || []).filter(
|
||||
(n) => (n?.comfyClass === NODE_NAME || n?.type === NODE_NAME) &&
|
||||
parseInt(getSettings(n)?.refiner?.batch_variants || 1, 10) > 1);
|
||||
api.addEventListener("execution_start", () => {
|
||||
for (const n of armedStudioNodes()) _fpBatchProcessing.add(n.id);
|
||||
refreshAllBatchButtons();
|
||||
});
|
||||
const onEnd = () => {
|
||||
if (_fpBatchProcessing.size === 0) return;
|
||||
for (const n of (app.graph?._nodes || [])) {
|
||||
if (!_fpBatchProcessing.has(n.id)) continue;
|
||||
_fpBatchProcessing.delete(n.id);
|
||||
const s = getSettings(n); // one-shot: disarm after the batch ran
|
||||
if (s?.refiner && parseInt(s.refiner.batch_variants || 1, 10) > 1) {
|
||||
s.refiner.batch_variants = 1;
|
||||
saveSettings(n, s);
|
||||
}
|
||||
}
|
||||
setTimeout(refreshAllBatchButtons, 400); // give the manifest a moment to land
|
||||
};
|
||||
api.addEventListener("execution_success", onEnd);
|
||||
api.addEventListener("execution_error", onEnd);
|
||||
},
|
||||
beforeRegisterNodeDef(nodeType, nodeData) {
|
||||
if (nodeData.name !== NODE_NAME) return;
|
||||
const orig = nodeType.prototype.onNodeCreated;
|
||||
|
||||
Reference in New Issue
Block a user