digital-garbageandClaude Opus 4.8
c8e4f2c404
UI: Batch Training tab gains Mode (Regular / Interactive Guessing), Direction, Range, and a Learning toggle; Arm sets refiner.guess_* + batch_learning; auto-disarm clears guess_mode; rating list shows each rung's [x factor]; Submit passes the learn flag. Learning (batch_training.py): submit honors learn=false (pure generation, no training). Guess batches detected via guess_factor -> _learn_guess_ceiling finds the furthest-pushed acceptable rung before quality broke and stores a per-key safe ceiling (absolute spread + factor) via _store_guess_ceiling. Dual auto-cap (refine_v2, normal runs only): soft pre-limit on absolute_strength + hard clamp of the output conditioning's video-channel spread to the learned-safe value in _v2_finalize_conditioning. Skipped during the guessing ladder itself. Helpers conditioning_spread / scale_conditioning_spread. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
315 lines
13 KiB
Python
315 lines
13 KiB
Python
"""Batch Training server routes + learning intake.
|
||
|
||
Serves the Studio batch-rating panel and, on submit, trains the value function on the
|
||
batch's (frozen) conditioning with each entry's rating — same cond, N rewards = the
|
||
variance-reduction RLHF signal. Batches live in ComfyUI's temp dir (wiped on restart).
|
||
"""
|
||
|
||
import os
|
||
import re
|
||
import json
|
||
import shutil
|
||
|
||
import torch
|
||
|
||
try:
|
||
from aiohttp import web
|
||
from server import PromptServer
|
||
except Exception: # pragma: no cover - only available inside ComfyUI
|
||
web = None
|
||
PromptServer = None
|
||
|
||
|
||
def _safe_key(key):
|
||
return re.sub(r"[^A-Za-z0-9_.-]", "_", str(key or "default").strip() or "default")
|
||
|
||
|
||
def _batches_root():
|
||
try:
|
||
import folder_paths
|
||
base = folder_paths.get_temp_directory()
|
||
except Exception:
|
||
base = os.path.join(os.path.dirname(os.path.abspath(__file__)), "refinements", "tmp")
|
||
return os.path.join(base, "funpack_batches")
|
||
|
||
|
||
def _key_dir(key):
|
||
return os.path.join(_batches_root(), _safe_key(key))
|
||
|
||
|
||
def _latest_batch(key):
|
||
"""Most recent batch stamp dir for a key (stamps sort lexically), or None."""
|
||
kd = _key_dir(key)
|
||
if not os.path.isdir(kd):
|
||
return None
|
||
stamps = sorted(
|
||
(d for d in os.listdir(kd) if os.path.isdir(os.path.join(kd, d))),
|
||
reverse=True,
|
||
)
|
||
for stamp in stamps:
|
||
d = os.path.join(kd, stamp)
|
||
if os.path.exists(os.path.join(d, "batch.json")):
|
||
return stamp, d
|
||
return None
|
||
|
||
|
||
def _load_manifest(d):
|
||
with open(os.path.join(d, "batch.json"), "r") as f:
|
||
return json.load(f)
|
||
|
||
|
||
def _rating_labels():
|
||
try:
|
||
try:
|
||
from .conditioning import V2_RATING_LABELS
|
||
except ImportError:
|
||
from conditioning import V2_RATING_LABELS
|
||
return list(V2_RATING_LABELS)
|
||
except Exception:
|
||
return ["Perfect", "Loved it", "Nailed it", "Missing details", "Missing action",
|
||
"Missing quality", "Awful", "-Just forget it-"]
|
||
|
||
|
||
def _profile_for(label):
|
||
try:
|
||
try:
|
||
from .conditioning import normalize_refiner_v2_rating
|
||
except ImportError:
|
||
from conditioning import normalize_refiner_v2_rating
|
||
return normalize_refiner_v2_rating(label)
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _train_batch(key, manifest, batch_dir, ratings):
|
||
"""Train the value function on (cond, reward) for each rated entry — each entry carries its
|
||
own conditioning (identical across a seed-only batch, distinct per shortcut variant)."""
|
||
try:
|
||
try:
|
||
from .value_function import OnlineValueFunction
|
||
from .conditioning import refinement_state_path, tensor_to_serializable, FunPackVideoRefinerV2
|
||
except ImportError:
|
||
from value_function import OnlineValueFunction
|
||
from conditioning import refinement_state_path, tensor_to_serializable, FunPackVideoRefinerV2
|
||
except Exception as e:
|
||
return 0, f"learning modules unavailable: {e}"
|
||
|
||
vf_path = refinement_state_path(key, "value_fn", prefix="refine_v2", extension="pt")
|
||
trained = 0
|
||
rated_items = [] # for Phase 3b axis/direction learning
|
||
with torch.inference_mode(False):
|
||
vf = None
|
||
for item in manifest.get("items", []):
|
||
label = ratings.get(str(item.get("id")))
|
||
if not label:
|
||
continue
|
||
profile = _profile_for(label)
|
||
if not profile or profile.get("skip_learning"):
|
||
continue # '-Just forget it-' and unknowns: skip
|
||
cond_file = item.get("cond")
|
||
if not cond_file:
|
||
continue
|
||
cond_path = os.path.join(batch_dir, cond_file)
|
||
if not os.path.exists(cond_path):
|
||
continue
|
||
cond = torch.load(cond_path, map_location="cpu")
|
||
if not isinstance(cond, torch.Tensor):
|
||
continue
|
||
if vf is None:
|
||
vf = OnlineValueFunction.load_or_create(vf_path, hidden_dim=int(cond.shape[-1]))
|
||
if vf is None:
|
||
return 0, "could not create value function"
|
||
reward = float(profile.get("reward", 0.0))
|
||
vf.train_on(cond.clone(), reward) # 3a: value function
|
||
try:
|
||
rated_items.append({"rating": label, "conditioning": tensor_to_serializable(cond)})
|
||
except Exception:
|
||
pass
|
||
trained += 1
|
||
if trained and vf is not None:
|
||
vf.save(vf_path)
|
||
if trained == 0:
|
||
return 0, "no rated entries with saved conditioning"
|
||
# 3b: feed the per-axis direction memory + persistent repair bias (axes consistently rated
|
||
# missing/wrong across the frozen batch = high-confidence signal for the next generation).
|
||
axis_note = None
|
||
try:
|
||
axis_note = FunPackVideoRefinerV2()._v2_learn_from_batch(key, rated_items)
|
||
print(f"[FunPack batch] {axis_note}")
|
||
except Exception as e:
|
||
axis_note = f"axis learning skipped: {e}"
|
||
return trained, axis_note
|
||
|
||
|
||
def _store_guess_ceiling(key, factor, spread, direction):
|
||
"""Write the learned Interactive-Guessing safe ceiling into the refiner state (global) for
|
||
`key`, preserving everything else. guess_safe_spread is the absolute cap the refiner clamps
|
||
future output conditioning to; guess_safe_factor is for display."""
|
||
try:
|
||
try:
|
||
from .conditioning import refinement_state_path
|
||
except ImportError:
|
||
from conditioning import refinement_state_path
|
||
path = refinement_state_path(key, "clip", prefix="refine_v2")
|
||
data = {}
|
||
if os.path.exists(path):
|
||
with open(path, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
if not isinstance(data, dict):
|
||
data = {}
|
||
data.setdefault("version", 2)
|
||
g = data.setdefault("global", {})
|
||
g["guess_safe_factor"] = float(factor)
|
||
if spread is not None:
|
||
g["guess_safe_spread"] = float(spread)
|
||
g["guess_direction"] = str(direction)
|
||
with open(path, "w", encoding="utf-8") as f:
|
||
json.dump(data, f)
|
||
return True
|
||
except Exception as e:
|
||
print(f"[FunPack guess] store ceiling failed: {e}")
|
||
return False
|
||
|
||
|
||
def _learn_guess_ceiling(key, manifest, batch_dir, ratings):
|
||
"""Interactive Guessing learning: from the rated ladder, find the safe steering ceiling — the
|
||
furthest-pushed rung still rated acceptable before quality broke — and store it (absolute
|
||
spread + factor) so the refiner auto-caps future steering. Reward >= 0 = acceptable."""
|
||
try:
|
||
from .conditioning import conditioning_spread
|
||
except ImportError:
|
||
try:
|
||
from conditioning import conditioning_spread
|
||
except Exception:
|
||
conditioning_spread = lambda c: None
|
||
rated = []
|
||
for item in manifest.get("items", []):
|
||
gf = item.get("guess_factor")
|
||
if gf is None:
|
||
continue
|
||
label = ratings.get(str(item.get("id")))
|
||
if not label:
|
||
continue
|
||
prof = _profile_for(label)
|
||
if not prof or prof.get("skip_learning"):
|
||
continue
|
||
rated.append((float(gf), float(prof.get("reward", 0.0)), item.get("cond")))
|
||
if not rated:
|
||
return "Interactive Guessing: no rated rungs — nothing learned."
|
||
up = sum(1 for f, _, _ in rated if f > 1.0) >= sum(1 for f, _, _ in rated if f < 1.0)
|
||
broke = [f for f, r, _ in rated if r < 0.0]
|
||
ok = [(f, c) for f, r, c in rated if r >= 0.0]
|
||
if up:
|
||
fail = min(broke) if broke else float("inf")
|
||
candidates = [(f, c) for f, c in ok if f < fail]
|
||
safe = max(candidates, key=lambda x: x[0]) if candidates else None
|
||
else:
|
||
fail = max(broke) if broke else float("-inf")
|
||
candidates = [(f, c) for f, c in ok if f > fail]
|
||
safe = min(candidates, key=lambda x: x[0]) if candidates else None
|
||
if safe is None:
|
||
# even the base broke — cap at base (1.0), no headroom
|
||
_store_guess_ceiling(key, 1.0, None, "up" if up else "down")
|
||
return "Interactive Guessing: every rung broke — capped steering to base (×1.00)."
|
||
safe_factor, safe_cond_file = safe
|
||
safe_spread = None
|
||
if safe_cond_file:
|
||
cond_path = os.path.join(batch_dir, safe_cond_file)
|
||
if os.path.exists(cond_path):
|
||
try:
|
||
with torch.inference_mode(False):
|
||
safe_spread = conditioning_spread(torch.load(cond_path, map_location="cpu"))
|
||
except Exception:
|
||
safe_spread = None
|
||
_store_guess_ceiling(key, safe_factor, safe_spread, "up" if up else "down")
|
||
return (f"Interactive Guessing: safe up to ×{safe_factor:.2f}"
|
||
+ (f" (spread {safe_spread:.3f}) — future steering auto-capped." if safe_spread
|
||
else " — future steering auto-capped."))
|
||
|
||
|
||
if PromptServer is not None and web is not None:
|
||
|
||
@PromptServer.instance.routes.get("/funpack/batch/list")
|
||
async def funpack_batch_list(request):
|
||
key = request.rel_url.query.get("key", "")
|
||
if not str(key).strip():
|
||
return web.json_response({"found": False, "reason": "no key", "labels": _rating_labels()})
|
||
latest = _latest_batch(key)
|
||
if latest is None:
|
||
return web.json_response({"found": False, "labels": _rating_labels()})
|
||
stamp, d = latest
|
||
try:
|
||
manifest = _load_manifest(d)
|
||
except Exception as e:
|
||
return web.json_response({"found": False, "reason": f"manifest read failed: {e}",
|
||
"labels": _rating_labels()})
|
||
subfolder = manifest.get("subfolder") or f"funpack_batches/{_safe_key(key)}/{stamp}"
|
||
items = [{
|
||
"id": it.get("id"),
|
||
"index": it.get("index"),
|
||
"seed": it.get("seed"),
|
||
"preview": it.get("preview"),
|
||
"variant": it.get("variant"),
|
||
"guess_factor": it.get("guess_factor"),
|
||
"prompt": it.get("prompt"),
|
||
"rating": it.get("rating"),
|
||
} for it in manifest.get("items", [])]
|
||
return web.json_response({
|
||
"found": True,
|
||
"key": key,
|
||
"stamp": stamp,
|
||
"subfolder": subfolder,
|
||
"created": manifest.get("created"),
|
||
"scene_prompts": manifest.get("scene_prompts", []),
|
||
"items": items,
|
||
"labels": _rating_labels(),
|
||
})
|
||
|
||
@PromptServer.instance.routes.post("/funpack/batch/submit")
|
||
async def funpack_batch_submit(request):
|
||
try:
|
||
body = await request.json()
|
||
except Exception:
|
||
return web.json_response({"error": "invalid JSON"}, status=400)
|
||
key = str(body.get("key", "")).strip()
|
||
stamp = str(body.get("stamp", "")).strip()
|
||
ratings = body.get("ratings", {}) or {}
|
||
learn = bool(body.get("learn", True))
|
||
if not key or not stamp:
|
||
return web.json_response({"error": "key and stamp required"}, status=400)
|
||
batch_dir = os.path.join(_key_dir(key), stamp)
|
||
if not os.path.exists(os.path.join(batch_dir, "batch.json")):
|
||
return web.json_response({"error": "batch not found"}, status=404)
|
||
manifest = _load_manifest(batch_dir)
|
||
ratings = {str(k): v for k, v in ratings.items()}
|
||
is_guess = any(it.get("guess_factor") is not None for it in manifest.get("items", []))
|
||
if not learn:
|
||
trained, note = 0, "learning off — generated only, nothing learned."
|
||
elif is_guess:
|
||
note = _learn_guess_ceiling(key, manifest, batch_dir, ratings)
|
||
trained = 1 if "safe" in (note or "") else 0
|
||
else:
|
||
trained, note = _train_batch(key, manifest, batch_dir, ratings)
|
||
# Submit clears the batch regardless (rated material is consumed).
|
||
try:
|
||
shutil.rmtree(batch_dir, ignore_errors=True)
|
||
except Exception:
|
||
pass
|
||
if trained == 0:
|
||
return web.json_response({"ok": True, "trained": 0, "warning": note or "nothing trained"})
|
||
return web.json_response({"ok": True, "trained": trained, "info": note})
|
||
|
||
@PromptServer.instance.routes.post("/funpack/batch/forget")
|
||
async def funpack_batch_forget(request):
|
||
try:
|
||
body = await request.json()
|
||
except Exception:
|
||
return web.json_response({"error": "invalid JSON"}, status=400)
|
||
key = str(body.get("key", "")).strip()
|
||
stamp = str(body.get("stamp", "")).strip()
|
||
if not key or not stamp:
|
||
return web.json_response({"error": "key and stamp required"}, status=400)
|
||
batch_dir = os.path.join(_key_dir(key), stamp)
|
||
shutil.rmtree(batch_dir, ignore_errors=True)
|
||
return web.json_response({"ok": True})
|