1from __future__ import annotations

2

3import argparse

4import hashlib

5import json

6import os

7import random

8import re

9import shutil

10import time

11from pathlib import Path

12from typing import Any, Dict, Sequence, Tuple

13

14import librosa

15import numpy as np

16import soundfile as sf

17import torch

18from transformers import AutoFeatureExtractor, Trainer, TrainerCallback

19

20

21TARGET_SR = 16000

22STEREO_CHANNELS = 2

23MONO_FEATURE_BINS = 80

24STEREO_FEATURE_BINS = MONO_FEATURE_BINS * STEREO_CHANNELS

25AUDIO_CHANNEL_MODE = "stereo_lr_logmel_concat"

26

27HF_CACHE_MODEL_RE = re.compile(r"models--([^/\\]+)--([^/\\]+)", flags=re.IGNORECASE)

28OPTIMIZER_STATE_NAME = "optimizer.pt"

29SCHEDULER_STATE_NAME = "scheduler.pt"

30SCALER_STATE_NAME = "scaler.pt"

31TRAINER_STATE_NAME = "trainer_state.json"

32RNG_STATE_NAME = "rng_state.pth"

33TRAINING_ARGS_NAME = "training_args.bin"

34

35

36def str2bool(value: str | bool) -> bool:

37 if isinstance(value, bool):

38 return value

39 normalized = value.strip().lower()

40 if normalized in {"1", "true", "t", "yes", "y"}:

41 return True

42 if normalized in {"0", "false", "f", "no", "n"}:

43 return False

44 raise argparse.ArgumentTypeError(f"Invalid boolean value: {value}")

45

46

47def sanitize_experiment_name(raw: str) -> str:

48 cleaned = re.sub(r"[^A-Za-z0-9._-]+", "_", str(raw).strip())

49 cleaned = cleaned.strip("._-")

50 return cleaned or "Phase2_ParamSeq_WhisperBase_Experiment"

51

52

53def resolve_pause_paths(args: argparse.Namespace, output_dir: Path) -> Tuple[Path, Path]:

54 pause_flag_file = (

55 Path(str(args.pause_flag_file)).expanduser().resolve()

56 if str(args.pause_flag_file).strip()

57 else (output_dir / "PAUSE").resolve()

58 )

59 pause_ckpt_dir = (

60 Path(str(args.pause_ckpt_dir)).expanduser().resolve()

61 if str(args.pause_ckpt_dir).strip()

62 else (output_dir / "pause_ckpt").resolve()

63 )

64 if pause_ckpt_dir == output_dir:

65 raise ValueError("pause_ckpt_dir must not be the same as output_dir.")

66 return pause_flag_file, pause_ckpt_dir

67

68

69def is_trainer_checkpoint(path: Path) -> bool:

70 return (

71 path.exists()

72 and path.is_dir()

73 and (path / TRAINER_STATE_NAME).exists()

74 and (path / TRAINING_ARGS_NAME).exists()

75 )

76

77

78def checkpoint_step(path: Path) -> int:

79 match = re.fullmatch(r"checkpoint-(\d+)", path.name)

80 if not match:

81 return -1

82 return int(match.group(1))

83

84

85def find_latest_trainer_checkpoint(output_dir: Path) -> Path | None:

86 if not output_dir.exists() or not output_dir.is_dir():

87 return None

88 candidates = [

89 child

90 for child in output_dir.iterdir()

91 if child.is_dir()

92 and checkpoint_step(child) >= 0

93 and is_trainer_checkpoint(child)

94 ]

95 if not candidates:

96 return None

97 return max(candidates, key=checkpoint_step)

98

99

100def checkpoint_input_feature_bins(path: Path) -> int:

101 config_path = path / "config.json"

102 if not config_path.exists():

103 return MONO_FEATURE_BINS

104 try:

105 payload = json.loads(config_path.read_text(encoding="utf-8"))

106 except (OSError, json.JSONDecodeError):

107 return MONO_FEATURE_BINS

108 return int(payload.get("num_mel_bins") or MONO_FEATURE_BINS)

109

110

111def is_stereo_checkpoint(path: Path) -> bool:

112 return checkpoint_input_feature_bins(path) == STEREO_FEATURE_BINS

113

114

115def resolve_resume_checkpoint(

116 args: argparse.Namespace,

117 pause_ckpt_dir: Path,

118) -> Tuple[Path | None, str]:

119 if str(args.resume_from_checkpoint).strip():

120 resume_path = Path(str(args.resume_from_checkpoint)).expanduser().resolve()

121 if not resume_path.exists():

122 raise FileNotFoundError(f"resume_from_checkpoint not found: {resume_path}")

123 if not is_stereo_checkpoint(resume_path):

124 raise ValueError(

125 f"Explicit resume checkpoint is not stereo-compatible: {resume_path} "

126 f"(num_mel_bins={checkpoint_input_feature_bins(resume_path)}, "

127 f"expected={STEREO_FEATURE_BINS})"

128 )

129 return resume_path, "explicit"

130

131 if not bool(args.auto_resume):

132 return None, "disabled"

133

134 if is_trainer_checkpoint(pause_ckpt_dir) and is_stereo_checkpoint(pause_ckpt_dir):

135 return pause_ckpt_dir, "pause_ckpt"

136 if is_trainer_checkpoint(pause_ckpt_dir):

137 print(

138 f"[Resume] skipping mono/incompatible pause checkpoint: {pause_ckpt_dir} "

139 f"(num_mel_bins={checkpoint_input_feature_bins(pause_ckpt_dir)})"

140 )

141

142 latest_ckpt = find_latest_trainer_checkpoint(Path(str(args.output_dir)))

143 if latest_ckpt is not None and is_stereo_checkpoint(latest_ckpt):

144 return latest_ckpt, "latest_checkpoint"

145 if latest_ckpt is not None:

146 print(

147 f"[Resume] skipping mono/incompatible latest checkpoint: {latest_ckpt} "

148 f"(num_mel_bins={checkpoint_input_feature_bins(latest_ckpt)})"

149 )

150 return None, "none"

151

152

153def load_audio_stereo_16k(path: Path, target_sr: int = TARGET_SR) -> np.ndarray:

154 audio, sr = sf.read(str(path), always_2d=True)

155 audio = np.asarray(audio, dtype=np.float32)

156 if audio.size == 0:

157 raise ValueError(f"Empty audio file: {path}")

158

159 if audio.shape[1] == 1:

160 audio = np.repeat(audio, STEREO_CHANNELS, axis=1)

161 elif audio.shape[1] < STEREO_CHANNELS:

162 raise ValueError(f"Expected at least one audio channel, got shape {audio.shape}: {path}")

163 elif audio.shape[1] > STEREO_CHANNELS:

164 audio = audio[:, :STEREO_CHANNELS]

165

166 stereo = audio.T

167 if int(sr) != int(target_sr):

168 stereo = np.stack(

169 [

170 librosa.resample(channel, orig_sr=int(sr), target_sr=int(target_sr))

171 for channel in stereo

172 ],

173 axis=0,

174 )

175 return np.asarray(stereo, dtype=np.float32)

176

177

178def fix_audio_length(waveform: np.ndarray, max_len: int) -> np.ndarray:

179 if waveform.ndim != 2 or waveform.shape[0] != STEREO_CHANNELS:

180 raise ValueError(f"Expected stereo waveform with shape [2, T], got {waveform.shape}")

181 if waveform.shape[1] > max_len:

182 waveform = waveform[:, :max_len]

183 elif waveform.shape[1] < max_len:

184 waveform = np.pad(

185 waveform,

186 ((0, 0), (0, max_len - waveform.shape[1])),

187 mode="constant",

188 )

189 return waveform.astype(np.float32, copy=False)

190

191

192def stereo_to_features(

193 feature_extractor: AutoFeatureExtractor,

194 waveform: np.ndarray,

195) -> torch.Tensor:

196 if waveform.ndim != 2 or waveform.shape[0] != STEREO_CHANNELS:

197 raise ValueError(f"Expected stereo waveform with shape [2, T], got {waveform.shape}")

198

199 channel_features: list[torch.Tensor] = []

200 for channel in waveform:

201 features = feature_extractor(

202 channel,

203 sampling_rate=TARGET_SR,

204 return_tensors="pt",

205 ).input_features

206 channel_features.append(features[0].cpu().to(torch.float32))

207

208 if len(channel_features) != STEREO_CHANNELS:

209 raise ValueError(f"Expected {STEREO_CHANNELS} channel feature tensors")

210 if channel_features[0].shape != channel_features[1].shape:

211 raise ValueError(

212 "Stereo channel feature shapes differ: "

213 f"{channel_features[0].shape} vs {channel_features[1].shape}"

214 )

215 return torch.cat(channel_features, dim=0).to(torch.float32)

216

217

218def whisper_encoder_first_conv(model: torch.nn.Module) -> torch.nn.Conv1d:

219 candidates = [

220 getattr(model, "encoder", None),

221 getattr(getattr(model, "model", None), "encoder", None),

222 getattr(getattr(model, "whisper", None), "encoder", None),

223 ]

224 for encoder in candidates:

225 conv1 = getattr(encoder, "conv1", None) if encoder is not None else None

226 if isinstance(conv1, torch.nn.Conv1d):

227 return conv1

228 raise RuntimeError("Could not locate Whisper encoder conv1 for stereo adaptation")

229

230

231def adapt_whisper_model_for_stereo(model: torch.nn.Module) -> None:

232 conv1 = whisper_encoder_first_conv(model)

233 if conv1.in_channels == STEREO_FEATURE_BINS:

234 setattr(model.config, "num_mel_bins", STEREO_FEATURE_BINS)

235 return

236 if conv1.in_channels != MONO_FEATURE_BINS:

237 raise RuntimeError(

238 f"Unsupported Whisper conv1 input channels: {conv1.in_channels}; "

239 f"expected {MONO_FEATURE_BINS} or {STEREO_FEATURE_BINS}"

240 )

241

242 new_conv = torch.nn.Conv1d(

243 in_channels=STEREO_FEATURE_BINS,

244 out_channels=conv1.out_channels,

245 kernel_size=conv1.kernel_size,

246 stride=conv1.stride,

247 padding=conv1.padding,

248 dilation=conv1.dilation,

249 groups=conv1.groups,

250 bias=conv1.bias is not None,

251 padding_mode=conv1.padding_mode,

252 device=conv1.weight.device,

253 dtype=conv1.weight.dtype,

254 )

255 with torch.no_grad():

256 new_conv.weight[:, :MONO_FEATURE_BINS, :] = conv1.weight * 0.5

257 new_conv.weight[:, MONO_FEATURE_BINS:, :] = conv1.weight * 0.5

258 if conv1.bias is not None and new_conv.bias is not None:

259 new_conv.bias.copy_(conv1.bias)

260

261 candidates = [

262 getattr(model, "encoder", None),

263 getattr(getattr(model, "model", None), "encoder", None),

264 getattr(getattr(model, "whisper", None), "encoder", None),

265 ]

266 for encoder in candidates:

267 if encoder is not None and getattr(encoder, "conv1", None) is conv1:

268 encoder.conv1 = new_conv

269 break

270 setattr(model.config, "num_mel_bins", STEREO_FEATURE_BINS)

271

272

273def normalize_model_name_for_cache(model_name: str) -> str:

274 raw = str(model_name).strip()

275 if not raw:

276 return "unknown_model"

277 normalized = raw.replace("\\", "/")

278 hit = HF_CACHE_MODEL_RE.search(normalized)

279 if hit:

280 return f"{hit.group(1)}/{hit.group(2)}"

281 return raw

282

283

284def normalize_audio_relpath_for_cache(rel: Path) -> Path | None:

285 parts = [part for part in rel.parts if part not in {"", "."}]

286 if not parts:

287 return None

288 lower_parts = [part.lower() for part in parts]

289 if "renders" in lower_parts:

290 render_idx = lower_parts.index("renders")

291 if render_idx + 1 < len(parts):

292 parts = parts[render_idx + 1 :]

293 else:

294 return None

295 return Path(*parts)

296

297

298def feature_cache_base_dir(

299 feature_cache_root: Path,

300 model_name: str,

301 max_seconds: float,

302) -> Path:

303 model_id = normalize_model_name_for_cache(model_name)

304 model_tag = re.sub(r"[^A-Za-z0-9._-]+", "_", model_id).strip("_")

305 max_tag = f"{max_seconds:.3f}".replace(".", "p")

306 return feature_cache_root / model_tag / f"sr{TARGET_SR}_max{max_tag}_{AUDIO_CHANNEL_MODE}"

307

308

309def feature_cache_path_for_audio(

310 wav_path: Path,

311 audio_root: Path,

312 cache_base: Path,

313) -> Path:

314 wav_path = wav_path.resolve()

315 audio_root = audio_root.resolve()

316 try:

317 rel = wav_path.relative_to(audio_root)

318 normalized_rel = normalize_audio_relpath_for_cache(rel)

319 if normalized_rel is None:

320 digest = hashlib.sha1(str(wav_path).encode("utf-8")).hexdigest()

321 return cache_base / f"{digest}.pt"

322 return (cache_base / normalized_rel).with_suffix(".pt")

323 except ValueError:

324 digest = hashlib.sha1(str(wav_path).encode("utf-8")).hexdigest()

325 return cache_base / f"{digest}.pt"

326

327

328def atomic_torch_save(tensor: torch.Tensor, path: Path) -> None:

329 path.parent.mkdir(parents=True, exist_ok=True)

330 tmp = path.with_suffix(path.suffix + f".tmp_{os.getpid()}")

331 torch.save(tensor.cpu(), tmp)

332 os.replace(tmp, path)

333

334

335def compute_input_features(

336 wav_path: Path,

337 feature_extractor: AutoFeatureExtractor,

338 max_len: int,

339) -> torch.Tensor:

340 waveform = load_audio_stereo_16k(wav_path, TARGET_SR)

341 waveform = fix_audio_length(waveform, max_len)

342 return stereo_to_features(feature_extractor, waveform)

343

344

345def load_or_compute_input_features(

346 wav_path: Path,

347 feature_extractor: AutoFeatureExtractor,

348 max_len: int,

349 use_feature_cache: bool,

350 cache_path: Path,

351 save_on_miss: bool,

352) -> torch.Tensor:

353 if use_feature_cache and cache_path.exists():

354 tensor = torch.load(cache_path, map_location="cpu")

355 if not isinstance(tensor, torch.Tensor):

356 raise ValueError(f"Invalid cached tensor format: {cache_path}")

357 if tensor.ndim < 2 or int(tensor.shape[0]) != STEREO_FEATURE_BINS:

358 raise ValueError(

359 f"Cached features are not stereo {STEREO_FEATURE_BINS}-bin tensors: "

360 f"{cache_path} shape={tuple(tensor.shape)}"

361 )

362 return tensor.to(torch.float32)

363

364 tensor = compute_input_features(

365 wav_path=wav_path,

366 feature_extractor=feature_extractor,

367 max_len=max_len,

368 )

369 if use_feature_cache and save_on_miss:

370 atomic_torch_save(tensor, cache_path)

371 return tensor

372

373

374def preextract_feature_cache(

375 rows: Sequence[Dict[str, Any]],

376 audio_root: Path,

377 feature_extractor: AutoFeatureExtractor,

378 max_len: int,

379 cache_base: Path,

380 report_every: int,

381) -> None:

382 if not rows:

383 return

384 report_every = max(1, int(report_every))

385 started_at = time.time()

386 cache_hit = 0

387 cache_miss = 0

388 total = len(rows)

389 print(f"[FeatureCache] preextract started: total={total}, cache_base={cache_base}")

390

391 for index, row in enumerate(rows, start=1):

392 wav_path = Path(str(row["audio"])).expanduser().resolve()

393 cache_path = feature_cache_path_for_audio(

394 wav_path=wav_path,

395 audio_root=audio_root,

396 cache_base=cache_base,

397 )

398 if cache_path.exists():

399 cache_hit += 1

400 else:

401 tensor = compute_input_features(

402 wav_path=wav_path,

403 feature_extractor=feature_extractor,

404 max_len=max_len,

405 )

406 atomic_torch_save(tensor, cache_path)

407 cache_miss += 1

408

409 if index % report_every == 0 or index == total:

410 elapsed = time.time() - started_at

411 print(

412 f"[FeatureCache] {index}/{total} done, "

413 f"hit={cache_hit}, miss={cache_miss}, elapsed={elapsed:.1f}s"

414 )

415

416 total_elapsed = time.time() - started_at

417 print(

418 f"[FeatureCache] preextract completed: hit={cache_hit}, miss={cache_miss}, "

419 f"time={total_elapsed:.1f}s"

420 )

421

422

423class PauseSaveAndStopCallback(TrainerCallback):

424 def __init__(

425 self,

426 pause_flag_file: Path,

427 pause_ckpt_dir: Path,

428 remove_pause_flag: bool = True,

429 ) -> None:

430 self.pause_flag_file = pause_flag_file

431 self.pause_ckpt_dir = pause_ckpt_dir

432 self.remove_pause_flag = remove_pause_flag

433 self._trainer: Trainer | None = None

434 self._pause_handled = False

435

436 def set_trainer(self, trainer: Trainer) -> None:

437 self._trainer = trainer

438

439 def _save_pause_checkpoint(self) -> None:

440 if self._trainer is None:

441 raise RuntimeError("Pause callback has no trainer reference.")

442

443 trainer = self._trainer

444 if self.pause_ckpt_dir.exists():

445 if self.pause_ckpt_dir.is_dir():

446 shutil.rmtree(self.pause_ckpt_dir)

447 else:

448 self.pause_ckpt_dir.unlink()

449 self.pause_ckpt_dir.mkdir(parents=True, exist_ok=True)

450

451 trainer.save_model(str(self.pause_ckpt_dir))

452 torch.save(trainer.args, self.pause_ckpt_dir / TRAINING_ARGS_NAME)

453 if trainer.optimizer is not None:

454 torch.save(trainer.optimizer.state_dict(), self.pause_ckpt_dir / OPTIMIZER_STATE_NAME)

455 if trainer.lr_scheduler is not None:

456 torch.save(

457 trainer.lr_scheduler.state_dict(),

458 self.pause_ckpt_dir / SCHEDULER_STATE_NAME,

459 )

460

461 accelerator = getattr(trainer, "accelerator", None)

462 scaler = getattr(accelerator, "scaler", None)

463 if scaler is not None and hasattr(scaler, "state_dict"):

464 torch.save(scaler.state_dict(), self.pause_ckpt_dir / SCALER_STATE_NAME)

465

466 trainer.state.save_to_json(str(self.pause_ckpt_dir / TRAINER_STATE_NAME))

467 rng_state: Dict[str, Any] = {

468 "python": random.getstate(),

469 "numpy": np.random.get_state(),

470 "cpu": torch.random.get_rng_state(),

471 }

472 if torch.cuda.is_available():

473 rng_state["cuda"] = torch.cuda.random.get_rng_state_all()

474 torch.save(rng_state, self.pause_ckpt_dir / RNG_STATE_NAME)

475

476 def on_step_end(self, args, state, control, **kwargs):

477 if self._pause_handled or not self.pause_flag_file.exists():

478 return control

479

480 print("[PAUSE] detected -> saving checkpoint then stopping.")

481 try:

482 if self._trainer is None:

483 raise RuntimeError("Pause callback has no trainer reference.")

484 if self._trainer.args.should_save:

485 self._save_pause_checkpoint()

486 print(f"[PAUSE] checkpoint saved: {self.pause_ckpt_dir}")

487 if self.remove_pause_flag and self.pause_flag_file.exists():

488 self.pause_flag_file.unlink()

489 print(f"[PAUSE] flag removed: {self.pause_flag_file}")

490 else:

491 print("[PAUSE] non-saving process received pause signal, stopping.")

492 except Exception as exc:

493 print(f"[PAUSE] failed to save pause checkpoint: {exc}")

494

495 self._pause_handled = True

496 control.should_training_stop = True

497 return control