train_disentangle.py

import datetime
import math
import os
from collections import defaultdict
import json
import random
import numpy as np

import librosa
import pytorch_lightning as pl
import torch
import torch.nn.functional as F
from peft import LoraConfig, get_peft_model
from torch import nn
# from torch.nn import TransformerEncoder, TransformerEncoderLayer
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader
from pytorch_lightning import Trainer, seed_everything
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.strategies.ddp import DDPStrategy
import wandb

from audioldm2.latent_diffusion.models.ddpm import LatentDiffusion
from audioldm2.latent_diffusion.modules.audiomae import models_mae
from audioldm2.utils import plot_mel, save_wave, default_audioldm_config
from calc_metrics import calc_metric, calc_cer
from codebook_clustering import create_reduced_codebook, vote_with_mode
from dataset import mel_spectrogram, read_jsonlines, Meg2MelDataset
from inference import load_models
from brain_module import BrainModule
from ns3_facodec import TransformerEncoder, ResidualVQ
from ns3_facodec.quantize.fvq import FactorizedVectorQuantize
from projector import FeatureProjector, ProjectorLoss
from train_proj import WarmupScheduler


def fix_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


class BrainFVQ(nn.Module):
    def __init__(self, fvq: FactorizedVectorQuantize, feature_dim=256, num_clusters=64, pool_scale=4, temperature=0.2):
        super().__init__()

        # 注册codebook为缓冲区(不参与梯度更新)
        reduced_codebook, mapping = create_reduced_codebook(
            fvq.codebook.weight, num_clusters=num_clusters
        )
        self.register_buffer('codebook', reduced_codebook)
        self.codebook_mapping = mapping

        self.in_proj = fvq.in_proj
        self.out_proj = fvq.out_proj
        for p in self.in_proj.parameters():
            p.requires_grad = False
        for p in self.out_proj.parameters():
            p.requires_grad = False

        self.codebook_size = self.codebook.shape[0]
        self.codebook_dim = self.codebook.shape[1]

        # self.sub_num = sub_num
        self.pool_scale = pool_scale
        self.temperature = temperature

    def forward(self, z):
        # z = self.encoder(z)
        z_e = self.in_proj(z)
        z_q, indices = self.vector_quantization(z_e)
        z_q = (
                z_e + (z_q - z_e).detach()
        )  # noop in forward pass, straight-through gradient estimator in backward pass
        # z_q = z_e
        z_out = self.out_proj(z_q)  # [B, S, H]
        return z_out, z_e

    def vector_quantization(self, z_e, normalize=True):
        bs, t, d = z_e.shape
        z_e = z_e.reshape(-1, d)

        # 计算输入向量与 codebook 向量之间的欧式距离
        if normalize:
            z_e = F.normalize(z_e)
            codebook = F.normalize(self.codebook)
        else:
            codebook = self.codebook
        distances = torch.cdist(z_e, codebook, p=2)  # (batch_size, codebook_size)

        # 找到最小距离对应的索引
        indices = torch.argmin(distances, dim=1)  # (batch_size,)

        # 根据索引从 codebook 中取出量化向量
        z_q = self.codebook[indices]  # (batch_size, d)
        z_q = z_q.reshape(bs, t, -1)
        indices = indices.reshape(bs, t)

        return z_q, indices

    def cal_contrastive_loss(self, z_pred, vq_id, normalize=True):
        """

        Args:
            z_pred: (bs*sub, t, d)
            vq_id: (bs, t)
            normalize: if True, use cosine similarity

        Returns:
            contrastive_loss

        """
        # (bs*sub, t, d) -> (bs*sub, d, t) -> (bs*t*sub, d)
        bs_sub = z_pred.shape[0]
        bs = vq_id.shape[0]
        sub_num = bs_sub // bs

        z_pred = z_pred.permute(0, 2, 1)
        z_pred = F.avg_pool1d(z_pred, kernel_size=self.pool_scale, stride=self.pool_scale)
        d, t = z_pred.shape[-2], z_pred.shape[-1]
        z_pred = z_pred.reshape(bs, sub_num, d, t).permute(0, 3, 1, 2).reshape(-1, d)

        # z_target == codebook  (codebook_size, d)
        if normalize:
            z_pred = F.normalize(z_pred)
            codebook = F.normalize(self.codebook)
        else:
            codebook = self.codebook

        sim_matrix = -torch.cdist(z_pred, codebook)  # 使用负欧氏距离作为相似度
        sim_matrix = sim_matrix / self.temperature

        # (bs, t) -> (bs*t) -> (bs*t*sub)
        mapped_id = [self.codebook_mapping[int(id)] for id in vq_id.reshape(-1)]
        label_id = vote_with_mode(mapped_id, pool=self.pool_scale)
        targets = torch.tensor(label_id, dtype=torch.int64, device=z_pred.device)\
            .unsqueeze(1).repeat(1, sub_num).reshape(-1)

        # 创建标签矩阵 (one-hot)
        target_matrix = F.one_hot(targets, num_classes=self.codebook_size).float()

        # 计算InfoNCE损失 (对数似然)
        log_probs = F.log_softmax(sim_matrix, dim=1)
        loss = -(target_matrix * log_probs).sum(dim=1).mean()

        # target_code = codebook[targets]
        # loss = F.mse_loss(z_pred, target_code)
        # # loss = 1 - F.cosine_similarity(z_pred, target_code, dim=1).mean()

        return loss


class BrainRVQ(nn.Module):
    def __init__(
            self,
            rvq: ResidualVQ,
            feature_dim=256,
            sub_num=1,
            num_clusters=64,
            pool_scale=4,
            temperature=0.2,
    ):
        super().__init__()

        self.layers = nn.ModuleList()
        for fvq in rvq.layers:
            self.layers.append(BrainFVQ(fvq=fvq,
                                        feature_dim=feature_dim,
                                        num_clusters=num_clusters,
                                        pool_scale=pool_scale, temperature=temperature))

    def forward(self, z):
        quantized_out = 0.0
        residual = z
        z_projected = []

        for layer in self.layers:
            z_out, z_e = layer(residual)
            residual = residual - z_out
            quantized_out = quantized_out + z_out
            z_projected.append(z_e.unsqueeze(0))

        z_projected = torch.vstack(z_projected)
        return quantized_out, z_projected

    def cal_contrastive_loss(self, z_pred, vq_id):
        loss = 0.0
        for i, layer in enumerate(self.layers):
            loss += layer.cal_contrastive_loss(z_pred[i], vq_id[i])

        return loss


class BrainRVQEncoders(nn.Module):
    def __init__(
            self,
            fa_decoder,
            feature_dim=256,
            sub_num=1,
            num_clusters=64,
            pool_scale=4,
            temperature=0.2,
            use_prosody=True,
            use_content=True,
            use_detail=True,
            use_timbre=True,
    ):
        super().__init__()

        self.use_prosody = use_prosody
        self.use_content = use_content
        self.use_detail = use_detail
        self.use_timbre = use_timbre

        if use_prosody:
            # self.prosody_encoder = TransformerEncoder(
            #     enc_emb_tokens=None,
            #     encoder_layer=4,
            #     encoder_hidden=feature_dim,
            #     encoder_head=4,
            #     conv_filter_size=1024,
            #     conv_kernel_size=5,
            #     encoder_dropout=0.1,
            #     use_cln=False,
            # )
            self.prosody_vq = BrainRVQ(rvq=fa_decoder.quantizer[0],
                                       feature_dim=feature_dim,
                                       sub_num=sub_num, num_clusters=num_clusters,
                                       pool_scale=pool_scale, temperature=temperature)
        if use_content:
            # self.content_encoder = TransformerEncoder(
            #     enc_emb_tokens=None,
            #     encoder_layer=4,
            #     encoder_hidden=feature_dim,
            #     encoder_head=4,
            #     conv_filter_size=1024,
            #     conv_kernel_size=5,
            #     encoder_dropout=0.1,
            #     use_cln=False,
            # )
            self.content_vq = BrainRVQ(rvq=fa_decoder.quantizer[1],
                                       feature_dim=feature_dim,
                                       sub_num=sub_num, num_clusters=num_clusters,
                                       pool_scale=pool_scale, temperature=temperature)
        if use_detail:
            self.detail_vq = BrainRVQ(rvq=fa_decoder.quantizer[2],
                                      feature_dim=feature_dim,
                                      sub_num=sub_num, num_clusters=num_clusters,
                                      pool_scale=1, temperature=temperature)
        if use_timbre:
            # self.timbre_encoder = TransformerEncoder(
            #     encoder_layer=TransformerEncoderLayer(d_model=vq_dim, nhead=4), num_layers=4)
            self.timbre_encoder = TransformerEncoder(
                enc_emb_tokens=None,
                encoder_layer=4,
                encoder_hidden=feature_dim,
                encoder_head=4,
                conv_filter_size=1024,
                conv_kernel_size=5,
                encoder_dropout=0.1,
                use_cln=False,
            )

    def forward(self, x, return_ortho_loss=False):
        outs = 0
        buf = {}
        z_projected_buf = {}

        if self.use_timbre:
            h_t = self.timbre_encoder(x).mean(1)
            buf['timbre'] = h_t

        p_input = x
        c_input = x
        if self.use_prosody:
            # p_hidden = self.prosody_encoder(p_input)
            z_p, z_projected_p = self.prosody_vq(p_input)
            outs += z_p
            x = x - z_p
            buf['prosody'] = z_p
            z_projected_buf['prosody'] = z_projected_p

        if self.use_content:
            # c_hidden = self.content_encoder(c_input)
            z_c, z_projected_c = self.content_vq(c_input)
            outs += z_c
            x = x - z_c
            buf['content'] = z_c
            z_projected_buf['content'] = z_projected_c

        if self.use_detail:
            z_d, z_projected_d = self.detail_vq(x)
            outs += z_d
            buf['detail'] = z_d
            z_projected_buf['detail'] = z_projected_d

        if not isinstance(outs, torch.Tensor):
            outs = x

        if return_ortho_loss and self.use_prosody and self.use_content:
            ortho_loss = self.cal_ortho_loss(p_hidden, c_hidden)
            return outs, buf, z_projected_buf, ortho_loss

        return outs, buf, z_projected_buf

    def cal_ortho_loss(self, prosody_features, content_features):
        n_batch, n_t, d = prosody_features.shape
        prosody_features = prosody_features.reshape(-1, d)
        content_features = content_features.reshape(-1, d)
        # 归一化特征向量
        p_norm = F.normalize(prosody_features, p=2, dim=1)
        c_norm = F.normalize(content_features, p=2, dim=1)

        # 计算批次内样本的内积矩阵
        similarity = torch.mm(p_norm, c_norm.t())

        # 最小化绝对内积值（促进正交性）
        ortho_loss = torch.mean(torch.abs(similarity))

        return ortho_loss

    def cal_contrastive_loss(self, z_projected_buf, vq_id):
        loss = 0.0
        log_dic = {}

        if self.use_prosody:
            z_projected_p = z_projected_buf['prosody']
            vq_id_p = vq_id[:1]
            p_loss = self.prosody_vq.cal_contrastive_loss(z_projected_p, vq_id_p)
            loss += p_loss
            log_dic.update({"p": p_loss})

        if self.use_content:
            z_projected_c = z_projected_buf['content']
            vq_id_c = vq_id[1:3]
            c_loss = self.content_vq.cal_contrastive_loss(z_projected_c, vq_id_c)
            loss += c_loss
            log_dic.update({"c": c_loss})

        if self.use_detail:
            z_projected_d = z_projected_buf['detail']
            vq_id_d = vq_id[3:]
            d_loss = self.detail_vq.cal_contrastive_loss(z_projected_d, vq_id_d)
            loss += d_loss
            log_dic.update({"d": d_loss})

        return loss, log_dic


class BrainFACodec(pl.LightningModule):
    def __init__(self,
                 modal='fmri',
                 facodecldm_ckpt_path=None,
                 d_t=112,
                 num_patches=56,
                 hidden_dim=768,
                 num_clusters=64,
                 use_prosody=True,
                 use_content=True,
                 use_detail=True,
                 use_timbre=True,
                 use_detail_cycle_loss=True,
                 use_timbre_cycle_loss=True,
                 use_cycle_loss=False,
                 sup_temperature=0.2,
                 cycle_temperature=0.5,
                 pool_scale=4,
                 lambda_cycle=5.0,
                 lambda_sup=5.0,
                 lambda_recon=0.1,
                 sub_num=1,
                 val_sub_list=None,
                 test_sub_list=None,
                 use_clip_loss=False,
                 log_dir=None,
                 val_wav_dir=None,
                 test_wav_dir=None,
                 val_trans_path=None,
                 test_trans_path=None,
                 layout=None
                 ):

        super().__init__()
        # self.multi_sub = multi_sub
        # self.d_t = d_t
        self.num_patches = num_patches

        self.use_prosody = use_prosody
        self.use_content = use_content
        self.use_detail = use_detail
        self.use_timbre = use_timbre
        self.use_detail_cycle_loss = use_detail_cycle_loss
        self.use_timbre_cycle_loss = use_timbre_cycle_loss
        self.use_cycle_loss = use_cycle_loss
        self.sup_temperature = sup_temperature
        self.cycle_temperature = cycle_temperature
        self.pool_scale = pool_scale
        self.lambda_cycle = lambda_cycle
        self.lambda_sup = lambda_sup

        self.sub_num = sub_num
        self.val_sub_list = val_sub_list
        self.test_sub_list = test_sub_list
        self.use_clip_loss = use_clip_loss

        self.modal = modal
        if self.modal == 'fmri':
            self.encoder = FmriEncoder(ch2=hidden_dim, d_t=d_t)
        else:
            self.encoder = BrainModule(out_channels=hidden_dim, loaded_layout=layout,
                                       subject_layers=True, n_subjects=sub_num)

        self.fa_encoder, self.fa_decoder = self.init_facodec()

        self.brain_rvq = BrainRVQEncoders(self.fa_decoder,
                                          feature_dim=hidden_dim,
                                          sub_num=sub_num,
                                          num_clusters=num_clusters,
                                          pool_scale=pool_scale,
                                          temperature=sup_temperature,
                                          use_prosody=use_prosody,
                                          use_content=use_content,
                                          use_detail=use_detail,
                                          use_timbre=use_timbre,
                                          )
        # self.crm = CRM(in_channels=hidden_dim)

        self.log_dir = log_dir
        self.val_wav_dir = val_wav_dir
        self.test_wav_dir = test_wav_dir
        self.val_trans_path = val_trans_path
        self.test_trans_path = test_trans_path
        self.val_metrics = ["similarity", "fpc", "fad", "clap_score", "per"]
        self.test_metrics = ["similarity", "fpc", "mcd", "fad", "clap_score", "per"]

        self.ldm = self.init_ldm(num_patches=num_patches)
        # self.audiomae = self.init_audiomae()
        self.fa_encoder, self.fa_decoder = self.init_facodec()

        # self.dynamic_weight = DynamicWeightedLoss(use_prosody, use_content, use_detail)
        self.lambda_recon = lambda_recon

        self.projector = FeatureProjector(input_seq_len=d_t, output_seq_len=num_patches, input_dim=hidden_dim)
        if facodecldm_ckpt_path is not None:
            self.init_projector(ckpt_path=facodecldm_ckpt_path)
        # self.projector_loss = ProjectorLoss()

        # target_modules = get_target_modules(self.generator.adaptor)
        # prepare_model_for_lora(self.generator.adaptor, target_modules)

    def configure_optimizers(self):
        # return torch.optim.AdamW(self.parameters(), lr=1e-4)
        # return torch.optim.Adam(self.parameters(), lr=2e-4, betas=(0.5, 0.9))
        optimizer = torch.optim.Adam(self.parameters(), lr=2e-4, betas=(0.5, 0.9), weight_decay=1e-4)
        scheduler = WarmupScheduler(optimizer, warmup_steps=1000)
        # scheduler = WarmupExponentialDecayScheduler(optimizer, warmup_steps=self.warmup_steps)
        scheduler_config = {
            "scheduler": scheduler,
            "interval": "step",  # 每一步更新学习率
            "frequency": 1,  # 每步执行一次
        }
        return {"optimizer": optimizer, "lr_scheduler": scheduler_config}

    def init_projector(self, ckpt_path):
        ckpt = torch.load(ckpt_path)["state_dict"]
        new_state_dict = {}
        for key, value in ckpt.items():
            if key.startswith('projector.'):
                new_key = key[len('projector.'):]  # 去掉前缀
                new_state_dict[new_key] = value

        self.projector.load_state_dict(new_state_dict)
        for p in self.projector.parameters():
            p.requires_grad = False

    def init_facodec(self):
        fa_encoder, fa_decoder = load_models()
        for p in fa_encoder.parameters():
            p.requires_grad = False
        for p in fa_decoder.parameters():
            p.requires_grad = False
        return fa_encoder, fa_decoder

    def init_audiomae(self, ckpt_path="checkpoint/audiomae/pretrained.pth"):
        audiomae = models_mae.__dict__["mae_vit_base_patch16"](
            in_chans=1, audio_exp=True, img_size=(1024, 128)
        )
        audiomae_ckpt = torch.load(ckpt_path)['model']
        audiomae.load_state_dict(audiomae_ckpt, strict=False)
        for p in audiomae.parameters():
            p.requires_grad = False
        return audiomae

    def init_ldm(self, num_patches=208, ckpt_path='checkpoint/audioldm2-speech-gigaspeech.pth'):
        model_config = default_audioldm_config(model_name='audioldm2-speech-gigaspeech')
        latent_diffusion = LatentDiffusion(**model_config["model"]["params"])
        latent_diffusion.latent_t_size = num_patches // 2
        ckpt = torch.load(ckpt_path)["state_dict"]
        latent_diffusion.load_state_dict(ckpt, strict=False)

        latent_diffusion.use_ema = False
        for p in latent_diffusion.parameters():
            p.requires_grad = False
        # for p in latent_diffusion.model.parameters():
        #     p.requires_grad = True
        return latent_diffusion

    def forward(self, x):
        h_d, disentangled_buf = self.disentangle(x)
        if self.use_timbre:
            c = self.crm(h_d, disentangled_buf['timbre'])
        else:
            c = h_d
        c = self.projector(c)
        recon_audio, recon_mel = self.ldm.recon(c)  # (B, F, T)
        return recon_audio, recon_mel

    def disentangle(self, x, subj_idx):
        latent_embed = self.encoder(x, subj_idx)
        h_d, disentangled_buf, _ = self.brain_rvq(latent_embed)
        return h_d, disentangled_buf

    def recon(self, h_d, h_t):
        # if self.use_timbre:
        #     c = self.crm(h_d, h_t)
        # else:
        #     c = h_d
        # c = self.projector(c)
        # recon_audio, recon_mel = self.ldm.recon(c)  # (B, F, T)
        # return recon_audio, recon_mel
        recon_audio, recon_mel = self.generator.recon(h_d.transpose(1, 2), h_t)
        return recon_audio, recon_mel

    def generate_gt_cond(self, ta_kaldi_fbank):
        emb_enc = self.audiomae.forward_encoder_no_mask(ta_kaldi_fbank)
        cond_audiomae = emb_enc[:, 1:1+self.num_patches, :].detach()  # remove cls
        return cond_audiomae

    def multi_sub_contrastive_loss(self, z_buf, audio, loss_fn, include_timbre=True):
        losses = []
        loss = 0.0
        log_dic = {}

        enc_out = self.fa_encoder(audio)
        vq_post_emb, _, _, quantized, spk_embs = self.fa_decoder(enc_out, eval_vq=False, vq=True)

        if self.use_prosody:
            z_p = z_buf['prosody']
            z_p_rec = quantized[0]
            p_cycle_loss = loss_fn(z_p, z_p_rec)
            loss += p_cycle_loss
            # losses.append(p_cycle_loss)
            log_dic.update({"p": p_cycle_loss})

        if self.use_content:
            z_c = z_buf['content']
            z_c_rec = quantized[1]
            c_cycle_loss = loss_fn(z_c, z_c_rec)
            loss += c_cycle_loss
            # losses.append(c_cycle_loss)
            log_dic.update({"c": c_cycle_loss})

        if self.use_detail and self.use_detail_cycle_loss:
            z_d = z_buf['detail']
            z_d_rec = quantized[2]
            d_cycle_loss = loss_fn(z_d, z_d_rec)
            loss += d_cycle_loss
            # losses.append(d_cycle_loss)
            log_dic.update({"d": d_cycle_loss})

        if include_timbre and self.use_timbre and self.use_timbre_cycle_loss:
            timbre_embed = z_buf['timbre']
            t_cycle_loss = loss_fn(timbre_embed, spk_embs, is_timbre=True)
            loss += t_cycle_loss
            # losses.append(t_cycle_loss)
            log_dic.update({"t": t_cycle_loss})

        # loss = self.auto_weight(losses)
        return loss, log_dic

    def compute_cross_cycle_loss(self, encoder, decoder, z_prosody, z_content, z_detail):
        """
        Compute cross-cycle consistency loss for disentangled representations.

        Args:
            encoder: Function to encode audio into (z_prosody, z_content, z_detail).
            decoder: Function to decode (z_prosody, z_content, z_detail) into audio.
            z_prosody: Tensor of shape (batch_size, time_steps, feature_dim) for prosody.
            z_content: Tensor of shape (batch_size, time_steps, feature_dim) for content.
            z_detail: Tensor of shape (batch_size, time_steps, feature_dim) for detail.

        Returns:
            cross_cycle_loss: Scalar tensor representing the cross-cycle loss.
        """
        batch_size = z_prosody.size(0)
        cross_cycle_loss = 0.0

        # Replace prosody
        for i in range(batch_size):
            z_prosody_mix = z_prosody.clone()
            z_prosody_mix[i] = z_prosody[(i + 1) % batch_size]  # Replace with next batch's prosody
            mixed_audio = decoder(z_prosody_mix, z_content, z_detail)
            z_p_mix, z_c_mix, z_d_mix = encoder(mixed_audio)
            cross_cycle_loss += F.mse_loss(z_prosody[i], z_p_mix[i])  # ||z_prosody_1 - z_p_mix||

        # Replace content
        for i in range(batch_size):
            z_content_mix = z_content.clone()
            z_content_mix[i] = z_content[(i + 1) % batch_size]  # Replace with next batch's content
            mixed_audio = decoder(z_prosody, z_content_mix, z_detail)
            z_p_mix, z_c_mix, z_d_mix = encoder(mixed_audio)
            cross_cycle_loss += F.mse_loss(z_content[i], z_c_mix[i])  # ||z_content_2 - z_c_mix||

        # Replace detail
        for i in range(batch_size):
            z_detail_mix = z_detail.clone()
            z_detail_mix[i] = z_detail[(i + 1) % batch_size]  # Replace with next batch's detail
            mixed_audio = decoder(z_prosody, z_content, z_detail_mix)
            z_p_mix, z_c_mix, z_d_mix = encoder(mixed_audio)
            cross_cycle_loss += F.mse_loss(z_detail[i], z_d_mix[i])  # ||z_detail_1 - z_d_mix||

        return cross_cycle_loss

    def compute_cosine_loss(self, z_pred, z_target, is_timbre=False, add_noise=False):
        """
        计算余弦相似度损失。

        参数：
        - z_pred: 形状 (bs * sub, t, d)
        - z_target: 形状 (bs, d, t)
        - sub: 子维度 sub 的大小

        返回：
        - 平均余弦损失
        """
        if is_timbre:
            bs_sub, d = z_pred.shape
            bs, d = z_target.shape
            sub_num = bs_sub // bs
            z_pred = z_pred.view(bs, sub_num, d).permute(1, 0, 2)
        else:
            bs_sub, t, d = z_pred.shape
            bs, d, t = z_target.shape
            sub_num = bs_sub // bs

            # 重新调整形状
            z_pred = z_pred.view(bs, sub_num, t, d).permute(1, 0, 2, 3).reshape(sub_num, bs * t, d)
            z_target = z_target.permute(0, 2, 1).reshape(bs * t, d)  # (bs*t, d)

        if add_noise:
            noise = torch.normal(mean=0.0, std=0.004, size=z_target.size(), device=z_target.device)
            z_target = z_target.detach() + noise  # * self.alpha

        # 计算余弦损失
        cosine_losses = []
        for i in range(sub_num):
            cosine_loss = 1 - F.cosine_similarity(z_pred[i], z_target, dim=1).mean()
            cosine_losses.append(cosine_loss)

        # 计算 sub 维度的均值
        return torch.stack(cosine_losses).mean()

    def sup_loss(self, z_pred, z_target):
        # (bs*sub, t, d) -> (bs*sub, d, t) -> (bs*t*sub, d)
        z_pred = z_pred.permute(0, 2, 1)
        z_pred = F.avg_pool1d(z_pred, kernel_size=self.pool_scale, stride=self.pool_scale)
        d, t = z_pred.shape[-2], z_pred.shape[-1]
        z_pred = z_pred.reshape(-1, self.sub_num, d, t).permute(0, 3, 1, 2).reshape(-1, d)

        # (bs, d, t) -> (bs*t, d)
        d, t = z_target.shape[-2], z_target.shape[-1]
        z_target = F.avg_pool1d(z_target, kernel_size=self.pool_scale, stride=self.pool_scale)
        z_target = z_target.permute(0, 2, 1).reshape(-1, d)

        noise = torch.normal(mean=0.0, std=0.004, size=z_target.size(), device=z_target.device)
        z_target = z_target.detach() + noise  # * self.alpha

        bs = z_pred.shape[0] // self.sub_num
        sub_num = self.sub_num
        similarity_matrix = torch.cosine_similarity(
            z_pred.unsqueeze(1),
            z_target.unsqueeze(0),
            dim=-1
        ) / self.sup_temperature
        similarity_matrix = torch.exp(similarity_matrix)

        pos_sim_index_1 = torch.arange(bs, device=similarity_matrix.device) \
            .view(bs, 1).unsqueeze(1) \
            .expand(bs, sub_num, 1).contiguous().view(bs * sub_num, 1)
        pos_sim_1 = torch.gather(input=similarity_matrix, dim=1, index=pos_sim_index_1)
        loss_1 = torch.mean(
            -torch.div(pos_sim_1.sum(dim=1), similarity_matrix.sum(dim=1)).log())

        pos_sim_index_2 = torch.arange(bs * sub_num, device=similarity_matrix.device).view(bs, sub_num).T
        pos_sim_2 = torch.gather(input=similarity_matrix, dim=0, index=pos_sim_index_2)
        loss_2 = torch.mean(
            -torch.div(pos_sim_2.sum(dim=0), similarity_matrix.sum(dim=0)).log())
        return loss_1 + loss_2

    def cycle_loss(self, z_pred, z_target):
        # (bs*sub, t, d) -> (bs*sub, d, t) -> (sub, bs*t, d)
        z_pred = z_pred.permute(0, 2, 1)
        z_pred = F.avg_pool1d(z_pred, kernel_size=self.pool_scale, stride=self.pool_scale)
        d, t = z_pred.shape[-2], z_pred.shape[-1]
        z_pred = z_pred.reshape(-1, self.sub_num, d, t).permute(1, 0, 3, 2).reshape(self.sub_num, -1, d)

        # (bs*sub, d, t) -> (sub, bs*t, d)
        z_target = F.avg_pool1d(z_target, kernel_size=self.pool_scale, stride=self.pool_scale)
        d, t = z_target.shape[-2], z_target.shape[-1]
        z_target = z_target.reshape(-1, self.sub_num, d, t).permute(1, 0, 3, 2).reshape(self.sub_num, -1, d)

        bs_t = z_pred.shape[1]
        cont_losses = []
        for i in range(self.sub_num):
            similarity_matrix = torch.cosine_similarity(
                z_pred[i].unsqueeze(1),
                z_target[i].unsqueeze(0),
                dim=-1
            ) / self.sup_temperature
            similarity_matrix = torch.exp(similarity_matrix)

            pos_sim_index_1 = torch.arange(bs_t, device=similarity_matrix.device).view(bs_t, 1)
            pos_sim_1 = torch.gather(input=similarity_matrix, dim=1, index=pos_sim_index_1)
            loss_1 = torch.mean(
                -torch.div(pos_sim_1.sum(dim=1), similarity_matrix.sum(dim=1)).log())
            cont_losses.append(loss_1)
        return torch.stack(cont_losses).mean()

    def training_step(self, batch, batch_idx):
        loss = 0.0
        log_dic = {}

        brain_data, audio, _, _, mel_padded, ta_kaldi_fbank, wav_name, subj = batch

        if self.modal == 'fmri':
            bs, sub_num, d, h, w = brain_data.shape
            assert sub_num == self.sub_num
            brain_data = brain_data.view(-1, d, h, w)
        else:
            bs, sub_num, ch, t = brain_data.shape
            assert sub_num == self.sub_num
            brain_data = brain_data.view(-1, ch, t)
            subj = subj.view(-1)
            subj_idx = {'subject_index': subj}
        mel_padded = mel_padded.unsqueeze(1).repeat(1, sub_num, 1, 1, 1) \
            .view(-1, mel_padded.shape[-3], mel_padded.shape[-2], mel_padded.shape[-1])

        latent_embed = self.encoder(brain_data, subj_idx)
        outs, out_buf, projected_buf = self.brain_rvq(latent_embed)

        if self.use_timbre:
            # h_d = self.crm(outs, out_buf['timbre'])
            style = self.fa_decoder.timbre_linear(out_buf['timbre']).unsqueeze(1)  # (B, 1, 2d)
            gamma, beta = style.chunk(2, dim=2)  # (B, 1, d)

            h_d = outs
            h_d = self.fa_decoder.timbre_norm(h_d)
            h_d = h_d * gamma + beta
        else:
            h_d = outs
        c = self.projector(h_d)

        recon_loss, recon_mel = self.ldm(mel=mel_padded, cond=c, return_mel_recon=True)
        recon_audio, recon_mel = self.ldm.mel_spectrogram_to_waveform(
            recon_mel, savepath='', to_numpy=False, name='', save=False
        )
        recon_audio = recon_audio[:, :, :audio.shape[-1]]
        loss += recon_loss * self.lambda_recon
        log_dic.update({"recon_loss": recon_loss})

        enc_out = self.fa_encoder(audio)
        vq_post_emb, vq_id, _, quantized, spk_embs = self.fa_decoder(enc_out, eval_vq=False, vq=True)

        # supervision loss
        cont_loss, cont_log_dict = self.brain_rvq.cal_contrastive_loss(projected_buf, vq_id)
        # cont_loss = self.dynamic_weight(cont_log_dict)
        # log_dic.update(self.dynamic_weight.log_vars())
        loss += cont_loss
        log_dic.update({("cont_loss/" + k): float(v) for k, v in cont_log_dict.items()})

        if self.use_timbre and self.use_timbre_cycle_loss:
            timbre_embed = out_buf['timbre']
            cosine_loss = self.compute_cosine_loss(timbre_embed, spk_embs, is_timbre=True, add_noise=True)
            loss += cosine_loss
            log_dic.update({"sim_loss/t": cosine_loss})

        latent_loss = self.sup_loss(latent_embed, enc_out)
        loss += latent_loss
        log_dic.update({"cont_loss/latent": latent_loss})

        self.log_dict(
            {("train/" + k): float(v) for k, v in log_dic.items()},
            prog_bar=True,
            logger=True,
            on_step=True,
            on_epoch=False,
            sync_dist=True
        )

        self.log(
            "global_step",
            float(self.global_step),
            prog_bar=True,
            logger=True,
            on_step=True,
            on_epoch=False,
        )

        lr = self.trainer.optimizers[0].param_groups[0]["lr"]
        self.log(
            "lr_abs",
            float(lr),
            prog_bar=True,
            logger=True,
            on_step=True,
            on_epoch=False,
        )

        return loss

    def get_validation_folder_name(self):
        now = datetime.datetime.now()
        timestamp = now.strftime("%m-%d-%H:%M")
        return "val_%s_%s" % (
            self.global_step,
            timestamp,
        )

    def on_validation_epoch_start(self) -> None:
        self.save_path = os.path.join(self.log_dir, self.get_validation_folder_name())
        self.save_path_list = [os.path.join(self.save_path, sub) for sub in self.val_sub_list]
        for save_path in self.save_path_list:
            os.makedirs(os.path.join(save_path, 'img'), exist_ok=True)
            os.makedirs(os.path.join(save_path, 'wav'), exist_ok=True)
        # print("Mel save path: ", self.img_save_path)
        # print("Wav save path: ", self.wav_save_path)

    @torch.no_grad()
    def validation_step(self, batch, batch_idx):
        log_dic = {}
        brain_data, audio, _, mel, _, _, wav_name, subj = batch
        if self.modal == 'fmri':
            bs, sub_num, d, h, w = brain_data.shape
            # assert sub_num == self.sub_num
            brain_data = brain_data.view(-1, d, h, w)
        else:
            bs, sub_num, ch, t = brain_data.shape
            # assert sub_num == self.sub_num
            brain_data = brain_data.view(-1, ch, t)
            subj = subj.view(-1)
            subj_idx = {'subject_index': subj}
        mel = mel.unsqueeze(1).repeat(1, sub_num, 1, 1, 1) \
            .view(-1, mel.shape[-3], mel.shape[-2], mel.shape[-1])

        latent_embed = self.encoder(brain_data, subj_idx)
        outs, out_buf, quantize_buf = self.brain_rvq(latent_embed)
        # recon_audio = self.fa_decoder.inference(outs.transpose(1, 2), quantized_buf['timbre'])

        if self.use_timbre:
            # h_d = self.crm(outs, out_buf['timbre'])
            style = self.fa_decoder.timbre_linear(out_buf['timbre']).unsqueeze(1)  # (B, 1, 2d)
            gamma, beta = style.chunk(2, dim=2)  # (B, 1, d)

            h_d = outs
            h_d = self.fa_decoder.timbre_norm(h_d)
            h_d = h_d * gamma + beta
        else:
            h_d = outs
        c = self.projector(h_d)

        recon_audio, recon_mel = self.ldm.recon(c)  # (B, F, T)
        recon_audio = recon_audio[:, :, :audio.shape[-1]]
        recon_mel = recon_mel[:, :, :mel.shape[-2]]

        # # supervision loss
        # # quantized_buf['timbre'] = timbre_embed
        # sup_loss, sup_log_dic = self.multi_sub_contrastive_loss(quantized_buf, audio, loss_fn=self.sup_loss)
        # log_dic.update({("sup_loss/" + k): float(v) for k, v in sup_log_dic.items()})
        #
        # # cycle loss
        # # quantized_buf['timbre'] = timbre_embed
        # cycle_loss, cycle_log_dic = self.multi_sub_contrastive_loss(quantized_buf, recon_audio, loss_fn=self.cycle_loss)
        # log_dic.update({("cycle_loss/" + k): float(v) for k, v in cycle_log_dic.items()})

        # sim loss
        sim_loss, sim_log_dic = self.multi_sub_contrastive_loss(out_buf, audio, loss_fn=self.compute_cosine_loss)
        log_dic.update({("sim_loss/" + k): float(v) for k, v in sim_log_dic.items()})

        self.log_dict(
            {("val/" + k): float(v) for k, v in log_dic.items()},
            prog_bar=True,
            logger=True,
            on_step=False,
            on_epoch=True,
            sync_dist=True
        )

        # recon_mel = torch.cat([mel_spectrogram(recon_audio_seg) for recon_audio_seg in recon_audio])
        recon_audio = recon_audio.cpu().detach().numpy()
        mel = mel[:, 0, :, :].permute(0, 2, 1).cpu().detach().numpy()  # (B, F, T)
        recon_mel = recon_mel.cpu().detach().numpy()  # (B, F, T)

        for i in range(len(wav_name)):
            sample_name = wav_name[i]
            for sub_id, save_path in enumerate(self.save_path_list):
                img_save_name = os.path.join(save_path, 'img', sample_name + '.png')
                plot_mel(mel[i * sub_num + sub_id], recon_mel[i * sub_num + sub_id], fig_path=img_save_name)
                wav_save_name = os.path.join(save_path, 'wav', sample_name + '.wav')
                save_wave(recon_audio[i * sub_num + sub_id], wav_path=wav_save_name)
        # else:
        #     sample_name = wav_name[i]
        #     img_save_name = os.path.join(self.img_save_path, sample_name + '.png')
        #     plot_mel(mel[i], recon_mel[i], fig_path=img_save_name)
        #     wav_save_name = os.path.join(self.wav_save_path, sample_name + '.wav')
        #     save_wave(recon_audio[i], wav_path=wav_save_name)

    def on_validation_epoch_end(self) -> None:
        values_dict = defaultdict(list)
        for save_path in self.save_path_list:
            print('Calculate Metrics:', save_path)
            result = calc_metric(
                self.val_wav_dir,
                os.path.join(save_path, 'wav'),
                save_path,
                self.val_metrics,
                fs=16000,
                method="dtw",
                db_scale=True,
                need_mean=True,
                model_name="wavlm",
                similarity_mode="pairwith",
                ltr_path="None",
                intelligibility_mode="gt_audio",
                language="english",
            )
            for key, value in result.items():
                values_dict[key].append(value)

        mean_dict = {key: sum(values) / len(values) for key, values in values_dict.items()}

        data = json.dumps(mean_dict, indent=4)
        with open(os.path.join(self.save_path, "result.json"), "w", newline="\n") as f:
            f.write(data)

        self.log_dict(
            {("val/metric/" + k): v for k, v in mean_dict.items()},
            prog_bar=True,
            logger=True,
            on_step=False,
            on_epoch=True,
            sync_dist=True
        )

    def on_test_epoch_start(self) -> None:
        self.test_save_path = os.path.join(self.log_dir, "test")
        self.test_save_path_list = [os.path.join(self.test_save_path, sub) for sub in self.test_sub_list]
        for save_path in self.test_save_path_list:
            os.makedirs(os.path.join(save_path, 'img'), exist_ok=True)
            os.makedirs(os.path.join(save_path, 'wav'), exist_ok=True)
        # print("Mel save path: ", self.test_img_save_path)
        # print("Wav save path: ", self.test_wav_save_path)

    @torch.no_grad()
    def test_step(self, batch, batch_idx):
        brain_data, audio, _, mel, _, _, wav_name, subj = batch
        if self.modal == 'fmri':
            bs, sub_num, d, h, w = brain_data.shape
            assert sub_num == len(self.test_sub_list)
            brain_data = brain_data.view(-1, d, h, w)
        else:
            bs, sub_num, ch, t = brain_data.shape
            assert sub_num == len(self.test_sub_list)
            brain_data = brain_data.view(-1, ch, t)
            subj = subj.view(-1)
            subj_idx = {'subject_index': subj}
        mel = mel.unsqueeze(1).repeat(1, sub_num, 1, 1, 1) \
            .view(-1, mel.shape[-3], mel.shape[-2], mel.shape[-1])

        latent_embed = self.encoder(brain_data, subj_idx)
        outs, out_buf, quantize_buf = self.brain_rvq(latent_embed)
        # recon_audio = self.fa_decoder.inference(outs.transpose(1, 2), quantized_buf['timbre'])

        if self.use_timbre:
            # h_d = self.crm(outs, out_buf['timbre'])
            style = self.fa_decoder.timbre_linear(out_buf['timbre']).unsqueeze(1)  # (B, 1, 2d)
            gamma, beta = style.chunk(2, dim=2)  # (B, 1, d)

            h_d = outs
            h_d = self.fa_decoder.timbre_norm(h_d)
            h_d = h_d * gamma + beta
        else:
            h_d = outs
        c = self.projector(h_d)

        recon_audio, recon_mel = self.ldm.recon(c)  # (B, F, T)
        recon_audio = recon_audio[:, :, :audio.shape[-1]]
        recon_mel = recon_mel[:, :, :mel.shape[-2]]

        # recon_mel = torch.cat([mel_spectrogram(recon_audio_seg) for recon_audio_seg in recon_audio])
        recon_audio = recon_audio.cpu().detach().numpy()
        mel = mel[:, 0, :, :].permute(0, 2, 1).cpu().detach().numpy()  # (B, F, T)
        recon_mel = recon_mel.cpu().detach().numpy()  # (B, F, T)

        for i in range(len(wav_name)):
            sample_name = wav_name[i]
            for sub_id, save_path in enumerate(self.test_save_path_list):
                img_save_name = os.path.join(save_path, 'img', sample_name + '.png')
                plot_mel(mel[i * sub_num + sub_id], recon_mel[i * sub_num + sub_id], fig_path=img_save_name)
                wav_save_name = os.path.join(save_path, 'wav', sample_name + '.wav')
                save_wave(recon_audio[i * sub_num + sub_id], wav_path=wav_save_name)

    def on_test_epoch_end(self) -> None:
        result_list = []
        for save_path in self.test_save_path_list:
            print('Calculate Metrics:', save_path)
            result = calc_metric(
                self.test_wav_dir,
                os.path.join(save_path, 'wav'),
                save_path,
                self.test_metrics,
                fs=16000,
                method="dtw",
                db_scale=True,
                need_mean=True,
                model_name="wavlm",
                similarity_mode="pairwith",
                ltr_path="None",
                intelligibility_mode="gt_audio",
                language="english",
            )
            result_list.append(result)

        self.log_dict(
            {("test/metric/" + k): v for k, v in result_list[0].items()},
            prog_bar=True,
            logger=True,
            on_step=False,
            on_epoch=True,
            sync_dist=True
        )
        # self.log_dict(
        #     {("test_cross/metric/" + k): v for k, v in result_list[-1].items()},
        #     prog_bar=True,
        #     logger=True,
        #     on_step=False,
        #     on_epoch=True,
        #     sync_dist=True
        # )


def train_meg_gw():
    fix_seed(42)
    torch.set_float32_matmul_precision("high")
    sub_id_list = [1, 2] + list(range(4, 12)) + list(range(13, 16)) + list(range(17, 20)) + list(range(22, 28))

    training_sub_num = 11
    train_sub_list = ['sub-%02d' % sub_id for sub_id in sub_id_list[:training_sub_num]]
    val_sub_list = ['sub-01', 'sub-02', 'sub-04']
    test_sub_list = ['sub-01', 'sub-02', 'sub-04']
    # test_sub_list = ['sub-%02d' % sub_id for sub_id in sub_id_list[-3:]]

    epoch_num = 10

    proj_name = 'Fine-Recon-Gw-%s' % train_sub_list[0] if training_sub_num == 1 else 'Fine-Recon-Gw-%d' % training_sub_num

    dataset_path = "path/to/MEG-MASC/derivatives/meg_sr_80/"
    train_data_list = []
    val_data_list = []
    test_data_list = []
    for sub in train_sub_list:
        train_path = dataset_path + sub + "/train.jsonl"
        train_data_list.append(read_jsonlines(train_path))

    for sub in val_sub_list:
        valid_path = dataset_path + sub + "/val.jsonl"
        val_data_list.append(read_jsonlines(valid_path))

    for sub in test_sub_list:
        test_path = dataset_path + sub + "/test.jsonl"
        test_data_list.append(read_jsonlines(test_path))

    num_patches = 208
    train_dataset = Meg2MelDataset(train_data_list, sub_id_list, num_patches=num_patches, duration=4.)
    print('Training Samples:', len(train_dataset))
    train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=30)

    val_dataset = Meg2MelDataset(val_data_list, sub_id_list, num_patches=num_patches, duration=4.)
    print('Valid Samples:', len(val_dataset))
    val_loader = DataLoader(val_dataset, batch_size=30, shuffle=False, num_workers=1)

    test_dataset = Meg2MelDataset(test_data_list, sub_id_list, num_patches=num_patches, duration=4.)
    print('Test Samples:', len(test_dataset))
    test_loader = DataLoader(test_dataset, batch_size=30, shuffle=False, num_workers=1)

    test_wav_dir = "path/to/MEG-MASC/derivatives/test_wav/"
    val_wav_dir = "path/to/MEG-MASC/derivatives/val_wav/"

    log_path = 'path/to/log/%s/DisBrain-Discrete-ProtoRVQ32-recon01-cosLoss-freezeG-warmup-clip-grad-fad-%depoch' % (proj_name, epoch_num)
    wandb_path = log_path
    checkpoint_path = os.path.join(log_path, "checkpoints")
    os.makedirs(checkpoint_path, exist_ok=True)

    checkpoint_callback = ModelCheckpoint(
        dirpath=checkpoint_path,
        monitor="val/metric/fad",
        mode="min",
        filename="checkpoint_step={global_step:.0f}-fad={val/metric/fad:.3f}",
        auto_insert_metric_name=False,
        # every_n_train_steps=save_checkpoint_every_n_steps,
        # every_n_epochs=10,
        save_top_k=2,
        save_last=False,
        save_weights_only=True
    )

    wandb_logger = WandbLogger(
        save_dir=wandb_path,
        project=proj_name,
        name="DisBrain-Discrete-ProtoRVQ32-recon01-cosLoss-freezeG-warmup-clip-grad-fad[%depoch]" % (epoch_num),
    )

    import pickle
    with open('layout.pkl', 'rb') as f:
        layout = pickle.load(f)

    facodecldm_ckpt_path = "path/to/log/Fine-Recon-Projector/checkpoint.ckpt"

    # Initialize the network
    net = BrainFACodec(modal='meg', sub_num=training_sub_num,
                       d_t=320, num_patches=num_patches, hidden_dim=256,
                       num_clusters=32,
                       lambda_recon=0.1,
                       layout=layout,
                       facodecldm_ckpt_path=facodecldm_ckpt_path,
                       use_detail=False,
                       val_sub_list=val_sub_list, test_sub_list=test_sub_list,
                       log_dir=log_path, val_wav_dir=val_wav_dir, test_wav_dir=test_wav_dir)
                       # val_trans_path=val_trans_path, test_trans_path=test_trans_path)

    trainer = Trainer(
        accelerator="gpu",
        devices=1,
        logger=wandb_logger,
        max_epochs=epoch_num,
        num_sanity_val_steps=0,  # 先运行一遍val
        # limit_val_batches=10,
        check_val_every_n_epoch=2,
        # val_check_interval=2736,
        strategy=DDPStrategy(find_unused_parameters=True),
        callbacks=[checkpoint_callback],
        gradient_clip_val=1.0
    )

    trainer.fit(net, train_loader, val_loader)
    trainer.test(net, test_loader, ckpt_path="best")
    wandb.finish()


def test_meg_gw():
    fix_seed(42)
    torch.set_float32_matmul_precision("high")
    sub_id_list = [1, 2] + list(range(4, 12)) + list(range(13, 16)) + list(range(17, 20)) + list(range(22, 28))

    training_sub_num = 11
    train_sub_list = ['sub-%02d' % sub_id for sub_id in sub_id_list[:training_sub_num]]

    test_sub_list = ['sub-%02d' % sub_id for sub_id in sub_id_list[3:11]]
    epoch_num = 10

    proj_name = 'Fine-Recon-Gw-%s' % train_sub_list[0] if training_sub_num == 1 else 'Fine-Recon-Gw-%d' % training_sub_num

    dataset_path = "path/to/MEG-MASC/derivatives/meg_sr_80/"
    test_data_list = []

    for sub in test_sub_list:
        test_path = dataset_path + sub + "/test.jsonl"
        test_data_list.append(read_jsonlines(test_path))

    num_patches = 208

    test_dataset = Meg2MelDataset(test_data_list, sub_id_list, num_patches=num_patches, duration=4.)
    print('Test Samples:', len(test_dataset))
    test_loader = DataLoader(test_dataset, batch_size=30, shuffle=False, num_workers=1)

    test_wav_dir = "path/to/MEG-MASC/derivatives/test_wav/"
    val_wav_dir = "path/to/MEG-MASC/derivatives/val_wav/"

    log_path = 'path/to/log/%s/DisBrain-Discrete-ProtoRVQ32-recon01-cosLoss-freezeG-warmup-clip-grad-fad-%depoch' % (proj_name, epoch_num)

    import pickle
    with open('layout.pkl', 'rb') as f:
        layout = pickle.load(f)

    facodecldm_ckpt_path = "path/to/log/Fine-Recon-Projector/checkpoint.ckpt"

    # Initialize the network
    net = BrainFACodec(modal='meg', sub_num=training_sub_num,
                       d_t=320, num_patches=num_patches, hidden_dim=256,
                       num_clusters=32,
                       lambda_recon=0.1,
                       layout=layout,
                       facodecldm_ckpt_path=facodecldm_ckpt_path,
                       use_detail=False,
                       test_sub_list=test_sub_list,
                       log_dir=log_path, val_wav_dir=val_wav_dir, test_wav_dir=test_wav_dir)
                       # val_trans_path=val_trans_path, test_trans_path=test_trans_path)

    ckpt_path = "path/to/log/Fine-Recon-Gw-11/DisBrain-Discrete-ProtoRVQ32-recon01-cosLoss-freezeG-warmup-clip-grad-fad-10epoch/checkpoints/checkpoint.ckpt"

    trainer = Trainer(
        accelerator="gpu",
        devices=1,
        # logger=wandb_logger,
        max_epochs=epoch_num,
        num_sanity_val_steps=0,  # 先运行一遍val
        # limit_val_batches=10,
        check_val_every_n_epoch=2,
        # val_check_interval=2736,
        strategy=DDPStrategy(find_unused_parameters=True),
        # callbacks=[checkpoint_callback],
        gradient_clip_val=1.0
    )

    trainer.test(net, test_loader, ckpt_path=ckpt_path)
    # wandb.finish()


def train_meg_sherlock():
    fix_seed(42)
    torch.set_float32_matmul_precision("high")

    subj = 'sub-001'
    sub_id_list = [1, 2]
    train_sub_list = [subj]
    val_sub_list = [subj]
    test_sub_list = ['sub-001', 'sub-002']
    epoch_num = 10
    proj_name = 'Fine-Recon-Sherlock-%s' % subj

    train_data_list = []
    val_data_list = []
    test_data_list = []
    for sub in train_sub_list:
        train_path = "path/to/Holmes/meg_80/%s/train.jsonl" % sub
        train_data_list.append(read_jsonlines(train_path))

    for sub in val_sub_list:
        valid_path = "path/to/Holmes/meg_80/%s/val.jsonl" % sub
        val_data_list.append(read_jsonlines(valid_path))

    for sub in test_sub_list:
        test_path = "path/to/Holmes/meg_80/%s/test.jsonl" % sub
        test_data_list.append(read_jsonlines(test_path))

    num_patches = 208
    train_dataset = Meg2MelDataset(train_data_list, sub_id_list, num_patches=num_patches, duration=4.)
    print('Training Samples:', len(train_dataset))
    train_loader = DataLoader(train_dataset, batch_size=25, shuffle=True, num_workers=30)

    val_dataset = Meg2MelDataset(val_data_list, sub_id_list, num_patches=num_patches, duration=4.)
    print('Valid Samples:', len(val_dataset))
    val_loader = DataLoader(val_dataset, batch_size=30, shuffle=False, num_workers=1)

    test_dataset = Meg2MelDataset(test_data_list, sub_id_list, num_patches=num_patches, duration=4.)
    print('Test Samples:', len(test_dataset))
    test_loader = DataLoader(test_dataset, batch_size=30, shuffle=False, num_workers=1)

    test_wav_dir = "path/to/Holmes/wav_4s/ses-010/"
    val_wav_dir = "path/to/Holmes/wav_4s/ses-009/"

    log_path = 'path/to/log/%s/DisBrain-Discrete-ProtoRVQ32-cosLoss-freezeG-warmup-clip-grad-fad-%depoch' % (proj_name, epoch_num)
    wandb_path = log_path
    checkpoint_path = os.path.join(log_path, "checkpoints")
    os.makedirs(checkpoint_path, exist_ok=True)

    checkpoint_callback = ModelCheckpoint(
        dirpath=checkpoint_path,
        monitor="val/metric/fad",
        mode="min",
        filename="checkpoint_step={global_step:.0f}-fad={val/metric/fad:.3f}",
        auto_insert_metric_name=False,
        # every_n_train_steps=save_checkpoint_every_n_steps,
        # every_n_epochs=10,
        save_top_k=1,
        save_last=False,
        save_weights_only=True
    )

    wandb_logger = WandbLogger(
        save_dir=wandb_path,
        project=proj_name,
        name="DisBrain-Discrete-ProtoRVQ32-cosLoss-freezeG-warmup-clip-grad-fad[%depoch]" % (epoch_num),
    )

    import pickle
    with open('layout_sherlock.pkl', 'rb') as f:
        layout = pickle.load(f)

    facodecldm_ckpt_path = "path/to/log/Fine-Recon-Projector/checkpoint.ckpt"

    # Initialize the network
    net = BrainFACodec(modal='meg', sub_num=len(train_sub_list), d_t=320, num_patches=num_patches, hidden_dim=256,
                       layout=layout,
                       num_clusters=32,
                       facodecldm_ckpt_path=facodecldm_ckpt_path,
                       use_detail=False,
                       val_sub_list=val_sub_list, test_sub_list=test_sub_list,
                       log_dir=log_path, val_wav_dir=val_wav_dir, test_wav_dir=test_wav_dir)
                       # val_trans_path=val_trans_path, test_trans_path=test_trans_path)

    trainer = Trainer(
        accelerator="gpu",
        devices=1,
        logger=wandb_logger,
        max_epochs=epoch_num,
        num_sanity_val_steps=0,  # 先运行一遍val
        # limit_val_batches=10,
        check_val_every_n_epoch=2,
        # val_check_interval=2736,
        strategy=DDPStrategy(find_unused_parameters=True),
        callbacks=[checkpoint_callback],
        gradient_clip_val=1.0
    )

    trainer.fit(net, train_loader, val_loader)
    trainer.test(net, test_loader, ckpt_path="best")
    wandb.finish()


if __name__ == '__main__':
    train_meg_gw()
    # test_meg_gw()
    # train_meg_sherlock()