From 97e101522c55834309111d8131908230d3acf303 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Tue, 10 Jun 2025 08:01:13 -0600 Subject: [PATCH] Increase ema feedback amount. Normalize the dfe 4 image embeds --- toolkit/ema.py | 3 ++- toolkit/models/diffusion_feature_extraction.py | 5 +++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/toolkit/ema.py b/toolkit/ema.py index e3b3a7ea..b34554bb 100644 --- a/toolkit/ema.py +++ b/toolkit/ema.py @@ -137,7 +137,8 @@ class ExponentialMovingAverage: update_param = False if self.use_feedback: - param_float.add_(tmp) + # make feedback 10x decay + param_float.add_(tmp * 10) update_param = True if self.param_multiplier != 1.0: diff --git a/toolkit/models/diffusion_feature_extraction.py b/toolkit/models/diffusion_feature_extraction.py index e911484a..01d7f278 100644 --- a/toolkit/models/diffusion_feature_extraction.py +++ b/toolkit/models/diffusion_feature_extraction.py @@ -469,8 +469,9 @@ class DiffusionFeatureExtractor4(nn.Module): ) # embeds = id_embeds['hidden_states'][-2] # penultimate layer - embeds = id_embeds['pooler_output'] - return embeds + image_embeds = id_embeds['pooler_output'] + image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True) + return image_embeds def forward( self,