h3-blackwell-runtime/tools/serve_hot_runtime.py

"""Serve a resident prompt-only H3 runtime over a small JSON HTTP API."""

from __future__ import annotations

import argparse
import json
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.parse import urlparse

import torch
from PIL import Image

from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, DEFAULT_ATTENTION_BACKEND
from h3_blackwell_runtime.runtime import H3HotRuntime, RuntimeConfig


def _load_image(value) -> torch.Tensor | None:
    """Accept a keyframe as on-disk path or base64 JPEG/PNG -> ``[1,3,H,W]`` float ``[0,1]``."""
    if value in (None, ""):
        return None
    if isinstance(value, (list, tuple)):
        value = value[0]
    if isinstance(value, dict):
        value = value.get("url") or value.get("path") or value.get("b64")
    if isinstance(value, str) and "\n" not in value and len(value) < 2048 and not value.startswith("data:"):
        path = Path(value)
        if path.exists():
            image = Image.open(path).convert("RGB")
            import numpy as np

            tensor = torch.from_numpy(np.array(image))[None].permute(0, 3, 1, 2).float() / 255.0
            return tensor
    data = value
    if isinstance(value, str) and value.startswith("data:"):
        data = value.split(",", 1)[1]
    if isinstance(data, str) and len(data) >= 1024:
        import base64
        import io

        raw = base64.b64decode(data)
        image = Image.open(io.BytesIO(raw)).convert("RGB")
        import numpy as np

        return torch.from_numpy(np.array(image))[None].permute(0, 3, 1, 2).float() / 255.0
    raise ValueError("first_frame/last_frame must be a path or a base64/data-URL image")


parser = argparse.ArgumentParser()
parser.add_argument("--host", default="0.0.0.0")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--attention", choices=AVAILABLE_BACKENDS, default=DEFAULT_ATTENTION_BACKEND, help="Initial attention backend. Requests can switch with the JSON attention field.")
parser.add_argument("--vae-dtype", choices=("float32", "float16", "bfloat16"), default="float16")
parser.add_argument("--vae-tile-size", type=int, default=256)
parser.add_argument("--mlp-chunks", type=int, default=1)
parser.add_argument("--mlp-chunk-threshold", type=int, default=4096)
parser.add_argument("--warmup", action="store_true", help="Run a tiny generation before accepting traffic.")
parser.add_argument("--warmup-output", type=Path, default=Path("/output/h3-blackwell-runtime/hot-runtime-warmup.mp4"))
args = parser.parse_args()

runtime = H3HotRuntime(RuntimeConfig(attention=args.attention, vae_dtype=args.vae_dtype, vae_tile_size=args.vae_tile_size, mlp_chunks=args.mlp_chunks, mlp_chunk_threshold=args.mlp_chunk_threshold))
runtime_lock = threading.Lock()
warmup_result = None
if args.warmup:
    warmup_result = runtime.generate(
        prompt="A small warmup cat blinks in soft light.",
        output=args.warmup_output,
        width=320,
        height=192,
        frames=22,
        steps=2,
        seed=440501,
        mux_audio=True,
    )


def service_status() -> dict:
    return {
        "ready": True,
        "attention_backends": list(AVAILABLE_BACKENDS),
        "runtime": runtime.status(),
    }


def write_json(handler: BaseHTTPRequestHandler, status: int, payload: dict) -> None:
    body = json.dumps(payload, indent=2).encode("utf-8")
    handler.send_response(status)
    handler.send_header("Content-Type", "application/json")
    handler.send_header("Content-Length", str(len(body)))
    handler.end_headers()
    handler.wfile.write(body)


class Handler(BaseHTTPRequestHandler):
    def log_message(self, format: str, *args) -> None:
        return

    def do_GET(self) -> None:
        path = urlparse(self.path).path
        if path in {"/health", "/ready"}:
            write_json(self, 200, {**service_status(), "warmup_result": warmup_result})
            return
        write_json(self, 404, {"error": "not found"})

    def do_POST(self) -> None:
        path = urlparse(self.path).path
        if path != "/generate":
            write_json(self, 404, {"error": "not found"})
            return
        try:
            length = int(self.headers.get("Content-Length", "0"))
            payload = json.loads(self.rfile.read(length).decode("utf-8")) if length else {}
            prompt = payload["prompt"]
            output = payload["output"]
            width = int(payload.get("width", 960))
            height = int(payload.get("height", 544))
            frames = int(payload.get("frames", 124))
            steps = int(payload.get("steps", 12))
            seed = int(payload.get("seed", 440407))
            attention = payload.get("attention")
            if attention is not None and attention not in AVAILABLE_BACKENDS:
                write_json(self, 400, {"error": "unsupported attention", "attention": attention, "available": list(AVAILABLE_BACKENDS)})
                return
            mux_audio = bool(payload.get("mux_audio", True))
            ffmpeg_loglevel = payload.get("ffmpeg_loglevel", "error")
            first_frame = _load_image(payload.get("first_frame"))
            last_frame = _load_image(payload.get("last_frame"))
            save_latent = payload.get("save_latent")
            cache_mode = payload.get("cache_mode")
            cache_threshold = float(payload.get("cache_threshold", 0.0))
            cache_start_percent = float(payload.get("cache_start_percent", 0.0))
            cache_end_percent = float(payload.get("cache_end_percent", 1.0))
            cache_subsample_factor = int(payload.get("cache_subsample_factor", 2))
            started = time.perf_counter()
            with runtime_lock:
                result = runtime.generate(
                    prompt=prompt,
                    output=output,
                    width=width,
                    height=height,
                    frames=frames,
                    steps=steps,
                    seed=seed,
                    attention=attention,
                    first_frame=first_frame,
                    last_frame=last_frame,
                    mux_audio=mux_audio,
                    ffmpeg_loglevel=ffmpeg_loglevel,
                    save_latent=save_latent,
                    cache_mode=cache_mode,
                    cache_threshold=cache_threshold,
                    cache_start_percent=cache_start_percent,
                    cache_end_percent=cache_end_percent,
                    cache_subsample_factor=cache_subsample_factor,
                )
            result["wall_seconds"] = time.perf_counter() - started
            write_json(self, 200, result)
        except Exception as exc:
            write_json(self, 500, {"error": type(exc).__name__, "message": str(exc)})


server = ThreadingHTTPServer((args.host, args.port), Handler)
print(json.dumps({"serving": True, "host": args.host, "port": args.port, **service_status(), "warmup_result": warmup_result}, indent=2), flush=True)
server.serve_forever()