Adjust the x0 of some DFEs

toolkit/models/diffusion_feature_extraction.py

@@ -537,8 +537,19 @@ class DiffusionFeatureExtractor4(nn.Module):
         if model is not None and hasattr(model, 'get_stepped_pred'):
             stepped_latents = model.get_stepped_pred(noise_pred, noise)
         else:
-            stepped_latents = self.step_latents(noise, noise_pred, noisy_latents, timesteps, scheduler)
-            
+            with torch.no_grad():
+                tv = timesteps.to(noise_pred.device).to(noise_pred.dtype) / 1000.0
+                # expand shape to match noise_pred
+                while len(tv.shape) < len(noise_pred.shape):
+                    tv = tv.unsqueeze(-1)
+                    # min 0.001
+                    tv = torch.clamp(tv, min=0.001)
+
+            # step latent
+            x0 = noisy_latents - tv * noise_pred
+
+            stepped_latents = x0
+
         latents = stepped_latents.to(self.vae.device, dtype=self.vae.dtype)
 
         scaling_factor = self.vae.config.scaling_factor if hasattr(self.vae.config, 'scaling_factor') else 1.0
@@ -595,7 +606,7 @@ class DiffusionFeatureExtractor4(nn.Module):
                 self.losses[key] /= self.log_every
                 # print in 2.000e-01 format
                 print(f" - {key}: {self.losses[key]:.3e}")
-            self.losses[key] = 0.0
+                self.losses[key] = 0.0
         
         # total_loss += mse_loss
         self.step += 1
@@ -684,6 +695,10 @@ class DiffusionFeatureExtractor6(nn.Module):
         if not torch.is_floating_point(x):
             x = x.float()
 
+        # VAE decode can overshoot [0, 1] slightly; clamp so the rescale
+        # heuristic below never mistakes an overshoot for 0..255 input
+        x = torch.clamp(x, 0.0, 1.0)
+
         # Resize
         # if not divisible by 16 or total pixels > max_res*max_res, resize to fit within 16 patches
         max_res = 512
@@ -794,7 +809,7 @@ class DiffusionFeatureExtractor6(nn.Module):
             self.losses['dinov3'] = dino_loss.item()
         else:
             self.losses['dinov3'] += dino_loss.item()
-        
+
         with torch.no_grad():
             if self.step % self.log_every == 0 and self.step > 0:
                 print(f"DFE losses:")
@@ -802,7 +817,7 @@ class DiffusionFeatureExtractor6(nn.Module):
                     self.losses[key] /= self.log_every
                     # print in 2.000e-01 format
                     print(f" - {key}: {self.losses[key]:.3e}")
-                self.losses[key] = 0.0
+                    self.losses[key] = 0.0
             
             # total_loss += mse_loss
             self.step += 1