handler.py

#!/usr/bin/env python3
"""
HTTP server for LingBot-World Base Cam NF4.

Runs a Flask server with an async job queue.

Endpoints:
    POST /generate        Submit a generation job (returns job ID immediately)
    GET  /status/<id>     Poll for job result
    GET /download/<id>    Download completed video (if available on server)
    GET  /health          Liveness check

Video output:
    If R2_ACCOUNT_ID / R2_ACCESS_KEY / R2_SECRET_KEY / R2_BUCKET env vars are set, the
    finished video is uploaded to Cloudflare R2 and a presigned URL is returned.
    Otherwise the video is saved to OUTPUT_DIR (default /app/outputs/) and the
    response includes a local file path.

Environment variables:
    MODEL_DIR       Path to model weights       (default: /app)
    OUTPUT_DIR      Local output directory       (default: /app/outputs)
    PORT            HTTP listen port             (default: 8080)
    R2_ACCOUNT_ID   Cloudflare account ID        (optional)
    R2_ACCESS_KEY   R2 API access key            (optional)
    R2_SECRET_KEY   R2 API secret key            (optional)
    R2_BUCKET       R2 bucket name               (default: lingbot-outputs)
    URL_EXPIRY      Presigned URL lifetime secs  (default: 86400)
"""

import base64
import io
import logging
import os
import shutil
import sys
import tempfile
import threading
import time
import uuid

import numpy as np
import torch
from flask import Flask, jsonify, request, send_file
from PIL import Image


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger(__name__)


MODEL_DIR = os.environ.get("MODEL_DIR", "/app")
OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/app/outputs")
os.makedirs(OUTPUT_DIR, exist_ok=True)

R2_CONFIGURED = all(
    os.environ.get(k) for k in ("R2_ACCOUNT_ID", "R2_ACCESS_KEY", "R2_SECRET_KEY")
)


sys.path.insert(0, MODEL_DIR)
from generate_prequant import WanI2V_PreQuant, save_video

logger.info("Initializing WanI2V_PreQuant pipeline ...")
_t0 = time.time()
pipeline = WanI2V_PreQuant(checkpoint_dir=MODEL_DIR, device_id=0, t5_cpu=True)
logger.info(f"Pipeline ready in {time.time() - _t0:.1f}s")


def _get_r2_client():
    import boto3
    from botocore.config import Config as BotoConfig

    return boto3.client(
        "s3",
        endpoint_url=f"https://{os.environ['R2_ACCOUNT_ID']}.r2.cloudflarestorage.com",
        aws_access_key_id=os.environ["R2_ACCESS_KEY"],
        aws_secret_access_key=os.environ["R2_SECRET_KEY"],
        config=BotoConfig(signature_version="s3v4"),
        region_name="auto",
    )


def _upload_to_r2(local_path: str, object_key: str) -> str:
    """Upload *local_path* to R2 and return a presigned download URL."""
    bucket = os.environ.get("R2_BUCKET", "lingbot-outputs")
    expiry = int(os.environ.get("URL_EXPIRY", 86400))
    client = _get_r2_client()
    client.upload_file(local_path, bucket, object_key, ExtraArgs={"ContentType": "video/mp4"})
    logger.info(f"Uploaded {object_key} -> R2 bucket '{bucket}'")
    return client.generate_presigned_url(
        "get_object", Params={"Bucket": bucket, "Key": object_key}, ExpiresIn=expiry,
    )


VALID_SIZES = {"480*832", "720*1280"}

def _validate(inp: dict) -> str | None:
    if not inp.get("image"):
        return "Missing required field: 'image' (base64-encoded JPEG/PNG)"
    if not inp.get("prompt"):
        return "Missing required field: 'prompt'"
    size = inp.get("size", "480*832")
    if size not in VALID_SIZES:
        return f"Invalid size '{size}'. Must be one of {VALID_SIZES}"
    frame_num = inp.get("frame_num", 81)
    if not isinstance(frame_num, int) or frame_num < 5:
        return f"frame_num must be an integer >= 5, got {frame_num}"
    return None


def generate_video(job_input: dict, job_id: str | None = None) -> dict:
    """
    Run the diffusion pipeline and deliver the result.

    Returns a dict with ``video_url`` (R2) or ``video_path`` (local) on
    success, or ``error`` on failure.
    """
    if job_id is None:
        job_id = str(uuid.uuid4())

    error = _validate(job_input)
    if error:
        return {"error": error}

    # Decode image
    try:
        image_bytes = base64.b64decode(job_input["image"])
        input_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    except Exception as exc:
        return {"error": f"Failed to decode image: {exc}"}

    # Parameters
    prompt = job_input["prompt"]
    frame_num = job_input.get("frame_num", 81)
    size = job_input.get("size", "480*832")
    seed = job_input.get("seed", -1)
    guide_scale = job_input.get("guide_scale", 5.0)
    sampling_steps = job_input.get("sampling_steps", 40)
    h, w = map(int, size.split("*"))

    # Optional camera poses
    action_path = None
    tmp_action_dir = None
    if job_input.get("action_poses") and job_input.get("action_intrinsics"):
        try:
            tmp_action_dir = tempfile.mkdtemp(prefix="lingbot_actions_")
            for name, key in [("poses.npy", "action_poses"),
                              ("intrinsics.npy", "action_intrinsics")]:
                with open(os.path.join(tmp_action_dir, name), "wb") as f:
                    f.write(base64.b64decode(job_input[key]))
            action_path = tmp_action_dir
        except Exception as exc:
            return {"error": f"Failed to decode camera actions: {exc}"}

    # Run pipeline
    logger.info(
        f"Job {job_id}: prompt='{prompt[:80]}...' frames={frame_num} "
        f"size={size} seed={seed} steps={sampling_steps}"
    )
    t0 = time.time()
    try:
        video = pipeline.generate(
            input_prompt=prompt,
            img=input_image,
            action_path=action_path,
            max_area=h * w,
            frame_num=frame_num,
            sampling_steps=sampling_steps,
            guide_scale=guide_scale,
            seed=seed,
        )
    except torch.cuda.OutOfMemoryError:
        return {"error": "CUDA out of memory. Try fewer frames or 480*832."}
    except Exception as exc:
        return {"error": f"Generation failed: {exc}"}
    finally:
        if tmp_action_dir:
            shutil.rmtree(tmp_action_dir, ignore_errors=True)

    gen_secs = time.time() - t0
    logger.info(f"Job {job_id}: done in {gen_secs:.1f}s")

    # Save video
    output_path = os.path.join(OUTPUT_DIR, f"{job_id}.mp4")
    try:
        save_video(video, output_path, fps=16)
    except Exception as exc:
        return {"error": f"Failed to save video: {exc}"}

    result = {
        "seed": seed,
        "duration_sec": round(gen_secs, 1),
        "frame_num": frame_num,
        "size": size,
    }

    # Upload to R2 (if configured) or keep local
    if R2_CONFIGURED:
        try:
            object_key = f"outputs/{job_id}.mp4"
            result["video_url"] = _upload_to_r2(output_path, object_key)
            os.remove(output_path)
        except Exception as exc:
            logger.warning(f"Job {job_id}: R2 upload failed ({exc}), keeping local file")
            result["video_path"] = output_path
    else:
        result["video_path"] = output_path

    return result


app = Flask(__name__)

_jobs: dict[str, dict] = {}
_jobs_lock = threading.Lock()


def _run_job(job_id: str, job_input: dict):
    """Background thread target."""
    try:
        result = generate_video(job_input, job_id)
    except Exception as exc:
        result = {"error": f"Unexpected failure: {exc}"}
    with _jobs_lock:
        _jobs[job_id]["status"] = "FAILED" if "error" in result else "COMPLETED"
        _jobs[job_id]["result"] = result


@app.route("/health", methods=["GET"])
def health():
    return jsonify({"status": "healthy", "gpu": torch.cuda.get_device_name(0)})


@app.route("/generate", methods=["POST"])
def generate():
    """Accept a job and return immediately with a job ID (HTTP 202)."""
    body = request.get_json(silent=True)
    if body is None:
        return jsonify({"error": "Request body must be JSON"}), 400
    job_input = body.get("input", body)

    error = _validate(job_input)
    if error:
        return jsonify({"error": error}), 400

    job_id = str(uuid.uuid4())
    with _jobs_lock:
        _jobs[job_id] = {"status": "IN_PROGRESS", "result": None, "submitted_at": time.time()}

    logger.info(f"Job {job_id}: accepted")
    threading.Thread(target=_run_job, args=(job_id, job_input), daemon=True).start()
    return jsonify({"id": job_id, "status": "IN_PROGRESS"}), 202


@app.route("/status/<job_id>", methods=["GET"])
def status(job_id):
    """Poll for job result."""
    with _jobs_lock:
        job = _jobs.get(job_id)
    if job is None:
        return jsonify({"error": f"Job {job_id} not found"}), 404
    resp = {"id": job_id, "status": job["status"]}
    if job["result"] is not None:
        resp["output"] = job["result"]
    return jsonify(resp)


@app.route("/download/<job_id>", methods=["GET"])
def download(job_id):
    """Download a completed video by job ID."""
    video_path = os.path.join(OUTPUT_DIR, f"{job_id}.mp4")
    if not os.path.isfile(video_path):
        return jsonify({"error": f"Video for job {job_id} not found (may have been uploaded to R2)"}), 404
    return send_file(video_path, mimetype="video/mp4", as_attachment=True, download_name=f"{job_id}.mp4")


if __name__ == "__main__":
    port = int(os.environ.get("PORT", 8080))
    logger.info(f"Starting server on port {port} ...")
    app.run(host="0.0.0.0", port=port)