From 4dd8fcd6741bf5d06a36c5e8e5a32a9b49db4fb9 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Fri, 29 May 2026 04:35:50 +0300 Subject: [PATCH] Fix value function gradient: move to MLP device before backward pass MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MLP lives on CPU, conditioning on GPU — device mismatch silently fell back to fixed direction. Now moves conditioning to MLP device for the forward/ backward, then returns gradient to original device. Also surface value function fallback errors instead of swallowing them. Co-Authored-By: Claude Sonnet 4.6 --- samplers.py | 3 ++- value_function.py | 5 +++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/samplers.py b/samplers.py index 127174f3..3d9f4a3d 100644 --- a/samplers.py +++ b/samplers.py @@ -1348,7 +1348,8 @@ class FunPackLTXAVSceneChainSampler: try: grad = _vf.gradient(cond) d = torch.nn.functional.normalize(grad.float(), dim=-1).to(cond.dtype) - except Exception: + except Exception as _e: + print(f"[FunPackSceneChain] embed_guidance: value function gradient failed ({_e}), using fixed direction") d = _fixed.to(cond.device, cond.dtype).expand_as(cond) else: d = _fixed.to(cond.device, cond.dtype).expand_as(cond) diff --git a/value_function.py b/value_function.py index 606cc60a..f53c92cc 100644 --- a/value_function.py +++ b/value_function.py @@ -73,11 +73,12 @@ class OnlineValueFunction(nn.Module): def gradient(self, conditioning): """∂reward/∂conditioning — same shape as input.""" + mlp_device = next(self.parameters()).device with torch.inference_mode(False), torch.enable_grad(): - c_in = conditioning.detach().clone().float().requires_grad_(True) + c_in = conditioning.detach().clone().float().to(mlp_device).requires_grad_(True) reward = self.forward(self.compress(c_in).unsqueeze(0)) reward.backward() - return c_in.grad.to(conditioning.dtype) + return c_in.grad.to(conditioning.device, conditioning.dtype) def is_ready(self): return len(self.buffer_c) >= self.MIN_SAMPLES