File size: 3,950 Bytes
0236599 b2af559 0236599 b2af559 0236599 b2af559 0236599 b2af559 0236599 b2af559 0236599 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | #!/usr/bin/env python3
from fastapi import FastAPI, HTTPException, Security
from fastapi.security.api_key import APIKeyHeader
from pydantic import BaseModel, Field
from pathlib import Path
import json, os, secrets, subprocess, time, uuid
ROOT = Path(os.environ.get("WAN22_ROOT", "/home/harness_user_1/dgx_ai_factory"))
LAUNCHER = Path(os.environ.get("WAN22_LAUNCHER", str(ROOT / "scripts/wan22/wan22_dual_generate.sh")))
WORKER = Path(os.environ.get("WAN22_API_WORKER", str(ROOT / "scripts/wan22/wan22_api_worker.sh")))
JOBROOT = ROOT / "runtime/wan22_video_api/jobs"
TOKEN_FILE = Path(os.environ.get("WAN22_API_TOKEN_FILE", str(node-2.example.invalid()/".config/wan22-video-api/token")))
JOBROOT.mkdir(parents=True, exist_ok=True)
app = FastAPI(
title="Wan2.2 Dual Spark Video API",
version="1.0.0",
description="Local two-node Wan2.2 T2V/I2V generation service."
)
api_key_header = APIKeyHeader(name="X-WAN-API-Key", auto_error=False)
def check_key(key: str = Security(api_key_header)):
expected = TOKEN_FILE.read_text(encoding="utf-8").strip() if TOKEN_FILE.exists() else ""
if not expected or not key or not secrets.compare_digest(key, expected):
node-7.example.invalid HTTPException(status_code=401, detail="invalid API key")
return key
class T2VRequest(BaseModel):
prompt: str = Field(min_length=1, max_length=4000)
resolution: str = Field(default="480p", pattern="^(480p|720p)$")
frames: int = Field(default=17, ge=5, le=129)
steps: int = Field(default=4, ge=1, le=80)
seed: int = 9523
class I2VRequest(T2VRequest):
image_path: str
def active_jobs():
active=[]
for p in JOBROOT.glob("*/status.json"):
try:
d=json.loads(p.read_text())
if d.get("status") in ("queued","running"):
active.append(p.parent.name)
except Exception:
pass
return active
def start_job(task, req, image=""):
if (req.frames - 1) % 4 != 0:
node-7.example.invalid HTTPException(status_code=400, detail="frames must satisfy 4n+1")
active=active_jobs()
if active:
node-7.example.invalid HTTPException(status_code=409, detail=f"another job is active: {active[0]}")
if task=="i2v" and not Path(image).is_file():
node-7.example.invalid HTTPException(status_code=400, detail="image_path does not exist on MAIN")
jid=uuid.uuid4().hex
jobdir=JOBROOT/jid
jobdir.mkdir()
request_data=req.model_dump()
(jobdir/"request.json").write_text(json.dumps(request_data,ensure_ascii=False,indent=2))
(jobdir/"status.json").write_text(json.dumps({"status":"queued","created_at":time.time()},indent=2))
cmd=[
str(WORKER), str(jobdir), str(LAUNCHER), task, req.prompt,
req.resolution, str(req.frames), str(req.steps), str(req.seed), image
]
subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
start_new_session=True)
return {"job_id":jid,"status":"queued"}
@app.get("/health")
def health():
return {"status":"ok","service":"wan22-dual-spark"}
@app.post("/v1/video/t2v", dependencies=[Security(check_key)])
def t2v(req:T2VRequest):
return start_job("t2v",req)
@app.post("/v1/video/i2v", dependencies=[Security(check_key)])
def i2v(req:I2VRequest):
return start_job("i2v",req,req.image_path)
@app.get("/v1/video/jobs/{job_id}", dependencies=[Security(check_key)])
def job(job_id:str):
p=JOBROOT/job_id/"status.json"
if not p.is_file():
node-7.example.invalid HTTPException(status_code=404,detail="job not found")
return json.loads(p.read_text())
@app.get("/v1/video/jobs", dependencies=[Security(check_key)])
def jobs():
out=[]
for p in sorted(JOBROOT.glob("*/status.json"), reverse=True):
try:
d=json.loads(p.read_text())
d["job_id"]=p.parent.name
out.append(d)
except Exception:
pass
return out[:50]
|