BlueWaveSemi45/DramaboxCPU
0
1#!/usr/bin/env python32"""3Audio-Only IC-LoRA Training for Voice Cloning on LTX-2.3.4 5Uses the IC-LoRA pattern: reference audio tokens are APPENDED to the end of6the target sequence using AudioConditionByReferenceLatent. Loss is computed7only on target tokens; reference tokens remain clean (denoise_mask=0).8 9This follows the official video-to-video IC-LoRA strategy closely, but adapted10for the audio-only modality path.11 12Usage (single GPU):13 CUDA_VISIBLE_DEVICES=0 python train_audio_iclora.py --data-dir ... --speaker-index ...14 15Usage (multi-GPU with accelerate):16 CUDA_VISIBLE_DEVICES=4,5,6,7 accelerate launch --num_processes=4 train_audio_iclora.py ...17"""18 19import argparse20import logging21import math22import os23import random24import shutil25import sys26import time27from collections import defaultdict28from pathlib import Path29 30import torch31import torch.nn.functional as F32from torch.utils.data import DataLoader, Dataset33 34REPO_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))35sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "ltx2"))36# ltx-pipelines already on path via ltx2/37 38MODEL_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))39 40# Import audio conditioning item from our module41sys.path.insert(0, MODEL_DIR)42from audio_conditioning import AudioConditionByReferenceLatent43 44 45# ─── Timestep Sampling ───46 47class DistilledTimestepSampler:48 """Sample timesteps from the distilled sigma schedule.49 50 The distilled model was trained to denoise at these specific sigma values.51 We sample uniformly from the intervals between consecutive sigmas,52 matching the distribution the model actually operates on.53 """54 55 # Distilled 8-step sigma values (boundaries of denoising intervals)56 SIGMAS = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0]57 58 def __init__(self, jitter: float = 0.02):59 self.jitter = jitter60 61 def sample(self, batch_size: int, seq_length: int = None, device: torch.device = None) -> torch.Tensor:62 n_intervals = len(self.SIGMAS) - 163 interval_idx = torch.randint(0, n_intervals, (batch_size,), device=device)64 t = torch.rand(batch_size, device=device)65 sigma_high = torch.tensor([self.SIGMAS[i] for i in interval_idx], device=device)66 sigma_low = torch.tensor([self.SIGMAS[i + 1] for i in interval_idx], device=device)67 sigma = sigma_low + t * (sigma_high - sigma_low)68 return sigma.clamp(0.01, 0.99)69 70 71class ShiftedLogitNormalTimestepSampler:72 """Shifted logit-normal distribution, shift depends on sequence length."""73 74 def __init__(self, std: float = 1.0, eps: float = 1e-3, uniform_prob: float = 0.1):75 self.std = std76 self.eps = eps77 self.uniform_prob = uniform_prob78 self.normal_999_percentile = 3.0902 * std79 self.normal_005_percentile = -2.5758 * std80 81 def sample(self, batch_size: int, seq_length: int, device: torch.device = None) -> torch.Tensor:82 mu = self._get_shift(seq_length)83 normal = torch.randn(batch_size, device=device) * self.std + mu84 logitnormal = torch.sigmoid(normal)85 86 p999 = torch.sigmoid(torch.tensor(mu + self.normal_999_percentile, device=device))87 p005 = torch.sigmoid(torch.tensor(mu + self.normal_005_percentile, device=device))88 stretched = (logitnormal - p005) / (p999 - p005)89 stretched = torch.where(stretched >= self.eps, stretched, 2 * self.eps - stretched)90 stretched = stretched.clamp(0, 1)91 92 uniform = (1 - self.eps) * torch.rand(batch_size, device=device) + self.eps93 prob = torch.rand(batch_size, device=device)94 return torch.where(prob > self.uniform_prob, stretched, uniform)95 96 @staticmethod97 def _get_shift(seq_length, min_tok=1024, max_tok=4096, min_s=0.95, max_s=2.05):98 m = (max_s - min_s) / (max_tok - min_tok)99 return m * seq_length + (min_s - m * min_tok)100 101 102# ─── Dataset ───103 104def build_speaker_map(index_paths, data_dirs):105 """Map speaker → [(data_dir, sample_idx)] from index file(s).106 107 The sample index comes from field 0 of the `~`-delimited row when it108 parses as int (allows subset indexes that keep original sample numbers),109 otherwise we fall back to the row's line number (legacy behaviour for110 string-keyed indexes like tts_training_data_podcast).111 """112 speaker_to_samples = defaultdict(list)113 for index_path, data_dir in zip(index_paths, data_dirs):114 with open(index_path) as f:115 for line_num, line in enumerate(f):116 parts = line.strip().split("~")117 if len(parts) < 7:118 continue119 try:120 idx = int(parts[0])121 except ValueError:122 idx = line_num123 speaker_id = parts[1]124 speaker_to_samples[speaker_id].append((data_dir, idx))125 return {k: v for k, v in speaker_to_samples.items() if len(v) >= 2}126 127 128class IDLoRADataset(Dataset):129 # Silence-latent reference loaded once, used to detect and strip any130 # leading silence frames baked into the preprocessed audio_latents. The131 # training loop ALREADY prepends 0-25 random silence frames, so we don't132 # want accidental silence in the source data compounding on top.133 _silence_ref = None134 135 @classmethod136 def _load_silence_ref(cls):137 if cls._silence_ref is None:138 p = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),139 "assets", "silence_latent_frame.pt")140 if os.path.exists(p):141 cls._silence_ref = torch.load(p, weights_only=True).float().squeeze() # [C, F]142 return cls._silence_ref143 144 def __init__(self, speaker_map):145 self.samples = []146 self.speaker_map = {}147 for speaker, entries in speaker_map.items():148 valid = []149 for data_dir, idx in entries:150 audio_path = Path(data_dir) / "audio_latents" / f"sample_{idx:06d}.pt"151 cond_path = Path(data_dir) / "conditions" / f"sample_{idx:06d}.pt"152 if audio_path.exists() and cond_path.exists():153 valid.append((data_dir, idx))154 if len(valid) >= 2:155 self.speaker_map[speaker] = valid156 for speaker, entries in self.speaker_map.items():157 for entry in entries:158 self.samples.append((entry, speaker))159 IDLoRADataset._load_silence_ref()160 161 def __len__(self):162 return len(self.samples)163 164 def _load_sample(self, data_dir, idx):165 base = Path(data_dir)166 audio = torch.load(base / "audio_latents" / f"sample_{idx:06d}.pt", weights_only=False)167 # Prefer prefix-stripped text embeddings if they exist (re-encoded with168 # just the quoted dialogue, dropping the "A woman says, " / "A man169 # speaks with X accent, " scene-description prefix).170 stripped = base / "conditions_stripped" / f"sample_{idx:06d}.pt"171 cond_path = stripped if stripped.exists() else base / "conditions" / f"sample_{idx:06d}.pt"172 cond = torch.load(cond_path, weights_only=False)173 if isinstance(audio, dict):174 audio = audio.get("audio_latent", audio.get("latent", list(audio.values())[0]))175 if audio.dim() == 2:176 audio = audio.unsqueeze(0)177 audio_feats = cond.get("audio_prompt_embeds", cond.get("prompt_embeds"))178 attn_mask = cond.get("prompt_attention_mask")179 # The audio_connector has num_learnable_registers=128 and asserts the180 # input sequence length is divisible by 128. Our new preprocessing181 # saved trimmed conditions (dropping left-padding to save disk), which182 # produces short/irregular sequence lengths. Left-pad back to the next183 # multiple of 128 with zeros (matching the tokenizer's left-padding184 # convention) so this assertion holds.185 REG = 128186 L = audio_feats.shape[0]187 target_L = ((L + REG - 1) // REG) * REG188 if target_L != L:189 pad_len = target_L - L190 pad_emb = torch.zeros(pad_len, audio_feats.shape[1],191 dtype=audio_feats.dtype)192 pad_mask = torch.zeros(pad_len, dtype=attn_mask.dtype)193 audio_feats = torch.cat([pad_emb, audio_feats], dim=0)194 attn_mask = torch.cat([pad_mask, attn_mask], dim=0)195 return audio, audio_feats, attn_mask196 197 def __getitem__(self, idx):198 (data_dir, tgt_idx), speaker = self.samples[idx]199 tgt_latent, audio_feats, attn_mask = self._load_sample(data_dir, tgt_idx)200 201 # Drop the reference entirely for non-voice-cloning categories:202 # - SFX samples (speaker starts with "sfx_"): descriptive sound events,203 # no speaker identity to clone.204 # - Song/music samples (suno dataset): prompts describe the music style,205 # reference audio doesn't transfer anything useful.206 # Return a zero-length ref so the model trains target-only for these.207 drop_ref = speaker.startswith("sfx_") or "preprocessed_ltx_suno" in str(data_dir)208 if drop_ref:209 C, F_dim = tgt_latent.shape[0], tgt_latent.shape[2]210 ref_latent = torch.zeros(C, 0, F_dim, dtype=tgt_latent.dtype)211 else:212 entries = self.speaker_map[speaker]213 ref_entry = random.choice([e for e in entries if e[1] != tgt_idx])214 ref_latent, _, _ = self._load_sample(*ref_entry)215 216 return {217 "tgt_latent": tgt_latent,218 "ref_latent": ref_latent,219 "audio_features": audio_feats,220 "attention_mask": attn_mask,221 }222 223 224# ─── Model building ───225 226def build_audio_only_model(checkpoint_path, device, dtype):227 from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder228 from ltx_core.loader.registry import DummyRegistry229 from ltx_core.loader.sd_ops import SDOps230 from ltx_core.model.transformer.model import LTXModel, LTXModelType231 from ltx_core.model.model_protocol import ModelConfigurator232 from ltx_core.model.transformer.attention import AttentionFunction233 from ltx_core.model.transformer.rope import LTXRopeType234 235 sd_ops = SDOps("AO").with_matching(prefix="model.diffusion_model.").with_replacement("model.diffusion_model.", "")236 237 class Cfg(ModelConfigurator[LTXModel]):238 @classmethod239 def from_config(cls, config):240 t = config.get("transformer", {})241 cp = None242 if not t.get("caption_proj_before_connector", False):243 from ltx_core.model.transformer.text_projection import create_caption_projection244 with torch.device("meta"):245 cp = create_caption_projection(t, audio=True)246 return LTXModel(247 model_type=LTXModelType.AudioOnly,248 audio_num_attention_heads=t.get("audio_num_attention_heads", 32),249 audio_attention_head_dim=t.get("audio_attention_head_dim", 64),250 audio_in_channels=t.get("audio_in_channels", 128),251 audio_out_channels=t.get("audio_out_channels", 128),252 num_layers=t.get("num_layers", 48),253 audio_cross_attention_dim=t.get("audio_cross_attention_dim", 2048),254 norm_eps=t.get("norm_eps", 1e-6),255 attention_type=AttentionFunction(t.get("attention_type", "default")),256 positional_embedding_theta=t.get("positional_embedding_theta", 10000.0),257 audio_positional_embedding_max_pos=t.get("audio_positional_embedding_max_pos", [20]),258 timestep_scale_multiplier=t.get("timestep_scale_multiplier", 1000),259 use_middle_indices_grid=t.get("use_middle_indices_grid", True),260 rope_type=LTXRopeType(t.get("rope_type", "interleaved")),261 double_precision_rope=t.get("frequencies_precision", False) == "float64",262 apply_gated_attention=t.get("apply_gated_attention", False),263 audio_caption_projection=cp,264 cross_attention_adaln=t.get("cross_attention_adaln", False),265 )266 267 builder = Builder(model_path=checkpoint_path, model_class_configurator=Cfg,268 model_sd_ops=sd_ops, registry=DummyRegistry())269 return builder.build(device=device, dtype=dtype)270 271 272def load_audio_connector(checkpoint_path, device, dtype):273 # ltx-trainer already on path via ltx2/274 from ltx_trainer.model_loader import load_embeddings_processor275 emb_proc = load_embeddings_processor(checkpoint_path, device=device, dtype=dtype)276 connector = emb_proc.audio_connector277 del emb_proc278 return connector279 280 281def apply_lora(model, rank, alpha, dropout=0.0):282 from peft import LoraConfig, get_peft_model283 config = LoraConfig(284 r=rank, lora_alpha=alpha, lora_dropout=dropout, bias="none",285 target_modules=[286 # Self-attention over audio tokens (voice-transfer pathway via ref).287 "audio_attn1.to_k", "audio_attn1.to_q", "audio_attn1.to_v", "audio_attn1.to_out.0",288 # Cross-attention (audio ↔ text context) NOT adapted — keep base289 # model's prompt→audio behaviour intact and rely on dataset balance290 # to drive expressiveness. (v15c tried this with adaLN unfreeze,291 # that proved too destructive; v16 tries it adaLN-frozen.)292 # FFN — non-linear capacity for style/phonetic adaptation.293 "audio_ff.net.0.proj", "audio_ff.net.2",294 ],295 )296 model = get_peft_model(model, config)297 trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)298 total = sum(p.numel() for p in model.parameters())299 logging.info(f"LoRA: {trainable:,} trainable / {total:,} total ({100*trainable/total:.1f}%)")300 return model301 302 303@torch.no_grad()304def prepare_audio_context(audio_connector, audio_features, attention_mask, device, dtype):305 from ltx_core.text_encoders.gemma.embeddings_processor import convert_to_additive_mask306 audio_features = audio_features.to(device=device, dtype=dtype)307 attention_mask = attention_mask.to(device=device)308 if audio_features.shape[0] > 1:309 results = []310 for i in range(audio_features.shape[0]):311 feat_i = audio_features[i:i+1]312 mask_i = attention_mask[i:i+1]313 additive = convert_to_additive_mask(mask_i, feat_i.dtype)314 enc_i, _ = audio_connector(feat_i, additive)315 results.append(enc_i)316 return torch.cat(results, dim=0)317 additive_mask = convert_to_additive_mask(attention_mask, audio_features.dtype)318 audio_encoded, _ = audio_connector(audio_features, additive_mask)319 return audio_encoded320 321 322# ─── Validation ───323 324def _unwrap_model_safe(model):325 """Strip DDP / peft wrappers without going through accelerate.unwrap_model,326 which imports deepspeed — broken in our env (torch API drift)."""327 while hasattr(model, "module"):328 model = model.module329 return model330 331 332def run_validation(lora_path, val_config_path, output_dir, step, lora_rank=128):333 """Call validate.py in a subprocess. It loads TTSServer (the same stack334 the warm server / Gradio app uses), attaches our LoRA, then iterates every335 entry in val_config with the same inference settings the user tests with.336 Single subprocess amortises the model-load cost across all val entries.337 338 Forces validation onto VAL_GPU (default "0") because training already339 occupies the rest. Override via TRAIN_VAL_GPU env var.340 """341 import subprocess342 val_dir = os.path.join(output_dir, "validation", f"step_{step:05d}")343 os.makedirs(val_dir, exist_ok=True)344 script = os.path.join(os.path.dirname(__file__), "validate.py")345 cmd = [346 sys.executable, script,347 "--val-config", val_config_path,348 "--output-dir", val_dir,349 "--lora", lora_path,350 "--lora-rank", str(lora_rank),351 # Use raw estimator output (no +10% buffer) so we can hear352 # whether the model needs more/less duration at current quality.353 "--duration-multiplier", "1.0",354 ]355 log_path = os.path.join(val_dir, "validate.log")356 env = os.environ.copy()357 # Validation needs its OWN GPU (training fills the others).358 env["CUDA_VISIBLE_DEVICES"] = os.environ.get("TRAIN_VAL_GPU", "0")359 try:360 with open(log_path, "w") as logf:361 result = subprocess.run(362 cmd, stdout=logf, stderr=subprocess.STDOUT, timeout=1800, env=env,363 )364 if result.returncode == 0:365 logging.info(f" Validation step {step}: OK → {val_dir}")366 else:367 logging.warning(f" Validation step {step} FAILED (see {log_path})")368 except subprocess.TimeoutExpired:369 logging.warning(f" Validation step {step} TIMEOUT (>30min)")370 371 372# ─── Args ───373 374def parse_args():375 # First pass: pull out --config so its values can become argparse defaults.376 cfg_parser = argparse.ArgumentParser(add_help=False)377 cfg_parser.add_argument("--config", default=None,378 help="YAML file with default values for any of the flags below. "379 "Explicit CLI flags still override the YAML.")380 cfg_args, remaining = cfg_parser.parse_known_args()381 yaml_defaults: dict = {}382 if cfg_args.config:383 import yaml as _yaml384 with open(cfg_args.config) as f:385 yaml_defaults = _yaml.safe_load(f) or {}386 # YAML keys are dashes-or-underscores → normalize to argparse dest (underscore).387 yaml_defaults = {k.replace("-", "_"): v for k, v in yaml_defaults.items()}388 389 def _yaml(name, fallback):390 return yaml_defaults.get(name, fallback)391 392 p = argparse.ArgumentParser(393 parents=[cfg_parser],394 description="Audio-Only IC-LoRA Training for Voice Cloning",395 )396 p.add_argument("--data-dir", required="data_dir" not in yaml_defaults,397 nargs="+", default=_yaml("data_dir", None))398 p.add_argument("--speaker-index", required="speaker_index" not in yaml_defaults,399 nargs="+", default=_yaml("speaker_index", None))400 p.add_argument("--output-dir", default=_yaml("output_dir", os.path.join(MODEL_DIR, "tts_iclora_v1")))401 p.add_argument("--checkpoint", default=_yaml("checkpoint", os.path.join(MODEL_DIR, "dramabox-dit-v1.safetensors")))402 p.add_argument("--full-checkpoint", default=_yaml("full_checkpoint", os.path.join(MODEL_DIR, "dramabox-audio-components.safetensors")))403 p.add_argument("--base-model", choices=["distilled", "dev"], default=_yaml("base_model", "dev"),404 help="Base model type: distilled uses DistilledTimestepSampler, dev uses ShiftedLogitNormal")405 p.add_argument("--lora-rank", type=int, default=_yaml("lora_rank", 128))406 p.add_argument("--lora-alpha", type=int, default=_yaml("lora_alpha", 128))407 p.add_argument("--lora-dropout", type=float, default=_yaml("lora_dropout", 0.0),408 help="Dropout applied to LoRA A/B matrices during training. "409 "Recommended ~0.1 for small datasets to regularize.")410 p.add_argument("--resume-lora", default=_yaml("resume_lora", None))411 p.add_argument("--resume-step-offset", type=int, default=_yaml("resume_step_offset", None),412 help="Step to add when naming saved checkpoints. If None, inferred "413 "from --resume-lora filename (e.g. lora_step_10000.safetensors → 10000). "414 "Set to 0 to start numbering at 0 regardless.")415 p.add_argument("--ref-ratio", type=float, default=_yaml("ref_ratio", 0.3),416 help="Fraction of target length to use as reference (default 0.3)")417 p.add_argument("--max-ref-tokens", type=int, default=_yaml("max_ref_tokens", 200),418 help="Maximum reference tokens after patchification (default 200)")419 p.add_argument("--text-dropout", type=float, default=_yaml("text_dropout", 0.0),420 help="Probability of dropping text conditioning (forces reliance on voice ref)")421 p.add_argument("--steps", type=int, default=_yaml("steps", 30000))422 p.add_argument("--lr", type=float, default=_yaml("lr", 3e-5))423 p.add_argument("--lr-scheduler", choices=["cosine", "linear", "constant"], default=_yaml("lr_scheduler", "cosine"))424 p.add_argument("--batch-size", type=int, default=_yaml("batch_size", 1))425 p.add_argument("--grad-accum", type=int, default=_yaml("grad_accum", 4))426 p.add_argument("--max-grad-norm", type=float, default=_yaml("max_grad_norm", 1.0))427 p.add_argument("--save-every", type=int, default=_yaml("save_every", 1000))428 p.add_argument("--log-every", type=int, default=_yaml("log_every", 50))429 p.add_argument("--seed", type=int, default=_yaml("seed", 42))430 p.add_argument("--warmup-steps", type=int, default=_yaml("warmup_steps", 100))431 p.add_argument("--val-config", default=_yaml("val_config", None))432 return p.parse_args(remaining)433 434 435# ─── Main ───436 437def main():438 from accelerate import Accelerator439 from accelerate.utils import set_seed440 441 args = parse_args()442 443 accelerator = Accelerator(444 gradient_accumulation_steps=args.grad_accum,445 mixed_precision="bf16",446 )447 448 is_main = accelerator.is_main_process449 if is_main:450 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")451 else:452 logging.basicConfig(level=logging.WARNING)453 454 set_seed(args.seed)455 device = accelerator.device456 dtype = torch.bfloat16457 458 os.makedirs(args.output_dir, exist_ok=True)459 460 # Save training args461 if is_main:462 import yaml463 args_dict = vars(args).copy()464 args_dict["_meta"] = {465 "world_size": accelerator.num_processes,466 "dtype": str(dtype),467 "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),468 "script": "train_audio_iclora.py",469 "pattern": "IC-LoRA (ref appended to end)",470 }471 with open(os.path.join(args.output_dir, "training_args.yaml"), "w") as f:472 yaml.dump(args_dict, f, default_flow_style=False, sort_keys=False)473 474 from ltx_core.components.patchifiers import AudioPatchifier475 from ltx_core.model.transformer.modality import Modality476 from ltx_core.guidance.perturbations import BatchedPerturbationConfig477 from ltx_core.tools import AudioLatentTools478 from ltx_core.types import AudioLatentShape, LatentState479 from ltx_pipelines.utils.helpers import modality_from_latent_state, timesteps_from_mask480 481 # Build speaker map482 if is_main:483 logging.info("Building speaker map...")484 speaker_map = build_speaker_map(args.speaker_index, args.data_dir)485 if is_main:486 logging.info(f"Speaker map: {len(speaker_map)} speakers, "487 f"{sum(len(v) for v in speaker_map.values())} samples")488 489 # Load model490 if is_main:491 logging.info("Loading audio-only model...")492 model = build_audio_only_model(args.checkpoint, device, dtype)493 494 if is_main:495 logging.info("Loading audio connector...")496 audio_connector = load_audio_connector(args.full_checkpoint, device, dtype)497 audio_connector.eval()498 for p in audio_connector.parameters():499 p.requires_grad = False500 501 if is_main:502 logging.info(f"Applying LoRA (rank={args.lora_rank}, alpha={args.lora_alpha})...")503 model = apply_lora(model, args.lora_rank, args.lora_alpha, args.lora_dropout)504 505 # Resume from checkpoint506 if args.resume_lora:507 from safetensors.torch import load_file as st_load508 if is_main:509 logging.info(f"Resuming from: {args.resume_lora}")510 lora_sd = st_load(args.resume_lora)511 mapped = {}512 for k, v in lora_sd.items():513 nk = k.replace(".lora_A.weight", ".lora_A.default.weight").replace(514 ".lora_B.weight", ".lora_B.default.weight")515 mapped[nk] = v516 model.load_state_dict(mapped, strict=False)517 518 # Determine step offset for save filenames. Without this, resuming a run519 # restarts step numbering at 0 and would overwrite earlier phase-1520 # checkpoints with the same save_every cadence.521 if args.resume_step_offset is None:522 resume_offset = 0523 if args.resume_lora:524 import re as _re525 m = _re.search(r"lora_step_(\d+)", os.path.basename(args.resume_lora))526 if m:527 resume_offset = int(m.group(1))528 args.resume_step_offset = resume_offset529 if is_main and args.resume_step_offset:530 logging.info(f"Save-step offset: +{args.resume_step_offset}")531 532 model.train()533 model.base_model.model.set_gradient_checkpointing(True)534 535 # Dataset & DataLoader536 dataset = IDLoRADataset(speaker_map)537 if is_main:538 logging.info(f"Dataset: {len(dataset)} samples, {len(dataset.speaker_map)} speakers")539 540 def collate_fn(batch):541 """Pad variable-length audio to max in batch, track real lengths for loss masking."""542 max_tgt_T = max(b["tgt_latent"].shape[1] for b in batch) # [C, T, F]543 max_ref_T = max(b["ref_latent"].shape[1] for b in batch)544 C = batch[0]["tgt_latent"].shape[0]545 F_dim = batch[0]["tgt_latent"].shape[2]546 547 tgt_list, ref_list, feat_list, mask_list = [], [], [], []548 tgt_lengths, ref_lengths = [], []549 550 for b in batch:551 tgt = b["tgt_latent"]552 ref = b["ref_latent"]553 tgt_lengths.append(tgt.shape[1])554 ref_lengths.append(ref.shape[1])555 556 if tgt.shape[1] < max_tgt_T:557 pad = torch.zeros(C, max_tgt_T - tgt.shape[1], F_dim, dtype=tgt.dtype)558 tgt = torch.cat([tgt, pad], dim=1)559 tgt_list.append(tgt)560 561 if ref.shape[1] < max_ref_T:562 pad = torch.zeros(C, max_ref_T - ref.shape[1], F_dim, dtype=ref.dtype)563 ref = torch.cat([ref, pad], dim=1)564 ref_list.append(ref)565 566 feat_list.append(b["audio_features"])567 mask_list.append(b["attention_mask"])568 569 return {570 "tgt_latent": torch.stack(tgt_list),571 "ref_latent": torch.stack(ref_list),572 "audio_features": torch.stack(feat_list),573 "attention_mask": torch.stack(mask_list),574 "tgt_lengths": torch.tensor(tgt_lengths),575 "ref_lengths": torch.tensor(ref_lengths),576 }577 578 dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=2,579 pin_memory=True, drop_last=True, collate_fn=collate_fn)580 581 # Optimizer & Scheduler582 optimizer = torch.optim.AdamW(583 [p for p in model.parameters() if p.requires_grad],584 lr=args.lr, betas=(0.9, 0.999), weight_decay=0.01,585 )586 587 from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR, ConstantLR588 warmup = LinearLR(optimizer, start_factor=0.01, end_factor=1.0, total_iters=args.warmup_steps)589 remaining = args.steps - args.warmup_steps590 if args.lr_scheduler == "cosine":591 # Warmup -> constant hold (20% of remaining) -> cosine decay592 hold_steps = max(remaining // 5, 0)593 decay_steps = max(remaining - hold_steps, 1)594 hold_sched = ConstantLR(optimizer, factor=1.0, total_iters=hold_steps)595 decay_sched = CosineAnnealingLR(optimizer, T_max=decay_steps, eta_min=1e-6)596 scheduler = SequentialLR(597 optimizer,598 [warmup, hold_sched, decay_sched],599 milestones=[args.warmup_steps, args.warmup_steps + hold_steps],600 )601 elif args.lr_scheduler == "linear":602 main_sched = LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=max(remaining, 1))603 scheduler = SequentialLR(optimizer, [warmup, main_sched], milestones=[args.warmup_steps])604 else:605 main_sched = ConstantLR(optimizer, factor=1.0, total_iters=max(remaining, 1))606 scheduler = SequentialLR(optimizer, [warmup, main_sched], milestones=[args.warmup_steps])607 608 # Prepare with Accelerate — but NOT the scheduler. AcceleratedScheduler609 # calls the underlying scheduler.step() `num_processes` times per sync,610 # which silently scales down our warmup/cosine spans by that factor.611 # We call scheduler.step() ourselves, gated on sync_gradients → exactly612 # one advance per optimizer step, as the yaml spec intends.613 model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)614 615 patchifier = AudioPatchifier(patch_size=1)616 617 # Select timestep sampler based on base model type618 if args.base_model == "distilled":619 timestep_sampler = DistilledTimestepSampler()620 if is_main:621 logging.info("Using DistilledTimestepSampler (matching distilled model sigmas)")622 else:623 timestep_sampler = ShiftedLogitNormalTimestepSampler()624 if is_main:625 logging.info("Using ShiftedLogitNormalTimestepSampler (dev model)")626 627 # Training loop628 if is_main:629 logging.info(f"Training: {args.steps} steps, lr={args.lr}, scheduler={args.lr_scheduler}, "630 f"batch={args.batch_size}, grad_accum={args.grad_accum}, "631 f"world_size={accelerator.num_processes}, "632 f"ref_ratio={args.ref_ratio}, max_ref_tokens={args.max_ref_tokens}")633 logging.info("IC-LoRA pattern: ref tokens APPENDED to target, loss on target only")634 635 data_iter = iter(dataloader)636 step = 0637 accum_loss = 0.0638 best_loss = float("inf")639 best_step = 0640 t0 = time.time()641 642 total_micro_steps = args.steps * args.grad_accum643 644 for micro_step in range(total_micro_steps):645 try:646 batch = next(data_iter)647 except StopIteration:648 data_iter = iter(dataloader)649 batch = next(data_iter)650 651 is_opt_step = (micro_step + 1) % args.grad_accum == 0652 if is_opt_step:653 step += 1654 655 with accelerator.accumulate(model):656 tgt_latent = batch["tgt_latent"].to(dtype=dtype) # [B, C, max_tgt_T, F]657 ref_latent = batch["ref_latent"].to(dtype=dtype) # [B, C, max_ref_T, F]658 tgt_lengths = batch["tgt_lengths"].to(device=device) # [B]659 B = tgt_latent.shape[0]660 661 # ── Random silence padding (0-1s) ── ltx_audio_tts baseline.662 # User observed reference-audio leak at end of generations when this663 # was reduced to 5 (v14) or 10 frames (v16/v17) — the model seemed664 # to use the extra target budget to regurgitate ref content. Full665 # 25 frames (0-1s avg 500ms) was apparently load-bearing for666 # regularising the boundary and reducing hallucinations.667 # Uses the real silence latent (not zeros) so the VAE decodes it as668 # true silence instead of static noise.669 max_pad_frames = 25 # ~1s at 25 latent frames/sec670 pad_frames = random.randint(0, max_pad_frames)671 if pad_frames > 0:672 C, F_dim = tgt_latent.shape[1], tgt_latent.shape[3]673 if not hasattr(args, '_silence_frame') or args._silence_frame is None:674 _sf_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "assets", "silence_latent_frame.pt")675 if os.path.exists(_sf_path):676 args._silence_frame = torch.load(_sf_path, weights_only=True) # [C, 1, F]677 if is_main:678 logging.info(f"Loaded silence latent from {_sf_path}")679 else:680 args._silence_frame = False # fallback to zeros681 if is_main:682 logging.warning(f"silence_latent_frame.pt not found, using zeros")683 if args._silence_frame is not False:684 sf = args._silence_frame.to(dtype=dtype, device=device) # [C, 1, F]685 silence_pad = sf.unsqueeze(0).expand(B, -1, pad_frames, -1) # [B, C, pad, F]686 else:687 silence_pad = torch.zeros(B, C, pad_frames, F_dim, dtype=dtype, device=device)688 tgt_latent = torch.cat([silence_pad, tgt_latent], dim=2)689 690 # Cap reference to max_ref_tokens (in latent frames, before patchification)691 # After patchification, ref_T tokens = ref frames (patch_size=1)692 ref_T_frames = min(ref_latent.shape[2], args.max_ref_tokens)693 ref_latent = ref_latent[:, :, :ref_T_frames, :]694 695 tgt_T_frames = tgt_latent.shape[2] # max (padded) target frames696 697 # ── Step 1: Create target AudioLatentShape and AudioLatentTools ──698 tgt_shape = AudioLatentShape(699 batch=B,700 channels=tgt_latent.shape[1], # 8701 frames=tgt_T_frames,702 mel_bins=tgt_latent.shape[3], # 16703 )704 705 audio_tools = AudioLatentTools(706 patchifier=patchifier,707 target_shape=tgt_shape,708 )709 710 # ── Step 2: Create initial state from target latent ──711 # create_initial_state patchifies: [B, C, T, F] -> [B, T, C*F]712 # Also creates denoise_mask=1 (all target tokens will be denoised)713 # and computes temporal positions714 state = audio_tools.create_initial_state(715 device=device,716 dtype=dtype,717 initial_latent=tgt_latent,718 )719 # state.latent: [B, tgt_T, 128], state.denoise_mask: [B, tgt_T, 1]720 # state.positions: [B, 1, tgt_T, 2]721 722 tgt_T = audio_tools.target_shape.token_count() # = tgt_T_frames723 724 # ── Step 3: Apply flow-matching noise to target BEFORE appending ref ──725 # Sample sigma726 total_tokens = tgt_T + ref_T_frames727 sigma = timestep_sampler.sample(B, total_tokens, device=device)728 sigma_exp = sigma.view(-1, 1, 1) # [B, 1, 1]729 730 noise = torch.randn_like(state.latent) # [B, tgt_T, 128]731 noisy_tgt = (1 - sigma_exp) * state.latent + sigma_exp * noise732 733 # Replace the latent in state with the noisy version734 # (clean_latent stays clean for post_process_latent pattern)735 state = LatentState(736 latent=noisy_tgt,737 denoise_mask=state.denoise_mask,738 positions=state.positions,739 clean_latent=state.clean_latent,740 attention_mask=state.attention_mask,741 )742 743 # ── Step 4: Append reference tokens using AudioConditionByReferenceLatent ──744 # This appends ref tokens to the END with denoise_mask=0 (frozen/clean)745 # Skip entirely when ref_T=0 (SFX / song samples): the model trains746 # target-only for those categories since there's no voice to clone.747 if ref_T_frames > 0:748 ref_conditioning = AudioConditionByReferenceLatent(749 latent=ref_latent,750 strength=1.0, # 1.0 = ref fully clean (denoise_mask=0)751 )752 state = ref_conditioning.apply_to(753 latent_state=state,754 latent_tools=audio_tools,755 )756 # state.latent: [B, tgt_T + ref_T, 128]757 # state.denoise_mask: [B, tgt_T + ref_T, 1]758 # target tokens: 1.0 (denoise), ref tokens: 0.0 (frozen)759 # state.positions: [B, 1, tgt_T + ref_T, 2]760 761 # ── Step 5: Build loss mask for target tokens (excluding padding) ──762 # loss_mask: 1 for real target tokens, 0 for padding and ref tokens763 loss_mask = torch.zeros(B, tgt_T, device=device)764 for b_idx in range(B):765 real_len = min(tgt_lengths[b_idx].item(), tgt_T)766 loss_mask[b_idx, :real_len] = 1.0767 768 # ── Step 6: Prepare text context ──769 # Text conditioning dropout: randomly zero out text context to force770 # the model to rely on the voice reference for identity/style.771 with torch.no_grad():772 audio_context = prepare_audio_context(773 audio_connector, batch["audio_features"],774 batch["attention_mask"], device, dtype)775 if args.text_dropout > 0 and random.random() < args.text_dropout:776 audio_context = torch.zeros_like(audio_context)777 778 # ── Step 7: Build Modality using modality_from_latent_state ──779 # timesteps = sigma * denoise_mask (ref gets 0, target gets sigma)780 audio_mod = modality_from_latent_state(781 state=state,782 context=audio_context,783 sigma=sigma,784 enabled=True,785 )786 787 # ── Step 8: Forward pass ──788 perturbations = BatchedPerturbationConfig.empty(B)789 with torch.autocast(device_type="cuda", dtype=dtype):790 _, velocity_pred = model(video=None, audio=audio_mod, perturbations=perturbations)791 792 # ── Step 9: Compute loss (IC-LoRA pattern) ──793 # Target is at the FRONT (indices 0..tgt_T), ref at the END794 # velocity target = noise - clean795 tgt_patchified = audio_tools.patchifier.patchify(tgt_latent) # [B, tgt_T, 128]796 target_velocity = noise - tgt_patchified797 798 # Extract target portion of prediction799 pred_tgt = velocity_pred[:, :tgt_T] # [B, tgt_T, 128]800 801 # MSE loss with mask: only on real target tokens (not padding or ref)802 per_token_mse = (pred_tgt - target_velocity).pow(2).mean(dim=-1) # [B, tgt_T]803 loss = per_token_mse.mul(loss_mask).div(loss_mask.mean().clamp(min=1e-6)).mean()804 805 accelerator.backward(loss)806 807 if accelerator.sync_gradients and args.max_grad_norm > 0:808 accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm)809 810 optimizer.step()811 optimizer.zero_grad()812 # Only advance the LR scheduler once per OPTIMIZER step (not per813 # micro-step). Mirrors AcceleratedOptimizer.step() which is814 # internally gated on sync_gradients.815 if accelerator.sync_gradients:816 scheduler.step()817 818 accum_loss += loss.item()819 820 # Logging & saving on optimization steps only821 if is_opt_step and step % args.log_every == 0 and is_main:822 avg_loss = accum_loss / (args.log_every * args.grad_accum)823 lr = optimizer.param_groups[0]["lr"]824 elapsed = time.time() - t0825 sps = step / elapsed if elapsed > 0 else 0826 eta = (args.steps - step) / sps if sps > 0 else 0827 logging.info(828 f"Step {step}/{args.steps} | loss={avg_loss:.4f} | lr={lr:.2e} | "829 f"tgt_T={tgt_T} ref_T={ref_T_frames} total={tgt_T + ref_T_frames} | "830 f"{sps:.1f} steps/s | ETA {eta/60:.0f}min"831 )832 833 # Save best whenever loss improves — no warmup gate, so we can834 # observe best checkpoints during warmup too.835 if avg_loss < best_loss:836 best_loss = avg_loss837 old_best = os.path.join(args.output_dir, f"best_step_{best_step:05d}.safetensors")838 best_step = step + args.resume_step_offset839 new_best = os.path.join(args.output_dir, f"best_step_{best_step:05d}.safetensors")840 unwrapped = _unwrap_model_safe(model)841 unwrapped.save_pretrained(args.output_dir)842 adapter = os.path.join(args.output_dir, "adapter_model.safetensors")843 if os.path.exists(adapter):844 shutil.copy(adapter, new_best)845 if old_best != new_best and os.path.exists(old_best):846 os.remove(old_best)847 logging.info(f"New best: loss={best_loss:.4f} at step {best_step}")848 849 accum_loss = 0.0850 851 if is_opt_step and step % args.save_every == 0 and is_main:852 global_step = step + args.resume_step_offset853 save_path = os.path.join(args.output_dir, f"lora_step_{global_step:05d}.safetensors")854 logging.info(f"Saving: {save_path}")855 unwrapped = _unwrap_model_safe(model)856 unwrapped.save_pretrained(args.output_dir)857 adapter = os.path.join(args.output_dir, "adapter_model.safetensors")858 if os.path.exists(adapter):859 shutil.copy(adapter, save_path)860 861 if args.val_config:862 logging.info(f"Running validation at step {global_step}...")863 model.eval()864 run_validation(save_path, args.val_config, args.output_dir, global_step,865 lora_rank=args.lora_rank)866 model.train()867 868 # Final save869 if is_main:870 unwrapped = _unwrap_model_safe(model)871 unwrapped.save_pretrained(args.output_dir)872 adapter = os.path.join(args.output_dir, "adapter_model.safetensors")873 global_step = step + args.resume_step_offset874 save_path = os.path.join(args.output_dir, f"lora_step_{global_step:05d}.safetensors")875 if os.path.exists(adapter):876 shutil.copy(adapter, save_path)877 logging.info(f"Training complete! {step} steps in {time.time()-t0:.0f}s")878 logging.info(f"Best loss: {best_loss:.4f} at step {best_step}")879 880 881if __name__ == "__main__":882 main()883