From 44e694695e6c5cdeb804de09e9d4b6af39e6a2d2 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Fri, 29 May 2026 03:45:18 +0300 Subject: [PATCH] Add online value function for reward-guided sampling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit value_function.py: OnlineValueFunction MLP that trains incrementally on (conditioning, reward) pairs. Uses replay buffer of 100 samples, 20 gradient steps per update. save/load to .pt file via refinement_state_path. Studio popup: new "Value guidance" toggle. When on, trains the value function after every rated generation (good and bad). Falls back to no-op if conditioning deserialization fails. chain sampler: _load_value_function reads the trained MLP if ready (5+ samples). _build_embed_guidance_wrapper uses value function gradient (∂reward/∂conditioning) when available, fixed liked direction as fallback. Log reports which mode is active and sample count. Co-Authored-By: Claude Sonnet 4.6 --- conditioning.py | 23 ++++++++- samplers.py | 48 +++++++++++++----- value_function.py | 112 ++++++++++++++++++++++++++++++++++++++++++ web/funpack_studio.js | 8 ++- 4 files changed, 177 insertions(+), 14 deletions(-) create mode 100644 value_function.py diff --git a/conditioning.py b/conditioning.py index c6c6a7cf..83bbe545 100644 --- a/conditioning.py +++ b/conditioning.py @@ -11120,7 +11120,8 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): refinement_key_input="", positive_conditioning=None, clip_vision_output=None, source_image=None, model=None, mode="Refine", advisor_mode="Off", advisor_thinking=True, advisor_clip=None, feedback_prompt="", prompt_repair=True, temporal_style="natural", - split_by_transitions=False, split_transition_placement="start", reference_injection=False, latent=None, seed_output_connected=False, + split_by_transitions=False, split_transition_placement="start", reference_injection=False, + value_guidance=False, latent=None, seed_output_connected=False, _seed=None, _seed_source="fresh seed", _scene_seeds=None): seed = int(_seed) if _seed is not None else random.randint(1, 0xffffffffffffffff) encode_cache = {} @@ -11224,6 +11225,24 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): ) memory_status = f"{memory_status}\n{seed_memory_status}\n{scene_sync_status}\n{intent_family_status}\n{intent_learning_status}" self._v2_update_conditioning_memory(global_state, previous_run, learning_profile, axis_feedback) + if value_guidance and has_previous_run and refinement_key and not learning_profile.get("skip_learning"): + try: + try: + from .value_function import OnlineValueFunction + except ImportError: + from value_function import OnlineValueFunction + payload = (previous_run or {}).get("conditioning") + if isinstance(payload, dict): + cond_tensor = serializable_to_tensor(payload) + reward = float(learning_profile.get("reward", 0.0)) + vf_path = refinement_state_path(refinement_key, "value_fn", prefix="refine_v2", extension="pt") + vf = OnlineValueFunction.load_or_create(vf_path, hidden_dim=cond_tensor.shape[-1]) + if vf is not None: + vf.train_on(cond_tensor, reward) + vf.save(vf_path) + print(f"[FunPackRefiner] Value function updated — {vf.n_trained} samples, buffer={len(vf.buffer_c)}") + except Exception as e: + print(f"[FunPackRefiner] Value function training failed: {e}") if has_previous_run and not learning_profile.get("skip_learning"): self._v2_update_streaks(global_state, learning_profile, update_conditioning_strength=not prompt_only_mode) repair_feedback, repair_persistence_status = self._v2_active_repair_feedback( @@ -12938,6 +12957,7 @@ class FunPackStudio: split_transition_placement = "start" reference_injection = bool(rf.get("reference_injection", False)) vision_conditioning = bool(rf.get("vision_conditioning", True)) + value_guidance = bool(rf.get("value_guidance", False)) # feedback_prompt: popup wins if override is on, else external wins popup_feedback = str(rf.get("feedback_prompt", "") or "") @@ -13065,6 +13085,7 @@ class FunPackStudio: split_by_transitions=split_by_transitions, split_transition_placement=split_transition_placement, reference_injection=reference_injection, + value_guidance=value_guidance, latent=latent, seed_output_connected=seed_output_connected, _seed=seed, diff --git a/samplers.py b/samplers.py index d63fb4ce..127174f3 100644 --- a/samplers.py +++ b/samplers.py @@ -1287,6 +1287,23 @@ class FunPackLTXAVSceneChainSampler: except Exception: return None + def _load_value_function(self, refinement_key): + """Load the online value function if trained and ready.""" + try: + try: + from .value_function import OnlineValueFunction + from .conditioning import refinement_state_path + except ImportError: + from value_function import OnlineValueFunction + from conditioning import refinement_state_path + path = refinement_state_path(refinement_key, "value_fn", prefix="refine_v2", extension="pt") + if not __import__("os").path.exists(path): + return None + vf = OnlineValueFunction.load(path) + return vf if vf.is_ready() else None + except Exception: + return None + def _load_liked_direction(self, refinement_key): """Read the liked conditioning direction from the Refiner's state file.""" try: @@ -1309,15 +1326,14 @@ class FunPackLTXAVSceneChainSampler: except Exception: return None - def _build_embed_guidance_wrapper(self, model, liked_dir, strength): + def _build_embed_guidance_wrapper(self, model, liked_dir, strength, value_fn=None): """Register a model_function_wrapper that nudges conditioning toward the - liked quality direction at each denoising step, scaling with sigma so - the effect is stronger during detail refinement (low sigma) than structure - formation (high sigma).""" + liked quality direction at each denoising step. Uses value function gradient + when available, falls back to the fixed liked direction otherwise.""" old_wrapper = model.model_options.get("model_function_wrapper") - direction = torch.nn.functional.normalize(liked_dir.float(), dim=-1) + fixed_dir = torch.nn.functional.normalize(liked_dir.float(), dim=-1) - def _embed_wrapper(apply_fn, args, _ew=old_wrapper, _dir=direction, _s=strength): + def _embed_wrapper(apply_fn, args, _ew=old_wrapper, _fixed=fixed_dir, _vf=value_fn, _s=strength): c = args.get("c") or {} cond = c.get("c_crossattn") if cond is not None: @@ -1326,13 +1342,18 @@ class FunPackLTXAVSceneChainSampler: sigma = float(ts.max().item()) if ts is not None else 1.0 except Exception: sigma = 1.0 - # Scale: no effect at sigma=1 (pure noise), full effect below sigma=0.5 scale = max(0.0, 1.0 - sigma * 2.0) if scale > 0: - d = _dir.to(cond.device, cond.dtype) - new_cond = cond + (_s * scale) * d.expand_as(cond) + if _vf is not None: + try: + grad = _vf.gradient(cond) + d = torch.nn.functional.normalize(grad.float(), dim=-1).to(cond.dtype) + except Exception: + d = _fixed.to(cond.device, cond.dtype).expand_as(cond) + else: + d = _fixed.to(cond.device, cond.dtype).expand_as(cond) new_c = dict(c) - new_c["c_crossattn"] = new_cond + new_c["c_crossattn"] = cond + (_s * scale) * d args = dict(args) args["c"] = new_c if _ew is not None: @@ -1436,12 +1457,15 @@ class FunPackLTXAVSceneChainSampler: # Load liked direction once for embed_guidance _liked_dir = None + _value_fn = None if embed_guidance and refinement_key_input: _liked_dir = self._load_liked_direction(refinement_key_input) if _liked_dir is None: print("[FunPackSceneChain] embed_guidance: no liked direction found (need 3+ liked generations)") else: - print(f"[FunPackSceneChain] embed_guidance: active, strength={embed_guidance_strength}, direction shape={list(_liked_dir.shape)}") + _value_fn = self._load_value_function(refinement_key_input) + mode = f"value function ({_value_fn.n_trained} samples)" if _value_fn else "fixed direction" + print(f"[FunPackSceneChain] embed_guidance: active via {mode}, strength={embed_guidance_strength}") first_scene_seed = self._scene_seed(scene_conditionings[0]) if first_scene_seed is None: @@ -1487,7 +1511,7 @@ class FunPackLTXAVSceneChainSampler: guide_tail = 0 if embed_guidance and _liked_dir is not None: - _eg_old_wrapper = self._build_embed_guidance_wrapper(model, _liked_dir, embed_guidance_strength) + _eg_old_wrapper = self._build_embed_guidance_wrapper(model, _liked_dir, embed_guidance_strength, value_fn=_value_fn) sampled = self._sample_chunk( model, sampler, sigmas, scene_seed, cfg, scene_positive, scene_negative, chunk, ) diff --git a/value_function.py b/value_function.py new file mode 100644 index 00000000..f2eb38e4 --- /dev/null +++ b/value_function.py @@ -0,0 +1,112 @@ +"""Online value function for reward-guided sampling. + +Trains a small MLP incrementally on user-rated generations. +At inference, provides ∂reward/∂conditioning as a per-step steering gradient. +""" + +import os +import random +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class OnlineValueFunction(nn.Module): + MIN_SAMPLES = 5 + BUFFER_SIZE = 100 + BATCH_SIZE = 16 + TRAIN_STEPS = 20 + + def __init__(self, hidden_dim): + super().__init__() + self.hidden_dim = hidden_dim + mid = min(256, hidden_dim // 4) + self.net = nn.Sequential( + nn.Linear(hidden_dim, mid), + nn.SiLU(), + nn.Linear(mid, 64), + nn.SiLU(), + nn.Linear(64, 1), + ) + self._optimizer = None + self.buffer_c = [] # compressed conditioning tensors [hidden_dim] + self.buffer_r = [] # reward floats + self.n_trained = 0 + + @property + def optimizer(self): + if self._optimizer is None: + self._optimizer = torch.optim.Adam(self.parameters(), lr=2e-3) + return self._optimizer + + @staticmethod + def compress(conditioning): + """[1, seq, D] or [seq, D] → [D] mean vector.""" + c = conditioning.float() + if c.dim() == 3: + c = c.squeeze(0) + return c.mean(dim=0) + + def forward(self, c): + return self.net(c) + + def train_on(self, conditioning, reward): + """Add a new (conditioning, reward) sample and do a few training steps.""" + c = self.compress(conditioning).detach().cpu() + self.buffer_c.append(c) + self.buffer_r.append(float(reward)) + if len(self.buffer_c) > self.BUFFER_SIZE: + self.buffer_c = self.buffer_c[-self.BUFFER_SIZE:] + self.buffer_r = self.buffer_r[-self.BUFFER_SIZE:] + if len(self.buffer_c) < 2: + return + device = next(self.parameters()).device + for _ in range(self.TRAIN_STEPS): + idx = random.sample(range(len(self.buffer_c)), min(self.BATCH_SIZE, len(self.buffer_c))) + c_b = torch.stack([self.buffer_c[i] for i in idx]).to(device) + r_b = torch.tensor([self.buffer_r[i] for i in idx], device=device).unsqueeze(1) + self.optimizer.zero_grad() + F.mse_loss(self.forward(c_b), r_b).backward() + self.optimizer.step() + self.n_trained += 1 + + def gradient(self, conditioning): + """∂reward/∂conditioning — same shape as input.""" + c_in = conditioning.detach().float().requires_grad_(True) + reward = self.forward(self.compress(c_in).unsqueeze(0)) + reward.backward() + return c_in.grad.to(conditioning.dtype) + + def is_ready(self): + return len(self.buffer_c) >= self.MIN_SAMPLES + + def save(self, path): + os.makedirs(os.path.dirname(path), exist_ok=True) + buf = torch.stack(self.buffer_c) if self.buffer_c else torch.zeros(0, self.hidden_dim) + torch.save({ + "hidden_dim": self.hidden_dim, + "state_dict": self.state_dict(), + "buffer_c": buf, + "buffer_r": list(self.buffer_r), + "n_trained": self.n_trained, + }, path) + + @classmethod + def load(cls, path): + data = torch.load(path, map_location="cpu", weights_only=False) + vf = cls(hidden_dim=data["hidden_dim"]) + vf.load_state_dict(data["state_dict"]) + buf = data["buffer_c"] + vf.buffer_c = [buf[i] for i in range(len(buf))] + vf.buffer_r = list(data["buffer_r"]) + vf.n_trained = data.get("n_trained", 0) + return vf + + @classmethod + def load_or_create(cls, path, hidden_dim=None): + if os.path.exists(path): + try: + return cls.load(path) + except Exception: + pass + return cls(hidden_dim=hidden_dim) if hidden_dim is not None else None diff --git a/web/funpack_studio.js b/web/funpack_studio.js index 90bf9790..6f75bbc8 100644 --- a/web/funpack_studio.js +++ b/web/funpack_studio.js @@ -69,7 +69,7 @@ function defaultSettings() { refinement_key: "", overrides: { refinement_key: false, feedback_prompt: false, user_intent_prompt: false, negative_prompt: false }, scene_builder: { mode: "Pass-through", scene: NONE_SENTINEL, scene_name: "", aliases: "", scene_positive: "", scene_negative: "" }, - refiner: { mode: "Refine", advisor_mode: "Off", advisor_thinking: true, prompt_repair: true, im_feeling_lucky: false, reset_session: false, feedback_prompt: "", user_intent_prompt_override: "", negative_prompt: "", temporal_style: "natural", split_by_transitions: false, split_transition_placement: "start", reference_injection: false, vision_conditioning: true }, + refiner: { mode: "Refine", advisor_mode: "Off", advisor_thinking: true, prompt_repair: true, im_feeling_lucky: false, reset_session: false, feedback_prompt: "", user_intent_prompt_override: "", negative_prompt: "", temporal_style: "natural", split_by_transitions: false, split_transition_placement: "start", reference_injection: false, vision_conditioning: true, value_guidance: false }, advisor_llm: { enabled: false, model_path: "huihui-ai/Huihui-Qwen3-8B-abliterated-v2", dtype: "bfloat16" }, loras: [], loras_config: { mode: "ltx2", per_block: false }, @@ -933,6 +933,12 @@ function openPanel(node) { body.append(el("div", "funpack-studio-hint", "Use the reference image as visual context when encoding the prompt with Gemma3. Disable to test without vision conditioning.")); body.append(row("Vision conditioning", visCondToggle.wrap)); + if (settings.refiner.value_guidance === undefined) settings.refiner.value_guidance = false; + const valueGuidanceToggle = toggleEl(!!settings.refiner.value_guidance, "Value guidance training"); + valueGuidanceToggle.inp.addEventListener("change", () => { settings.refiner.value_guidance = valueGuidanceToggle.inp.checked; }); + body.append(el("div", "funpack-studio-hint", "Train an online value function on rated generations. After 5+ ratings (good and bad), the chain sampler can use it for per-step gradient steering. Enable 'embed_guidance' in the sampler to use.")); + body.append(row("Value guidance", valueGuidanceToggle.wrap)); + body.append(sectionTitle("Negative prompt")); body.append(el("div", "funpack-studio-hint", "Encoded via CLIP and output as negative conditioning. Skipped when negative_conditioning input is connected."));