#!/usr/bin/env python3 """ Audio-Only IC-LoRA Training for Voice Cloning on LTX-2.3. Uses the IC-LoRA pattern: reference audio tokens are APPENDED to the end of the target sequence using AudioConditionByReferenceLatent. Loss is computed only on target tokens; reference tokens remain clean (denoise_mask=0). This follows the official video-to-video IC-LoRA strategy closely, but adapted for the audio-only modality path. Usage (single GPU): CUDA_VISIBLE_DEVICES=0 python train_audio_iclora.py --data-dir ... --speaker-index ... Usage (multi-GPU with accelerate): CUDA_VISIBLE_DEVICES=4,5,6,7 accelerate launch --num_processes=4 train_audio_iclora.py ... """ import argparse import logging import math import os import random import shutil import sys import time from collections import defaultdict from pathlib import Path import torch import torch.nn.functional as F from torch.utils.data import DataLoader, Dataset REPO_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "ltx2")) # ltx-pipelines already on path via ltx2/ MODEL_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) # Import audio conditioning item from our module sys.path.insert(0, MODEL_DIR) from audio_conditioning import AudioConditionByReferenceLatent # ─── Timestep Sampling ─── class DistilledTimestepSampler: """Sample timesteps from the distilled sigma schedule. The distilled model was trained to denoise at these specific sigma values. We sample uniformly from the intervals between consecutive sigmas, matching the distribution the model actually operates on. """ # Distilled 8-step sigma values (boundaries of denoising intervals) SIGMAS = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0] def __init__(self, jitter: float = 0.02): self.jitter = jitter def sample(self, batch_size: int, seq_length: int = None, device: torch.device = None) -> torch.Tensor: n_intervals = len(self.SIGMAS) - 1 interval_idx = torch.randint(0, n_intervals, (batch_size,), device=device) t = torch.rand(batch_size, device=device) sigma_high = torch.tensor([self.SIGMAS[i] for i in interval_idx], device=device) sigma_low = torch.tensor([self.SIGMAS[i + 1] for i in interval_idx], device=device) sigma = sigma_low + t * (sigma_high - sigma_low) return sigma.clamp(0.01, 0.99) class ShiftedLogitNormalTimestepSampler: """Shifted logit-normal distribution, shift depends on sequence length.""" def __init__(self, std: float = 1.0, eps: float = 1e-3, uniform_prob: float = 0.1): self.std = std self.eps = eps self.uniform_prob = uniform_prob self.normal_999_percentile = 3.0902 * std self.normal_005_percentile = -2.5758 * std def sample(self, batch_size: int, seq_length: int, device: torch.device = None) -> torch.Tensor: mu = self._get_shift(seq_length) normal = torch.randn(batch_size, device=device) * self.std + mu logitnormal = torch.sigmoid(normal) p999 = torch.sigmoid(torch.tensor(mu + self.normal_999_percentile, device=device)) p005 = torch.sigmoid(torch.tensor(mu + self.normal_005_percentile, device=device)) stretched = (logitnormal - p005) / (p999 - p005) stretched = torch.where(stretched >= self.eps, stretched, 2 * self.eps - stretched) stretched = stretched.clamp(0, 1) uniform = (1 - self.eps) * torch.rand(batch_size, device=device) + self.eps prob = torch.rand(batch_size, device=device) return torch.where(prob > self.uniform_prob, stretched, uniform) @staticmethod def _get_shift(seq_length, min_tok=1024, max_tok=4096, min_s=0.95, max_s=2.05): m = (max_s - min_s) / (max_tok - min_tok) return m * seq_length + (min_s - m * min_tok) # ─── Dataset ─── def build_speaker_map(index_paths, data_dirs): """Map speaker → [(data_dir, sample_idx)] from index file(s). The sample index comes from field 0 of the `~`-delimited row when it parses as int (allows subset indexes that keep original sample numbers), otherwise we fall back to the row's line number (legacy behaviour for string-keyed indexes like tts_training_data_podcast). """ speaker_to_samples = defaultdict(list) for index_path, data_dir in zip(index_paths, data_dirs): with open(index_path) as f: for line_num, line in enumerate(f): parts = line.strip().split("~") if len(parts) < 7: continue try: idx = int(parts[0]) except ValueError: idx = line_num speaker_id = parts[1] speaker_to_samples[speaker_id].append((data_dir, idx)) return {k: v for k, v in speaker_to_samples.items() if len(v) >= 2} class IDLoRADataset(Dataset): # Silence-latent reference loaded once, used to detect and strip any # leading silence frames baked into the preprocessed audio_latents. The # training loop ALREADY prepends 0-25 random silence frames, so we don't # want accidental silence in the source data compounding on top. _silence_ref = None @classmethod def _load_silence_ref(cls): if cls._silence_ref is None: p = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "assets", "silence_latent_frame.pt") if os.path.exists(p): cls._silence_ref = torch.load(p, weights_only=True).float().squeeze() # [C, F] return cls._silence_ref def __init__(self, speaker_map): self.samples = [] self.speaker_map = {} for speaker, entries in speaker_map.items(): valid = [] for data_dir, idx in entries: audio_path = Path(data_dir) / "audio_latents" / f"sample_{idx:06d}.pt" cond_path = Path(data_dir) / "conditions" / f"sample_{idx:06d}.pt" if audio_path.exists() and cond_path.exists(): valid.append((data_dir, idx)) if len(valid) >= 2: self.speaker_map[speaker] = valid for speaker, entries in self.speaker_map.items(): for entry in entries: self.samples.append((entry, speaker)) IDLoRADataset._load_silence_ref() def __len__(self): return len(self.samples) def _load_sample(self, data_dir, idx): base = Path(data_dir) audio = torch.load(base / "audio_latents" / f"sample_{idx:06d}.pt", weights_only=False) # Prefer prefix-stripped text embeddings if they exist (re-encoded with # just the quoted dialogue, dropping the "A woman says, " / "A man # speaks with X accent, " scene-description prefix). stripped = base / "conditions_stripped" / f"sample_{idx:06d}.pt" cond_path = stripped if stripped.exists() else base / "conditions" / f"sample_{idx:06d}.pt" cond = torch.load(cond_path, weights_only=False) if isinstance(audio, dict): audio = audio.get("audio_latent", audio.get("latent", list(audio.values())[0])) if audio.dim() == 2: audio = audio.unsqueeze(0) audio_feats = cond.get("audio_prompt_embeds", cond.get("prompt_embeds")) attn_mask = cond.get("prompt_attention_mask") # The audio_connector has num_learnable_registers=128 and asserts the # input sequence length is divisible by 128. Our new preprocessing # saved trimmed conditions (dropping left-padding to save disk), which # produces short/irregular sequence lengths. Left-pad back to the next # multiple of 128 with zeros (matching the tokenizer's left-padding # convention) so this assertion holds. REG = 128 L = audio_feats.shape[0] target_L = ((L + REG - 1) // REG) * REG if target_L != L: pad_len = target_L - L pad_emb = torch.zeros(pad_len, audio_feats.shape[1], dtype=audio_feats.dtype) pad_mask = torch.zeros(pad_len, dtype=attn_mask.dtype) audio_feats = torch.cat([pad_emb, audio_feats], dim=0) attn_mask = torch.cat([pad_mask, attn_mask], dim=0) return audio, audio_feats, attn_mask def __getitem__(self, idx): (data_dir, tgt_idx), speaker = self.samples[idx] tgt_latent, audio_feats, attn_mask = self._load_sample(data_dir, tgt_idx) # Drop the reference entirely for non-voice-cloning categories: # - SFX samples (speaker starts with "sfx_"): descriptive sound events, # no speaker identity to clone. # - Song/music samples (suno dataset): prompts describe the music style, # reference audio doesn't transfer anything useful. # Return a zero-length ref so the model trains target-only for these. drop_ref = speaker.startswith("sfx_") or "preprocessed_ltx_suno" in str(data_dir) if drop_ref: C, F_dim = tgt_latent.shape[0], tgt_latent.shape[2] ref_latent = torch.zeros(C, 0, F_dim, dtype=tgt_latent.dtype) else: entries = self.speaker_map[speaker] ref_entry = random.choice([e for e in entries if e[1] != tgt_idx]) ref_latent, _, _ = self._load_sample(*ref_entry) return { "tgt_latent": tgt_latent, "ref_latent": ref_latent, "audio_features": audio_feats, "attention_mask": attn_mask, } # ─── Model building ─── def build_audio_only_model(checkpoint_path, device, dtype): from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder from ltx_core.loader.registry import DummyRegistry from ltx_core.loader.sd_ops import SDOps from ltx_core.model.transformer.model import LTXModel, LTXModelType from ltx_core.model.model_protocol import ModelConfigurator from ltx_core.model.transformer.attention import AttentionFunction from ltx_core.model.transformer.rope import LTXRopeType sd_ops = SDOps("AO").with_matching(prefix="model.diffusion_model.").with_replacement("model.diffusion_model.", "") class Cfg(ModelConfigurator[LTXModel]): @classmethod def from_config(cls, config): t = config.get("transformer", {}) cp = None if not t.get("caption_proj_before_connector", False): from ltx_core.model.transformer.text_projection import create_caption_projection with torch.device("meta"): cp = create_caption_projection(t, audio=True) return LTXModel( model_type=LTXModelType.AudioOnly, audio_num_attention_heads=t.get("audio_num_attention_heads", 32), audio_attention_head_dim=t.get("audio_attention_head_dim", 64), audio_in_channels=t.get("audio_in_channels", 128), audio_out_channels=t.get("audio_out_channels", 128), num_layers=t.get("num_layers", 48), audio_cross_attention_dim=t.get("audio_cross_attention_dim", 2048), norm_eps=t.get("norm_eps", 1e-6), attention_type=AttentionFunction(t.get("attention_type", "default")), positional_embedding_theta=t.get("positional_embedding_theta", 10000.0), audio_positional_embedding_max_pos=t.get("audio_positional_embedding_max_pos", [20]), timestep_scale_multiplier=t.get("timestep_scale_multiplier", 1000), use_middle_indices_grid=t.get("use_middle_indices_grid", True), rope_type=LTXRopeType(t.get("rope_type", "interleaved")), double_precision_rope=t.get("frequencies_precision", False) == "float64", apply_gated_attention=t.get("apply_gated_attention", False), audio_caption_projection=cp, cross_attention_adaln=t.get("cross_attention_adaln", False), ) builder = Builder(model_path=checkpoint_path, model_class_configurator=Cfg, model_sd_ops=sd_ops, registry=DummyRegistry()) return builder.build(device=device, dtype=dtype) def load_audio_connector(checkpoint_path, device, dtype): # ltx-trainer already on path via ltx2/ from ltx_trainer.model_loader import load_embeddings_processor emb_proc = load_embeddings_processor(checkpoint_path, device=device, dtype=dtype) connector = emb_proc.audio_connector del emb_proc return connector def apply_lora(model, rank, alpha, dropout=0.0): from peft import LoraConfig, get_peft_model config = LoraConfig( r=rank, lora_alpha=alpha, lora_dropout=dropout, bias="none", target_modules=[ # Self-attention over audio tokens (voice-transfer pathway via ref). "audio_attn1.to_k", "audio_attn1.to_q", "audio_attn1.to_v", "audio_attn1.to_out.0", # Cross-attention (audio ↔ text context) NOT adapted — keep base # model's prompt→audio behaviour intact and rely on dataset balance # to drive expressiveness. (v15c tried this with adaLN unfreeze, # that proved too destructive; v16 tries it adaLN-frozen.) # FFN — non-linear capacity for style/phonetic adaptation. "audio_ff.net.0.proj", "audio_ff.net.2", ], ) model = get_peft_model(model, config) trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) total = sum(p.numel() for p in model.parameters()) logging.info(f"LoRA: {trainable:,} trainable / {total:,} total ({100*trainable/total:.1f}%)") return model @torch.no_grad() def prepare_audio_context(audio_connector, audio_features, attention_mask, device, dtype): from ltx_core.text_encoders.gemma.embeddings_processor import convert_to_additive_mask audio_features = audio_features.to(device=device, dtype=dtype) attention_mask = attention_mask.to(device=device) if audio_features.shape[0] > 1: results = [] for i in range(audio_features.shape[0]): feat_i = audio_features[i:i+1] mask_i = attention_mask[i:i+1] additive = convert_to_additive_mask(mask_i, feat_i.dtype) enc_i, _ = audio_connector(feat_i, additive) results.append(enc_i) return torch.cat(results, dim=0) additive_mask = convert_to_additive_mask(attention_mask, audio_features.dtype) audio_encoded, _ = audio_connector(audio_features, additive_mask) return audio_encoded # ─── Validation ─── def _unwrap_model_safe(model): """Strip DDP / peft wrappers without going through accelerate.unwrap_model, which imports deepspeed — broken in our env (torch API drift).""" while hasattr(model, "module"): model = model.module return model def run_validation(lora_path, val_config_path, output_dir, step, lora_rank=128): """Call validate.py in a subprocess. It loads TTSServer (the same stack the warm server / Gradio app uses), attaches our LoRA, then iterates every entry in val_config with the same inference settings the user tests with. Single subprocess amortises the model-load cost across all val entries. Forces validation onto VAL_GPU (default "0") because training already occupies the rest. Override via TRAIN_VAL_GPU env var. """ import subprocess val_dir = os.path.join(output_dir, "validation", f"step_{step:05d}") os.makedirs(val_dir, exist_ok=True) script = os.path.join(os.path.dirname(__file__), "validate.py") cmd = [ sys.executable, script, "--val-config", val_config_path, "--output-dir", val_dir, "--lora", lora_path, "--lora-rank", str(lora_rank), # Use raw estimator output (no +10% buffer) so we can hear # whether the model needs more/less duration at current quality. "--duration-multiplier", "1.0", ] log_path = os.path.join(val_dir, "validate.log") env = os.environ.copy() # Validation needs its OWN GPU (training fills the others). env["CUDA_VISIBLE_DEVICES"] = os.environ.get("TRAIN_VAL_GPU", "0") try: with open(log_path, "w") as logf: result = subprocess.run( cmd, stdout=logf, stderr=subprocess.STDOUT, timeout=1800, env=env, ) if result.returncode == 0: logging.info(f" Validation step {step}: OK → {val_dir}") else: logging.warning(f" Validation step {step} FAILED (see {log_path})") except subprocess.TimeoutExpired: logging.warning(f" Validation step {step} TIMEOUT (>30min)") # ─── Args ─── def parse_args(): # First pass: pull out --config so its values can become argparse defaults. cfg_parser = argparse.ArgumentParser(add_help=False) cfg_parser.add_argument("--config", default=None, help="YAML file with default values for any of the flags below. " "Explicit CLI flags still override the YAML.") cfg_args, remaining = cfg_parser.parse_known_args() yaml_defaults: dict = {} if cfg_args.config: import yaml as _yaml with open(cfg_args.config) as f: yaml_defaults = _yaml.safe_load(f) or {} # YAML keys are dashes-or-underscores → normalize to argparse dest (underscore). yaml_defaults = {k.replace("-", "_"): v for k, v in yaml_defaults.items()} def _yaml(name, fallback): return yaml_defaults.get(name, fallback) p = argparse.ArgumentParser( parents=[cfg_parser], description="Audio-Only IC-LoRA Training for Voice Cloning", ) p.add_argument("--data-dir", required="data_dir" not in yaml_defaults, nargs="+", default=_yaml("data_dir", None)) p.add_argument("--speaker-index", required="speaker_index" not in yaml_defaults, nargs="+", default=_yaml("speaker_index", None)) p.add_argument("--output-dir", default=_yaml("output_dir", os.path.join(MODEL_DIR, "tts_iclora_v1"))) p.add_argument("--checkpoint", default=_yaml("checkpoint", os.path.join(MODEL_DIR, "dramabox-dit-v1.safetensors"))) p.add_argument("--full-checkpoint", default=_yaml("full_checkpoint", os.path.join(MODEL_DIR, "dramabox-audio-components.safetensors"))) p.add_argument("--base-model", choices=["distilled", "dev"], default=_yaml("base_model", "dev"), help="Base model type: distilled uses DistilledTimestepSampler, dev uses ShiftedLogitNormal") p.add_argument("--lora-rank", type=int, default=_yaml("lora_rank", 128)) p.add_argument("--lora-alpha", type=int, default=_yaml("lora_alpha", 128)) p.add_argument("--lora-dropout", type=float, default=_yaml("lora_dropout", 0.0), help="Dropout applied to LoRA A/B matrices during training. " "Recommended ~0.1 for small datasets to regularize.") p.add_argument("--resume-lora", default=_yaml("resume_lora", None)) p.add_argument("--resume-step-offset", type=int, default=_yaml("resume_step_offset", None), help="Step to add when naming saved checkpoints. If None, inferred " "from --resume-lora filename (e.g. lora_step_10000.safetensors → 10000). " "Set to 0 to start numbering at 0 regardless.") p.add_argument("--ref-ratio", type=float, default=_yaml("ref_ratio", 0.3), help="Fraction of target length to use as reference (default 0.3)") p.add_argument("--max-ref-tokens", type=int, default=_yaml("max_ref_tokens", 200), help="Maximum reference tokens after patchification (default 200)") p.add_argument("--text-dropout", type=float, default=_yaml("text_dropout", 0.0), help="Probability of dropping text conditioning (forces reliance on voice ref)") p.add_argument("--steps", type=int, default=_yaml("steps", 30000)) p.add_argument("--lr", type=float, default=_yaml("lr", 3e-5)) p.add_argument("--lr-scheduler", choices=["cosine", "linear", "constant"], default=_yaml("lr_scheduler", "cosine")) p.add_argument("--batch-size", type=int, default=_yaml("batch_size", 1)) p.add_argument("--grad-accum", type=int, default=_yaml("grad_accum", 4)) p.add_argument("--max-grad-norm", type=float, default=_yaml("max_grad_norm", 1.0)) p.add_argument("--save-every", type=int, default=_yaml("save_every", 1000)) p.add_argument("--log-every", type=int, default=_yaml("log_every", 50)) p.add_argument("--seed", type=int, default=_yaml("seed", 42)) p.add_argument("--warmup-steps", type=int, default=_yaml("warmup_steps", 100)) p.add_argument("--val-config", default=_yaml("val_config", None)) return p.parse_args(remaining) # ─── Main ─── def main(): from accelerate import Accelerator from accelerate.utils import set_seed args = parse_args() accelerator = Accelerator( gradient_accumulation_steps=args.grad_accum, mixed_precision="bf16", ) is_main = accelerator.is_main_process if is_main: logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") else: logging.basicConfig(level=logging.WARNING) set_seed(args.seed) device = accelerator.device dtype = torch.bfloat16 os.makedirs(args.output_dir, exist_ok=True) # Save training args if is_main: import yaml args_dict = vars(args).copy() args_dict["_meta"] = { "world_size": accelerator.num_processes, "dtype": str(dtype), "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"), "script": "train_audio_iclora.py", "pattern": "IC-LoRA (ref appended to end)", } with open(os.path.join(args.output_dir, "training_args.yaml"), "w") as f: yaml.dump(args_dict, f, default_flow_style=False, sort_keys=False) from ltx_core.components.patchifiers import AudioPatchifier from ltx_core.model.transformer.modality import Modality from ltx_core.guidance.perturbations import BatchedPerturbationConfig from ltx_core.tools import AudioLatentTools from ltx_core.types import AudioLatentShape, LatentState from ltx_pipelines.utils.helpers import modality_from_latent_state, timesteps_from_mask # Build speaker map if is_main: logging.info("Building speaker map...") speaker_map = build_speaker_map(args.speaker_index, args.data_dir) if is_main: logging.info(f"Speaker map: {len(speaker_map)} speakers, " f"{sum(len(v) for v in speaker_map.values())} samples") # Load model if is_main: logging.info("Loading audio-only model...") model = build_audio_only_model(args.checkpoint, device, dtype) if is_main: logging.info("Loading audio connector...") audio_connector = load_audio_connector(args.full_checkpoint, device, dtype) audio_connector.eval() for p in audio_connector.parameters(): p.requires_grad = False if is_main: logging.info(f"Applying LoRA (rank={args.lora_rank}, alpha={args.lora_alpha})...") model = apply_lora(model, args.lora_rank, args.lora_alpha, args.lora_dropout) # Resume from checkpoint if args.resume_lora: from safetensors.torch import load_file as st_load if is_main: logging.info(f"Resuming from: {args.resume_lora}") lora_sd = st_load(args.resume_lora) mapped = {} for k, v in lora_sd.items(): nk = k.replace(".lora_A.weight", ".lora_A.default.weight").replace( ".lora_B.weight", ".lora_B.default.weight") mapped[nk] = v model.load_state_dict(mapped, strict=False) # Determine step offset for save filenames. Without this, resuming a run # restarts step numbering at 0 and would overwrite earlier phase-1 # checkpoints with the same save_every cadence. if args.resume_step_offset is None: resume_offset = 0 if args.resume_lora: import re as _re m = _re.search(r"lora_step_(\d+)", os.path.basename(args.resume_lora)) if m: resume_offset = int(m.group(1)) args.resume_step_offset = resume_offset if is_main and args.resume_step_offset: logging.info(f"Save-step offset: +{args.resume_step_offset}") model.train() model.base_model.model.set_gradient_checkpointing(True) # Dataset & DataLoader dataset = IDLoRADataset(speaker_map) if is_main: logging.info(f"Dataset: {len(dataset)} samples, {len(dataset.speaker_map)} speakers") def collate_fn(batch): """Pad variable-length audio to max in batch, track real lengths for loss masking.""" max_tgt_T = max(b["tgt_latent"].shape[1] for b in batch) # [C, T, F] max_ref_T = max(b["ref_latent"].shape[1] for b in batch) C = batch[0]["tgt_latent"].shape[0] F_dim = batch[0]["tgt_latent"].shape[2] tgt_list, ref_list, feat_list, mask_list = [], [], [], [] tgt_lengths, ref_lengths = [], [] for b in batch: tgt = b["tgt_latent"] ref = b["ref_latent"] tgt_lengths.append(tgt.shape[1]) ref_lengths.append(ref.shape[1]) if tgt.shape[1] < max_tgt_T: pad = torch.zeros(C, max_tgt_T - tgt.shape[1], F_dim, dtype=tgt.dtype) tgt = torch.cat([tgt, pad], dim=1) tgt_list.append(tgt) if ref.shape[1] < max_ref_T: pad = torch.zeros(C, max_ref_T - ref.shape[1], F_dim, dtype=ref.dtype) ref = torch.cat([ref, pad], dim=1) ref_list.append(ref) feat_list.append(b["audio_features"]) mask_list.append(b["attention_mask"]) return { "tgt_latent": torch.stack(tgt_list), "ref_latent": torch.stack(ref_list), "audio_features": torch.stack(feat_list), "attention_mask": torch.stack(mask_list), "tgt_lengths": torch.tensor(tgt_lengths), "ref_lengths": torch.tensor(ref_lengths), } dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=2, pin_memory=True, drop_last=True, collate_fn=collate_fn) # Optimizer & Scheduler optimizer = torch.optim.AdamW( [p for p in model.parameters() if p.requires_grad], lr=args.lr, betas=(0.9, 0.999), weight_decay=0.01, ) from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR, ConstantLR warmup = LinearLR(optimizer, start_factor=0.01, end_factor=1.0, total_iters=args.warmup_steps) remaining = args.steps - args.warmup_steps if args.lr_scheduler == "cosine": # Warmup -> constant hold (20% of remaining) -> cosine decay hold_steps = max(remaining // 5, 0) decay_steps = max(remaining - hold_steps, 1) hold_sched = ConstantLR(optimizer, factor=1.0, total_iters=hold_steps) decay_sched = CosineAnnealingLR(optimizer, T_max=decay_steps, eta_min=1e-6) scheduler = SequentialLR( optimizer, [warmup, hold_sched, decay_sched], milestones=[args.warmup_steps, args.warmup_steps + hold_steps], ) elif args.lr_scheduler == "linear": main_sched = LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=max(remaining, 1)) scheduler = SequentialLR(optimizer, [warmup, main_sched], milestones=[args.warmup_steps]) else: main_sched = ConstantLR(optimizer, factor=1.0, total_iters=max(remaining, 1)) scheduler = SequentialLR(optimizer, [warmup, main_sched], milestones=[args.warmup_steps]) # Prepare with Accelerate — but NOT the scheduler. AcceleratedScheduler # calls the underlying scheduler.step() `num_processes` times per sync, # which silently scales down our warmup/cosine spans by that factor. # We call scheduler.step() ourselves, gated on sync_gradients → exactly # one advance per optimizer step, as the yaml spec intends. model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader) patchifier = AudioPatchifier(patch_size=1) # Select timestep sampler based on base model type if args.base_model == "distilled": timestep_sampler = DistilledTimestepSampler() if is_main: logging.info("Using DistilledTimestepSampler (matching distilled model sigmas)") else: timestep_sampler = ShiftedLogitNormalTimestepSampler() if is_main: logging.info("Using ShiftedLogitNormalTimestepSampler (dev model)") # Training loop if is_main: logging.info(f"Training: {args.steps} steps, lr={args.lr}, scheduler={args.lr_scheduler}, " f"batch={args.batch_size}, grad_accum={args.grad_accum}, " f"world_size={accelerator.num_processes}, " f"ref_ratio={args.ref_ratio}, max_ref_tokens={args.max_ref_tokens}") logging.info("IC-LoRA pattern: ref tokens APPENDED to target, loss on target only") data_iter = iter(dataloader) step = 0 accum_loss = 0.0 best_loss = float("inf") best_step = 0 t0 = time.time() total_micro_steps = args.steps * args.grad_accum for micro_step in range(total_micro_steps): try: batch = next(data_iter) except StopIteration: data_iter = iter(dataloader) batch = next(data_iter) is_opt_step = (micro_step + 1) % args.grad_accum == 0 if is_opt_step: step += 1 with accelerator.accumulate(model): tgt_latent = batch["tgt_latent"].to(dtype=dtype) # [B, C, max_tgt_T, F] ref_latent = batch["ref_latent"].to(dtype=dtype) # [B, C, max_ref_T, F] tgt_lengths = batch["tgt_lengths"].to(device=device) # [B] B = tgt_latent.shape[0] # ── Random silence padding (0-1s) ── ltx_audio_tts baseline. # User observed reference-audio leak at end of generations when this # was reduced to 5 (v14) or 10 frames (v16/v17) — the model seemed # to use the extra target budget to regurgitate ref content. Full # 25 frames (0-1s avg 500ms) was apparently load-bearing for # regularising the boundary and reducing hallucinations. # Uses the real silence latent (not zeros) so the VAE decodes it as # true silence instead of static noise. max_pad_frames = 25 # ~1s at 25 latent frames/sec pad_frames = random.randint(0, max_pad_frames) if pad_frames > 0: C, F_dim = tgt_latent.shape[1], tgt_latent.shape[3] if not hasattr(args, '_silence_frame') or args._silence_frame is None: _sf_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "assets", "silence_latent_frame.pt") if os.path.exists(_sf_path): args._silence_frame = torch.load(_sf_path, weights_only=True) # [C, 1, F] if is_main: logging.info(f"Loaded silence latent from {_sf_path}") else: args._silence_frame = False # fallback to zeros if is_main: logging.warning(f"silence_latent_frame.pt not found, using zeros") if args._silence_frame is not False: sf = args._silence_frame.to(dtype=dtype, device=device) # [C, 1, F] silence_pad = sf.unsqueeze(0).expand(B, -1, pad_frames, -1) # [B, C, pad, F] else: silence_pad = torch.zeros(B, C, pad_frames, F_dim, dtype=dtype, device=device) tgt_latent = torch.cat([silence_pad, tgt_latent], dim=2) # Cap reference to max_ref_tokens (in latent frames, before patchification) # After patchification, ref_T tokens = ref frames (patch_size=1) ref_T_frames = min(ref_latent.shape[2], args.max_ref_tokens) ref_latent = ref_latent[:, :, :ref_T_frames, :] tgt_T_frames = tgt_latent.shape[2] # max (padded) target frames # ── Step 1: Create target AudioLatentShape and AudioLatentTools ── tgt_shape = AudioLatentShape( batch=B, channels=tgt_latent.shape[1], # 8 frames=tgt_T_frames, mel_bins=tgt_latent.shape[3], # 16 ) audio_tools = AudioLatentTools( patchifier=patchifier, target_shape=tgt_shape, ) # ── Step 2: Create initial state from target latent ── # create_initial_state patchifies: [B, C, T, F] -> [B, T, C*F] # Also creates denoise_mask=1 (all target tokens will be denoised) # and computes temporal positions state = audio_tools.create_initial_state( device=device, dtype=dtype, initial_latent=tgt_latent, ) # state.latent: [B, tgt_T, 128], state.denoise_mask: [B, tgt_T, 1] # state.positions: [B, 1, tgt_T, 2] tgt_T = audio_tools.target_shape.token_count() # = tgt_T_frames # ── Step 3: Apply flow-matching noise to target BEFORE appending ref ── # Sample sigma total_tokens = tgt_T + ref_T_frames sigma = timestep_sampler.sample(B, total_tokens, device=device) sigma_exp = sigma.view(-1, 1, 1) # [B, 1, 1] noise = torch.randn_like(state.latent) # [B, tgt_T, 128] noisy_tgt = (1 - sigma_exp) * state.latent + sigma_exp * noise # Replace the latent in state with the noisy version # (clean_latent stays clean for post_process_latent pattern) state = LatentState( latent=noisy_tgt, denoise_mask=state.denoise_mask, positions=state.positions, clean_latent=state.clean_latent, attention_mask=state.attention_mask, ) # ── Step 4: Append reference tokens using AudioConditionByReferenceLatent ── # This appends ref tokens to the END with denoise_mask=0 (frozen/clean) # Skip entirely when ref_T=0 (SFX / song samples): the model trains # target-only for those categories since there's no voice to clone. if ref_T_frames > 0: ref_conditioning = AudioConditionByReferenceLatent( latent=ref_latent, strength=1.0, # 1.0 = ref fully clean (denoise_mask=0) ) state = ref_conditioning.apply_to( latent_state=state, latent_tools=audio_tools, ) # state.latent: [B, tgt_T + ref_T, 128] # state.denoise_mask: [B, tgt_T + ref_T, 1] # target tokens: 1.0 (denoise), ref tokens: 0.0 (frozen) # state.positions: [B, 1, tgt_T + ref_T, 2] # ── Step 5: Build loss mask for target tokens (excluding padding) ── # loss_mask: 1 for real target tokens, 0 for padding and ref tokens loss_mask = torch.zeros(B, tgt_T, device=device) for b_idx in range(B): real_len = min(tgt_lengths[b_idx].item(), tgt_T) loss_mask[b_idx, :real_len] = 1.0 # ── Step 6: Prepare text context ── # Text conditioning dropout: randomly zero out text context to force # the model to rely on the voice reference for identity/style. with torch.no_grad(): audio_context = prepare_audio_context( audio_connector, batch["audio_features"], batch["attention_mask"], device, dtype) if args.text_dropout > 0 and random.random() < args.text_dropout: audio_context = torch.zeros_like(audio_context) # ── Step 7: Build Modality using modality_from_latent_state ── # timesteps = sigma * denoise_mask (ref gets 0, target gets sigma) audio_mod = modality_from_latent_state( state=state, context=audio_context, sigma=sigma, enabled=True, ) # ── Step 8: Forward pass ── perturbations = BatchedPerturbationConfig.empty(B) with torch.autocast(device_type="cuda", dtype=dtype): _, velocity_pred = model(video=None, audio=audio_mod, perturbations=perturbations) # ── Step 9: Compute loss (IC-LoRA pattern) ── # Target is at the FRONT (indices 0..tgt_T), ref at the END # velocity target = noise - clean tgt_patchified = audio_tools.patchifier.patchify(tgt_latent) # [B, tgt_T, 128] target_velocity = noise - tgt_patchified # Extract target portion of prediction pred_tgt = velocity_pred[:, :tgt_T] # [B, tgt_T, 128] # MSE loss with mask: only on real target tokens (not padding or ref) per_token_mse = (pred_tgt - target_velocity).pow(2).mean(dim=-1) # [B, tgt_T] loss = per_token_mse.mul(loss_mask).div(loss_mask.mean().clamp(min=1e-6)).mean() accelerator.backward(loss) if accelerator.sync_gradients and args.max_grad_norm > 0: accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm) optimizer.step() optimizer.zero_grad() # Only advance the LR scheduler once per OPTIMIZER step (not per # micro-step). Mirrors AcceleratedOptimizer.step() which is # internally gated on sync_gradients. if accelerator.sync_gradients: scheduler.step() accum_loss += loss.item() # Logging & saving on optimization steps only if is_opt_step and step % args.log_every == 0 and is_main: avg_loss = accum_loss / (args.log_every * args.grad_accum) lr = optimizer.param_groups[0]["lr"] elapsed = time.time() - t0 sps = step / elapsed if elapsed > 0 else 0 eta = (args.steps - step) / sps if sps > 0 else 0 logging.info( f"Step {step}/{args.steps} | loss={avg_loss:.4f} | lr={lr:.2e} | " f"tgt_T={tgt_T} ref_T={ref_T_frames} total={tgt_T + ref_T_frames} | " f"{sps:.1f} steps/s | ETA {eta/60:.0f}min" ) # Save best whenever loss improves — no warmup gate, so we can # observe best checkpoints during warmup too. if avg_loss < best_loss: best_loss = avg_loss old_best = os.path.join(args.output_dir, f"best_step_{best_step:05d}.safetensors") best_step = step + args.resume_step_offset new_best = os.path.join(args.output_dir, f"best_step_{best_step:05d}.safetensors") unwrapped = _unwrap_model_safe(model) unwrapped.save_pretrained(args.output_dir) adapter = os.path.join(args.output_dir, "adapter_model.safetensors") if os.path.exists(adapter): shutil.copy(adapter, new_best) if old_best != new_best and os.path.exists(old_best): os.remove(old_best) logging.info(f"New best: loss={best_loss:.4f} at step {best_step}") accum_loss = 0.0 if is_opt_step and step % args.save_every == 0 and is_main: global_step = step + args.resume_step_offset save_path = os.path.join(args.output_dir, f"lora_step_{global_step:05d}.safetensors") logging.info(f"Saving: {save_path}") unwrapped = _unwrap_model_safe(model) unwrapped.save_pretrained(args.output_dir) adapter = os.path.join(args.output_dir, "adapter_model.safetensors") if os.path.exists(adapter): shutil.copy(adapter, save_path) if args.val_config: logging.info(f"Running validation at step {global_step}...") model.eval() run_validation(save_path, args.val_config, args.output_dir, global_step, lora_rank=args.lora_rank) model.train() # Final save if is_main: unwrapped = _unwrap_model_safe(model) unwrapped.save_pretrained(args.output_dir) adapter = os.path.join(args.output_dir, "adapter_model.safetensors") global_step = step + args.resume_step_offset save_path = os.path.join(args.output_dir, f"lora_step_{global_step:05d}.safetensors") if os.path.exists(adapter): shutil.copy(adapter, save_path) logging.info(f"Training complete! {step} steps in {time.time()-t0:.0f}s") logging.info(f"Best loss: {best_loss:.4f} at step {best_step}") if __name__ == "__main__": main()