From d423d313fe1c094ec8446ff903c675436bbbab02 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Tue, 2 Jun 2026 17:49:52 +0300 Subject: [PATCH] Batch Training: sampler consumes Studio variant-packed conditioning (Phase 4 contract) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per the Studio/sampler boundary law (Studio owns conditioning; sampler only samples), the sampler no longer encodes anything. It now auto-detects when Studio has packed N shortcut VARIANTS into positive — each scene entry tagged with 'funpack_batch_variant' = variant index — via _split_batch_variants: if present, it samples one chain per variant (variant count wins); otherwise it falls back to seed-only repeats of batch_iterations. Decisions: 'sampler loops, Studio varies' + 'pack into positive + markers'. Each entry now saves its OWN conditioning (.cond.pt) — identical across a seed-only batch, distinct per variant — and the manifest records variant index + prompt per item. Submit (batch_training.py) trains the VF per entry from its own cond, so variant batches teach the model which prompt variant rated better (not just seed-quality variance). Studio panel shows the variant tag + prompt per row. Studio-side variant PRODUCTION (resolve shortcuts N times, encode, pack markers) is the remaining half — a Studio/refiner feature, tracked next. Sampler side is ready for it. Co-Authored-By: Claude Opus 4.8 --- batch_training.py | 37 +++++++++++++---------- samplers.py | 69 ++++++++++++++++++++++++++++++++----------- web/funpack_studio.js | 7 +++-- 3 files changed, 77 insertions(+), 36 deletions(-) diff --git a/batch_training.py b/batch_training.py index 4eccd972..66eec12f 100644 --- a/batch_training.py +++ b/batch_training.py @@ -82,13 +82,8 @@ def _profile_for(label): def _train_batch(key, manifest, batch_dir, ratings): - """Train the value function on (cond, reward) for each rated entry. Returns count trained.""" - cond_file = manifest.get("cond") - if not cond_file: - return 0, "no conditioning saved for this batch — cannot train" - cond_path = os.path.join(batch_dir, cond_file) - if not os.path.exists(cond_path): - return 0, "conditioning file missing" + """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 @@ -99,15 +94,10 @@ def _train_batch(key, manifest, batch_dir, ratings): 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 with torch.inference_mode(False): - cond = torch.load(cond_path, map_location="cpu") - if not isinstance(cond, torch.Tensor): - return 0, "saved conditioning is not a tensor" - vf_path = refinement_state_path(key, "value_fn", prefix="refine_v2", extension="pt") - vf = OnlineValueFunction.load_or_create(vf_path, hidden_dim=int(cond.shape[-1])) - if vf is None: - return 0, "could not create value function" - trained = 0 + vf = None for item in manifest.get("items", []): label = ratings.get(str(item.get("id"))) if not label: @@ -115,11 +105,26 @@ def _train_batch(key, manifest, batch_dir, ratings): 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) trained += 1 - if trained: + if trained and vf is not None: vf.save(vf_path) + if trained == 0: + return 0, "no rated entries with saved conditioning" return trained, None diff --git a/samplers.py b/samplers.py index 63cd0c72..dfa4ba52 100644 --- a/samplers.py +++ b/samplers.py @@ -2419,6 +2419,24 @@ class FunPackLTXAVSceneChainSampler: os.makedirs(d, exist_ok=True) return d, safe, rel + def _split_batch_variants(self, positive): + """If Studio packed shortcut variants into positive (each scene entry tagged with + 'funpack_batch_variant' = variant index), return a list of per-variant positive lists in + variant order. None when there are no markers or only one variant (= normal seed-only).""" + groups = {} + for entry in positive: + meta = entry[1] if isinstance(entry, (list, tuple)) and len(entry) > 1 and isinstance(entry[1], dict) else None + vi = meta.get("funpack_batch_variant") if isinstance(meta, dict) else None + if vi is None: + return None + try: + groups.setdefault(int(vi), []).append(entry) + except (TypeError, ValueError): + return None + if len(groups) <= 1: + return None + return [groups[k] for k in sorted(groups)] + def _save_batch_preview(self, decoded, path, max_frames=16, width=256): """Save a decoded video tensor [T,H,W,C] in 0..1 as a downscaled animated webp.""" try: @@ -2488,24 +2506,28 @@ class FunPackLTXAVSceneChainSampler: 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, "iterations": int(batch_n), + manifest = {"key": key, "created": stamp, "subfolder": rel.replace(os.sep, "/"), "scene_prompts": scene_prompts, "items": []} - # Save the (frozen) conditioning that produced this batch so Submit can train the value - # function on (cond, rating) for every entry — same cond, N rewards = variance reduction. - try: - cond0 = positive[0][0] if isinstance(positive[0], (list, tuple)) else None - if isinstance(cond0, torch.Tensor): - torch.save(cond0.detach().cpu(), os.path.join(batch_dir, "cond.pt")) - manifest["cond"] = "cond.pt" - except Exception as e: - print(f"[FunPackSceneChain] batch cond save failed: {e}") + # 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" + manifest["iterations"] = len(entries) + manifest["mode"] = mode last = None base_seed = int(seed) - for i in range(batch_n): - iter_seed = base_seed + i - print(f"[FunPackSceneChain] Batch Training {i + 1}/{batch_n} (seed={iter_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})") out = self.sample( - model, vae, positive, negative, sampler, sigmas, iter_seed, latent_template, + 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, @@ -2515,9 +2537,17 @@ class FunPackLTXAVSceneChainSampler: ) last = out out_latent = out[0] - iid = f"{safe}_{stamp}_{i:02d}" + iid = f"{safe}_{stamp}_{idx:02d}" latent_path = os.path.join(batch_dir, f"{iid}.latent.pt") preview_path = os.path.join(batch_dir, f"{iid}.webp") + cond_name = None + try: + cond_i = pos_i[0][0] if pos_i and isinstance(pos_i[0], (list, tuple)) else None + if isinstance(cond_i, torch.Tensor): + cond_name = f"{iid}.cond.pt" + torch.save(cond_i.detach().cpu(), os.path.join(batch_dir, cond_name)) + except Exception as e: + print(f"[FunPackSceneChain] batch cond save failed: {e}") try: torch.save({"samples": self._latent_tensors(out_latent)[0].detach().cpu()}, latent_path) except Exception as e: @@ -2531,9 +2561,12 @@ class FunPackLTXAVSceneChainSampler: except Exception as e: print(f"[FunPackSceneChain] batch decode failed: {e}") manifest["items"].append({ - "index": i, "id": iid, "seed": iter_seed, + "index": idx, "id": iid, "seed": iter_seed, + "variant": idx if mode == "variant" else None, + "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, + "cond": cond_name, "rating": None, }) try: @@ -2541,8 +2574,8 @@ 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: {batch_n} generations in temp/{rel.replace(os.sep, '/')} " - f"(cleared on restart) — rate them in Studio.") + status = (f"Batch Training complete ({mode}): {manifest['iterations']} generations in " + f"temp/{rel.replace(os.sep, '/')} (cleared on restart) — rate them in Studio.") print(f"[FunPackSceneChain] {status}") if last is None: raise RuntimeError("Batch Training produced no output.") diff --git a/web/funpack_studio.js b/web/funpack_studio.js index 0ee58dbc..623632f6 100644 --- a/web/funpack_studio.js +++ b/web/funpack_studio.js @@ -1096,8 +1096,11 @@ function openPanel(node) { data.items.forEach((it) => { const r = el("div", ""); r.style.cssText = "display:flex;align-items:center;gap:6px;margin:2px 0;"; - const name = el("span", "", `#${(it.index ?? 0)} ${it.preview || it.id} (seed ${it.seed})`); - name.style.cssText = "flex:1;font-size:12px;opacity:0.85;"; + const vtag = (it.variant != null) ? `[v${it.variant}] ` : ""; + const label = it.prompt ? it.prompt : (it.preview || it.id); + const name = el("span", "", `#${(it.index ?? 0)} ${vtag}${label} (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", () => {