kick-gen-v1 / training /train_diffusion.py

Upload 9 files

a3b520a verified 9 months ago

10.7 kB

"""Training script for the latent diffusion model."""

import copy
import csv
import sys
from pathlib import Path

import torch
from torch.utils.data import DataLoader, Dataset, random_split
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

sys.path.insert(0, str(Path(__file__).parent.parent))

from models.autoencoder import KickVAE
from models.diffusion import LatentUNet, NoiseScheduler
from models.text_encoder import KeywordEncoder, build_vocab
from training.config import DiffusionConfig


# ---------------------------------------------------------------------------
# Dataset
# ---------------------------------------------------------------------------

class LatentDataset(Dataset):
    """Dataset of pre-encoded VAE latents with keyword token IDs."""

    def __init__(
        self,
        latents_dir: Path,
        metadata_csv: Path,
        vocab: list[str],
    ) -> None:
        self.latent_files = sorted(
            f for f in latents_dir.glob("*.pt") if not f.name.startswith("._")
        )
        if not self.latent_files:
            raise FileNotFoundError(f"No .pt files in {latents_dir}")

        # Build keyword lookup: filename_stem -> list of keyword strings
        self.kw_to_idx = {kw: i for i, kw in enumerate(vocab)}
        self.keywords: dict[str, list[int]] = {}
        with open(metadata_csv) as f:
            reader = csv.DictReader(f)
            for row in reader:
                stem = Path(row["filename"]).stem
                ids = []
                for kw in row["keywords"].split(","):
                    kw = kw.strip().lower()
                    if kw in self.kw_to_idx:
                        ids.append(self.kw_to_idx[kw])
                self.keywords[stem] = ids

    def __len__(self) -> int:
        return len(self.latent_files)

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, list[int]]:
        path = self.latent_files[idx]
        latent = torch.load(path, weights_only=False)
        # latent filename matches mel filename stem
        stem = path.stem
        token_ids = self.keywords.get(stem, [])
        return latent, token_ids


def collate_fn(
    batch: list[tuple[torch.Tensor, list[int]]],
) -> tuple[torch.Tensor, list[list[int]]]:
    """Custom collate to handle variable-length keyword lists."""
    latents = torch.stack([b[0] for b in batch])
    token_ids = [b[1] for b in batch]
    return latents, token_ids


# ---------------------------------------------------------------------------
# Pre-encode latents
# ---------------------------------------------------------------------------

def pre_encode_latents(cfg: DiffusionConfig) -> None:
    """Encode all mel spectrograms to latents using frozen VAE."""
    cfg.latents_dir.mkdir(parents=True, exist_ok=True)

    # Check if already done
    existing = list(cfg.latents_dir.glob("*.pt"))
    if len(existing) > 100:
        print(f"Latents dir already has {len(existing)} files, skipping encoding.")
        return

    device = torch.device(
        "cuda" if torch.cuda.is_available()
        else "mps" if torch.backends.mps.is_available()
        else "cpu"
    )

    # Load VAE
    checkpoint = torch.load(cfg.vae_checkpoint, weights_only=False)
    vae = KickVAE(latent_dim=cfg.latent_dim).to(device)
    vae.load_state_dict(checkpoint["model_state_dict"])
    vae.eval()

    mel_files = sorted(
        f for f in cfg.data_dir.glob("*.pt") if not f.name.startswith("._")
    )
    print(f"Encoding {len(mel_files)} mel spectrograms to latents...")

    with torch.no_grad():
        for f in tqdm(mel_files):
            out_path = cfg.latents_dir / f.name
            if out_path.exists():
                continue
            mel = torch.load(f, weights_only=False).unsqueeze(0).to(device)
            latent = vae.encode(mel).squeeze(0).cpu()
            torch.save(latent, out_path)

    print("Latent encoding complete.")


# ---------------------------------------------------------------------------
# EMA
# ---------------------------------------------------------------------------

class EMA:
    """Exponential moving average of model parameters."""

    def __init__(self, model: torch.nn.Module, decay: float = 0.9999) -> None:
        self.decay = decay
        self.shadow = copy.deepcopy(model)
        self.shadow.eval()
        for p in self.shadow.parameters():
            p.requires_grad_(False)

    @torch.no_grad()
    def update(self, model: torch.nn.Module) -> None:
        for s, p in zip(self.shadow.parameters(), model.parameters()):
            s.data.mul_(self.decay).add_(p.data, alpha=1 - self.decay)


# ---------------------------------------------------------------------------
# Training
# ---------------------------------------------------------------------------

def train(cfg: DiffusionConfig | None = None) -> None:
    if cfg is None:
        cfg = DiffusionConfig()

    # Pre-encode latents
    pre_encode_latents(cfg)

    device = torch.device(
        "cuda" if torch.cuda.is_available()
        else "mps" if torch.backends.mps.is_available()
        else "cpu"
    )
    print(f"Using device: {device}")

    # Build vocab and dataset
    vocab = build_vocab(cfg.metadata_csv)
    print(f"Vocabulary size: {len(vocab)}")

    dataset = LatentDataset(cfg.latents_dir, cfg.metadata_csv, vocab)
    val_size = int(len(dataset) * cfg.val_split)
    train_size = len(dataset) - val_size
    train_set, val_set = random_split(
        dataset, [train_size, val_size],
        generator=torch.Generator().manual_seed(42),
    )

    train_loader = DataLoader(
        train_set,
        batch_size=cfg.batch_size,
        shuffle=True,
        num_workers=cfg.num_workers,
        pin_memory=True,
        collate_fn=collate_fn,
    )
    val_loader = DataLoader(
        val_set,
        batch_size=cfg.batch_size,
        shuffle=False,
        num_workers=cfg.num_workers,
        pin_memory=True,
        collate_fn=collate_fn,
    )
    print(f"Train: {train_size}, Val: {val_size}")

    # Model
    model = LatentUNet(
        latent_dim=cfg.latent_dim,
        base_channels=cfg.base_channels,
        cond_dim=cfg.cond_dim,
    ).to(device)
    text_enc = KeywordEncoder(
        vocab_size=len(vocab),
        embed_dim=cfg.text_embed_dim,
        cond_dim=cfg.cond_dim,
    ).to(device)

    scheduler = NoiseScheduler(cfg.timesteps, cfg.beta_start, cfg.beta_end).to(device)
    ema = EMA(model, cfg.ema_decay)

    optimizer = torch.optim.AdamW(
        list(model.parameters()) + list(text_enc.parameters()),
        lr=cfg.learning_rate,
    )
    scaler = torch.amp.GradScaler(enabled=cfg.use_amp and device.type == "cuda")

    # Logging
    cfg.log_dir.mkdir(parents=True, exist_ok=True)
    cfg.checkpoint_dir.mkdir(parents=True, exist_ok=True)
    writer = SummaryWriter(cfg.log_dir)

    # Training loop (iteration-based)
    global_step = 0
    model.train()
    text_enc.train()

    print(f"Training for {cfg.iterations} iterations...")

    while global_step < cfg.iterations:
        for latents, token_ids in train_loader:
            if global_step >= cfg.iterations:
                break

            latents = latents.to(device)
            batch_size = latents.shape[0]

            # Classifier-free guidance dropout: replace keywords with empty list
            dropped_ids = []
            for ids in token_ids:
                if torch.rand(1).item() < cfg.cfg_dropout:
                    dropped_ids.append([])
                else:
                    dropped_ids.append(ids)

            # Sample timesteps and noise
            t = torch.randint(0, cfg.timesteps, (batch_size,), device=device)
            noise = torch.randn_like(latents)
            noisy = scheduler.add_noise(latents, noise, t)

            with torch.amp.autocast(
                device_type=device.type,
                enabled=cfg.use_amp and device.type == "cuda",
            ):
                cond = text_enc(dropped_ids, device)
                pred_noise = model(noisy, t, cond)
                loss = torch.nn.functional.mse_loss(pred_noise, noise)
                loss = loss / cfg.gradient_accumulation

            scaler.scale(loss).backward()

            if (global_step + 1) % cfg.gradient_accumulation == 0:
                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad()
                ema.update(model)

            # Logging
            if global_step % 50 == 0:
                writer.add_scalar(
                    "train/loss", loss.item() * cfg.gradient_accumulation, global_step
                )

            if global_step % 500 == 0:
                print(
                    f"Step {global_step}/{cfg.iterations} "
                    f"loss={loss.item() * cfg.gradient_accumulation:.6f}"
                )

            # Validation
            if global_step % 1000 == 0 and global_step > 0:
                model.eval()
                text_enc.eval()
                val_loss_sum = 0.0
                val_count = 0
                with torch.no_grad():
                    for vl, vt in val_loader:
                        vl = vl.to(device)
                        vt_step = torch.randint(
                            0, cfg.timesteps, (vl.shape[0],), device=device
                        )
                        vn = torch.randn_like(vl)
                        vnoisy = scheduler.add_noise(vl, vn, vt_step)
                        vcond = text_enc(vt, device)
                        vpred = model(vnoisy, vt_step, vcond)
                        val_loss_sum += torch.nn.functional.mse_loss(vpred, vn).item()
                        val_count += 1
                avg_val = val_loss_sum / max(val_count, 1)
                writer.add_scalar("val/loss", avg_val, global_step)
                print(f"  val_loss={avg_val:.6f}")
                model.train()
                text_enc.train()

            # Checkpoint
            if (global_step + 1) % cfg.checkpoint_every == 0:
                path = cfg.checkpoint_dir / f"diffusion_step_{global_step+1}.pt"
                torch.save({
                    "step": global_step + 1,
                    "model_state_dict": model.state_dict(),
                    "ema_state_dict": ema.shadow.state_dict(),
                    "text_enc_state_dict": text_enc.state_dict(),
                    "optimizer_state_dict": optimizer.state_dict(),
                    "vocab": vocab,
                    "config": cfg,
                }, path)
                print(f"Saved checkpoint: {path}")

            global_step += 1

    writer.close()
    print("Diffusion training complete.")


if __name__ == "__main__":
    train()

	"""Training script for the latent diffusion model."""

	import copy
	import csv
	import sys
	from pathlib import Path

	import torch
	from torch.utils.data import DataLoader, Dataset, random_split
	from torch.utils.tensorboard import SummaryWriter
	from tqdm import tqdm

	sys.path.insert(0, str(Path(__file__).parent.parent))

	from models.autoencoder import KickVAE
	from models.diffusion import LatentUNet, NoiseScheduler
	from models.text_encoder import KeywordEncoder, build_vocab
	from training.config import DiffusionConfig


	# ---------------------------------------------------------------------------
	# Dataset
	# ---------------------------------------------------------------------------

	class LatentDataset(Dataset):
	"""Dataset of pre-encoded VAE latents with keyword token IDs."""

	def __init__(
	self,
	latents_dir: Path,
	metadata_csv: Path,
	vocab: list[str],
	) -> None:
	self.latent_files = sorted(
	f for f in latents_dir.glob("*.pt") if not f.name.startswith("._")
	)
	if not self.latent_files:
	raise FileNotFoundError(f"No .pt files in {latents_dir}")

	# Build keyword lookup: filename_stem -> list of keyword strings
	self.kw_to_idx = {kw: i for i, kw in enumerate(vocab)}
	self.keywords: dict[str, list[int]] = {}
	with open(metadata_csv) as f:
	reader = csv.DictReader(f)
	for row in reader:
	stem = Path(row["filename"]).stem
	ids = []
	for kw in row["keywords"].split(","):
	kw = kw.strip().lower()
	if kw in self.kw_to_idx:
	ids.append(self.kw_to_idx[kw])
	self.keywords[stem] = ids

	def __len__(self) -> int:
	return len(self.latent_files)

	def __getitem__(self, idx: int) -> tuple[torch.Tensor, list[int]]:
	path = self.latent_files[idx]
	latent = torch.load(path, weights_only=False)
	# latent filename matches mel filename stem
	stem = path.stem
	token_ids = self.keywords.get(stem, [])
	return latent, token_ids


	def collate_fn(
	batch: list[tuple[torch.Tensor, list[int]]],
	) -> tuple[torch.Tensor, list[list[int]]]:
	"""Custom collate to handle variable-length keyword lists."""
	latents = torch.stack([b[0] for b in batch])
	token_ids = [b[1] for b in batch]
	return latents, token_ids


	# ---------------------------------------------------------------------------
	# Pre-encode latents
	# ---------------------------------------------------------------------------

	def pre_encode_latents(cfg: DiffusionConfig) -> None:
	"""Encode all mel spectrograms to latents using frozen VAE."""
	cfg.latents_dir.mkdir(parents=True, exist_ok=True)

	# Check if already done
	existing = list(cfg.latents_dir.glob("*.pt"))
	if len(existing) > 100:
	print(f"Latents dir already has {len(existing)} files, skipping encoding.")
	return

	device = torch.device(
	"cuda" if torch.cuda.is_available()
	else "mps" if torch.backends.mps.is_available()
	else "cpu"
	)

	# Load VAE
	checkpoint = torch.load(cfg.vae_checkpoint, weights_only=False)
	vae = KickVAE(latent_dim=cfg.latent_dim).to(device)
	vae.load_state_dict(checkpoint["model_state_dict"])
	vae.eval()

	mel_files = sorted(
	f for f in cfg.data_dir.glob("*.pt") if not f.name.startswith("._")
	)
	print(f"Encoding {len(mel_files)} mel spectrograms to latents...")

	with torch.no_grad():
	for f in tqdm(mel_files):
	out_path = cfg.latents_dir / f.name
	if out_path.exists():
	continue
	mel = torch.load(f, weights_only=False).unsqueeze(0).to(device)
	latent = vae.encode(mel).squeeze(0).cpu()
	torch.save(latent, out_path)

	print("Latent encoding complete.")


	# ---------------------------------------------------------------------------
	# EMA
	# ---------------------------------------------------------------------------

	class EMA:
	"""Exponential moving average of model parameters."""

	def __init__(self, model: torch.nn.Module, decay: float = 0.9999) -> None:
	self.decay = decay
	self.shadow = copy.deepcopy(model)
	self.shadow.eval()
	for p in self.shadow.parameters():
	p.requires_grad_(False)

	@torch.no_grad()
	def update(self, model: torch.nn.Module) -> None:
	for s, p in zip(self.shadow.parameters(), model.parameters()):
	s.data.mul_(self.decay).add_(p.data, alpha=1 - self.decay)


	# ---------------------------------------------------------------------------
	# Training
	# ---------------------------------------------------------------------------

	def train(cfg: DiffusionConfig \| None = None) -> None:
	if cfg is None:
	cfg = DiffusionConfig()

	# Pre-encode latents
	pre_encode_latents(cfg)

	device = torch.device(
	"cuda" if torch.cuda.is_available()
	else "mps" if torch.backends.mps.is_available()
	else "cpu"
	)
	print(f"Using device: {device}")

	# Build vocab and dataset
	vocab = build_vocab(cfg.metadata_csv)
	print(f"Vocabulary size: {len(vocab)}")

	dataset = LatentDataset(cfg.latents_dir, cfg.metadata_csv, vocab)
	val_size = int(len(dataset) * cfg.val_split)
	train_size = len(dataset) - val_size
	train_set, val_set = random_split(
	dataset, [train_size, val_size],
	generator=torch.Generator().manual_seed(42),
	)

	train_loader = DataLoader(
	train_set,
	batch_size=cfg.batch_size,
	shuffle=True,
	num_workers=cfg.num_workers,
	pin_memory=True,
	collate_fn=collate_fn,
	)
	val_loader = DataLoader(
	val_set,
	batch_size=cfg.batch_size,
	shuffle=False,
	num_workers=cfg.num_workers,
	pin_memory=True,
	collate_fn=collate_fn,
	)
	print(f"Train: {train_size}, Val: {val_size}")

	# Model
	model = LatentUNet(
	latent_dim=cfg.latent_dim,
	base_channels=cfg.base_channels,
	cond_dim=cfg.cond_dim,
	).to(device)
	text_enc = KeywordEncoder(
	vocab_size=len(vocab),
	embed_dim=cfg.text_embed_dim,
	cond_dim=cfg.cond_dim,
	).to(device)

	scheduler = NoiseScheduler(cfg.timesteps, cfg.beta_start, cfg.beta_end).to(device)
	ema = EMA(model, cfg.ema_decay)

	optimizer = torch.optim.AdamW(
	list(model.parameters()) + list(text_enc.parameters()),
	lr=cfg.learning_rate,
	)
	scaler = torch.amp.GradScaler(enabled=cfg.use_amp and device.type == "cuda")

	# Logging
	cfg.log_dir.mkdir(parents=True, exist_ok=True)
	cfg.checkpoint_dir.mkdir(parents=True, exist_ok=True)
	writer = SummaryWriter(cfg.log_dir)

	# Training loop (iteration-based)
	global_step = 0
	model.train()
	text_enc.train()

	print(f"Training for {cfg.iterations} iterations...")

	while global_step < cfg.iterations:
	for latents, token_ids in train_loader:
	if global_step >= cfg.iterations:
	break

	latents = latents.to(device)
	batch_size = latents.shape[0]

	# Classifier-free guidance dropout: replace keywords with empty list
	dropped_ids = []
	for ids in token_ids:
	if torch.rand(1).item() < cfg.cfg_dropout:
	dropped_ids.append([])
	else:
	dropped_ids.append(ids)

	# Sample timesteps and noise
	t = torch.randint(0, cfg.timesteps, (batch_size,), device=device)
	noise = torch.randn_like(latents)
	noisy = scheduler.add_noise(latents, noise, t)

	with torch.amp.autocast(
	device_type=device.type,
	enabled=cfg.use_amp and device.type == "cuda",
	):
	cond = text_enc(dropped_ids, device)
	pred_noise = model(noisy, t, cond)
	loss = torch.nn.functional.mse_loss(pred_noise, noise)
	loss = loss / cfg.gradient_accumulation

	scaler.scale(loss).backward()

	if (global_step + 1) % cfg.gradient_accumulation == 0:
	scaler.step(optimizer)
	scaler.update()
	optimizer.zero_grad()
	ema.update(model)

	# Logging
	if global_step % 50 == 0:
	writer.add_scalar(
	"train/loss", loss.item() * cfg.gradient_accumulation, global_step
	)

	if global_step % 500 == 0:
	print(
	f"Step {global_step}/{cfg.iterations} "
	f"loss={loss.item() * cfg.gradient_accumulation:.6f}"
	)

	# Validation
	if global_step % 1000 == 0 and global_step > 0:
	model.eval()
	text_enc.eval()
	val_loss_sum = 0.0
	val_count = 0
	with torch.no_grad():
	for vl, vt in val_loader:
	vl = vl.to(device)
	vt_step = torch.randint(
	0, cfg.timesteps, (vl.shape[0],), device=device
	)
	vn = torch.randn_like(vl)
	vnoisy = scheduler.add_noise(vl, vn, vt_step)
	vcond = text_enc(vt, device)
	vpred = model(vnoisy, vt_step, vcond)
	val_loss_sum += torch.nn.functional.mse_loss(vpred, vn).item()
	val_count += 1
	avg_val = val_loss_sum / max(val_count, 1)
	writer.add_scalar("val/loss", avg_val, global_step)
	print(f" val_loss={avg_val:.6f}")
	model.train()
	text_enc.train()

	# Checkpoint
	if (global_step + 1) % cfg.checkpoint_every == 0:
	path = cfg.checkpoint_dir / f"diffusion_step_{global_step+1}.pt"
	torch.save({
	"step": global_step + 1,
	"model_state_dict": model.state_dict(),
	"ema_state_dict": ema.shadow.state_dict(),
	"text_enc_state_dict": text_enc.state_dict(),
	"optimizer_state_dict": optimizer.state_dict(),
	"vocab": vocab,
	"config": cfg,
	}, path)
	print(f"Saved checkpoint: {path}")

	global_step += 1

	writer.close()
	print("Diffusion training complete.")


	if __name__ == "__main__":
	train()