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

gligen/ldm/models/diffusion/loss.py

@@ -566,120 +566,16 @@ 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_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
-
-
-
-def caculate_loss_LAC(attn_maps_mid, attn_maps_up, attn_maps_down, bboxes, object_positions, t, res=16, smooth_att = True,sigma=0.5,kernel_size=3 ):
-    attn16 = get_all_attention(attn_maps_mid, attn_maps_up, attn_maps_down, res)
-    all_attn = [attn16]
-
-
-    loss = 0.
-    pad_loss = 0.
-    total_fg_map = torch.zeros(size=(16, 16)).cuda()
-
-    # alpha是pad loss的权重
-    # beta是pad loss内部的权重 例如 beta是SOT的 1 - beta是EOT的
-    alpha = 0.2
-    beta = 0.8
-
-    object_number = len(bboxes)
-    if object_number == 0:
-        return torch.tensor(0).float().cuda() if torch.cuda.is_available() else torch.tensor(0).float()
-    # attn16 = get_all_attention(attn_maps_down[-1], attn_maps_mid, attn_maps_up[0], 16)
-    # all_attn = [attn16]
-    max_loss = 0
+    return loss #+ alpha * pad_loss
 
 
-    for attn_map in all_attn:
-        # print(attn_map.shape)
-        # 原来是[8, 64, 77] 现在只取后一半 attn_map [4, 64, 77]
-        sum_in = 0.
-        sum_out = 0.
-
-        i, j, k = attn_map.shape
-        H = W = i # 在这里是8
-        for obj_idx in range(object_number): # 对于每个box
-            obj_loss = 0
-            mask = torch.zeros(size=(H, W)).cuda() if torch.cuda.is_available() else torch.zeros(size=(H, W))
-            for obj_box in bboxes[obj_idx]:
 
-                x_min, y_min, x_max, y_max = int(obj_box[0] * W), \
-                    int(obj_box[1] * H), int(obj_box[2] * W), int(obj_box[3] * H)
-                mask[y_min: y_max, x_min: x_max] = 1 # mask是一个全0矩阵 当前物体box的位置设为1
-                total_fg_map[y_min: y_max, x_min: x_max] = 1
 
-            # 选中obj在token中的位置（即token对应的map） reshape到[4, 16, 16]
-            for obj_position in [object_positions[obj_idx]]: # 注意，object_positions是一个list，形如[[6], [10]] 代表第一个物体在第6个token，第二个物体在第10个token
-                # 选中物体对应位置（例如[6]）的map，然后reshape到[4, 16, 16]
-                
-                # print(attn_map[:, :, obj_position].shape)
-                ca_map_obj = attn_map[:, :, obj_position].sum(-1)
-                
-                print(ca_map_obj.shape)
-                if smooth_att:
-                    smoothing = GaussianSmoothing(channels=1, kernel_size=kernel_size, sigma=sigma, dim=2).cuda()
-                    input = F.pad(ca_map_obj.unsqueeze(0).unsqueeze(0), (1, 1, 1, 1), mode='reflect')
-                    ca_map_obj = smoothing(input).squeeze(0).squeeze(0)
-
-                ca_map_obj = ca_map_obj.reshape(H, W)
-                norm_ca_map_obj = ca_map_obj / ca_map_obj.max()
-
-                norm_ca_map_obj = norm_ca_map_obj.reshape(H, W)
-
-                # avg_fg_value = torch.mean(ca_map_obj * mask)
-                # print('avg_fg_value', avg_fg_value) 
-
-                sum_in += (norm_ca_map_obj * mask).sum()
-                sum_out += (norm_ca_map_obj * (1 - mask)).sum()
-
-            # 在这里每个物体对应1个box，所以len是1
-            loss += (obj_loss/len(object_positions[obj_idx]))
-
-        # get pad_loss
-        #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()
-
-
-        #pad_map = beta * norm_sot_map + (1 - beta) * norm_eot_map
-
-    # pad_map = pad_map.to(torch.float64)
-
-    #total_fg_mask = total_fg_map#.to(torch.float64)
-    #fg_map = pad_map * total_fg_mask
-
-    # print(fg_map.shape)
-    # print(pad_map.shape)
-    # 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))
-    # print('mse_loss', mse_loss)
-    # print('bce_loss', bce_loss)
-    #bce_loss = torch.clamp(bce_loss, max=0.99)
-    # pad_loss += mse_loss
-    #pad_loss += bce_loss
-    #pad_loss += (1 - torch.mean((pad_map * total_fg_map).reshape(-1).sum(dim=-1) / pad_map.reshape(-1).sum(dim=-1)) ) **2
-    
-    # print('该步优化结束')
-
-
-    loss += (1 - sum_in / (sum_in + sum_out)) ** 2
-    # loss += max_loss
-    # print('loss', loss)
-    # print('pad_loss', alpha * pad_loss)
-
-
-    return loss + alpha * pad_loss