train_seq.py

import multiprocessing
from functools import partial
from pathlib import Path
from args import FineDance_parse_train_opt, save_arguments_to_yaml
import sys

import torch
import torch.nn.functional as F
import wandb
from accelerate import Accelerator, DistributedDataParallelKwargs
from accelerate.state import AcceleratorState
from torch.utils.data import DataLoader
from tqdm import tqdm

from dataset.FineDance_dataset import FineDance_Smpl
from dataset.preprocess import increment_path
from model.adan import Adan
from model.diffusion import GaussianDiffusion
from model.model import DanceDecoder
from vis import SMPLX_Skeleton, SMPLSkeleton
from train_w_test import test
import os

from dataset.humanml3d_dataset import HumanML3D
from torch.utils.data.dataloader import default_collate

from model.pseudo_pair_training import AlignedBanks, FeatureStore
from typing import Dict, Any, Optional, Tuple
import torch.nn.functional  as Fnn  # 避免与上面 F 混淆（不过两者相同）

import pdb



def wrap(x):
    return {f"module.{key}": value for key, value in x.items()}
def maybe_wrap(x, num):
    return x if num == 1 else wrap(x)

def collate_skip_none(batch):
    batch = [b for b in batch if b is not None]
    return default_collate(batch) if batch else None

class EDGE:
    def __init__(
        self,
        opt,
        feature_type,
        checkpoint_path="",
        normalizer=None,
        EMA=True,
        learning_rate=2e-4,
        weight_decay=0.02,
    ):
        self.opt = opt
        
        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
        self.accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])
        state = AcceleratorState()
        
        self.repr_dim = repr_dim = opt.nfeats
        feature_dim = 512 + 193
   
        self.horizon = horizon = opt.full_seq_len
        
        self.accelerator.wait_for_everyone()

        self.resume_num = 0
        self.normalizer = None
        
        model = DanceDecoder(
            nfeats=repr_dim,
            seq_len=horizon,
            latent_dim=512,
            ff_size=1024,
            num_layers=8,
            num_heads=8,
            dropout=0.1,
            cond_feature_dim=feature_dim,
            activation=F.gelu,
            use_rotary=True,
            finetune=opt.finetune,
            contrastive = opt.contrastive
        ) 
        
        
        if opt.nfeats == 139 or opt.nfeats == 135:
            smplx_fk = SMPLSkeleton(device=self.accelerator.device)
        else:
            smplx_fk = SMPLX_Skeleton(device=self.accelerator.device, batch=512000)
        
        diffusion = GaussianDiffusion(
            model,
            opt,
            horizon,
            repr_dim,
            smplx_model = smplx_fk,
            schedule="cosine",
            n_timestep=1000,
            predict_epsilon=False,
            loss_type="l2",
            use_p2=False,
            cond_drop_prob=0.25,
            guidance_weight=3,
            do_normalize = opt.do_normalize,
            lambda_prior_preserv = 1,
            lora_finetune=opt.finetune, 
        )

        print(
            "Model has {} parameters".format(sum(y.numel() for y in model.parameters()))
        )

        self.model = self.accelerator.prepare(model)
        diffusion.model = self.model 
        self.diffusion = diffusion.to(self.accelerator.device)
        
        self.smplx_fk = smplx_fk   
        self.trainable = [p for p in model.parameters() if p.requires_grad]
        n_total = sum(p.numel() for p in model.parameters())
        n_train = sum(p.numel() for p in model.parameters() if p.requires_grad)
        print("trainable ratio:", n_train / n_total)
        print("text_encoder requires_grad 统计：")
        for name, p in model.named_parameters():
            if p.requires_grad:
                print("  ✓", name)
                
                
        optim = Adan(self.trainable, lr=1e-3, weight_decay=1e-2)
        self.optim = self.accelerator.prepare(optim)

        if checkpoint_path != "":
            self.resume_num = int(os.path.basename(checkpoint_path).split("-")[1].split(".")[0])
            ckpt = torch.load(checkpoint_path, map_location="cpu")

            def _has_module_prefix(sd): 
                return any(k.startswith("module.") for k in sd.keys())

            def _align_prefix(sd, target_module):
                tkeys = list(target_module.state_dict().keys())
                if len(tkeys) == 0:
                    return sd
                tgt_has_module = tkeys[0].startswith("module.")
                src_has_module = _has_module_prefix(sd)
                if tgt_has_module and not src_has_module:
                    return { f"module.{k}": v for k, v in sd.items() }
                if (not tgt_has_module) and src_has_module:
                    return { (k[len("module."):] if k.startswith("module.") else k): v for k, v in sd.items() }
                return sd

            raw_sd = ckpt.get("raw_state_dict", None)
            ema_sd = ckpt.get("ema_state_dict", None)

            # 1) 先把 student（可训练的）恢复到 raw
            if raw_sd is not None:
                raw_sd = _align_prefix(raw_sd, self.diffusion.model)
                self.diffusion.model.load_state_dict(raw_sd, strict=False)
            else:
                # 兼容老 ckpt：退而求其次用 ema
                ema_sd_tmp = ckpt.get("model_state_dict") or ckpt
                ema_sd_tmp = _align_prefix(ema_sd_tmp, self.diffusion.model)
                self.diffusion.model.load_state_dict(ema_sd_tmp, strict=False)

            # 2) 再把 teacher 恢复到 ema（对比学习的关键）
            if hasattr(self.diffusion, "teacher") and isinstance(self.diffusion.teacher, torch.nn.Module):
                base_sd_for_teacher = ema_sd if ema_sd is not None else raw_sd
                if base_sd_for_teacher is not None:
                    base_sd_for_teacher = _align_prefix(base_sd_for_teacher, self.diffusion.teacher)
                    self.diffusion.teacher.load_state_dict(base_sd_for_teacher, strict=False)
                self.diffusion.teacher.eval()
                for p in self.diffusion.teacher.parameters():
                    p.requires_grad_(False)

            # 3) master_model（用于非对比/采样）——可用 ema 或 raw，同步即可
            if hasattr(self.diffusion, "master_model") and isinstance(self.diffusion.master_model, torch.nn.Module):
                mm_sd = ema_sd if ema_sd is not None else (raw_sd if raw_sd is not None else None)
                if mm_sd is not None:
                    mm_sd = _align_prefix(mm_sd, self.diffusion.master_model)
                    self.diffusion.master_model.load_state_dict(mm_sd, strict=False)

            # 4) 恢复优化器/调度器（避免 LR 重置引发的 loss 抖动）
            if hasattr(self, "optimizer") and self.optimizer is not None and ("optimizer" in ckpt):
                try:
                    self.optimizer.load_state_dict(ckpt["optimizer"])
                except Exception as e:
                    print(f"[warn] optimizer state not loaded: {e}")
            if hasattr(self, "lr_scheduler") and self.lr_scheduler is not None and ("lr_scheduler" in ckpt):
                try:
                    self.lr_scheduler.load_state_dict(ckpt["lr_scheduler"])
                except Exception as e:
                    print(f"[warn] lr_scheduler state not loaded: {e}")
            

    def eval(self):
        self.diffusion.eval()

    def train(self):
        self.diffusion.train()

    def prepare(self, objects):
        return self.accelerator.prepare(*objects)
    
    def train_loop(self, opt):

      
        num_cpus = multiprocessing.cpu_count()
        print("batchsize=:", opt.batch_size)
        train_set = HumanML3D(args=opt, split="train")
        hml_loader = DataLoader(
            train_set, 
            batch_size=opt.batch_size, 
            shuffle=True, 
            num_workers=min(int(num_cpus * 0.5), 8),  
            pin_memory=True,
            drop_last=True,
            collate_fn=collate_skip_none
            )
        
        print("train_dataset: FineDance_Dataset ")  
        train_dataset = FineDance_Smpl(
            args=opt,            # data/
            istrain=True,
        )
        fd_loader = DataLoader(
            train_dataset,
            batch_size=opt.batch_size,
            shuffle=True,
            num_workers=min(int(num_cpus * 0.5), 8),      # num_workers=min(int(num_cpus * 0.75), 32), 
            pin_memory=True,
            drop_last=True,
        )
        hml_loader, fd_loader = self.accelerator.prepare(hml_loader, fd_loader)
        
        # ===================== 加载 预构建的伪配对环境 =====================
        if not opt.contrastive:
            cache_path = os.path.join(opt.project, "pseudo_env_cache.pt")
            if not os.path.isfile(cache_path):
                if self.accelerator.is_main_process:
                    print(f"错误：缓存文件未找到: {cache_path}")
                self.accelerator.wait_for_everyone()
                raise FileNotFoundError(f"Pseudo-env cache not found at {cache_path}。先运行构建脚本生成缓存。")

            blob = torch.load(cache_path, map_location="cpu")
            # —— 强校验：只认新版本键 —— #
            required_keys = ["pairs_fd2hml", "pairs_hml2fd", "fd_names", "hml_names"]
            for k in required_keys:
                assert k in blob, f"[pseudo-env] 缺少必要键: {k}"

            pairs_fd2hml = list(blob["pairs_fd2hml"])
            pairs_hml2fd = list(blob["pairs_hml2fd"])
            fd_names     = list(blob["fd_names"])
            hml_names    = list(blob["hml_names"])
            # 文本映射优先来自 hml_text_map；若无则由 hml_texts 构造
            hml_text_map = blob.get("hml_text_map", None)
            if hml_text_map is None:
                hml_texts = list(blob.get("hml_texts", []))
                assert len(hml_texts) == len(hml_names), "[pseudo-env] 既无 hml_text_map 也无等长 hml_texts"
                hml_text_map = { str(hml_names[i]): str(hml_texts[i]) for i in range(len(hml_names)) }

            # 可选附件（若存在就带上，便于训练侧直接查）
            fd_music_feats_map = blob.get("fd_music_feats_map", {})
            fd_text_raw_map    = blob.get("fd_text_raw_map",  {})

            # 组装查询用 dict（新版本只用这两张表）
            fd2hml = {}
            for p in pairs_fd2hml:
                fd_name  = str(p["fd_name"])
                hml_name = str(p["hml_name"])
                fd2hml[fd_name] = {
                    "hml_name":  hml_name,
                    "sim":       float(p.get("sim", 0.0)),
                    "hml_text":  str(p.get("hml_text", hml_text_map.get(hml_name, ""))),
                    "music_feats": p.get("music_feats", None),  # 若构建时已带
                    "text_raw":    p.get("text_raw",    None),
                }

            hml2fd = {}
            for p in pairs_hml2fd:
                hml_name = str(p["hml_name"])
                fd_name  = str(p["fd_name"])
                hml2fd[hml_name] = {
                    "fd_name":  fd_name,
                    "sim":      float(p.get("sim", 0.0)),
                    "hml_text": str(p.get("hml_text", hml_text_map.get(hml_name, ""))),
                    "music_feats": p.get("music_feats", None),
                    "text_raw":    p.get("text_raw",    None),
                }

            motion_bank = {
                "fd2hml": fd2hml,
                "hml2fd": hml2fd,
                "hml_text_map": hml_text_map,
                # 附件映射（可空）
                "fd_music_feats_map": fd_music_feats_map,
                "fd_text_raw_map":    fd_text_raw_map,
                "meta": {"source_file": cache_path, "format": "new_only"},
            }

            # 注入 diffusion（仅保留新接口/直挂属性）
            injected = False
            if hasattr(self.diffusion, "set_motion_pair_bank"):
                self.diffusion.set_motion_pair_bank(motion_bank); injected = True
            elif hasattr(self.diffusion, "set_pseudo_env"):
                # 新版 set_pseudo_env 支持 motion_bank 参数
                try:
                    self.diffusion.set_pseudo_env(None, None, motion_bank=motion_bank); injected = True
                except TypeError:
                    injected = False
            if not injected:
                setattr(self.diffusion, "motion_pair_bank", motion_bank)

            if self.accelerator.is_main_process:
                eg = next(iter(fd2hml.items()), None)
                eg_msg = f"{eg[0]} → {eg[1]['hml_name']} (sim={eg[1]['sim']:.3f})" if eg else "无有效配对"
                print(f"[pseudo-env] 新版本缓存就绪 | FD→HML={len(fd2hml)} | HML→FD={len(hml2fd)} | 示例: {eg_msg}")
            self.accelerator.wait_for_everyone()
        # ============================================================================

        
        load_loop = (
            partial(tqdm, position=1, desc="Batch")
            if self.accelerator.is_main_process
            else lambda x: x
        )
        
        if self.accelerator.is_main_process:
            save_dir = str(increment_path(Path(opt.project) / opt.exp_name))
            opt.exp_name = save_dir.split("/")[-1]
            if opt.wandb:
                wandb.init(project=opt.wandb_pj_name, name=opt.exp_name, resume=False)
                wandb.save("params.yaml")
            save_dir = Path(save_dir)
            wdir = save_dir / "weights"
            wdir.mkdir(parents=True, exist_ok=True)
            yaml_path = os.path.join(wdir, 'parameters.yaml')
            save_arguments_to_yaml(opt, yaml_path)


        self.accelerator.wait_for_everyone()
        for epoch in range(1, opt.epochs + 1):
            print("epoch:", epoch+self.resume_num)
            avg_loss = 0
            avg_vloss = 0
            avg_fkloss = 0
            avg_footloss = 0
            avg_info_mus = 0.0
            avg_info_txt = 0.0

            # train
            self.train()
            if opt.finetune:
                fd_iter  = iter(fd_loader)
                hml_iter = iter(hml_loader)
                steps_per_epoch = len(fd_loader)

                for step in load_loop(range(steps_per_epoch)):
                    try:
                        fd_batch = next(fd_iter)
                    except StopIteration:
                        fd_iter = iter(fd_loader)
                        fd_batch = next(fd_iter)
                    try:
                        hml_batch = next(hml_iter)
                    except StopIteration:
                        hml_iter = iter(hml_loader)
                        hml_batch = next(hml_iter)

                    # HumanML3D
                    x, filename_hml, hd_text = hml_batch              # x: HML motion, hd_text: list[str]
                    # FineDance
                    motion_fd, cond, filename_fd, fd_text, text_raw = fd_batch        # motion_fd: FD motion, cond: FD music cond
                    
                    device = self.accelerator.device
                    x = x.to(device)
                    cond = cond.to(device)
                    motion_fd = motion_fd.to(device)
                    
                    # ===========================================================
                    if getattr(opt, "contrastive", False):
                        
                        # 1) 计算所有对比项（返回 dict）
                        outs = self.diffusion._contrastive_losses(
                            x_hml=x, cond_fd=cond, motion_fd=motion_fd, hd_text=hd_text
                        )

                        # 2) 线性 warmup（只给“跨域 motion 对齐项”用）
                        step = int(getattr(self.diffusion, 'global_step', 0))
                        warm_steps = int(getattr(self.diffusion, 'mm_align_warmup_steps', 0))
                        warm = float(min(step / max(1, warm_steps), 1.0)) if warm_steps > 0 else 1.0

                        # 3) 权重（带默认）
                        w_mus    = float(getattr(self.diffusion, 'contrastive_w_music', 1.0))
                        w_txt    = float(getattr(self.diffusion, 'contrastive_w_text',  1.0))
                        w_center = float(getattr(self.diffusion, 'w_mm_center', 0.05))
                        w_coral  = float(getattr(self.diffusion, 'w_mm_coral',  0.05))
                        w_mmd    = float(getattr(self.diffusion, 'w_mm_mmd',    0.00))
                        w_mnn    = float(getattr(self.diffusion, 'w_mm_mnn',    0.00))

                        # 4) 组装总 loss
                        total_loss = 0.0
                        if outs.get('L_mus') is not None:
                            total_loss = total_loss + w_mus * outs['L_mus']
                            try:   avg_info_mus += float(outs['L_mus'].detach().cpu().item())
                            except: pass
                        if outs.get('L_txt') is not None:
                            total_loss = total_loss + w_txt * outs['L_txt']
                            try:   avg_info_txt += float(outs['L_txt'].detach().cpu().item())
                            except: pass

                        # 跨域 Motion 的温和对齐（带 warmup）
                        if outs.get('L_mm_center') is not None:
                            total_loss = total_loss + warm * w_center * outs['L_mm_center']
                            try:   avg_mm_center += float(outs['L_mm_center'].detach().cpu().item())
                            except: pass
                        if outs.get('L_mm_coral') is not None:
                            total_loss = total_loss + warm * w_coral  * outs['L_mm_coral']
                            try:   avg_mm_coral  += float(outs['L_mm_coral'].detach().cpu().item())
                            except: pass
                        if (outs.get('L_mm_mmd') is not None) and (w_mmd > 0):
                            total_loss = total_loss + warm * w_mmd * outs['L_mm_mmd']
                            try:   avg_mm_mmd   += float(outs['L_mm_mmd'].detach().cpu().item())
                            except: pass
                        if (outs.get('L_mm_mnn') is not None) and (w_mnn > 0):
                            total_loss = total_loss + warm * w_mnn * outs['L_mm_mnn']
                            try:   avg_mm_mnn   += float(outs['L_mm_mnn'].detach().cpu().item())
                            except: pass

                        # 5) 反传与优化
                        self.optim.zero_grad(set_to_none=True)
                        self.accelerator.backward(total_loss)
                        self.accelerator.clip_grad_norm_(self.trainable, 1.0)
                        self.optim.step()

                        # 6) 分域入队（取代 _safe_update_queue）
                        if hasattr(self.diffusion, "_pending_queue_update"):
                            with torch.no_grad():
                                for k, domain in self.diffusion._pending_queue_update:  # [(k_fd,"fd"), (k_hml,"hml"), ...]
                                    self.diffusion._dequeue_and_enqueue(k, domain=domain)
                            self.diffusion._pending_queue_update = []

                        # 7) EMA：Teacher 跟随 Student
                        self.diffusion._ema_update_teacher()

                        # 8) 自增 step（如果 diffusion 里维护这个计数）
                        if hasattr(self.diffusion, 'global_step'):
                            self.diffusion.global_step += 1

                    else:
                        # ===== 常规联合损失（重建 + InfoNCE）=====
                        total_loss, (loss, v_loss, fk_loss, foot_loss) = self.diffusion(
                            x, cond, motion_fd, hd_text, fd_text, text_raw, names_hml=filename_hml, names_fd=filename_fd,  t_override=None
                        )
                        
                        self.optim.zero_grad()
                        self.accelerator.backward(total_loss)
                        self.accelerator.clip_grad_norm_(self.trainable, 1.0)
                        self.optim.step()

                        if self.accelerator.is_main_process:
                            avg_loss     += loss.detach().cpu().item()
                            avg_vloss    += v_loss.detach().cpu().item()
                            avg_fkloss   += fk_loss.detach().cpu().item()
                            avg_footloss += foot_loss.detach().cpu().item()

                        # 这里保留 master_model 的 EMA（供采样阶段用）
                        if step % opt.ema_interval == 0:
                            self.diffusion.ema.update_model_average(
                                self.diffusion.master_model, self.diffusion.model
                            )
            else:
                for step, (x, cond, filename, text) in enumerate(
                    load_loop(fd_loader)
                ):
                    if opt.nfeats == 139 or opt.nfeats==135:
                        x = x[:, :, :139]
                    total_loss, (loss, v_loss, fk_loss, foot_loss)  = self.diffusion(
                        x, cond, text, t_override=None
                    )
                    self.optim.zero_grad()
                    self.accelerator.backward(total_loss)
                    self.optim.step()
                    # contrastive accum
                    if getattr(self.diffusion, 'last_contrastive', None):
                        Lm = self.diffusion.last_contrastive.get('L_mus', None)
                        Lt = self.diffusion.last_contrastive.get('L_txt', None)
                        if Lm is not None:
                            avg_info_mus += float(Lm.detach().cpu().item())
                        if Lt is not None:
                            avg_info_txt += float(Lt.detach().cpu().item())
                    # ema update and train loss update only on main
                    if self.accelerator.is_main_process:
                        avg_loss += loss.detach().cpu().numpy()
                        avg_vloss += v_loss.detach().cpu().numpy()
                        avg_fkloss += fk_loss.detach().cpu().numpy()
                        avg_footloss += foot_loss.detach().cpu().numpy()
                    if step % opt.ema_interval == 0:
                        self.diffusion.ema.update_model_average(
                            self.diffusion.master_model, self.diffusion.model
                        )
                        
            #-----------------------------------------------------------------------------------------------------------
            # test
            # Save model
        
            if ((epoch+self.resume_num) % opt.save_interval) == 0  or epoch<=1 :
            #if ((epoch+self.resume_num) % opt.save_interval) == 0 :
                self.accelerator.wait_for_everyone()
                self.eval()    
                if self.accelerator.is_main_process:
                    if getattr(opt, "contrastive", False):
                        # —— 计算 warmup 系数（只给跨域 motion 对齐项用）
                        step = iter_idx if 'iter_idx' in locals() else int(getattr(self.diffusion, 'global_step', 0))
                        warm_steps = int(getattr(self.diffusion, 'mm_align_warmup_steps', 0))
                        warm = float(min(step / max(1, warm_steps), 1.0)) if warm_steps > 0 else 1.0

                        # —— 累计量（若没定义则初始化）
                        if 'avg_mm_center' not in locals(): avg_mm_center = 0.0
                        if 'avg_mm_coral'  not in locals(): avg_mm_coral  = 0.0
                        if 'avg_mm_mmd'    not in locals(): avg_mm_mmd    = 0.0
                        if 'avg_mm_mnn'    not in locals(): avg_mm_mnn    = 0.0

                        # outs 来自前面 self.diffusion._contrastive_losses(...)
                        def _tofloat(x):
                            return float(x.detach().cpu().item()) if (x is not None) else None

                        # 累计（只在有值时加）
                        v_mus    = _tofloat(outs.get('L_mus'))
                        v_txt    = _tofloat(outs.get('L_txt'))
                        v_center = _tofloat(outs.get('L_mm_center'))
                        v_coral  = _tofloat(outs.get('L_mm_coral'))
                        v_mmd    = _tofloat(outs.get('L_mm_mmd'))
                        v_mnn    = _tofloat(outs.get('L_mm_mnn'))

                        if v_mus   is not None: avg_info_mus += v_mus
                        if v_txt   is not None: avg_info_txt += v_txt
                        if v_center is not None: avg_mm_center += v_center
                        if v_coral  is not None: avg_mm_coral  += v_coral
                        if v_mmd    is not None: avg_mm_mmd    += v_mmd
                        if v_mnn    is not None: avg_mm_mnn    += v_mnn

                        # —— 打印（按 finetune 决定是否除以 steps_per_epoch）
                        denom = max(1, steps_per_epoch) if opt.finetune else 1
                        log_dict = {
                            "InfoNCE Music↔Motion": (avg_info_mus / denom),
                            "InfoNCE Text↔Motion":  (avg_info_txt / denom),
                            "Center (FD↔HML)":      (avg_mm_center / denom),
                            "CORAL  (FD↔HML)":      (avg_mm_coral  / denom),
                        }

                        # 可选项：只有打开对应 loss 才打印，避免 None
                        if v_mmd is not None:
                            log_dict["MMD    (FD↔HML)"] = (avg_mm_mmd / denom)
                        if v_mnn is not None:
                            log_dict["MNN-InfoNCE(FD↔HML)"] = (avg_mm_mnn / denom)

                        # 辅助观测：当前 warmup 系数
                        log_dict["warmup(mm-align)"] = warm

                        print(log_dict)
                    else:    
                        avg_loss /= step+1
                        avg_vloss /= step+1
                        avg_fkloss /= step+1
                        avg_footloss /= step+1
                        #test the fid : delete it later
                        if opt.test_fid:
                            test(self,opt,epoch+self.resume_num)
                        log_dict = {
                        "Train Loss": avg_loss,
                        "V Loss": avg_vloss,
                        "FK Loss": avg_fkloss,
                        "Foot Loss": avg_footloss,
                        }
                        print(log_dict)
                    if opt.wandb:
                        wandb.log(log_dict)
                    
                    
                    if (self.accelerator.is_main_process 
                        and ((epoch + self.resume_num) % (opt.save_interval * 5) == 0 or epoch == 1)):
                        # 同步 master/teacher
                        if hasattr(self.diffusion, "ema"):
                            self.diffusion.ema.update_model_average(self.diffusion.master_model, self.diffusion.model)

                        to_save = {
                            "epoch": epoch + self.resume_num,
                            # 对比学习：保存 teacher（EMA）和 student（raw）
                            "ema_state_dict": (
                                self.diffusion.teacher.state_dict() if opt.contrastive
                                else self.diffusion.master_model.state_dict()
                            ),
                            "raw_state_dict": self.diffusion.model.state_dict(),
                            "opt": vars(opt),
                        }
                        # 强烈建议：一并保存优化器/调度器
                        if hasattr(self, "optimizer") and self.optimizer is not None:
                            to_save["optimizer"] = self.optimizer.state_dict()
                        if hasattr(self, "lr_scheduler") and self.lr_scheduler is not None:
                            to_save["lr_scheduler"] = self.lr_scheduler.state_dict()

                        torch.save(to_save, os.path.join(wdir, f"train-{epoch+self.resume_num}.pt"))
                                        
                    
            #-----------------------------------------------------------------------------------------------------------
            
            
        if self.accelerator.is_main_process:
            if opt.wandb:
                wandb.run.finish()
            else:
                pass
          

    def render_sample(
        self, data_tuple, label, render_dir, render_count=-1, mode='normal', fk_out=None, render=True,vis_attn_ma = False
    ):
        _, cond, wavname, text,cls_t = data_tuple
        if render_count < 0:
            render_count = len(cond)
        shape = (render_count, self.horizon, self.repr_dim)
        cond = cond.to(self.accelerator.device).float()
        cls_t = cls_t.to(self.accelerator.device).float()
        motion= []
        text = [t[0] for t in text]


        self.diffusion.render_sample(
            shape,
            motion,
            cond[:render_count],
            text[:render_count],
            cls_t[:render_count],
            self.normalizer,
            label,
            render_dir,
            name=wavname[:render_count],
            sound=False,
            mode=mode,           
            fk_out=fk_out,
            render=render,
            vis_attn_ma = vis_attn_ma
        )
        
        
    def render_sample_inpaint(
        self, data_tuple, label, render_dir, render_count=-1, mode='inpaint', fk_out=None, render=True,
    ):
        motion, cond, wavname, text = data_tuple
        assert len(cond.shape) == 3
        if render_count < 0:
            render_count = len(cond)
        shape = (render_count, self.horizon, self.repr_dim)
        cond = cond.to(self.accelerator.device).float()
        text = text.to(self.accelerator.device).float()
        motion = motion.to(self.accelerator.device).float()


        self.diffusion.render_sample(
            shape,
            motion,
            cond[:render_count],
            text[:render_count],
            self.normalizer,
            label,
            render_dir,
            name=wavname[:render_count],
            sound=False,
            mode=mode,           
            fk_out=fk_out,
            render=render
        )


def train(opt):
    model = EDGE(opt, opt.feature_type,opt.checkpoint)
    model.train_loop(opt)
    
if __name__ == "__main__":
    
    opt = FineDance_parse_train_opt()
    command = ' '.join(sys.argv)
    import time
    import random
    time.sleep(random.randint(0, 2))
    if not os.path.exists(os.path.join(opt.project, opt.exp_name)):
        os.makedirs(os.path.join(opt.project, opt.exp_name), exist_ok=False)
    with open(os.path.join(opt.project, opt.exp_name, 'command.txt'), 'w') as f:
        f.write(command)

    yaml_path = os.path.join(opt.project, opt.exp_name, 'parameters.yaml')
    save_arguments_to_yaml(opt, yaml_path)
            
    train(opt)