Update app.py

app.py

@@ -2,176 +2,145 @@ import os
 import subprocess
 import requests
 import gradio as gr
-import srt
 import datetime
 from googletrans import Translator
 
-# -----------------------
-# Load HF Token safely
-# -----------------------
-HF_TOKEN = os.environ.get("HF_TOKEN")  
-HF_STT_MODEL = "openai/whisper-small"  
+HF_TOKEN = os.environ.get("HF_TOKEN")
+HF_MODEL = "openai/whisper-small"
 
 translator = Translator()
 
-# -----------------------
-# FFmpeg helpers
-# -----------------------
-def extract_audio(video_path, out_audio="audio.wav"):
+# ---------------------------
+# Extract audio with FFmpeg
+# ---------------------------
+def extract_audio(video_path, audio_out="audio.wav"):
     cmd = [
-        "ffmpeg", "-y", "-i", str(video_path),
-        "-vn", "-ac", "1", "-ar", "16000", "-f", "wav", out_audio
+        "ffmpeg", "-y", "-i", video_path,
+        "-vn", "-ac", "1", "-ar", "16000",
+        "-f", "wav", audio_out
     ]
-    proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-    if proc.returncode != 0:
-        raise RuntimeError(proc.stderr.decode())
-    return out_audio
+    subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    return audio_out
 
-def get_audio_duration(path):
+# ---------------------------
+# Get audio duration
+# ---------------------------
+def get_duration(path):
     cmd = [
         "ffprobe", "-v", "error", "-show_entries",
         "format=duration", "-of",
         "default=noprint_wrappers=1:nokey=1", path
     ]
-    proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-    if proc.returncode != 0:
-        return None
+    p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
     try:
-        return float(proc.stdout.decode().strip())
+        return float(p.stdout.decode().strip())
     except:
-        return None
+        return 0.0
 
-# -----------------------
-# HF API Call
-# -----------------------
-def hf_asr_transcribe(audio_file_path, model=HF_STT_MODEL, hf_token=HF_TOKEN):
-    if not hf_token:
-        raise RuntimeError("HF_TOKEN missing! Add it in your Space settings.")
+# ---------------------------
+# HF API → Speech-to-text
+# ---------------------------
+def transcribe_audio(audio_path):
+    if HF_TOKEN is None:
+        raise RuntimeError("HF_TOKEN missing in Space settings.")
 
-    url = f"https://api-inference.huggingface.co/models/{model}"
-    headers = {"Authorization": f"Bearer {hf_token}"}
+    url = f"https://api-inference.huggingface.co/models/{HF_MODEL}"
+    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
 
-    with open(audio_file_path, "rb") as f:
+    with open(audio_path, "rb") as f:
         audio_bytes = f.read()
 
-    response = requests.post(url, headers=headers, data=audio_bytes, timeout=120)
+    r = requests.post(url, headers=headers, data=audio_bytes)
+    out = r.json()
 
-    if response.status_code != 200:
-        raise RuntimeError(f"HF API Error {response.status_code}: {response.text}")
+    if "text" in out:
+        return out["text"]
+    return ""
 
-    try:
-        data = response.json()
-    except:
-        return response.text
-
-    if isinstance(data, dict) and "text" in data:
-        return data["text"]
-
-    if isinstance(data, list):
-        texts = []
-        for item in data:
-            if isinstance(item, dict) and "text" in item:
-                texts.append(item["text"])
-        if texts:
-            return " ".join(texts)
-
-    return str(data)
-
-# -----------------------
-# Generate SRT
-# -----------------------
-def transcript_to_srt(transcript, duration_seconds, avg_segment_len=4.0):
-    words = transcript.strip().split()
+# ---------------------------
+# Build SRT manually
+# ---------------------------
+def build_srt(text, duration):
+    words = text.split()
     if not words:
         return ""
 
-    if duration_seconds and duration_seconds > 0:
-        wps = len(words) / duration_seconds
-    else:
-        wps = 2.5
+    words_per_second = len(words) / max(duration, 1)
+    seg_len = 4  # seconds per subtitle segment
+    w_per_seg = max(1, int(seg_len * words_per_second))
+
+    srt_text = ""
+    start_sec = 0
+    index = 1
+    i = 0
 
-    words_per_segment = max(1, int(round(wps * avg_segment_len)))
-    segments = []
-    seg_index = 1
-    idx = 0
-    current_time = 0.0
+    while i < len(words):
+        segment_words = words[i:i + w_per_seg]
+        end_sec = start_sec + seg_len
 
-    while idx < len(words):
-        chunk = words[idx: idx + words_per_segment]
-        text = " ".join(chunk)
+        start_time = str(datetime.timedelta(seconds=start_sec))
+        end_time = str(datetime.timedelta(seconds=end_sec))
+        
+        # Ensure hh:mm:ss,ms format
+        if "." not in start_time:
+            start_time += ",000"
+        else:
+            start_time = start_time.replace(".", ",")[:12]
 
-        start = datetime.timedelta(seconds=current_time)
-        end = datetime.timedelta(seconds=current_time + (len(chunk) / wps))
+        if "." not in end_time:
+            end_time += ",000"
+        else:
+            end_time = end_time.replace(".", ",")[:12]
 
-        segments.append(
-            srt.Subtitle(
-                index=seg_index,
-                start=start,
-                end=end,
-                content=text
-            )
-        )
+        srt_text += f"{index}\n{start_time} --> {end_time}\n{' '.join(segment_words)}\n\n"
 
-        seg_index += 1
-        idx += words_per_segment
-        current_time += len(chunk) / wps
+        start_sec = end_sec
+        index += 1
+        i += w_per_seg
 
-    return srt.compose(segments)
+    return srt_text
 
-# -----------------------
-# Main Function
-# -----------------------
-def process_video(video_file, target_lang):
-    audio_path = "audio.wav"
-    extract_audio(video_file, audio_path)
-    duration = get_audio_duration(audio_path) or 0.0
+# ---------------------------
+# Master function
+# ---------------------------
+def process_video(video, lang):
+    audio_path = extract_audio(video)
+    duration = get_duration(audio_path)
 
-    transcript = hf_asr_transcribe(audio_path)
+    text = transcribe_audio(audio_path)
 
-    if target_lang != "original":
-        try:
-            transcript = translator.translate(transcript, dest=target_lang).text
-        except:
-            pass
+    if lang != "original":
+        text = translator.translate(text, dest=lang).text
 
-    srt_text = transcript_to_srt(transcript, duration)
-    srt_path = "subtitles.srt"
+    srt_output = build_srt(text, duration)
 
-    with open(srt_path, "w", encoding="utf-8") as f:
-        f.write(srt_text)
+    output_path = "subtitles.srt"
+    with open(output_path, "w", encoding="utf-8") as f:
+        f.write(srt_output)
 
-    return srt_path
+    return output_path
 
-# -----------------------
+# ---------------------------
 # Gradio UI
-# -----------------------
+# ---------------------------
 languages = {
-    "original": "No translation",
+    "original": "Original",
     "en": "English",
     "ur": "Urdu",
     "hi": "Hindi",
     "ps": "Pashto",
-    "ar": "Arabic",
+    "ar": "Arabic"
 }
 
-def gradio_process(video, lang):
-    if video is None:
-        return None
-    try:
-        return process_video(video, lang)
-    except Exception as e:
-        raise gr.Error(str(e))
-
 demo = gr.Interface(
-    fn=gradio_process,
+    fn=process_video,
     inputs=[
         gr.Video(label="Upload Video"),
         gr.Dropdown(choices=list(languages.keys()), value="original", label="Translate Into")
     ],
     outputs=gr.File(label="Download SRT"),
-    title="AI Subtitle Generator (Hugging Face API)",
-    description="Upload video → Auto subtitles → Optional translation → Download .srt"
+    title="AI Subtitle Generator",
+    description="Upload a video → auto subtitles → optional translation → download SRT."
 )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()