128 lines
5.4 KiB
Python
128 lines
5.4 KiB
Python
"""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
|
|
|
|
from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, DEFAULT_ATTENTION_BACKEND
|
|
from h3_blackwell_runtime.runtime import H3HotRuntime, RuntimeConfig
|
|
|
|
|
|
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")
|
|
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,
|
|
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()
|