"""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 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), "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)) 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()