digital-garbageandClaude Sonnet 5
d91388922e
New LatentValueFunction (value_function.py) mirrors OnlineValueFunction but is trained on a pooled snapshot of the sampler's own predicted output (x0_hat) instead of the input conditioning, borrowing NoiseTilt's core idea (reward gradient from the model's own prediction) without its SDE noise-term trick, which has no analogue in this deterministic-ODE sampler. New output_guidance / output_guidance_strength toggle on the Chain Sampler (off by default): per-step correction installed as the outermost model_function_wrapper (on top of embed_guidance/score_slider), same near-zero-forward-pass shape as embed_guidance. End of run writes a small pooled snapshot per refinement key; the next rating cycle trains the value function on it via a new call site in conditioning.py. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
220 lines
9.3 KiB
Python
220 lines
9.3 KiB
Python
"""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 = 10
|
|
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."""
|
|
with torch.inference_mode(False), torch.enable_grad():
|
|
c = self.compress(conditioning.clone()).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."""
|
|
mlp_device = next(self.parameters()).device
|
|
orig_device = conditioning.device
|
|
with torch.inference_mode(False), torch.enable_grad():
|
|
# new_empty inside inference_mode(False) creates a normal (non-inference) tensor.
|
|
# copy_ fills it from the inference tensor without inheriting the flag.
|
|
c_fresh = torch.empty(conditioning.shape, dtype=torch.float32, device=mlp_device)
|
|
c_fresh.copy_(conditioning)
|
|
c_fresh.requires_grad_(True)
|
|
reward = self.forward(self.compress(c_fresh).unsqueeze(0))
|
|
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
|
|
# Cap total displacement at 8% of original norm — prevents stripping action/content
|
|
max_displacement = conditioning.float().norm().item() * 0.08
|
|
prev_reward = -float("inf")
|
|
with torch.inference_mode(False), torch.enable_grad():
|
|
c_orig = torch.empty(conditioning.shape, dtype=torch.float32, device=mlp_device)
|
|
c_orig.copy_(conditioning)
|
|
c = c_orig.clone()
|
|
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
|
|
# Project back if displacement exceeds cap
|
|
delta = c - c_orig
|
|
disp = delta.norm()
|
|
if disp > max_displacement:
|
|
c = c_orig + delta * (max_displacement / disp)
|
|
if reward_val - prev_reward < 1e-3:
|
|
break
|
|
prev_reward = reward_val
|
|
return c.detach().to(orig_device, orig_dtype)
|
|
|
|
def search(self, conditioning, n=16, eps_scale=1.0):
|
|
"""Monte Carlo conditioning search: score N random perturbations, return the best.
|
|
|
|
eps_scale widens (>1) or tightens (<1) the perturbation radius. The path planner uses
|
|
a wider radius to push a disliked config region toward a different conditioning, and the
|
|
caller skips this search entirely to LOCK a liked region's conditioning in place."""
|
|
if not self.is_ready():
|
|
return conditioning
|
|
mlp_device = next(self.parameters()).device
|
|
orig_device = conditioning.device
|
|
orig_dtype = conditioning.dtype
|
|
eps = conditioning.float().norm().item() * 0.025 * float(eps_scale)
|
|
with torch.inference_mode(False), torch.no_grad():
|
|
c_base = torch.empty(conditioning.shape, dtype=torch.float32, device=mlp_device)
|
|
c_base.copy_(conditioning)
|
|
best = c_base
|
|
best_score = float(self.forward(self.compress(c_base).unsqueeze(0)).item())
|
|
for _ in range(n):
|
|
noise = F.normalize(torch.randn_like(c_base).flatten(), dim=0).reshape(c_base.shape)
|
|
candidate = c_base + eps * noise
|
|
score = float(self.forward(self.compress(candidate).unsqueeze(0)).item())
|
|
if score > best_score:
|
|
best_score = score
|
|
best = candidate
|
|
return best.detach().to(orig_device, orig_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
|
|
|
|
|
|
def compress_packed_latent(x, hidden_dim):
|
|
"""Pool a packed AV latent [B,1,N] (LTXAV's flattened video/audio stream layout — see
|
|
samplers._get_latent_shapes) down to a fixed-size [hidden_dim] vector, differentiably.
|
|
|
|
N varies per generation (scene length/resolution), unlike text conditioning's fixed
|
|
encoder width, so adaptive pooling (not a plain mean) is what makes a single value
|
|
function usable across differently-shaped runs. Idempotent on already-compressed input
|
|
(shape [hidden_dim]) so the same call works whether fed a live raw latent (inference,
|
|
where autograd needs to flow through this op to produce a full-shape gradient) or a
|
|
pre-compressed snapshot loaded back from disk (training, where the raw latent was never
|
|
persisted — only this vector was, to keep the sidecar file small)."""
|
|
if x.dim() == 1 and x.shape[0] == hidden_dim:
|
|
return x
|
|
c = x.float()
|
|
if c.dim() == 1:
|
|
c = c.view(1, 1, -1)
|
|
elif c.dim() == 2:
|
|
c = c.unsqueeze(1)
|
|
elif c.dim() > 3:
|
|
c = c.reshape(c.shape[0], 1, -1)
|
|
pooled = F.adaptive_avg_pool1d(c, hidden_dim) # [B,1,hidden_dim]
|
|
return pooled.mean(dim=0).reshape(-1) # average over batch (usually B=1, a no-op) -> [hidden_dim]
|
|
|
|
|
|
class LatentValueFunction(OnlineValueFunction):
|
|
"""OnlineValueFunction trained on the SAMPLER'S OWN predicted output (a packed-latent
|
|
x0_hat estimate) instead of the input text conditioning — same tiny-MLP mechanism, just
|
|
fed a different, more direct signal. See [[research_2026_training_free_candidates]] /
|
|
NoiseTilt for the idea this borrows (reward gradient computed from the model's own
|
|
prediction) and [[project_embed_guidance]] for the conditioning-space sibling this
|
|
complements. Cost stays the same as embed_guidance: one backward pass through a few
|
|
hundred params, no extra diffusion-model forward pass."""
|
|
|
|
DEFAULT_HIDDEN_DIM = 512
|
|
|
|
def compress(self, x):
|
|
return compress_packed_latent(x, self.hidden_dim)
|
|
|
|
|