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, attention_backend_status
from h3_blackwell_runtime.runtime import H3HotRuntime, RuntimeConfig, TURBO_VARIANTS, normalize_upscale


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("--turbo-lora", action="append", default=[], metavar="NAME=PATH", help="Load a resident 4step or 8step Turbo adapter.")
parser.add_argument("--latent-upscaler", help="Load the optional H3 3D latent upscaler for request-level spatial upscaling.")
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()

turbo_loras = []
for value in args.turbo_lora:
    if "=" not in value:
        parser.error("--turbo-lora must use NAME=PATH")
    name, path = value.split("=", 1)
    turbo_loras.append((name, path))

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, turbo_loras=tuple(turbo_loras), latent_upscaler_path=args.latent_upscaler))
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),
        "attention_backend_status": attention_backend_status(),
        "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))
            turbo = payload.get("turbo")
            if turbo in {"", "none"}:
                turbo = None
            if turbo is not None and turbo not in runtime.turbo_loras:
                write_json(self, 400, {"error": "unsupported turbo", "turbo": turbo, "available": sorted(runtime.turbo_loras)})
                return
            steps = int(payload.get("steps", TURBO_VARIANTS[turbo]["steps"] if turbo else 12))
            if turbo is not None and steps != TURBO_VARIANTS[turbo]["steps"]:
                write_json(self, 400, {"error": "invalid turbo steps", "turbo": turbo, "required_steps": TURBO_VARIANTS[turbo]["steps"]})
                return
            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))
            keep_intermediates = bool(payload.get("keep_intermediates", False))
            try:
                upscale = normalize_upscale(payload.get("upscale"))
            except ValueError as exc:
                write_json(self, 400, {"error": "invalid upscale", "message": str(exc)})
                return
            if upscale is not None and runtime.latent_upscaler is None:
                write_json(self, 400, {"error": "H3 latent upscaler is not loaded"})
                return
            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")
            if turbo is not None and cache_mode not in {None, "", "disabled", "none"}:
                write_json(self, 400, {"error": "turbo does not support denoiser caching", "turbo": turbo})
                return
            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,
                    turbo=turbo,
                    upscale=upscale,
                    first_frame=first_frame,
                    last_frame=last_frame,
                    mux_audio=mux_audio,
                    keep_intermediates=keep_intermediates,
                    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()