| |
| |
| |
|
|
| import argparse, json, torch, torch.nn.functional as F |
| from pathlib import Path |
| from torch.utils.data import Dataset, DataLoader |
| from torch.optim import AdamW |
| from torch.optim.lr_scheduler import CosineAnnealingLR |
| try: |
| from vspark_model import VSparkModel, VSparkConfig, DDPMScheduler |
| except ImportError: |
| from vspark_build_fixed import VSparkModel, VSparkConfig, DDPMScheduler |
|
|
|
|
| class DummyVideoDataset(Dataset): |
| def __init__(self, cfg, size=200): |
| self.cfg = cfg |
| self.size = size |
|
|
| def __len__(self): |
| return self.size |
|
|
| def __getitem__(self, _): |
| cfg = self.cfg |
| Hl = cfg.latent_size |
| return dict( |
| video_latents = torch.randn(cfg.num_frames, cfg.in_channels, Hl, Hl), |
| input_ids = torch.randint(0, 49408, (cfg.max_text_len,)), |
| audio_mel = torch.randn(cfg.mel_frames, cfg.n_mels), |
| ) |
|
|
|
|
| def train(args): |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| cfg = VSparkConfig() |
| model = VSparkModel(cfg).to(device) |
| scheduler = DDPMScheduler(cfg).to(device) |
| dataset = DummyVideoDataset(cfg) |
| loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=0) |
| optim = AdamW(model.parameters(), lr=args.lr, weight_decay=1e-2) |
| lr_sched = CosineAnnealingLR(optim, T_max=args.epochs * len(loader), eta_min=1e-6) |
| out_dir = Path(args.output_dir); out_dir.mkdir(parents=True, exist_ok=True) |
|
|
| step = 0 |
| for epoch in range(args.epochs): |
| model.train() |
| for batch in loader: |
| lat = batch["video_latents"].to(device) |
| ids = batch["input_ids"].to(device) |
| mel = batch["audio_mel"].to(device) |
| B = lat.shape[0] |
| t = torch.randint(0, cfg.num_timesteps, (B,), device=device) |
| noisy_lat, noise_lat = scheduler.add_noise(lat, t) |
| noisy_mel, noise_mel = scheduler.add_noise(mel, t) |
| pred_lat = model(noisy_lat, t, ids) |
| pred_mel = model.predict_audio(noisy_mel, t, ids) |
| loss = F.mse_loss(pred_lat, noise_lat) + 0.5 * F.mse_loss(pred_mel, noise_mel) |
| optim.zero_grad() |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| optim.step(); lr_sched.step() |
| step += 1 |
| if step % args.log_every == 0: |
| print(f" step={step} loss={loss.item():.4f}") |
| ckpt = out_dir / f"ckpt_epoch{epoch+1:03d}.pt" |
| torch.save({"model_state_dict": model.state_dict(), "epoch": epoch + 1}, ckpt) |
|
|
|
|
| if __name__ == "__main__": |
| p = argparse.ArgumentParser() |
| p.add_argument("--output_dir", default="./checkpoints") |
| p.add_argument("--epochs", type=int, default=100) |
| p.add_argument("--batch_size", type=int, default=1) |
| p.add_argument("--lr", type=float, default=1e-4) |
| p.add_argument("--log_every", type=int, default=10) |
| train(p.parse_args()) |
|
|
