From f85bf065bf4a427ec0a0a76013174b451a0ea29e Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Fri, 27 Mar 2026 07:02:07 -0600
Subject: [PATCH] Use pooler embeddings for DFE v6 with dino v3

---
 toolkit/models/diffusion_feature_extraction.py | 4 ++--
 version.py                                     | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/toolkit/models/diffusion_feature_extraction.py b/toolkit/models/diffusion_feature_extraction.py
index c1effd80..d50867c1 100644
--- a/toolkit/models/diffusion_feature_extraction.py
+++ b/toolkit/models/diffusion_feature_extraction.py
@@ -775,11 +775,11 @@ class DiffusionFeatureExtractor6(nn.Module):
             # go from -1 to 1 to 0 to 1
             target_img = (target_img + 1) / 2
             target_dino_input = self.prepare_inputs(target_img)
-            target_dino_output = self.model(**target_dino_input, output_hidden_states=True)['hidden_states'][-1].detach()
+            target_dino_output = self.model(**target_dino_input).pooler_output.detach()
             # normalize
             target_dino_output = (target_dino_output - target_dino_output.mean()) / (target_dino_output.std() + 1e-6)
         pred_dino_input = self.prepare_inputs(pred_images)
-        pred_dino_output = self.model(**pred_dino_input, output_hidden_states=True)['hidden_states'][-1]
+        pred_dino_output = self.model(**pred_dino_input).pooler_output
         # normalize
         pred_dino_output = (pred_dino_output - pred_dino_output.mean()) / (pred_dino_output.std() + 1e-6)
         dino_loss = torch.nn.functional.mse_loss(
diff --git a/version.py b/version.py
index c5482cfc..256107c3 100644
--- a/version.py
+++ b/version.py
@@ -1 +1 @@
-VERSION = "0.7.31"
+VERSION = "0.7.32"