methods.py

from __future__ import annotations

import time
from typing import Callable

import torch

from models import FLUX
from wgnc import project_wgnc

Tensor = torch.Tensor


def optimize_gradient_preconditioning(
    *,
    flux: FLUX,
    reward_model: Callable[[Tensor], Tensor],
    prompt: str,
    seed: int,
    num_iters: int,
    lr: float,
    grad_clip: float,
    guidance: float,
    project_latent: bool,
) -> dict[str, object]:
    """ICML 2026 gradient preconditioning on the initial FLUX noise."""
    condition = flux.prepare_prompt(prompt)
    generator = torch.Generator(device=flux.device).manual_seed(seed)
    x = torch.randn(condition.latent_shape, device=flux.device, dtype=torch.float32, generator=generator)
    if project_latent:
        x = project_wgnc(x)
    x = torch.nn.Parameter(x)
    opt = torch.optim.Adam([x], lr=lr)
    best_reward = -float("inf")
    best_image: Tensor | None = None
    start = time.perf_counter()

    for _ in range(num_iters):
        signal = flux.reward_gradient(x, reward_model, condition=condition, guidance=guidance)
        reward = signal.reward.mean()
        gradient = signal.gradient.detach().to(torch.float32)
        grad_rms = gradient.square().mean().sqrt()
        if torch.isfinite(grad_rms) and grad_rms > 0:
            gradient = project_wgnc(gradient / grad_rms.clamp_min(1e-12))
            gradient = gradient * min(float(grad_rms.item()), grad_clip)
        else:
            gradient = torch.zeros_like(gradient)

        opt.zero_grad(set_to_none=True)
        x.grad = -gradient
        opt.step()

        with torch.no_grad():
            if project_latent:
                x.copy_(project_wgnc(x.detach()))
            if float(reward.item()) > best_reward:
                best_reward = float(reward.item())
                best_image = signal.image.detach()

    with torch.no_grad():
        final_image = flux.decode(x.detach())
        final_reward = float(reward_model(final_image).mean().item())
    if final_reward > best_reward:
        best_reward = final_reward
        best_image = final_image.detach()
    if best_image is None:
        raise RuntimeError("gradient preconditioning did not produce a finite-reward image")

    return {
        "image": best_image,
        "target_reward": best_reward,
        "seconds": time.perf_counter() - start,
    }