logloss = nn.BCELoss()
ssim_loss = pytorch_ssim.SSIM()
def cosine_loss(a, v, y):
    d = nn.functional.cosine_similarity(a, v)
    loss = logloss(d.unsqueeze(1), y)

    return loss

class LaplacianLoss(nn.Module):
    def __init__(self):
        super(LaplacianLoss, self).__init__()
        # 定义一个拉普拉斯核，注意要符合PyTorch的卷积权重形状(C_out, C_in, H, W)
        self.kernel = torch.tensor([[-1.0, -1.0, -1.0],
                                    [-1.0,  8.0, -1.0],
                                    [-1.0, -1.0, -1.0]], dtype=torch.float32).view(1, 1, 3, 3).requires_grad_(False)
        self.sobel_x = nn.Parameter(torch.tensor([[-1., 0., 1.],
                                          [-2., 0., 2.],
                                          [-1., 0., 1.]], dtype=torch.float32).view(1, 1, 3, 3), requires_grad=False)
        self.sobel_y = nn.Parameter(torch.tensor([[-1., 0., 1.],
                                          [-2., 0., 2.],
                                          [-1., 0., 1.]], dtype=torch.float32).view(1, 1, 3, 3), requires_grad=False)

    def forward(self, output, target):
        # 假设output和target都是[N, C, H, W]的Tensor
        
        # 检查设备，确保卷积核和输入在同一设备上
        device = output.device
        #[batch_size, 3, 5, 288, 288]
        output = output.view(output.shape[0], -1, 288, 288)
        target = target.view(target.shape[0], -1, 288, 288)
        # 如果是RGB图像，需要转化为单通道灰度图，因为拉普拉斯核定义在单通道上
        if output.shape[1] > 3:
            output = output.mean(dim=1, keepdim=True)
            target = target.mean(dim=1, keepdim=True)
        
        # 将拉普拉斯核移动到当前设备上
        laplace_kernel = self.kernel.to(device)
        kernel_x = self.sobel_x.to(device)
        kernel_y = self.sobel_y.to(device)
        
        # 使用F.conv2d应用拉普拉斯滤波
        # filtered_output = F.conv2d(output, laplace_kernel, padding=1)
        # filtered_target = F.conv2d(target, laplace_kernel, padding=1)
        G_x_output = F.conv2d(output, kernel_x, padding=1)
        G_y_output = F.conv2d(output, kernel_y, padding=1)
        filtered_output = torch.sqrt(G_x_output + G_y_output)

        G_x_target = F.conv2d(target, kernel_x, padding=1)
        G_y_target = F.conv2d(target, kernel_y, padding=1)
        filtered_target = torch.sqrt(G_x_target + G_y_target)
        
        #a = filtered_output.detach().cpu().numpy()[0].transpose(1,2,0)
        #b = filtered_target.detach().cpu().numpy()[0].transpose(1,2,0)
        #cv2.imwrite("output.jpg", (a[144:, :, :]*255).astype(np.uint8))
        #cv2.imwrite("target.jpg", (b[144:, :, :]*255).astype(np.uint8))
        #import pdb
        #pdb.set_trace()
        loss = F.l1_loss(filtered_output[:, :, 144:, :], filtered_target[:, :, 144:, :])
        
        return loss

laplacian_loss = LaplacianLoss()
device = torch.device(config["other"]["device"] if use_cuda else "cpu")
syncnet = SyncNet().to(device)
for p in syncnet.parameters():
    p.requires_grad = False

recon_loss = nn.L1Loss()
def get_sync_loss(mel, g):
    g = g[:, :, :, g.size(3)//2:]
    g = torch.cat([g[:, :, i] for i in range(syncnet_T)], dim=1)
    # B, 3 * T, H//2, W
    a, v = syncnet(mel, g)
    y = torch.ones(g.size(0), 1).float().to(device)
    return cosine_loss(a, v, y)

def train(device, model, disc, train_data_loader, test_data_loader, optimizer, disc_optimizer,
          checkpoint_dir=None, checkpoint_interval=None, nepochs=None):
    disc.train()
    model.train()
    
    global global_step, global_epoch
    resumed_step = global_step

    while global_epoch < nepochs:
        print('Starting Epoch: {}'.format(global_epoch))
        running_sync_loss, running_l1_loss, disc_loss, running_perceptual_loss, running_edge_loss, running_ssim_loss = 0., 0., 0., 0., 0., 0.
        running_disc_real_loss, running_disc_fake_loss = 0., 0.
        prog_bar = tqdm(enumerate(train_data_loader))
        for step, (x, indiv_mels, mel, gt) in prog_bar:

            x = x.to(device)
            mel = mel.to(device)
            indiv_mels = indiv_mels.to(device)
            gt = gt.to(device)

            ### Train generator now. Remove ALL grads. 
            optimizer.zero_grad()
            disc_optimizer.zero_grad()

            g = model(indiv_mels, x)
            if hparams.ssim_wt > 0.:
                ssim = -ssim_loss(g.reshape(-1,3,288,288), gt.reshape(-1, 3, 288, 288))
            else:
                ssim = 0.

            if hparams.syncnet_wt > 0.:
                sync_loss = get_sync_loss(mel, g)
            else:
                sync_loss = 0.

            if hparams.disc_wt > 0.:
                perceptual_loss = disc.perceptual_forward(g)
            else:
                perceptual_loss = 0.

            l1loss = recon_loss(g, gt)
            edge_loss = laplacian_loss(g, gt)

            loss = hparams.syncnet_wt * sync_loss + hparams.disc_wt * perceptual_loss + \
                   0.3 * ssim  + 0.4 * edge_loss + (1. - hparams.syncnet_wt - hparams.disc_wt - hparams.ssim_wt - 0.4) * l1loss
            # loss = hparams.syncnet_wt * sync_loss + hparams.disc_wt * perceptual_loss + \
            #        hparams.ssim_wt * ssim  + 0.5 * edge_loss + (1. - hparams.syncnet_wt - hparams.disc_wt - hparams.ssim_wt - 0.1) * l1loss

            loss.backward()
            optimizer.step()
            

            disc_real_loss = 0
            disc_fake_loss = 0
            running_disc_real_loss += disc_real_loss
            running_disc_fake_loss += disc_fake_loss
            
            # # Remove all gradients before Training disc
            # disc_optimizer.zero_grad()

            # pred = disc(gt)
            # disc_real_loss = F.binary_cross_entropy(pred, torch.ones((len(pred), 1)).to(device))
            # disc_real_loss.backward()

            # pred = disc(g.detach())
            # disc_fake_loss = F.binary_cross_entropy(pred, torch.zeros((len(pred), 1)).to(device))
            # disc_fake_loss.backward()

            # disc_optimizer.step()

            # running_disc_real_loss += disc_real_loss.item()
            # running_disc_fake_loss += disc_fake_loss.item()


            # Logs
            global_step += 1
            cur_session_steps = global_step - resumed_step

            running_l1_loss += l1loss.item()
            if hparams.syncnet_wt > 0.:
                running_sync_loss += sync_loss.item()
            else:
                running_sync_loss += 0.

            if hparams.disc_wt > 0.:
                running_perceptual_loss += perceptual_loss.item()
            else:
                running_perceptual_loss += 0.

            if hparams.ssim_wt > 0.:
                running_ssim_loss += -ssim.item()
            else:
                running_ssim_loss += 0.
                
            running_edge_loss += edge_loss.item()


            if global_step % config["gan"]["gan_eval_interval"] == 0:
                with torch.no_grad():
                    average_sync_loss = eval_model(test_data_loader, global_step, device, model, disc, checkpoint_dir)
                    disc.train()
                    model.train()

                    if average_sync_loss < 0.5:
                        hparams.set_hparam('syncnet_wt', 0.03)

            # if global_step % checkpoint_interval == 0:
            #     save_sample_images(x, g, gt, global_step, checkpoint_dir, average_sync_loss)
            
            if  global_step % checkpoint_interval == 0:
                save_checkpoint(
                    model, optimizer, global_step, checkpoint_dir, global_epoch, '', average_sync_loss)
                save_checkpoint(disc, disc_optimizer, global_step, checkpoint_dir, global_epoch, prefix='disc_', val_loss=average_sync_loss)

            prog_bar.set_description('HQ Wav2lip Epoch: {}, L1: {}, Sync: {}, Percep: {} Edge: {}| Fake: {}, Real: {} | SSIM: {}'.format(
                                                                                        global_epoch, 
                                                                                        round(running_l1_loss / (step + 1), 5),
                                                                                        round(running_sync_loss / (step + 1), 5),
                                                                                        round(running_perceptual_loss / (step + 1), 5),
                                                                                        round(running_edge_loss / (step + 1), 5),
                                                                                        round(running_disc_fake_loss / (step + 1), 5),
                                                                                        round(running_disc_real_loss / (step + 1), 5),
                                                                                        round(running_ssim_loss / (step + 1), 5)))
        
        logg.info(f"Epoch{str(global_epoch)} L1 Loss: {round(running_l1_loss / (step + 1), 5)} | Sync Loss: {round(running_sync_loss/(step+1), 5)}")
        global_epoch += 1这个代码，在加入edge loss后 l1和ssim loss全部变成了nan，是不是哪儿有问题