119 lines
4.3 KiB
Python
119 lines
4.3 KiB
Python
import json
|
|
from typing import List, Optional
|
|
|
|
from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException
|
|
from fastapi.responses import FileResponse
|
|
from sqlalchemy.orm import Session, selectinload
|
|
|
|
from app.core.deps import get_current_user
|
|
from app.db.session import get_db
|
|
from app.models import Asset, Job, JobOutput, User
|
|
from app.schemas import JobCreateRequest, JobListResponse, JobResponse
|
|
from app.services.orchestrator import reconcile_job_outputs_if_missing, run_job
|
|
from app.services.storage import output_storage
|
|
|
|
router = APIRouter(prefix="/api/jobs", tags=["jobs"])
|
|
|
|
|
|
@router.post("/", response_model=JobResponse, status_code=201)
|
|
async def create_job(
|
|
payload: JobCreateRequest,
|
|
background_tasks: BackgroundTasks,
|
|
db: Session = Depends(get_db),
|
|
current_user: User = Depends(get_current_user),
|
|
):
|
|
def assert_owns(asset_id: Optional[str], label: str):
|
|
if asset_id:
|
|
asset = (
|
|
db.query(Asset)
|
|
.filter(Asset.id == asset_id, Asset.owner_id == current_user.id, Asset.is_trashed.is_(False))
|
|
.first()
|
|
)
|
|
if not asset:
|
|
raise HTTPException(400, f"{label} asset not found or not owned by user")
|
|
|
|
assert_owns(payload.ground_truth_asset_id, "ground_truth")
|
|
assert_owns(payload.motion_asset_id, "motion")
|
|
assert_owns(payload.audio_asset_id, "audio")
|
|
assert_owns(payload.pose_asset_id, "pose_sheet")
|
|
for ref_id in payload.reference_asset_ids or []:
|
|
assert_owns(ref_id, f"reference {ref_id}")
|
|
|
|
job = Job(
|
|
owner_id=current_user.id,
|
|
mode=payload.mode,
|
|
submode=payload.submode,
|
|
prompt=payload.prompt,
|
|
negative_prompt=payload.negative_prompt,
|
|
status="created",
|
|
ground_truth_asset_id=payload.ground_truth_asset_id,
|
|
motion_asset_id=payload.motion_asset_id,
|
|
audio_asset_id=payload.audio_asset_id,
|
|
pose_asset_id=payload.pose_asset_id,
|
|
reference_asset_ids_json=json.dumps(payload.reference_asset_ids) if payload.reference_asset_ids else None,
|
|
settings_json=json.dumps(payload.settings) if payload.settings else None,
|
|
)
|
|
db.add(job)
|
|
db.commit()
|
|
db.refresh(job)
|
|
|
|
background_tasks.add_task(run_job, job.id)
|
|
return job
|
|
|
|
|
|
@router.get("/", response_model=List[JobListResponse])
|
|
def list_jobs(
|
|
mode: Optional[str] = None,
|
|
status: Optional[str] = None,
|
|
db: Session = Depends(get_db),
|
|
current_user: User = Depends(get_current_user),
|
|
):
|
|
query = db.query(Job).filter(Job.owner_id == current_user.id)
|
|
query = query.options(selectinload(Job.outputs))
|
|
if mode:
|
|
query = query.filter(Job.mode == mode)
|
|
if status:
|
|
query = query.filter(Job.status == status)
|
|
return query.order_by(Job.created_at.desc()).limit(100).all()
|
|
|
|
|
|
@router.get("/{job_id}", response_model=JobResponse)
|
|
async def get_job(job_id: str, db: Session = Depends(get_db), current_user: User = Depends(get_current_user)):
|
|
job = (
|
|
db.query(Job)
|
|
.options(selectinload(Job.outputs), selectinload(Job.events))
|
|
.filter(Job.id == job_id, Job.owner_id == current_user.id)
|
|
.first()
|
|
)
|
|
if not job:
|
|
raise HTTPException(404, "Job not found")
|
|
if job.status == "completed" and not job.outputs and job.comfy_prompt_id:
|
|
await reconcile_job_outputs_if_missing(job.id)
|
|
db.expire_all()
|
|
job = (
|
|
db.query(Job)
|
|
.options(selectinload(Job.outputs), selectinload(Job.events))
|
|
.filter(Job.id == job_id, Job.owner_id == current_user.id)
|
|
.first()
|
|
)
|
|
return job
|
|
|
|
|
|
@router.get("/{job_id}/outputs/{output_id}/download")
|
|
def download_output(
|
|
job_id: str,
|
|
output_id: str,
|
|
db: Session = Depends(get_db),
|
|
current_user: User = Depends(get_current_user),
|
|
):
|
|
job = db.query(Job).filter(Job.id == job_id, Job.owner_id == current_user.id).first()
|
|
if not job:
|
|
raise HTTPException(404, "Job not found")
|
|
output = db.query(JobOutput).filter(JobOutput.id == output_id, JobOutput.job_id == job_id).first()
|
|
if not output:
|
|
raise HTTPException(404, "Output not found")
|
|
abs_path = output_storage.absolute_path(output.file_path)
|
|
if not abs_path.exists():
|
|
raise HTTPException(404, "Output file not found")
|
|
return FileResponse(str(abs_path))
|