Add persistent hot runtime service

PLAN.md

@@ -54,7 +54,7 @@ Prompt-only FL2VA is now at warm Comfy parity with the direct Sage2 baseline. Fe
 
 1. Validate and benchmark the existing `sage3` backend against the same cat prompt, seed, dimensions, and FP16 VAE runtime path used for Sage2 parity.
    - First cat benchmark result: Sage3 runs successfully but is slower than Sage2 in this direct path. Sampling was `123.675s` versus Sage2 `114.414s`; warm after text conditioning was `158.111s` versus Sage2 `149.304s`. Same-seed MP4 frame diff versus Sage2 was mean `46.563`, max `255`, so keep Sage3 experimental pending human visual review and stricter tensor gates.
-2. Build a persistent hot runtime service instead of measuring only process-per-run CLIs. A warm container must preload and retain Qwen, H3, video VAE, and audio VAE in GPU memory for the selected attention backend, then accept video jobs without model-load latency. Add explicit startup warmup, readiness reporting, backend selection, and timing fields that separate resident-model request latency from cold startup.
+2. Build a persistent hot runtime service instead of measuring only process-per-run CLIs. A warm container must preload and retain Qwen, H3, video VAE, and audio VAE in GPU memory, then accept video jobs without model-load latency. Add explicit startup warmup, readiness reporting, request-level attention selection (`sage2`, `sdpa`, `sage3` initially), and timing fields that separate resident-model request latency from cold startup.
 3. Add exact memory/lifetime optimizations next: `kj_head_sliced` and `kj_chunked_ffn`. These must preserve the validated direct outputs before being kept.
 4. Evaluate prior H3-tested attention candidates as standalone adapters: `sol_attn` and `kj_sage`.
 5. Evaluate approximate denoiser caches only after exact baselines are recorded: `easycache` and `h3_cache`.

README.md

@@ -49,3 +49,13 @@ Generation and latent-decode tools are quiet by default: they suppress ffmpeg ba
 - `--vae-tile-size 256`: set the direct video VAE spatial tile size. `tools/direct_t2v_preview.py` also accepts `H3_VAE_TILE_SIZE`.
 
 Standalone `tools/compare_*`, `tools/trace_*`, `tools/inspect_*`, and `tools/patch_comfy_*` scripts are debugging utilities and remain opt-in by being separate commands.
+
+## Hot Runtime Service
+
+`tools/serve_hot_runtime.py` keeps Qwen, H3, video VAE, and audio VAE resident in one process. Start the optional Spark service with:
+
+```bash
+docker compose -f compose.spark.yml up -d h3-hot-runtime
+```
+
+Use `GET /ready` to confirm resident model readiness. Use `POST /generate` with JSON fields like `prompt`, `output`, `width`, `height`, `frames`, `steps`, `seed`, and optional `attention`. Supported request-level attention values are `sage2`, `sdpa`, and `sage3`; switching attention does not reload model weights.

compose.spark.yml

@@ -12,3 +12,20 @@ services:
       - /home/daniel/aeon-spark-test/h3/h3-blackwell-runtime/artifacts:/artifacts:ro
       - /home/daniel/StoryStudioAssets/H3-output:/output
     command: ["sleep", "infinity"]
+  h3-hot-runtime:
+    build:
+      context: .
+      dockerfile: Dockerfile.spark
+    image: h3-blackwell-runtime:dev
+    gpus: all
+    volumes:
+      - /home/daniel/aeon-spark-test/h3/comfy-models/diffusion_models:/models:ro
+      - /home/daniel/aeon-spark-test/h3/comfy-models/text_encoders:/text-encoders:ro
+      - /home/daniel/aeon-spark-test/h3/comfy-models/vae:/vae:ro
+      - /home/daniel/aeon-spark-test/h3/h3-blackwell-runtime/artifacts:/artifacts:ro
+      - /home/daniel/StoryStudioAssets/H3-output:/output
+    ports:
+      - "8001:8000"
+    environment:
+      H3_DISABLE_MMAP: "1"
+    command: ["python", "/opt/h3-blackwell-runtime/tools/serve_hot_runtime.py", "--host", "0.0.0.0", "--port", "8000", "--attention", "sage2", "--warmup"]

src/h3_blackwell_runtime/runtime.py

@@ -0,0 +1,235 @@
+"""Resident prompt-only H3 runtime used by the hot service and tests."""
+
+from __future__ import annotations
+
+import subprocess
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+
+from .audio_vae_decoder import MiniMaxH3AudioVAE
+from .attention import AVAILABLE_BACKENDS
+from .checkpoint import H3Checkpoint
+from .denoiser import H3PackedDenoiser
+from .packing import H3PromptPacker
+from .qwen3vl_text import Qwen3VLPromptConditioner
+from .sampler import sample_video_res_multistep
+from .t2v import random_av_latents
+from .token_refiner import H3TokenRefiner
+from .vae_decoder import MiniMaxH3VideoVAE, dtype_from_name
+
+
+@dataclass(frozen=True)
+class RuntimeConfig:
+    model_path: str = "/models/minimax_h3_fl2va_pruned_nvfp4.safetensors"
+    text_encoder_path: str = "/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors"
+    tokenizer_path: str = "/opt/h3-blackwell-runtime/src/h3_blackwell_runtime/qwen25_tokenizer"
+    video_vae_path: str = "/vae/minimax_h3_video_vae_fp16.safetensors"
+    audio_vae_path: str = "/vae/minimax_h3_audio_vae_fp32.safetensors"
+    attention: str = "sage2"
+    vae_dtype: str = "float16"
+    vae_tile_size: int = 256
+    device: str = "cuda"
+
+
+def _sync() -> None:
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+
+
+def _ffmpeg_command(loglevel: str, *parts: str) -> list[str]:
+    return ["ffmpeg", "-hide_banner", "-loglevel", loglevel, *parts]
+
+
+class H3HotRuntime:
+    """Keep all prompt-only H3 models resident for repeated requests."""
+
+    def __init__(self, config: RuntimeConfig):
+        self.config = config
+        self.attention = config.attention
+        self.loaded_at = time.time()
+        self.load_stages: list[dict] = []
+
+        self.checkpoint = H3Checkpoint(config.model_path, device=config.device)
+        self.conditioner = self._timed_load(
+            "qwen_loaded",
+            lambda: Qwen3VLPromptConditioner(config.text_encoder_path, config.tokenizer_path),
+        )
+        self.model = self._timed_load(
+            "h3_loaded",
+            lambda: H3PackedDenoiser.from_checkpoint(self.checkpoint, attention_backend=config.attention).eval(),
+        )
+        self.refiner = self._timed_load(
+            "token_refiner_loaded",
+            lambda: H3TokenRefiner(self.checkpoint, attention_backend=config.attention).eval(),
+        )
+        self.packer = H3PromptPacker(self.checkpoint)
+        self.video_vae = self._timed_load(
+            "video_vae_loaded",
+            self._load_video_vae,
+        )
+        self.audio_vae = self._timed_load(
+            "audio_vae_loaded",
+            lambda: MiniMaxH3AudioVAE.from_safetensors(config.audio_vae_path, device=config.device).eval(),
+        )
+
+    def _timed_load(self, stage: str, fn):
+        _sync()
+        start = time.perf_counter()
+        value = fn()
+        _sync()
+        self.load_stages.append({"stage": stage, "seconds": time.perf_counter() - start})
+        return value
+
+    def _load_video_vae(self) -> MiniMaxH3VideoVAE:
+        vae = MiniMaxH3VideoVAE.from_safetensors(
+            self.config.video_vae_path,
+            device=self.config.device,
+            dtype=dtype_from_name(self.config.vae_dtype),
+        ).eval()
+        vae.tile_size = self.config.vae_tile_size
+        return vae
+
+    def status(self) -> dict:
+        return {
+            "ready": True,
+            "initial_attention": self.config.attention,
+            "current_attention": self.attention,
+            "vae_dtype": self.config.vae_dtype,
+            "vae_tile_size": self.config.vae_tile_size,
+            "loaded_at": self.loaded_at,
+            "load_stages": self.load_stages,
+        }
+
+    @torch.inference_mode()
+    def set_attention(self, attention: str) -> None:
+        if attention not in AVAILABLE_BACKENDS:
+            raise ValueError(f"Unsupported attention backend: {attention}")
+        if attention == self.attention:
+            return
+        for module in self.model.modules():
+            if hasattr(module, "backend"):
+                module.backend = attention
+        for block in self.refiner.blocks:
+            block.attention_backend = attention
+        self.attention = attention
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        *,
+        prompt: str,
+        output: str | Path,
+        width: int,
+        height: int,
+        frames: int,
+        steps: int,
+        seed: int,
+        attention: str | None = None,
+        mux_audio: bool = True,
+        ffmpeg_loglevel: str = "error",
+        save_latent: str | Path | None = None,
+    ) -> dict:
+        stages: list[dict] = []
+
+        def timed(stage: str, fn):
+            _sync()
+            start = time.perf_counter()
+            value = fn()
+            _sync()
+            stages.append({"stage": stage, "seconds": time.perf_counter() - start})
+            return value
+
+        output = Path(output)
+        output.parent.mkdir(parents=True, exist_ok=True)
+        if attention is not None:
+            self.set_attention(attention)
+
+        video, audio, aligned_frames = timed(
+            "latents_initialized",
+            lambda: random_av_latents(width, height, frames, seed, device=self.config.device),
+        )
+        text = timed("text_conditioned", lambda: self.refiner(self.conditioner(prompt)))
+        sampled = timed(
+            "sampled",
+            lambda: sample_video_res_multistep(
+                self.model,
+                self.packer,
+                text,
+                video,
+                audio,
+                steps=steps,
+                return_audio=mux_audio,
+            ),
+        )
+        if mux_audio:
+            latent, audio_latent = sampled
+        else:
+            latent, audio_latent = sampled, None
+        if save_latent is not None:
+            latent_path = Path(save_latent)
+            latent_path.parent.mkdir(parents=True, exist_ok=True)
+            state = {"latent": latent.detach().cpu(), "frames": aligned_frames, "width": width, "height": height, "prompt": prompt, "seed": seed}
+            if audio_latent is not None:
+                state["audio_latent"] = audio_latent.detach().cpu()
+            torch.save(state, latent_path)
+        pixels = timed("vae_decoded", lambda: self.video_vae.decode(latent.to(next(self.video_vae.parameters()).dtype))[:, :, :aligned_frames])
+        pixels = timed("pixels_cpu", lambda: ((pixels[0].permute(1, 2, 3, 0).clamp(-1, 1) + 1) * 127.5).to(torch.uint8).cpu())
+
+        raw = output.with_suffix(".rgb")
+        video_output = output.with_name(output.stem + ".video.mp4") if mux_audio else output
+        timed("raw_write", lambda: pixels.numpy().tofile(raw))
+        timed(
+            "video_encode",
+            lambda: subprocess.run(
+                _ffmpeg_command(
+                    ffmpeg_loglevel,
+                    "-y",
+                    "-f",
+                    "rawvideo",
+                    "-pixel_format",
+                    "rgb24",
+                    "-video_size",
+                    f"{pixels.shape[2]}x{pixels.shape[1]}",
+                    "-framerate",
+                    "24",
+                    "-i",
+                    str(raw),
+                    "-an",
+                    "-c:v",
+                    "libx264",
+                    "-pix_fmt",
+                    "yuv420p",
+                    str(video_output),
+                ),
+                check=True,
+            ),
+        )
+        raw.unlink()
+
+        audio_output = None
+        if mux_audio:
+            audio_output = output.with_suffix(".wav")
+            waveform = timed("audio_decoded", lambda: self.audio_vae.decode(audio_latent.to(next(self.audio_vae.parameters()).dtype)).clamp(-1, 1).cpu()[0])
+            audio_raw = audio_output.with_suffix(".f32le")
+            timed("audio_raw_write", lambda: waveform.transpose(0, 1).contiguous().numpy().tofile(audio_raw))
+            timed("audio_encode", lambda: subprocess.run(_ffmpeg_command(ffmpeg_loglevel, "-y", "-f", "f32le", "-ar", "32000", "-ac", "2", "-i", str(audio_raw), str(audio_output)), check=True))
+            audio_raw.unlink()
+            timed("mux", lambda: subprocess.run(_ffmpeg_command(ffmpeg_loglevel, "-y", "-i", str(video_output), "-i", str(audio_output), "-c:v", "copy", "-c:a", "aac", "-shortest", str(output)), check=True))
+            video_output.unlink()
+
+        return {
+            "output": str(output),
+            "audio_output": str(audio_output) if audio_output is not None else None,
+            "frames": aligned_frames,
+            "width": width,
+            "height": height,
+            "seed": seed,
+            "attention": self.attention,
+            "vae_dtype": self.config.vae_dtype,
+            "vae_tile_size": self.config.vae_tile_size,
+            "stages": stages,
+            "request_seconds": sum(stage["seconds"] for stage in stages),
+        }

tools/serve_hot_runtime.py

@@ -0,0 +1,115 @@
+"""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.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=("sage2", "sdpa", "sage3"), default="sage2", 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("--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))
+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": ["sage2", "sdpa", "sage3"],
+        "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 {"sage2", "sdpa", "sage3"}:
+                write_json(self, 400, {"error": "unsupported attention", "attention": attention, "available": ["sage2", "sdpa", "sage3"]})
+                return
+            mux_audio = bool(payload.get("mux_audio", True))
+            ffmpeg_loglevel = payload.get("ffmpeg_loglevel", "error")
+            save_latent = payload.get("save_latent")
+            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,
+                )
+            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()