time stamp sync

app.py

@@ -16,14 +16,21 @@ if hasattr(torch, "load"):
 import tempfile
 from pathlib import Path
 import numpy as np
+import soundfile as sf
 import gradio as gr
-import spaces
 import whisper
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from TTS.api import TTS
 import yt_dlp
 from moviepy import VideoFileClip, AudioFileClip, AudioClip, concatenate_audioclips
 
+
+def gpu_decorator(fn):
+    if os.environ.get("SPACE_ID"):
+        import spaces
+        return spaces.GPU(duration=120)(fn)
+    return fn
+
 # ---------------------------------------------------------------------------
 # B. Global model loading on CPU (ZeroGPU has no CUDA at import time)
 # ---------------------------------------------------------------------------
@@ -36,6 +43,14 @@ trans_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled
 tts_model = TTS("tts_models/multilingual/multi-dataset/xtts_v2")
 print("All models loaded on CPU.")
 
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+XTTS_SAMPLE_RATE = 24000
+MIN_SEGMENT_DURATION = 1.5  # seconds
+MAX_STRETCH_RATE = 2.0
+MIN_STRETCH_RATE = 0.5
+
 # ---------------------------------------------------------------------------
 # C. Helper functions
 # ---------------------------------------------------------------------------
@@ -54,18 +69,97 @@ def download_youtube_video(url: str, output_dir: str) -> str:
         return ydl.prepare_filename(info)
 
 
-def translate_uk_to_en(text: str, device: str) -> str:
-    """Translate Ukrainian text to English using NLLB-200."""
+def merge_short_segments(segments):
+    """Merge consecutive short segments to avoid garbage TTS output."""
+    if not segments:
+        return []
+
+    merged = []
+    current = {
+        "start": segments[0]["start"],
+        "end": segments[0]["end"],
+        "text": segments[0]["text"].strip(),
+    }
+
+    for seg in segments[1:]:
+        text = seg["text"].strip()
+        if not text:
+            continue
+        duration = current["end"] - current["start"]
+        if duration < MIN_SEGMENT_DURATION:
+            current["end"] = seg["end"]
+            current["text"] += " " + text
+        else:
+            if current["text"]:
+                merged.append(current)
+            current = {"start": seg["start"], "end": seg["end"], "text": text}
+
+    if current["text"]:
+        merged.append(current)
+
+    return merged
+
+
+def translate_segments_uk_to_en(segments, device):
+    """Batch translation of segments using NLLB-200 with proper tokenizer batching."""
+    texts = [seg["text"] for seg in segments]
+
     trans_tokenizer.src_lang = "ukr_Cyrl"
-    inputs = trans_tokenizer(text, return_tensors="pt").to(device)
+    inputs = trans_tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
     translated_tokens = trans_model.generate(
         **inputs,
         forced_bos_token_id=trans_tokenizer.convert_tokens_to_ids("eng_Latn"),
-        max_length=256,
+        max_length=512,
         num_beams=5,
         repetition_penalty=1.5,
     )
-    return trans_tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
+    return [t.strip() for t in trans_tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)]
+
+
+def build_audio_canvas(segments, translated_texts, ref_audio_path, video_duration, tmp_dir, progress):
+    """Generate per-segment TTS, time-stretch to fit, and assemble onto a silent canvas."""
+    canvas = np.zeros(int(video_duration * XTTS_SAMPLE_RATE))
+    total = len(segments)
+
+    for i, (seg, text) in enumerate(zip(segments, translated_texts)):
+        progress(0.55 + 0.30 * (i / total), desc=f"Synthesizing segment {i+1}/{total}...")
+
+        if len(text) < 5:
+            continue
+
+        seg_start = seg["start"]
+        seg_end = seg["end"]
+        target_duration = seg_end - seg_start
+
+        # First pass: generate at natural speed to measure duration
+        tts_audio = tts_model.tts(text=text, speaker_wav=ref_audio_path, language="en")
+        tts_audio = np.array(tts_audio, dtype=np.float32)
+        tts_duration = len(tts_audio) / XTTS_SAMPLE_RATE
+
+        # If duration is off, regenerate with speed parameter
+        speed = tts_duration / target_duration
+        speed = max(MIN_STRETCH_RATE, min(MAX_STRETCH_RATE, speed))
+        if abs(speed - 1.0) >= 0.05:
+            tts_audio = tts_model.tts(text=text, speaker_wav=ref_audio_path, language="en", speed=speed)
+            tts_audio = np.array(tts_audio, dtype=np.float32)
+
+        # Truncate if it would overlap the next segment
+        if i + 1 < total:
+            max_samples = int((segments[i + 1]["start"] - seg_start) * XTTS_SAMPLE_RATE)
+            if len(tts_audio) > max_samples:
+                tts_audio = tts_audio[:max_samples]
+
+        # Place on canvas
+        start_sample = int(seg_start * XTTS_SAMPLE_RATE)
+        end_sample = start_sample + len(tts_audio)
+        if end_sample > len(canvas):
+            tts_audio = tts_audio[:len(canvas) - start_sample]
+            end_sample = len(canvas)
+        canvas[start_sample:end_sample] = tts_audio
+
+    canvas_path = os.path.join(tmp_dir, "dubbed_canvas.wav")
+    sf.write(canvas_path, canvas, XTTS_SAMPLE_RATE)
+    return canvas_path
 
 
 def swap_audio_in_video(video_path: str, audio_path: str, output_path: str):
@@ -101,7 +195,7 @@ def swap_audio_in_video(video_path: str, audio_path: str, output_path: str):
 # D. Main processing function
 # ---------------------------------------------------------------------------
 
-@spaces.GPU(duration=120)
+@gpu_decorator
 def process_video(youtube_url, video_file, progress=gr.Progress()):
     if not youtube_url and video_file is None:
         raise gr.Error("Please provide a YouTube URL or upload a video file.")
@@ -131,29 +225,30 @@ def process_video(youtube_url, video_file, progress=gr.Progress()):
         with VideoFileClip(video_path) as video:
             video.audio.write_audiofile(ref_audio_path, logger=None)
 
-        # Step 3: Transcribe with Whisper
+        # Step 3: Transcribe with Whisper (segment-level)
         progress(0.40, desc="Transcribing Ukrainian audio...")
         result = whisper_model.transcribe(ref_audio_path, task="transcribe", language="uk")
+        raw_segments = result["segments"]
         ukrainian_text = result["text"]
 
-        # Step 4: Translate with NLLB
-        progress(0.55, desc="Translating to English...")
-        english_text = translate_uk_to_en(ukrainian_text, device)
-
-        # Step 5: Voice clone with XTTS
-        progress(0.70, desc="Cloning voice...")
-        cloned_audio_path = os.path.join(tmp_dir, "cloned_audio.wav")
-        tts_model.tts_to_file(
-            text=english_text,
-            speaker_wav=ref_audio_path,
-            language="en",
-            file_path=cloned_audio_path,
-        )
-
-        # Step 6: Swap audio
-        progress(0.85, desc="Combining video and dubbed audio...")
+        # Step 4: Merge short segments
+        merged = merge_short_segments(raw_segments)
+
+        # Step 5: Context-aware translate
+        progress(0.50, desc="Translating to English...")
+        translated_texts = translate_segments_uk_to_en(merged, device)
+        english_text = " ".join(translated_texts)
+
+        # Step 6: Per-segment TTS + time-stretch + canvas
+        with VideoFileClip(video_path) as v:
+            video_duration = v.duration
+        canvas_path = build_audio_canvas(merged, translated_texts, ref_audio_path, video_duration, tmp_dir, progress)
+
+        # Step 7: Combine video and audio
         output_path = os.path.join(tmp_dir, "dubbed_output.mp4")
-        swap_audio_in_video(video_path, cloned_audio_path, output_path)
+
+        progress(0.90, desc="Combining video and dubbed audio...")
+        swap_audio_in_video(video_path, canvas_path, output_path)
 
         progress(1.0, desc="Done!")
         return output_path, ukrainian_text, english_text