Update gligen/ldm/models/diffusion/loss.py

gligen/ldm/models/diffusion/loss.py

@@ -554,27 +554,27 @@ def caculate_loss_LoCo_V2(attn_maps_mid, attn_maps_up, attn_maps_down, bboxes, o
 
             loss += (obj_loss) # /len(object_positions[obj_idx])
 
-    #     sot_map = attn_map[:, :, 0].reshape(H, W)
-    #     eot_map = attn_map[:, :, -1].reshape(H, W)
+        sot_map = attn_map[:, :, 0].reshape(H, W)
+        eot_map = attn_map[:, :, -1].reshape(H, W)
 
-    #     norm_sot_map = (1 - sot_map) / (1 - sot_map).max()
-    #     norm_eot_map = eot_map / eot_map.max()
+        norm_sot_map = (1 - sot_map) / (1 - sot_map).max()
+        norm_eot_map = eot_map / eot_map.max()
 
 
-    #     pad_map = beta * norm_sot_map + (1 - beta) * norm_eot_map
+        pad_map = beta * norm_sot_map + (1 - beta) * norm_eot_map
 
-    # total_fg_mask = total_fg_map
-    # fg_map = pad_map * total_fg_mask
+    total_fg_mask = total_fg_map
+    fg_map = pad_map * total_fg_mask
 
-    # bce_loss = F.binary_cross_entropy_with_logits(pad_map.to(torch.float16).reshape(-1), fg_map.to(torch.float16).reshape(-1))
+    bce_loss = F.binary_cross_entropy_with_logits(pad_map.to(torch.float16).reshape(-1), fg_map.to(torch.float16).reshape(-1))
 
-    # pad_loss += bce_loss
+    pad_loss += bce_loss
     if sum_in + sum_out == 0:
         return torch.tensor(0).float().cuda() if torch.cuda.is_available() else torch.tensor(0).float()
     # loss += (1 - sum_in / (sum_in + sum_out)) ** 2
     # print('loss', loss)
     # return loss
-    return loss #+ alpha * pad_loss
+    return loss + alpha * pad_loss