[ADD] better names in loss weight

Author: nviolante25

sgm/modules/diffusionmodules/loss_weighting.py

@@ -34,9 +34,9 @@ def __call__(self, sigma: torch.Tensor) -> torch.Tensor:
 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)
+        batch_size, N = bools.shape
+        indices = torch.arange(N, device=bools.device).unsqueeze(0).expand(batch_size, N)
+        weights = torch.full((batch_size, N), max_weight, dtype=torch.float, device=bools.device)
 
         for b in range(batch_size):
             true_idx = indices[b][bools[b]]

sgm/modules/diffusionmodules/wrappers.py

@@ -57,15 +57,12 @@ def forward(
         x = rearrange(x, "b f c h w -> (b f) c h w")
         dense_y=rearrange(c["plucker"], "b f c h w -> (b f) c h w")
 
-        #TODO: remove
-        c = torch.zeros((b, 1, 1024)).type_as(x).to(x.device)
-        c = repeat(c, "b 1 c -> (b f) 1 c", f=f)
         t = repeat(t, "b -> (b f)", f=f)
 
         out = self.diffusion_model(
             x,
             t=t,
-            y=c, # c["crossattn"]
+            y=c["crossattn"],
             dense_y=dense_y,
             num_frames=f,
             **kwargs,