Update app.py

app.py

@@ -3,18 +3,18 @@ import tempfile
 from pathlib import Path
 import gradio as gr
 import ffmpeg
-from faster_whisper import WhisperModel
+import whisper
 
 # -------- 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 %= 3600000
+    minutes = ms // 60000
+    ms %= 60000
+    seconds = ms // 1000
+    millis = ms % 1000
+    return f"{hours:02d}:{minutes:02d}:{seconds:02d},{millis:03d}"
 
 
 def segments_to_srt(segments: list) -> str:
@@ -22,71 +22,98 @@ def segments_to_srt(segments: list) -> str:
     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)
+        text = seg["text"].strip().replace("\n", " ")
+        if text:
+            lines.append(f"{i}\n{start_ts} --> {end_ts}\n{text}\n")
     return "\n".join(lines)
 
 
 # -------- Config --------
-MODEL_NAME = "Systran/faster-whisper-small"  # good for HF CPU
-DEVICE = "cpu"
+MODEL_NAME = "base"
 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.")
+print(f"Loading Whisper model '{MODEL_NAME}'...")
+model = whisper.load_model(MODEL_NAME)
+print("Model loaded successfully!")
 
 
 # -------- Core functions --------
 def extract_audio(input_path: str, out_path: str):
-    """Extracts mono 16 kHz WAV using ffmpeg"""
-    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"):
-    """Transcribe uploaded file to SRT; compatible with HF Spaces"""
-    tmp_dir = Path(tempfile.mkdtemp(prefix="subgen_"))
+    """Extract mono 16 kHz WAV using ffmpeg"""
+    (
+        ffmpeg
+        .input(input_path)
+        .output(out_path, format="wav", acodec="pcm_s16le", ac=1, ar="16000")
+        .overwrite_output()
+        .run(quiet=True)
+    )
 
-    # Handle file object
+
+def transcribe_file_to_srt(file_obj, language="en"):
+    tmp_dir = Path(tempfile.mkdtemp(prefix="subgen_"))
     input_path = tmp_dir / Path(file_obj.name).name
     with open(input_path, "wb") as f:
         f.write(file_obj.read())
 
-    # Extract audio and transcribe
     audio_path = tmp_dir / "audio.wav"
     extract_audio(str(input_path), str(audio_path))
-    segments, _ = model.transcribe(str(audio_path), language=language)
-    segs = [{"start": s.start, "end": s.end, "text": s.text} for s in segments]
-    srt_text = segments_to_srt(segs)
 
-    # Save .srt file
+    result = model.transcribe(str(audio_path), language=language)
+    segments = []
+    for i, seg in enumerate(result["segments"]):
+        segments.append({
+            "start": seg["start"],
+            "end": seg["end"],
+            "text": seg["text"]
+        })
+
+    srt_text = segments_to_srt(segments)
     output_path = OUTPUT_DIR / f"{Path(file_obj.name).stem}.srt"
-    with open(output_path, "w", encoding="utf-8") as f:
-        f.write(srt_text)
+    output_path.write_text(srt_text, encoding="utf-8")
+
     return str(output_path), "✅ Subtitles generated successfully!"
 
 
+# -------- UI Styling --------
+def build_style(theme="light"):
+    if theme == "dark":
+        bg = "#0f2027"
+        color = "#ffffff"
+        button = "#00adb5"
+    else:
+        bg = "#f0f2f5"
+        color = "#000000"
+        button = "#0077ff"
+
+    return f"""
+    <style>
+    body {{
+        background: {bg};
+        color: {color};
+        font-family: 'Poppins', sans-serif;
+        transition: background 0.5s, color 0.5s;
+    }}
+    .gr-button {{
+        background-color: {button} !important;
+        color: white !important;
+        font-weight: bold;
+        border-radius: 10px !important;
+    }}
+    .gr-button:hover {{
+        filter: brightness(1.2);
+    }}
+    </style>
+    """
+
+
 # -------- Gradio UI --------
 with gr.Blocks(title="AI Subtitle Generator") as demo:
+    theme_state = gr.State("light")
+
+    style_html = gr.HTML(build_style("light"))
     gr.HTML("<h1 style='text-align:center;'>🎬 AI Subtitle Generator</h1>")
-    gr.HTML(
-        "<p style='text-align:center;'>Upload a video or audio file to generate English <b>.srt</b> subtitles.</p>"
-    )
+    gr.HTML("<p style='text-align:center;'>Upload a video or audio file to generate English <b>.srt</b> subtitles.</p>")
 
     with gr.Row():
         input_file = gr.File(label="Upload video/audio file")
@@ -94,16 +121,30 @@ with gr.Blocks(title="AI Subtitle Generator") as demo:
 
     status_box = gr.Textbox(label="Status", interactive=False)
 
-    def on_click(file):
+    with gr.Row():
+        generate_btn = gr.Button("🚀 Generate Subtitles")
+        clear_btn = gr.Button("🧹 Clear Chat")
+        theme_btn = gr.Button("🌗 Toggle Theme")
+
+    # Button logic
+    def on_generate(file):
+        if not file:
+            return None, "⚠️ Please upload a file first!"
         srt_path, msg = transcribe_file_to_srt(file)
         return srt_path, msg
 
-    generate_btn = gr.Button("Generate Subtitles")
-    generate_btn.click(on_click, inputs=[input_file], outputs=[output_file, status_box])
+    def on_clear():
+        return None, None, ""
 
-    gr.HTML(
-        "<p style='text-align:center;font-size:14px;opacity:0.7;'>Powered by Faster-Whisper + Gradio</p>"
-    )
+    def on_theme(current):
+        new_theme = "dark" if current == "light" else "light"
+        return new_theme, gr.update(value=build_style(new_theme))
+
+    generate_btn.click(on_generate, inputs=[input_file], outputs=[output_file, status_box])
+    clear_btn.click(on_clear, outputs=[input_file, output_file, status_box])
+    theme_btn.click(on_theme, inputs=[theme_state], outputs=[theme_state, style_html])
+
+    gr.HTML("<p style='text-align:center;font-size:14px;opacity:0.6;'>✨ Built with OpenAI Whisper + Gradio</p>")
 
 if __name__ == "__main__":
     demo.launch()