dac-vae/train.py

from comet_ml import Experiment
import argparse
import os
import time
import typing
from dataclasses import dataclass
from datetime import timedelta
from typing import Dict, Union

import numpy as np
import soundfile as sf
import torch
import yaml

from model import Discriminator
from model import DACVAE as VAE
from loss import (GANLoss, L1Loss, MelSpectrogramLoss,
                         MultiScaleSTFTLoss, kl_loss)
from torch import nn
from torch.distributed import destroy_process_group, init_process_group
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim import Adam, AdamW
from torch.optim.lr_scheduler import ConstantLR, LinearLR, SequentialLR
from torch.utils.data.distributed import DistributedSampler

from audiotools import AudioSignal
from audiotools.core import util
from audiotools.data import transforms
from audiotools.data.datasets import AudioDataset, AudioLoader, ConcatDataset
from audiotools.ml.decorators import Tracker, timer, when


def ddp_setup():
    print("Setting up DDP")
    init_process_group(backend="nccl", timeout=timedelta(seconds=7200))
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))


def build_transform(
    augment_prob=1.0,
    preprocess=["Identity"],
    augment=["Identity"],
    postprocess=["Identity", "RescaleAudio", "ShiftPhase"],
):
    to_tfm = lambda l: [getattr(transforms, x)() for x in l]
    preprocess = transforms.Compose(*to_tfm(preprocess), name="preprocess")
    augment = transforms.Compose(*to_tfm(augment), name="augment", prob=augment_prob)
    postprocess = transforms.Compose(*to_tfm(postprocess), name="postprocess")
    return transforms.Compose(preprocess, augment, postprocess)


def build_dataset(sample_rate, folders=None, **kwargs):
    if folders is None:
        folders = {}
    datasets = []
    for _, v in folders.items():
        loader = AudioLoader(sources=v)
        transform = build_transform()
        dataset = AudioDataset(
            loader, sample_rate, num_channels=2, transform=transform, **kwargs
        )
        datasets.append(dataset)
    dataset = ConcatDataset(datasets)
    dataset.transform = transform
    return dataset


@dataclass
class State:
    generator: DDP
    optimizer_g: Union[AdamW, Adam]
    scheduler_g: torch.optim.lr_scheduler._LRScheduler

    discriminator: DDP
    optimizer_d: Union[AdamW, Adam]
    scheduler_d: torch.optim.lr_scheduler._LRScheduler

    stft_loss: MultiScaleSTFTLoss
    mel_loss: MelSpectrogramLoss
    gan_loss: GANLoss
    waveform_loss: L1Loss

    train_dataset: AudioDataset
    val_dataset: AudioDataset

    tracker: Tracker
    lambdas: Dict[str, float]

    # ema: EMA  # Add EMA to State


class ResumableDistributedSampler(DistributedSampler):  # pragma: no cover
    """Distributed sampler that can be resumed from a given start index."""

    def __init__(self, dataset, start_idx: int = 0, **kwargs):
        super().__init__(dataset, **kwargs)
        # Start index, allows to resume an experiment at the index it was
        self.start_idx = start_idx // self.num_replicas if start_idx is not None else 0

    def __iter__(self):
        for i, idx in enumerate(super().__iter__()):
            if i >= self.start_idx:
                yield idx
        self.start_idx = 0  # set the index back to 0 so for the next epoch


def prepare_dataloader(
    dataset: AudioDataset,
    world_size: int,
    local_rank: int,
    start_idx: int = 0,
    shuffle: bool = True,
    **kwargs,
):
    sampler = ResumableDistributedSampler(
        dataset,
        start_idx,
        num_replicas=world_size,
        rank=local_rank,
        shuffle=shuffle,
    )

    sampler = None
    if start_idx > 0:
        # Create a simple resumable sampler
        indices = list(range(start_idx, len(dataset))) + list(range(start_idx))
        sampler = torch.utils.data.SubsetRandomSampler(indices)

    if "num_workers" in kwargs:
        kwargs["num_workers"] = max(kwargs["num_workers"] // world_size, 1)
    kwargs["batch_size"] = max(kwargs["batch_size"] // world_size, 1)
    dataloader = torch.utils.data.DataLoader(dataset, sampler=sampler, **kwargs)
    dataloader = torch.utils.data.DataLoader(
        dataset,
        sampler=sampler,
        shuffle=(sampler is None),  # Only shuffle if no sampler
        num_workers=24,  # Can use more workers since no distribution
        pin_memory=True,
        persistent_workers=True,
        prefetch_factor=8,  # Can be higher for single GPU
        drop_last=True,
        **kwargs
    )
    return dataloader


class Trainer:
    def __init__(self, args) -> None:
        self.local_rank = int(os.environ["LOCAL_RANK"])
        self.global_rank = int(os.environ["RANK"])
        self.world_size = int(os.environ["WORLD_SIZE"])
        torch.backends.cudnn.benchmark = True
        torch.cuda.set_device(self.local_rank)
        torch.cuda.empty_cache()

        configs = yaml.safe_load(open(args.config_path, "r"))
        print("configs: ", configs)
        self.configs = configs

        self.gan_start_step = configs.get("gan_start_step", 0)

        self.num_iters = configs.get("num_iters", 250000)

        self.generator = VAE(**configs["vae"])

        self.discriminator = Discriminator(**configs["discriminator"])

        total_steps = configs["num_samples"] // configs["batch_size"]

        if configs["optimizer"]["scheduler"] == "linearlr":
            self.optimizer_g, self.scheduler_g = self.get_scheduler(
                self.generator, total_steps, configs["optimizer"]
            )
        else:
            self.optimizer_g, self.scheduler_g = self.get_constant_scheduler(
                self.generator, total_steps
            )

        if configs["disc_optimizer"]["scheduler"] == "constantlr":
            self.optimizer_d, self.scheduler_d = self.get_constant_scheduler(
                self.discriminator, total_steps
            )
        else:
            self.optimizer_d, self.scheduler_d = self.get_scheduler(
                self.discriminator, total_steps, configs["disc_optimizer"]
            )

        save_path = args.save_path
        os.makedirs(save_path, exist_ok=True)
        self.save_path = save_path

        if self.local_rank == 0:
            print(f"Rank {self.local_rank}: Initializing Comet.ml")
            experiment = Experiment(
                api_key=os.environ.get(
                    "COMET_API_KEY"
                ),  # Set COMET_API_KEY in your environment
                project_name="DACVAE",
                workspace=os.environ.get("COMET_WORKSPACE"),  # Optional: Set workspace
                # experiment_key=args.run_id,  # Use run_id as experiment key
            )
            experiment.log_parameters(configs)  # Log configuration
            writer = experiment
        else:
            writer = None

        print(f"Rank {self.local_rank}: Setting up tracker")
        self.tracker = Tracker(
            writer=writer, log_file=f"{save_path}/log.txt", rank=self.local_rank
        )
        self.val_idx = configs.get("val_idx", [0, 1, 2, 3, 4, 5, 6, 7])
        self.save_iters = configs.get("save_iters", 1000)
        self.sample_freq = configs.get("sample_freq", 10000)
        self.valid_freq = configs.get("valid_freq", 1000)

        self.tracker.print(self.generator)
        self.tracker.print(self.discriminator)

        self.waveform_loss = L1Loss()
        self.stft_loss = MultiScaleSTFTLoss(**configs["MultiScaleSTFTLoss"])
        self.mel_loss = MelSpectrogramLoss(**configs["MelSpectrogramLoss"])

        print(f"{self.global_rank} Loading datasets...")
        sample_rate = configs["vae"]["sample_rate"]
        train_folders = {k: v for k, v in configs.get("train_folders", {}).items()}
        val_folders = {k: v for k, v in configs.get("val_folders", {}).items()}
        self.batch_size = configs["batch_size"]
        self.val_batch_size = configs["val_batch_size"]
        self.num_workers = configs["num_workers"]

        print(f"Rank {self.local_rank}: Validating train dataset")
        self.train_dataset = build_dataset(
            sample_rate, train_folders, **configs["train"]
        )
        print(f"Rank {self.local_rank}: Validating val dataset")
        self.val_dataset = build_dataset(sample_rate, val_folders, **configs["val"])

        self.lambdas = configs["lambdas"]

        if args.resume:
            checkpoint_dir = os.path.join(args.save_path, args.tag)
            self.resume_from_checkpoint(checkpoint_dir)

        self.gan_loss = GANLoss(self.discriminator)
        print("self.tracker.step: ", self.tracker.step)

        self.generator = self.generator.to(self.local_rank)
        self.discriminator = self.discriminator.to(self.local_rank)
        #
        self.generator = nn.SyncBatchNorm.convert_sync_batchnorm(self.generator)
        self.discriminator = nn.SyncBatchNorm.convert_sync_batchnorm(self.discriminator)

        # Wrap models with DDP
        self.generator = DDP(self.generator, device_ids=[self.local_rank])
        self.discriminator = DDP(self.discriminator, device_ids=[self.local_rank])

        # ema_decay = self.configs.get("ema_decay", 0.999)  # Add to your config YAML or set default
        # self.ema = EMA(self.unwrap(self.generator), decay=ema_decay, device=self.local_rank)

        self.state = State(
            generator=self.generator,
            optimizer_g=self.optimizer_g,
            scheduler_g=self.scheduler_g,
            discriminator=self.discriminator,
            optimizer_d=self.optimizer_d,
            scheduler_d=self.scheduler_d,
            tracker=self.tracker,
            train_dataset=self.train_dataset,
            val_dataset=self.val_dataset,
            stft_loss=self.stft_loss.to(self.local_rank),
            mel_loss=self.mel_loss.to(self.local_rank),
            gan_loss=self.gan_loss.to(self.local_rank),
            waveform_loss=self.waveform_loss.to(self.local_rank),
            lambdas=self.lambdas,
            # ema=self.ema,  # Add EMA to state
        )
        train_dataloader = prepare_dataloader(
            self.train_dataset,
            world_size=self.world_size,
            local_rank=self.local_rank,
            start_idx=self.state.tracker.step,  # Use step directly
            batch_size=self.batch_size,
            collate_fn=self.state.train_dataset.collate,
        )

        self.len_train = len(train_dataloader)

        self.train_dataloader = self.get_infinite_loader(train_dataloader)

        if self.global_rank == 0:
            self.val_dataloader = prepare_dataloader(
                self.state.val_dataset,
                world_size=1,
                local_rank=0,
                start_idx=0,
                shuffle=False,
                batch_size=self.val_batch_size,
                collate_fn=self.state.val_dataset.collate,
            )

        self.seed = 0
        self.val_real_audio = []
        self.val_gen_audio = []
        self.initial_norm = configs.get("initial_norm", float("inf"))
        self.max_norm = configs.get("max_norm", float("inf"))
        self.initial_norm_d = configs.get("initial_norm_d", float("inf"))
        self.max_norm_d = configs.get("max_norm_d", float("inf"))

        self.init_logs_penalty = self.state.lambdas["logs_penalty"]
        self.init_lipschitz_penalty = self.state.lambdas["lipschitz_penalty"]
        self.kl_max_beta = self.state.lambdas["kl/loss"]
        self.hold_base_steps = configs.get("hold_base_steps", 200000)

    def get_scheduler(self, model, total_steps, configs):
        warmup_steps = configs.get("warmup_steps", 0)
        if configs["type"] == "Adamw":
            optimizer = AdamW(
                model.parameters(),
                lr=configs["lr"],
                weight_decay=configs["weight_decay"],
            )
        else:
            optimizer = Adam(
                model.parameters(),
                lr=configs["lr"],
                weight_decay=configs["weight_decay"],
            )

        # Warmup from near-zero to max_lr
        warmup = LinearLR(
            optimizer,
            start_factor=1e-9,
            end_factor=1.0,  # Go up to max_lr
            total_iters=warmup_steps,
        )
        remaining_iters = total_steps - warmup_steps
        constant = ConstantLR(
            optimizer,
            factor=1.0,  # Keep the learning rate constant at max_lr
            total_iters=remaining_iters,
        )

        scheduler = SequentialLR(
            optimizer, schedulers=[warmup, constant], milestones=[warmup_steps]
        )
        return optimizer, scheduler

    def get_constant_scheduler(self, model, total_steps):
        if self.configs["optimizer"]["type"] == "adamw":
            optimizer = AdamW(
                model.parameters(),
                lr=self.configs["optimizer"]["lr"],
                weight_decay=self.configs["optimizer"]["weight_decay"],
            )
        else:
            optimizer = Adam(
                model.parameters(),
                lr=self.configs["optimizer"]["lr"],
                weight_decay=self.configs["optimizer"]["weight_decay"],
            )
        scheduler = ConstantLR(
            optimizer,
            factor=1.0,  # Keep the learning rate constant at max_lr
            total_iters=total_steps,
        )
        return optimizer, scheduler

    def get_infinite_loader(self, dataset):
        print(
            f"Rank {torch.distributed.get_rank() if torch.distributed.is_initialized() else 0}: Starting infinite loader"
        )
        # Skip iterations if resuming
        iterator = iter(dataset)
        steps_to_skip = self.state.tracker.step
        while True:
            try:
                batch = next(iterator)
                if batch is None:
                    print(f"Rank {torch.distributed.get_rank()}: Skipping None batch")
                    continue
                yield batch
            except StopIteration:
                iterator = iter(dataset)  # Reset iterator at the end of the dataset

    def log_grad_norms(self, output, norm_threshold=1.0):
        """
        Log gradient norms for key DACVAE components to aid debugging.
        Tracks pre-clipping norms for encoder, decoder, and selected blocks.

        Args:
            output (dict): Dictionary to store gradient norm logs.
            norm_threshold (float): Log norms above this threshold to reduce noise.
        """
        # Initialize dictionaries for norms
        submodule_norms = {
            "en_conv_post": 0.0,
            "de_conv_pre": 0.0,
            "encoder_initial_conv": 0.0,
            "encoder_final_conv": 0.0,
            "encoder_snake1d_alpha": 0.0,
            "decoder_initial_conv": 0.0,
            "decoder_final_conv": 0.0,
            "decoder_snake1d_alpha": 0.0,
        }
        norm_values = []  # For distributional statistics

        # Initialize norms for a few representative blocks (e.g., first and last)
        num_enc_blocks = len(self.state.generator.module.encoder_rates)
        num_dec_blocks = len(self.state.generator.module.decoder_rates)
        for i in [0, num_enc_blocks - 1]:  # First and last encoder blocks
            submodule_norms.update(
                {
                    f"encoder_block_{i}": 0.0,
                    f"encoder_block_{i}_snake1d": 0.0,
                    f"encoder_block_{i}_conv1d": 0.0,
                }
            )
        for i in [0, num_dec_blocks - 1]:  # First and last decoder blocks
            submodule_norms.update(
                {
                    f"decoder_block_{i}": 0.0,
                    f"decoder_block_{i}_snake1d": 0.0,
                    f"decoder_block_{i}_conv_transpose": 0.0,
                }
            )

        # Calculate indices for final layers
        enc_final_conv_idx = num_enc_blocks + 2
        dec_final_conv_idx = num_dec_blocks * 2 + 1

        # Iterate through parameters
        for name, param in self.state.generator.named_parameters():
            if param.grad is not None:
                norm = param.grad.norm().item()
                norm_values.append(norm)

                # DACVAE layers
                if "en_conv_post" in name:
                    submodule_norms["en_conv_post"] += norm**2
                elif "de_conv_pre" in name:
                    submodule_norms["de_conv_pre"] += norm**2

                # Encoder components
                if "encoder.block.0" in name:
                    submodule_norms["encoder_initial_conv"] += norm**2
                elif f"encoder.block.{enc_final_conv_idx}" in name:
                    submodule_norms["encoder_final_conv"] += norm**2
                elif "encoder" in name and "alpha" in name:
                    submodule_norms["encoder_snake1d_alpha"] += norm**2
                for i in [0, num_enc_blocks - 1]:
                    block_idx = i + 1
                    if f"encoder.block.{block_idx}" in name:
                        submodule_norms[f"encoder_block_{i}"] += norm**2
                        if "block.3" in name:  # Snake1d
                            submodule_norms[f"encoder_block_{i}_snake1d"] += norm**2
                        elif "block.4" in name:  # WNConv1d
                            submodule_norms[f"encoder_block_{i}_conv1d"] += norm**2

                # Decoder components
                if "decoder.model.0" in name:
                    submodule_norms["decoder_initial_conv"] += norm**2
                elif f"decoder.model.{dec_final_conv_idx}" in name:
                    submodule_norms["decoder_final_conv"] += norm**2
                elif "decoder" in name and "alpha" in name:
                    submodule_norms["decoder_snake1d_alpha"] += norm**2
                for i in [0, num_dec_blocks - 1]:
                    block_idx = i * 2 + 1
                    if f"decoder.model.{block_idx}" in name:
                        submodule_norms[f"decoder_block_{i}"] += norm**2
                        if "block.0" in name:  # Snake1d
                            submodule_norms[f"decoder_block_{i}_snake1d"] += norm**2
                        elif "block.1" in name:  # WNConvTranspose1d
                            submodule_norms[f"decoder_block_{i}_conv_transpose"] += (
                                norm**2
                            )

        # Compute square root of summed norms and log if above threshold
        for key in submodule_norms:
            norm = submodule_norms[key] ** 0.5
            if norm > norm_threshold:
                output[f"grad_norm/{key}"] = norm

        # Log pre-clipping norm statistics
        if norm_values:
            output["grad_norm/pre_clip_max"] = max(norm_values)
            output["grad_norm/pre_clip_mean"] = sum(norm_values) / len(norm_values)
            output["grad_norm/pre_clip_95th_percentile"] = (
                torch.tensor(norm_values).quantile(0.95).item()
            )

    def compute_lipschitz_penalty(self, lambda_lip=0.01):
        penalty = 0.0
        for name, param in self.state.generator.named_parameters():
            if (
                ("decoder" in name or "de_conv_pre" in name)
                and param.grad is not None
                and "weight" in name
            ):
                grad_norm = param.grad.norm(2)
                penalty += grad_norm**2
        return lambda_lip * penalty

    def compute_gradient_penalty(self, recons, z):
        # Compute gradients of decoder output w.r.t. latents
        grads = torch.autograd.grad(
            outputs=recons,
            inputs=z,
            grad_outputs=torch.ones_like(recons),
            create_graph=True,
            retain_graph=True,
        )[0]
        grad_norm = grads.norm(2, dim=[1, 2]).mean()
        return 0.1 * grad_norm  # Weight for penalty

    def cosine_decay_with_warmup(
        self,
        cur_step,
        base_value,
        total_steps,
        final_value,
        warmup_value=0.0,
        warmup_steps=0,
        hold_base_steps=0,
    ):
        """Cosine schedule with warmup, adapted from R3GAN."""
        # Ensure cur_step is a tensor
        cur_step = torch.tensor(cur_step, dtype=torch.float32)

        # Compute decay term
        denom = float(total_steps - warmup_steps - hold_base_steps)
        if denom <= 0:
            raise ValueError(
                "total_steps must be greater than warmup_steps + hold_base_steps"
            )
        phase = torch.pi * (cur_step - warmup_steps - hold_base_steps) / denom
        decay = 0.5 * (1 + torch.cos(phase))

        # Compute current value
        cur_value = base_value + (1 - decay) * (final_value - base_value)

        # Apply hold_base_steps condition
        if hold_base_steps > 0:
            cur_value = torch.where(
                cur_step > warmup_steps + hold_base_steps,
                cur_value,
                torch.tensor(base_value, dtype=torch.float32),
            )

        # Apply warmup_steps condition
        if warmup_steps > 0:
            slope = (base_value - warmup_value) / warmup_steps
            warmup_v = slope * cur_step + warmup_value
            cur_value = torch.where(cur_step < warmup_steps, warmup_v, cur_value)

        # Apply total_steps cap
        cur_value = torch.where(
            cur_step > total_steps,
            torch.tensor(final_value, dtype=torch.float32),
            cur_value,
        )

        return cur_value.item()  # Return as float

    def smooth_increase(
        self,
        step: int,
        initial_beta: float = 0.01,
        final_beta: float = 0.0,
        total_steps: int = 50000,
    ) -> float:
        """Compute a linear decrease for beta."""
        progress = min(step / total_steps, 1.0)
        beta = initial_beta + progress * (final_beta - initial_beta)
        return beta

    @timer()
    def train_loop(self, batch):
        print(f"Rank {self.local_rank}: Starting train_loop")

        self.max_gen_norm = self.cosine_decay_with_warmup(
            cur_step=self.tracker.step,
            base_value=self.initial_norm,  # e.g., 100
            total_steps=self.num_iters,  # e.g., 250000
            final_value=self.max_norm,
            warmup_value=self.initial_norm,
            warmup_steps=0,
            hold_base_steps=self.hold_base_steps,
        )

        self.max_d_norm = self.cosine_decay_with_warmup(
            cur_step=self.tracker.step,
            base_value=self.initial_norm_d,  # e.g., 100
            total_steps=self.num_iters,  # e.g., 250000
            final_value=self.max_norm_d,
            warmup_value=self.initial_norm_d,
            warmup_steps=0,
            hold_base_steps=self.hold_base_steps,
        )

        self.state.generator.train()
        if self.tracker.step >= self.gan_start_step:
            self.state.discriminator.train()
            print(
                f"Rank {self.local_rank}: Discriminator training mode: {self.state.discriminator.training}"
            )
        output = {}
        output = {}
        timing_logs = {}

        output["max_gen_norm"] = self.max_gen_norm
        output["max_d_norm"] = self.max_d_norm

        train_loop_start = time.time()

        # Batch preparation
        batch_prepare_start = time.time()
        batch = util.prepare_batch(batch, self.local_rank)
        timing_logs["batch_prepare"] = time.time() - batch_prepare_start

        # Data transformation
        transform_start = time.time()
        with torch.no_grad():
            signal = self.train_dataset.transform(
                batch["signal"].clone(), **batch["transform_args"]
            )
            signal.audio_data = torch.clamp(signal.audio_data, -1.0, 1.0)
        timing_logs["transform"] = time.time() - transform_start

        # Generator forward
        gen_forward_start = time.time()
        out = self.state.generator(signal.audio_data, signal.sample_rate)
        recons = AudioSignal(out["audio"], signal.sample_rate)
        timing_logs["gen_forward"] = time.time() - gen_forward_start
        z, mu, logs = out["z"], out["mu"], out["logs"]
        z.requires_grad_(True)
        logs_reg = torch.mean(logs.abs())  # Penalize large logs

        output["kl/loss"] = kl_loss(logs, mu)
        output["logs_penalty"] = logs_reg

        kl_beta = self.cosine_decay_with_warmup(
            cur_step=self.tracker.step,
            base_value=self.kl_max_beta,  # e.g., 100
            total_steps=self.num_iters,  # e.g., 250000
            final_value=0.1,  # 0.1,
            warmup_value=self.initial_norm,
            warmup_steps=0,
            hold_base_steps=self.hold_base_steps,
        )

        output["kl/beta"] = kl_beta

        logs_penalty_weight = self.cosine_decay_with_warmup(
            cur_step=self.tracker.step,
            base_value=self.init_logs_penalty,  # Initial weight for logs_penalty
            total_steps=self.num_iters,  # e.g., 250000
            final_value=self.init_logs_penalty
            * 0.01,  # * 0.0001,  # 10% of initial weight
            warmup_value=self.init_logs_penalty,
            warmup_steps=0,
            hold_base_steps=self.hold_base_steps,
        )
        lipschitz_penalty_weight = self.cosine_decay_with_warmup(
            cur_step=self.tracker.step,
            base_value=self.init_lipschitz_penalty,  # Initial weight for lipschitz_penalty
            total_steps=self.num_iters,  # e.g., 250000
            final_value=self.init_lipschitz_penalty
            * 0.01,  # * 0.0001,  # 10% of initial weight
            warmup_value=self.init_lipschitz_penalty,
            warmup_steps=0,
            hold_base_steps=self.hold_base_steps,
        )

        # Discriminator loss
        if self.tracker.step >= self.gan_start_step:
            print(f"Rank {self.local_rank}: Discriminator loss")
            disc_loss_start = time.time()
            output["adv/disc_loss"] = self.state.gan_loss.discriminator_loss(
                recons, signal
            )
            timing_logs["disc_loss"] = time.time() - disc_loss_start

            # Discriminator backward
            disc_backward_start = time.time()
            self.state.optimizer_d.zero_grad(set_to_none=True)
            output["adv/disc_loss"].backward()
            output["other/grad_norm_d"] = torch.nn.utils.clip_grad_norm_(
                self.state.discriminator.parameters(), self.max_d_norm
            )
            self.state.optimizer_d.step()
            self.state.scheduler_d.step()
            timing_logs["disc_backward"] = time.time() - disc_backward_start

            # DDP synchronization for discriminator
            disc_ddp_sync_start = time.time()
            # if torch.distributed.is_initialized():
            #     torch.distributed.barrier()
            timing_logs["disc_ddp_sync"] = time.time() - disc_ddp_sync_start
            (
                output["adv/gen_loss"],
                output["adv/feat_loss"],
            ) = self.state.gan_loss.generator_loss(recons, signal)

        # Generator losses
        gen_loss_start = time.time()
        output["stft/loss"] = self.state.stft_loss(recons, signal)

        output["mel/loss"] = self.state.mel_loss(recons, signal)
        output["waveform/loss"] = self.state.waveform_loss(recons, signal)

        output["lipschitz_penalty"] = self.compute_lipschitz_penalty(lambda_lip=0.01)
        output["grad_penalty"] = self.compute_gradient_penalty(recons.audio_data, z)

        loss_keys = [
            "stft/loss",
            "mel/loss",
            "waveform/loss",
            "kl/loss",
            "logs_penalty",
            "lipschitz_penalty",
            "grad_penalty",
        ]
        # print("self.tracker.step >= self.gan_start_step: ", self.tracker.step >= self.gan_start_step)
        if self.tracker.step >= self.gan_start_step:
            loss_keys.extend(["adv/gen_loss", "adv/feat_loss"])

        loss_weights = {k: self.state.lambdas.get(k, 1.0) for k in loss_keys}
        loss_weights["kl/loss"] = kl_beta
        loss_weights["logs_penalty"] = logs_penalty_weight
        loss_weights["lipschitz_penalty"] = lipschitz_penalty_weight

        # log the loss weights
        output.update({f"loss_weight/{k}": v for k, v in loss_weights.items()})

        output["loss"] = sum(
            [loss_weights[k] * output[k] for k in loss_keys if k in output]
        )
        timing_logs["gen_loss"] = time.time() - gen_loss_start

        # Generator backward
        print(f"Rank {self.local_rank}: Updating generator")
        gen_backward_start = time.time()
        self.state.optimizer_g.zero_grad(set_to_none=True)
        output["loss"].backward()

        encoder_grad_norm = torch.nn.utils.clip_grad_norm_(
            self.state.generator.module.encoder.parameters(), self.max_gen_norm
        )
        decoder_grad_norm = torch.nn.utils.clip_grad_norm_(
            self.state.generator.module.decoder.parameters(), self.max_gen_norm
        )

        if self.tracker.step % 2 == 0:  # Log every 100 iterations
            self.log_grad_norms(output, norm_threshold=0.0)

        output["other/grad_norm"] = torch.nn.utils.clip_grad_norm_(
            self.state.generator.parameters(), self.max_gen_norm
        )

        # Log gradient norms
        output["other/grad_norm_encoder"] = (
            encoder_grad_norm.item()
            if torch.is_tensor(encoder_grad_norm)
            else encoder_grad_norm
        )
        output["other/grad_norm_decoder"] = (
            decoder_grad_norm.item()
            if torch.is_tensor(decoder_grad_norm)
            else decoder_grad_norm
        )

        self.state.optimizer_g.step()
        self.state.scheduler_g.step()
        timing_logs["gen_backward"] = time.time() - gen_backward_start

        # self.state.ema.update()

        # DDP synchronization for generator
        gen_ddp_sync_start = time.time()
        # if torch.distributed.is_initialized():
        #     torch.distributed.barrier()
        timing_logs["gen_ddp_sync"] = time.time() - gen_ddp_sync_start

        # Other metrics
        output["other/learning_rate"] = self.state.optimizer_g.param_groups[0]["lr"]
        output["other/batch_size"] = signal.batch_size * self.world_size

        # Total train_loop time
        timing_logs["total_train_loop"] = time.time() - train_loop_start
        output.update({f"time/{k}": v for k, v in timing_logs.items()})

        print(f"Rank {self.local_rank}: train_loop complete")
        return {k: v for k, v in sorted(output.items())}

    def checkpoint(self):
        from datetime import datetime

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        step = self.state.tracker.step
        tags = ["latest"]
        if step % self.save_iters == 0:
            tags.append(f"{step // 1000}k")

        self.state.tracker.print(f"Saving checkpoint at step {step}")

        # Prepare everything for saving
        checkpoint = {
            "generator": self.unwrap(self.state.generator).state_dict(),
            "discriminator": self.unwrap(self.state.discriminator).state_dict(),
            "optimizer_g": self.state.optimizer_g.state_dict(),
            "optimizer_d": self.state.optimizer_d.state_dict(),
            "scheduler_g": self.state.scheduler_g.state_dict(),
            "scheduler_d": self.state.scheduler_d.state_dict(),
            "tracker": self.state.tracker.state_dict(),
            # "ema": self.state.ema.state_dict(),  # Save EMA state
            "step": step,
            "config": self.configs,
            "metadata": {
                "logs": self.state.tracker.history,
                "step": step,
                "config": self.configs,
            },
        }

        # Save for each tag (latest, 120k, etc)
        for tag in tags:
            save_folder = f"{self.save_path}/{tag}_{timestamp}"
            os.makedirs(save_folder, exist_ok=True)
            save_path = os.path.join(save_folder, "checkpoint.pt")
            torch.save(checkpoint, save_path)
            self.state.tracker.print(f"Checkpoint saved: {save_path}")

    def resume_from_checkpoint(self, load_folder):
        checkpoint_path = os.path.join(load_folder, "checkpoint.pt")
        assert os.path.exists(
            checkpoint_path
        ), f"Checkpoint {checkpoint_path} not found!"
        checkpoint = torch.load(checkpoint_path, map_location="cpu")

        # Load model state dicts
        self.unwrap(self.generator).load_state_dict(checkpoint["generator"])
        self.unwrap(self.discriminator).load_state_dict(checkpoint["discriminator"])

        # Load optimizer and scheduler state dicts **after** model is on device
        self.optimizer_g.load_state_dict(checkpoint["optimizer_g"])
        self.optimizer_d.load_state_dict(checkpoint["optimizer_d"])

        for state in self.optimizer_g.state.values():
            for k, v in state.items():
                if torch.is_tensor(v):
                    state[k] = v.to(self.local_rank)
        for state in self.optimizer_d.state.values():
            for k, v in state.items():
                if torch.is_tensor(v):
                    state[k] = v.to(self.local_rank)

        self.scheduler_g.load_state_dict(checkpoint["scheduler_g"])
        self.scheduler_d.load_state_dict(checkpoint["scheduler_d"])

        # Load EMA state
        # if "ema" in checkpoint:
        #     self.ema.load_state_dict(checkpoint["ema"])

        # Load tracker/logs/step
        self.tracker.load_state_dict(checkpoint["tracker"])
        self.tracker.step = checkpoint.get("step", 0)

        self.tracker.print(
            f"Checkpoint loaded from {checkpoint_path} at step {self.tracker.step}"
        )

    def unwrap(self, model):
        if hasattr(model, "module"):
            return model.module
        return model

    @torch.no_grad()
    def save_samples(self, val_idx):
        print(f"Rank {self.local_rank}: Starting save_samples")
        self.state.tracker.print("Saving audio samples to WandB")
        self.state.generator.eval()

        # Apply EMA weights
        # self.state.ema.apply_shadow()

        samples = [self.val_dataset[idx] for idx in val_idx]
        batch = self.val_dataset.collate(samples)
        batch = util.prepare_batch(batch, self.local_rank)
        signal = self.val_dataset.transform(
            batch["signal"].clone(), **batch["transform_args"]
        )

        out = self.state.generator(signal.audio_data, signal.sample_rate)
        recons = AudioSignal(out["audio"], signal.sample_rate)

        # Restore original weights
        # self.state.ema.restore()

        audio_dict = {"recons": recons}
        # if self.state.tracker.step == 0:
        audio_dict["signal"] = signal

        audio_logs = {}
        for k, v in audio_dict.items():
            for nb in range(v.batch_size):
                audio_data = v[nb].cpu().audio_data
                if audio_data.dim() == 3:
                    audio_data = audio_data.squeeze(0)
                elif audio_data.dim() == 1:
                    audio_data = audio_data.unsqueeze(0)

                audio_data = audio_data.numpy().astype(np.float32)
                if audio_data.max() > 1.0 or audio_data.min() < -1.0:
                    audio_data /= np.abs(audio_data).max()

                sample_rate = int(v[nb].sample_rate)
                if sample_rate <= 0:
                    raise ValueError(f"Invalid sample rate: {sample_rate}")

                # Save audio to a temporary file
                temp_file = f"temp_audio_{k}_{nb}.wav"
                sf.write(temp_file, audio_data.T, sample_rate)

                self.state.tracker.writer.log_audio(
                    temp_file,
                    metadata={
                        "caption": f"{k} sample {nb}",
                        "sample_rate": sample_rate,
                        "step": self.state.tracker.step,
                    },
                    step=self.state.tracker.step,
                )
                
                # Clean up temporary file
                os.remove(temp_file)

    def train(self):
        print(f"Rank {self.local_rank}: Starting train ")
        util.seed(self.seed)

        max_iters = self.num_iters
        train_loop = self.tracker.log("train", "value", history=False)(
            self.tracker.track("train", max_iters, completed=self.state.tracker.step)(
                self.train_loop
            )
        )

        save_samples = when(lambda: self.local_rank == 0)(self.save_samples)
        checkpoint = when(lambda: self.global_rank == 0)(self.checkpoint)

        with self.tracker.live:
            for self.tracker.step, batch in enumerate(
                self.train_dataloader, start=self.state.tracker.step
            ):
                self.tracker.print(
                    f"Rank {self.global_rank}: Iteration {self.tracker.step}/{max_iters} "
                )
                output = train_loop(batch)

                if self.global_rank == 0:
                    for k, v in output.items():
                        value = v.item() if torch.is_tensor(v) else v
                        self.tracker.writer.log_metric(k, value, step=self.tracker.step)

                last_iter = self.tracker.step == max_iters - 1

                if self.tracker.step % self.sample_freq == 0 or last_iter:
                    # torch.distributed.barrier()
                    save_samples(self.val_idx)
                    checkpoint()

                if last_iter:
                    break


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Distributed DAC training")
    parser.add_argument(
        "--config_path",
        type=str,
        default="configs/configx2.yml",
        help="Path to config YAML",
    )
    parser.add_argument("--run_id", type=str, required=True, help="Run ID for wandb")
    parser.add_argument(
        "--resume", action="store_true", help="Resume training from checkpoint"
    )
    parser.add_argument(
        "--load_weights", action="store_true", help="Load weights from checkpoint"
    )
    parser.add_argument(
        "--save_path", type=str, default="ckpts", help="Path to save checkpoints"
    )
    parser.add_argument("--tag", type=str, default="latest", help="Tag for checkpoint")
    args = parser.parse_args()

    ddp_setup()
    trainer = Trainer(args)
    trainer.train()
    destroy_process_group()