diff --git a/samplers.py b/samplers.py index b00afa0d..22e14df9 100644 --- a/samplers.py +++ b/samplers.py @@ -1467,7 +1467,11 @@ class FunPackLTXAVSceneChainSampler: print("[FunPackSceneChain] embed_guidance: no liked direction found (need 3+ liked generations)") else: _value_fn = self._load_value_function(refinement_key_input) - mode = f"value function ({_value_fn.n_trained} samples)" if _value_fn else "fixed direction" + if _value_fn: + ready = _value_fn.is_ready() + mode = f"value function ({_value_fn.n_trained} samples, ascent {'on' if ready else 'pending'})" + else: + mode = "fixed direction" print(f"[FunPackSceneChain] embed_guidance: active via {mode}, strength={embed_guidance_strength}") first_scene_seed = self._scene_seed(scene_conditionings[0]) @@ -1513,6 +1517,10 @@ class FunPackLTXAVSceneChainSampler: else: guide_tail = 0 + if embed_guidance and _value_fn is not None and _value_fn.is_ready(): + orig_cond, orig_extra = scene_positive[0][0], scene_positive[0][1] + ascended = _value_fn.ascend(orig_cond) + scene_positive = [[ascended, orig_extra]] + list(scene_positive[1:]) if embed_guidance and _liked_dir is not None: _eg_old_wrapper = self._build_embed_guidance_wrapper(model, _liked_dir, embed_guidance_strength, value_fn=_value_fn) sampled = self._sample_chunk( diff --git a/value_function.py b/value_function.py index 8609f230..7f0683d4 100644 --- a/value_function.py +++ b/value_function.py @@ -85,6 +85,30 @@ class OnlineValueFunction(nn.Module): reward.backward() return c_fresh.grad.to(orig_device, conditioning.dtype) + def ascend(self, conditioning): + """Gradient ascent on conditioning until reward plateaus. Self-terminating, no user params.""" + if not self.is_ready(): + return conditioning + mlp_device = next(self.parameters()).device + orig_device = conditioning.device + orig_dtype = conditioning.dtype + step_size = conditioning.float().norm().item() * 0.005 + prev_reward = -float("inf") + with torch.inference_mode(False), torch.enable_grad(): + c = torch.empty(conditioning.shape, dtype=torch.float32, device=mlp_device) + c.copy_(conditioning) + for _ in range(50): + c = c.detach().requires_grad_(True) + reward = self.forward(self.compress(c).unsqueeze(0)) + reward.backward() + reward_val = reward.item() + grad = F.normalize(c.grad.float(), dim=-1) + c = c.detach() + step_size * grad + if reward_val - prev_reward < 1e-3: + break + prev_reward = reward_val + return c.detach().to(orig_device, orig_dtype) + def is_ready(self): return len(self.buffer_c) >= self.MIN_SAMPLES