From 4ca896c73a3d553f14decb54885d1989c64a4fab Mon Sep 17 00:00:00 2001
From: digital-garbage <133010426+digital-garbage@users.noreply.github.com>
Date: Sat, 30 May 2026 15:47:28 +0300
Subject: [PATCH] Add Monte Carlo conditioning search after gradient ascent
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

search() generates 16 random unit-direction perturbations (2.5% of
conditioning norm each), scores each with the VF, returns the best.
Runs after ascend() in _v2_ascend_conditioning — ascent exploits the
known gradient direction, search explores the neighbourhood around it
for any better nearby region. Total cost: 16 VF forward passes ≈ free.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 conditioning.py   |  8 ++++++--
 value_function.py | 22 ++++++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/conditioning.py b/conditioning.py
index 98f1c85c..e5560896 100644
--- a/conditioning.py
+++ b/conditioning.py
@@ -12142,8 +12142,12 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
             ascended = []
             for entry in conditioning_list:
                 tensor, extra = entry[0], entry[1]
-                ascended.append((vf.ascend(tensor), extra))
-            print(f"[FunPackRefiner] Conditioning ascent applied ({vf.n_trained} samples, {len(ascended)} scene(s))")
+                # Gradient ascent first (exploit known gradient direction)
+                c = vf.ascend(tensor)
+                # Monte Carlo search around ascended point (explore neighbourhood)
+                c = vf.search(c)
+                ascended.append((c, extra))
+            print(f"[FunPackRefiner] Conditioning ascent+search applied ({vf.n_trained} samples, {len(ascended)} scene(s))")
             return ascended
         except Exception as e:
             print(f"[FunPackRefiner] Conditioning ascent failed: {e}")
diff --git a/value_function.py b/value_function.py
index 4c88d817..2da426d5 100644
--- a/value_function.py
+++ b/value_function.py
@@ -117,6 +117,28 @@ class OnlineValueFunction(nn.Module):
                 prev_reward = reward_val
         return c.detach().to(orig_device, orig_dtype)
 
+    def search(self, conditioning, n=16):
+        """Monte Carlo conditioning search: score N random perturbations, return the best."""
+        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
+        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