main.py

from fastapi import FastAPI, HTTPException, Request
from pydantic import BaseModel
from youtubeaudio import YoutubeMedia
import subprocess
import os
import uuid
import logging
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
import ffmpeg

class Url(BaseModel):
    url: str
format = "%(asctime)s: %(message)s"

logging.basicConfig(format=format, level=logging.DEBUG,
                    datefmt="%H:%M:%S")
MODEL=os.environ['model']
BASE_DIR="/tmp/holosubs"
VIDEO_DIR=os.path.join(BASE_DIR,"videos")

app = FastAPI()
app.mount("/videos", StaticFiles(directory=VIDEO_DIR), name="videos")

# CORS
origins = ["*"]
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.post("/captions/")
def read_root(url: Url, request: Request):
    requestID=str(uuid.uuid4())
    # Download wav file and get filename
    ytaudio=YoutubeMedia(url, os.path.join(BASE_DIR,requestID))
    ytaudio.download()
    audio_filename, video_filename = ytaudio.audio_filename, ytaudio.video_filename
    # Resample file
    ytaudio.resample('16k')
    
    # Generate subtitles
    subtitle_filename=os.path.join("/tmp/holosubs/tracks", requestID)
    logging.info(f'Output will be writen to {subtitle_filename}.srt')
    cmd=['/usr/local/bin/whisper','-m',f'/home/user/root/models/ggml-{MODEL}.bin'
         ,'-f',audio_filename, '-of', subtitle_filename, '-tr', '-osrt', '-mc', '0']
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    (output, err) = p.communicate()
    p_status = p.wait()
    if err:
        logging.error("Translation failed with error",err)
        raise HTTPException(status_code=500, detail="Whisper translation failed")
    media=ffmpeg.input(video_filename)
    ffmpeg.concat(media.video. \
        filter('subtitles', subtitle_filename+".srt"), \
        media.audio, v=1, a=1). \
        output(os.path.join(VIDEO_DIR,requestID+".mp4")).run()
    
    # Return video url delivered
    return {"video": f'/videos/{requestID+".mp4"}'}