olivv-cs-ComfyUI-FunPack/value_function.py

"""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."""
        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
        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

    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_latent(latent):
    """[B, C, T, H, W] or [C, T, H, W] → [C] mean over spatial/temporal dims."""
    x = latent.float()
    if x.dim() == 5:
        x = x[0]  # [C, T, H, W]
    if x.dim() == 4:
        return x.mean(dim=[1, 2, 3])  # [C]
    return x