publisher/tasks/subtitles.py

"""
BASYX V11 — Subtitles Generator
Creates SRT subtitles from transcription
"""

import os
import tempfile
import asyncio
from faster_whisper import WhisperModel
import httpx


# --------------------------------------------------
# GLOBAL MODEL (shared)
# --------------------------------------------------

MODEL = WhisperModel(
    model_size_or_path="base",
    device="cpu",
    compute_type="int8"
)


# --------------------------------------------------
# HELPERS
# --------------------------------------------------

async def save_upload(file):
    suffix = os.path.splitext(file.filename)[-1]
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
    tmp.write(await file.read())
    tmp.close()
    return tmp.name


async def download_url(url: str):
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")

    async with httpx.AsyncClient(timeout=300) as client:
        async with client.stream("GET", url) as r:
            r.raise_for_status()
            async for chunk in r.aiter_bytes():
                tmp.write(chunk)

    tmp.close()
    return tmp.name


def format_time(seconds: float):
    h = int(seconds // 3600)
    m = int((seconds % 3600) // 60)
    s = int(seconds % 60)
    ms = int((seconds - int(seconds)) * 1000)

    return f"{h:02}:{m:02}:{s:02},{ms:03}"


def build_srt(segments):
    lines = []

    for i, seg in enumerate(segments, start=1):
        lines.append(str(i))
        lines.append(
            f"{format_time(seg.start)} --> {format_time(seg.end)}"
        )
        lines.append(seg.text.strip())
        lines.append("")

    return "\n".join(lines)


# --------------------------------------------------
# MAIN TASK
# --------------------------------------------------

async def run(context):

    media_path = None

    try:
        # ---------------- INPUT ----------------

        if context.input_file:
            media_path = await save_upload(context.input_file)

        elif context.url_input:
            media_path = await download_url(context.url_input)

        else:
            return {
                "status": "error",
                "message": "No input provided"
            }

        # ---------------- TRANSCRIBE ----------------

        segments, info = await asyncio.to_thread(
            MODEL.transcribe,
            media_path,
            beam_size=5
        )

        # ---------------- BUILD SRT ----------------

        srt_text = build_srt(list(segments))

        # Save file
        srt_path = tempfile.NamedTemporaryFile(
            delete=False,
            suffix=".srt"
        ).name

        with open(srt_path, "w", encoding="utf-8") as f:
            f.write(srt_text)

        return {
            "status": "success",
            "language": info.language,
            "srt_path": srt_path,
            "preview": srt_text[:1000]
        }

    except Exception as e:
        return {
            "status": "error",
            "message": str(e)
        }

    finally:
        if media_path and os.path.exists(media_path):
            os.remove(media_path)