engine / api.py
VeuReu's picture
Upload 8 files
3d74263 verified
raw
history blame
6.26 kB
# api.py — versión corregida (orden de definición)
import os
import uuid
from fastapi import FastAPI, UploadFile, File, Form, Depends, Header, HTTPException, APIRouter
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from typing import Optional
from models_job import JobCreate, JobStatus, JobResult
from queue_manager import job_store, job_queue, start_worker, UPLOAD_DIR
from worker import process_job
from pydantic import BaseModel
import subprocess
import tempfile
import base64
import requests
API_SHARED_TOKEN = os.environ.get("API_SHARED_TOKEN")
UI_SPACE_URL = os.environ.get("UI_SPACE_URL") # ej: https://org-tu--ui--space.hf.space
# ---------- Matxa - Alvocat (router) ----------
router = APIRouter()
HF_TOKEN = os.getenv("HF_TOKEN", "")
MATXA_TTS_URL = os.getenv("MATXA_TTS_URL", "").strip()
INFERENCE_URL = "https://api-inference.huggingface.co/models/projecte-aina/matxa-alvocat"
class TTSRequest(BaseModel):
text: str
@router.post("/tts/matxa")
def tts_matxa(req: TTSRequest):
text = (req.text or "").strip()
if not text:
raise HTTPException(status_code=400, detail="Empty text")
try:
if MATXA_TTS_URL:
headers = {}
if HF_TOKEN:
headers["Authorization"] = f"Bearer {HF_TOKEN}"
resp = requests.post(
MATXA_TTS_URL,
headers=headers,
json={"text": text},
timeout=60,
)
if resp.status_code != 200:
raise HTTPException(status_code=502, detail=f"Space TTS error: {resp.text}")
if resp.headers.get("content-type", "").startswith("audio/"):
audio_bytes = resp.content
b64 = base64.b64encode(audio_bytes).decode("utf-8")
return {"mp3_data_url": f"data:audio/mpeg;base64,{b64}"}
else:
data = resp.json()
if "audio" in data and isinstance(data["audio"], str) and data["audio"].startswith("data:audio"):
return {"mp3_data_url": data["audio"]}
elif "audio_b64" in data:
audio_bytes = base64.b64decode(data["audio_b64"])
b64 = base64.b64encode(audio_bytes).decode("utf-8")
return {"mp3_data_url": f"data:audio/mpeg;base64,{b64}"}
else:
audio_bytes = data.get("bytes")
if isinstance(audio_bytes, str):
audio_bytes = base64.b64decode(audio_bytes)
b64 = base64.b64encode(audio_bytes).decode("utf-8")
return {"mp3_data_url": f"data:audio/mpeg;base64,{b64}"}
else:
if not HF_TOKEN:
raise HTTPException(status_code=500, detail="HF_TOKEN not set")
headers = {
"Authorization": f"Bearer {HF_TOKEN}",
"Accept": "audio/mpeg",
}
resp = requests.post(
INFERENCE_URL,
headers=headers,
json={"inputs": text},
timeout=60,
)
if resp.status_code != 200:
raise HTTPException(status_code=502, detail=f"Inference API error: {resp.text}")
audio_bytes = resp.content
b64 = base64.b64encode(audio_bytes).decode("utf-8")
return {"mp3_data_url": f"data:audio/mpeg;base64,{b64}"}
except HTTPException:
raise
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# ---------- FastAPI app principal ----------
app = FastAPI(title="Veureu AD – API Space")
# CORS (restringe a tu UI Space si pasas UI_SPACE_URL)
app.add_middleware(
CORSMiddleware,
allow_origins=[UI_SPACE_URL] if UI_SPACE_URL else ["*"],
allow_credentials=False,
allow_methods=["*"],
allow_headers=["*"],
)
# Lanza el worker al arrancar
start_worker(process_job)
# -------- Auth sencilla por token compartido --------
def check_auth(authorization: Optional[str] = Header(None)):
if not API_SHARED_TOKEN:
return True
if not authorization or not authorization.startswith("Bearer "):
raise HTTPException(401, "Missing token")
if authorization.split(" ", 1)[1] != API_SHARED_TOKEN:
raise HTTPException(403, "Invalid token")
return True
# -------- Rutas "jobs" --------
@app.get("/")
def read_root():
return {"message": "Hello World"}
@app.post("/jobs")
async def create_job(
mode: str = Form(default="both"),
video_file: Optional[UploadFile] = File(default=None),
video_url: Optional[str] = Form(default=None),
_auth=Depends(check_auth),
):
if not video_file and not video_url:
raise HTTPException(400, "Debe enviarse un 'video_file' o un 'video_url'.")
job_id = str(uuid.uuid4())
local_path = None
if video_file:
os.makedirs(UPLOAD_DIR, exist_ok=True)
save_path = os.path.join(UPLOAD_DIR, f"{job_id}_{video_file.filename}")
with open(save_path, "wb") as f:
f.write(await video_file.read())
local_path = save_path
st = JobStatus(job_id=job_id, status="queued", progress=0, message="En cola")
job_store.set_status(job_id, st)
job_queue.put({"job_id": job_id, "mode": mode, "local_path": local_path, "video_url": video_url})
return {"job_id": job_id}
@app.get("/jobs/{job_id}/status", response_model=JobStatus)
def get_status(job_id: str, _auth=Depends(check_auth)):
st = job_store.get_status(job_id)
if not st:
raise HTTPException(404, "Job no encontrado")
return st
@app.get("/jobs/{job_id}/result", response_model=JobResult)
def get_result(job_id: str, _auth=Depends(check_auth)):
res = job_store.get_result(job_id)
if not res:
st = job_store.get_status(job_id)
if st and st.status != "completed":
raise HTTPException(409, "El job no ha terminado")
raise HTTPException(404, "Resultado no encontrado")
return res
# <<< AHORA SÍ >>>
app.include_router(router)