app.py

"""
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("""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8"/>
<meta name="viewport" content="width=device-width,initial-scale=1"/>
<title>Wan2.2 TI2V 5B — Video Generator</title>
<style>
:root {
  --bg: #0a0a14; --surface: #13132a; --surface2: #1c1c3a;
  --border: #2a2a4a; --text: #e2e2f0; --text2: #8888aa;
  --accent: #7c4dff; --accent2: #b388ff; --danger: #ff5252;
  --radius: 14px;
}
*{box-sizing:border-box;margin:0;padding:0}
body{font-family:-apple-system,'Segoe UI',Roboto,sans-serif;background:var(--bg);color:var(--text);min-height:100vh}
.app{max-width:1100px;margin:0 auto;padding:1.5rem}
header{text-align:center;padding:1.5rem 0}
header h1{font-size:2rem;background:linear-gradient(135deg,var(--accent),var(--accent2));-webkit-background-clip:text;-webkit-text-fill-color:transparent}
header .tagline{color:var(--text2);margin-top:.4rem;font-size:.95rem}
.layout{display:grid;grid-template-columns:1fr 1fr;gap:1.5rem;align-items:start}
@media(max-width:800px){.layout{grid-template-columns:1fr}}
.card{background:var(--surface);border:1px solid var(--border);border-radius:var(--radius);padding:1.5rem}
.card h2{font-size:1.15rem;margin-bottom:1rem}
.field{margin-bottom:1rem}
.field label{display:block;font-weight:600;margin-bottom:.35rem;font-size:.88rem;color:var(--text2)}
.field textarea,.field input,.field select{width:100%;padding:.7rem .85rem;border-radius:10px;border:1px solid var(--border);background:var(--surface2);color:var(--text);font-size:.93rem;font-family:inherit}
.field textarea{resize:vertical;min-height:90px}
.field textarea:focus,.field input:focus,.field select:focus{outline:none;border-color:var(--accent);box-shadow:0 0 0 3px rgba(124,77,255,.15)}
.row{display:grid;grid-template-columns:1fr 1fr;gap:1rem}
.btn{display:inline-flex;align-items:center;justify-content:center;gap:.5rem;padding:.75rem 1.5rem;border:none;border-radius:10px;font-weight:700;font-size:.95rem;cursor:pointer;transition:all .2s}
.btn-primary{background:linear-gradient(135deg,var(--accent),#5e3fcc);color:#fff;width:100%}
.btn-primary:hover{opacity:.92;transform:translateY(-1px)}
.btn-primary:disabled{opacity:.45;cursor:not-allowed;transform:none}
.file-upload{border:2px dashed var(--border);border-radius:var(--radius);padding:1.5rem;text-align:center;cursor:pointer;transition:border-color .2s;position:relative}
.file-upload:hover{border-color:var(--accent)}
.file-upload .icon{font-size:2rem;margin-bottom:.5rem}
.file-upload .hint{color:var(--text2);font-size:.85rem}
.file-upload input[type=file]{position:absolute;inset:0;opacity:0;cursor:pointer}
.file-upload.has-image{padding:.5rem}
.file-upload img{max-height:140px;border-radius:8px}
.preview-name{font-size:.85rem;color:var(--accent2);margin-top:.3rem}
.gallery{display:grid;grid-template-columns:repeat(auto-fill,minmax(200px,1fr));gap:.75rem;max-height:70vh;overflow-y:auto}
.job-card{background:var(--surface2);border:1px solid var(--border);border-radius:10px;overflow:hidden}
.job-card video{width:100%;display:block;border-radius:10px 10px 0 0;background:#000}
.job-card .meta{padding:.6rem .8rem;font-size:.8rem}
.job-card .prompt{white-space:nowrap;overflow:hidden;text-overflow:ellipsis;margin-bottom:.2rem}
.job-card .info{color:var(--text2);display:flex;justify-content:space-between}
.job-card .status{font-weight:600}
.job-card .status.running{color:#ffab40}
.job-card .status.done{color:#69f0ae}
.job-card .status.error{color:var(--danger)}
.job-card .status.queued{color:var(--text2)}
.empty-state{text-align:center;padding:2.5rem 1rem;color:var(--text2)}
.empty-state .icon{font-size:3rem;margin-bottom:.8rem}
.spinner{display:inline-block;width:1rem;height:1rem;border:2px solid var(--text2);border-top-color:var(--accent);border-radius:50%;animation:spin .7s linear infinite}
@keyframes spin{to{transform:rotate(360deg)}}
.progress-bar{height:4px;background:var(--border);border-radius:2px;margin-top:.5rem;overflow:hidden}
.progress-bar .fill{height:100%;background:linear-gradient(90deg,var(--accent),var(--accent2));border-radius:2px;transition:width .3s}
#status-msg{font-size:.85rem;color:var(--accent2);margin-top:.75rem;min-height:1.2em}
.tabs{display:flex;gap:.25rem;margin-bottom:1rem;background:var(--surface2);border-radius:10px;padding:3px}
.tab{padding:.5rem 1rem;border-radius:8px;cursor:pointer;font-size:.88rem;font-weight:600;color:var(--text2);transition:.15s}
.tab.active{background:var(--accent);color:#fff}
.tab:hover:not(.active){color:var(--text)}
footer{text-align:center;padding:2rem 1rem;color:var(--text2);font-size:.8rem}
footer a{color:var(--accent2)}
</style>
</head>
<body>
<div class="app">
<header>
  <h1>🎬 Wan2.2 TI2V 5B</h1>
  <p class="tagline">Text-to-Video &amp; Image-to-Video · 720p @ 24fps</p>
</header>

<div class="layout">
  <div class="card">
    <h2>✨ Create Video</h2>
    <form id="gen-form">
      <div class="field">
        <label for="prompt">Prompt</label>
        <textarea id="prompt" name="prompt" placeholder="Describe what you want to see…" required></textarea>
      </div>
      <div class="field">
        <label>Reference Image <span style="color:var(--text2);font-weight:400">(optional — enables image-to-video)</span></label>
        <div class="file-upload" id="file-area">
          <div id="file-placeholder">
            <div class="icon">🖼️</div>
            <p class="hint">Click or drag to upload</p>
          </div>
          <img id="preview-img" style="display:none" alt=""/>
          <div class="preview-name" id="preview-name" style="display:none"></div>
          <input type="file" id="image" name="image" accept="image/*"/>
        </div>
      </div>
      <div class="row">
        <div class="field">
          <label for="steps">Steps (20–50)</label>
          <input type="number" id="steps" value="30" min="20" max="50"/>
        </div>
        <div class="field">
          <label for="size">Resolution</label>
          <select id="size" name="size">
            <option value="1280*704">1280×704 (16:9)</option>
            <option value="704*1280">704×1280 (9:16)</option>
            <option value="960*960">960×960 (1:1)</option>
          </select>
        </div>
      </div>
      <div class="field">
        <label for="guidance">Guidance Scale</label>
        <input type="number" id="guidance" value="5" min="1" max="10" step="0.5"/>
      </div>
      <button type="submit" class="btn btn-primary" id="submit-btn">🎬 Generate Video</button>
      <div id="status-msg"></div>
      <div class="progress-bar" id="progress-bar" style="display:none"><div class="fill" style="width:0%"></div></div>
    </form>
  </div>

  <div class="card">
    <div class="tabs">
      <div class="tab active" data-tab="gallery">📽️ Gallery</div>
      <div class="tab" data-tab="api">📡 API</div>
    </div>
    <div id="tab-gallery">
      <div class="gallery" id="gallery"></div>
      <div class="empty-state" id="empty-state"><div class="icon">🎞️</div><p>Generated videos appear here</p></div>
    </div>
    <div id="tab-api" style="display:none">
      <h2 style="font-size:1rem;margin-bottom:.5rem">API Usage</h2>
      <pre style="background:var(--surface2);padding:1rem;border-radius:10px;overflow-x:auto;font-size:.8rem;line-height:1.5"><b># Text-to-Video</b>
curl -X POST <span id="api-url">…</span>/generate \\
  -F "prompt=A cat astronaut in space" \\
  -F "steps=30" -o video.mp4

<b># Image-to-Video</b>
curl -X POST <span id="api-url2">…</span>/generate \\
  -F "prompt=Animate this scene" \\
  -F "image=@cat.jpg" -o video.mp4</pre>
    </div>
  </div>
</div>

<footer>
  Powered by <a href="https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B" target="_blank">Wan-AI/Wan2.2-TI2V-5B</a>
  · Official inference · Apache 2.0
</footer>
</div>

<script>
const API = window.location.origin;
document.querySelectorAll('#api-url,#api-url2').forEach(e=>e.textContent=API);

const form = document.getElementById('gen-form');
const submitBtn = document.getElementById('submit-btn');
const statusMsg = document.getElementById('status-msg');
const progressBar = document.getElementById('progress-bar');
const progressFill = progressBar.querySelector('.fill');
const gallery = document.getElementById('gallery');
const emptyState = document.getElementById('empty-state');

const imageInput = document.getElementById('image');
const previewImg = document.getElementById('preview-img');
const previewName = document.getElementById('preview-name');
const filePlaceholder = document.getElementById('file-placeholder');
const fileArea = document.getElementById('file-area');

imageInput.addEventListener('change',()=>{
  const f=imageInput.files[0];
  if(!f){previewImg.style.display='none';previewName.style.display='none';filePlaceholder.style.display='';fileArea.classList.remove('has-image');return}
  fileArea.classList.add('has-image');filePlaceholder.style.display='none';
  previewImg.style.display='';previewName.style.display='';previewName.textContent=f.name;
  const r=new FileReader();r.onload=e=>{previewImg.src=e.target.result};r.readAsDataURL(f);
});

document.querySelectorAll('.tab').forEach(t=>t.addEventListener('click',()=>{
  document.querySelectorAll('.tab').forEach(x=>x.classList.remove('active'));
  t.classList.add('active');
  document.getElementById('tab-gallery').style.display=t.dataset.tab==='gallery'?'':'none';
  document.getElementById('tab-api').style.display=t.dataset.tab==='api'?'':'none';
}));

form.addEventListener('submit',async e=>{
  e.preventDefault();
  submitBtn.disabled=true;
  statusMsg.textContent='⏳ Generating… (~8 min for 81 frames on T4)';
  progressBar.style.display='';progressFill.style.width='15%';

  const fd=new FormData(form);
  try{
    const res=await fetch(API+'/generate',{method:'POST',body:fd});
    if(!res.ok){const t=await res.text();throw new Error(t||res.statusText)}
    progressFill.style.width='95%';
    statusMsg.textContent='✅ Done!';
    loadGallery();
  }catch(err){
    statusMsg.textContent='❌ '+err.message;
    progressBar.style.display='none';
  }finally{
    submitBtn.disabled=false;
    setTimeout(()=>{progressBar.style.display='none';statusMsg.textContent=''},5000);
  }
});

async function loadGallery(){
  try{
    const res=await fetch(API+'/jobs');
    const jobs=await res.json();
    if(!jobs||!jobs.length){gallery.innerHTML='';emptyState.style.display='';return}
    emptyState.style.display='none';
    gallery.innerHTML=jobs.reverse().map(j=>{
      const dur=j.duration?j.duration+'s':'';
      const cls=j.status;
      if(j.status==='done'&&j.output){
        return `<div class="job-card"><video src="${j.output}" controls preload="metadata"></video><div class="meta"><div class="prompt">${esc(j.prompt)}</div><div class="info"><span class="status ${cls}">✓ Done</span><span>${dur}</span></div></div></div>`;
      }else if(j.status==='running'){
        return `<div class="job-card" style="padding:2rem;text-align:center"><div class="spinner"></div><div class="meta"><div class="prompt">${esc(j.prompt)}</div><div class="info"><span class="status running">Generating…</span></div></div></div>`;
      }else if(j.status==='queued'){
        return `<div class="job-card" style="padding:2rem;text-align:center"><div style="font-size:2rem">⏳</div><div class="meta"><div class="prompt">${esc(j.prompt)}</div><div class="info"><span class="status queued">Queued</span></div></div></div>`;
      }else{
        return `<div class="job-card" style="padding:1.5rem;text-align:center"><div style="font-size:1.5rem">⚠️</div><div class="meta"><div class="prompt">${esc(j.prompt)}</div><div class="info"><span class="status error">Error</span><span>${esc(j.error||'')}</span></div></div></div>`;
      }
    }).join('');
  }catch(e){console.error(e)}
}
function esc(s){const d=document.createElement('div');d.textContent=s||'';return d.innerHTML}

setInterval(loadGallery,5000);
loadGallery();
</script>
</body>
</html>""")


# ---------- 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)