Ni2 / main.py
Hong-2's picture
Create main.py
57b7bf6 verified
Raw
History Blame Contribute Delete
2.81 kB
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
import subprocess
import os
import uuid
import shutil
app = FastAPI()
# Mở CORS để Frontend HTML của bạn có thể gọi API mà không bị chặn
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Tạo thư mục chứa video nếu chưa có
os.makedirs("media", exist_ok=True)
# Cung cấp đường dẫn tĩnh để xem video
app.mount("/media", StaticFiles(directory="media"), name="media")
class RenderRequest(BaseModel):
code: str
scene_name: str
quality: str = "m" # l: 480p, m: 720p, h: 1080p
@app.post("/render")
async def render_video(req: RenderRequest):
# Tạo ID duy nhất cho mỗi lần render để không bị trùng file
run_id = str(uuid.uuid4())
py_file = f"{run_id}.py"
# Ghi code AI sinh ra vào file python
with open(py_file, "w", encoding="utf-8") as f:
f.write(req.code)
q_flag = f"-q{req.quality}"
out_dir = f"temp_{run_id}"
# Lệnh gọi thư viện manim từ terminal
cmd = ["manim", q_flag, "--media_dir", out_dir, py_file, req.scene_name]
try:
# Chạy manim, giới hạn thời gian render tối đa 120 giây
result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
if result.returncode != 0:
raise Exception(result.stderr or result.stdout)
# Tìm file mp4 được sinh ra trong thư mục kết quả
video_path = None
for root, dirs, files in os.walk(out_dir):
for file in files:
if file.endswith(".mp4"):
video_path = os.path.join(root, file)
break
if not video_path:
raise Exception("Render xong nhưng không tìm thấy file MP4.")
# Copy video ra thư mục /media công khai
final_url_path = f"/media/{run_id}.mp4"
shutil.copy(video_path, f".{final_url_path}")
return {"video_url": final_url_path}
except subprocess.TimeoutExpired:
raise HTTPException(status_code=408, detail={"stderr": "Render quá thời gian (Timeout 120s)."})
except Exception as e:
raise HTTPException(status_code=400, detail={"stderr": str(e)})
finally:
# Dọn dẹp rác sau khi render xong (hoặc lỗi)
if os.path.exists(py_file):
os.remove(py_file)
if os.path.exists(out_dir):
shutil.rmtree(out_dir)
@app.get("/")
def read_root():
return {"status": "Manim Backend is running!"}