Update app.py

app.py

@@ -9,69 +9,125 @@ from faster_whisper import WhisperModel
 
 # --- Helper functions ---
 def _format_timestamp(seconds: float) -> str:
-ms = int(round(seconds * 1000))
-hours = ms // 3600000
-ms_rem = ms % 3600000
-minutes = ms_rem // 60000
-ms_rem = ms_rem % 60000
-secs = ms_rem // 1000
-millis = ms_rem % 1000
-return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"
-
-
+    ms = int(round(seconds * 1000))
+    hours = ms // 3600000
+    ms_rem = ms % 3600000
+    minutes = ms_rem // 60000
+    ms_rem = ms_rem % 60000
+    secs = ms_rem // 1000
+    millis = ms_rem % 1000
+    return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"
 
 
 def segments_to_srt(segments: list) -> str:
-lines = []
-for i, seg in enumerate(segments, start=1):
-start_ts = _format_timestamp(seg['start'])
-end_ts = _format_timestamp(seg['end'])
-text = seg['text'].replace('\n', ' ').strip()
-if not text:
-continue
-block = f"{i}\n{start_ts} --> {end_ts}\n{text}\n"
-lines.append(block)
-return "\n".join(lines)
-
-
+    lines = []
+    for i, seg in enumerate(segments, start=1):
+        start_ts = _format_timestamp(seg['start'])
+        end_ts = _format_timestamp(seg['end'])
+        text = seg['text'].replace('\n', ' ').strip()
+        if not text:
+            continue
+        block = f"{i}\n{start_ts} --> {end_ts}\n{text}\n"
+        lines.append(block)
+    return "\n".join(lines)
 
 
 # --- Configuration ---
-MODEL_NAME = "Systran/faster-whisper-small" # small model for HF free CPU
+MODEL_NAME = "Systran/faster-whisper-small"  # small model for HF free CPU
 DEVICE = "cpu"
 OUTPUT_DIR = Path("outputs/subtitles")
 OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
 
-
 print(f"Loading model {MODEL_NAME} on {DEVICE} ...")
 model = WhisperModel(MODEL_NAME, device=DEVICE)
 print("Model loaded.")
 
 
+def extract_audio(input_path: str, out_path: str):
+    try:
+        (
+            ffmpeg
+            .input(input_path)
+            .output(out_path, format='wav', acodec='pcm_s16le', ac=1, ar='16000')
+            .overwrite_output()
+            .run(quiet=True)
+        )
+    except ffmpeg.Error as e:
+        stderr = getattr(e, 'stderr', None)
+        msg = stderr.decode() if stderr else str(e)
+        raise RuntimeError(f"ffmpeg error: {msg}")
 
 
-def extract_audio(input_path: str, out_path: str):
-try:
-_ = (
-ffmpeg
-.input(input_path)
-.output(out_path, format='wav', acodec='pcm_s16le', ac=1, ar='16000')
-.overwrite_output()
-.run(quiet=True)
-)
-except ffmpeg.Error as e:
-stderr = getattr(e, 'stderr', None)
-msg = stderr.decode() if stderr else str(e)
-raise RuntimeError(f"ffmpeg error: {msg}")
+def transcribe_file_to_srt(file_obj, language: str = "en"):
+    filename = getattr(file_obj, 'name', getattr(file_obj, 'filename', f"upload_{uuid.uuid4()}.bin"))
+    input_filepath = Path(filename)
 
+    tmp_dir = Path(tempfile.mkdtemp(prefix="subgen_"))
+    saved_input = tmp_dir / (str(uuid.uuid4()) + input_filepath.suffix)
 
+    file_obj.seek(0)
+    with open(saved_input, 'wb') as f:
+        f.write(file_obj.read())
 
+    wav_path = tmp_dir / (saved_input.stem + ".wav")
+    extract_audio(str(saved_input), str(wav_path))
 
-def transcribe_file_to_srt(file_obj, language: str = "en"):
-filename = getattr(file_obj, 'name', getattr(file_obj, 'filename', f"upload_{uuid.uuid4()}.bin"))
-input_filepath = Path(filename)
+    print("Starting transcription... This may take a while depending on file length and model.")
+    segments, info = model.transcribe(str(wav_path), language=language)
+
+    segs = []
+    for seg in segments:
+        start = getattr(seg, 'start', seg.get('start') if isinstance(seg, dict) else None)
+        end = getattr(seg, 'end', seg.get('end') if isinstance(seg, dict) else None)
+        text = getattr(seg, 'text', seg.get('text') if isinstance(seg, dict) else '')
+        if start is None or end is None:
+            continue
+        segs.append({
+            'start': float(start),
+            'end': float(end),
+            'text': str(text).strip()
+        })
+
+    srt_text = segments_to_srt(segs)
+    out_name = f"subtitles_{saved_input.stem}.srt"
+    out_path = OUTPUT_DIR / out_name
+
+    with open(out_path, 'w', encoding='utf-8') as f:
+        f.write(srt_text)
+
+    print(f"Saved SRT to {out_path}")
+    return out_path
+
+
+# --- Gradio UI ---
+def generate_and_return(file):
+    try:
+        srt_path = transcribe_file_to_srt(file, language="en")
+        return srt_path
+    except Exception as e:
+        return None, str(e)
+
+
+with gr.Blocks(title="AI Subtitle Generator — English (.srt)") as demo:
+    gr.Markdown("# 🎬 AI Subtitle Generator (English)\nUpload a video or audio file and download the generated .srt subtitles.")
+
+    with gr.Row():
+        inp = gr.File(label="Upload video or audio file (.mp4, .mkv, .mp3, ...)")
+        out = gr.File(label="Download generated .srt file")
+
+    generate_btn = gr.Button("Generate Subtitles")
+    status = gr.Textbox(label="Status", interactive=False)
+
+    def on_click(file):
+        status.value = "Processing..."
+        path = transcribe_file_to_srt(file, language='en')
+        status.value = f"Done — saved: {path}"
+        return path
+
+    generate_btn.click(on_click, inputs=[inp], outputs=[out, status])
+
+    gr.Markdown("---\n⚙️ **Note:** Make sure ffmpeg is installed. This may take time on CPU. Use a smaller model for faster processing.")
 
 
-tmp_dir = Path(tempfile.mkdtemp(prefix="subgen_"))
-saved_input = tmp_dir / (str(uuid.uuid4()) + input_filepath.suffix)
-demo.launch(share=True)
+if __name__ == "__main__":
+    demo.launch(share=True)