Add pre-generation conditioning ascent to value function guidance
Before each denoising run, when the value function is ready, gradient-ascend the text conditioning until predicted reward plateaus (self-terminating: step size = 0.5% of conditioning norm, stops when delta < 1e-3, caps at 50 steps). No user-facing params — fully automated. Per-step wrapper still runs on top. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Sonnet 4.6
parent
e999a96d9b
commit
2f28d7eb39
+9
-1
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user