Initial Animatrix import

This commit is contained in:
Sagnik
2026-04-17 19:11:57 +05:30
commit c7994d17a9
60 changed files with 8516 additions and 0 deletions

View File

@@ -0,0 +1 @@

View File

@@ -0,0 +1,28 @@
from fastapi import APIRouter, Depends
from sqlalchemy.orm import Session
from app.core.deps import get_current_user
from app.db.session import get_db
from app.models import Job, User
from app.services.comfy_client import comfy_client
router = APIRouter(prefix="/api/admin", tags=["admin"])
@router.get("/health")
async def health(_: User = Depends(get_current_user)):
return {"api": "ok", "comfyui": await comfy_client.health_check()}
@router.get("/queue")
async def queue(_: User = Depends(get_current_user)):
return await comfy_client.get_queue()
@router.get("/jobs-summary")
def jobs_summary(db: Session = Depends(get_db), _: User = Depends(get_current_user)):
total = db.query(Job).count()
active = db.query(Job).filter(Job.status.in_(["validating", "uploading_assets", "queued", "executing", "collecting_outputs"])).count()
completed = db.query(Job).filter(Job.status == "completed").count()
failed = db.query(Job).filter(Job.status == "failed").count()
return {"total": total, "active": active, "completed": completed, "failed": failed}

View File

@@ -0,0 +1,120 @@
from datetime import datetime, timedelta, timezone
from typing import List, Optional
from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile
from sqlalchemy.orm import Session
from app.core.deps import get_current_user
from app.db.session import get_db
from app.models import Asset, User
from app.schemas import AssetResponse, AssetTrashRequest
from app.services.storage import asset_storage
router = APIRouter(prefix="/api/assets", tags=["assets"])
ALLOWED_TYPES = {
"image": ["image/jpeg", "image/png", "image/webp"],
"video": ["video/mp4", "video/webm", "video/quicktime"],
"audio": ["audio/mpeg", "audio/mp4", "audio/wav", "audio/ogg", "audio/x-wav"],
"pose_sheet": ["image/jpeg", "image/png", "image/webp"],
}
MAX_SIZE_BYTES = 500 * 1024 * 1024
@router.post("/upload", response_model=AssetResponse, status_code=201)
async def upload_asset(
file: UploadFile = File(...),
asset_type: str = Form(...),
db: Session = Depends(get_db),
current_user: User = Depends(get_current_user),
):
if asset_type not in ALLOWED_TYPES:
raise HTTPException(400, f"asset_type must be one of {list(ALLOWED_TYPES.keys())}")
mime = file.content_type or ""
if mime not in ALLOWED_TYPES[asset_type]:
raise HTTPException(400, f"Unsupported mime type {mime} for {asset_type}")
subfolder = f"{current_user.id}/{asset_type}"
storage_path, size_bytes = await asset_storage.save_upload(file, subfolder)
if size_bytes > MAX_SIZE_BYTES:
raise HTTPException(413, "File too large (max 500MB)")
thumbnail_path = None
width = height = None
duration_seconds = None
if asset_type in ("image", "pose_sheet"):
thumbnail_path = asset_storage.generate_thumbnail(storage_path, f"{current_user.id}/thumbs")
try:
from PIL import Image
abs_path = asset_storage.absolute_path(storage_path)
with Image.open(abs_path) as image:
width, height = image.size
except Exception:
pass
elif asset_type == "video":
thumbnail_path = asset_storage.generate_video_thumbnail(storage_path, f"{current_user.id}/thumbs")
duration_seconds = asset_storage.detect_duration_seconds(storage_path)
else:
duration_seconds = asset_storage.detect_duration_seconds(storage_path)
asset = Asset(
owner_id=current_user.id,
asset_type=asset_type,
mime_type=mime,
original_filename=file.filename or "upload",
storage_path=storage_path,
thumbnail_path=thumbnail_path,
size_bytes=size_bytes,
width=width,
height=height,
duration_seconds=duration_seconds,
)
db.add(asset)
db.commit()
db.refresh(asset)
return asset
@router.get("/", response_model=List[AssetResponse])
def list_assets(
asset_type: Optional[str] = None,
include_trashed: bool = False,
db: Session = Depends(get_db),
current_user: User = Depends(get_current_user),
):
query = db.query(Asset).filter(Asset.owner_id == current_user.id)
if not include_trashed:
query = query.filter(Asset.is_trashed.is_(False))
if asset_type:
query = query.filter(Asset.asset_type == asset_type)
return query.order_by(Asset.created_at.desc()).all()
@router.post("/trash", status_code=200)
def move_assets_to_trash(
payload: AssetTrashRequest,
db: Session = Depends(get_db),
current_user: User = Depends(get_current_user),
):
if not payload.asset_ids:
raise HTTPException(400, "No asset ids provided")
assets = (
db.query(Asset)
.filter(Asset.owner_id == current_user.id, Asset.id.in_(payload.asset_ids))
.all()
)
if not assets:
raise HTTPException(404, "No matching assets found")
delete_after_at = datetime.now(timezone.utc) + timedelta(days=30)
for asset in assets:
asset.is_trashed = True
asset.delete_after_at = delete_after_at
db.commit()
return {
"moved_to_trash": len(assets),
"delete_after_at": delete_after_at.isoformat(),
}

View File

@@ -0,0 +1,61 @@
from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
from sqlalchemy.orm import Session
from app.core.config import settings
from app.core.deps import get_current_user
from app.core.security import create_access_token, hash_password, verify_password
from app.db.session import get_db
from app.models import User
from app.schemas import LoginRequest, RegisterRequest, UserResponse
router = APIRouter(prefix="/api/auth", tags=["auth"])
@router.post("/register", response_model=UserResponse, status_code=status.HTTP_201_CREATED)
def register(payload: RegisterRequest, db: Session = Depends(get_db)):
existing = db.query(User).filter(User.email == payload.email).first()
if existing:
raise HTTPException(status_code=400, detail="Email already registered")
user = User(email=payload.email, password_hash=hash_password(payload.password))
db.add(user)
db.commit()
db.refresh(user)
return user
def _is_secure_request(request: Request) -> bool:
forwarded_proto = request.headers.get("x-forwarded-proto", "")
if "https" in forwarded_proto.lower():
return True
if request.url.scheme == "https":
return True
return settings.BACKEND_BASE_URL.startswith("https://")
@router.post("/login")
def login(payload: LoginRequest, request: Request, response: Response, db: Session = Depends(get_db)):
user = db.query(User).filter(User.email == payload.email).first()
if not user or not verify_password(payload.password, user.password_hash):
raise HTTPException(status_code=401, detail="Invalid credentials")
token = create_access_token(subject=user.id)
response.set_cookie(
key="access_token",
value=token,
httponly=True,
samesite="lax",
secure=_is_secure_request(request),
max_age=60 * 60 * 24 * 7,
)
return {"message": "Logged in", "user": UserResponse.model_validate(user)}
@router.post("/logout")
def logout(response: Response):
response.delete_cookie("access_token")
return {"message": "Logged out"}
@router.get("/me", response_model=UserResponse)
def me(current_user: User = Depends(get_current_user)):
return current_user

View File

@@ -0,0 +1,118 @@
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))