diff --git a/toolkit/models/diffusion_feature_extraction.py b/toolkit/models/diffusion_feature_extraction.py index cc53225e..1aae0319 100644 --- a/toolkit/models/diffusion_feature_extraction.py +++ b/toolkit/models/diffusion_feature_extraction.py @@ -689,7 +689,12 @@ class DiffusionFeatureExtractor6(nn.Module): self.losses = {} self.log_every = 100 self.step = 0 - + + # cache normalization constants once so prepare_inputs doesn't rebuild + # them (CPU tensor construction + H2D copy) on every call + self.image_mean = torch.tensor(self.processor.image_mean, device=device, dtype=dtype).view(1, 3, 1, 1) + self.image_std = torch.tensor(self.processor.image_std, device=device, dtype=dtype).view(1, 3, 1, 1) + def prepare_inputs(self, tensor_0_1: torch.Tensor): """ tensor_0_1: (bs, 3, h, w), float, values in [0, 1] @@ -703,8 +708,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 + # VAE decode can overshoot [0, 1] slightly; clamp back into range. + # Inputs are documented as 0..1 and clamped here, so no 0..255 rescale + # heuristic is needed (the old max().item() check forced a GPU sync + # every call and could never trigger after this clamp anyway). x = torch.clamp(x, 0.0, 1.0) # Resize @@ -722,17 +729,10 @@ class DiffusionFeatureExtractor6(nn.Module): target_w = (target_w // p) * p x = F.interpolate(x, size=(target_h, target_w), mode="bilinear", align_corners=False) - # Rescale (HF processors usually assume uint8 0..255 inputs; your inputs are already 0..1) - if self.processor.do_rescale: - # If it looks like [0..1], skip to avoid double-scaling. - # If user accidentally passed 0..255 floats, this will fix it. - if x.detach().max().item() > 1.0 + 1e-6: - x = x * float(self.processor.rescale_factor or 1.0 / 255.0) - # Normalize if self.processor.do_normalize: - mean = torch.tensor(self.processor.image_mean, device=x.device, dtype=x.dtype).view(1, 3, 1, 1) - std = torch.tensor(self.processor.image_std, device=x.device, dtype=x.dtype).view(1, 3, 1, 1) + mean = self.image_mean.to(device=x.device, dtype=x.dtype) + std = self.image_std.to(device=x.device, dtype=x.dtype) x = (x - mean) / std return {"pixel_values": x} @@ -814,18 +814,22 @@ class DiffusionFeatureExtractor6(nn.Module): pred_dino_output.float(), target_dino_output.float() ) + # accumulate on-GPU; .item() every step forces a full pipeline sync, + # so only sync when we actually log if 'dinov3' not in self.losses: - self.losses['dinov3'] = dino_loss.item() + self.losses['dinov3'] = dino_loss.detach() else: - self.losses['dinov3'] += dino_loss.item() + self.losses['dinov3'] = self.losses['dinov3'] + dino_loss.detach() with torch.no_grad(): if self.step % self.log_every == 0 and self.step > 0: print(f"DFE losses:") for key in self.losses: - self.losses[key] /= self.log_every + avg = self.losses[key] / self.log_every + if torch.is_tensor(avg): + avg = avg.item() # print in 2.000e-01 format - print(f" - {key}: {self.losses[key]:.3e}") + print(f" - {key}: {avg:.3e}") self.losses[key] = 0.0 # total_loss += mse_loss