[ADD] loss weighting

Author: nviolante25

configs/example_training/seva-clipl_dl3dv.yaml

@@ -111,7 +111,7 @@ model:
       target: sgm.modules.diffusionmodules.loss.StandardDiffusionLoss
       params:
         loss_weighting_config:
-          target: sgm.modules.diffusionmodules.loss_weighting.EpsWeighting
+          target: sgm.modules.diffusionmodules.loss_weighting.SevaWeighting
         sigma_sampler_config:
           target: sgm.modules.diffusionmodules.sigma_sampling.DiscreteSampling
           params:

sgm/modules/diffusionmodules/loss.py

@@ -86,7 +86,7 @@ def _forward(
         model_output = denoiser(
             network, noised_input, sigmas, cond, **additional_model_inputs
         )
-        w = append_dims(self.loss_weighting(sigmas), input.ndim)
+        w = append_dims(self.loss_weighting(sigmas, cond["mask"]), input.ndim)
         return self.get_loss(model_output, input, w)
 
     def get_loss(self, model_output, target, w):
@@ -103,3 +103,19 @@ def get_loss(self, model_output, target, w):
             return loss
         else:
             raise NotImplementedError(f"Unknown loss type {self.loss_type}")
+
+def interpolate_weights_batch(bools: torch.Tensor, max_weight=5.0) -> torch.Tensor:
+    B, N = bools.shape
+    indices = torch.arange(N, device=bools.device).unsqueeze(0).expand(B, N)
+    weights = torch.full((B, N), max_weight, dtype=torch.float, device=bools.device)
+    
+    for b in range(B):
+        true_idx = indices[b][bools[b]]
+        if len(true_idx) > 0:
+            dists = torch.stack([torch.abs(indices[b] - t) for t in true_idx]).min(dim=0).values
+            dists[bools[b]] = 0
+            weights[b] = dists / dists.max() * max_weight
+        else:
+            weights[b] = max_weight
+
+    return weights

sgm/modules/diffusionmodules/loss_weighting.py

@@ -30,3 +30,22 @@ def __init__(self):
 class EpsWeighting(DiffusionLossWeighting):
     def __call__(self, sigma: torch.Tensor) -> torch.Tensor:
         return sigma**-2.0
+    
+class SevaWeighting(DiffusionLossWeighting):
+    def __call__(self, sigma: torch.Tensor, mask, max_weight=5.0) -> torch.Tensor:
+        bools = mask.to(torch.bool)
+        batch_size, num_frames = bools.shape
+        indices = torch.arange(num_frames, device=bools.device).unsqueeze(0).expand(batch_size, num_frames)
+        weights = torch.full((batch_size, num_frames), max_weight, dtype=torch.float, device=bools.device)
+        
+        for b in range(batch_size):
+            true_idx = indices[b][bools[b]]
+            if len(true_idx) > 0:
+                dists = torch.stack([torch.abs(indices[b] - t) for t in true_idx]).min(dim=0).values
+                dists[bools[b]] = 0
+                weights[b] = dists / dists.max() * max_weight
+            else:
+                weights[b] = max_weight
+
+        return weights
+