""" Wan2.2 TI2V 5B — Text-to-Video & Image-to-Video Uses official Wan2.2 inference code (generate.py) for native T2V + I2V support. FastAPI wrapper for Hugging Face Docker Space. """ import os import io import json import uuid import time import asyncio import subprocess from pathlib import Path from contextlib import asynccontextmanager from datetime import datetime from typing import Optional from fastapi import FastAPI, Form, UploadFile, File, HTTPException from fastapi.responses import FileResponse, HTMLResponse # ---------- Config ---------- CKPT_DIR = os.getenv("CKPT_DIR", "/data/models/Wan2.2-TI2V-5B") WAN22_DIR = Path("/app/Wan2.2") OUTPUT_DIR = Path("/tmp/outputs") OUTPUT_DIR.mkdir(parents=True, exist_ok=True) JOBS_FILE = OUTPUT_DIR / "jobs.json" # Job store _jobs: dict = {} def _load_jobs(): global _jobs if JOBS_FILE.exists(): try: _jobs = json.loads(JOBS_FILE.read_text()) except Exception: _jobs = {} def _save_jobs(): JOBS_FILE.write_text(json.dumps(_jobs, default=str)) @asynccontextmanager async def lifespan(app: FastAPI): _load_jobs() ckpt_ok = Path(CKPT_DIR).exists() print(f"[startup] CKPT_DIR={CKPT_DIR} exists={ckpt_ok}") print(f"[startup] Wan2.2 dir={WAN22_DIR} exists={WAN22_DIR.exists()}") if not ckpt_ok: print("[startup] Downloading model (~54 GB) — this happens once, cached on /data …") try: subprocess.run( ["huggingface-cli", "download", "Wan-AI/Wan2.2-TI2V-5B", "--local-dir", CKPT_DIR, "--repo-type", "model"], check=True, timeout=3600 ) print("[startup] Model downloaded.") except Exception as e: print(f"[startup] Model download failed: {e}") else: print("[startup] Model already cached.") yield app = FastAPI(lifespan=lifespan, title="Wan2.2 TI2V 5B") def _check_model(): if not Path(CKPT_DIR).exists(): raise HTTPException(503, "Model checkpoint not found. Building/downloading…") async def _run_generation(job_id: str, prompt: str, image_path: Optional[str], steps: int, guidance_scale: float, size: str): _jobs[job_id]["status"] = "running" _jobs[job_id]["started_at"] = time.time() _save_jobs() out_path = OUTPUT_DIR / f"{job_id}.mp4" cmd = [ "python", str(WAN22_DIR / "generate.py"), "--task", "ti2v-5B", "--size", size, "--ckpt_dir", CKPT_DIR, "--offload_model", "True", "--convert_model_dtype", "--t5_cpu", "--prompt", prompt, "--sample_steps", str(steps), ] if image_path: cmd += ["--image", image_path] print(f"[generate {job_id}] {'I2V' if image_path else 'T2V'} prompt={prompt[:80]!r}") try: proc = await asyncio.create_subprocess_exec( *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, cwd=str(WAN22_DIR), ) stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=1200) if proc.returncode != 0: err = stderr.decode()[-500:] raise RuntimeError(f"generate.py failed (code {proc.returncode}): {err}") # The official generate.py saves output to the working dir. # Find the generated file and move it. generated = sorted(WAN22_DIR.glob("*.mp4"), key=os.path.getmtime, reverse=True) if generated: import shutil shutil.move(str(generated[0]), str(out_path)) else: raise RuntimeError("No output video produced. stdout: " + stdout.decode()[-200:]) duration = round(time.time() - _jobs[job_id]["started_at"], 1) _jobs[job_id]["status"] = "done" _jobs[job_id]["duration"] = duration _jobs[job_id]["output"] = f"/output/{job_id}.mp4" except asyncio.TimeoutError: _jobs[job_id]["status"] = "error" _jobs[job_id]["error"] = "Generation timed out after 20 minutes." except Exception as e: _jobs[job_id]["status"] = "error" _jobs[job_id]["error"] = str(e) finally: _save_jobs() # ---------- UI ---------- @app.get("/", response_class=HTMLResponse) def index(): return HTMLResponse(""" Wan2.2 TI2V 5B — Video Generator

🎬 Wan2.2 TI2V 5B

Text-to-Video & Image-to-Video · 720p @ 24fps

✨ Create Video

🖼️

Click or drag to upload

📽️ Gallery
📡 API
""") # ---------- Generate ---------- @app.post("/generate") async def generate( prompt: str = Form(...), image: UploadFile = File(None), steps: int = Form(30, ge=20, le=50), size: str = Form("1280*704"), guidance_scale: float = Form(5.0), ): _check_model() # Save uploaded image if present image_path = None if image and image.filename: try: from PIL import Image as PILImage contents = await image.read() img = PILImage.open(io.BytesIO(contents)).convert("RGB") img_path = OUTPUT_DIR / f"{uuid.uuid4().hex[:8]}.jpg" img.save(str(img_path), "JPEG", quality=95) image_path = str(img_path) except Exception as e: raise HTTPException(400, f"Failed to read image: {e}") job_id = uuid.uuid4().hex[:12] _jobs[job_id] = { "id": job_id, "prompt": prompt[:300], "status": "queued", "steps": steps, "size": size, "guidance_scale": guidance_scale, "has_image": image_path is not None, "created_at": datetime.utcnow().isoformat(), } _save_jobs() asyncio.create_task(_run_generation(job_id, prompt, image_path, steps, guidance_scale, size)) # Poll until done (max 20 minutes on T4 with CPU offload) for _ in range(1200): await asyncio.sleep(1) status = _jobs.get(job_id, {}).get("status") if status == "done": return FileResponse( str(OUTPUT_DIR / f"{job_id}.mp4"), media_type="video/mp4", filename=f"wan2.2_{job_id}.mp4" ) if status == "error": raise HTTPException(500, _jobs[job_id].get("error", "Unknown error")) raise HTTPException(504, "Generation timed out after 20 minutes.") # ---------- Gallery / Jobs ---------- @app.get("/jobs") async def list_jobs(): return sorted(_jobs.values(), key=lambda j: j.get("created_at", ""), reverse=True) @app.get("/output/{filename}") async def serve_output(filename: str): fp = OUTPUT_DIR / filename if not fp.exists(): raise HTTPException(404, "File not found") return FileResponse(str(fp), media_type="video/mp4") @app.get("/health") def health(): ckpt_ok = Path(CKPT_DIR).exists() return { "status": "ok" if ckpt_ok else "building", "model_dir": CKPT_DIR, "model_exists": ckpt_ok, "wan22_dir": str(WAN22_DIR), "jobs_count": len(_jobs), } if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860)