diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index 480d4e5..759b8ff 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/extensions_built_in/sd_trainer/SDTrainer.py @@ -403,11 +403,11 @@ class SDTrainer(BaseSDTrainProcess): sd=self.sd ) self.dfe.to(self.device_torch) - if hasattr(self.dfe, 'vision_encoder') and self.train_config.gradient_checkpointing: + if hasattr(self.dfe, 'vision_encoder'): # must be set to train for gradient checkpointing to work self.dfe.vision_encoder.train() self.dfe.vision_encoder.gradient_checkpointing = True - elif hasattr(self.dfe, 'model') and self.train_config.gradient_checkpointing: + elif hasattr(self.dfe, 'model'): if hasattr(self.dfe.model, 'enable_gradient_checkpointing'): self.dfe.model.train() self.dfe.model.enable_gradient_checkpointing() @@ -750,6 +750,8 @@ class SDTrainer(BaseSDTrainProcess): batch=batch, scheduler=self.sd.noise_scheduler ) + dfe_loss = dfe_loss.mean() + self.additional_logs['loss/dfe'] = dfe_loss.item() additional_loss += dfe_loss * self.train_config.diffusion_feature_extractor_weight else: raise ValueError(f"Unknown diffusion feature extractor version {self.dfe.version}") diff --git a/toolkit/models/diffusion_feature_extraction.py b/toolkit/models/diffusion_feature_extraction.py index d77c866..b1ba570 100644 --- a/toolkit/models/diffusion_feature_extraction.py +++ b/toolkit/models/diffusion_feature_extraction.py @@ -168,9 +168,10 @@ class DiffusionFeatureExtractor(nn.Module): class DiffusionFeatureExtractor3(nn.Module): - def __init__(self, device=torch.device("cuda"), dtype=torch.bfloat16, vae=None): + def __init__(self, device=torch.device("cuda"), dtype=torch.bfloat16, vae=None, sd: BaseModel = None): super().__init__() self.version = 3 + self.sd_ref = weakref.ref(sd) if sd is not None else None if vae is None: vae = AutoencoderTiny.from_pretrained( "madebyollin/taef1", torch_dtype=torch.bfloat16) @@ -272,6 +273,8 @@ class DiffusionFeatureExtractor3(nn.Module): if model is not None and hasattr(model, 'get_stepped_pred'): stepped_latents = model.get_stepped_pred(noise_pred, noise) + elif self.sd_ref is not None and getattr(self.sd_ref(), "x0_pred", False): + stepped_latents = noise_pred else: # stepped_latents = noise - noise_pred # first we step the scheduler from current timestep to the very end for a full denoise @@ -369,9 +372,10 @@ class DiffusionFeatureExtractor3(nn.Module): return total_loss class DiffusionFeatureExtractor4(nn.Module): - def __init__(self, device=torch.device("cuda"), dtype=torch.bfloat16, vae=None): + def __init__(self, device=torch.device("cuda"), dtype=torch.bfloat16, vae=None, sd: BaseModel = None): super().__init__() self.version = 4 + self.sd_ref = weakref.ref(sd) if sd is not None else None if vae is None: raise ValueError("vae must be provided for DFE4") self.vae = vae @@ -548,8 +552,11 @@ class DiffusionFeatureExtractor4(nn.Module): # min 0.001 tv = torch.clamp(tv, min=0.001) - # step latent - x0 = noisy_latents - tv * noise_pred + if self.sd_ref is not None and getattr(self.sd_ref(), "x0_pred", False): + x0 = noise_pred + else: + # step latent + x0 = noisy_latents - tv * noise_pred stepped_latents = x0 @@ -616,8 +623,8 @@ class DiffusionFeatureExtractor4(nn.Module): return total_loss class DiffusionFeatureExtractor5(DiffusionFeatureExtractor4): - def __init__(self, device=torch.device("cuda"), dtype=torch.bfloat16, vae=None): - super().__init__(device=device, dtype=dtype, vae=vae) + def __init__(self, device=torch.device("cuda"), dtype=torch.bfloat16, vae=None, sd: BaseModel = None): + super().__init__(device=device, dtype=dtype, vae=vae, sd=sd) self.version = 5 def step_latents(self, noise, noise_pred, noisy_latents, timesteps, scheduler, total_steps: int = 1000, eps: float = 1e-6): @@ -663,9 +670,10 @@ class DiffusionFeatureExtractor5(DiffusionFeatureExtractor4): class DiffusionFeatureExtractor6(nn.Module): - def __init__(self, device=torch.device("cuda"), dtype=torch.bfloat16, vae=None): + def __init__(self, device=torch.device("cuda"), dtype=torch.bfloat16, vae=None, sd: BaseModel = None): super().__init__() self.version = 6 + self.sd_ref = weakref.ref(sd) if sd is not None else None if vae is None: raise ValueError("vae must be provided for DFE4") self.vae = vae @@ -689,6 +697,15 @@ class DiffusionFeatureExtractor6(nn.Module): 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) + # last_hidden_state token order: [CLS, registers..., patches] + self.num_prefix_tokens = 1 + getattr(self.model.config, "num_register_tokens", 0) + # CLS is a global anchor only; patch tokens carry the spatial signal + self.cls_weight = 0.1 + + def _dino_features(self, inputs): + hidden = self.model(**inputs).last_hidden_state + return hidden[:, 0], hidden[:, self.num_prefix_tokens:] + def prepare_inputs(self, tensor_0_1: torch.Tensor): """ tensor_0_1: (bs, 3, h, w), float, values in [0, 1] @@ -764,8 +781,11 @@ class DiffusionFeatureExtractor6(nn.Module): # min 0.001 tv = torch.clamp(tv, min=0.001) - # step latent - x0 = noisy_latents - tv * noise_pred + if self.sd_ref is not None and getattr(self.sd_ref(), "x0_pred", False): + x0 = noise_pred + else: + # step latent + x0 = noisy_latents - tv * noise_pred stepped_latents = x0 @@ -791,23 +811,29 @@ 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).pooler_output.detach() - # normalize - target_dino_output = (target_dino_output - target_dino_output.mean()) / (target_dino_output.std() + 1e-6) + target_cls, target_patches = self._dino_features(target_dino_input) + target_cls = target_cls.detach() + target_patches = target_patches.detach() pred_dino_input = self.prepare_inputs(pred_images) - 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( - pred_dino_output.float(), target_dino_output.float() - ) + pred_cls, pred_patches = self._dino_features(pred_dino_input) + + # per-token cosine so no single bright region dominates + patch_loss = (1.0 - F.cosine_similarity(pred_patches.float(), target_patches.float(), dim=-1)).mean(dim=1) + cls_loss = 1.0 - F.cosine_similarity(pred_cls.float(), target_cls.float(), dim=-1) + # (B*T,) -> (B,) so video frames fold back into their source sample + bs = noise_pred.shape[0] + patch_loss = patch_loss.view(bs, -1).mean(dim=1) + cls_loss = cls_loss.view(bs, -1).mean(dim=1) + # per-sample (B,); trainer applies the model's per-sample loss scaling and reduces + dino_loss = patch_loss + self.cls_weight * cls_loss # 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.detach() - else: - self.losses['dinov3'] = self.losses['dinov3'] + dino_loss.detach() + for key, val in (('dinov3_patch', patch_loss.mean()), ('dinov3_cls', cls_loss.mean())): + if key not in self.losses: + self.losses[key] = val.detach() + else: + self.losses[key] = self.losses[key] + val.detach() with torch.no_grad(): if self.step % self.log_every == 0 and self.step > 0: @@ -1376,19 +1402,19 @@ class DiffusionFeatureExtractor10(nn.Module): def load_dfe(model_path, vae=None, sd: 'BaseModel' = None) -> DiffusionFeatureExtractor: if model_path == "v3": - dfe = DiffusionFeatureExtractor3(vae=vae) + dfe = DiffusionFeatureExtractor3(vae=vae, sd=sd) dfe.eval() return dfe if model_path == "v4": - dfe = DiffusionFeatureExtractor4(vae=vae) + dfe = DiffusionFeatureExtractor4(vae=vae, sd=sd) dfe.eval() return dfe if model_path == "v5": - dfe = DiffusionFeatureExtractor5(vae=vae) + dfe = DiffusionFeatureExtractor5(vae=vae, sd=sd) dfe.eval() return dfe if model_path == "v6": - dfe = DiffusionFeatureExtractor6(vae=vae) + dfe = DiffusionFeatureExtractor6(vae=vae, sd=sd) dfe.eval() return dfe if model_path == "v7":