From 2c2e493df84f7640325a820502d995f1c733d08b Mon Sep 17 00:00:00 2001 From: Palash Debnath Date: Fri, 26 Jun 2026 14:29:38 +0530 Subject: [PATCH] fix(dub): stream segments to disk to stop long-video RAM spikes (#639) (#709) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Takes over and completes #639 (original work by @trungthanh1288). Dub generation held every segment's audio in RAM until final mix, so long/feature-length dubs and big batches could exhaust memory. Segments now stream to disk as rendered; the final track assembles from those files via a 30s-chunk memmap writer, so peak memory stays flat regardless of length. Completed on top of the original PR: - Watermarking: keep the project's 'every OmniVoice audio carries the signature' guarantee without double-marking. Since seg_.wav is BOTH the downloadable file AND the assembly input, mark each fresh segment once at synthesis and drop the per-chunk embed in the memmap writer (the final mix inherits the mark) — main's proven policy. Verified with real AudioSeal: 0.9999 detect confidence on the final track and on seg WAVs; cached/silence not re-marked. - Fix a crash regression: zero/negative-duration segments returned an in-memory zero-length entry instead of writing empty audio (which raised). Regression test added. - Perf: drop per-segment gc.collect(); throttle empty_cache() to every 16th call (the replaced code batched I/O to keep this off the hot path). - Clean up the mix_ temp WAVs after assembly. - Rewrite the watermark test for the multi-chunk (>30s) path; assert both the final track and the seg WAV are marked, with no double-mark. 212 passed / 1 skipped; route inventory clean. Co-authored-by: mergetest Co-authored-by: trungthanh1288 Co-authored-by: Claude Opus 4.8 (1M context) --- CHANGELOG.md | 8 + backend/api/routers/dub_generate.py | 575 +++++++++++++++++++--------- tests/test_smart_fit_generate.py | 98 +++++ 3 files changed, 503 insertions(+), 178 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 546c8963..e81d8c39 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -124,6 +124,14 @@ across dub, generate, and design (a corrupt-binary failure no longer poses as "get a free token" link. (#657, #669) ### Fixed +- **Long-video dubs no longer spike RAM during assembly.** Dub generation used + to hold every segment's audio in memory until the whole track was mixed, so a + 50-video batch or a single feature-length dub could exhaust RAM and crash. Each + segment now streams to disk as it's rendered and the final track is assembled + from those files via a 30s-chunk memmap writer, keeping memory flat regardless + of video length. Per-segment download WAVs and the final track stay correctly + watermarked (marked once at synthesis, no double-mark), and zero/negative-length + segments no longer crash the run. (#639) - **A corrupt or wrong-architecture native component no longer masquerades as "out of memory."** A synth failure caused by a bad `.dll`/`.pyd`/`.exe` on Windows (`[WinError 193] %1 is not a valid Win32 application` — e.g. torch, diff --git a/backend/api/routers/dub_generate.py b/backend/api/routers/dub_generate.py index 5c913558..bf9d936a 100644 --- a/backend/api/routers/dub_generate.py +++ b/backend/api/routers/dub_generate.py @@ -107,6 +107,123 @@ async def dub_generate(job_id: str, req: DubRequest): all_segment_wavs = [] sync_scores = [] + # Throttle the device cache flush. empty_cache() is a synchronous + # device stall, so calling it every segment (as the old code did) + # serialised the GPU loop; the batched-I/O design it replaced kept + # it off the hot path on purpose. Flush every ~16 releases instead — + # frequent enough to bound VRAM, rare enough to stay invisible. + _RELEASE_FLUSH_EVERY = 16 + _release_count = {"n": 0} + + def _release_audio_tensors(*objs) -> None: + """Best-effort VRAM cleanup after a segment is safely on disk. + + Tensors are freed by the callers' own ``del`` once they fall out + of scope; this only throttles the device cache flush. ``*objs`` is + kept for call-site compatibility but intentionally unused — a local + ``del`` here would only unbind the parameter, never the caller's + reference. + """ + _release_count["n"] += 1 + if _release_count["n"] % _RELEASE_FLUSH_EVERY != 0: + return + try: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + torch.mps.empty_cache() + except Exception: + pass + + # mix_ scratch WAVs written for silence/cached-fail/error slots are + # pure assembly inputs (no preview/regen contract), so they're deleted + # once the final track is written. + _mix_temp_paths: list[str] = [] + + def _store_mix_wav(start: float, end: float, wav: torch.Tensor, sr: int, seg_key: str): + """Write one segment to disk and keep only its path in the mix manifest. + + A zero/negative-length buffer is never written (``atomic_save_wav`` + raises on empty audio); instead a harmless zero-length in-memory + entry is returned, which the assembly tolerates via its ``e > s`` + guard. + """ + if wav.shape[-1] <= 0: + return (start, end, torch.zeros(1, 0), sr) + path = dub_seg_path(job_id, seg_key) + os.makedirs(os.path.dirname(path), exist_ok=True) + atomic_save_wav(path, wav.detach().cpu(), sr) + if seg_key.startswith("mix_"): + _mix_temp_paths.append(path) + _release_audio_tensors(wav) + return (start, end, path, sr) + + def _entry_num_samples(entry) -> int: + # Zero/negative-duration slots are kept as in-memory tensors (never + # written to disk); report their length directly. + if isinstance(entry[2], torch.Tensor): + return int(entry[2].shape[-1]) + try: + info = torchaudio.info(entry[2]) + return int(info.num_frames) + except Exception: + wav, _sr = torchaudio.load(entry[2]) + n = int(wav.shape[-1]) + _release_audio_tensors(wav) + return n + + def _load_entry_wav(entry, target_sr: int) -> torch.Tensor: + if isinstance(entry[2], torch.Tensor): + return entry[2] + wav, loaded_sr = torchaudio.load(entry[2]) + if loaded_sr != target_sr: + import torchaudio.functional as AF + wav = AF.resample(wav, loaded_sr, target_sr) + return wav + + def _write_memmap_wav_atomic(target_path: str, samples, sample_rate: int) -> None: + """Write a mono float32 memmap to int16 WAV without loading it all. + + Intentionally does NOT watermark: the final track is assembled from + per-segment WAVs that were already watermarked once at synthesis + time (see the seg-write path below), exactly as ``main`` does. + Re-marking here would double-mark every segment in the final mix. + """ + import tempfile + import wave + import numpy as np + + target_dir = os.path.dirname(target_path) or "." + target_base = os.path.basename(target_path) + fd, tmp_path = tempfile.mkstemp( + prefix=f".{target_base}.", + suffix=".wav", + dir=target_dir, + ) + os.close(fd) + chunk_samples = max(sample_rate * 30, 1) + try: + with wave.open(tmp_path, "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(sample_rate) + total_len = int(samples.shape[0]) + for off in range(0, total_len, chunk_samples): + chunk = np.array(samples[off: off + chunk_samples], dtype=np.float32, copy=True) + if chunk.size == 0: + continue + np.nan_to_num(chunk, copy=False, nan=0.0, posinf=1.0, neginf=-1.0) + chunk = np.clip(chunk, -1.0, 1.0) + pcm = (chunk * 32767.0).astype(" target_samples: cached_wav = cached_wav[..., :target_samples] - all_segment_wavs.append((seg.start, seg.end, cached_wav, _model.sampling_rate)) + all_segment_wavs.append(_store_mix_wav(seg.start, seg.end, cached_wav, _model.sampling_rate, f"mix_{seg_id}")) + try: + del cached_wav + except Exception: + pass + _release_audio_tensors() sync_scores.append(getattr(seg, 'sync_ratio', None) or 1.0) _t_cache += time.perf_counter() - _t_cache_0 continue @@ -191,8 +321,13 @@ async def dub_generate(job_id: str, req: DubRequest): # is broken — cleaner than aborting the whole mix. yield f"data: {json.dumps({'type': 'warning', 'segment': i, 'message': f'cached seg lost, padding silence: {str(e)[:120]}'})}\n\n" sr = _model.sampling_rate - silence = torch.zeros(1, int(seg_duration * sr)) - all_segment_wavs.append((seg.start, seg.end, silence, sr)) + silence = torch.zeros(1, max(0, int(seg_duration * sr))) + all_segment_wavs.append(_store_mix_wav(seg.start, seg.end, silence, sr, f"mix_{seg_id}")) + try: + del silence + except Exception: + pass + _release_audio_tensors() sync_scores.append(1.0) continue @@ -457,7 +592,7 @@ async def dub_generate(job_id: str, req: DubRequest): except Exception as e: logger.debug("seg fingerprint skipped for %s: %s", seg_id, e) - _pending_seg_writes.append((i, audio_tensor, _model.sampling_rate, seg_id, _seg_fp, _num_step)) + _pending_seg_writes.append((i, _model.sampling_rate, seg_id, _seg_fp, _num_step)) # RVC needs the WAV on disk, so write it immediately only # when RVC is active (uncommon path). @@ -479,32 +614,54 @@ async def dub_generate(job_id: str, req: DubRequest): except Exception as e: yield f"data: {json.dumps({'type': 'warning', 'segment': i, 'message': f'RVC skipped: {str(e)[:120]}'})}\n\n" - all_segment_wavs.append((seg.start, seg.end, audio_tensor, _model.sampling_rate)) + # Watermark this FRESH TTS output exactly once, right before it + # is persisted. The same seg_.wav is BOTH the downloadable + # per-segment file AND the assembly input for the final track, + # so marking it here (and nowhere else) gives the downloadable + # WAV its mark back and the final mix inherits it — no double- + # mark. Cached-reuse audio is already marked; silence/zero slots + # carry no speech to mark, so neither is re-watermarked. + audio_tensor = embed_watermark(audio_tensor, _model.sampling_rate) + + seg_wav_path = dub_seg_path(job_id, seg_id) + try: + # Keep the existing per-segment WAV contract for previews + # and partial regeneration, but do not keep the tensor in RAM. + atomic_save_wav(seg_wav_path, audio_tensor, _model.sampling_rate) + except Exception as e: + logger.warning("seg write failed for %s: %s", seg_id, e) + # If the durable segment write fails, still preserve a mix + # copy so this generation can finish. + all_segment_wavs.append(_store_mix_wav(seg.start, seg.end, audio_tensor, _model.sampling_rate, f"mix_{seg_id}")) + try: + del audio_tensor + except Exception: + pass + _release_audio_tensors() + else: + all_segment_wavs.append((seg.start, seg.end, seg_wav_path, _model.sampling_rate)) + try: + del audio_tensor + except Exception: + pass + _release_audio_tensors() except Exception as e: yield f"data: {json.dumps({'type': 'error', 'segment': i, 'error': str(e)})}\n\n" sr = _model.sampling_rate - all_segment_wavs.append((seg.start, seg.end, torch.zeros(1, int(seg_duration * sr)), sr)) + all_segment_wavs.append(_store_mix_wav(seg.start, seg.end, torch.zeros(1, max(0, int(seg_duration * sr))), sr, f"mix_{seg_id}")) sync_scores.append(1.0) _t_loop_end = time.perf_counter() yield f"data: {json.dumps({'type': 'assembling'})}\n\n" - # ── Batch disk-write phase ──────────────────────────────────── - # Flush all per-segment WAVs and fingerprints in one burst now - # that the GPU-hot loop is done. This keeps I/O off the critical - # path and cuts ~200ms × N_segments of latency. + # ── Batch metadata phase ────────────────────────────────────── + # Per-segment WAVs were written during the loop to keep RAM bounded. + # Flush only lightweight fingerprints/quality metadata here. _t_diskw_0 = time.perf_counter() hashes = job.setdefault("seg_hashes", {}) quality_map = job.setdefault("seg_num_step", {}) - for (_si, _wav, _sr, _sid, _fp, _nstep) in _pending_seg_writes: - seg_wav_path = dub_seg_path(job_id, _sid) - try: - # Apply invisible watermark before writing to disk - _wav = embed_watermark(_wav, _sr) - atomic_save_wav(seg_wav_path, _wav, _sr) - except Exception as e: - logger.warning("deferred seg write failed for %s: %s", _sid, e) + for (_si, _sr, _sid, _fp, _nstep) in _pending_seg_writes: if _fp is not None: hashes[_sid] = _fp quality_map[_sid] = _nstep @@ -532,8 +689,8 @@ async def dub_generate(job_id: str, req: DubRequest): if strategy == "stretch_video": cursor = 0.0 - for i, (orig_start, orig_end, wav, _) in enumerate(all_segment_wavs): - wl_i = wav.shape[-1] + for i, (orig_start, orig_end, wav_path, _) in enumerate(all_segment_wavs): + wl_i = _entry_num_samples((orig_start, orig_end, wav_path, sr)) natural_dur = (wl_i / sr) if wl_i > 0 else max(0.0, orig_end - orig_start) if i == 0: # Preserve the pre-roll (silence before the first seg). @@ -583,9 +740,9 @@ async def dub_generate(job_id: str, req: DubRequest): "start": s, "end": e, } - for i, (s, e, _w, _) in enumerate(all_segment_wavs) + for i, (s, e, _path, _) in enumerate(all_segment_wavs) ], - [w.shape[-1] / sr for (_s, _e, w, _) in all_segment_wavs], + [_entry_num_samples(entry) / sr for entry in all_segment_wavs], orig_total_dur, fit_params, ) @@ -598,183 +755,245 @@ async def dub_generate(job_id: str, req: DubRequest): # not from the plan — so subtitles land exactly on the audio. fitted_cues: list[dict] = [] - full_audio = torch.zeros(1, total_samples) + lang_code = req.language_code or "und" + track_path = os.path.join(DUB_DIR, job_id, f"dubbed_{lang_code}.wav") + os.makedirs(os.path.dirname(track_path), exist_ok=True) - for i, (start, end, wav, _) in enumerate(all_segment_wavs): - seg_ref = req.segments[i] if i < len(req.segments) else None - seg_gain = getattr(seg_ref, "gain", None) if seg_ref is not None else None - seg_gain = seg_gain if seg_gain is not None else 1.0 - seg_gain = max(0.0, min(2.0, seg_gain)) - adjusted = wav * seg_gain - wl = adjusted.shape[-1] - natural_dur = wl / sr if wl > 0 else 0.0 - orig_dur = max(0.0, end - start) + import gc + import tempfile + import numpy as np - if strategy == "stretch_video": - # Mode B: audio at natural rate, placed on the stretched - # timeline. No trim, no atempo. dub_export handles the video. - new_start, _new_end = new_layout[i] - place_at = new_start - fit_status.append({ - "status": "video_stretched", - "stretch_ratio": round(natural_dur / max(orig_dur, 1e-3), 3), - }) + mix_samples = max(total_samples, 1) + fd, mix_path = tempfile.mkstemp( + prefix=f".{os.path.basename(track_path)}.mix.", + suffix=".f32", + dir=os.path.dirname(track_path), + ) + os.close(fd) + try: + with open(mix_path, "r+b") as mix_file: + mix_file.truncate(mix_samples * 4) + mix_audio = np.memmap(mix_path, dtype=np.float32, mode="r+", shape=(mix_samples,)) - elif strategy == "smart_fit": - # Smart Fit: apply the planner's audio_rate via the same - # pitch-preserving atempo pipe strict_slot uses, place the - # result at the planned new_start, and hard-trim whatever - # the caps couldn't absorb. The video side (video_ratio per - # chunk) is persisted below for the export pipeline. - sf = fit_plan.segments[i] - place_at = sf.new_start - if sf.audio_rate > 1.0 + 1e-6 and wl > 0: - target = max(1, int(round(wl / sf.audio_rate))) - try: - adjusted = await _pitch_preserving_stretch( - adjusted, target, sr, - ) - except Exception as e: - logger.warning( - "atempo stretch failed for seg %d (%.2f×), " - "falling back to linear interp: %s", - i, sf.audio_rate, e, - ) - adjusted = torch.nn.functional.interpolate( - adjusted.unsqueeze(0), - size=target, - mode='linear', - align_corners=False, - ).squeeze(0) - wl = adjusted.shape[-1] - # Residual overflow → hard-trim to the segment's new video - # slot (fade below keeps the cut pop-free). - new_slot_samples = int(max(0.0, sf.new_end - sf.new_start) * sr) - if new_slot_samples > 0 and wl > new_slot_samples: - adjusted = adjusted[..., :new_slot_samples] - wl = adjusted.shape[-1] - # Truthful per-segment verdict for the UI badge. - entry = {"status": sf.status} - if sf.audio_rate > 1.0 + 1e-6: - entry["audio_rate"] = round(sf.audio_rate, 3) - if sf.video_ratio > 1.0 + 1e-6: - entry["video_ratio"] = round(sf.video_ratio, 3) - if sf.overflow_s > 0: - entry["overflow_s"] = round(sf.overflow_s, 3) - fit_status.append(entry) - # Cue times from the ACTUAL stretched sample positions. - fitted_cues.append({ - "id": sf.seg_id, - "start": round(place_at, 4), - "end": round(place_at + wl / sr, 4), - }) + for i, (start, end, wav_path, _) in enumerate(all_segment_wavs): + seg_ref = req.segments[i] if i < len(req.segments) else None + seg_gain = getattr(seg_ref, "gain", None) if seg_ref is not None else None + seg_gain = seg_gain if seg_gain is not None else 1.0 + seg_gain = max(0.0, min(2.0, seg_gain)) + wav = _load_entry_wav((start, end, wav_path, sr), sr) + adjusted = wav * seg_gain + if adjusted.ndim == 2 and adjusted.shape[0] > 1: + adjusted = adjusted.mean(dim=0, keepdim=True) + wl = adjusted.shape[-1] + natural_dur = wl / sr if wl > 0 else 0.0 + orig_dur = max(0.0, end - start) - elif strategy == "concise": - # Mode A: never compress. Allow the audio to extend into the - # silent gap before the next seg (existing heuristic) plus - # any extra `overflow_budget_s`. Beyond that, hard-trim with - # a short fade so we never overlap the next speaker. - place_at = start - effective_end = end - if i + 1 < len(all_segment_wavs): - next_start = all_segment_wavs[i + 1][0] - gap = next_start - end - if gap > GAP_OVERFLOW_BUFFER_S: - effective_end = end + min( - gap - GAP_OVERFLOW_BUFFER_S, GAP_OVERFLOW_MAX_S, - ) - effective_end += overflow_budget_s - slot_samples_eff = int(max(0.0, (effective_end - start)) * sr) - if slot_samples_eff > 0 and wl > slot_samples_eff: - overflow_s = (wl - slot_samples_eff) / sr - adjusted = adjusted[..., :slot_samples_eff] - wl = adjusted.shape[-1] + if strategy == "stretch_video": + # Mode B: audio at natural rate, placed on the stretched + # timeline. No trim, no atempo. dub_export handles the video. + new_start, _new_end = new_layout[i] + place_at = new_start fit_status.append({ - "status": "overflows", - "overflow_s": round(overflow_s, 3), + "status": "video_stretched", + "stretch_ratio": round(natural_dur / max(orig_dur, 1e-3), 3), }) - else: - fit_status.append({"status": "fits"}) - else: - # strict_slot (legacy): preserve the previous atempo / trim / - # off semantics so existing callers and back-compat tests - # keep passing. - place_at = start - effective_end = end - if i + 1 < len(all_segment_wavs): - next_start = all_segment_wavs[i + 1][0] - gap = next_start - end - if gap > GAP_OVERFLOW_BUFFER_S: - effective_end = end + min( - gap - GAP_OVERFLOW_BUFFER_S, GAP_OVERFLOW_MAX_S, - ) - slot_samples = int(max(0.0, (effective_end - start)) * sr) - if slot_fit != "off" and slot_samples > 0 and wl > slot_samples: - if slot_fit == "time_stretch": - ratio = wl / slot_samples - capped_ratio = min(ratio, MAX_STRETCH_RATIO) - capped_target = int(wl / capped_ratio) + elif strategy == "smart_fit": + # Smart Fit: apply the planner's audio_rate via the same + # pitch-preserving atempo pipe strict_slot uses, place the + # result at the planned new_start, and hard-trim whatever + # the caps couldn't absorb. The video side (video_ratio per + # chunk) is persisted below for the export pipeline. + sf = fit_plan.segments[i] + place_at = sf.new_start + if sf.audio_rate > 1.0 + 1e-6 and wl > 0: + target = max(1, int(round(wl / sf.audio_rate))) try: adjusted = await _pitch_preserving_stretch( - adjusted, capped_target, sr, + adjusted, target, sr, ) - if adjusted.shape[-1] > slot_samples: - adjusted = adjusted[..., :slot_samples] - if ratio > MAX_STRETCH_RATIO: - logger.info( - "seg %d compression %.2f× exceeded cap; " - "stretched to %.2f×, tail trimmed", - i, ratio, capped_ratio, - ) except Exception as e: logger.warning( "atempo stretch failed for seg %d (%.2f×), " "falling back to linear interp: %s", - i, ratio, e, + i, sf.audio_rate, e, ) adjusted = torch.nn.functional.interpolate( adjusted.unsqueeze(0), - size=slot_samples, + size=target, mode='linear', align_corners=False, ).squeeze(0) - else: # "trim" - adjusted = adjusted[..., :slot_samples] + wl = adjusted.shape[-1] + # Residual overflow → hard-trim to the segment's new video + # slot (fade below keeps the cut pop-free). + new_slot_samples = int(max(0.0, sf.new_end - sf.new_start) * sr) + if new_slot_samples > 0 and wl > new_slot_samples: + adjusted = adjusted[..., :new_slot_samples] + wl = adjusted.shape[-1] + # Truthful per-segment verdict for the UI badge. + entry = {"status": sf.status} + if sf.audio_rate > 1.0 + 1e-6: + entry["audio_rate"] = round(sf.audio_rate, 3) + if sf.video_ratio > 1.0 + 1e-6: + entry["video_ratio"] = round(sf.video_ratio, 3) + if sf.overflow_s > 0: + entry["overflow_s"] = round(sf.overflow_s, 3) + fit_status.append(entry) + # Cue times from the ACTUAL stretched sample positions. + fitted_cues.append({ + "id": sf.seg_id, + "start": round(place_at, 4), + "end": round(place_at + wl / sr, 4), + }) + + elif strategy == "concise": + # Mode A: never compress. Allow the audio to extend into the + # silent gap before the next seg (existing heuristic) plus + # any extra `overflow_budget_s`. Beyond that, hard-trim with + # a short fade so we never overlap the next speaker. + place_at = start + effective_end = end + if i + 1 < len(all_segment_wavs): + next_start = all_segment_wavs[i + 1][0] + gap = next_start - end + if gap > GAP_OVERFLOW_BUFFER_S: + effective_end = end + min( + gap - GAP_OVERFLOW_BUFFER_S, GAP_OVERFLOW_MAX_S, + ) + effective_end += overflow_budget_s + slot_samples_eff = int(max(0.0, (effective_end - start)) * sr) + if slot_samples_eff > 0 and wl > slot_samples_eff: + overflow_s = (wl - slot_samples_eff) / sr + adjusted = adjusted[..., :slot_samples_eff] + wl = adjusted.shape[-1] + fit_status.append({ + "status": "overflows", + "overflow_s": round(overflow_s, 3), + }) + else: + fit_status.append({"status": "fits"}) + + else: + # strict_slot (legacy): preserve the previous atempo / trim / + # off semantics so existing callers and back-compat tests + # keep passing. + place_at = start + effective_end = end + if i + 1 < len(all_segment_wavs): + next_start = all_segment_wavs[i + 1][0] + gap = next_start - end + if gap > GAP_OVERFLOW_BUFFER_S: + effective_end = end + min( + gap - GAP_OVERFLOW_BUFFER_S, GAP_OVERFLOW_MAX_S, + ) + slot_samples = int(max(0.0, (effective_end - start)) * sr) + if slot_fit != "off" and slot_samples > 0 and wl > slot_samples: + if slot_fit == "time_stretch": + ratio = wl / slot_samples + capped_ratio = min(ratio, MAX_STRETCH_RATIO) + capped_target = int(wl / capped_ratio) + try: + adjusted = await _pitch_preserving_stretch( + adjusted, capped_target, sr, + ) + if adjusted.shape[-1] > slot_samples: + adjusted = adjusted[..., :slot_samples] + if ratio > MAX_STRETCH_RATIO: + logger.info( + "seg %d compression %.2f× exceeded cap; " + "stretched to %.2f×, tail trimmed", + i, ratio, capped_ratio, + ) + except Exception as e: + logger.warning( + "atempo stretch failed for seg %d (%.2f×), " + "falling back to linear interp: %s", + i, ratio, e, + ) + adjusted = torch.nn.functional.interpolate( + adjusted.unsqueeze(0), + size=slot_samples, + mode='linear', + align_corners=False, + ).squeeze(0) + else: # "trim" + adjusted = adjusted[..., :slot_samples] + wl = adjusted.shape[-1] + fit_status.append({ + "status": "fits", + "compression_applied": (slot_fit == "time_stretch" + and wl != int(natural_dur * sr)), + }) + + # Common: short fades to avoid pops, then mix into disk-backed audio. + fade_ms = 15 + fade_samples = int((fade_ms / 1000.0) * sr) + if wl > fade_samples * 2: + ramp_up = torch.linspace(0, 1, fade_samples, device=adjusted.device) + ramp_down = torch.linspace(1, 0, fade_samples, device=adjusted.device) + adjusted[0, :fade_samples] *= ramp_up + adjusted[0, -fade_samples:] *= ramp_down + + s = int(place_at * sr) + if s < 0: + adjusted = adjusted[..., -s:] wl = adjusted.shape[-1] - fit_status.append({ - "status": "fits", - "compression_applied": (slot_fit == "time_stretch" - and wl != int(natural_dur * sr)), - }) + s = 0 + e = min(s + wl, total_samples) + if s < total_samples and e > s: + mix_len = e - s + seg_np = ( + adjusted[:, :mix_len] + .detach() + .cpu() + .to(torch.float32) + .clamp(-1.0, 1.0) + .squeeze(0) + .numpy() + ) + mix_audio[s:e] += seg_np + try: + del wav, adjusted + except Exception: + pass + _release_audio_tensors() - # Common: short fades to avoid pops, then mix into full_audio. - fade_ms = 15 - fade_samples = int((fade_ms / 1000.0) * sr) - if wl > fade_samples * 2: - ramp_up = torch.linspace(0, 1, fade_samples, device=adjusted.device) - ramp_down = torch.linspace(1, 0, fade_samples, device=adjusted.device) - adjusted[0, :fade_samples] *= ramp_up - adjusted[0, -fade_samples:] *= ramp_down - - s = int(place_at * sr) - e = min(s + wl, total_samples) - if s < total_samples: - full_audio[:, s:e] += adjusted[:, :e - s] - - lang_code = req.language_code or "und" - track_path = os.path.join(DUB_DIR, job_id, f"dubbed_{lang_code}.wav") - _t_save_0 = time.perf_counter() - # Apply invisible watermark to the final assembled track - full_audio = embed_watermark(full_audio, sr) - atomic_save_wav(track_path, full_audio, sr) - _t_save = time.perf_counter() - _t_save_0 - _t_mix = _t_save_0 - _t_loop_end + _t_save_0 = time.perf_counter() + mix_audio.flush() + _write_memmap_wav_atomic(track_path, mix_audio[:mix_samples], sr) + _t_save = time.perf_counter() - _t_save_0 + _t_mix = _t_save_0 - _t_loop_end + finally: + try: + mix_audio.flush() + mix_mmap = getattr(mix_audio, "_mmap", None) + if mix_mmap is not None: + mix_mmap.close() + except Exception: + pass + try: + del mix_audio + except Exception: + pass + gc.collect() + try: + os.unlink(mix_path) + except OSError: + pass + # The final track is written; the mix_ scratch WAVs (silence / + # cached-fail / error slots) have served their only purpose as + # assembly inputs and would otherwise leak into the job dir. + for _mp in _mix_temp_paths: + try: + os.unlink(_mp) + except OSError: + pass # Per-track metadata. For stretch_video, the dub wav is at the new # (longer) timeline, so we record its actual duration here too — the # mux step needs this to know whether to use the original video as-is # or stretch it per the plan. - track_dur = full_audio.shape[-1] / sr if full_audio.shape[-1] > 0 else 0.0 + track_dur = total_samples / sr if total_samples > 0 else 0.0 job["dubbed_tracks"][lang_code] = { "path": track_path, "language": req.language, diff --git a/tests/test_smart_fit_generate.py b/tests/test_smart_fit_generate.py index 7d9ecaf6..0974a1a6 100644 --- a/tests/test_smart_fit_generate.py +++ b/tests/test_smart_fit_generate.py @@ -146,6 +146,104 @@ def _track_samples(job_dir): # ── Audio-only stretch ───────────────────────────────────────────────── +def test_final_dub_track_and_seg_wav_are_watermarked(patched_generate, monkeypatch): + """Watermarking policy after the streaming-to-disk rewrite (#639). + + Fresh TTS output is watermarked exactly ONCE, right before its + per-segment WAV is written. That seg WAV is BOTH the downloadable file + and the assembly input, so: + - the downloadable seg_.wav carries the mark, and + - the final assembled track inherits it (the streaming memmap writer + does NOT re-watermark, so there's no double-mark). + + The fixture uses a >30s track so the multi-chunk memmap write path runs, + and a marker planted deep in the second 30s chunk so the test proves + that path preserves the watermark across the chunk boundary. + """ + run, model, job, job_dir = patched_generate + import api.routers.dub_generate as dg + from services import watermark + import torchaudio + + # The fake TTS emits a constant 0.25; the mix fades ramp that through + # [0, 0.25], so a positive marker could be forged by the fade. Use a + # negative marker the fades can never produce ⇒ only the planted window + # ever matches. -0.5 survives the int16 PCM round-trip. + marker = -0.5 + watermark_calls: list[int] = [] + # Plant the mark ~80k samples before the buffer end: clear of the 15ms + # mix fades AND (once the seg is placed at start=1.0s) inside the second + # 30s chunk of the final memmap write. + mark_back_off = 80_000 + mark_len = 256 + + def fake_embed(wav, sr): + watermark_calls.append(int(wav.shape[-1])) + out = wav.clone() + n = out.shape[-1] + off = max(0, n - mark_back_off) + out[..., off: off + mark_len] = marker + return out + + def fake_detect(wav, sr): + mark = torch.full_like(wav, marker) + hit = bool(torch.any(torch.isclose(wav, mark, atol=2e-3))) + return {"is_watermarked": hit, "confidence": 1.0 if hit else 0.0} + + monkeypatch.setattr(dg, "embed_watermark", fake_embed) + monkeypatch.setattr(watermark, "detect_watermark", fake_detect) + + job["duration"] = 35.0 + # 33s of natural speech placed at 1.0s → track is 35s (>30s ⇒ 2 chunks), + # the seg ends at 34s so its tail (and the planted mark) lives in chunk 2. + segs = [{"start": 1.0, "end": 34.0, "text": "33:hola"}] + _done(run(_body(segs, timing_strategy="concise"))) + + # Final assembled track is watermarked (mark survived the multi-chunk + # int16 memmap write). + final_wav, sr = torchaudio.load(str(job_dir / "dubbed_es.wav")) + assert final_wav.shape[-1] == int(35.0 * SR) + assert watermark.detect_watermark(final_wav, sr)["is_watermarked"] is True + + # Downloadable per-segment WAV is watermarked too. + seg_wav, seg_sr = torchaudio.load(str(job_dir / "seg_0.wav")) + assert watermark.detect_watermark(seg_wav, seg_sr)["is_watermarked"] is True + + # Marked exactly once, on the FRESH seg (33s natural length) — NOT on the + # 35s assembled track. One call ⇒ no double-mark. + assert watermark_calls == [int(33.0 * SR)] + assert int(33.0 * SR) != int(35.0 * SR) + + +def test_zero_and_negative_duration_segments_dont_crash(patched_generate): + """Zero/negative-duration slots must not feed a negative length to + torch.zeros (raises) nor write an empty WAV (atomic_save_wav raises). + They become harmless in-memory entries the assembly tolerates, and the + positive-duration silence's mix_ scratch WAV is cleaned up (#639).""" + run, model, job, job_dir = patched_generate + import torchaudio + + job["duration"] = 5.0 + segs = [ + {"start": 0.0, "end": 1.0, "text": "0.5:hola"}, # normal → seg_0.wav + {"start": 1.0, "end": 1.0, "text": ""}, # zero duration + {"start": 2.0, "end": 2.8, "text": " "}, # positive silence → mix temp + {"start": 4.0, "end": 3.5, "text": "boom"}, # negative duration + ] + parsed = run(_body(segs, timing_strategy="concise")) + + # Completed without raising and produced a track. + _done(parsed) + track = job_dir / "dubbed_es.wav" + assert track.exists() + n, sr = torchaudio.load(str(track))[0].shape[-1], SR + assert n == int(5.0 * SR) + + # The mix_ scratch WAV for the positive-duration silence is gone. + leftovers = [p.name for p in job_dir.glob("seg_mix_*.wav")] + assert leftovers == [], f"mix scratch WAVs leaked: {leftovers}" + + def test_smart_fit_audio_only_stretch_keeps_original_duration(patched_generate): run, model, job, job_dir = patched_generate # seg0 [0,1] natural 0.5s → fits. seg1 [2,3] is last → slot extends to