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

gligen/ldm/models/diffusion/loss.py

@@ -354,7 +354,7 @@ def caculate_loss_LoCo(attn_maps_mid, attn_maps_up, attn_maps_down, bboxes, obje
     # fg_map = torch.sigmoid(fg_map)
 
     # mse_loss = F.mse_loss(pad_map.reshape(-1), fg_map.reshape(-1))
-    bce_loss = F.binary_cross_entropy(torch.sigmoid(pad_map.reshape(-1)), fg_map.reshape(-1))
+    bce_loss = F.binary_cross_entropy_with_logits(pad_map.reshape(-1), fg_map.reshape(-1))
     # print('mse_loss', mse_loss)
     # print('bce_loss', bce_loss)
     #bce_loss = torch.clamp(bce_loss, max=0.99)
@@ -566,7 +566,7 @@ def caculate_loss_LoCo_V2(attn_maps_mid, attn_maps_up, attn_maps_down, bboxes, o
     total_fg_mask = total_fg_map
     fg_map = pad_map * total_fg_mask
 
-    bce_loss = F.binary_cross_entropy(torch.sigmoid(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
     if sum_in + sum_out == 0: