Tberox / app.py
understanding's picture
Update app.py
b34dcc3 verified
import httpx
import ffmpeg
import os
import sqlite3
import uvicorn
import uuid
from fastapi import FastAPI, BackgroundTasks
from pydantic import BaseModel
from starlette.responses import FileResponse, JSONResponse
from typing import Optional
# --- Configuration ---
app = FastAPI()
DOWNLOAD_DIR = "downloads"
ENCODED_DIR = "encoded"
DB_FILE = "jobs.db"
os.makedirs(DOWNLOAD_DIR, exist_ok=True)
os.makedirs(ENCODED_DIR, exist_ok=True)
# ---------------------
# --- Database Setup ---
def get_db():
conn = sqlite3.connect(DB_FILE)
conn.row_factory = sqlite3.Row
return conn
def init_db():
with get_db() as conn:
conn.execute('''
CREATE TABLE IF NOT EXISTS jobs (
task_id TEXT PRIMARY KEY,
status TEXT NOT NULL,
file_name TEXT,
error_message TEXT,
duration INTEGER
)
''')
conn.commit()
# --- Background Re-encoding Task ---
def process_video_task(url: str, cookie: Optional[str], task_id: str):
db = get_db()
temp_in_path = os.path.join(DOWNLOAD_DIR, f"{task_id}_in.mp4")
temp_out_path = os.path.join(ENCODED_DIR, f"{task_id}_out.mp4")
try:
# 1. Update DB: Downloading
db.execute("UPDATE jobs SET status = 'downloading' WHERE task_id = ?", (task_id,))
db.commit()
headers = {
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/135.0.0.0 Safari/537.36",
"Referer": "https://terabox.com/"
}
if cookie:
headers["Cookie"] = cookie
with httpx.stream("GET", url, headers=headers, follow_redirects=True, timeout=600.0) as r:
r.raise_for_status()
with open(temp_in_path, 'wb') as f:
for chunk in r.iter_bytes(chunk_size=8192):
f.write(chunk)
# 2. Update DB: "Fixing" (not encoding)
db.execute("UPDATE jobs SET status = 'encoding' WHERE task_id = ?", (task_id,))
db.commit()
#
# βœ…βœ…βœ… THE FIX IS HERE βœ…βœ…βœ…
# We changed this from vcodec='libx264' to 'c=copy'
# This is the FAST (0% CPU) "re-package"
#
(
ffmpeg
.input(temp_in_path)
.output(temp_out_path, c='copy', movflags='+faststart') # <-- CHANGED
.run(capture_stdout=True, capture_stderr=True)
)
# 3. Get duration using ffprobe
probe = ffmpeg.probe(temp_out_path)
duration = int(float(probe['format']['duration'])) # Get duration in seconds
# 4. Update DB: Complete
final_file_name = f"{task_id}_out.mp4"
db.execute(
"UPDATE jobs SET status = 'complete', file_name = ?, duration = ? WHERE task_id = ?",
(final_file_name, duration, task_id)
)
db.commit()
except Exception as e:
error_msg = str(e)
if hasattr(e, 'stderr'):
error_msg = e.stderr.decode()
print(f"Error processing task {task_id}: {error_msg}")
db.execute("UPDATE jobs SET status = 'error', error_message = ? WHERE task_id = ?", (error_msg, task_id))
db.commit()
finally:
if os.path.exists(temp_in_path):
os.remove(temp_in_path)
db.close()
# --- API Endpoints ---
class VideoRequest(BaseModel):
url: str
cookie: Optional[str] = None
@app.post("/process")
async def start_processing(request: VideoRequest, background_tasks: BackgroundTasks):
""" START the encoding job. """
task_id = str(uuid.uuid4())
with get_db() as db:
db.execute("INSERT INTO jobs (task_id, status) VALUES (?, 'queued')", (task_id,))
db.commit()
background_tasks.add_task(process_video_task, request.url, request.cookie, task_id)
return {"status": "queued", "task_id": task_id}
@app.get("/status/{task_id}")
async def get_status(task_id: str):
""" CHECK the status of the encoding job. """
with get_db() as db:
job = db.execute("SELECT * FROM jobs WHERE task_id = ?", (task_id,)).fetchone()
if not job:
return JSONResponse(status_code=404, content={"status": "not_found"})
return dict(job) # This will now include 'duration'
@app.get("/download/{file_name}")
async def download_file(file_name: str):
""" DOWNLOAD the final re-encoded file. """
file_path = os.path.join(ENCODED_DIR, file_name)
if not os.path.exists(file_path):
return JSONResponse(status_code=404, content={"status": "file_not_found"})
return FileResponse(file_path, media_type='video/mp4', filename=file_name)
@app.get("/")
async def health_check():
""" Health check endpoint. """
return {"status": "ok"}
# --- Server Start ---
@app.on_event("startup")
async def startup_event():
init_db() # Create the database and table on startup
print("Database initialized.")
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860)