| |
| 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] |
|
|
