Initial direct H3 runtime

.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+*.py[cod]
+artifacts/
+*.pt
+*.mp4
+*.rgb
+.env

Dockerfile.spark

@@ -0,0 +1,20 @@
+# GB10/Grace Blackwell development image. It does not start inference by default.
+FROM ghcr.io/aeon-7/comfyui-aeon-spark:slim
+
+WORKDIR /opt/h3-blackwell-runtime
+COPY . .
+
+# Build the upstream GB10 Blackwell SageAttention3 package; the base image supplies CUDA 13 and Torch.
+RUN git clone --depth 1 https://github.com/thu-ml/SageAttention.git /tmp/SageAttention \
+    && cd /tmp/SageAttention/sageattention3_blackwell \
+    && sed -i 's/cc_major, cc_minor = torch.cuda.get_device_capability()/cc_major, cc_minor = (12, 1)/' setup.py \
+    && TORCH_CUDA_ARCH_LIST=12.1a NVCC_APPEND_FLAGS='-gencode=arch=compute_121a,code=sm_121a' python -m pip install --no-cache-dir --no-build-isolation . \
+    && rm -rf /tmp/SageAttention
+
+RUN python -m pip install --no-cache-dir --no-deps -e . \
+    && python -c "import comfy_kitchen, torch; from sageattn3 import sageattn3_blackwell; print(torch.__version__, torch.version.cuda)"
+
+ENV H3_MODEL_PATH=/models/minimax_h3_ref2va_pruned_nvfp4.safetensors
+ENV TORCH_COMPILE_DISABLE=0 TORCHDYNAMO_DISABLE=0
+ENTRYPOINT []
+CMD ["bash"]

PLAN.md

@@ -0,0 +1,57 @@
+# H3 Blackwell Runtime Plan
+
+## Goal
+
+Build a direct MiniMax H3 Ref2VA runtime for Blackwell and Grace Blackwell that consumes the current Comfy safetensors checkpoints while removing ComfyUI and Raylight from the denoising critical path.
+
+The runtime must support one GPU first, then correct 2/4/6/8 GPU execution. It must retain the current NVFP4 model artifacts and use SageAttention3 where quality validation permits.
+
+## Future LTX 2.5 Track
+
+Add a separate LTX 2.5 direct-runtime adapter after H3 single-GPU parity is stable. Target the gated Lightricks `ltx-2.5-22b-distilled-transformer-nvfp4.safetensors` artifact (18.7 GB, release commit `dd53cc2cd45bbeaa3563dfb575cba3f49cf44761`).
+
+- Keep LTX model loading, conditioning, scheduler, VAE, and validation isolated from H3; this is a second model family, not an H3 checkpoint variant.
+- Inspect the safetensors header and published architecture/configuration before sharing H3 modules or kernels.
+- Establish a LTX SDPA/Sage2 correctness baseline before evaluating Sage3, FlashAttention-4, Sol-Attn, cache methods, or distributed layouts.
+- Respect the LTX 2 Community License Agreement and gated-access requirements; do not automate downloads without authorized access.
+
+## Reference Baseline
+
+The first acceptance target is the clean one-GPU ComfyUI baseline in `../h3-lab/h3-raylight-usp2-results.json`:
+
+- RTX PRO 6000 Blackwell, 96 GB
+- Ref2VA, 960x544, 124 frames, 24 fps
+- 12 steps, `beta`, `res_multistep`, seed `440202`
+- SageAttention3, 1 GPU
+- ComfyUI execution time: `49.893s`
+
+The direct runner must first match the model contract and output quality. Beating this timing comes after correctness is established.
+
+## Architecture
+
+1. Checkpoint adapter: read Comfy safetensors metadata, preserve packed low-precision weights and scale tensors, and map them into a canonical H3 state dictionary.
+2. Conditioning service: execute and cache Qwen layer-50 embeddings, modality tags, and reference VAE latents once per request.
+3. H3 denoiser: implement the packed Ref2VA DiT, 3-axis RoPE, AdaLN, dual audio/video schedule, and RES multistep solver without node-graph orchestration.
+4. Kernel layer: retain the known-good NVFP4 linear path initially; add explicit SageAttention3 and CUDA-graph buckets after exact single-GPU output validation.
+5. Distributed layer: use ragged Ulysses all-to-all for Q/K/V head exchange, Sage3 on full packed tokens per local head shard, then inverse exchange. Add tensor parallelism only after sequence parallel correctness is proven.
+
+## Milestones
+
+1. Inspect the actual local `pruned_nvfp4` checkpoint header and classify every tensor/scale layout.
+2. Create a direct single-GPU denoiser step matching ComfyUI for a fixed captured payload.
+3. Implement full single-GPU Ref2VA and compare per-step tensors plus final AV output against ComfyUI.
+4. Apply SageAttention3 and CUDA graphs; benchmark against the 49.893s reference.
+   - Optional backends and execution strategies to evaluate behind the same per-step quality gate: FlashAttention-4 (the Blackwell successor to Hopper-only FlashAttention-3), EasyCache/H3-Cache, Sol-Attn, and KJ exact memory-lifetime patches.
+   - Keep backend selection explicit per run; retain only candidates that match the validated direct correctness path and improve the measured denoising bottleneck.
+   - Current correctness baseline: SageAttention2 (`sage2`), which exactly matches the captured ComfyUI `--use-sage-attention` output. SDPA is a fallback; SageAttention3 remains experimental and must pass the same quality gate.
+   - Sol-Attn and KJ Sage have prior H3 test evidence and are supported experimental candidates. Integrate each as an isolated standalone adapter, record the exact mode/version, and gate it against the Sage2 per-step reference before combining it with caches or other approximation strategies.
+5. Implement ragged Ulysses Sage3 with transport-identity and distributed-versus-single-Sage3 tests.
+6. Sweep Ulysses/tensor-parallel layouts on 2/4/6/8 GPUs in an NVLink/NVSwitch domain.
+
+## Non-Negotiable Validation
+
+- Never silently pad semantic H3 tokens for unmasked attention.
+- Compare distributed output against the identical single-GPU Sage3 path before comparing to SDPA.
+- Validate denoiser outputs at each scheduler step, not only encoded video.
+- Record attention, GEMM, communication, VAE, and end-to-end timings separately.
+- Treat SageAttention3 as an experimental quality-gated kernel for H3.

README.md

@@ -0,0 +1,25 @@
+# H3 Blackwell Runtime
+
+Direct MiniMax H3 Ref2VA runtime research project. ComfyUI is the checkpoint and correctness oracle, not the target runtime.
+
+## First Gate
+
+Inspect the mounted H3 NVFP4 safetensors headers before designing an importer:
+
+```powershell
+python .\tools\inspect_safetensors.py /runpod-volume/ComfyUI/models/diffusion_models/minimax_h3_ref2va_pruned_nvfp4.safetensors
+```
+
+Write the output to `artifacts/checkpoints/` on the mounted volume. The result must identify packed weights, scales, and tensor naming before any kernel conversion work begins.
+
+## Benchmark Contract
+
+`benchmarks/ref2va-960x544-124f.json` is the single-GPU performance contract. Record direct-runner results as JSON and compare them with:
+
+```powershell
+python .\tools\compare_benchmark.py --result direct-result.json
+```
+
+## DGX Spark
+
+`Dockerfile.spark` and `compose.spark.yml` prepare an ARM64 GB10 development image using the existing AEON CUDA 13/SageAttention3 base. The compose target opens a shell only; it does not start inference.

benchmarks/ref2va-960x544-124f.json

@@ -0,0 +1,22 @@
+{
+  "name": "ref2va-960x544-124f-12step",
+  "workflow": "../../h3-lab/ref2va-three-anchor-dragon-raylight-2gpu.json",
+  "single_gpu_overrides": {
+    "GPU": 1,
+    "GPU_SELECT": "0",
+    "ulysses_degree": 1,
+    "ring_degree": 1,
+    "cfg_degree": 1,
+    "dp_degree": 1
+  },
+  "model": "minimax_h3_ref2va_pruned_nvfp4.safetensors",
+  "resolution": [960, 544],
+  "frames": 124,
+  "fps": 24,
+  "steps": 12,
+  "scheduler": "beta",
+  "sampler": "res_multistep",
+  "seed": 440202,
+  "reference_comfy_sage3_seconds": 49.893,
+  "measurement": "ComfyUI prompt execution time after warm-up"
+}

compose.spark-comfy-lab.yml

@@ -0,0 +1,77 @@
+services:
+  comfyui:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    image: aeon-h3-comfy:gb10-h3-lab
+    container_name: aeon-h3-comfy-gb10-h3-lab
+    runtime: nvidia
+    ports:
+      - "8202:8188"
+    environment:
+      SKIP_MODEL_DOWNLOAD: "1"
+      SKIP_ABLITERATED: "1"
+      COMFYUI_FLAGS: >-
+        --listen 0.0.0.0
+        --port 8188
+        --use-sage-attention
+        --disable-pinned-memory
+        --reserve-vram 2.0
+        --preview-method auto
+        --enable-cors-header
+        --enable-manager
+        --enable-assets
+        --extra-model-paths-config /etc/comfyui/h3-ltx-model-paths.yaml
+    volumes:
+      - ./workspace:/workspace/ComfyUI
+      - /home/daniel/StoryStudioAssets/H3-output:/workspace/ComfyUI/output
+      - ./h3-ltx-model-paths.yaml:/etc/comfyui/h3-ltx-model-paths.yaml:ro
+      - /home/daniel/aeon-spark-test/h3/comfy-models:/h3-models:ro
+      - /home/daniel/aeon-spark-test/h3/models/MiniMax-H3-Turbo-FL2VA:/turbo:ro
+      - /home/daniel/comfy-data/models/vae:/ltx-video-vae:ro
+    shm_size: "32gb"
+    ipc: host
+    ulimits:
+      memlock: -1
+      stack: 67108864
+
+  comfyui031:
+    build:
+      context: .
+      dockerfile: Dockerfile.031
+    image: aeon-h3-comfy:gb10-h3-lab-031-capture
+    container_name: aeon-h3-comfy-gb10-h3-lab-031
+    runtime: nvidia
+    ports:
+      - "8203:8188"
+    environment:
+      SKIP_MODEL_DOWNLOAD: "1"
+      SKIP_ABLITERATED: "1"
+      H3_CAPTURE_DIR: /h3-capture
+      H3_SAMPLER_CAPTURE_DIR: /h3-sampler-capture
+      COMFYUI_FLAGS: >-
+        --listen 0.0.0.0
+        --port 8188
+        --use-sage-attention
+        --disable-pinned-memory
+        --reserve-vram 2.0
+        --preview-method auto
+        --enable-cors-header
+        --enable-manager
+        --enable-assets
+        --database-url sqlite:////workspace/ComfyUI/user/comfyui-v031.db
+        --extra-model-paths-config /etc/comfyui/h3-ltx-model-paths.yaml
+    volumes:
+      - ./workspace:/workspace/ComfyUI
+      - /home/daniel/StoryStudioAssets/H3-output:/workspace/ComfyUI/output
+      - ./h3-ltx-model-paths.yaml:/etc/comfyui/h3-ltx-model-paths.yaml:ro
+      - /home/daniel/aeon-spark-test/h3/comfy-models:/h3-models:ro
+      - /home/daniel/aeon-spark-test/h3/models/MiniMax-H3-Turbo-FL2VA:/turbo:ro
+      - /home/daniel/comfy-data/models/vae:/ltx-video-vae:ro
+      - /home/daniel/aeon-spark-test/h3/h3-blackwell-runtime/artifacts/capture:/h3-capture
+      - /home/daniel/aeon-spark-test/h3/h3-blackwell-runtime/artifacts/fl2va-sampler-reference:/h3-sampler-capture
+    shm_size: "32gb"
+    ipc: host
+    ulimits:
+      memlock: -1
+      stack: 67108864

compose.spark.yml

@@ -0,0 +1,14 @@
+services:
+  h3-blackwell-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
+    command: ["sleep", "infinity"]

pyproject.toml

@@ -0,0 +1,14 @@
+[project]
+name = "h3-blackwell-runtime"
+version = "0.1.0"
+description = "Direct MiniMax H3 Blackwell inference research runtime"
+requires-python = ">=3.12"
+dependencies = [
+  "comfy-kitchen==0.2.28",
+  "safetensors>=0.5.0",
+  "torch==2.9.1+cu130",
+  "transformers>=4.51,<5"
+]
+
+[tool.setuptools.packages.find]
+where = ["src"]

src/h3_blackwell_runtime/__init__.py

@@ -0,0 +1,6 @@
+"""Direct MiniMax H3 runtime components."""
+
+from .qwen3vl_text import Qwen3VL32BTextEncoder, Qwen3VLPromptConditioner
+from .vae_decoder import MiniMaxH3VideoVAE
+
+__all__ = ["MiniMaxH3VideoVAE", "Qwen3VL32BTextEncoder", "Qwen3VLPromptConditioner"]

src/h3_blackwell_runtime/adaln.py

@@ -0,0 +1,39 @@
+"""Curve-form AdaLN used by the pruned H3 NVFP4 checkpoint."""
+
+import torch
+import torch.nn.functional as functional
+from torch import nn
+
+from .checkpoint import H3Checkpoint
+
+
+class H3CurveAdaLN(nn.Module):
+    """Interpolate the H3 timestep curve and emit six modality-specific AdaLN tensors."""
+
+    def __init__(self, curve_table: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, *, hidden_size: int = 5376):
+        super().__init__()
+        if curve_table.shape != (1025, 8):
+            raise ValueError(f"Unexpected H3 AdaLN curve shape: {tuple(curve_table.shape)}")
+        if weight.shape != (18 * hidden_size, curve_table.shape[1]) or bias.shape != (18 * hidden_size,):
+            raise ValueError("Unexpected H3 AdaLN projection dimensions.")
+        self.hidden_size = hidden_size
+        self.register_buffer("curve_table", curve_table.to(torch.float32), persistent=False)
+        self.register_buffer("weight", weight.to(torch.float32), persistent=False)
+        self.register_buffer("bias", bias.to(torch.float32), persistent=False)
+
+    @classmethod
+    def from_checkpoint(cls, checkpoint: H3Checkpoint, prefix: str):
+        return cls(
+            checkpoint.tensor("adaln_t_table", dtype=torch.float32),
+            checkpoint.tensor(f"{prefix}.linear.weight", dtype=torch.bfloat16),
+            checkpoint.tensor(f"{prefix}.linear.bias", dtype=torch.bfloat16),
+        )
+
+    def forward(self, timesteps: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        """Return shift/scale/gate tensors ordered as MSA then MLP for all modalities."""
+        table = self.curve_table
+        position = timesteps.float().clamp(0, 1) * (table.shape[0] - 1)
+        lower = position.floor().long().clamp(max=table.shape[0] - 2)
+        embedding = torch.lerp(table[lower], table[lower + 1], (position - lower).unsqueeze(1))
+        values = functional.linear(embedding, self.weight, self.bias)
+        return values.view(values.shape[0] * 3, 6 * self.hidden_size).chunk(6, dim=-1)

src/h3_blackwell_runtime/attention.py

@@ -0,0 +1,110 @@
+"""Direct H3 self-attention using packed NVFP4 linears and SageAttention3."""
+
+import torch
+import torch.nn.functional as functional
+from torch import nn
+
+from .checkpoint import H3Checkpoint
+from .nvfp4 import Nvfp4Linear
+
+
+AVAILABLE_BACKENDS = ("sage2", "sdpa", "sage3")
+PLANNED_BACKENDS = ("flash4", "easycache", "h3_cache", "sol_attn", "kj_sage", "kj_chunked_ffn", "kj_head_sliced")
+
+
+def attention_backend_status() -> dict[str, str]:
+    """Report direct-runtime attention choices without importing ComfyUI nodes."""
+    status = {name: "available" for name in AVAILABLE_BACKENDS}
+    status.update({"flash4": "planned: exact Blackwell kernel adapter"})
+    status.update({"easycache": "planned: approximate denoiser cache"})
+    status.update({"h3_cache": "planned: approximate H3-specific cache"})
+    status.update({"sol_attn": "experimental: prior H3-tested sparse Triton attention; standalone adapter pending"})
+    status.update({"kj_sage": "experimental: prior H3-tested Sage patch; standalone adapter pending"})
+    status.update({"kj_chunked_ffn": "planned: exact memory-lifetime adapter"})
+    status.update({"kj_head_sliced": "planned: exact memory-lifetime adapter"})
+    return status
+
+
+def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+    return torch.nn.functional.rms_norm(x, (x.shape[-1],), weight.to(x.dtype), eps)
+
+
+def apply_split_half_rope(x: torch.Tensor, rotation: torch.Tensor) -> torch.Tensor:
+    """Apply H3's split-half rotary table to `[batch, sequence, heads, dim]`."""
+    rotated_width = rotation.shape[-3] * 2
+    half = rotated_width // 2
+    if rotated_width > x.shape[-1]:
+        raise ValueError("RoPE rotation width exceeds the attention head dimension.")
+
+    pair = torch.stack((x[..., :half], x[..., half:rotated_width]), dim=-1)
+    pair = torch.matmul(rotation.to(x.dtype), pair.unsqueeze(-1)).squeeze(-1)
+    return torch.cat((pair[..., 0], pair[..., 1], x[..., rotated_width:]), dim=-1)
+
+
+class H3SageAttention(nn.Module):
+    """One MiniMax H3 attention module, independent of ComfyUI and Raylight."""
+
+    def __init__(
+        self,
+        qkv_proj: Nvfp4Linear,
+        out_proj: Nvfp4Linear,
+        q_norm_weight: torch.Tensor,
+        k_norm_weight: torch.Tensor,
+        *,
+        heads: int = 56,
+        head_dim: int = 128,
+        eps: float = 1e-5,
+        backend: str = "sage2",
+    ):
+        super().__init__()
+        self.qkv_proj = qkv_proj
+        self.out_proj = out_proj
+        self.heads = heads
+        self.head_dim = head_dim
+        self.eps = eps
+        if backend in PLANNED_BACKENDS:
+            raise ValueError(f"H3 attention backend '{backend}' needs a standalone adapter and is not installed.")
+        if backend not in AVAILABLE_BACKENDS:
+            raise ValueError(f"Unsupported H3 attention backend: {backend}")
+        self.backend = backend
+        self.register_buffer("q_norm_weight", q_norm_weight, persistent=False)
+        self.register_buffer("k_norm_weight", k_norm_weight, persistent=False)
+
+    @classmethod
+    def from_checkpoint(cls, checkpoint: H3Checkpoint, prefix: str, *, output_dtype=torch.bfloat16, backend: str = "sage2"):
+        return cls(
+            checkpoint.nvfp4_linear(f"{prefix}.qkv_proj", output_dtype=output_dtype),
+            checkpoint.nvfp4_linear(f"{prefix}.out_proj", output_dtype=output_dtype),
+            checkpoint.tensor(f"{prefix}.q_norm.weight", dtype=output_dtype),
+            checkpoint.tensor(f"{prefix}.k_norm.weight", dtype=output_dtype),
+            backend=backend,
+        )
+
+    def forward(self, x: torch.Tensor, rope_rotation: torch.Tensor) -> torch.Tensor:
+        if x.ndim != 2:
+            raise ValueError("H3 attention expects `[sequence, hidden]` input.")
+        sequence = x.shape[0]
+        inner = self.heads * self.head_dim
+        q, k, v = self.qkv_proj(x).split(inner, dim=-1)
+        q = q.view(1, sequence, self.heads, self.head_dim)
+        k = k.view(1, sequence, self.heads, self.head_dim)
+        v = v.view(1, sequence, self.heads, self.head_dim)
+
+        from comfy_kitchen import rms_rope_split_half_
+
+        rms_rope_split_half_(q, k, rope_rotation, self.q_norm_weight, self.k_norm_weight, epsilon=self.eps, rot_dim=rope_rotation.shape[-3] * 2)
+        q = q.transpose(1, 2).contiguous()
+        k = k.transpose(1, 2).contiguous()
+        v = v.transpose(1, 2).contiguous()
+
+        if self.backend == "sage2":
+            from sageattention import sageattn
+
+            out = sageattn(q, k, v, is_causal=False, tensor_layout="HND", smooth_k=False)
+        elif self.backend == "sage3":
+            from sageattn3 import sageattn3_blackwell
+
+            out = sageattn3_blackwell(q, k, v, is_causal=False)
+        else:
+            out = functional.scaled_dot_product_attention(q, k, v, is_causal=False)
+        return self.out_proj(out.transpose(1, 2).reshape(sequence, inner).contiguous())

src/h3_blackwell_runtime/backbone.py

@@ -0,0 +1,41 @@
+"""Direct 50-block H3 denoiser backbone with standalone Sage3 attention."""
+
+import torch
+from torch import nn
+
+from .adaln import H3CurveAdaLN
+from .block import H3DiTBlock
+from .checkpoint import H3Checkpoint
+from .rope import h3_rope_rotation
+
+
+class H3DenoiserBackbone(nn.Module):
+    """Execute H3 transformer blocks over an already packed Ref2VA hidden sequence."""
+
+    def __init__(self, blocks: list[H3DiTBlock], adaln: list[H3CurveAdaLN], inv_freq: torch.Tensor):
+        super().__init__()
+        if len(blocks) != 50 or len(adaln) != 50:
+            raise ValueError("The released H3 denoiser has exactly 50 transformer blocks.")
+        self.blocks = nn.ModuleList(blocks)
+        self.adaln = nn.ModuleList(adaln)
+        self.register_buffer("inv_freq", inv_freq.to(torch.float32), persistent=False)
+
+    @classmethod
+    def from_checkpoint(cls, checkpoint: H3Checkpoint, *, output_dtype=torch.bfloat16, attention_backend: str = "sage2"):
+        return cls(
+            [H3DiTBlock.from_checkpoint(checkpoint, index, output_dtype=output_dtype, attention_backend=attention_backend) for index in range(50)],
+            [H3CurveAdaLN.from_checkpoint(checkpoint, f"blocks.{index}.adaln_proj") for index in range(50)],
+            checkpoint.tensor("rope.inv_freq", dtype=torch.float32),
+        )
+
+    def forward(
+        self,
+        hidden: torch.Tensor,
+        timesteps: torch.Tensor,
+        position_ids: torch.Tensor,
+        segments: list[tuple[int, int, int]],
+    ) -> torch.Tensor:
+        rotation = h3_rope_rotation(position_ids.to(hidden.device), self.inv_freq, hidden.dtype)
+        for block, adaln in zip(self.blocks, self.adaln, strict=True):
+            hidden = block(hidden, rotation, *adaln(timesteps), segments)
+        return hidden

src/h3_blackwell_runtime/block.py

@@ -0,0 +1,99 @@
+"""Direct MiniMax H3 DiT block over the standalone Sage3 attention unit."""
+
+import torch
+from torch import nn
+
+from .attention import H3SageAttention, rms_norm
+from .checkpoint import H3Checkpoint
+from .nvfp4 import Nvfp4Linear
+
+
+def modulate_segments(
+    x: torch.Tensor,
+    shift: torch.Tensor,
+    scale: torch.Tensor,
+    segments: list[tuple[int, int, int]],
+) -> torch.Tensor:
+    """Apply H3's per-modality/per-timestep AdaLN parameters to contiguous rows."""
+    output = torch.empty_like(x)
+    for start, stop, row in segments:
+        output[start:stop] = x[start:stop] * (1 + scale[row].to(x.dtype)) + shift[row].to(x.dtype)
+    return output
+
+
+def gate_segments(
+    residual: torch.Tensor,
+    update: torch.Tensor,
+    gate: torch.Tensor,
+    segments: list[tuple[int, int, int]],
+) -> torch.Tensor:
+    """Add a gated fresh sublayer output without altering the source residual."""
+    output = residual.clone()
+    for start, stop, row in segments:
+        output[start:stop].addcmul_(update[start:stop], gate[row].to(update.dtype))
+    return output
+
+
+class H3SwiGLU(nn.Module):
+    def __init__(self, fc1: Nvfp4Linear, fc2: Nvfp4Linear):
+        super().__init__()
+        self.fc1 = fc1
+        self.fc2 = fc2
+
+    @classmethod
+    def from_checkpoint(cls, checkpoint: H3Checkpoint, prefix: str, *, output_dtype=torch.bfloat16):
+        return cls(
+            checkpoint.nvfp4_linear(f"{prefix}.fc1", output_dtype=output_dtype),
+            checkpoint.nvfp4_linear(f"{prefix}.fc2", output_dtype=output_dtype),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate, up = self.fc1(x).chunk(2, dim=-1)
+        return self.fc2(torch.nn.functional.silu(gate).mul_(up))
+
+
+class H3DiTBlock(nn.Module):
+    """One H3 transformer block with externally supplied AdaLN tensors."""
+
+    def __init__(
+        self,
+        norm1_weight: torch.Tensor,
+        norm2_weight: torch.Tensor,
+        attention: H3SageAttention,
+        mlp: H3SwiGLU,
+        *,
+        norm_eps: float = 1e-5,
+    ):
+        super().__init__()
+        self.attention = attention
+        self.mlp = mlp
+        self.norm_eps = norm_eps
+        self.register_buffer("norm1_weight", norm1_weight, persistent=False)
+        self.register_buffer("norm2_weight", norm2_weight, persistent=False)
+
+    @classmethod
+    def from_checkpoint(cls, checkpoint: H3Checkpoint, index: int, *, output_dtype=torch.bfloat16, attention_backend: str = "sage2"):
+        prefix = f"blocks.{index}"
+        return cls(
+            checkpoint.tensor(f"{prefix}.norm1.weight", dtype=output_dtype),
+            checkpoint.tensor(f"{prefix}.norm2.weight", dtype=output_dtype),
+            H3SageAttention.from_checkpoint(checkpoint, f"{prefix}.attn", output_dtype=output_dtype, backend=attention_backend),
+            H3SwiGLU.from_checkpoint(checkpoint, f"{prefix}.mlp", output_dtype=output_dtype),
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope_rotation: torch.Tensor,
+        shift_msa: torch.Tensor,
+        scale_msa: torch.Tensor,
+        gate_msa: torch.Tensor,
+        shift_mlp: torch.Tensor,
+        scale_mlp: torch.Tensor,
+        gate_mlp: torch.Tensor,
+        segments: list[tuple[int, int, int]],
+    ) -> torch.Tensor:
+        h = modulate_segments(rms_norm(x, self.norm1_weight, self.norm_eps), shift_msa, scale_msa, segments)
+        x = gate_segments(x, self.attention(h, rope_rotation), gate_msa, segments)
+        h = modulate_segments(rms_norm(x, self.norm2_weight, self.norm_eps), shift_mlp, scale_mlp, segments)
+        return gate_segments(x, self.mlp(h), gate_mlp, segments)

src/h3_blackwell_runtime/checkpoint.py

@@ -0,0 +1,35 @@
+"""Lazy loading for the current Comfy-format H3 safetensors checkpoint."""
+
+from pathlib import Path
+
+import torch
+
+from .nvfp4 import Nvfp4Linear, load_nvfp4_linear
+
+
+class H3Checkpoint:
+    """Load individual tensors/modules without materializing the whole checkpoint."""
+
+    def __init__(self, path: str | Path, device: str | torch.device = "cuda"):
+        self.path = str(path)
+        self.device = str(device)
+
+    def tensor(self, name: str, *, dtype: torch.dtype | None = None) -> torch.Tensor:
+        from safetensors import safe_open
+
+        with safe_open(self.path, framework="pt", device=self.device) as checkpoint:
+            value = checkpoint.get_tensor(name)
+        return value.to(dtype=dtype) if dtype is not None else value
+
+    def nvfp4_linear(self, prefix: str, *, output_dtype=torch.bfloat16) -> Nvfp4Linear:
+        names = ("comfy_quant", "weight", "weight_scale", "weight_scale_2", "bias")
+        tensors = {}
+        from safetensors import safe_open
+
+        with safe_open(self.path, framework="pt", device=self.device) as checkpoint:
+            available = set(checkpoint.keys())
+            for suffix in names:
+                name = f"{prefix}.{suffix}"
+                if name in available:
+                    tensors[name] = checkpoint.get_tensor(name)
+        return load_nvfp4_linear(tensors, prefix, output_dtype=output_dtype)

src/h3_blackwell_runtime/conditioning.py

@@ -0,0 +1,20 @@
+"""Prompt-only conditioning primitives independent of ComfyUI's node API."""
+
+from pathlib import Path
+
+import torch
+
+
+class H3PromptTokenizer:
+    """Tokenize raw H3 prompt text without Qwen chat-template tokens."""
+
+    def __init__(self, tokenizer_dir: str | Path):
+        from transformers import AutoTokenizer
+
+        self.tokenizer = AutoTokenizer.from_pretrained(str(tokenizer_dir), local_files_only=True)
+
+    def __call__(self, prompt: str, *, device: torch.device | str = "cuda") -> torch.Tensor:
+        if not prompt:
+            prompt = " "
+        encoded = self.tokenizer(prompt, add_special_tokens=False, return_tensors="pt")
+        return encoded.input_ids.to(device)

src/h3_blackwell_runtime/denoiser.py

@@ -0,0 +1,36 @@
+"""Packed-input H3 transformer core, independent of ComfyUI node execution."""
+
+import torch
+from torch import nn
+
+from .backbone import H3DenoiserBackbone
+from .checkpoint import H3Checkpoint
+from .final import H3FinalLayer
+
+
+class H3PackedDenoiser(nn.Module):
+    """Run the H3 transformer once its Ref2VA payload has been packed into hidden rows."""
+
+    def __init__(self, backbone: H3DenoiserBackbone, final_layer: H3FinalLayer):
+        super().__init__()
+        self.backbone = backbone
+        self.final_layer = final_layer
+
+    @classmethod
+    def from_checkpoint(cls, checkpoint: H3Checkpoint, *, output_dtype=torch.bfloat16, attention_backend: str = "sage2"):
+        return cls(
+            H3DenoiserBackbone.from_checkpoint(checkpoint, output_dtype=output_dtype, attention_backend=attention_backend),
+            H3FinalLayer.from_checkpoint(checkpoint, output_dtype=output_dtype),
+        )
+
+    def forward(
+        self,
+        hidden: torch.Tensor,
+        timesteps: torch.Tensor,
+        position_ids: torch.Tensor,
+        segments: list[tuple[int, int, int]],
+        video_segment: tuple[int, int, int],
+        audio_segment: tuple[int, int, int],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        hidden = self.backbone(hidden, timesteps, position_ids, segments)
+        return self.final_layer(hidden, timesteps, video_segment, audio_segment)

src/h3_blackwell_runtime/final.py

@@ -0,0 +1,73 @@
+"""Final H3 curve-AdaLN and video/audio output heads."""
+
+import torch
+import torch.nn.functional as functional
+from torch import nn
+
+from .attention import rms_norm
+from .checkpoint import H3Checkpoint
+
+
+class H3FinalLayer(nn.Module):
+    def __init__(
+        self,
+        curve_table: torch.Tensor,
+        norm_weight: torch.Tensor,
+        adaln_weight: torch.Tensor,
+        adaln_bias: torch.Tensor,
+        video_weight: torch.Tensor,
+        video_bias: torch.Tensor,
+        audio_weight: torch.Tensor,
+        audio_bias: torch.Tensor,
+        *,
+        hidden_size: int = 5376,
+        eps: float = 1e-5,
+    ):
+        super().__init__()
+        if adaln_weight.shape != (2 * hidden_size, curve_table.shape[1]):
+            raise ValueError("Unexpected final H3 AdaLN projection dimensions.")
+        self.hidden_size = hidden_size
+        self.eps = eps
+        self.register_buffer("curve_table", curve_table.to(torch.float32), persistent=False)
+        self.register_buffer("norm_weight", norm_weight, persistent=False)
+        self.register_buffer("adaln_weight", adaln_weight.to(torch.float32), persistent=False)
+        self.register_buffer("adaln_bias", adaln_bias.to(torch.float32), persistent=False)
+        self.register_buffer("video_weight", video_weight.to(torch.float32), persistent=False)
+        self.register_buffer("video_bias", video_bias.to(torch.float32), persistent=False)
+        self.register_buffer("audio_weight", audio_weight.to(torch.float32), persistent=False)
+        self.register_buffer("audio_bias", audio_bias.to(torch.float32), persistent=False)
+
+    @classmethod
+    def from_checkpoint(cls, checkpoint: H3Checkpoint, *, output_dtype=torch.bfloat16):
+        return cls(
+            checkpoint.tensor("adaln_t_table", dtype=torch.float32),
+            checkpoint.tensor("final_layer.norm.weight", dtype=output_dtype),
+            checkpoint.tensor("final_layer.adaln_proj.linear.weight", dtype=torch.float32),
+            checkpoint.tensor("final_layer.adaln_proj.linear.bias", dtype=torch.float32),
+            checkpoint.tensor("final_layer.video_out.weight", dtype=torch.float32),
+            checkpoint.tensor("final_layer.video_out.bias", dtype=torch.float32),
+            checkpoint.tensor("final_layer.audio_out.weight", dtype=torch.float32),
+            checkpoint.tensor("final_layer.audio_out.bias", dtype=torch.float32),
+        )
+
+    def forward(
+        self,
+        hidden: torch.Tensor,
+        timesteps: torch.Tensor,
+        video_segment: tuple[int, int, int],
+        audio_segment: tuple[int, int, int],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        position = timesteps.float().clamp(0, 1) * (self.curve_table.shape[0] - 1)
+        lower = position.floor().long().clamp(max=self.curve_table.shape[0] - 2)
+        embedding = torch.lerp(self.curve_table[lower], self.curve_table[lower + 1], (position - lower).unsqueeze(1))
+        shift, scale = functional.linear(embedding, self.adaln_weight, self.adaln_bias).chunk(2, dim=-1)
+
+        video_start, video_stop, video_row = video_segment
+        audio_start, audio_stop, audio_row = audio_segment
+        normalized = rms_norm(hidden, self.norm_weight, self.eps)
+        video_hidden = normalized[video_start:video_stop] * (1 + scale[video_row].to(hidden.dtype)) + shift[video_row].to(hidden.dtype)
+        audio_hidden = normalized[audio_start:audio_stop] * (1 + scale[audio_row].to(hidden.dtype)) + shift[audio_row].to(hidden.dtype)
+        return (
+            functional.linear(video_hidden.float(), self.video_weight, self.video_bias),
+            functional.linear(audio_hidden.float(), self.audio_weight, self.audio_bias),
+        )

src/h3_blackwell_runtime/nvfp4.py

@@ -0,0 +1,88 @@
+"""Standalone Blackwell NVFP4 linear adapter for Comfy-format checkpoints."""
+
+import json
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as functional
+from torch import nn
+
+
+@dataclass(frozen=True)
+class Nvfp4LinearTensors:
+    weight: torch.Tensor
+    weight_scale: torch.Tensor
+    weight_scale_2: torch.Tensor
+    bias: torch.Tensor | None
+    in_features: int
+    out_features: int
+
+
+def parse_quant_sidecar(sidecar: torch.Tensor) -> dict:
+    """Validate the small JSON descriptor stored beside each packed weight."""
+    metadata = json.loads(bytes(sidecar.cpu().tolist()))
+    if metadata.get("format") != "nvfp4":
+        raise ValueError(f"Unsupported quantization metadata: {metadata!r}")
+    return metadata
+
+
+class Nvfp4Linear(nn.Module):
+    """Execute a packed Comfy NVFP4 linear with Comfy Kitchen's CUDA 13 kernel."""
+
+    def __init__(self, tensors: Nvfp4LinearTensors, output_dtype=torch.bfloat16):
+        super().__init__()
+        if tensors.weight.dtype != torch.uint8 or tensors.weight.ndim != 2:
+            raise ValueError("NVFP4 weights must be a rank-2 packed uint8 tensor.")
+        if tensors.weight.shape != (tensors.out_features, tensors.in_features // 2):
+            raise ValueError("Packed NVFP4 dimensions do not match logical linear dimensions.")
+        if tensors.in_features % 32:
+            raise ValueError("Blackwell NVFP4 GEMM requires input width divisible by 32.")
+
+        self.in_features = tensors.in_features
+        self.out_features = tensors.out_features
+        self.output_dtype = output_dtype
+        self.register_buffer("weight", tensors.weight.contiguous(), persistent=False)
+        self.register_buffer("weight_scale", tensors.weight_scale.view(torch.float8_e4m3fn).contiguous(), persistent=False)
+        self.register_buffer("weight_scale_2", tensors.weight_scale_2.to(torch.float32).contiguous(), persistent=False)
+        self.register_buffer("bias", tensors.bias.contiguous() if tensors.bias is not None else None, persistent=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.shape[-1] != self.in_features:
+            raise ValueError(f"Expected feature width {self.in_features}, received {x.shape[-1]}.")
+        if x.dtype not in (torch.float16, torch.bfloat16):
+            raise ValueError("NVFP4 linear accepts FP16 or BF16 activations.")
+
+        from comfy_kitchen.tensor import QuantizedTensor, TensorCoreNVFP4Layout
+
+        original_shape = x.shape[:-1]
+        flat_x = x.reshape(-1, self.in_features).contiguous()
+        packed_x = QuantizedTensor.from_float(flat_x, "TensorCoreNVFP4Layout")
+        packed_weight = QuantizedTensor(
+            self.weight,
+            "TensorCoreNVFP4Layout",
+            TensorCoreNVFP4Layout.Params(
+                scale=self.weight_scale_2,
+                block_scale=self.weight_scale,
+                orig_dtype=self.output_dtype,
+                orig_shape=(self.out_features, self.in_features),
+            ),
+        )
+        output = functional.linear(packed_x, packed_weight, self.bias)
+        return output[:flat_x.shape[0], :self.out_features].reshape(*original_shape, self.out_features)
+
+
+def load_nvfp4_linear(tensors: dict[str, torch.Tensor], prefix: str, *, output_dtype=torch.bfloat16) -> Nvfp4Linear:
+    """Load one Comfy-format NVFP4 linear from a safetensors tensor mapping."""
+    sidecar_key = f"{prefix}.comfy_quant"
+    parse_quant_sidecar(tensors[sidecar_key])
+    weight = tensors[f"{prefix}.weight"]
+    in_features = weight.shape[1] * 2
+    packed = Nvfp4LinearTensors(
+        weight=weight,
+        weight_scale=tensors[f"{prefix}.weight_scale"],
+        weight_scale_2=tensors[f"{prefix}.weight_scale_2"],
+        bias=tensors.get(f"{prefix}.bias"),
+        in_features=in_features,
+        out_features=weight.shape[0],
+    )
+    return Nvfp4Linear(packed, output_dtype=output_dtype)

src/h3_blackwell_runtime/packing.py

@@ -0,0 +1,92 @@
+"""Direct prompt-only H3 packed-token construction."""
+
+import math
+
+import torch
+import torch.nn.functional as functional
+
+
+FRAME_PER_TOKEN = (1, 4, 4, 4, 4)
+FRAME_RESCALE = 5.0 / 3.0
+
+
+def patchify_video(latent: torch.Tensor) -> torch.Tensor:
+    batch, channels, frames, height, width = latent.shape
+    if batch != 1:
+        raise ValueError("H3 supports batch size one.")
+    if height % 2 or width % 2:
+        raise ValueError("H3 video latent dimensions must be divisible by two.")
+    return latent.reshape(batch, channels, frames, 1, height // 2, 2, width // 2, 2).permute(0, 2, 4, 6, 1, 3, 5, 7).reshape(-1, channels * 4)
+
+
+def pack_audio(latent: torch.Tensor) -> torch.Tensor:
+    return latent[0].permute(1, 2, 0).reshape(-1, latent.shape[1])
+
+
+def unpatchify_video(rows: torch.Tensor, frames: int, latent_height: int, latent_width: int) -> torch.Tensor:
+    height, width = latent_height // 2, latent_width // 2
+    x = rows.reshape(1, frames, height, width, 24, 1, 2, 2).permute(0, 4, 1, 5, 2, 6, 3, 7)
+    return x.reshape(1, 24, frames, latent_height, latent_width)
+
+
+def _axis(dim: int, area: float) -> torch.Tensor:
+    ratio, count = dim / area, dim // 2
+    return (torch.arange(count, dtype=torch.float64) * ratio / count + (1 - ratio) / 2) * 32
+
+
+def _video_positions(frames: int, height: int, width: int, offset: float) -> torch.Tensor:
+    area = math.sqrt(height * width)
+    ys, xs = torch.meshgrid(_axis(height, area), _axis(width, area), indexing="ij")
+    spatial = torch.stack((ys.flatten(), xs.flatten()), dim=-1)
+    spans = torch.tensor([FRAME_RESCALE * FRAME_PER_TOKEN[index % 5] for index in range(frames)], dtype=torch.float64)
+    times = offset + torch.cat((torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)))
+    result = torch.empty(frames, spatial.shape[0], 3, dtype=torch.float64)
+    result[:, :, 0], result[:, :, 1:] = times[:, None], spatial[None]
+    return result.reshape(-1, 3)
+
+
+def _audio_positions(steps: int, offset: float, width: int, height: int) -> torch.Tensor:
+    area = math.sqrt(height * width)
+    x_axis = _axis(width, area)
+    result = torch.zeros(steps * 2, 3, dtype=torch.float64)
+    result[:, 0] = (offset + torch.arange(steps, dtype=torch.float64)).repeat(2)
+    result[:steps, 2], result[steps:, 2] = x_axis[0], x_axis[-1]
+    return result
+
+
+class H3PromptPacker:
+    """Build `[text | audio | video]` tokens for prompt-only H3 T2V."""
+
+    def __init__(self, checkpoint):
+        self.video_weight = checkpoint.tensor("video_patch_proj.weight", dtype=torch.float32)
+        self.video_bias = checkpoint.tensor("video_patch_proj.bias", dtype=torch.float32)
+        self.audio_weight = checkpoint.tensor("audio_patch_proj.weight", dtype=torch.float32)
+        self.audio_bias = checkpoint.tensor("audio_patch_proj.bias", dtype=torch.float32)
+        self.text_weight = checkpoint.tensor("condition_proj.weight", dtype=torch.bfloat16)
+        self.text_bias = checkpoint.tensor("condition_proj.bias", dtype=torch.bfloat16)
+
+    def __call__(self, text: torch.Tensor, video: torch.Tensor, audio: torch.Tensor, sigma: float) -> tuple[torch.Tensor, torch.Tensor, list[tuple[int, int, int]], tuple[int, int, int], tuple[int, int, int]]:
+        if text.shape[-1] == 5120:
+            text_rows = functional.linear(text[0].to(self.text_weight.dtype), self.text_weight, self.text_bias).to(torch.bfloat16)
+        elif text.shape[-1] == 5376:
+            text_rows = text[0].to(torch.bfloat16)
+        else:
+            raise ValueError("H3 text states must be Qwen 5120-wide or refined 5376-wide.")
+        video_rows = functional.linear(patchify_video(video).float(), self.video_weight, self.video_bias).to(torch.bfloat16)
+        audio_rows = functional.linear(pack_audio(audio).float(), self.audio_weight, self.audio_bias).to(torch.bfloat16)
+        text_length, audio_length = text_rows.shape[0], audio_rows.shape[0]
+        hidden = torch.cat((text_rows, audio_rows, video_rows))
+        video_sigma = torch.tensor(float(sigma), device=hidden.device).clamp(min=1e-6)
+        base = video_sigma / (12.0 + video_sigma * (1.0 - 12.0))
+        audio_sigma = 3.0 * base / (1.0 + (3.0 - 1.0) * base)
+        video_time, audio_time = 1 - float(video_sigma), 1 - float(audio_sigma)
+        unique_times = sorted({video_time, audio_time})
+        row = {value: index for index, value in enumerate(unique_times)}
+        video_row, audio_row = row[video_time] * 3, row[audio_time] * 3
+        times = torch.tensor(unique_times, device=hidden.device, dtype=torch.float32)
+        positions = torch.cat((torch.stack((torch.arange(text_length, dtype=torch.float64), torch.zeros(text_length), torch.zeros(text_length)), dim=-1), _audio_positions(audio.shape[-1], float(text_length), video.shape[-1], video.shape[-2]), _video_positions(video.shape[2], video.shape[-2], video.shape[-1], float(text_length) + audio.shape[-1])))
+        block_video_segment = (text_length + audio_length, hidden.shape[0], video_row)
+        block_audio_segment = (text_length, text_length + audio_length, audio_row + 2)
+        final_video_segment = (text_length + audio_length, hidden.shape[0], row[video_time])
+        final_audio_segment = (text_length, text_length + audio_length, row[audio_time])
+        return hidden, times, [(0, text_length, video_row + 1), block_audio_segment, block_video_segment], positions, final_video_segment, final_audio_segment

src/h3_blackwell_runtime/qwen25_tokenizer/tokenizer_config.json

@@ -0,0 +1,241 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "processor_class": "Qwen2_5_VLProcessor",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

src/h3_blackwell_runtime/qwen3vl_text.py

@@ -0,0 +1,166 @@
+"""Direct prompt-only Qwen3-VL-32B conditioning for MiniMax H3.
+
+The MiniMax checkpoint contains the first 50 Qwen3-VL decoder layers.  H3
+uses their unnormalized final output, not the language-model head or final
+RMSNorm.  This module intentionally has no ComfyUI imports and does not load
+the Qwen vision encoder; image/video prompt construction remains a separate
+feature.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from .checkpoint import H3Checkpoint
+
+
+@dataclass(frozen=True)
+class Qwen3VL32BTextConfig:
+    vocab_size: int = 151936
+    hidden_size: int = 5120
+    intermediate_size: int = 25600
+    num_layers: int = 50
+    num_attention_heads: int = 64
+    num_key_value_heads: int = 8
+    head_dim: int = 128
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 5_000_000.0
+
+
+class _RMSNorm(nn.Module):
+    def __init__(self, weight: torch.Tensor, eps: float):
+        super().__init__()
+        self.eps = eps
+        self.register_buffer("weight", weight, persistent=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        variance = x.float().square().mean(dim=-1, keepdim=True)
+        return (x * torch.rsqrt(variance + self.eps)).to(x.dtype) * self.weight.to(x.dtype)
+
+
+def _rope(query: torch.Tensor, key: torch.Tensor, theta: float) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply Qwen's split-half rotary embedding to [B, H, S, D] Q/K tensors."""
+    positions = torch.arange(query.shape[-2], device=query.device, dtype=torch.float32)
+    dimensions = torch.arange(0, query.shape[-1], 2, device=query.device, dtype=torch.float32)
+    frequencies = positions[:, None] / theta ** (dimensions / query.shape[-1])
+    angles = torch.cat((frequencies, frequencies), dim=-1)
+    cos = angles.cos()[None, None].to(query.dtype)
+    sin = angles.sin()[None, None].to(query.dtype)
+
+    def rotate_half(value: torch.Tensor) -> torch.Tensor:
+        first, second = value.chunk(2, dim=-1)
+        return torch.cat((-second, first), dim=-1)
+
+    return query * cos + rotate_half(query) * sin, key * cos + rotate_half(key) * sin
+
+
+class _Qwen3VLBlock(nn.Module):
+    def __init__(self, checkpoint: H3Checkpoint, prefix: str, config: Qwen3VL32BTextConfig, dtype: torch.dtype):
+        super().__init__()
+        self.config = config
+        self.input_layernorm = _RMSNorm(checkpoint.tensor(f"{prefix}.input_layernorm.weight", dtype=dtype), config.rms_norm_eps)
+        self.post_attention_layernorm = _RMSNorm(checkpoint.tensor(f"{prefix}.post_attention_layernorm.weight", dtype=dtype), config.rms_norm_eps)
+        self.q_norm = _RMSNorm(checkpoint.tensor(f"{prefix}.self_attn.q_norm.weight", dtype=dtype), config.rms_norm_eps)
+        self.k_norm = _RMSNorm(checkpoint.tensor(f"{prefix}.self_attn.k_norm.weight", dtype=dtype), config.rms_norm_eps)
+        self.q_proj = checkpoint.nvfp4_linear(f"{prefix}.self_attn.q_proj", output_dtype=dtype)
+        self.k_proj = checkpoint.nvfp4_linear(f"{prefix}.self_attn.k_proj", output_dtype=dtype)
+        self.v_proj = checkpoint.nvfp4_linear(f"{prefix}.self_attn.v_proj", output_dtype=dtype)
+        self.o_proj = checkpoint.nvfp4_linear(f"{prefix}.self_attn.o_proj", output_dtype=dtype)
+        self.gate_proj = checkpoint.nvfp4_linear(f"{prefix}.mlp.gate_proj", output_dtype=dtype)
+        self.up_proj = checkpoint.nvfp4_linear(f"{prefix}.mlp.up_proj", output_dtype=dtype)
+        self.down_proj = checkpoint.nvfp4_linear(f"{prefix}.mlp.down_proj", output_dtype=dtype)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        residual = hidden_states
+        x = self.input_layernorm(hidden_states)
+        batch, sequence, _ = x.shape
+        query = self.q_proj(x).view(batch, sequence, self.config.num_attention_heads, self.config.head_dim).transpose(1, 2)
+        key = self.k_proj(x).view(batch, sequence, self.config.num_key_value_heads, self.config.head_dim).transpose(1, 2)
+        value = self.v_proj(x).view(batch, sequence, self.config.num_key_value_heads, self.config.head_dim).transpose(1, 2)
+        query = self.q_norm(query)
+        key = self.k_norm(key)
+        query, key = _rope(query, key, self.config.rope_theta)
+        attention = F.scaled_dot_product_attention(query, key, value, is_causal=True, enable_gqa=True)
+        hidden_states = residual + self.o_proj(attention.transpose(1, 2).reshape(batch, sequence, -1))
+        residual = hidden_states
+        x = self.post_attention_layernorm(hidden_states)
+        return residual + self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Qwen3VL32BTextEncoder(nn.Module):
+    """Mounted-checkpoint Qwen3-VL prompt conditioner returning layer-50 states."""
+    config = Qwen3VL32BTextConfig()
+
+    def __init__(self, checkpoint_path: str | Path, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.bfloat16):
+        super().__init__()
+        self.checkpoint_path = str(checkpoint_path)
+        self.device_name = str(device)
+        self.dtype = dtype
+        checkpoint = H3Checkpoint(checkpoint_path, device=device)
+        self._validate_checkpoint(checkpoint)
+        self.register_buffer("embed_tokens", checkpoint.tensor("model.embed_tokens.weight", dtype=dtype), persistent=False)
+        self.layers = nn.ModuleList(
+            _Qwen3VLBlock(checkpoint, f"model.layers.{index}", self.config, dtype)
+            for index in range(self.config.num_layers)
+        )
+
+    @classmethod
+    def _validate_checkpoint(cls, checkpoint: H3Checkpoint) -> None:
+        from safetensors import safe_open
+
+        required = {"model.embed_tokens.weight"}
+        linear_names = (
+            "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj",
+            "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj",
+        )
+        for index in range(cls.config.num_layers):
+            prefix = f"model.layers.{index}"
+            required.update({
+                f"{prefix}.input_layernorm.weight",
+                f"{prefix}.post_attention_layernorm.weight",
+                f"{prefix}.self_attn.q_norm.weight",
+                f"{prefix}.self_attn.k_norm.weight",
+            })
+            for linear in linear_names:
+                required.update(
+                    f"{prefix}.{linear}.{suffix}"
+                    for suffix in ("comfy_quant", "weight", "weight_scale", "weight_scale_2")
+                )
+        with safe_open(checkpoint.path, framework="pt", device="cpu") as file:
+            names = set(file.keys())
+        missing = sorted(required - names)
+        if missing:
+            raise ValueError(
+                "Not a supported MiniMax Qwen3-VL-32B NVFP4 checkpoint; missing "
+                + ", ".join(missing)
+            )
+
+    @torch.inference_mode()
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """Return unnormalized `[batch, tokens, 5120]` output after decoder layer 50."""
+        if input_ids.ndim != 2:
+            raise ValueError(f"input_ids must have shape [batch, tokens], got {tuple(input_ids.shape)}")
+        if input_ids.numel() == 0:
+            raise ValueError("input_ids must contain at least one token")
+        hidden_states = F.embedding(input_ids.to(self.embed_tokens.device), self.embed_tokens).to(self.dtype)
+        for layer in self.layers:
+            hidden_states = layer(hidden_states)
+        return hidden_states
+
+
+class Qwen3VLPromptConditioner:
+    """Tokenize raw H3 prompt text and produce Qwen layer-50 conditioning."""
+    def __init__(self, checkpoint_path: str | Path, tokenizer_dir: str | Path, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.bfloat16):
+        from .conditioning import H3PromptTokenizer
+
+        self.tokenizer = H3PromptTokenizer(tokenizer_dir)
+        self.encoder = Qwen3VL32BTextEncoder(checkpoint_path, device=device, dtype=dtype)
+        self.device = device
+
+    def __call__(self, prompt: str) -> torch.Tensor:
+        return self.encoder(self.tokenizer(prompt, device=self.device))

src/h3_blackwell_runtime/rope.py

@@ -0,0 +1,15 @@
+"""H3 three-axis split-half rotary position embeddings."""
+
+import torch
+
+
+def h3_rope_rotation(position_ids: torch.Tensor, inv_freq: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    """Build H3's `[1, sequence, 1, 48, 2, 2]` rotation table."""
+    positions = position_ids.to(torch.float32)
+    frequencies = inv_freq.to(device=positions.device, dtype=torch.float32)
+    per_axis = positions.unsqueeze(-1) * frequencies.view(1, 1, -1)
+    half_angles = torch.cat(per_axis.unbind(dim=1), dim=-1)
+    cosine, sine = torch.cos(half_angles), torch.sin(half_angles)
+    return torch.stack((cosine, -sine, sine, cosine), dim=-1).reshape(
+        1, positions.shape[0], 1, half_angles.shape[-1], 2, 2
+    ).to(dtype)

src/h3_blackwell_runtime/sampler.py

@@ -0,0 +1,85 @@
+"""Minimal direct prompt-only H3 video sampler for visual smoke previews."""
+
+import torch
+
+from .packing import H3PromptPacker, unpatchify_video
+
+
+def shifted_sigma(base: torch.Tensor, shift: float) -> torch.Tensor:
+    """H3 flow-SNR shift used by Comfy's ModelSamplingDiscreteFlow."""
+    return shift * base / (1 + (shift - 1) * base)
+
+
+def beta_sigmas(steps: int, *, device: torch.device | str, alpha: float = 0.6, beta: float = 0.6) -> torch.Tensor:
+    """Comfy's discrete beta scheduler over H3's 1,000-entry shift-12 table."""
+    from scipy.stats import beta as beta_distribution
+
+    table = shifted_sigma(torch.arange(1, 1001, device=device, dtype=torch.float32) / 1000, 12.0)
+    fractions = 1.0 - torch.arange(steps, device=device, dtype=torch.float64).cpu().numpy() / steps
+    indices = torch.from_numpy((999 * beta_distribution.ppf(fractions, alpha, beta)).round().astype("int64")).to(device)
+    indices = torch.unique_consecutive(indices)
+    return torch.cat((table[indices], table.new_zeros(1)))
+
+
+def res_multistep_update(x: torch.Tensor, denoised: torch.Tensor, sigma: torch.Tensor, sigma_down: torch.Tensor, old_denoised: torch.Tensor | None, old_sigma_down: torch.Tensor | None, previous_sigma: torch.Tensor | None) -> torch.Tensor:
+    """Deterministic H3/Comfy RES multistep update."""
+    if sigma_down == 0 or old_denoised is None:
+        return x + (x - denoised) / sigma * (sigma_down - sigma)
+    t, t_old, t_next, t_prev = sigma.log().neg(), old_sigma_down.log().neg(), sigma_down.log().neg(), previous_sigma.log().neg()
+    h = t_next - t
+    c2 = (t_prev - t_old) / h
+    phi1 = torch.expm1(-h) / -h
+    phi2 = (phi1 - 1.0) / -h
+    b1 = torch.nan_to_num(phi1 - phi2 / c2, nan=0.0)
+    b2 = torch.nan_to_num(phi2 / c2, nan=0.0)
+    return torch.exp(-h) * x + h * (b1 * denoised + b2 * old_denoised)
+
+
+def _unpack_audio(rows: torch.Tensor) -> torch.Tensor:
+    steps = rows.shape[0] // 2
+    return rows.reshape(2, steps, 32).permute(2, 0, 1).unsqueeze(0)
+
+
+def _audio_sigma(video_sigma: torch.Tensor) -> torch.Tensor:
+    base = video_sigma / (12.0 + video_sigma * (1.0 - 12.0))
+    return 3.0 * base / (1.0 + (3.0 - 1.0) * base)
+
+
+@torch.inference_mode()
+def sample_video_res_multistep(model, packer: H3PromptPacker, text: torch.Tensor, video: torch.Tensor, audio: torch.Tensor, *, steps: int = 12) -> torch.Tensor:
+    """Direct H3 beta/RES sampling with Comfy-equivalent joint AV carry semantics."""
+    sigmas = beta_sigmas(steps, device=video.device)
+    audio_carried = audio * 4.0
+    video_history = audio_history = None
+    video_history_sigma = audio_history_sigma = None
+    for index, sigma in enumerate(sigmas[:-1]):
+        sigma_down = sigmas[index + 1]
+        sigma_audio = _audio_sigma(sigma)
+        native_audio = audio_carried * (sigma_audio / sigma)
+        hidden, times, segments, positions, video_segment, audio_segment = packer(text, video, native_audio, float(sigma))
+        raw_video, raw_audio = model(hidden, times, positions, segments, video_segment, audio_segment)
+        velocity_video = -unpatchify_video(raw_video, video.shape[2], video.shape[-2], video.shape[-1])
+        velocity_audio_native = -_unpack_audio(raw_audio)
+        carry = sigma_audio / sigma
+        velocity_audio = (1.0 - 4.0) * (audio_carried * carry) + (1.0 + 3.0 * sigma_audio) * velocity_audio_native
+        video_denoised = video - sigma * velocity_video
+        audio_denoised = audio_carried - sigma * velocity_audio
+        previous_sigma = sigmas[index - 1] if index else None
+        video = res_multistep_update(video, video_denoised, sigma, sigma_down, video_history, video_history_sigma, previous_sigma)
+        audio_carried = res_multistep_update(audio_carried, audio_denoised, sigma, sigma_down, audio_history, audio_history_sigma, previous_sigma)
+        video_history, audio_history = video_denoised, audio_denoised
+        video_history_sigma = audio_history_sigma = sigma_down
+    return video
+
+
+@torch.inference_mode()
+def sample_video_euler(model, packer: H3PromptPacker, text: torch.Tensor, video: torch.Tensor, audio: torch.Tensor, *, steps: int = 2) -> torch.Tensor:
+    """Use Euler updates to obtain a visual-only H3 preview, not parity sampling."""
+    sigmas = beta_sigmas(steps, device=video.device)
+    for index in range(steps):
+        sigma = sigmas[index]
+        hidden, timesteps, segments, positions, video_segment, audio_segment = packer(text, video, audio, float(sigma))
+        velocity, _ = model(hidden, timesteps, positions, segments, video_segment, audio_segment)
+        velocity = unpatchify_video(velocity, video.shape[2], video.shape[-2], video.shape[-1])
+        video.add_(velocity.to(video.dtype), alpha=float(sigmas[index + 1] - sigma))
+    return video

src/h3_blackwell_runtime/t2v.py

@@ -0,0 +1,30 @@
+"""Direct prompt-only MiniMax H3 T2V shape and latent helpers."""
+
+import torch
+
+
+FPS = 24
+AUDIO_LATENT_FPS = 40
+
+
+def align_frame_count(frames: int) -> int:
+    """Snap to H3's valid 17k+5 temporal grid."""
+    frames = max(5, frames)
+    return frames + (5 - frames) % 17
+
+
+def temporal_shape(frames: int) -> tuple[int, int, int]:
+    """Return output frames, video-latent frames, and joint audio-latent steps."""
+    frames = align_frame_count(frames)
+    video_steps = 2 if frames <= 5 else ((frames - 5) // 17) * 5 + 2
+    return frames, video_steps, round(frames / FPS * AUDIO_LATENT_FPS)
+
+
+def empty_av_latents(width: int, height: int, frames: int, *, device: torch.device | str = "cuda") -> tuple[torch.Tensor, torch.Tensor, int]:
+    """Allocate H3's joint video/audio sampling state without ComfyUI objects."""
+    if width % 32 or height % 32:
+        raise ValueError("H3 T2V dimensions must be multiples of 32.")
+    frames, video_steps, audio_steps = temporal_shape(frames)
+    video = torch.zeros((1, 24, video_steps, height // 16, width // 16), device=device)
+    audio = torch.zeros((1, 32, 2, audio_steps), device=device)
+    return video, audio, frames

src/h3_blackwell_runtime/token_refiner.py

@@ -0,0 +1,58 @@
+"""Direct two-block H3 text token refiner."""
+
+import torch
+import torch.nn.functional as functional
+from torch import nn
+
+from .attention import rms_norm
+from .checkpoint import H3Checkpoint
+
+
+class _Linear(nn.Module):
+    def __init__(self, checkpoint: H3Checkpoint, prefix: str, dtype: torch.dtype):
+        super().__init__()
+        self.register_buffer("weight", checkpoint.tensor(f"{prefix}.weight", dtype=dtype), persistent=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return functional.linear(x, self.weight)
+
+
+class _RefinerBlock(nn.Module):
+    def __init__(self, checkpoint: H3Checkpoint, prefix: str, dtype: torch.dtype):
+        super().__init__()
+        self.qkv = _Linear(checkpoint, f"{prefix}.attn.qkv_proj", dtype)
+        self.out = _Linear(checkpoint, f"{prefix}.attn.out_proj", dtype)
+        self.fc1 = _Linear(checkpoint, f"{prefix}.mlp.fc1", dtype)
+        self.fc2 = _Linear(checkpoint, f"{prefix}.mlp.fc2", dtype)
+        self.register_buffer("norm1", checkpoint.tensor(f"{prefix}.norm1.weight", dtype=dtype), persistent=False)
+        self.register_buffer("norm2", checkpoint.tensor(f"{prefix}.norm2.weight", dtype=dtype), persistent=False)
+        self.register_buffer("q_norm", checkpoint.tensor(f"{prefix}.attn.q_norm.weight", dtype=dtype), persistent=False)
+        self.register_buffer("k_norm", checkpoint.tensor(f"{prefix}.attn.k_norm.weight", dtype=dtype), persistent=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        normalized = rms_norm(x, self.norm1, 1e-5)
+        sequence = x.shape[0]
+        q, k, v = self.qkv(normalized).split(7168, dim=-1)
+        q = rms_norm(q.view(1, sequence, 56, 128), self.q_norm, 1e-5).transpose(1, 2)
+        k = rms_norm(k.view(1, sequence, 56, 128), self.k_norm, 1e-5).transpose(1, 2)
+        v = v.view(1, sequence, 56, 128).transpose(1, 2)
+        x = x + self.out(functional.scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(sequence, 7168))
+        gate, up = self.fc1(rms_norm(x, self.norm2, 1e-5)).chunk(2, dim=-1)
+        return x + self.fc2(functional.silu(gate) * up)
+
+
+class H3TokenRefiner(nn.Module):
+    """Project Qwen layer-50 states and refine them for H3 T2V."""
+    def __init__(self, checkpoint: H3Checkpoint, dtype: torch.dtype = torch.bfloat16):
+        super().__init__()
+        self.register_buffer("condition_weight", checkpoint.tensor("condition_proj.weight", dtype=dtype), persistent=False)
+        self.register_buffer("condition_bias", checkpoint.tensor("condition_proj.bias", dtype=dtype), persistent=False)
+        self.blocks = nn.ModuleList(_RefinerBlock(checkpoint, f"token_refiner.blocks.{index}", dtype) for index in range(2))
+        self.register_buffer("final_norm", checkpoint.tensor("token_refiner.final_norm.weight", dtype=dtype), persistent=False)
+
+    @torch.inference_mode()
+    def forward(self, qwen_states: torch.Tensor) -> torch.Tensor:
+        x = functional.linear(qwen_states[0].to(self.condition_weight.dtype), self.condition_weight, self.condition_bias)
+        for block in self.blocks:
+            x = block(x)
+        return rms_norm(x, self.final_norm, 1e-5).unsqueeze(0)

src/h3_blackwell_runtime/upstream_model.py

@@ -0,0 +1,694 @@
+"""MiniMax H3 audio-video DiT.
+
+Single-stream packed-token transformer denoising video (24ch, patch 1x2x2) and
+stereo audio (32ch, 40 Hz) latents jointly, conditioned on Qwen3-VL layer-50 hidden states.
+The packed sequence is:
+[text | cond rows | audio | video] for t2va/fl2va
+[text | reference blocks | audio | video] for ref2va
+
+Timestep domain: the model receives the *video* sigma from the sampler and
+derives per-token timesteps t = 1 - sigma internally; the audio stream runs on
+its own shifted schedule (sigma_shift video 12.0 / audio 3.0), mapped from the
+video sigma in closed form. The sampler carries the audio latent scaled onto the
+video schedule (ModelSamplingAV); forward() undoes that scale and converts the
+velocity back, so _forward only ever sees the stream's own latent.
+"""
+
+import math
+import os
+
+import torch
+import torch.nn as nn
+
+import comfy.ldm.common_dit
+import comfy.model_management
+import comfy.model_prefetch
+import comfy.ops
+import comfy.patcher_extension
+import comfy.quant_ops
+from comfy.ldm.modules.attention import optimized_attention
+
+FRAME_PER_TOKEN = (1, 4, 4, 4, 4)
+FRAME_RESCALE = 5.0 / 3.0
+VISUAL_COND_TIMESTEP = 0.999
+H3_CAPTURE_ACTIVE = False
+AUDIO_COND_TIMESTEP = 1.0
+
+
+def time_shift_sigma(sigma, from_shift, to_shift):
+    # invert sigma = s*b/(1+(s-1)*b) to the base grid, re-apply the other shift
+    base = sigma / (from_shift + sigma * (1.0 - from_shift))
+    return to_shift * base / (1.0 + (to_shift - 1.0) * base)
+
+
+def patchify_video(latent, patch_size=(1, 2, 2)):
+    # [B, C, T, H, W] -> [B*t*h*w, C*pt*ph*pw]
+    b, c, t_full, h_full, w_full = latent.shape
+    pt, ph, pw = patch_size
+    t, h, w = t_full // pt, h_full // ph, w_full // pw
+    x = latent.reshape(b, c, t, pt, h, ph, w, pw)
+    x = torch.einsum("nctrhpwq->nthwcrpq", x)
+    return x.reshape(b * t * h * w, c * pt * ph * pw)
+
+
+def unpatchify_video(rows, t, h, w, c=24, patch_size=(1, 2, 2)):
+    pt, ph, pw = patch_size
+    x = rows.reshape(-1, t, h, w, c, pt, ph, pw)
+    x = torch.einsum("nthwcrpq->nctrhpwq", x)
+    return x.reshape(-1, c, t * pt, h * ph, w * pw)
+
+
+def pack_audio(latent):
+    # [B, C=32, ch=2, T] -> [ch*T, 32] channel-major (ch0 t0..T-1, ch1 t0..T-1)
+    b, c, ch, t = latent.shape
+    return latent[0].permute(1, 2, 0).reshape(ch * t, c)
+
+
+def unpack_audio(rows, ch=2):
+    t = rows.shape[0] // ch
+    return rows.reshape(ch, t, rows.shape[-1]).permute(2, 0, 1).unsqueeze(0)
+
+
+def _axis_from_sqrt_area(dim, patch, sqrt_area):
+    # linspace((1 - ratio) / 2, (1 + ratio) / 2, dim // patch, endpoint=False) * 32
+    ratio = dim / sqrt_area
+    n = dim // patch
+    return (torch.arange(n, dtype=torch.float64) * (ratio / n) + (1.0 - ratio) / 2.0) * 32.0
+
+
+def _frame_grid(h, w):
+    # area-normalized (h, w) coordinates of one latent frame's 2x2-patch rows
+    area = math.sqrt(h * w)
+    hh, ww = torch.meshgrid(_axis_from_sqrt_area(h, 2, area), _axis_from_sqrt_area(w, 2, area), indexing="ij")
+    return torch.stack([hh.reshape(-1), ww.reshape(-1)], dim=-1), _axis_from_sqrt_area(w, 2, area)
+
+
+def _video_t_spans(n):
+    return [FRAME_RESCALE * FRAME_PER_TOKEN[k % 5] for k in range(n)]
+
+
+def _video_t_grid(n, origin):
+    # origin + exclusive cumsum
+    spans = torch.tensor(_video_t_spans(n), dtype=torch.float64)
+    return float(origin) + torch.cat([torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)])
+
+
+def _audio_grid(cursor, t, w_low, w_high):
+    # channel-major stereo rows: t advances per latent frame, w pinned to the grid extremes per stereo channel, h stays 0
+    g = torch.zeros(t * 2, 3, dtype=torch.float64)
+    g[:, 0] = (cursor + torch.arange(t, dtype=torch.float64)).repeat(2)
+    g[:t, 2] = w_low
+    g[t:, 2] = w_high
+    return g
+
+
+def _video_grid(vt, frame, cursor):
+    g = torch.empty(vt, frame.shape[0], 3, dtype=torch.float64)
+    g[:, :, 0] = _video_t_grid(vt, cursor)[:, None]
+    g[:, :, 1:] = frame[None]
+    return g.reshape(-1, 3)
+
+
+class TimeEmbedder(nn.Module):
+    def __init__(self, freq_dim, hidden, out, dtype=None, device=None, operations=None):
+        super().__init__()
+        self.freq_dim = freq_dim
+        self.proj_in = operations.Linear(freq_dim, hidden, bias=True, dtype=dtype, device=device)
+        self.proj_out = operations.Linear(hidden, out, bias=True, dtype=dtype, device=device)
+
+    def forward(self, t):
+        # t: [M] in [0, 1]; fp32 throughout, cos before sin
+        half = self.freq_dim // 2
+        freqs = torch.exp(-math.log(10000.0) * torch.arange(half, dtype=torch.float32, device=t.device) / half)
+        args = t.to(torch.float32)[:, None] * freqs[None]
+        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        return self.proj_out(nn.functional.silu(self.proj_in(emb)))
+
+
+def rope_rotation_table(angles, dtype):
+    """[S, rot_dim] pair angles -> [1, S, 1, rot_dim/2, 2, 2] rotation matrices."""
+    half = angles.shape[-1] // 2
+    ang = angles[:, :half]  # duplicated halves: [:, :half] == [:, half:]
+    c, s = torch.cos(ang), torch.sin(ang)
+    table = torch.stack([c, -s, s, c], dim=-1).reshape(1, angles.shape[0], 1, half, 2, 2)
+    return table.to(dtype)
+
+
+class Attention(nn.Module):
+    def __init__(self, hidden, heads, head_dim, eps, dtype=None, device=None, operations=None):
+        super().__init__()
+        self.heads = heads
+        self.head_dim = head_dim
+        inner = heads * head_dim
+        self.qkv_proj = operations.Linear(hidden, inner * 3, bias=False, dtype=dtype, device=device)
+        self.q_norm = operations.RMSNorm(head_dim, eps=eps, dtype=dtype, device=device)
+        self.k_norm = operations.RMSNorm(head_dim, eps=eps, dtype=dtype, device=device)
+        self.out_proj = operations.Linear(inner, hidden, bias=False, dtype=dtype, device=device)
+
+    def forward(self, x, rope_freqs=None, transformer_options={}):
+        s = x.shape[0]
+        q, k, v = self.qkv_proj(x).split(self.heads * self.head_dim, dim=-1)
+        if getattr(self, "_h3_capture_index", -1) == 0 and H3_CAPTURE_ACTIVE:
+            capture_dir = os.getenv("H3_CAPTURE_DIR")
+            if capture_dir:
+                torch.save({"q": q.detach().cpu(), "k": k.detach().cpu(), "v": v.detach().cpu()}, os.path.join(capture_dir, "block0_qkv_raw.pt"))
+        v = v.view(s, self.heads, self.head_dim)
+        if rope_freqs is not None:
+            # fused per-head RMSNorm + partial split-half rope, in place on the qkv buffer
+            q = q.view(1, s, self.heads, self.head_dim)
+            k = k.view(1, s, self.heads, self.head_dim)
+            qw = comfy.model_management.cast_to(self.q_norm.weight, device=x.device)
+            kw = comfy.model_management.cast_to(self.k_norm.weight, device=x.device)
+            rot = rope_freqs.shape[-3] * 2
+            if comfy.model_management.in_training:
+                q, k = comfy.quant_ops.ck.rms_rope_split_half(
+                    q, k, rope_freqs, qw, kw, epsilon=self.q_norm.eps, rot_dim=rot)
+            else:
+                comfy.quant_ops.ck.rms_rope_split_half_(
+                    q, k, rope_freqs, qw, kw, epsilon=self.q_norm.eps, rot_dim=rot)
+            q = q[0]
+            k = k[0]
+        else:
+            q = self.q_norm(q.view(s, self.heads, self.head_dim))
+            k = self.k_norm(k.view(s, self.heads, self.head_dim))
+        q = q.transpose(0, 1).unsqueeze(0)
+        k = k.transpose(0, 1).unsqueeze(0)
+        v = v.transpose(0, 1).unsqueeze(0)
+        if getattr(self, "_h3_capture_index", -1) == 0 and H3_CAPTURE_ACTIVE:
+            capture_dir = os.getenv("H3_CAPTURE_DIR")
+            if capture_dir:
+                torch.save({"q": q.detach().cpu(), "k": k.detach().cpu(), "v": v.detach().cpu()}, os.path.join(capture_dir, "block0_qkv_prepared.pt"))
+        out = optimized_attention(q, k, v, self.heads, mask=None, skip_reshape=True, transformer_options=transformer_options)
+        return self.out_proj(out.squeeze(0))
+
+
+class MLP(nn.Module):
+    def __init__(self, hidden, ffn, dtype=None, device=None, operations=None):
+        super().__init__()
+        self.fc1 = operations.Linear(hidden, ffn * 2, bias=False, dtype=dtype, device=device)
+        self.fc2 = operations.Linear(ffn, hidden, bias=False, dtype=dtype, device=device)
+
+    def forward(self, x):
+        return comfy.ops.linear_input_act(self.fc2, self.fc1(x), "swiglu")
+
+
+class AdalnProj(nn.Module):
+    def __init__(self, t_dim, hidden, expand, modalities, apply_silu=True,
+                 dtype=None, device=None, operations=None):
+        super().__init__()
+        self.expand = expand
+        self.modalities = modalities
+        self.hidden = hidden
+        self.apply_silu = apply_silu
+        self.linear = operations.Linear(t_dim, expand * hidden * modalities, bias=True, dtype=dtype, device=device)
+
+    def forward(self, t_emb):
+        # [M, t_dim] -> expand tensors of [M*modalities, hidden]
+        x = self.linear(nn.functional.silu(t_emb) if self.apply_silu else t_emb)
+        x = x.view(x.shape[0] * self.modalities, self.expand * self.hidden)
+        return x.chunk(self.expand, dim=-1)
+
+
+def _mod_scale_shift(h, shift, scale, segments):
+    # segments: [(start, stop, mod_row)] covering h contiguously.
+    for a, b, row in segments:
+        h[a:b].mul_(1.0 + scale[row].to(h.dtype)).add_(shift[row].to(h.dtype))
+    return h
+
+
+def _mod_gate(x, gate, other, segments):
+    # other is the fresh attn/mlp output: accumulate the gated residual into the stream in place, one fused kernel per segment
+    for a, b, row in segments:
+        x[a:b].addcmul_(other[a:b], gate[row].to(x.dtype))
+    return x
+
+
+class RefinerBlock(nn.Module):
+    def __init__(self, hidden, heads, head_dim, ffn, eps, qk_eps, dtype=None, device=None, operations=None):
+        super().__init__()
+        self.norm1 = operations.RMSNorm(hidden, eps=eps, dtype=dtype, device=device)
+        self.norm2 = operations.RMSNorm(hidden, eps=eps, dtype=dtype, device=device)
+        self.attn = Attention(hidden, heads, head_dim, qk_eps, dtype=dtype, device=device, operations=operations)
+        self.mlp = MLP(hidden, ffn, dtype=dtype, device=device, operations=operations)
+
+    def forward(self, x, transformer_options={}):
+        # attn/mlp outputs are fresh: accumulate residuals in place
+        x = self.attn(self.norm1(x), transformer_options=transformer_options).add_(x)
+        return self.mlp(self.norm2(x)).add_(x)
+
+
+class TokenRefiner(nn.Module):
+    def __init__(self, num_layers, hidden, heads, head_dim, ffn, eps, qk_eps, final_eps,
+                 dtype=None, device=None, operations=None):
+        super().__init__()
+        self.blocks = nn.ModuleList([
+            RefinerBlock(hidden, heads, head_dim, ffn, eps, qk_eps, dtype=dtype, device=device, operations=operations)
+            for _ in range(num_layers)])
+        self.final_norm = operations.RMSNorm(hidden, eps=final_eps, dtype=dtype, device=device)
+
+    def forward(self, x, transformer_options={}):
+        for block in self.blocks:
+            x = block(x, transformer_options=transformer_options)
+        return self.final_norm(x)
+
+
+class DiTBlock(nn.Module):
+    def __init__(self, hidden, heads, head_dim, ffn, t_dim, eps, qk_eps,
+                 apply_silu=True, adaln_dtype=None, dtype=None, device=None, operations=None):
+        super().__init__()
+        self.norm1 = operations.RMSNorm(hidden, eps=eps, dtype=dtype, device=device)
+        self.norm2 = operations.RMSNorm(hidden, eps=eps, dtype=dtype, device=device)
+        self.attn = Attention(hidden, heads, head_dim, qk_eps, dtype=dtype, device=device, operations=operations)
+        self.mlp = MLP(hidden, ffn, dtype=dtype, device=device, operations=operations)
+        self.adaln_proj = AdalnProj(t_dim, hidden, 6, 3, apply_silu=apply_silu,
+                                    dtype=adaln_dtype if adaln_dtype is not None else dtype,
+                                    device=device, operations=operations)
+
+    def forward(self, x, t_emb, mod_segments, rope_freqs, transformer_options={}):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(t_emb)
+        capture_dir = os.getenv("H3_CAPTURE_DIR") if getattr(self, "_h3_capture_index", -1) == 0 and H3_CAPTURE_ACTIVE else None
+        norm1 = self.norm1(x)
+        if capture_dir:
+            effective_weight, effective_bias, effective_stream = comfy.ops.cast_bias_weight(self.adaln_proj.linear, t_emb, offloadable=True)
+            torch.save({"norm": norm1.detach().cpu(), "shift": shift_msa.detach().cpu(), "scale": scale_msa.detach().cpu(), "effective_weight": effective_weight.detach().cpu(), "effective_bias": effective_bias.detach().cpu()}, os.path.join(capture_dir, "block0_norm1_adaln.pt"))
+            comfy.ops.uncast_bias_weight(self.adaln_proj.linear, effective_weight, effective_bias, effective_stream)
+        h = _mod_scale_shift(norm1, shift_msa, scale_msa, mod_segments)
+        if capture_dir: torch.save(h.detach().cpu(), os.path.join(capture_dir, "block0_norm1.pt"))
+        attn_out = self.attn(h, rope_freqs=rope_freqs, transformer_options=transformer_options)
+        if capture_dir: torch.save(attn_out.detach().cpu(), os.path.join(capture_dir, "block0_attention.pt"))
+        x = _mod_gate(x, gate_msa, attn_out, mod_segments)
+        if capture_dir: torch.save(x.detach().cpu(), os.path.join(capture_dir, "block0_post_attention.pt"))
+        h = _mod_scale_shift(self.norm2(x), shift_mlp, scale_mlp, mod_segments)
+        if capture_dir: torch.save(h.detach().cpu(), os.path.join(capture_dir, "block0_norm2.pt"))
+        mlp_out = self.mlp(h)
+        if capture_dir: torch.save(mlp_out.detach().cpu(), os.path.join(capture_dir, "block0_mlp.pt"))
+        x = _mod_gate(x, gate_mlp, mlp_out, mod_segments)
+        if capture_dir: torch.save(x.detach().cpu(), os.path.join(capture_dir, "block0_post_mlp.pt"))
+        return x
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, hidden, t_dim, video_dim, audio_dim, eps, apply_silu=True, adaln_dtype=None,
+                 dtype=None, device=None, operations=None):
+        super().__init__()
+        self.norm = operations.RMSNorm(hidden, eps=eps, dtype=dtype, device=device)
+        self.adaln_proj = AdalnProj(t_dim, hidden, 2, 1, apply_silu=apply_silu,
+                                    dtype=adaln_dtype if adaln_dtype is not None else dtype,
+                                    device=device, operations=operations)
+        # output heads are the checkpoint's fp32 island; norm/adaln are stored at model dtype
+        self.video_out = operations.Linear(hidden, video_dim, bias=True, dtype=torch.float32, device=device)
+        self.audio_out = operations.Linear(hidden, audio_dim, bias=True, dtype=torch.float32, device=device)
+
+    def forward(self, x, t_emb, video_seg, audio_seg):
+        # video_seg / audio_seg: (start, stop, timestep_row) of the target streams
+        shift, scale = self.adaln_proj(t_emb)
+        va, vb, vrow = video_seg
+        aa, ab, arow = audio_seg
+        hv = (self.norm(x[va:vb]) * (1.0 + scale[vrow]) + shift[vrow]).to(torch.float32)
+        ha = (self.norm(x[aa:ab]) * (1.0 + scale[arow]) + shift[arow]).to(torch.float32)
+        return self.video_out(hv), self.audio_out(ha)
+
+
+class PackedLayout:
+    """Static packed-sequence structure for one shape/conditioning signature."""
+
+    def __init__(self, text_len, latent_t, latent_h, latent_w, audio_t, keyframes=None, refs=None, frame_count=None):
+        frame, w_grid = _frame_grid(latent_h, latent_w)
+        frame_rows = frame.shape[0]
+
+        segments = [("text", text_len)]  # (kind, n_rows)
+        g = torch.zeros(text_len, 3, dtype=torch.float64)
+        g[:, 0] = torch.arange(text_len, dtype=torch.float64)
+        pos = [g]  # per segment: [n, 3] float64 (t, h, w)
+
+        img_pos, img_update = [], []
+        audio_pos, audio_update = [], []
+        cursor = text_len
+        row = text_len
+
+        if keyframes:
+            # fl2va: keyframe cond rows right after text, sharing the target spatial grid
+            for kf in keyframes:
+                pixel_index = kf["resolved_frame_index"]
+                if pixel_index == 0:
+                    cond_t = float(text_len)
+                elif frame_count is not None and pixel_index == frame_count - 1:
+                    cond_t = float(text_len) + sum(_video_t_spans(latent_t)) - FRAME_RESCALE
+                else:
+                    raise ValueError("only first/last keyframe anchors are supported")
+                g = torch.empty(frame_rows, 3, dtype=torch.float64)
+                g[:, 0] = cond_t
+                g[:, 1:] = frame
+                segments.append(("cond", frame_rows))
+                pos.append(g)
+                img_pos.append(torch.arange(row, row + frame_rows))
+                img_update.append(torch.zeros(frame_rows, dtype=torch.bool))
+                row += frame_rows
+
+        target_audio_w = (float(w_grid[0]), float(w_grid[-1]))
+        if refs:
+            cursor = float(text_len)
+            for blk in refs:
+                kind = blk["kind"]
+                if kind == "image":
+                    r_frame, _ = _frame_grid(blk["latent_h"], blk["latent_w"])
+                    n = r_frame.shape[0]
+                    g = torch.empty(n, 3, dtype=torch.float64)
+                    g[:, 0] = cursor
+                    g[:, 1:] = r_frame
+                    segments.append(("ref_img", n))
+                    pos.append(g)
+                    img_pos.append(torch.arange(row, row + n))
+                    img_update.append(torch.zeros(n, dtype=torch.bool))
+                    row += n
+                    cursor += 1.0
+                elif kind == "audio":
+                    rt = blk["ref_audio_t"]
+                    if rt > 0:
+                        segments.append(("ref_audio", rt * 2))
+                        pos.append(_audio_grid(cursor, rt, *target_audio_w))
+                        audio_pos.append(torch.arange(row, row + rt * 2))
+                        audio_update.append(torch.zeros(rt * 2, dtype=torch.bool))
+                        row += rt * 2
+                    cursor += float(rt)
+                elif kind in ("video", "video_audio"):
+                    # the block's audio rows pack immediately before its video
+                    # rows, both sharing the cursor origin
+                    rt = blk["ref_audio_t"]
+                    vt = blk["latent_t"]
+                    r_frame, r_w_grid = _frame_grid(blk["latent_h"], blk["latent_w"])
+                    if rt > 0:
+                        segments.append(("ref_audio", rt * 2))
+                        pos.append(_audio_grid(cursor, rt, float(r_w_grid[0]), float(r_w_grid[-1])))
+                        audio_pos.append(torch.arange(row, row + rt * 2))
+                        audio_update.append(torch.zeros(rt * 2, dtype=torch.bool))
+                        row += rt * 2
+                    n = vt * r_frame.shape[0]
+                    segments.append(("ref_img", n))
+                    pos.append(_video_grid(vt, r_frame, cursor))
+                    img_pos.append(torch.arange(row, row + n))
+                    img_update.append(torch.zeros(n, dtype=torch.bool))
+                    row += n
+                    cursor += max(float(rt), sum(_video_t_spans(vt)))
+
+        # target audio then target video, always the last two segments
+        segments.append(("audio", audio_t * 2))
+        pos.append(_audio_grid(cursor, audio_t, *target_audio_w))
+        audio_pos.append(torch.arange(row, row + audio_t * 2))
+        audio_update.append(torch.ones(audio_t * 2, dtype=torch.bool))
+        row += audio_t * 2
+
+        n_video = latent_t * frame_rows
+        segments.append(("video", n_video))
+        pos.append(_video_grid(latent_t, frame, cursor))
+        img_pos.append(torch.arange(row, row + n_video))
+        img_update.append(torch.ones(n_video, dtype=torch.bool))
+        row += n_video
+
+        self.seq_len = row
+        self.position_ids = torch.cat(pos)  # [S, 3] float64
+        self.img_pos = torch.cat(img_pos)
+        self.img_update = torch.cat(img_update)
+        self.audio_pos = torch.cat(audio_pos)
+        self.audio_update = torch.cat(audio_update)
+        self.signature = (text_len, latent_t, latent_h, latent_w, audio_t)
+        # contiguous segment table (start, stop, kind)
+        # kinds: text / cond / ref_img / ref_audio / audio / video
+        # the packed sequence is uniform per segment in (modality tag, timestep class),
+        # except the text span (tag runs resolved at forward time from the presentation tags)
+        seg_abs = []
+        off = 0
+        for kind, n in segments:
+            seg_abs.append((off, off + n, kind))
+            off += n
+        self.segments = seg_abs
+
+
+class MiniMaxH3Model(nn.Module):
+    def __init__(self, hidden_size=5376, num_layers=50, token_refiner_num_layers=2,
+                 num_attention_heads=56, attention_head_dim=128, ffn_hidden_size=14336,
+                 latents_dim=24, audio_latents_dim=32, patch_size=(1, 2, 2), text_dim=5120,
+                 timestep_input_dim=256, time_embed_hidden_size=5376, time_embed_dim=2688,
+                 rope_inv_freq_len=16, norm_eps=1e-5, qk_norm_eps=1e-5, final_norm_eps=1e-5,
+                 sigma_shift_video=12.0, sigma_shift_audio=3.0,
+                 adaln_curve_grid=None,
+                 image_model=None, dtype=None, device=None, operations=None, **kwargs):
+        super().__init__()
+        self.dtype = dtype
+        self.hidden_size = hidden_size
+        self.patch_size = tuple(patch_size)
+        self.latents_dim = latents_dim
+        self.audio_latents_dim = audio_latents_dim
+        self.sigma_shift_video = sigma_shift_video
+        self.sigma_shift_audio = sigma_shift_audio
+        self.use_adaln_curves = adaln_curve_grid is not None
+        # curve-form checkpoints replace the time embedder and full-width adaln weights with a small shared basis of the time-embedding curve
+        curve = {"apply_silu": not self.use_adaln_curves,
+                 "adaln_dtype": torch.float32 if self.use_adaln_curves else dtype}
+        video_patch_dim = latents_dim * self.patch_size[0] * self.patch_size[1] * self.patch_size[2]
+
+        self.video_patch_proj = operations.Linear(video_patch_dim, hidden_size, bias=True, dtype=torch.float32, device=device)
+        self.audio_patch_proj = operations.Linear(audio_latents_dim, hidden_size, bias=True, dtype=torch.float32, device=device)
+        self.condition_proj = operations.Linear(text_dim, hidden_size, bias=True, dtype=dtype, device=device)
+        if self.use_adaln_curves:
+            self.register_buffer("adaln_t_table", torch.empty(adaln_curve_grid, time_embed_dim, dtype=torch.float32))
+        else:
+            self.time_embedder = TimeEmbedder(timestep_input_dim, time_embed_hidden_size, time_embed_dim,
+                                              dtype=torch.float32, device=device, operations=operations)
+        self.rope = nn.Module()
+        self.rope.register_buffer("inv_freq", torch.empty(rope_inv_freq_len, dtype=torch.float32))
+        self.token_refiner = TokenRefiner(token_refiner_num_layers, hidden_size, num_attention_heads,
+                                          attention_head_dim, ffn_hidden_size, norm_eps, qk_norm_eps,
+                                          final_norm_eps, dtype=dtype, device=device, operations=operations)
+        self.blocks = nn.ModuleList([
+            DiTBlock(hidden_size, num_attention_heads, attention_head_dim, ffn_hidden_size,
+                     time_embed_dim, norm_eps, qk_norm_eps, **curve, dtype=dtype, device=device, operations=operations)
+            for _ in range(num_layers)])
+        self.final_layer = FinalLayer(hidden_size, time_embed_dim, video_patch_dim, audio_latents_dim,
+                                      final_norm_eps, **curve, dtype=dtype, device=device, operations=operations)
+
+    def preprocess_text_embeds(self, text_states):
+        """[B, L, text_dim] Qwen states -> [B, L, hidden] refined text embeds."""
+        if text_states.shape[-1] == self.hidden_size:
+            return text_states
+        return self.token_refiner(self.condition_proj(text_states[0])).unsqueeze(0)
+
+    def rope_freqs(self, position_ids, device):
+        # [S, 3] float64 -> [S, 96] fp32
+        pos = position_ids.to(torch.float32).to(device)
+        inv = comfy.model_management.cast_to(self.rope.inv_freq, device=device)
+        per_axis = pos.unsqueeze(-1) * inv.view(1, 1, -1)      # [S, 3, 16]
+        t_f, h_f, w_f = per_axis.unbind(dim=1)
+        half = torch.cat((t_f, h_f, w_f), dim=-1)              # [S, 48]
+        return torch.cat((half, half), dim=-1)                 # [S, 96]
+
+    def _cond_video_rows(self, payload, device):
+        """Concatenated visual condition rows (normalized latents -> patchified), with condition noise augmentation."""
+        rows = []
+        aug = payload.get("visual_cond_noise_aug", VISUAL_COND_TIMESTEP)
+        seed = int(payload.get("seed", 0))
+        # every condition intentionally restarts the same RNG stream
+        for z in payload.get("cond_video_latents", []):
+            r = patchify_video(z.to(torch.float32), self.patch_size)
+            if aug < 1.0:
+                gen = torch.Generator("cpu").manual_seed(seed)
+                noise = torch.randn(r.shape, generator=gen, dtype=torch.float32)
+                r = aug * r + (1.0 - aug) * noise.to(r.device)
+            rows.append(r.to(device))
+        return torch.cat(rows, dim=0) if rows else None
+
+    def _cond_audio_rows(self, payload, device):
+        rows = []
+        aug = payload.get("audio_cond_noise_aug", AUDIO_COND_TIMESTEP)
+        seed = int(payload.get("seed", 0)) + 1
+        for z in payload.get("cond_audio_latents", []):
+            r = pack_audio(z.to(torch.float32))
+            if aug < 1.0:
+                gen = torch.Generator("cpu").manual_seed(seed)
+                noise = torch.randn(r.shape, generator=gen, dtype=torch.float32)
+                r = aug * r + (1.0 - aug) * noise.to(r.device)
+            rows.append(r.to(device))
+        return torch.cat(rows, dim=0) if rows else None
+
+    def forward(self, x, timestep, context, transformer_options={}, minimax_payload=None, **kwargs):
+        # the sampler carries the audio as (sigma_v / sigma_a) * x_audio; undo it outside
+        # the wrappers so they and the network see the stream's own latent and velocity
+        scale = float((minimax_payload or {}).get("audio_scale", 1.0))
+        audio_src = x[1]
+        if scale != 1.0:
+            shift_v = float(transformer_options.get("minimax_h3_sigma_shift_video", self.sigma_shift_video))
+            shift_a = float(transformer_options.get("minimax_h3_sigma_shift_audio", self.sigma_shift_audio))
+            sigma_v = (timestep.flatten()[0] / 1000.0).float().clamp(min=1e-6)
+            sigma_a = time_shift_sigma(sigma_v, shift_v, shift_a)
+            carry = (sigma_a / sigma_v).to(audio_src.dtype)
+            x = [x[0], audio_src * carry]
+
+        out = comfy.patcher_extension.WrapperExecutor.new_class_executor(
+            self._forward,
+            self,
+            comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
+        ).execute(x, timestep, context, transformer_options, minimax_payload=minimax_payload, **kwargs)
+
+        if scale != 1.0:
+            # d/d(sigma_v) of the carried variable
+            out[1] = ((1.0 - scale) * (audio_src * carry)
+                      + (1.0 + (scale - 1.0) * sigma_a).to(out[1].dtype) * out[1])
+        return out
+
+    def _forward(self, x, timestep, context, transformer_options={}, minimax_payload=None, **kwargs):
+        video_x, audio_x = x[0], x[1]
+        orig_t, orig_h, orig_w = video_x.shape[2], video_x.shape[3], video_x.shape[4]
+        video_x = comfy.ldm.common_dit.pad_to_patch_size(video_x, self.patch_size)
+        if video_x.shape[0] != 1:
+            raise ValueError("MiniMax H3 supports batch size 1")
+        payload = minimax_payload or {}
+        device = video_x.device
+        dtype = context.dtype  # compute dtype
+
+        latent_t, lat_h, lat_w = video_x.shape[2], video_x.shape[3], video_x.shape[4]
+        audio_t = audio_x.shape[-1]
+        text_len = context.shape[1]
+        # extra_conds prebuilds the layout once per sampling run
+        layout = payload.get("layout")
+        if layout is None or layout.signature != (text_len, latent_t, lat_h, lat_w, audio_t):
+            layout = PackedLayout(text_len, latent_t, lat_h, lat_w, audio_t,
+                                  keyframes=payload.get("keyframes"),
+                                  refs=payload.get("refs"),
+                                  frame_count=payload.get("frame_count"))
+
+        # model_base passes model_sampling.timestep(sigma) = sigma * 1000
+        shift_v = float(transformer_options.get("minimax_h3_sigma_shift_video", self.sigma_shift_video))
+        shift_a = float(transformer_options.get("minimax_h3_sigma_shift_audio", self.sigma_shift_audio))
+        sigma_v = (timestep.flatten()[0] / 1000.0).float().clamp(min=1e-6)
+        t_v = float(1.0 - sigma_v)
+        t_a = float(1.0 - time_shift_sigma(sigma_v, shift_v, shift_a))
+
+        # distinct timesteps are known analytically: text/pad follow video, cond rows pin near 1
+        vis_aug = float(payload.get("visual_cond_noise_aug", VISUAL_COND_TIMESTEP))
+        aud_aug = float(payload.get("audio_cond_noise_aug", AUDIO_COND_TIMESTEP))
+        has_vis_cond = any(k in ("cond", "ref_img") for _, _, k in layout.segments)
+        has_aud_cond = any(k == "ref_audio" for _, _, k in layout.segments)
+        seg_t = {"text": t_v, "video": t_v, "audio": t_a,
+                 "cond": max(t_v, vis_aug), "ref_img": max(t_v, vis_aug),
+                 "ref_audio": max(t_a, aud_aug)}
+        unique_t = sorted({t_v, t_a} | ({seg_t["cond"]} if has_vis_cond else set())
+                          | ({seg_t["ref_audio"]} if has_aud_cond else set()))
+        t_row = {t: i for i, t in enumerate(unique_t)}
+        seg_tag = {"text": 1, "video": 0, "audio": 2, "cond": 0, "ref_img": 0, "ref_audio": 2}
+
+        text_tags = payload.get("text_token_tags")
+        mod_segments = []
+        for a, b, kind in layout.segments:
+            row_base = t_row[seg_t[kind]] * 3
+            if kind == "text" and text_tags is not None:
+                # the presentation text span mixes tags (vision pads carry the video modality) split into tag runs
+                tags = text_tags.view(-1).tolist()
+                run_start = 0
+                for i in range(1, b - a + 1):
+                    if i == b - a or tags[i] != tags[run_start]:
+                        mod_segments.append((a + run_start, a + i, row_base + int(tags[run_start])))
+                        run_start = i
+            else:
+                mod_segments.append((a, b, row_base + seg_tag[kind]))
+
+        # embed
+        img_update = layout.img_update.to(device)
+        audio_update = layout.audio_update.to(device)
+        video_rows = patchify_video(video_x.to(torch.float32), self.patch_size)
+        audio_rows = pack_audio(audio_x.to(torch.float32))
+        cond_video_rows = self._cond_video_rows(payload, device)
+        cond_audio_rows = self._cond_audio_rows(payload, device)
+
+        all_video_rows = video_rows
+        if cond_video_rows is not None:
+            all_video_rows = torch.empty(img_update.shape[0], video_rows.shape[1], dtype=torch.float32, device=device)
+            all_video_rows[~img_update] = cond_video_rows
+            all_video_rows[img_update] = video_rows
+        all_audio_rows = audio_rows
+        if cond_audio_rows is not None:
+            all_audio_rows = torch.empty(audio_update.shape[0], audio_rows.shape[1], dtype=torch.float32, device=device)
+            all_audio_rows[~audio_update] = cond_audio_rows
+            all_audio_rows[audio_update] = audio_rows
+
+        video_embed = self.video_patch_proj(all_video_rows).to(dtype)
+        audio_embed = self.audio_patch_proj(all_audio_rows).to(dtype)
+        text_states = context[0]
+        if text_states.shape[-1] != self.hidden_size:
+            text_states = self.token_refiner(self.condition_proj(text_states),
+                                             transformer_options=transformer_options)
+
+        # segments are contiguous: assemble by slices, embed rows follow segment order
+        h = torch.empty(layout.seq_len, self.hidden_size, dtype=dtype, device=device)
+        voff = aoff = 0
+        for a, b, kind in layout.segments:
+            n = b - a
+            if kind == "text":
+                h[a:b] = text_states
+            elif kind in ("cond", "ref_img", "video"):
+                h[a:b] = video_embed[voff:voff + n]
+                voff += n
+            else:  # ref_audio / audio
+                h[a:b] = audio_embed[aoff:aoff + n]
+                aoff += n
+
+        t_vals = torch.tensor(unique_t, dtype=torch.float32, device=device)
+        if self.use_adaln_curves:
+            # adaln projections consume interpolated coordinates of the time-embedding curve
+            table = comfy.model_management.cast_to(self.adaln_t_table, device=device)
+            pos = t_vals.clamp(0.0, 1.0) * (table.shape[0] - 1)     # t in [0,1] -> fractional grid index, out-of-range t clamps to the curve ends
+            i0 = pos.floor().long().clamp(max=table.shape[0] - 2)   # lower grid row, max-clamp keeps t=1.0 on the last interval instead of reading past the table
+            t_emb = torch.lerp(table[i0], table[i0 + 1], (pos - i0).unsqueeze(1))  # blend the two rows by the fractional part
+        else:
+            t_emb = self.time_embedder(t_vals).to(dtype)
+
+        # rotation table computed once per forward, consumed by the kitchen split-half rope
+        rope_freqs = rope_rotation_table(self.rope_freqs(layout.position_ids, device), dtype)
+
+        # Capture a single fully assembled payload before the first transformer block.
+        global H3_CAPTURE_ACTIVE
+        capture_dir = os.getenv("H3_CAPTURE_DIR")
+        if capture_dir and not H3_CAPTURE_ACTIVE and not os.path.exists(os.path.join(capture_dir, "blocks_complete")):
+            H3_CAPTURE_ACTIVE = True
+            os.makedirs(capture_dir, exist_ok=True)
+            torch.save({"hidden": h.detach().cpu(), "timesteps": t_vals.detach().cpu(), "position_ids": layout.position_ids, "segments": mod_segments}, os.path.join(capture_dir, "input.pt"))
+
+        # blocks
+        patches_replace = transformer_options.get("patches_replace", {})
+        blocks_replace = patches_replace.get("dit", {})
+        prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.blocks), device, transformer_options)
+        for i, block in enumerate(self.blocks):
+            block._h3_capture_index = i
+            block.attn._h3_capture_index = i
+            comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)
+            if ("double_block", i) in blocks_replace:
+                def block_wrap(args):
+                    return {"img": block(args["img"], args["t_emb"], args["mod_segments"], args["rope_freqs"],
+                                         transformer_options=args["transformer_options"])}
+                h = blocks_replace[("double_block", i)](
+                    {"img": h, "t_emb": t_emb, "mod_segments": mod_segments, "rope_freqs": rope_freqs,
+                     "transformer_options": transformer_options},
+                    {"original_block": block_wrap})["img"]
+            else:
+                h = block(h, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options)
+            if capture_dir and H3_CAPTURE_ACTIVE:
+                block_dir = os.path.join(capture_dir, "blocks")
+                os.makedirs(block_dir, exist_ok=True)
+                torch.save(h.detach().cpu(), os.path.join(block_dir, f"{i:02d}.pt"))
+        if capture_dir and H3_CAPTURE_ACTIVE:
+            open(os.path.join(capture_dir, "blocks_complete"), "a").close()
+        if prefetch_queue is not None:
+            comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, None)
+
+        # target streams are single contiguous segments (audio then video, last two)
+        video_seg = next((a, b, t_row[seg_t["video"]]) for a, b, k in layout.segments if k == "video")
+        audio_seg = next((a, b, t_row[seg_t["audio"]]) for a, b, k in layout.segments if k == "audio")
+        v, a = self.final_layer(h, t_emb, video_seg, audio_seg)
+
+        video_out = unpatchify_video(v, latent_t, lat_h // 2, lat_w // 2, self.latents_dim, self.patch_size)
+        video_out = video_out[:, :, :orig_t, :orig_h, :orig_w]
+        audio_out = unpack_audio(a)
+
+        if capture_dir and H3_CAPTURE_ACTIVE:
+            torch.save({"video": video_out.detach().cpu(), "audio": audio_out.detach().cpu(), "video_segment": video_seg, "audio_segment": audio_seg}, os.path.join(capture_dir, "output.pt"))
+            H3_CAPTURE_ACTIVE = False
+        return [-video_out.to(video_x.dtype), -audio_out.to(audio_x.dtype)]

src/h3_blackwell_runtime/upstream_nodes.py

@@ -0,0 +1,338 @@
+"""MiniMax H3 nodes: AV latent creation and task conditioning (t2va / fl2va / ref2va).
+
+The H3 packed-DiT consumes, via conditioning:
+- Qwen3-VL-32B hidden states with per-token modality tags (from the minimax CLIP)
+- keyframe / reference condition latents, re-injected every step (never denoised)
+
+Latents are NestedTensor pairs (video [B,24,T,H/16,W/16], audio [B,32,2,T40]);
+sampling runs on the flat pack with any stock sampler (the model handles the
+audio stream's shifted schedule internally).
+"""
+
+import math
+
+import torch
+import torchaudio
+
+import nodes
+import comfy.model_management
+import comfy.model_sampling
+import comfy.nested_tensor
+import comfy.utils
+import node_helpers
+from comfy_api.latest import ComfyExtension, io
+
+CANVAS_MULTIPLE = 32
+BASE_SHORT_EDGE = 768
+MAX_PIXELS = 768 * 1344
+REF_IMAGE_SHORT_EDGE = 2048
+FPS = 24
+AUDIO_LATENT_FPS = 40
+
+
+def align_frame_count(n):
+    while n % 17 != 5:
+        n += 1
+    return n
+
+
+def video_latent_t(frame_count):
+    return 2 if frame_count <= 5 else ((frame_count - 5) // 17) * 5 + 2
+
+
+def temporal_shape(length):
+    frame_count = align_frame_count(max(5, length))
+    duration = frame_count / FPS
+    return frame_count, video_latent_t(frame_count), round(duration * AUDIO_LATENT_FPS)
+
+
+def adapt_canvas(width, height):
+    """768-short-edge canvas with 768*1344 area cap, per-axis round to 32."""
+    ratio = width / height
+    if ratio >= 1.0:
+        nom_w, nom_h = BASE_SHORT_EDGE * ratio, BASE_SHORT_EDGE
+    else:
+        nom_w, nom_h = BASE_SHORT_EDGE, BASE_SHORT_EDGE / ratio
+    if nom_w * nom_h > MAX_PIXELS:
+        s = math.sqrt(MAX_PIXELS / (nom_w * nom_h))
+        nom_w, nom_h = nom_w * s, nom_h * s
+    return (max(CANVAS_MULTIPLE, round(nom_w / CANVAS_MULTIPLE) * CANVAS_MULTIPLE),
+            max(CANVAS_MULTIPLE, round(nom_h / CANVAS_MULTIPLE) * CANVAS_MULTIPLE))
+
+
+def _resize(image, width, height, crop):
+    # image [B, H, W, C] -> [B, height, width, 3]
+    samples = image[..., :3].movedim(-1, 1)
+    samples = comfy.utils.common_upscale(samples, width, height, "lanczos", crop)
+    return samples.movedim(1, -1)
+
+
+def _empty_av_latent(width, height, length, batch_size=1):
+    frame_count, latent_t, audio_t = temporal_shape(length)
+    video = torch.zeros([batch_size, 24, latent_t, height // 16, width // 16],
+                        device=comfy.model_management.intermediate_device())
+    audio = torch.zeros([batch_size, 32, 2, audio_t],
+                        device=comfy.model_management.intermediate_device())
+    return {"samples": comfy.nested_tensor.NestedTensor((video, audio))}, frame_count
+
+
+class EmptyMiniMaxH3LatentAV(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="EmptyMiniMaxH3LatentAV",
+            display_name="Empty MiniMax H3 AV Latent",
+            category="model/latent/minimax",
+            description="Joint video+audio latent for MiniMax H3. Duration snaps to the model's 17k+5 frame grid at 24 fps.",
+            inputs=[
+                io.Int.Input("width", default=1344, min=32, max=nodes.MAX_RESOLUTION, step=32),
+                io.Int.Input("height", default=768, min=32, max=nodes.MAX_RESOLUTION, step=32),
+                io.Int.Input("length", default=124, min=5, max=3600, step=17, tooltip="Frame count at 24 fps, snapped up to the model's 17k+5 grid (124 = ~5s; trained range is ~124-362, longer is untested)"),
+            ],
+            outputs=[io.Latent.Output()],
+        )
+
+    @classmethod
+    def execute(cls, width, height, length) -> io.NodeOutput:
+        latent, _ = _empty_av_latent(width, height, length)
+        return io.NodeOutput(latent)
+
+
+class MiniMaxH3ImageToVideo(io.ComfyNode):
+    """t2va and fl2va: prompt (+ optional first/last keyframes) -> conditioning + AV latent."""
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="MiniMaxH3ImageToVideo",
+            display_name="MiniMax H3 Image to Video",
+            category="model/conditioning/minimax",
+            inputs=[
+                io.Clip.Input("clip"),
+                io.Vae.Input("vae"),
+                io.String.Input("prompt", multiline=True, dynamic_prompts=True),
+                io.Int.Input("width", default=1344, min=32, max=nodes.MAX_RESOLUTION, step=32),
+                io.Int.Input("height", default=768, min=32, max=nodes.MAX_RESOLUTION, step=32),
+                io.Int.Input("length", default=124, min=5, max=3600, step=17, tooltip="Frame count at 24 fps, snapped up to the model's 17k+5 grid (124 = ~5s; trained range is ~124-362, longer is untested)"),
+                io.Image.Input("first_frame", optional=True),
+                io.Image.Input("last_frame", optional=True),
+            ],
+            outputs=[io.Conditioning.Output(display_name="positive"), io.Latent.Output()],
+        )
+
+    @classmethod
+    def execute(cls, clip, vae, prompt, width, height, length,
+                first_frame=None, last_frame=None) -> io.NodeOutput:
+        latent, frame_count = _empty_av_latent(width, height, length)
+
+        images = []
+        keyframes = []
+        if first_frame is not None:
+            # geometry anchor: plain stretch to canvas
+            img = _resize(first_frame[:1], width, height, "disabled")
+            images.append(img)
+            keyframes.append({"resolved_frame_index": 0, "image": img})
+        if last_frame is not None:
+            # follower: aspect-preserving cover-crop
+            img = _resize(last_frame[:1], width, height, "center")
+            images.append(img)
+            keyframes.append({"resolved_frame_index": frame_count - 1, "image": img})
+
+        tokens = clip.tokenize(prompt, images=images)
+        cond = clip.encode_from_tokens_scheduled(tokens)
+
+        if keyframes:
+            for kf in keyframes:
+                kf["latent"] = vae.encode(kf.pop("image"))
+            cond = node_helpers.conditioning_set_values(cond, {
+                "minimax_keyframes": keyframes,
+                "minimax_frame_count": frame_count,
+            })
+        return io.NodeOutput(cond, latent)
+
+
+class MiniMaxH3ReferenceToVideo(io.ComfyNode):
+    """ref2va: prompt + reference images / videos / audio -> conditioning + AV latent.
+
+    References enter the presentation in fixed order: images, then videos (each
+    soundtrack's <Audio j> label right before its <Video k>), then standalone
+    audio. Ordinals are 1-based per type, so the prompt refers to them as
+    <Picture i> / <Video k> / <Audio j>.
+    """
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="MiniMaxH3ReferenceToVideo",
+            description="<Picture i> / <Video k> / <Audio j> reference conditioning for MiniMax H3. Use the same tags when prompting.",
+            display_name="MiniMax H3 Reference to Video",
+            category="model/conditioning/minimax",
+            inputs=[
+                io.Clip.Input("clip"),
+                io.Vae.Input("vae"),
+                io.Vae.Input("audio_vae"),
+                io.String.Input("prompt", multiline=True, dynamic_prompts=True),
+                io.Int.Input("width", default=1344, min=32, max=nodes.MAX_RESOLUTION, step=32),
+                io.Int.Input("height", default=768, min=32, max=nodes.MAX_RESOLUTION, step=32),
+                io.Int.Input("length", default=124, min=5, max=3600, step=17, tooltip="Frame count at 24 fps, (124 = ~5s, trained range is ~124-362)"),
+                io.Combo.Input("ref_image_size", options=["match", "max"], default="match",
+                    tooltip="Reference image sizing. 'match' scales each ref (down only, keeping aspect) to the generation's pixel area; 'max' uses the reference pipeline's 2048px short edge for best identity fidelity. Reference tokens ride through every sampling step, so 'max' can be several times slower."),
+                io.Autogrow.Input("ref_images", optional=True,
+                    template=io.Autogrow.TemplatePrefix(
+                        input=io.Image.Input("ref_image", tooltip="Reference image (downscaled to 2048 short edge if larger, never upscaled)"),
+                        prefix="ref_image_", min=0, max=9)),
+                io.Autogrow.Input("ref_videos", optional=True,
+                    template=io.Autogrow.TemplatePrefix(
+                        input=io.Image.Input("ref_video", tooltip="Reference video frames at 24 fps (2-15s)"),
+                        prefix="ref_video_", min=0, max=3)),
+                io.Autogrow.Input("ref_video_audios", optional=True,
+                    template=io.Autogrow.TemplatePrefix(
+                        input=io.Audio.Input("ref_video_audio", tooltip="Soundtrack of the same-numbered reference video"),
+                        prefix="ref_video_audio_", min=0, max=3)),
+                io.Autogrow.Input("ref_audios", optional=True,
+                    template=io.Autogrow.TemplatePrefix(
+                        input=io.Audio.Input("ref_audio", tooltip="Standalone reference audio"),
+                        prefix="ref_audio_", min=0, max=3)),
+            ],
+            outputs=[io.Conditioning.Output(display_name="positive"), io.Latent.Output()],
+        )
+
+    @staticmethod
+    def _encode_ref_audio(audio_vae, audio):
+        waveform = audio["waveform"]  # [B, C, L]
+        sr = audio["sample_rate"]
+        vae_sr = getattr(audio_vae, "audio_sample_rate", 32000)
+        if sr != vae_sr:
+            waveform = torchaudio.functional.resample(waveform, sr, vae_sr)
+        z = audio_vae.encode(waveform[:1].movedim(1, -1))  # [1, 32, 2, T]
+        return z, z.shape[-1]
+
+    @classmethod
+    def execute(cls, clip, vae, audio_vae, prompt, width, height, length, ref_image_size="match",
+                ref_images=None, ref_videos=None, ref_video_audios=None, ref_audios=None) -> io.NodeOutput:
+        latent, frame_count = _empty_av_latent(width, height, length)
+
+        ref_items = []   # for the tokenizer presentation, in request order
+        ref_blocks = []  # for the DiT payload, same order
+
+        for img in (ref_images or {}).values():
+            if img is None:
+                continue
+            h, w = img.shape[1], img.shape[2]
+            if ref_image_size == "match":
+                # aspect-preserving scale (down only) to the generation's pixel area
+                scale = min(1.0, math.sqrt((width * height) / (w * h)))
+            else:
+                scale = min(1.0, REF_IMAGE_SHORT_EDGE / min(w, h))
+            tw = max(CANVAS_MULTIPLE, round(w * scale / CANVAS_MULTIPLE) * CANVAS_MULTIPLE)
+            th = max(CANVAS_MULTIPLE, round(h * scale / CANVAS_MULTIPLE) * CANVAS_MULTIPLE)
+            resized = _resize(img[:1], tw, th, "disabled")
+            z = vae.encode(resized)
+            ref_items.append({"type": "image", "data": resized})
+            ref_blocks.append({"kind": "image", "latent_h": th // 16, "latent_w": tw // 16, "latent": z})
+
+        ref_video_audios = ref_video_audios or {}
+        for name, video_frames in (ref_videos or {}).items():
+            if video_frames is None:
+                continue
+            # index-paired soundtrack: ref_video_audio_N belongs to ref_video_N
+            soundtrack = ref_video_audios.get("ref_video_audio_" + name.rsplit("_", 1)[-1])
+            vh, vw = video_frames.shape[1], video_frames.shape[2]
+            cw, ch = adapt_canvas(vw, vh)
+            if vw * vh < cw * ch:
+                cw = max(CANVAS_MULTIPLE, round(vw / CANVAS_MULTIPLE) * CANVAS_MULTIPLE)
+                ch = max(CANVAS_MULTIPLE, round(vh / CANVAS_MULTIPLE) * CANVAS_MULTIPLE)
+            frames = _resize(video_frames, cw, ch, "disabled")
+            if frames.shape[0] > frame_count:
+                frames = frames[:frame_count]
+            n = frames.shape[0]
+            if n < 5:
+                raise ValueError("MiniMax H3 reference videos need at least 5 frames (~0.2s at 24 fps)")
+            while n % 17 != 5:
+                n -= 1
+            frames = frames[:n]
+            z = vae.encode(frames)
+            audio_latent, ref_audio_t = (None, 0)
+            if soundtrack is not None:
+                audio_latent, ref_audio_t = cls._encode_ref_audio(audio_vae, soundtrack)
+                # the soundtrack gets its own <Audio j> label, emitted before <Video k>
+                ref_items.append({"type": "audio"})
+            # Qwen sees the video at 2 fps with timestamps
+            sample_idx = list(range(0, frames.shape[0], FPS // 2))
+            qwen_frames = frames[sample_idx]
+            ref_items.append({"type": "video", "data": qwen_frames,
+                              "timestamps": [i / 2.0 for i in range(len(sample_idx))]})
+            ref_blocks.append({"kind": "video_audio" if ref_audio_t else "video",
+                               "latent_t": z.shape[2], "latent_h": ch // 16, "latent_w": cw // 16,
+                               "ref_audio_t": ref_audio_t, "latent": z, "audio_latent": audio_latent})
+
+        for audio in (ref_audios or {}).values():
+            if audio is None:
+                continue
+            audio_latent, ref_audio_t = cls._encode_ref_audio(audio_vae, audio)
+            ref_items.append({"type": "audio"})
+            ref_blocks.append({"kind": "audio", "ref_audio_t": ref_audio_t, "audio_latent": audio_latent})
+
+        tokens = clip.tokenize(prompt, minimax_ref_items=ref_items)
+        cond = clip.encode_from_tokens_scheduled(tokens)
+        if ref_blocks:
+            cond = node_helpers.conditioning_set_values(cond, {"minimax_refs": ref_blocks})
+        return io.NodeOutput(cond, latent)
+
+
+class MiniMaxH3SigmaShift(io.ComfyNode):
+    """Set the video/audio flow shifts coherently.
+
+    The video shift drives the sampler's sigma schedule (ModelSamplingAV); both
+    values are also handed to the DiT, which inverts the video schedule to the
+    shared base grid and derives the audio schedule from it.
+    """
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="MiniMaxH3SigmaShift",
+            description="Set the video/audio flow shifts.",
+            display_name="ModelSamplingMiniMaxH3",
+            search_aliases=["sigma shift", "minimax shift"],
+            category="model/patch/minimax",
+            inputs=[
+                io.Model.Input("model"),
+                io.Float.Input("shift_video", default=12.0, min=0.01, max=100.0, step=0.01),
+                io.Float.Input("shift_audio", default=3.0, min=0.01, max=100.0, step=0.01),
+            ],
+            outputs=[io.Model.Output()],
+        )
+
+    @classmethod
+    def execute(cls, model, shift_video, shift_audio) -> io.NodeOutput:
+        m = model.clone()
+
+        class ModelSamplingAdvanced(comfy.model_sampling.ModelSamplingAV, comfy.model_sampling.CONST):
+            pass
+
+        original = m.get_model_object("model_sampling")
+        model_sampling = ModelSamplingAdvanced(model.model.model_config)
+        model_sampling.set_parameters(shift=shift_video, audio_shift=shift_audio)
+        if hasattr(original, "noise_scale"):
+            model_sampling.set_noise_scale(original.noise_scale)
+        m.add_object_patch("model_sampling", model_sampling)
+
+        to = m.model_options["transformer_options"] = m.model_options.get("transformer_options", {}).copy()
+        to["minimax_h3_sigma_shift_video"] = shift_video
+        to["minimax_h3_sigma_shift_audio"] = shift_audio
+        return io.NodeOutput(m)
+
+
+class MiniMaxH3Extension(ComfyExtension):
+    async def get_node_list(self):
+        return [
+            EmptyMiniMaxH3LatentAV,
+            MiniMaxH3ImageToVideo,
+            MiniMaxH3ReferenceToVideo,
+            MiniMaxH3SigmaShift,
+            ]
+
+
+async def comfy_entrypoint() -> MiniMaxH3Extension:
+    return MiniMaxH3Extension()

src/h3_blackwell_runtime/upstream_qwen3vl.py

@@ -0,0 +1,215 @@
+import os
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import Qwen2Tokenizer
+
+from comfy import sd1_clip
+import comfy.text_encoders.qwen_vl
+from .qwen35 import Qwen35VisionModel
+from .llama import BaseLlama, BaseQwen3, BaseGenerate, Llama2_, Qwen3VL_4BConfig, Qwen3VL_8BConfig, Qwen3VL_32BConfig
+
+
+QWEN3VL_VISION = {
+    "qwen3vl_4b": dict(hidden_size=1024, intermediate_size=4096, depth=24, deepstack_visual_indexes=[5, 11, 17]),
+    "qwen3vl_8b": dict(hidden_size=1152, intermediate_size=4304, depth=27, deepstack_visual_indexes=[8, 16, 24]),
+    "qwen3vl_32b": dict(hidden_size=1152, intermediate_size=4304, depth=27, deepstack_visual_indexes=[8, 16, 24]),
+}
+QWEN3VL_VISION_COMMON = dict(num_heads=16, patch_size=16, temporal_patch_size=2, in_channels=3,
+                             spatial_merge_size=2, num_position_embeddings=2304)
+
+QWEN3VL_CONFIGS = {"qwen3vl_4b": Qwen3VL_4BConfig, "qwen3vl_8b": Qwen3VL_8BConfig, "qwen3vl_32b": Qwen3VL_32BConfig}
+
+
+class Qwen3VLDeepstackMerger(nn.Module):
+    # DeepStack merger: postshuffle LayerNorm (applied after spatial merge), unlike the main merger.
+    def __init__(self, hidden_size, spatial_merge_size, out_hidden_size, device=None, dtype=None, ops=None):
+        super().__init__()
+        self.merge_dim = hidden_size * (spatial_merge_size ** 2)
+        self.norm = ops.LayerNorm(self.merge_dim, eps=1e-6, device=device, dtype=dtype)
+        self.linear_fc1 = ops.Linear(self.merge_dim, self.merge_dim, device=device, dtype=dtype)
+        self.linear_fc2 = ops.Linear(self.merge_dim, out_hidden_size, device=device, dtype=dtype)
+
+    def forward(self, x):
+        x = self.norm(x.view(-1, self.merge_dim))
+        return self.linear_fc2(F.gelu(self.linear_fc1(x)))
+
+
+class Qwen3VLVisionModel(Qwen35VisionModel):
+    # Qwen3.5 vision + DeepStack
+    def __init__(self, config, device=None, dtype=None, ops=None):
+        super().__init__(config, device=device, dtype=dtype, ops=ops)
+        self.deepstack_visual_indexes = config["deepstack_visual_indexes"]
+        self.deepstack_merger_list = nn.ModuleList([
+            Qwen3VLDeepstackMerger(self.hidden_size, self.spatial_merge_size, config["out_hidden_size"], device=device, dtype=dtype, ops=ops)
+            for _ in self.deepstack_visual_indexes
+        ])
+
+
+class Qwen3VL(BaseLlama, BaseQwen3, BaseGenerate, torch.nn.Module):
+    model_type = "qwen3vl_8b"
+
+    def __init__(self, config_dict, dtype, device, operations):
+        super().__init__()
+        config = QWEN3VL_CONFIGS[self.model_type](**config_dict)
+        self.num_layers = config.num_hidden_layers
+        self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
+        vision_config = {**QWEN3VL_VISION_COMMON, **QWEN3VL_VISION[self.model_type], "out_hidden_size": config.hidden_size}
+        self.visual = Qwen3VLVisionModel(vision_config, device=device, dtype=dtype, ops=operations)
+        self.dtype = dtype
+
+    def preprocess_embed(self, embed, device):
+        if embed["type"] == "image":
+            # Qwen3-VL normalizes to [-1, 1] (mean/std 0.5), unlike Qwen2.5-VL's CLIP normalization.
+            image, grid = comfy.text_encoders.qwen_vl.process_qwen2vl_images(embed["data"], patch_size=16, image_mean=[0.5, 0.5, 0.5], image_std=[0.5, 0.5, 0.5])
+            merged, deepstack = self.visual(image.to(device, dtype=torch.float32), grid)
+            return merged, {"grid": grid, "deepstack": deepstack}
+        return None, None
+
+    def build_image_inputs(self, embeds, embeds_info):
+        # Returns (position_ids, visual_pos_masks, deepstack) for the prompt
+        images = sorted([e for e in embeds_info if e.get("type") == "image"], key=lambda e: e["index"])
+        if len(images) == 0:
+            return None, None, None
+
+        device = embeds.device
+        seq = embeds.shape[1]
+        position_ids = comfy.text_encoders.qwen_vl.qwen2vl_mrope_position_ids(embeds_info, seq, device)
+
+        # DeepStack: mask of image positions + per-vision-layer features to inject there.
+        visual_pos_masks = torch.zeros((1, seq), dtype=torch.bool, device=device)
+        deepstack = None
+        for e in images:
+            start = e["index"]
+            end = e["size"] + start
+            visual_pos_masks[0, start:end] = True
+            ds = e["extra"]["deepstack"]
+            if deepstack is None:
+                deepstack = [d for d in ds]
+            else:
+                deepstack = [torch.cat([deepstack[i], ds[i]], dim=0) for i in range(len(ds))]
+        return position_ids, visual_pos_masks, deepstack
+
+    def forward(self, input_ids, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=None, embeds_info=[], **kwargs):
+        position_ids = kwargs.pop("position_ids", None)
+        visual_pos_masks = kwargs.pop("visual_pos_masks", None)
+        deepstack_embeds = kwargs.pop("deepstack_embeds", None)
+        if embeds is not None and position_ids is None:
+            position_ids, visual_pos_masks, deepstack_embeds = self.build_image_inputs(embeds, embeds_info)
+        return self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            embeds=embeds,
+            num_tokens=num_tokens,
+            intermediate_output=intermediate_output,
+            final_layer_norm_intermediate=final_layer_norm_intermediate,
+            dtype=dtype,
+            position_ids=position_ids,
+            embeds_info=embeds_info,
+            visual_pos_masks=visual_pos_masks,
+            deepstack_embeds=deepstack_embeds,
+            **kwargs,
+        )
+
+
+def _make_qwen3vl_model(model_type):
+    class Qwen3VL_(Qwen3VL):
+        pass
+    Qwen3VL_.model_type = model_type
+    return Qwen3VL_
+
+
+class Qwen3VLClipModel(sd1_clip.SDClipModel):
+    def __init__(self, device="cpu", layer="hidden", layer_idx=-1, dtype=None, attention_mask=True, model_options={}, model_type="qwen3vl_8b"):
+        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={},
+                         dtype=dtype, special_tokens={"pad": 151643}, layer_norm_hidden_state=False,
+                         model_class=_make_qwen3vl_model(model_type), enable_attention_masks=attention_mask,
+                         return_attention_masks=attention_mask, model_options=model_options)
+
+    def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty=0.0):
+        if isinstance(tokens, dict):
+            tokens = next(iter(tokens.values()))
+        tokens_only = [[t[0] for t in b] for b in tokens]
+        embeds, _, _, embeds_info = self.process_tokens(tokens_only, self.execution_device)
+        position_ids, visual_pos_masks, deepstack = self.transformer.build_image_inputs(embeds, embeds_info)
+        return self.transformer.generate(embeds, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed,
+                                         presence_penalty=presence_penalty, position_ids=position_ids,
+                                         visual_pos_masks=visual_pos_masks, deepstack_embeds=deepstack)
+
+
+class Qwen3VLTEModel(sd1_clip.SD1ClipModel):
+    def __init__(self, device="cpu", dtype=None, model_options={}, model_type="qwen3vl_8b"):
+        clip_model = lambda **kw: Qwen3VLClipModel(**kw, model_type=model_type)
+        super().__init__(device=device, dtype=dtype, name=model_type, clip_model=clip_model, model_options=model_options)
+
+
+class Qwen3VLSDTokenizer(sd1_clip.SDTokenizer):
+    def __init__(self, embedding_directory=None, tokenizer_data={}, embedding_size=4096, embedding_key="qwen3vl_8b"):
+        tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "qwen25_tokenizer")
+        super().__init__(tokenizer_path, pad_with_end=False, embedding_directory=embedding_directory, embedding_size=embedding_size, embedding_key=embedding_key, tokenizer_class=Qwen2Tokenizer,
+                         has_start_token=False, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=1, pad_token=151643, tokenizer_data=tokenizer_data)
+
+
+class Qwen3VLTokenizer(sd1_clip.SD1Tokenizer):
+    def __init__(self, embedding_directory=None, tokenizer_data={}, model_type="qwen3vl_8b"):
+        embedding_size = 2560 if model_type == "qwen3vl_4b" else 4096
+        tokenizer = lambda *a, **kw: Qwen3VLSDTokenizer(*a, **kw, embedding_size=embedding_size, embedding_key=model_type)
+        super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, name=model_type, tokenizer=tokenizer)
+        self.llama_template = "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+        self.llama_template_images = "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n"
+
+    def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=[], prevent_empty_text=False, thinking=False, skip_template=False, **kwargs):
+        image = kwargs.get("image", None)
+        if image is not None and len(images) == 0:
+            images = [image[i:i + 1] for i in range(image.shape[0])]
+
+        skip_template = skip_template or text.startswith('<|im_start|>')
+        if prevent_empty_text and text == '':
+            text = ' '
+
+        if skip_template:
+            llama_text = text
+        else:
+            if llama_template is not None:
+                template = llama_template
+            elif len(images) == 0:
+                template = self.llama_template
+            else:
+                template = self.llama_template_images
+                if len(images) > 1:
+                    vision_block = "<|vision_start|><|image_pad|><|vision_end|>"
+                    template = template.replace(vision_block, vision_block * len(images), 1)
+            llama_text = template.format(text)
+            if not thinking:  # Qwen3 convention: empty think block suppresses reasoning
+                llama_text += "<think>\n\n</think>\n\n"
+
+        tokens = super().tokenize_with_weights(llama_text, return_word_ids=return_word_ids, disable_weights=True, **kwargs)
+        key_name = next(iter(tokens))
+        embed_count = 0
+        for r in tokens[key_name]:
+            for i in range(len(r)):
+                if isinstance(r[i][0], (int, float)) and r[i][0] == 151655:  # <|image_pad|>
+                    if len(images) > embed_count:
+                        r[i] = ({"type": "image", "data": images[embed_count], "original_type": "image"},) + r[i][1:]
+                        embed_count += 1
+        return tokens
+
+
+def tokenizer(model_type="qwen3vl_8b"):
+    class Qwen3VLTokenizer_(Qwen3VLTokenizer):
+        def __init__(self, embedding_directory=None, tokenizer_data={}):
+            super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, model_type=model_type)
+    return Qwen3VLTokenizer_
+
+
+def te(dtype_llama=None, llama_quantization_metadata=None, model_type="qwen3vl_8b"):
+    class Qwen3VLTEModel_(Qwen3VLTEModel):
+        def __init__(self, device="cpu", dtype=None, model_options={}):
+            if dtype_llama is not None:
+                dtype = dtype_llama
+            if llama_quantization_metadata is not None:
+                model_options = model_options.copy()
+                model_options["quantization_metadata"] = llama_quantization_metadata
+            super().__init__(device=device, dtype=dtype, model_options=model_options, model_type=model_type)
+    return Qwen3VLTEModel_

src/h3_blackwell_runtime/upstream_text.py

@@ -0,0 +1,201 @@
+"""MiniMax H3 text/vision conditioning: Qwen3-VL-32B (truncated to 50 layers).
+
+The H3 presentation is NOT chat-templated: token ids are raw prompt/label text
+(no special tokens) with explicit vision blocks spliced in:
+
+  t2va:   <prompt>
+  fl2va:  "<Picture 1>: " <vision block> ["<Picture 2>: " <vision block>] <prompt>
+  ref2va: per condition in request order (1-based ordinals per type):
+            image -> "<Picture i>: " <vision block>
+            audio -> "<Audio j>: "              (audio never enters Qwen)
+            video -> "<Video k>: " then per 2-frame temporal block
+                     "<T.T seconds>" <vision block(2 frames)>
+          then <prompt>
+
+The conditioning is the unnormalized hidden state after LM layer 50 (the
+converted checkpoint is truncated there, so this is simply the last-layer
+output with no final norm). Vision-pad positions carry adaLN token tag 0
+(video modality) in the DiT; text positions carry tag 1 — the tags are
+returned alongside the embeddings as "minimax_token_tags".
+"""
+
+import math
+
+import torch
+import torch.nn.functional as F
+import comfy.sd1_clip
+from .qwen3vl import Qwen3VL, Qwen3VLSDTokenizer
+
+VISION_START = 151652
+VISION_END = 151653
+QWEN_IMAGE_MEAN = [0.5, 0.5, 0.5]
+QWEN_IMAGE_STD = [0.5, 0.5, 0.5]
+
+
+def process_video_block(frames, patch_size=16, temporal_patch_size=2, merge_size=2,
+                        min_pixels=3136, max_pixels=12845056):
+    """[2, H, W, C] frame pair -> (flatten_patches, grid_thw) with grid_t=1.
+
+    Same resize/normalize policy as process_qwen2vl_images, but the two frames
+    fill the temporal patch instead of repeating a single frame.
+    """
+    t, height, width, _ = frames.shape
+    imgs = frames.permute(0, 3, 1, 2)
+    factor = patch_size * merge_size
+    h_bar = round(height / factor) * factor
+    w_bar = round(width / factor) * factor
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = max(factor, math.floor(height / beta / factor) * factor)
+        w_bar = max(factor, math.floor(width / beta / factor) * factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = math.ceil(height * beta / factor) * factor
+        w_bar = math.ceil(width * beta / factor) * factor
+
+    imgs = F.interpolate(imgs, size=(h_bar, w_bar), mode="bilinear", align_corners=False)
+    mean = torch.tensor(QWEN_IMAGE_MEAN, device=imgs.device).view(1, 3, 1, 1)
+    std = torch.tensor(QWEN_IMAGE_STD, device=imgs.device).view(1, 3, 1, 1)
+    imgs = (imgs - mean) / std
+
+    grid_h = h_bar // patch_size
+    grid_w = w_bar // patch_size
+    patches = imgs.reshape(1, temporal_patch_size, 3, grid_h // merge_size, merge_size,
+                           patch_size, grid_w // merge_size, merge_size, patch_size)
+    patches = patches.permute(0, 3, 6, 4, 7, 2, 1, 5, 8)
+    flatten = patches.reshape(grid_h * grid_w, 3 * temporal_patch_size * patch_size * patch_size)
+    grid_thw = torch.stack([torch.tensor([1, grid_h, grid_w], device=frames.device, dtype=torch.long)])
+    return flatten, grid_thw
+
+
+def token_tags_from_embeds_info(seq_len, embeds_info):
+    # whole vision block VIDEO(0), including the flanking <|vision_start|>/<|vision_end|> tokens
+    # embeds_info spans cover only the expanded embeddings, so widen by one on each side.
+    tags = torch.ones(seq_len, dtype=torch.long)
+    for e in embeds_info:
+        if e.get("type") == "image":
+            tags[max(0, e["index"] - 1):e["index"] + e["size"] + 1] = 0
+    return tags
+
+
+class MiniMaxQwen3VL(Qwen3VL):
+    model_type = "qwen3vl_32b"
+
+    def preprocess_embed(self, embed, device):
+        if embed["type"] == "image" and embed.get("minimax_video_block", False):
+            flatten, grid = process_video_block(embed["data"])
+            merged, deepstack = self.visual(flatten.to(device, dtype=torch.float32), grid)
+            return merged, {"grid": grid, "deepstack": deepstack}
+        return super().preprocess_embed(embed, device)
+
+    def forward(self, input_ids, attention_mask=None, embeds=None, num_tokens=None,
+                intermediate_output=None, final_layer_norm_intermediate=True,
+                dtype=None, embeds_info=[], **kwargs):
+        seq = embeds.shape[1] if embeds is not None else input_ids.shape[1]
+        self.last_token_tags = token_tags_from_embeds_info(seq, embeds_info)
+        return super().forward(input_ids, attention_mask=attention_mask, embeds=embeds,
+                               num_tokens=num_tokens, intermediate_output=intermediate_output,
+                               final_layer_norm_intermediate=final_layer_norm_intermediate,
+                               dtype=dtype, embeds_info=embeds_info, **kwargs)
+
+
+class MiniMaxH3ClipModel(comfy.sd1_clip.SDClipModel):
+    def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None, model_options={}):
+        super().__init__(device=device, layer="last", layer_idx=None, textmodel_json_config={},
+                         dtype=dtype, special_tokens={"pad": 151643}, layer_norm_hidden_state=False,
+                         model_class=MiniMaxQwen3VL, enable_attention_masks=False,
+                         return_attention_masks=False, model_options=model_options)
+
+    def encode_token_weights(self, token_weight_pairs):
+        out = super().encode_token_weights(token_weight_pairs)
+        tags = getattr(self.transformer, "last_token_tags", None)
+        if tags is not None:
+            extra = out[2] if len(out) > 2 and isinstance(out[2], dict) else {}
+            extra["minimax_token_tags"] = tags
+            out = (out[0], out[1], extra)
+        return out
+
+
+class MiniMaxH3TEModel(comfy.sd1_clip.SD1ClipModel):
+    def __init__(self, device="cpu", dtype=None, model_options={}):
+        super().__init__(device=device, dtype=dtype, name="qwen3vl_32b",
+                         clip_model=MiniMaxH3ClipModel, model_options=model_options)
+
+
+class MiniMaxH3Tokenizer(comfy.sd1_clip.SD1Tokenizer):
+    def __init__(self, embedding_directory=None, tokenizer_data={}):
+        tokenizer = lambda *a, **kw: Qwen3VLSDTokenizer(*a, **kw, embedding_size=5120, embedding_key="qwen3vl_32b")
+        super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, name="qwen3vl_32b", tokenizer=tokenizer)
+
+    def _text_ids(self, text):
+        tok = self.qwen3vl_32b.tokenizer
+        return tok(text, add_special_tokens=False)["input_ids"]
+
+    @staticmethod
+    def _vision_entry(data, video_block=False):
+        emb = {"type": "image", "data": data, "original_type": "image"}
+        if video_block:
+            emb["minimax_video_block"] = True
+        return emb
+
+    def tokenize_with_weights(self, text, return_word_ids=False, images=[],
+                              minimax_ref_items=None, **kwargs):
+        entries = []
+
+        def add_text(s):
+            entries.extend((tid, 1.0) for tid in self._text_ids(s))
+
+        def add_vision(data, video_block=False):
+            entries.append((VISION_START, 1.0))
+            entries.append((self._vision_entry(data, video_block), 1.0))
+            entries.append((VISION_END, 1.0))
+
+        if minimax_ref_items:
+            counters = {"image": 0, "audio": 0, "video": 0}
+            for item in minimax_ref_items:
+                kind = item["type"]
+                counters[kind] += 1
+                if kind == "image":
+                    add_text("<Picture %d>: " % counters["image"])
+                    add_vision(item["data"])
+                elif kind == "audio":
+                    add_text("<Audio %d>: " % counters["audio"])
+                elif kind == "video":
+                    frames = item["data"]  # [T, H, W, C], sampled at 2 fps
+                    timestamps = item.get("timestamps")
+                    if timestamps is None:
+                        timestamps = [i / 2.0 for i in range(frames.shape[0])]
+                    if frames.shape[0] % 2 == 1:  # repeat-pad to temporal patch of 2
+                        frames = torch.cat([frames, frames[-1:]], dim=0)
+                        timestamps = list(timestamps) + [timestamps[-1]]
+                    add_text("<Video %d>: " % counters["video"])
+                    for i in range(0, frames.shape[0], 2):
+                        block_ts = (timestamps[i] + timestamps[i + 1]) / 2.0
+                        add_text("<%.1f seconds>" % block_ts)
+                        add_vision(frames[i:i + 2], video_block=True)
+        else:
+            for i, img in enumerate(images):
+                add_text("<Picture %d>: " % (i + 1))
+                add_vision(img)
+
+        add_text(text)
+        if len(entries) == 0:
+            entries.append((151643, 1.0))
+        if return_word_ids:
+            entries = [t + (0,) for t in entries]
+        return {"qwen3vl_32b": [entries]}
+
+    def untokenize(self, token_weight_pair):
+        return self.qwen3vl_32b.untokenize(token_weight_pair)
+
+
+def te(dtype_llama=None, llama_quantization_metadata=None, **kwargs):
+    class MiniMaxH3TEModel_(MiniMaxH3TEModel):
+        def __init__(self, device="cpu", dtype=None, model_options={}):
+            if dtype_llama is not None:
+                dtype = dtype_llama
+            if llama_quantization_metadata is not None:
+                model_options = model_options.copy()
+                model_options["quantization_metadata"] = llama_quantization_metadata
+            super().__init__(device=device, dtype=dtype, model_options=model_options)
+    return MiniMaxH3TEModel_

src/h3_blackwell_runtime/upstream_vae.py

@@ -0,0 +1,696 @@
+# MiniMax H3 video VAE: 3D causal CNN encoder + ViT3D decoder.
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import comfy.ops
+import comfy.quant_ops
+import comfy.rmsnorm
+from comfy.ldm.modules.attention import optimized_attention
+
+ops = comfy.ops.disable_weight_init
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+LATENTS_MEAN = [
+    0.858090341091156, -0.9606591463088989, 1.0661640167236328, -0.5090325474739075,
+    -0.2727581858634949, -1.3675414323806763, -0.2553254961967468, -0.26907554268836975,
+    -0.5376840829849243, -0.0464097298681736, 0.6657370328903198, 0.19690127670764923,
+    -0.5460608005523682, -0.4035342037677765, -0.23683024942874908, 0.25928452610969543,
+    -0.30133944749832153, 0.211341992020607, -1.1206848621368408, 0.3581933379173279,
+    -0.04225143790245056, 0.2604829967021942, 0.22864092886447906, 0.7056031823158264,
+]
+
+LATENTS_STD = [
+    1.2223774194717407, 1.2767263650894165, 1.68317747116088865, 1.7549455165863037,
+    1.5636216402053833, 2.194143533706665, 0.96531379222869875, 1.05698859691619875,
+    0.841948926448822, 0.7729952931404114, 1.8955937623977661, 0.946841835975647,
+    0.7996809482574463, 0.44988900423049925, 0.7197399735450745, 0.69362932443618775,
+    2.961095094680786, 2.7694199085235595, 3.0496184825897215, 2.1088054180145265,
+    3.276226282119751, 3.1627357006073, 2.28168129920959475, 2.6127843856811525,
+]
+
+
+# 3D causal CNN encoder
+
+class CausalConv3d(ops.Conv3d):
+    # Reflect spatial padding, causal (zeros, front-only) temporal padding.
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
+        super().__init__(in_channels, out_channels, kernel_size=kernel_size, stride=stride)
+        self.causal_padding = (padding,) * 3 if isinstance(padding, int) else tuple(padding)
+
+    def forward(self, x):
+        if sum(self.causal_padding) == 0:
+            return super().forward(x)
+
+        x = F.pad(x, (self.causal_padding[2], self.causal_padding[2], self.causal_padding[1], self.causal_padding[1], 0, 0),  mode="reflect")
+        if x.shape[2] == 1:
+            # single frame: the causal front padding is all zeros truncate the temporal taps instead of convolving zero frames
+            return super().forward(x, autopad="causal_zero")
+        x = F.pad(x, (0, 0, 0, 0, self.causal_padding[0] * 2, 0), mode="constant")
+        return super().forward(x)
+
+
+class TemporalIsolatedGroupNorm(ops.GroupNorm):
+    # GroupNorm with statistics computed per frame (time merged into batch).
+    def forward(self, x):
+        if x.dim() == 5:
+            b, c, t, h, w = x.shape
+            x = x.permute(0, 2, 1, 3, 4).contiguous().view(b * t, c, 1, h, w)
+            x = super().forward(x)
+            return x.view(b, t, c, h, w).permute(0, 2, 1, 3, 4).contiguous()
+        return super().forward(x)
+
+
+def group_norm_3d(num_channels):
+    return TemporalIsolatedGroupNorm(num_groups=32, num_channels=num_channels, eps=1e-6, affine=True)
+
+
+class Downsample3D(nn.Module):
+    def __init__(self, in_channels, out_channels, time_stride=1, space_stride=2):
+        super().__init__()
+        self.space_stride = space_stride
+        self.conv = CausalConv3d(
+            in_channels,
+            out_channels,
+            kernel_size=3,
+            padding=(1, 0, 0),
+            stride=(time_stride, space_stride, space_stride),
+        )
+
+    def forward(self, x):
+        if self.space_stride == 2:
+            x = F.pad(x, (0, 1, 0, 1, 0, 0), mode="reflect")
+        return self.conv(x)
+
+
+class ResnetBlock3D(nn.Module):
+    def __init__(self, in_channels, out_channels=None):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+
+        self.norm1 = group_norm_3d(in_channels)
+        self.norm2 = group_norm_3d(out_channels)
+        self.conv1 = CausalConv3d(in_channels, out_channels, kernel_size=3, padding=1)
+        self.conv2 = CausalConv3d(out_channels, out_channels, kernel_size=3, padding=1)
+        if in_channels != out_channels:
+            self.nin_shortcut = CausalConv3d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        h = self.conv1(F.silu(self.norm1(x), inplace=True))
+        h = self.conv2(F.silu(self.norm2(h), inplace=True))
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+        return h.add_(x)
+
+
+class EncoderFCN3D(nn.Module):
+    def __init__(self, ch, ch_mult, space_down, time_down, num_res_blocks, in_channels, z_channels, double_z=True):
+        super().__init__()
+        self.num_levels = len(ch_mult)
+        if isinstance(num_res_blocks, int):
+            num_res_blocks = [num_res_blocks] * self.num_levels
+        self.num_res_blocks = num_res_blocks
+
+        block_mid = [ch * ch_mult[i] for i in range(self.num_levels)]
+        block_in = [block_mid[0]] + block_mid[:-1]
+        block_out = block_mid
+
+        self.conv_in = CausalConv3d(in_channels, block_in[0], kernel_size=3, padding=1)
+
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_levels):
+            down = nn.Module()
+            down.block = nn.ModuleList()
+            for i in range(self.num_res_blocks[i_level]):
+                down.block.append(
+                    ResnetBlock3D(
+                        in_channels=block_in[i_level] if i == 0 else block_mid[i_level],
+                        out_channels=block_mid[i_level],
+                    )
+                )
+            if space_down[i_level] * time_down[i_level] > 1:
+                down.downsample = Downsample3D(
+                    block_mid[i_level],
+                    block_out[i_level],
+                    time_stride=time_down[i_level],
+                    space_stride=space_down[i_level],
+                )
+            self.down.append(down)
+
+        self.norm_out = group_norm_3d(block_out[-1])
+        self.conv_out = CausalConv3d(
+            block_out[-1],
+            2 * z_channels if double_z else z_channels,
+            kernel_size=3,
+            padding=1,
+        )
+
+    def forward(self, x):
+        h = self.conv_in(x)
+        for i_level in range(self.num_levels):
+            for i_block in range(self.num_res_blocks[i_level]):
+                h = self.down[i_level].block[i_block](h)
+            if hasattr(self.down[i_level], "downsample"):
+                h = self.down[i_level].downsample(h)
+        h = F.silu(self.norm_out(h))
+        return self.conv_out(h)
+
+
+# ViT3D decoder
+
+def create_token_ids(patch_dims, device, dtype):
+    coords_list = []
+    for dim_size in patch_dims:
+        coords = torch.arange(0.5, dim_size, dtype=dtype, device=device)
+        coords = coords / dim_size
+        coords = 2.0 * coords - 1.0
+        coords_list.append(coords)
+    coords = torch.stack(torch.meshgrid(*coords_list, indexing="ij"), dim=-1)
+    return coords.flatten(0, len(patch_dims) - 1).unsqueeze(0)
+
+
+class RotaryEmbeddingND(nn.Module):
+    def __init__(self, dim, rotary_base=100.0, n_dim=3):
+        super().__init__()
+        self.n_dim = n_dim
+        self.angle_scale = 2.0 * math.pi
+        inv_freq = 1 / rotary_base ** torch.arange(0, 1, 2 * n_dim / dim, dtype=torch.float32)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, img_ids):
+        # [B, S, n_dim] -> [B, S, 1, pairs, 2, 2] rotation table for the kitchen split-half rope
+        angles = (
+            self.angle_scale
+            * img_ids[:, :, :, None].float()
+            * self.inv_freq.to(img_ids.device)[None, None, None, :]
+        )
+        angles = angles.flatten(2, 3)
+        c, s = torch.cos(angles), torch.sin(angles)
+        table = torch.stack([c, -s, s, c], dim=-1).reshape(*angles.shape[:2], 1, angles.shape[-1], 2, 2)
+        return table.to(img_ids.dtype)
+
+
+class FeedForward(nn.Module):
+    # Gated SiLU FFN.
+    def __init__(self, dim, mult=4, bias=True, operations=ops):
+        super().__init__()
+        inner_dim = dim * mult
+        self.w1 = operations.Linear(dim, inner_dim * 2, bias=bias)
+        self.w2 = operations.Linear(inner_dim, dim, bias=bias)
+
+    def forward(self, x):
+        gate, x = self.w1(x).chunk(2, dim=-1)
+        return self.w2(F.silu(gate).mul_(x))
+
+
+class Attention(nn.Module):
+    def __init__(self, heads, dim_head, bias=True, eps=1e-5, operations=ops):
+        super().__init__()
+        self.dim_head = dim_head
+        self.heads = heads
+        inner_dim = dim_head * heads
+        self.norm_q = ops.RMSNorm(dim_head, eps=eps, elementwise_affine=False)
+        self.norm_k = ops.RMSNorm(dim_head, eps=eps, elementwise_affine=False)
+        self.to_qkv = operations.Linear(inner_dim, inner_dim * 3, bias=bias)
+        self.to_out = operations.Linear(inner_dim, inner_dim, bias=bias)
+
+    def forward(self, x, rotary_pos_emb=None):
+        batch_size, seq_len, _ = x.shape
+
+        qkv = self.to_qkv(x)
+        qkv = qkv.view(batch_size, seq_len, -1, 3 * self.dim_head)
+        query, key, value = torch.chunk(qkv, 3, dim=-1)
+
+        query = comfy.rmsnorm.rms_norm(query, self.norm_q.weight, self.norm_q.eps)
+        key = comfy.rmsnorm.rms_norm(key, self.norm_k.weight, self.norm_k.eps)
+
+        if rotary_pos_emb is not None:
+            rot = rotary_pos_emb.shape[-3] * 2
+            query[..., :rot], key[..., :rot] = comfy.quant_ops.ck.apply_rope_split_half(
+                query[..., :rot], key[..., :rot], rotary_pos_emb)
+
+        out = optimized_attention(query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2),
+                                  self.heads, skip_reshape=True).nan_to_num_(0.0)
+        return self.to_out(out)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, heads, dim_head, bias=True, eps=1e-5, operations=ops):
+        super().__init__()
+        dim = heads * dim_head
+        self.norm1 = ops.RMSNorm(dim, elementwise_affine=True, eps=eps)
+        self.attn = Attention(heads=heads, dim_head=dim_head, bias=bias, eps=eps, operations=operations)
+        self.scale1 = nn.Parameter(torch.empty(dim))
+        self.norm2 = ops.RMSNorm(dim, elementwise_affine=True, eps=eps)
+        self.ff = FeedForward(dim=dim, bias=bias, operations=operations)
+        self.scale2 = nn.Parameter(torch.empty(dim))
+
+    def forward(self, x, rotary_pos_emb=None):
+        x = x.addcmul_(self.attn(comfy.rmsnorm.rms_norm(x, self.norm1.weight, self.norm1.eps), rotary_pos_emb), comfy.ops.cast_to_input(self.scale1, x))
+        return x.addcmul_(self.ff(comfy.rmsnorm.rms_norm(x, self.norm2.weight, self.norm2.eps)), comfy.ops.cast_to_input(self.scale2, x))
+
+
+class ViT3DDecoder(nn.Module):
+    def __init__(self, patch_size=16, patch_size_t=4, in_channels=24, out_channels=3, num_layers=36, heads=32, dim_head=64, rope_theta=100.0,
+                 rope_dim_ratio=0.75, bias=True, eps=1e-5, num_register_tokens=4, operations=ops):
+        super().__init__()
+        dim = heads * dim_head
+        self.patch_size = patch_size
+        self.patch_size_t = patch_size_t
+        self.out_channels = out_channels
+        self.num_register_tokens = num_register_tokens
+
+        self.pos_embed = RotaryEmbeddingND(int(dim_head * rope_dim_ratio), rope_theta, n_dim=3)
+        self.x_embedder = ops.Linear(in_channels, dim)
+        self.register_tokens = nn.Parameter(torch.empty(1, num_register_tokens, dim))
+        # unused at inference; kept so the checkpoint loads without leftover keys
+        self.register_buffer("mask_token", torch.empty(1, 1, dim))
+
+        self.transformer_blocks = nn.ModuleList(
+            [TransformerBlock(heads=heads, dim_head=dim_head, bias=bias, eps=eps, operations=operations)
+             for _ in range(num_layers)]
+        )
+
+        self.norm_out = ops.LayerNorm(dim, elementwise_affine=True, eps=eps)
+        self.proj_out = ops.Linear(dim, out_channels * patch_size_t * patch_size * patch_size)
+
+    def forward(self, x):
+        B, C, latent_T, latent_H, latent_W = x.shape
+
+        h = self.x_embedder(x.flatten(2).transpose(1, 2))  # [B, T*H*W, C]
+
+        num_patches = h.shape[1]
+        num_suffix = 1 + self.num_register_tokens
+
+        h = torch.cat([h, comfy.ops.cast_to_input(self.register_tokens, h).expand(B, -1, -1), torch.zeros_like(h[:, 0:1, :])], dim=1)
+
+        img_ids = create_token_ids((latent_T, latent_H, latent_W), x.device, x.dtype).expand(B, -1, -1)
+        suffix_ids = torch.zeros((B, num_suffix, 3), device=x.device, dtype=img_ids.dtype)
+        img_ids = torch.cat([img_ids, suffix_ids], dim=1)
+
+        rotary_pos_emb = self.pos_embed(img_ids)
+
+        for block in self.transformer_blocks:
+            h = block(h, rotary_pos_emb)
+
+        output = self.proj_out(self.norm_out(h))
+
+        output = output[:, :num_patches, :]
+
+        output = output.view(
+            B, latent_T, latent_H, latent_W,
+            self.out_channels, self.patch_size_t, self.patch_size, self.patch_size,
+        )
+        output = output.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous()
+        output = output.reshape(
+            B, self.out_channels,
+            latent_T * self.patch_size_t,
+            latent_H * self.patch_size,
+            latent_W * self.patch_size,
+        )
+        return output
+
+
+# Full VAE
+
+class MiniMaxH3VideoVAE(nn.Module):
+    def __init__(
+        self,
+        in_channels=3,
+        out_ch=3,
+        ch=128,
+        embed_dim=24,
+        z_channels=24,
+        ch_mult=(1, 2, 2, 4, 4, 8),
+        num_res_blocks=2,
+        space_down=(2, 2, 2, 2, 1, 1),
+        time_down=(1, 2, 2, 1, 1, 1),
+        clip_length=17,
+        token_drop=3,
+        tile_size=256,
+        tile_overlap_min=64,
+        tiling=True,
+        operations=ops,
+    ):
+        super().__init__()
+        self.vae_ratio = int(math.prod(space_down))
+        self.vae_ratio_t = int(math.prod(time_down))
+
+        # temporal chunking parameters
+        self.clip_length = clip_length
+        self.token_drop = token_drop
+        self.frame_pre_padding = (-clip_length) % self.vae_ratio_t
+        self.tokens_chunk_size = math.ceil(clip_length / self.vae_ratio_t)
+        self.token_overlap = (-token_drop) % self.tokens_chunk_size
+        self.frame_overlap = max(self.token_overlap * self.vae_ratio_t - self.frame_pre_padding, 0)
+
+        # spatial tiling parameters
+        self.tiling = tiling
+        self.tile_size = tile_size
+        self.tile_overlap_min = tile_overlap_min
+
+        self.encoder = EncoderFCN3D(
+            ch=ch,
+            ch_mult=list(ch_mult),
+            space_down=list(space_down),
+            time_down=list(time_down),
+            num_res_blocks=num_res_blocks,
+            in_channels=in_channels,
+            z_channels=z_channels,
+            double_z=True,
+        )
+        self.quant_conv = ops.Conv3d(z_channels * 2, 2 * embed_dim, 1)
+        self.post_quant_conv = ops.Conv3d(embed_dim, z_channels, 1)
+        self.decoder = ViT3DDecoder(
+            patch_size=self.vae_ratio,
+            patch_size_t=self.vae_ratio_t,
+            in_channels=z_channels,
+            out_channels=out_ch,
+            operations=operations,
+        )
+
+        self.register_buffer("latents_mean", torch.tensor(LATENTS_MEAN))
+        self.register_buffer("latents_std", torch.tensor(LATENTS_STD))
+        self.register_buffer("pixel_mean", torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1, 1), persistent=False)
+        self.register_buffer("pixel_std", torch.tensor(IMAGENET_STD).view(1, 3, 1, 1, 1), persistent=False)
+
+    # single-shot forward
+
+    def _encode_moments(self, x):
+        return self.quant_conv(self.encoder(x))
+
+    def _decode_pixels(self, z):
+        return self.decoder(self.post_quant_conv(z))
+
+    def _adaptive_encode(self, x):
+        if self.tiling:
+            return self.tiled_encode(x)
+        return self._encode_moments(x)
+
+    def _adaptive_decode(self, z):
+        if self.tiling:
+            return self.tiled_decode(z)
+        return self._decode_pixels(z)
+
+    # spatial tiling
+
+    def split_tiles(self, input_len):
+        tile_size = self.tile_size
+        if tile_size >= input_len:
+            return [0], [input_len], []
+
+        N = math.ceil(input_len / tile_size)
+        while True:
+            overlaps = [self.tile_overlap_min] * (N - 1)
+            remaining = tile_size * N - sum(overlaps) - input_len
+            if remaining < 0:
+                N += 1
+            else:
+                break
+
+        remaining_units = remaining // self.vae_ratio
+        for i in range(remaining_units):
+            overlaps[i % (N - 1)] += self.vae_ratio
+
+        tile_start_idx = [0]
+        for i in range(N - 1):
+            tile_start_idx.append(tile_start_idx[-1] + tile_size - overlaps[i])
+
+        return tile_start_idx, [tile_size] * N, overlaps
+
+    def blend(self, a, b, blend_extent, dim):
+        blend_extent = min(a.shape[dim], b.shape[dim], blend_extent)
+
+        positions = torch.arange(blend_extent, device=b.device, dtype=b.dtype)
+        weight_a = 1 - positions / blend_extent
+        weight_b = positions / blend_extent
+
+        shape = [1] * a.ndim
+        shape[dim] = blend_extent
+        weight_a = weight_a.view(shape)
+        weight_b = weight_b.view(shape)
+
+        slice_a = [slice(None)] * a.ndim
+        slice_a[dim] = slice(-blend_extent, None)
+        slice_b = [slice(None)] * b.ndim
+        slice_b[dim] = slice(0, blend_extent)
+
+        blended = a[tuple(slice_a)] * weight_a + b[tuple(slice_b)] * weight_b
+
+        if blend_extent < b.shape[dim]:
+            slice_b_rest = [slice(None)] * b.ndim
+            slice_b_rest[dim] = slice(blend_extent, None)
+            return torch.cat([blended, b[tuple(slice_b_rest)]], dim=dim)
+        return blended
+
+    def tiled_encode(self, x):
+        height, width = x.shape[-2], x.shape[-1]
+        y_idx, y_len, y_overlap = self.split_tiles(height)
+        x_idx, x_len, x_overlap = self.split_tiles(width)
+
+        rows = []
+        for i_pos, i_len in zip(y_idx, y_len):
+            row = []
+            for j_pos, j_len in zip(x_idx, x_len):
+                tile = x[..., i_pos:i_pos + i_len, j_pos:j_pos + j_len]
+                row.append(self._encode_moments(tile))
+            rows.append(row)
+
+        latent_y_overlap = [o // self.vae_ratio for o in y_overlap]
+        latent_x_overlap = [o // self.vae_ratio for o in x_overlap]
+
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                if i > 0:
+                    tile = self.blend(rows[i - 1][j], tile, latent_y_overlap[i - 1], dim=-2)
+                if j > 0:
+                    tile = self.blend(row[j - 1], tile, latent_x_overlap[j - 1], dim=-1)
+                if i < len(rows) - 1:
+                    tile = tile[..., :-latent_y_overlap[i], :]
+                if j < len(row) - 1:
+                    tile = tile[..., :, :-latent_x_overlap[j]]
+                result_row.append(tile)
+            result_rows.append(torch.cat(result_row, dim=-1))
+        return torch.cat(result_rows, dim=-2)
+
+    def tiled_decode(self, z):
+        height, width = z.shape[-2] * self.vae_ratio, z.shape[-1] * self.vae_ratio
+        y_idx, y_len, y_overlap = self.split_tiles(height)
+        x_idx, x_len, x_overlap = self.split_tiles(width)
+
+        # Blended tiles are written straight into a pre-allocated canvas.
+        canvas = None
+        row_tails = []
+        out_y = 0
+        for i, (i_pos, i_len) in enumerate(zip(y_idx, y_len)):
+            zi, zl = i_pos // self.vae_ratio, i_len // self.vae_ratio
+            new_tails = []
+            left_tail = None
+            out_x = 0
+            for j, (j_pos, j_len) in enumerate(zip(x_idx, x_len)):
+                zj, zw = j_pos // self.vae_ratio, j_len // self.vae_ratio
+                tile = self._decode_pixels(z[..., zi:zi + zl, zj:zj + zw])
+                if i < len(y_idx) - 1:
+                    new_tails.append(tile[..., -y_overlap[i]:, :].clone())
+                next_left_tail = tile[..., :, -x_overlap[j]:].clone() if j < len(x_idx) - 1 else None
+                if i > 0:
+                    tile = self.blend(row_tails[j], tile, y_overlap[i - 1], dim=-2)
+                if j > 0:
+                    tile = self.blend(left_tail, tile, x_overlap[j - 1], dim=-1)
+                left_tail = next_left_tail
+                if i < len(y_idx) - 1:
+                    tile = tile[..., :-y_overlap[i], :]
+                if j < len(x_idx) - 1:
+                    tile = tile[..., :, :-x_overlap[j]]
+                if canvas is None:
+                    canvas = torch.empty(*tile.shape[:-2], height, width, dtype=tile.dtype, device=tile.device)
+                canvas[..., out_y:out_y + tile.shape[-2], out_x:out_x + tile.shape[-1]].copy_(tile)
+                out_x += tile.shape[-1]
+            row_tails = new_tails
+            out_y += tile.shape[-2]
+        return canvas
+
+    # temporal chunking
+
+    def encode_temporal(self, x):
+        if x.shape[2] % self.clip_length != 0:
+            pad_size = (-x.shape[2]) % self.clip_length
+            pad_frames = x[:, :, -1:].repeat(1, 1, pad_size, 1, 1)
+            x = torch.cat([x, pad_frames], dim=2)
+
+        num_chunks = x.shape[2] // self.clip_length
+
+        z_list = []
+        for i in range(num_chunks):
+            clip_x = x[:, :, i * self.clip_length:(i + 1) * self.clip_length, :, :]
+            z_list.append(self._adaptive_encode(clip_x))
+
+        z = torch.cat(z_list, dim=2)
+        if self.token_drop > 0:
+            z = z[:, :, :-self.token_drop]
+        return z
+
+    def _decode_temporal_pad_frames(self, z_len, pad_tokens):
+        if pad_tokens <= 0:
+            return 0
+        intra_tail = self.clip_length % self.vae_ratio_t
+        if intra_tail == 0:
+            return pad_tokens * self.vae_ratio_t
+
+        z_len_before_pad = z_len - pad_tokens
+        return sum(
+            (intra_tail if (z_len_before_pad + k) % self.tokens_chunk_size == 0
+             else self.vae_ratio_t)
+            for k in range(pad_tokens)
+        )
+
+    def _decode_temporal_frame_plan(self, z_len, num_chunks, pad_tokens):
+        chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
+        split_count = int(self.token_drop > 0) + 1
+        total_frames = 0
+        final_overlap_frames = 0
+
+        for i in range(num_chunks):
+            t_start_idx = i * self.tokens_chunk_size
+            t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap
+            clip_token_len = max(0, min(t_end_idx, z_len) - min(t_start_idx, z_len))
+            clip_frame_len = clip_token_len * self.vae_ratio_t
+
+            for j in range(split_count):
+                f_start_idx = j * chunk_dec
+                f_end_idx = min(f_start_idx + chunk_dec, clip_frame_len)
+                chunk_frames = max(0, f_end_idx - f_start_idx - self.frame_pre_padding)
+                if j == 0:
+                    total_frames += chunk_frames
+                else:
+                    final_overlap_frames = chunk_frames
+
+        total_frames += final_overlap_frames
+        return total_frames - self._decode_temporal_pad_frames(z_len, pad_tokens)
+
+    def decode_temporal(self, z):
+        chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
+        split_count = int(self.token_drop > 0) + 1
+
+        pseudo_total_tokens = z.shape[2] + self.token_drop
+
+        pad_tokens = 0
+        remainder = pseudo_total_tokens % self.tokens_chunk_size
+        if remainder != 0:
+            pad_tokens = self.tokens_chunk_size - remainder
+            pseudo_total_tokens += pad_tokens
+
+        num_chunks = pseudo_total_tokens // self.tokens_chunk_size - int(self.token_drop > 0)
+        if num_chunks < 1:
+            # too few tokens for one chunk (e.g. T_lat == 2): pad one extra chunk
+            pad_tokens += self.tokens_chunk_size
+            num_chunks += 1
+
+        if pad_tokens > 0:
+            pad_z = z[:, :, -1:, :, :].repeat(1, 1, pad_tokens, 1, 1)
+            z = torch.cat([z, pad_z], dim=2)
+
+        output_frames = self._decode_temporal_frame_plan(z.shape[2], num_chunks, pad_tokens)
+
+        dec = None
+        dec_overlap = None
+        write_pos = 0
+
+        def write_part(part):
+            nonlocal dec, write_pos
+            part_frames = part.shape[2]
+            if part_frames <= 0:
+                return
+            if dec is None:
+                out_shape = list(part.shape)
+                out_shape[2] = output_frames
+                dec = torch.empty(out_shape, dtype=part.dtype, device=part.device)
+            copy_frames = min(part_frames, max(0, dec.shape[2] - write_pos))
+            if copy_frames > 0:
+                dec[:, :, write_pos:write_pos + copy_frames, :, :].copy_(
+                    part[:, :, :copy_frames, :, :]
+                )
+                write_pos += copy_frames
+
+        for i in range(num_chunks):
+            t_start_idx = i * self.tokens_chunk_size
+            t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap
+            clip_z = z[:, :, t_start_idx:t_end_idx, :, :]
+
+            clip_dec = self._adaptive_decode(clip_z)
+
+            for j in range(split_count):
+                f_start_idx = j * chunk_dec
+                f_end_idx = min(f_start_idx + chunk_dec, clip_dec.shape[2])
+                clip_dec_chunk = clip_dec[:, :, f_start_idx:f_end_idx, :, :]
+                clip_dec_chunk = clip_dec_chunk[:, :, self.frame_pre_padding:, :, :]
+
+                if j == 0:
+                    if dec_overlap is not None:
+                        clip_dec_chunk = self.blend(
+                            dec_overlap, clip_dec_chunk, self.frame_overlap, dim=-3
+                        )
+                        dec_overlap = None
+                    write_part(clip_dec_chunk)
+                else:
+                    dec_overlap = clip_dec_chunk.contiguous()
+
+            if i == num_chunks - 1 and dec_overlap is not None:
+                write_part(dec_overlap)
+                dec_overlap = None
+
+            del clip_dec, clip_z
+
+        return dec
+
+
+    def encode(self, x):
+        # x: [B, 3, T, H, W] in [-1, 1] -> normalized latents [B, 24, T_lat, H/16, W/16]
+        if x.ndim == 4:
+            x = x.unsqueeze(2)
+
+        x = x.add(1.0).mul_(0.5).sub_(self.pixel_mean.to(x)).div_(self.pixel_std.to(x))
+
+        if x.shape[2] == 1:
+            moments = self._adaptive_encode(x)
+            moments = moments[:, :, -1:, :, :]
+        else:
+            moments = self.encode_temporal(x)
+
+        mean = torch.chunk(moments.float(), 2, dim=1)[0]
+
+        latents_mean = self.latents_mean.view(1, -1, 1, 1, 1).to(mean)
+        latents_std = self.latents_std.view(1, -1, 1, 1, 1).to(mean)
+        return (mean - latents_mean) / latents_std
+
+    def encode_tiled(self, x, **kwargs):
+        # tiling is always on internally with the reference's semantic tile sizes, ignore tiling fallbacks
+        return self.encode(x)
+
+    def decode_tiled(self, z, **kwargs):
+        return self.decode(z)
+
+    def decode(self, z):
+        # z: [B, 24, T_lat, H_lat, W_lat] normalized latents -> pixels [B, 3, T, H, W] in [-1, 1]
+        latents_mean = self.latents_mean.view(1, -1, 1, 1, 1).to(z)
+        latents_std = self.latents_std.view(1, -1, 1, 1, 1).to(z)
+        z = z * latents_std + latents_mean
+
+        if z.shape[2] == 1:
+            dec = self._adaptive_decode(z)
+            dec = dec[:, :, -1:, :, :]
+        else:
+            dec = self.decode_temporal(z)
+
+        dec = dec.float()
+        dec.mul_(self.pixel_std.to(dec)).add_(self.pixel_mean.to(dec)).clamp_(0.0, 1.0).mul_(2.0).sub_(1.0)
+        return dec

src/h3_blackwell_runtime/vae_decoder.py

@@ -0,0 +1,281 @@
+"""Direct, decoder-only MiniMax H3 video VAE implementation."""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+from torch import nn
+from torch.nn import functional as F
+
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+LATENTS_MEAN = (0.858090341091156, -0.9606591463088989, 1.0661640167236328, -0.5090325474739075, -0.2727581858634949, -1.3675414323806763, -0.2553254961967468, -0.26907554268836975, -0.5376840829849243, -0.0464097298681736, 0.6657370328903198, 0.19690127670764923, -0.5460608005523682, -0.4035342037677765, -0.23683024942874908, 0.25928452610969543, -0.30133944749832153, 0.211341992020607, -1.1206848621368408, 0.3581933379173279, -0.04225143790245056, 0.2604829967021942, 0.22864092886447906, 0.7056031823158264)
+LATENTS_STD = (1.2223774194717407, 1.2767263650894165, 1.68317747116088865, 1.7549455165863037, 1.5636216402053833, 2.194143533706665, 0.96531379222869875, 1.05698859691619875, 0.841948926448822, 0.7729952931404114, 1.8955937623977661, 0.946841835975647, 0.7996809482574463, 0.44988900423049925, 0.7197399735450745, 0.69362932443618775, 2.961095094680786, 2.7694199085235595, 3.0496184825897215, 2.1088054180145265, 3.276226282119751, 3.1627357006073, 2.28168129920959475, 2.6127843856811525)
+
+
+def _rms_norm(x: torch.Tensor, weight: torch.Tensor | None, eps: float) -> torch.Tensor:
+    result = x * torch.rsqrt(x.float().square().mean(dim=-1, keepdim=True) + eps).to(x.dtype)
+    return result if weight is None else result * weight.to(dtype=x.dtype)
+
+
+def create_token_ids(patch_dims: tuple[int, int, int], device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+    coords = [2.0 * (torch.arange(0.5, size, dtype=dtype, device=device) / size) - 1.0 for size in patch_dims]
+    return torch.stack(torch.meshgrid(*coords, indexing="ij"), dim=-1).flatten(0, 2).unsqueeze(0)
+
+
+class RotaryEmbeddingND(nn.Module):
+    def __init__(self, dim: int, rotary_base: float = 100.0, n_dim: int = 3, *, device=None):
+        super().__init__()
+        self.rotary_base = rotary_base
+        self.step = 2 * n_dim / dim
+        inv_freq = 1 / rotary_base ** torch.arange(0, 1, self.step, dtype=torch.float32, device=device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.angle_scale = 2.0 * math.pi
+
+    def forward(self, img_ids: torch.Tensor) -> torch.Tensor:
+        inv_freq = self.inv_freq
+        if inv_freq.device.type == "meta":
+            inv_freq = 1 / self.rotary_base ** torch.arange(0, 1, self.step, dtype=torch.float32, device=img_ids.device)
+        else:
+            inv_freq = inv_freq.to(img_ids.device)
+        angles = self.angle_scale * img_ids[:, :, :, None].float() * inv_freq[None, None, None, :]
+        angles = angles.flatten(2, 3)
+        cos, sin = torch.cos(angles), torch.sin(angles)
+        return torch.stack((cos, -sin, sin, cos), dim=-1).reshape(*angles.shape[:2], 1, -1, 2, 2).to(img_ids.dtype)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float, affine: bool, *, device=None):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.empty(dim, device=device)) if affine else None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return _rms_norm(x, self.weight, self.eps)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim: int, bias: bool = True, *, device=None):
+        super().__init__()
+        self.w1 = nn.Linear(dim, dim * 8, bias=bias, device=device)
+        self.w2 = nn.Linear(dim * 4, dim, bias=bias, device=device)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate, value = self.w1(x).chunk(2, dim=-1)
+        return self.w2(F.silu(gate) * value)
+
+
+def _apply_rope_split_half(x: torch.Tensor, table: torch.Tensor) -> torch.Tensor:
+    """Apply the reference split-half RoPE layout to leading rotary channels."""
+    pairs = table.shape[-3]
+    rot = pairs * 2
+    first, second = x[..., :pairs], x[..., pairs:rot]
+    cos, neg_sin, sin = table[..., 0, 0], table[..., 0, 1], table[..., 1, 0]
+    rotated = torch.cat((first * cos + second * neg_sin, first * sin + second * cos), dim=-1)
+    return torch.cat((rotated, x[..., rot:]), dim=-1)
+
+
+class Attention(nn.Module):
+    def __init__(self, heads: int, dim_head: int, bias: bool = True, eps: float = 1e-5, *, device=None):
+        super().__init__()
+        dim = heads * dim_head
+        self.heads, self.dim_head = heads, dim_head
+        self.norm_q = RMSNorm(dim_head, eps, False, device=device)
+        self.norm_k = RMSNorm(dim_head, eps, False, device=device)
+        self.to_qkv = nn.Linear(dim, dim * 3, bias=bias, device=device)
+        self.to_out = nn.Linear(dim, dim, bias=bias, device=device)
+
+    def forward(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor) -> torch.Tensor:
+        batch, sequence, _ = x.shape
+        qkv = self.to_qkv(x).view(batch, sequence, self.heads, 3 * self.dim_head)
+        query, key, value = qkv.chunk(3, dim=-1)
+        query, key = self.norm_q(query), self.norm_k(key)
+        query, key = _apply_rope_split_half(query, rotary_pos_emb), _apply_rope_split_half(key, rotary_pos_emb)
+        output = F.scaled_dot_product_attention(query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2))
+        return self.to_out(output.transpose(1, 2).reshape(batch, sequence, -1).nan_to_num_(0.0))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, heads: int, dim_head: int, bias: bool = True, eps: float = 1e-5, *, device=None):
+        super().__init__()
+        dim = heads * dim_head
+        self.norm1 = RMSNorm(dim, eps, True, device=device)
+        self.attn = Attention(heads, dim_head, bias, eps, device=device)
+        self.scale1 = nn.Parameter(torch.empty(dim, device=device))
+        self.norm2 = RMSNorm(dim, eps, True, device=device)
+        self.ff = FeedForward(dim, bias, device=device)
+        self.scale2 = nn.Parameter(torch.empty(dim, device=device))
+
+    def forward(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.norm1(x), rotary_pos_emb) * self.scale1.to(x.dtype)
+        return x + self.ff(self.norm2(x)) * self.scale2.to(x.dtype)
+
+
+class ViT3DDecoder(nn.Module):
+    def __init__(self, patch_size: int = 16, patch_size_t: int = 4, in_channels: int = 24, out_channels: int = 3, num_layers: int = 36, heads: int = 32, dim_head: int = 64, rope_theta: float = 100.0, rope_dim_ratio: float = 0.75, bias: bool = True, eps: float = 1e-5, num_register_tokens: int = 4, *, device=None):
+        super().__init__()
+        dim = heads * dim_head
+        self.patch_size, self.patch_size_t, self.out_channels = patch_size, patch_size_t, out_channels
+        self.num_register_tokens = num_register_tokens
+        self.pos_embed = RotaryEmbeddingND(int(dim_head * rope_dim_ratio), rope_theta, device=device)
+        self.x_embedder = nn.Linear(in_channels, dim, device=device)
+        self.register_tokens = nn.Parameter(torch.empty(1, num_register_tokens, dim, device=device))
+        self.register_buffer("mask_token", torch.empty(1, 1, dim, device=device))
+        self.transformer_blocks = nn.ModuleList([TransformerBlock(heads, dim_head, bias, eps, device=device) for _ in range(num_layers)])
+        self.norm_out = nn.LayerNorm(dim, eps=eps, device=device)
+        self.proj_out = nn.Linear(dim, out_channels * patch_size_t * patch_size * patch_size, device=device)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch, _, latent_t, latent_h, latent_w = x.shape
+        h = self.x_embedder(x.flatten(2).transpose(1, 2))
+        patches = h.shape[1]
+        h = torch.cat((h, self.register_tokens.to(h).expand(batch, -1, -1), torch.zeros_like(h[:, :1])), dim=1)
+        ids = create_token_ids((latent_t, latent_h, latent_w), x.device, x.dtype).expand(batch, -1, -1)
+        ids = torch.cat((ids, torch.zeros(batch, 1 + self.num_register_tokens, 3, device=x.device, dtype=x.dtype)), dim=1)
+        rope = self.pos_embed(ids)
+        for block in self.transformer_blocks:
+            h = block(h, rope)
+        output = self.proj_out(self.norm_out(h))[:, :patches]
+        output = output.view(batch, latent_t, latent_h, latent_w, self.out_channels, self.patch_size_t, self.patch_size, self.patch_size)
+        return output.permute(0, 4, 1, 5, 2, 6, 3, 7).reshape(batch, self.out_channels, latent_t * self.patch_size_t, latent_h * self.patch_size, latent_w * self.patch_size)
+
+
+class MiniMaxH3VideoVAE(nn.Module):
+    """Decoder-only H3 VAE. The public ``decode`` contract matches upstream_vae.py."""
+    def __init__(self, *, device=None, tiling: bool = True):
+        super().__init__()
+        self.vae_ratio, self.vae_ratio_t = 16, 4
+        self.clip_length, self.token_drop = 17, 3
+        self.tokens_chunk_size, self.token_overlap = 5, 3
+        self.frame_pre_padding, self.frame_overlap = 3, 9
+        self.tiling, self.tile_size, self.tile_overlap_min = tiling, 256, 64
+        self.post_quant_conv = nn.Conv3d(24, 24, 1, device=device)
+        self.decoder = ViT3DDecoder(device=device)
+        self.register_buffer("latents_mean", torch.tensor(LATENTS_MEAN, device=device))
+        self.register_buffer("latents_std", torch.tensor(LATENTS_STD, device=device))
+        self.register_buffer("pixel_mean", torch.tensor(IMAGENET_MEAN, device=device).view(1, 3, 1, 1, 1), persistent=False)
+        self.register_buffer("pixel_std", torch.tensor(IMAGENET_STD, device=device).view(1, 3, 1, 1, 1), persistent=False)
+
+    @classmethod
+    def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", tiling: bool = True) -> "MiniMaxH3VideoVAE":
+        model = cls(device="meta", tiling=tiling)
+        expected = model.state_dict()
+        with safe_open(str(path), framework="pt", device=str(device)) as checkpoint:
+            available = set(checkpoint.keys())
+            missing = sorted(set(expected) - available)
+            shape_errors = [(name, tuple(expected[name].shape), tuple(checkpoint.get_slice(name).get_shape())) for name in expected if name in available and tuple(expected[name].shape) != tuple(checkpoint.get_slice(name).get_shape())]
+            if missing or shape_errors:
+                details = ([f"missing: {', '.join(missing)}"] if missing else []) + ([f"shape mismatch: {shape_errors}"] if shape_errors else [])
+                raise ValueError("incompatible H3 VAE checkpoint; " + "; ".join(details))
+            weights = {name: checkpoint.get_tensor(name) for name in expected}
+        model.load_state_dict(weights, strict=True, assign=True)
+        return model
+
+    def _decode_pixels(self, z: torch.Tensor) -> torch.Tensor:
+        return self.decoder(self.post_quant_conv(z))
+
+    def split_tiles(self, length: int) -> tuple[list[int], list[int], list[int]]:
+        if self.tile_size >= length:
+            return [0], [length], []
+        count = math.ceil(length / self.tile_size)
+        while self.tile_size * count - self.tile_overlap_min * (count - 1) < length:
+            count += 1
+        overlaps = [self.tile_overlap_min] * (count - 1)
+        for index in range((self.tile_size * count - sum(overlaps) - length) // self.vae_ratio):
+            overlaps[index % len(overlaps)] += self.vae_ratio
+        starts = [0]
+        for overlap in overlaps:
+            starts.append(starts[-1] + self.tile_size - overlap)
+        return starts, [self.tile_size] * count, overlaps
+
+    @staticmethod
+    def blend(a: torch.Tensor, b: torch.Tensor, extent: int, dim: int) -> torch.Tensor:
+        extent = min(a.shape[dim], b.shape[dim], extent)
+        shape = [1] * a.ndim
+        shape[dim] = extent
+        position = torch.arange(extent, device=b.device, dtype=b.dtype).view(shape)
+        blended = a.narrow(dim, a.shape[dim] - extent, extent) * (1 - position / extent) + b.narrow(dim, 0, extent) * (position / extent)
+        return torch.cat((blended, b.narrow(dim, extent, b.shape[dim] - extent)), dim=dim)
+
+    def tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
+        y_starts, y_lengths, y_overlaps = self.split_tiles(z.shape[-2] * self.vae_ratio)
+        x_starts, x_lengths, x_overlaps = self.split_tiles(z.shape[-1] * self.vae_ratio)
+        rows: list[list[torch.Tensor]] = []
+        for y, height in zip(y_starts, y_lengths):
+            row = []
+            for x, width in zip(x_starts, x_lengths):
+                row.append(self._decode_pixels(z[..., y // 16:(y + height) // 16, x // 16:(x + width) // 16]))
+            rows.append(row)
+        result_rows = []
+        for row_index, row in enumerate(rows):
+            result_row = []
+            for index, tile in enumerate(row):
+                if row_index:
+                    tile = self.blend(rows[row_index - 1][index], tile, y_overlaps[row_index - 1], -2)
+                if result_row:
+                    tile = self.blend(row[index - 1], tile, x_overlaps[index - 1], -1)
+                if row_index < len(rows) - 1:
+                    tile = tile[..., :-y_overlaps[row_index], :]
+                if index < len(row) - 1:
+                    tile = tile[..., :, :-x_overlaps[index]]
+                result_row.append(tile)
+            result_rows.append(torch.cat(result_row, dim=-1))
+        return torch.cat(result_rows, dim=-2)
+
+    def _decode_temporal_pad_frames(self, z_len: int, pad_tokens: int) -> int:
+        if pad_tokens <= 0:
+            return 0
+        return sum(1 if (z_len - pad_tokens + index) % self.tokens_chunk_size == 0 else self.vae_ratio_t for index in range(pad_tokens))
+
+    def _decode_temporal_frame_plan(self, z_len: int, chunks: int, pad_tokens: int) -> int:
+        total, final_overlap = 0, 0
+        for index in range(chunks):
+            tokens = max(0, min(index * 5 + 8, z_len) - min(index * 5, z_len))
+            frames = tokens * self.vae_ratio_t
+            for split in range(2):
+                part = max(0, min((split + 1) * 20, frames) - split * 20 - self.frame_pre_padding)
+                if split == 0:
+                    total += part
+                else:
+                    final_overlap = part
+        return total + final_overlap - self._decode_temporal_pad_frames(z_len, pad_tokens)
+
+    def decode_temporal(self, z: torch.Tensor) -> torch.Tensor:
+        pseudo_tokens = z.shape[2] + self.token_drop
+        pad_tokens = (-pseudo_tokens) % self.tokens_chunk_size
+        pseudo_tokens += pad_tokens
+        chunks = pseudo_tokens // self.tokens_chunk_size - 1
+        if chunks < 1:
+            pad_tokens += self.tokens_chunk_size
+            chunks += 1
+        if pad_tokens:
+            z = torch.cat((z, z[:, :, -1:].expand(-1, -1, pad_tokens, -1, -1)), dim=2)
+        output_frames = self._decode_temporal_frame_plan(z.shape[2], chunks, pad_tokens)
+        output, overlap = [], None
+        for index in range(chunks):
+            clip = self._adaptive_decode(z[:, :, index * 5:index * 5 + 8])
+            first = clip[:, :, :20, :, :][:, :, self.frame_pre_padding:]
+            tail = clip[:, :, 20:40, :, :][:, :, self.frame_pre_padding:]
+            if overlap is not None:
+                first = self.blend(overlap, first, self.frame_overlap, -3)
+            output.append(first)
+            overlap = tail
+        if overlap is not None:
+            output.append(overlap)
+        return torch.cat(output, dim=2)[:, :, :output_frames]
+
+    def _adaptive_decode(self, z: torch.Tensor) -> torch.Tensor:
+        return self.tiled_decode(z) if self.tiling else self._decode_pixels(z)
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        z = z * self.latents_std.view(1, -1, 1, 1, 1).to(z) + self.latents_mean.view(1, -1, 1, 1, 1).to(z)
+        decoded = self._adaptive_decode(z) if z.shape[2] == 1 else self.decode_temporal(z)
+        if z.shape[2] == 1:
+            decoded = decoded[:, :, -1:]
+        pixel_std = torch.tensor(IMAGENET_STD, device=decoded.device, dtype=decoded.dtype).view(1, 3, 1, 1, 1)
+        pixel_mean = torch.tensor(IMAGENET_MEAN, device=decoded.device, dtype=decoded.dtype).view(1, 3, 1, 1, 1)
+        return (decoded.float() * pixel_std.float() + pixel_mean.float()).clamp_(0, 1).mul_(2).sub_(1)

tools/Inspect-H3Nvfp4Checkpoint.ps1

@@ -0,0 +1,27 @@
+$ErrorActionPreference = 'Stop'
+
+$envPath = Join-Path $PSScriptRoot '..\..\..\.env'
+Get-Content -LiteralPath $envPath | ForEach-Object {
+    if ($_ -match '^\s*([^#=\s]+)\s*=\s*(.*)\s*$') {
+        [Environment]::SetEnvironmentVariable($matches[1], $matches[2].Trim('"').Trim("'"), 'Process')
+    }
+}
+if (-not $env:RUNPOD_API_KEY) { throw 'RUNPOD_API_KEY is not set in .env.' }
+
+$checkpoint = '/runpod-volume/ComfyUI/models/diffusion_models/minimax_h3_ref2va_pruned_nvfp4.safetensors'
+$output = '/runpod-volume/h3-blackwell-runtime/artifacts/checkpoints/minimax_h3_ref2va_pruned_nvfp4.header.json'
+$probe = "import collections,json,os,struct; p='$checkpoint'; o='$output'; f=open(p,'rb'); n=struct.unpack('<Q',f.read(8))[0]; h=json.loads(f.read(n)); t={k:v for k,v in h.items() if k!='__metadata__'}; r={'checkpoint':p,'metadata':h.get('__metadata__',{}),'tensor_count':len(t),'dtypes':dict(sorted(collections.Counter(v['dtype'] for v in t.values()).items())),'tensors':t}; os.makedirs(os.path.dirname(o),exist_ok=True); open(o,'w').write(json.dumps(r,indent=2,sort_keys=True)); print(json.dumps({'output':o,'header_bytes':n,'tensor_count':len(t),'dtypes':r['dtypes'],'metadata':r['metadata']},sort_keys=True))"
+$encodedProbe = [Convert]::ToBase64String([Text.Encoding]::UTF8.GetBytes($probe))
+$payload = @{
+    name = 'h3-nvfp4-checkpoint-inspect'
+    image = 'runpod/pytorch:1.1.0-rc.154-cu1300-torch291-ubuntu2404'
+    cloud = 'SECURE'
+    gpu = @{ id = 'NVIDIA RTX PRO 6000 Blackwell Server Edition'; count = 1 }
+    disk = 20
+    dataCenterIds = @('EUR-IS-1')
+    mounts = @{ network = @(@{ volumeId = 'g8r2uqufz5'; path = '/runpod-volume' }) }
+    args = 'bash -lc ''set -eu; python -c "import base64;exec(base64.b64decode(\"' + $encodedProbe + '\"))"'''
+}
+
+$headers = @{ Authorization = "Bearer $env:RUNPOD_API_KEY" }
+Invoke-RestMethod -Method Post -Uri 'https://api.runpod.io/v2/pods' -Headers $headers -ContentType 'application/json' -Body ($payload | ConvertTo-Json -Depth 8)

tools/Probe-H3Nvfp4Linear.ps1

@@ -0,0 +1,63 @@
+$ErrorActionPreference = 'Stop'
+
+$envPath = Join-Path $PSScriptRoot '..\..\..\.env'
+Get-Content -LiteralPath $envPath | ForEach-Object {
+    if ($_ -match '^\s*([^#=\s]+)\s*=\s*(.*)\s*$') {
+        [Environment]::SetEnvironmentVariable($matches[1], $matches[2].Trim('"').Trim("'"), 'Process')
+    }
+}
+if (-not $env:RUNPOD_API_KEY) { throw 'RUNPOD_API_KEY is not set in .env.' }
+
+$probe = @'
+import json
+import torch
+from safetensors import safe_open
+import comfy_kitchen as ck
+from comfy_kitchen.tensor import TensorCoreNVFP4Layout
+
+path = "/runpod-volume/ComfyUI/models/diffusion_models/minimax_h3_ref2va_pruned_nvfp4.safetensors"
+prefix = "blocks.0.attn.qkv_proj"
+with safe_open(path, framework="pt", device="cuda") as checkpoint:
+    sidecar = checkpoint.get_tensor(prefix + ".comfy_quant")
+    assert json.loads(bytes(sidecar.cpu().tolist())) == {"format": "nvfp4"}
+    weight = checkpoint.get_tensor(prefix + ".weight").contiguous()
+    weight_scale = checkpoint.get_tensor(prefix + ".weight_scale").view(torch.float8_e4m3fn).contiguous()
+    weight_scale_2 = checkpoint.get_tensor(prefix + ".weight_scale_2").to(torch.float32).contiguous()
+
+x = torch.randn((1, weight.shape[1] * 2), device="cuda", dtype=torch.bfloat16)
+packed_x, x_params = TensorCoreNVFP4Layout.quantize(x)
+y = ck.scaled_mm_nvfp4(
+    packed_x, weight,
+    tensor_scale_a=x_params.scale,
+    tensor_scale_b=weight_scale_2,
+    block_scale_a=x_params.block_scale,
+    block_scale_b=weight_scale,
+    out_dtype=torch.bfloat16,
+)
+torch.cuda.synchronize()
+print(json.dumps({
+    "input_shape": list(x.shape),
+    "packed_weight_shape": list(weight.shape),
+    "weight_scale_shape": list(weight_scale.shape),
+    "weight_scale_2_shape": list(weight_scale_2.shape),
+    "output_shape": list(y.shape),
+    "output_dtype": str(y.dtype),
+    "finite": bool(torch.isfinite(y).all()),
+    "gpu": torch.cuda.get_device_name(),
+}, sort_keys=True))
+'@
+$encodedProbe = [Convert]::ToBase64String([Text.Encoding]::UTF8.GetBytes($probe))
+$setup = 'python -m pip install --no-cache-dir comfy-kitchen==0.2.28 "safetensors>=0.5.0"; python -c "import base64;exec(base64.b64decode(\"' + $encodedProbe + '\"))"'
+$payload = @{
+    name = 'h3-nvfp4-linear-probe'
+    image = 'runpod/pytorch:1.1.0-rc.154-cu1300-torch291-ubuntu2404'
+    cloud = 'SECURE'
+    gpu = @{ id = 'NVIDIA RTX PRO 6000 Blackwell Server Edition'; count = 1 }
+    disk = 20
+    dataCenterIds = @('EUR-IS-1')
+    mounts = @{ network = @(@{ volumeId = 'g8r2uqufz5'; path = '/runpod-volume' }) }
+    args = 'bash -lc ''set -eu; ' + $setup + ''''
+}
+
+$headers = @{ Authorization = "Bearer $env:RUNPOD_API_KEY" }
+Invoke-RestMethod -Method Post -Uri 'https://api.runpod.io/v2/pods' -Headers $headers -ContentType 'application/json' -Body ($payload | ConvertTo-Json -Depth 8)

tools/Summarize-H3Nvfp4Checkpoint.ps1

@@ -0,0 +1,26 @@
+$ErrorActionPreference = 'Stop'
+
+$envPath = Join-Path $PSScriptRoot '..\..\..\.env'
+Get-Content -LiteralPath $envPath | ForEach-Object {
+    if ($_ -match '^\s*([^#=\s]+)\s*=\s*(.*)\s*$') {
+        [Environment]::SetEnvironmentVariable($matches[1], $matches[2].Trim('"').Trim("'"), 'Process')
+    }
+}
+if (-not $env:RUNPOD_API_KEY) { throw 'RUNPOD_API_KEY is not set in .env.' }
+
+$header = '/runpod-volume/h3-blackwell-runtime/artifacts/checkpoints/minimax_h3_ref2va_pruned_nvfp4.header.json'
+$probe = "import collections,json; h=json.load(open('$header')); t=h['tensors']; groups=collections.Counter('.'.join(k.split('.')[:3]) for k in t); u8={k:v for k,v in t.items() if v['dtype']=='U8'}; scale=[k for k,v in t.items() if v['dtype']=='F8_E4M3']; print(json.dumps({'u8_tensors':u8,'fp8_scale_tensors':scale,'module_prefix_counts':groups.most_common(80)},sort_keys=True))"
+$encodedProbe = [Convert]::ToBase64String([Text.Encoding]::UTF8.GetBytes($probe))
+$payload = @{
+    name = 'h3-nvfp4-checkpoint-summary'
+    image = 'runpod/pytorch:1.1.0-rc.154-cu1300-torch291-ubuntu2404'
+    cloud = 'SECURE'
+    gpu = @{ id = 'NVIDIA RTX PRO 6000 Blackwell Server Edition'; count = 1 }
+    disk = 20
+    dataCenterIds = @('EUR-IS-1')
+    mounts = @{ network = @(@{ volumeId = 'g8r2uqufz5'; path = '/runpod-volume' }) }
+    args = 'bash -lc ''set -eu; python -c "import base64;exec(base64.b64decode(\"' + $encodedProbe + '\"))"'''
+}
+
+$headers = @{ Authorization = "Bearer $env:RUNPOD_API_KEY" }
+Invoke-RestMethod -Method Post -Uri 'https://api.runpod.io/v2/pods' -Headers $headers -ContentType 'application/json' -Body ($payload | ConvertTo-Json -Depth 8)

tools/compare_adaln_dispatch.py

@@ -0,0 +1,24 @@
+"""Compare ComfyUI's effective AdaLN dispatch tensors with direct FP32 linear."""
+
+import torch
+import torch.nn.functional as functional
+
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+
+
+capture_dir = "/artifacts/capture"
+inputs = torch.load(f"{capture_dir}/input.pt", map_location="cuda", weights_only=False)
+expected = torch.load(f"{capture_dir}/block0_norm1_adaln.pt", map_location="cuda", weights_only=False)
+weight = expected["effective_weight"]
+bias = expected["effective_bias"]
+curve = H3Checkpoint("/models/minimax_h3_ref2va_pruned_nvfp4.safetensors").tensor("adaln_t_table", dtype=torch.float32)
+position = inputs["timesteps"].float().clamp(0, 1) * (curve.shape[0] - 1)
+lower = position.floor().long().clamp(max=curve.shape[0] - 2)
+embedding = torch.lerp(curve[lower], curve[lower + 1], (position - lower).unsqueeze(1))
+values = functional.linear(embedding.to(weight.dtype), weight, bias).float().view(embedding.shape[0] * 3, 6 * expected["shift"].shape[-1])
+shift, scale, *_ = values.chunk(6, dim=-1)
+
+print(f"effective_weight_dtype={weight.dtype} effective_bias_dtype={bias.dtype}")
+for name, actual, reference in (("shift", shift, expected["shift"]), ("scale", scale, expected["scale"])):
+    delta = (actual.float() - reference.float()).abs()
+    print(f"{name} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")

tools/compare_attention_backends.py

@@ -0,0 +1,34 @@
+"""Compare direct attention kernels against captured Comfy block-0 output."""
+
+import time
+
+import torch
+import torch.nn.functional as functional
+
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.denoiser import H3PackedDenoiser
+
+
+payload = torch.load("/artifacts/capture/block0_qkv_prepared.pt", map_location="cuda", weights_only=False)
+expected = torch.load("/artifacts/capture/block0_attention.pt", map_location="cuda", weights_only=False)
+q, k, v = payload["q"], payload["k"], payload["v"]
+model = H3PackedDenoiser.from_checkpoint(H3Checkpoint("/models/minimax_h3_ref2va_pruned_nvfp4.safetensors")).eval()
+out_proj = model.backbone.blocks[0].attention.out_proj
+
+for name in ("sdpa", "sage2", "sage3"):
+    torch.cuda.synchronize()
+    start = time.perf_counter()
+    if name == "sdpa":
+        output = functional.scaled_dot_product_attention(q, k, v, is_causal=False)
+    elif name == "sage2":
+        from sageattention import sageattn
+
+        output = sageattn(q, k, v, is_causal=False, tensor_layout="HND", smooth_k=False)
+    else:
+        from sageattn3 import sageattn3_blackwell
+
+        output = sageattn3_blackwell(q, k, v, is_causal=False)
+    torch.cuda.synchronize()
+    output = out_proj(output.transpose(1, 2).reshape(q.shape[0], q.shape[2], -1).reshape(q.shape[2], -1).contiguous())
+    delta = (output.float() - expected.float()).abs()
+    print(f"{name} elapsed_s={time.perf_counter() - start:.3f} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")

tools/compare_benchmark.py

@@ -0,0 +1,28 @@
+"""Compare one direct-runner timing record to the fixed Comfy Sage3 baseline."""
+
+import argparse
+import json
+from pathlib import Path
+
+
+ROOT = Path(__file__).resolve().parents[1]
+CONTRACT = json.loads((ROOT / "benchmarks" / "ref2va-960x544-124f.json").read_text(encoding="utf-8"))
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--result", required=True, type=Path)
+    args = parser.parse_args()
+    result = json.loads(args.result.read_text(encoding="utf-8"))
+    actual = float(result["wall_time_seconds"])
+    baseline = float(CONTRACT["reference_comfy_sage3_seconds"])
+    print(json.dumps({
+        "baseline_seconds": baseline,
+        "direct_runner_seconds": actual,
+        "delta_seconds": round(actual - baseline, 3),
+        "speedup": round(baseline / actual, 4),
+    }, indent=2))
+
+
+if __name__ == "__main__":
+    main()

tools/compare_block0_adaln.py

@@ -0,0 +1,19 @@
+"""Compare direct block-0 RMSNorm and curve AdaLN parameters with ComfyUI."""
+
+import torch
+
+from h3_blackwell_runtime.attention import rms_norm
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.denoiser import H3PackedDenoiser
+
+
+capture_dir = "/artifacts/capture"
+inputs = torch.load(f"{capture_dir}/input.pt", map_location="cuda", weights_only=False)
+expected = torch.load(f"{capture_dir}/block0_norm1_adaln.pt", map_location="cuda", weights_only=False)
+model = H3PackedDenoiser.from_checkpoint(H3Checkpoint("/models/minimax_h3_ref2va_pruned_nvfp4.safetensors")).eval()
+shift, scale, *_ = model.backbone.adaln[0](inputs["timesteps"])
+norm = rms_norm(inputs["hidden"], model.backbone.blocks[0].norm1_weight, model.backbone.blocks[0].norm_eps)
+
+for name, actual, reference in (("norm", norm, expected["norm"]), ("shift", shift, expected["shift"]), ("scale", scale, expected["scale"])):
+    delta = (actual.float() - reference.float()).abs()
+    print(f"{name} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")

tools/compare_block0_norm_variants.py

@@ -0,0 +1,22 @@
+"""Identify the ComfyUI-compatible block-0 RMSNorm precision path."""
+
+import torch
+
+from h3_blackwell_runtime.block import modulate_segments
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.denoiser import H3PackedDenoiser
+
+
+capture_dir = "/artifacts/capture"
+inputs = torch.load(f"{capture_dir}/input.pt", map_location="cuda", weights_only=False)
+expected = torch.load(f"{capture_dir}/block0_norm1.pt", map_location="cuda", weights_only=False)
+checkpoint = H3Checkpoint("/models/minimax_h3_ref2va_pruned_nvfp4.safetensors")
+model = H3PackedDenoiser.from_checkpoint(checkpoint).eval()
+shift, scale, *_ = model.backbone.adaln[0](inputs["timesteps"])
+raw_weight = checkpoint.tensor("blocks.0.norm1.weight")
+
+for epsilon in (1e-6, 1e-5, 1e-4):
+    normalized = torch.nn.functional.rms_norm(inputs["hidden"], (inputs["hidden"].shape[-1],), raw_weight, epsilon)
+    actual = modulate_segments(normalized, shift, scale, inputs["segments"])
+    delta = (actual.float() - expected.float()).abs()
+    print(f"epsilon={epsilon:.0e} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")

tools/compare_block0_qkv.py

@@ -0,0 +1,39 @@
+"""Compare prepared block-0 QKV tensors and Sage3 output with ComfyUI."""
+
+import torch
+
+from h3_blackwell_runtime.attention import apply_split_half_rope, rms_norm
+from h3_blackwell_runtime.block import modulate_segments
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.denoiser import H3PackedDenoiser
+from h3_blackwell_runtime.rope import h3_rope_rotation
+
+
+capture_dir = "/artifacts/capture"
+inputs = torch.load(f"{capture_dir}/input.pt", map_location="cuda", weights_only=False)
+expected_qkv = torch.load(f"{capture_dir}/block0_qkv_prepared.pt", map_location="cuda", weights_only=False)
+expected_raw = torch.load(f"{capture_dir}/block0_qkv_raw.pt", map_location="cuda", weights_only=False)
+expected_norm1 = torch.load(f"{capture_dir}/block0_norm1.pt", map_location="cuda", weights_only=False)
+expected_attention = torch.load(f"{capture_dir}/block0_attention.pt", map_location="cuda", weights_only=False)
+checkpoint = H3Checkpoint("/models/minimax_h3_ref2va_pruned_nvfp4.safetensors")
+model = H3PackedDenoiser.from_checkpoint(checkpoint).eval()
+block = model.backbone.blocks[0]
+shift_msa, scale_msa, *_ = model.backbone.adaln[0](inputs["timesteps"])
+hidden = modulate_segments(rms_norm(inputs["hidden"], block.norm1_weight, block.norm_eps), shift_msa, scale_msa, inputs["segments"])
+sequence = hidden.shape[0]
+inner = block.attention.heads * block.attention.head_dim
+rotation = h3_rope_rotation(inputs["position_ids"], model.backbone.inv_freq, hidden.dtype)
+
+with torch.inference_mode():
+    raw_q, raw_k, raw_v = block.attention.qkv_proj(hidden).split(inner, dim=-1)
+    comfy_raw_q, comfy_raw_k, comfy_raw_v = block.attention.qkv_proj(expected_norm1).split(inner, dim=-1)
+    q, k, v = raw_q, raw_k, raw_v
+    q = apply_split_half_rope(rms_norm(q.view(1, sequence, 56, 128), block.attention.q_norm_weight, block.attention.eps), rotation).transpose(1, 2).contiguous()
+    k = apply_split_half_rope(rms_norm(k.view(1, sequence, 56, 128), block.attention.k_norm_weight, block.attention.eps), rotation).transpose(1, 2).contiguous()
+    v = v.view(1, sequence, 56, 128).transpose(1, 2).contiguous()
+    from sageattn3 import sageattn3_blackwell
+    attention = block.attention.out_proj(sageattn3_blackwell(q, k, v, is_causal=False).transpose(1, 2).reshape(sequence, inner).contiguous())
+
+for name, actual, expected in (("raw_q", raw_q, expected_raw["q"]), ("raw_k", raw_k, expected_raw["k"]), ("raw_v", raw_v, expected_raw["v"]), ("comfy_norm_raw_q", comfy_raw_q, expected_raw["q"]), ("comfy_norm_raw_k", comfy_raw_k, expected_raw["k"]), ("comfy_norm_raw_v", comfy_raw_v, expected_raw["v"]), ("q", q, expected_qkv["q"]), ("k", k, expected_qkv["k"]), ("v", v, expected_qkv["v"]), ("attention", attention, expected_attention)):
+    delta = (actual.float() - expected.float()).abs()
+    print(f"{name} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")

tools/compare_curve_embedding.py

@@ -0,0 +1,23 @@
+"""Compare ComfyUI curve interpolation and block-0 AdaLN projection exactly."""
+
+import torch
+import torch.nn.functional as functional
+
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.denoiser import H3PackedDenoiser
+
+
+capture_dir = "/artifacts/capture"
+inputs = torch.load(f"{capture_dir}/input.pt", map_location="cuda", weights_only=False)
+expected = torch.load(f"{capture_dir}/block0_norm1_adaln.pt", map_location="cuda", weights_only=False)
+model = H3PackedDenoiser.from_checkpoint(H3Checkpoint("/models/minimax_h3_ref2va_pruned_nvfp4.safetensors")).eval()
+adaln = model.backbone.adaln[0]
+position = inputs["timesteps"].float().clamp(0, 1) * (adaln.curve_table.shape[0] - 1)
+lower = position.floor().long().clamp(max=adaln.curve_table.shape[0] - 2)
+embedding = torch.lerp(adaln.curve_table[lower], adaln.curve_table[lower + 1], (position - lower).unsqueeze(1))
+values = functional.linear(embedding, adaln.weight, adaln.bias).view(embedding.shape[0] * 3, 6 * adaln.hidden_size)
+shift, scale, *_ = values.chunk(6, dim=-1)
+
+for name, actual, reference in (("embedding", embedding, inputs["t_emb"]), ("weight", adaln.weight, expected["weight"]), ("bias", adaln.bias, expected["bias"]), ("shift", shift, expected["shift"]), ("scale", scale, expected["scale"])):
+    delta = (actual.float() - reference.float()).abs()
+    print(f"{name} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")

tools/compare_direct_adaln_weight.py

@@ -0,0 +1,12 @@
+import torch
+
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.denoiser import H3PackedDenoiser
+
+
+expected = torch.load("/artifacts/capture/block0_norm1_adaln.pt", map_location="cuda", weights_only=False)
+model = H3PackedDenoiser.from_checkpoint(H3Checkpoint("/models/minimax_h3_ref2va_pruned_nvfp4.safetensors")).eval()
+for name, actual, reference in (("weight", model.backbone.adaln[0].weight, expected["effective_weight"]), ("bias", model.backbone.adaln[0].bias, expected["effective_bias"])):
+    for precision, value in (("fp32", actual), ("bf16", actual.bfloat16().float()), ("fp16", actual.half().float())):
+        delta = (value.float() - reference.float()).abs()
+        print(f"{name}_{precision} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")

tools/compare_fl2va_steps.py

@@ -0,0 +1,44 @@
+"""Compare direct FL2VA beta/RES steps with captured Comfy sampler state."""
+
+import glob
+
+import torch
+
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.denoiser import H3PackedDenoiser
+from h3_blackwell_runtime.packing import H3PromptPacker, unpatchify_video
+from h3_blackwell_runtime.qwen3vl_text import Qwen3VLPromptConditioner
+from h3_blackwell_runtime.sampler import _audio_sigma, _unpack_audio, res_multistep_update
+from h3_blackwell_runtime.token_refiner import H3TokenRefiner
+
+
+root = "/artifacts/fl2va-sampler-reference"
+initial = torch.load(f"{root}/initial.pt", map_location="cuda", weights_only=False)
+steps = [torch.load(path, map_location="cuda", weights_only=False) for path in sorted(glob.glob(f"{root}/step_*.pt"))]
+sigmas = initial["sigmas"].to("cuda")
+checkpoint = H3Checkpoint("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
+text = H3TokenRefiner(checkpoint)(Qwen3VLPromptConditioner("/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", "/opt/h3-blackwell-runtime/qwen25_tokenizer")("A brass-and-paper dragon flies above a rain-washed old city at blue hour."))
+model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend="sage2").eval()
+packer = H3PromptPacker(checkpoint)
+
+video, audio_carried = initial["initial_x"]
+video, audio_carried = video.to("cuda"), audio_carried.to("cuda")
+old_video = old_audio = old_sigma = None
+for index, reference in enumerate(steps):
+    sigma, sigma_down = sigmas[index], sigmas[index + 1]
+    native_audio = audio_carried * (_audio_sigma(sigma) / sigma)
+    hidden, times, segments, positions, video_segment, audio_segment = packer(text, video, native_audio, float(sigma))
+    raw_video, raw_audio = model(hidden, times, positions, segments, video_segment, audio_segment)
+    velocity_video = -unpatchify_video(raw_video, video.shape[2], video.shape[-2], video.shape[-1])
+    velocity_audio = -_unpack_audio(raw_audio)
+    carry = _audio_sigma(sigma) / sigma
+    velocity_audio = (1.0 - 4.0) * (audio_carried * carry) + (1.0 + 3.0 * _audio_sigma(sigma)) * velocity_audio
+    denoised = (video - sigma * velocity_video, audio_carried - sigma * velocity_audio)
+    reference_denoised = reference["denoised"]
+    denoised_delta = (denoised[0].float() - reference_denoised[0].float()).abs()
+    previous_sigma = sigmas[index - 1] if index else None
+    video = res_multistep_update(video, denoised[0], sigma, sigma_down, old_video, old_sigma, previous_sigma)
+    audio_carried = res_multistep_update(audio_carried, denoised[1], sigma, sigma_down, old_audio, old_sigma, previous_sigma)
+    latent_delta = (video.float() - reference["x"][0].float()).abs()
+    print(f"step={index:02d} x0_video_mean={denoised_delta.mean().item():.6g} x0_video_max={denoised_delta.max().item():.6g} latent_video_mean={latent_delta.mean().item():.6g} latent_video_max={latent_delta.max().item():.6g}")
+    old_video, old_audio, old_sigma = denoised[0], denoised[1], sigma_down

tools/direct_t2v_preview.py

@@ -0,0 +1,49 @@
+"""Generate a minimal direct, video-only H3 T2V preview without ComfyUI."""
+
+import argparse
+from pathlib import Path
+import subprocess
+import warnings
+
+warnings.filterwarnings("ignore", message="Found GPU0 NVIDIA GB10 which is of cuda capability 12.1.*", category=UserWarning)
+
+import torch
+
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.denoiser import H3PackedDenoiser
+from h3_blackwell_runtime.packing import H3PromptPacker
+from h3_blackwell_runtime.qwen3vl_text import Qwen3VLPromptConditioner
+from h3_blackwell_runtime.sampler import sample_video_res_multistep
+from h3_blackwell_runtime.t2v import empty_av_latents
+from h3_blackwell_runtime.token_refiner import H3TokenRefiner
+from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--prompt", default="A brass-and-paper dragon flies above a rain-washed old city at blue hour.")
+parser.add_argument("--output", type=Path, default=Path("/output/direct-h3-preview.mp4"))
+parser.add_argument("--width", type=int, default=320)
+parser.add_argument("--height", type=int, default=192)
+parser.add_argument("--frames", type=int, default=22)
+parser.add_argument("--steps", type=int, default=12)
+parser.add_argument("--seed", type=int, default=440204)
+parser.add_argument("--attention", choices=("sage2", "sdpa", "sage3"), default="sage2")
+args = parser.parse_args()
+
+torch.manual_seed(args.seed)
+checkpoint = H3Checkpoint("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
+conditioner = Qwen3VLPromptConditioner("/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", "/opt/h3-blackwell-runtime/qwen25_tokenizer")
+video, audio, frames = empty_av_latents(args.width, args.height, args.frames)
+video.normal_()
+model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend=args.attention).eval()
+text = H3TokenRefiner(checkpoint)(conditioner(args.prompt))
+latent = sample_video_res_multistep(model, H3PromptPacker(checkpoint), text, video, audio, steps=args.steps)
+vae = MiniMaxH3VideoVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda").eval()
+pixels = vae.decode(latent.to(next(vae.parameters()).dtype))[:, :, :frames]
+pixels = ((pixels[0].permute(1, 2, 3, 0).clamp(-1, 1) + 1) * 127.5).to(torch.uint8).cpu()
+args.output.parent.mkdir(parents=True, exist_ok=True)
+raw = args.output.with_suffix(".rgb")
+pixels.numpy().tofile(raw)
+subprocess.run(["ffmpeg", "-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(args.output)], check=True)
+raw.unlink()
+print({"output": str(args.output), "frames": frames, "shape": tuple(pixels.shape)})

tools/inspect_capture.py

@@ -0,0 +1,19 @@
+"""Print tensor and segment metadata from an H3 direct-runner capture."""
+
+import argparse
+
+import torch
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("capture_dir")
+args = parser.parse_args()
+
+for name in ("input.pt", "output.pt"):
+    payload = torch.load(f"{args.capture_dir}/{name}", map_location="cpu", weights_only=False)
+    print(name)
+    for key, value in payload.items():
+        if isinstance(value, torch.Tensor):
+            print(f"  {key}: shape={tuple(value.shape)} dtype={value.dtype} finite={torch.isfinite(value).all().item()}")
+        else:
+            print(f"  {key}: {value}")

tools/inspect_checkpoint_tensors.py

@@ -0,0 +1,10 @@
+"""List selected H3 checkpoint tensor shapes."""
+
+from safetensors import safe_open
+
+
+path = "/models/minimax_h3_ref2va_pruned_nvfp4.safetensors"
+with safe_open(path, framework="pt", device="cpu") as checkpoint:
+    for name in checkpoint.keys():
+        if name == "adaln_t_table" or name.startswith("blocks.0.adaln_proj") or name.startswith("final_layer.adaln_proj"):
+            print(name, tuple(checkpoint.get_tensor(name).shape))

tools/inspect_comfy_db.py

@@ -0,0 +1,13 @@
+"""List persisted ComfyUI SQLite tables and recent rows for capture recovery."""
+
+import sqlite3
+
+
+connection = sqlite3.connect("/workspace/ComfyUI/user/comfyui-v031.db")
+tables = connection.execute(
+    "SELECT name FROM sqlite_master WHERE type = 'table' ORDER BY name"
+).fetchall()
+print(tables)
+for (table,) in tables:
+    count = connection.execute(f'SELECT COUNT(*) FROM "{table}"').fetchone()[0]
+    print(table, count)

tools/inspect_fl2va_refiner.py

@@ -0,0 +1,7 @@
+from safetensors import safe_open
+
+
+with safe_open("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors", framework="pt", device="cpu") as checkpoint:
+    for name in checkpoint.keys():
+        if name.startswith("condition_proj") or name.startswith("token_refiner"):
+            print(name, tuple(checkpoint.get_tensor(name).shape))

tools/inspect_safetensors.py

@@ -0,0 +1,40 @@
+"""Report safetensors metadata without materializing model weights."""
+
+import argparse
+import json
+import struct
+from collections import Counter
+from pathlib import Path
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("checkpoint", type=Path)
+    parser.add_argument("--output", type=Path)
+    args = parser.parse_args()
+
+    with args.checkpoint.open("rb") as file:
+        header_size = struct.unpack("<Q", file.read(8))[0]
+        header = json.loads(file.read(header_size))
+
+    tensors = {name: spec for name, spec in header.items() if name != "__metadata__"}
+    report = {
+        "checkpoint": str(args.checkpoint),
+        "metadata": header.get("__metadata__", {}),
+        "tensor_count": len(tensors),
+        "dtypes": dict(sorted(Counter(spec["dtype"] for spec in tensors.values()).items())),
+        "tensors": {
+            name: {key: spec[key] for key in ("dtype", "shape", "data_offsets") if key in spec}
+            for name, spec in tensors.items()
+        },
+    }
+    rendered = json.dumps(report, indent=2, sort_keys=True)
+    if args.output:
+        args.output.parent.mkdir(parents=True, exist_ok=True)
+        args.output.write_text(rendered + "\n", encoding="utf-8")
+    else:
+        print(rendered)
+
+
+if __name__ == "__main__":
+    main()

tools/inspect_sampler_reference.py

@@ -0,0 +1,8 @@
+import glob
+
+import torch
+
+
+for path in sorted(glob.glob("/artifacts/fl2va-sampler-reference/*.pt")):
+    state = torch.load(path, map_location="cpu", weights_only=False)
+    print(path.rsplit("/", 1)[-1], float(state["sigma"]))

tools/list_attention_backends.py

@@ -0,0 +1,5 @@
+from h3_blackwell_runtime.attention import attention_backend_status
+
+
+for name, status in attention_backend_status().items():
+    print(f"{name}: {status}")

tools/localize_block0_sublayers.py

@@ -0,0 +1,32 @@
+"""Compare direct block-0 intermediates with ComfyUI captures."""
+
+import torch
+
+from h3_blackwell_runtime.block import gate_segments, modulate_segments
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.denoiser import H3PackedDenoiser
+from h3_blackwell_runtime.attention import rms_norm
+from h3_blackwell_runtime.rope import h3_rope_rotation
+
+
+capture_dir = "/artifacts/capture"
+inputs = torch.load(f"{capture_dir}/input.pt", map_location="cuda", weights_only=False)
+checkpoint = H3Checkpoint("/models/minimax_h3_ref2va_pruned_nvfp4.safetensors")
+model = H3PackedDenoiser.from_checkpoint(checkpoint).eval()
+block = model.backbone.blocks[0]
+adaln = model.backbone.adaln[0]
+rotation = h3_rope_rotation(inputs["position_ids"], model.backbone.inv_freq, inputs["hidden"].dtype)
+shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = adaln(inputs["timesteps"])
+
+with torch.inference_mode():
+    norm1 = modulate_segments(rms_norm(inputs["hidden"], block.norm1_weight, block.norm_eps), shift_msa, scale_msa, inputs["segments"])
+    attention = block.attention(norm1, rotation)
+    post_attention = gate_segments(inputs["hidden"], attention, gate_msa, inputs["segments"])
+    norm2 = modulate_segments(rms_norm(post_attention, block.norm2_weight, block.norm_eps), shift_mlp, scale_mlp, inputs["segments"])
+    mlp = block.mlp(norm2)
+    post_mlp = gate_segments(post_attention, mlp, gate_mlp, inputs["segments"])
+
+for name, actual in (("norm1", norm1), ("attention", attention), ("post_attention", post_attention), ("norm2", norm2), ("mlp", mlp), ("post_mlp", post_mlp)):
+    expected = torch.load(f"{capture_dir}/block0_{name}.pt", map_location="cuda", weights_only=False)
+    delta = (actual.float() - expected.float()).abs()
+    print(f"{name} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")

tools/localize_block_mismatch.py

@@ -0,0 +1,22 @@
+"""Compare every direct H3 block output with one ComfyUI per-block capture."""
+
+import torch
+
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.denoiser import H3PackedDenoiser
+from h3_blackwell_runtime.rope import h3_rope_rotation
+
+
+capture_dir = "/artifacts/capture"
+inputs = torch.load(f"{capture_dir}/input.pt", map_location="cuda", weights_only=False)
+checkpoint = H3Checkpoint("/models/minimax_h3_ref2va_pruned_nvfp4.safetensors")
+model = H3PackedDenoiser.from_checkpoint(checkpoint).eval()
+
+hidden = inputs["hidden"]
+rotation = h3_rope_rotation(inputs["position_ids"], model.backbone.inv_freq, hidden.dtype)
+with torch.inference_mode():
+    for index, (block, adaln) in enumerate(zip(model.backbone.blocks, model.backbone.adaln, strict=True)):
+        hidden = block(hidden, rotation, *adaln(inputs["timesteps"]), inputs["segments"])
+        expected = torch.load(f"{capture_dir}/blocks/{index:02d}.pt", map_location="cuda", weights_only=False)
+        delta = (hidden.float() - expected.float()).abs()
+        print(f"block={index:02d} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")

tools/patch_comfy_h3_adaln_dispatch.py

@@ -0,0 +1,15 @@
+"""Capture ComfyUI's effective dynamically cast AdaLN linear parameters."""
+
+from pathlib import Path
+
+
+model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
+source = model.read_text(encoding="utf-8")
+old = 'if capture_dir: torch.save({"norm": norm1.detach().cpu(), "shift": shift_msa.detach().cpu(), "scale": scale_msa.detach().cpu()}, os.path.join(capture_dir, "block0_norm1_adaln.pt"))'
+new = 'if capture_dir:\n            effective_weight, effective_bias, effective_stream = comfy.ops.cast_bias_weight(self.adaln_proj.linear, t_emb, offloadable=True)\n            torch.save({"norm": norm1.detach().cpu(), "shift": shift_msa.detach().cpu(), "scale": scale_msa.detach().cpu(), "effective_weight": effective_weight.detach().cpu(), "effective_bias": effective_bias.detach().cpu()}, os.path.join(capture_dir, "block0_norm1_adaln.pt"))\n            comfy.ops.uncast_bias_weight(self.adaln_proj.linear, effective_weight, effective_bias, effective_stream)'
+if source.count(old) == 1:
+    source = source.replace(old, new)
+elif new not in source:
+    raise RuntimeError("Unable to locate AdaLN capture payload.")
+model.write_text(source, encoding="utf-8")
+print("Applied H3 AdaLN dispatch capture patch.")

tools/patch_comfy_h3_adaln_weights.py

@@ -0,0 +1,15 @@
+"""Add block-0 in-memory AdaLN projection weights to the capture."""
+
+from pathlib import Path
+
+
+model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
+source = model.read_text(encoding="utf-8")
+old = 'torch.save({"norm": norm1.detach().cpu(), "shift": shift_msa.detach().cpu(), "scale": scale_msa.detach().cpu()}, os.path.join(capture_dir, "block0_norm1_adaln.pt"))'
+new = 'torch.save({"norm": norm1.detach().cpu(), "shift": shift_msa.detach().cpu(), "scale": scale_msa.detach().cpu(), "weight": self.adaln_proj.linear.weight.detach().cpu(), "bias": self.adaln_proj.linear.bias.detach().cpu()}, os.path.join(capture_dir, "block0_norm1_adaln.pt"))'
+if source.count(old) == 1:
+    source = source.replace(old, new)
+elif new not in source:
+    raise RuntimeError("Unable to locate AdaLN capture payload.")
+model.write_text(source, encoding="utf-8")
+print("Applied H3 AdaLN weight capture patch.")

tools/patch_comfy_h3_block0_adaln.py

@@ -0,0 +1,23 @@
+"""Capture block-0 RMSNorm and AdaLN components before modulation."""
+
+from pathlib import Path
+
+
+model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
+source = model.read_text(encoding="utf-8")
+old = (
+    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_capture_index\", -1) == 0 and H3_CAPTURE_ACTIVE else None\n"
+    "        h = _mod_scale_shift(self.norm1(x), shift_msa, scale_msa, mod_segments)\n"
+)
+new = (
+    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_capture_index\", -1) == 0 and H3_CAPTURE_ACTIVE else None\n"
+    "        norm1 = self.norm1(x)\n"
+    "        if capture_dir: torch.save({\"norm\": norm1.detach().cpu(), \"shift\": shift_msa.detach().cpu(), \"scale\": scale_msa.detach().cpu()}, os.path.join(capture_dir, \"block0_norm1_adaln.pt\"))\n"
+    "        h = _mod_scale_shift(norm1, shift_msa, scale_msa, mod_segments)\n"
+)
+if source.count(old) == 1:
+    source = source.replace(old, new)
+elif new not in source:
+    raise RuntimeError("Unable to locate block-0 norm1 capture point.")
+model.write_text(source, encoding="utf-8")
+print("Applied H3 block-0 RMSNorm/AdaLN capture patch.")

tools/patch_comfy_h3_block0_qkv.py

@@ -0,0 +1,31 @@
+"""Capture block-0 QKV before and after ComfyUI attention preparation."""
+
+from pathlib import Path
+
+
+model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
+source = model.read_text(encoding="utf-8")
+
+old = "            block._h3_capture_index = i\n            comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)\n"
+new = "            block._h3_capture_index = i\n            block.attn._h3_capture_index = i\n            comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)\n"
+if source.count(old) == 1:
+    source = source.replace(old, new)
+elif new not in source:
+    raise RuntimeError("Unable to locate the H3 block loop.")
+
+old = "        v = v.transpose(0, 1).unsqueeze(0)\n        out = optimized_attention(q, k, v, self.heads, mask=None, skip_reshape=True, transformer_options=transformer_options)\n"
+new = (
+    "        v = v.transpose(0, 1).unsqueeze(0)\n"
+    "        if getattr(self, \"_h3_capture_index\", -1) == 0 and H3_CAPTURE_ACTIVE:\n"
+    "            capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
+    "            if capture_dir:\n"
+    "                torch.save({\"q\": q.detach().cpu(), \"k\": k.detach().cpu(), \"v\": v.detach().cpu()}, os.path.join(capture_dir, \"block0_qkv_prepared.pt\"))\n"
+    "        out = optimized_attention(q, k, v, self.heads, mask=None, skip_reshape=True, transformer_options=transformer_options)\n"
+)
+if source.count(old) == 1:
+    source = source.replace(old, new)
+elif new not in source:
+    raise RuntimeError("Unable to locate the prepared QKV call.")
+
+model.write_text(source, encoding="utf-8")
+print("Applied H3 block-0 QKV capture patch.")

tools/patch_comfy_h3_block0_raw_qkv.py

@@ -0,0 +1,22 @@
+"""Capture block-0 QKV immediately after ComfyUI's quantized projection."""
+
+from pathlib import Path
+
+
+model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
+source = model.read_text(encoding="utf-8")
+old = "        q, k, v = self.qkv_proj(x).split(self.heads * self.head_dim, dim=-1)\n        v = v.view(s, self.heads, self.head_dim)\n"
+new = (
+    "        q, k, v = self.qkv_proj(x).split(self.heads * self.head_dim, dim=-1)\n"
+    "        if getattr(self, \"_h3_capture_index\", -1) == 0 and H3_CAPTURE_ACTIVE:\n"
+    "            capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
+    "            if capture_dir:\n"
+    "                torch.save({\"q\": q.detach().cpu(), \"k\": k.detach().cpu(), \"v\": v.detach().cpu()}, os.path.join(capture_dir, \"block0_qkv_raw.pt\"))\n"
+    "        v = v.view(s, self.heads, self.head_dim)\n"
+)
+if source.count(old) == 1:
+    source = source.replace(old, new)
+elif new not in source:
+    raise RuntimeError("Unable to locate the raw QKV projection.")
+model.write_text(source, encoding="utf-8")
+print("Applied H3 raw QKV capture patch.")

tools/patch_comfy_h3_block0_sublayers.py

@@ -0,0 +1,46 @@
+"""Capture block-0 AdaLN, attention, and MLP intermediates once."""
+
+from pathlib import Path
+
+
+model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
+source = model.read_text(encoding="utf-8")
+
+old = "        for i, block in enumerate(self.blocks):\n            comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)\n"
+new = "        for i, block in enumerate(self.blocks):\n            block._h3_capture_index = i\n            comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)\n"
+if source.count(old) == 1:
+    source = source.replace(old, new)
+elif new not in source:
+    raise RuntimeError("Unable to locate the H3 block loop.")
+
+old = (
+    "        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(t_emb)\n"
+    "        h = _mod_scale_shift(self.norm1(x), shift_msa, scale_msa, mod_segments)\n"
+    "        x = _mod_gate(x, gate_msa, self.attn(h, rope_freqs=rope_freqs, transformer_options=transformer_options), mod_segments)\n"
+    "        h = _mod_scale_shift(self.norm2(x), shift_mlp, scale_mlp, mod_segments)\n"
+    "        return _mod_gate(x, gate_mlp, self.mlp(h), mod_segments)\n"
+)
+new = (
+    "        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(t_emb)\n"
+    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_capture_index\", -1) == 0 and H3_CAPTURE_ACTIVE else None\n"
+    "        h = _mod_scale_shift(self.norm1(x), shift_msa, scale_msa, mod_segments)\n"
+    "        if capture_dir: torch.save(h.detach().cpu(), os.path.join(capture_dir, \"block0_norm1.pt\"))\n"
+    "        attn_out = self.attn(h, rope_freqs=rope_freqs, transformer_options=transformer_options)\n"
+    "        if capture_dir: torch.save(attn_out.detach().cpu(), os.path.join(capture_dir, \"block0_attention.pt\"))\n"
+    "        x = _mod_gate(x, gate_msa, attn_out, mod_segments)\n"
+    "        if capture_dir: torch.save(x.detach().cpu(), os.path.join(capture_dir, \"block0_post_attention.pt\"))\n"
+    "        h = _mod_scale_shift(self.norm2(x), shift_mlp, scale_mlp, mod_segments)\n"
+    "        if capture_dir: torch.save(h.detach().cpu(), os.path.join(capture_dir, \"block0_norm2.pt\"))\n"
+    "        mlp_out = self.mlp(h)\n"
+    "        if capture_dir: torch.save(mlp_out.detach().cpu(), os.path.join(capture_dir, \"block0_mlp.pt\"))\n"
+    "        x = _mod_gate(x, gate_mlp, mlp_out, mod_segments)\n"
+    "        if capture_dir: torch.save(x.detach().cpu(), os.path.join(capture_dir, \"block0_post_mlp.pt\"))\n"
+    "        return x\n"
+)
+if source.count(old) == 1:
+    source = source.replace(old, new)
+elif new not in source:
+    raise RuntimeError("Unable to locate H3 DiTBlock.forward.")
+
+model.write_text(source, encoding="utf-8")
+print("Applied H3 block-0 sublayer capture patch.")

tools/patch_comfy_h3_block_capture.py

@@ -0,0 +1,38 @@
+"""Extend the H3 capture hook to persist every block output once."""
+
+from pathlib import Path
+
+
+model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
+source = model.read_text(encoding="utf-8")
+
+old = 'if capture_dir and not H3_CAPTURE_ACTIVE:'
+new = 'if capture_dir and not H3_CAPTURE_ACTIVE and not os.path.exists(os.path.join(capture_dir, "blocks_complete")):'
+if source.count(old) == 1:
+    source = source.replace(old, new)
+elif new not in source:
+    raise RuntimeError("Unable to locate the H3 input capture condition.")
+
+old = (
+    '            else:\n'
+    '                h = block(h, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options)\n'
+    '        if prefetch_queue is not None:\n'
+)
+new = (
+    '            else:\n'
+    '                h = block(h, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options)\n'
+    '            if capture_dir and H3_CAPTURE_ACTIVE:\n'
+    '                block_dir = os.path.join(capture_dir, "blocks")\n'
+    '                os.makedirs(block_dir, exist_ok=True)\n'
+    '                torch.save(h.detach().cpu(), os.path.join(block_dir, f"{i:02d}.pt"))\n'
+    '        if capture_dir and H3_CAPTURE_ACTIVE:\n'
+    '            open(os.path.join(capture_dir, "blocks_complete"), "a").close()\n'
+    '        if prefetch_queue is not None:\n'
+)
+if source.count(old) == 1:
+    source = source.replace(old, new)
+elif new not in source:
+    raise RuntimeError("Unable to locate the H3 block loop.")
+
+model.write_text(source, encoding="utf-8")
+print("Applied H3 per-block capture patch.")

tools/patch_comfy_h3_capture.py

@@ -0,0 +1,39 @@
+"""Instrument ComfyUI v0.31.1 H3 once to capture direct-runner inputs/outputs."""
+
+from pathlib import Path
+
+
+def replace_once(path, old, new):
+    source = path.read_text(encoding="utf-8")
+    if new in source:
+        return
+    if source.count(old) != 1:
+        raise RuntimeError(f"Expected one matching block in {path}.")
+    path.write_text(source.replace(old, new), encoding="utf-8")
+
+
+model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
+replace_once(model, "import math\n", "import math\nimport os\n")
+replace_once(model, "VISUAL_COND_TIMESTEP = 0.999\n", "VISUAL_COND_TIMESTEP = 0.999\nH3_CAPTURE_ACTIVE = False\n")
+replace_once(
+    model,
+    "        # blocks\n        patches_replace = transformer_options.get(\"patches_replace\", {})\n",
+    "        # Capture a single fully assembled payload before the first transformer block.\n"
+    "        global H3_CAPTURE_ACTIVE\n"
+    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
+    "        if capture_dir and not H3_CAPTURE_ACTIVE:\n"
+    "            H3_CAPTURE_ACTIVE = True\n"
+    "            os.makedirs(capture_dir, exist_ok=True)\n"
+    "            torch.save({\"hidden\": h.detach().cpu(), \"timesteps\": t_vals.detach().cpu(), \"position_ids\": layout.position_ids, \"segments\": mod_segments}, os.path.join(capture_dir, \"input.pt\"))\n\n"
+    "        # blocks\n        patches_replace = transformer_options.get(\"patches_replace\", {})\n",
+)
+replace_once(
+    model,
+    "        return [-video_out.to(video_x.dtype), -audio_out.to(audio_x.dtype)]\n",
+    "        if capture_dir and H3_CAPTURE_ACTIVE:\n"
+    "            torch.save({\"video\": video_out.detach().cpu(), \"audio\": audio_out.detach().cpu(), \"video_segment\": video_seg, \"audio_segment\": audio_seg}, os.path.join(capture_dir, \"output.pt\"))\n"
+    "            H3_CAPTURE_ACTIVE = False\n"
+    "        return [-video_out.to(video_x.dtype), -audio_out.to(audio_x.dtype)]\n",
+)
+
+print("Applied H3 direct-runner capture patch.")

tools/patch_comfy_h3_curve_embedding.py

@@ -0,0 +1,15 @@
+"""Add the interpolated H3 curve embedding to the direct-runner input capture."""
+
+from pathlib import Path
+
+
+model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
+source = model.read_text(encoding="utf-8")
+old = 'torch.save({"hidden": h.detach().cpu(), "timesteps": t_vals.detach().cpu(), "position_ids": layout.position_ids, "segments": mod_segments}, os.path.join(capture_dir, "input.pt"))'
+new = 'torch.save({"hidden": h.detach().cpu(), "timesteps": t_vals.detach().cpu(), "t_emb": t_emb.detach().cpu(), "position_ids": layout.position_ids, "segments": mod_segments}, os.path.join(capture_dir, "input.pt"))'
+if source.count(old) == 1:
+    source = source.replace(old, new)
+elif new not in source:
+    raise RuntimeError("Unable to locate the H3 input capture payload.")
+model.write_text(source, encoding="utf-8")
+print("Applied H3 curve embedding capture patch.")

tools/patch_comfy_sampler_capture.py

@@ -0,0 +1,27 @@
+"""Persist compact per-step sampler state for one direct FL2VA parity run."""
+
+from pathlib import Path
+
+
+path = Path("/opt/ComfyUI/comfy/samplers.py")
+source = path.read_text(encoding="utf-8")
+if "import os\n" not in source:
+    source = source.replace("import logging\n", "import logging\nimport os\n")
+
+old = "            if callback is not None:\n                callback(x[\"i\"], x[\"denoised\"], x[\"x\"], total_steps)\n"
+new = (
+    "            capture_dir = os.getenv(\"H3_SAMPLER_CAPTURE_DIR\")\n"
+    "            if capture_dir:\n"
+    "                os.makedirs(capture_dir, exist_ok=True)\n"
+    "                if x[\"i\"] == 0:\n"
+    "                    torch.save({\"sigmas\": sigmas.detach().cpu(), \"initial_x\": x[\"x\"].detach().cpu()}, os.path.join(capture_dir, \"initial.pt\"))\n"
+    "                torch.save({\"sigma\": x[\"sigma\"].detach().cpu(), \"denoised\": x[\"denoised\"].detach().cpu(), \"x\": x[\"x\"].detach().cpu()}, os.path.join(capture_dir, f\"step_{x[\"i\"]:02d}.pt\"))\n"
+    "            if callback is not None:\n"
+    "                callback(x[\"i\"], x[\"denoised\"], x[\"x\"], total_steps)\n"
+)
+if source.count(old) == 1:
+    source = source.replace(old, new)
+elif new not in source:
+    raise RuntimeError("Unable to locate Comfy sampler callback.")
+path.write_text(source, encoding="utf-8")
+print("Applied H3 sampler reference capture patch.")

tools/replay_comfy_history.py

@@ -0,0 +1,23 @@
+"""Replay an exact completed ComfyUI prompt by history id."""
+
+import argparse
+import json
+from urllib.request import Request, urlopen
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("history_id")
+parser.add_argument("--endpoint", default="http://localhost:8188")
+args = parser.parse_args()
+
+with urlopen(f"{args.endpoint}/history") as response:
+    history = json.load(response)
+
+prompt = history[args.history_id]["prompt"][2]
+request = Request(
+    f"{args.endpoint}/prompt",
+    data=json.dumps({"prompt": prompt}).encode(),
+    headers={"Content-Type": "application/json"},
+)
+with urlopen(request) as response:
+    print(response.read().decode())

tools/run_qwen3vl_text.py

@@ -0,0 +1,14 @@
+"""Run direct prompt-only Qwen3-VL H3 conditioning."""
+
+import torch
+
+from h3_blackwell_runtime.qwen3vl_text import Qwen3VLPromptConditioner
+
+
+conditioner = Qwen3VLPromptConditioner(
+    "/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors",
+    "/opt/h3-blackwell-runtime/qwen25_tokenizer",
+)
+states = conditioner("A brass-and-paper dragon flies above a rain-washed old city at blue hour.")
+torch.cuda.synchronize()
+print({"shape": tuple(states.shape), "dtype": str(states.dtype), "finite": torch.isfinite(states).all().item()})

tools/smoke_h3_packing.py

@@ -0,0 +1,12 @@
+"""Smoke-test direct H3 prompt-only packed input construction."""
+
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.packing import H3PromptPacker
+from h3_blackwell_runtime.t2v import empty_av_latents
+
+
+checkpoint = H3Checkpoint("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
+video, audio, frames = empty_av_latents(960, 544, 124)
+text = video.new_zeros((1, 8, 5120), dtype=video.dtype)
+hidden, timesteps, segments, positions, video_segment, audio_segment = H3PromptPacker(checkpoint)(text, video, audio, 1.0)
+print({"frames": frames, "hidden": tuple(hidden.shape), "timesteps": timesteps.tolist(), "positions": tuple(positions.shape), "segments": segments, "video_segment": video_segment, "audio_segment": audio_segment})

tools/smoke_h3_t2v_denoiser.py

@@ -0,0 +1,23 @@
+"""Run one direct FL2VA H3 denoiser call from raw prompt conditioning."""
+
+import torch
+
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.denoiser import H3PackedDenoiser
+from h3_blackwell_runtime.packing import H3PromptPacker
+from h3_blackwell_runtime.qwen3vl_text import Qwen3VLPromptConditioner
+from h3_blackwell_runtime.t2v import empty_av_latents
+
+
+prompt = "A brass-and-paper dragon flies above a rain-washed old city at blue hour."
+text_encoder = Qwen3VLPromptConditioner("/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", "/opt/h3-blackwell-runtime/qwen25_tokenizer")
+checkpoint = H3Checkpoint("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
+video, audio, _ = empty_av_latents(960, 544, 124)
+video.normal_()
+text = text_encoder(prompt)
+hidden, timesteps, segments, positions, video_segment, audio_segment = H3PromptPacker(checkpoint)(text, video, audio, 1.0)
+model = H3PackedDenoiser.from_checkpoint(checkpoint).eval()
+with torch.inference_mode():
+    velocity_video, velocity_audio = model(hidden, timesteps, positions, segments, video_segment, audio_segment)
+torch.cuda.synchronize()
+print({"video": tuple(velocity_video.shape), "audio": tuple(velocity_audio.shape), "finite": torch.isfinite(velocity_video).all().item() and torch.isfinite(velocity_audio).all().item()})

tools/smoke_h3_token_refiner.py

@@ -0,0 +1,11 @@
+import torch
+
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.qwen3vl_text import Qwen3VLPromptConditioner
+from h3_blackwell_runtime.token_refiner import H3TokenRefiner
+
+
+states = Qwen3VLPromptConditioner("/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", "/opt/h3-blackwell-runtime/qwen25_tokenizer")("A brass dragon flies through rain.")
+refined = H3TokenRefiner(H3Checkpoint("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors"))(states)
+torch.cuda.synchronize()
+print({"shape": tuple(refined.shape), "dtype": str(refined.dtype), "finite": torch.isfinite(refined).all().item()})

tools/smoke_h3_vae_decoder.py

@@ -0,0 +1,28 @@
+"""Validate and optionally execute the direct H3 video VAE decoder."""
+
+import argparse
+
+import torch
+
+from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("checkpoint", nargs="?", default="/vae/minimax_h3_video_vae_fp16.safetensors")
+    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    parser.add_argument("--validate-only", action="store_true")
+    args = parser.parse_args()
+
+    vae = MiniMaxH3VideoVAE.from_safetensors(args.checkpoint, device=args.device).eval()
+    print(f"validated checkpoint keys and shapes: {args.checkpoint}")
+    if args.validate_only:
+        return
+    with torch.inference_mode():
+        latent = torch.zeros(1, 24, 1, 1, 1, device=args.device, dtype=next(vae.parameters()).dtype)
+        decoded = vae.decode(latent)
+    print(f"decoded latent {tuple(latent.shape)} -> {tuple(decoded.shape)} ({decoded.dtype}, {decoded.device})")
+
+
+if __name__ == "__main__":
+    main()

tools/smoke_qwen3vl_text.py

@@ -0,0 +1,31 @@
+"""Spark smoke test for standalone Qwen3-VL prompt-only conditioning."""
+
+import argparse
+from pathlib import Path
+
+import torch
+
+from h3_blackwell_runtime.conditioning import H3PromptTokenizer
+from h3_blackwell_runtime.qwen3vl_text import Qwen3VL32BTextEncoder
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", type=Path)
+    parser.add_argument("--tokenizer-dir", type=Path, default=Path("src/h3_blackwell_runtime/qwen25_tokenizer"))
+    parser.add_argument("--prompt", default="A brass-and-paper dragon in a rainy clockmaker workshop.")
+    parser.add_argument("--construct-model", action="store_true")
+    args = parser.parse_args()
+
+    tokenizer = H3PromptTokenizer(args.tokenizer_dir)
+    token_ids = tokenizer(args.prompt, device="cuda" if torch.cuda.is_available() else "cpu")
+    print({"token_shape": tuple(token_ids.shape), "token_ids": token_ids[0].tolist()})
+    if args.construct_model:
+        if args.checkpoint is None:
+            parser.error("--construct-model requires --checkpoint")
+        model = Qwen3VL32BTextEncoder(args.checkpoint)
+        print({"layers": len(model.layers), "embedding_shape": tuple(model.embed_tokens.shape), "dtype": str(model.dtype)})
+
+
+if __name__ == "__main__":
+    main()

tools/submit_fl2va_reference.py

@@ -0,0 +1,22 @@
+"""Submit a compact FL2VA prompt-only sampler reference graph."""
+
+import json
+from urllib.request import Request, urlopen
+
+
+prompt = {
+    "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "minimax_h3_fl2va_pruned_nvfp4.safetensors", "weight_dtype": "default"}},
+    "3": {"class_type": "CLIPLoader", "inputs": {"clip_name": "qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", "type": "minimax"}},
+    "4": {"class_type": "VAELoader", "inputs": {"vae_name": "minimax_h3_video_vae_fp16.safetensors"}},
+    "8": {"class_type": "MiniMaxH3ImageToVideo", "inputs": {"clip": ["3", 0], "vae": ["4", 0], "prompt": "A brass-and-paper dragon flies above a rain-washed old city at blue hour.", "width": 320, "height": 192, "length": 22}},
+    "9": {"class_type": "BasicGuider", "inputs": {"model": ["1", 0], "conditioning": ["8", 0]}},
+    "10": {"class_type": "RandomNoise", "inputs": {"noise_seed": 440204}},
+    "11": {"class_type": "KSamplerSelect", "inputs": {"sampler_name": "res_multistep"}},
+    "12": {"class_type": "BasicScheduler", "inputs": {"model": ["1", 0], "scheduler": "beta", "steps": 12, "denoise": 1.0}},
+    "13": {"class_type": "SamplerCustomAdvanced", "inputs": {"noise": ["10", 0], "guider": ["9", 0], "sampler": ["11", 0], "sigmas": ["12", 0], "latent_image": ["8", 1]}},
+    "14": {"class_type": "VAEDecode", "inputs": {"samples": ["13", 0], "vae": ["4", 0]}},
+    "15": {"class_type": "SaveImage", "inputs": {"images": ["14", 0], "filename_prefix": "h3-blackwell-runtime/fl2va-reference"}},
+}
+request = Request("http://localhost:8188/prompt", data=json.dumps({"prompt": prompt}).encode(), headers={"Content-Type": "application/json"})
+with urlopen(request) as response:
+    print(response.read().decode())

tools/submit_h3_capture_workflow.py

@@ -0,0 +1,42 @@
+"""Submit a fixed 960x544 Ref2VA graph for direct-runner fixture capture."""
+
+import json
+import os
+from urllib.request import Request, urlopen
+
+
+prompt = {
+    "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "minimax_h3_ref2va_pruned_nvfp4.safetensors", "weight_dtype": "default"}},
+    "3": {"class_type": "CLIPLoader", "inputs": {"clip_name": "qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", "type": "minimax"}},
+    "4": {"class_type": "VAELoader", "inputs": {"vae_name": "minimax_h3_video_vae_fp16.safetensors"}},
+    "5": {"class_type": "VAELoader", "inputs": {"vae_name": "minimax_h3_audio_vae_fp32.safetensors"}},
+    "10": {"class_type": "RandomNoise", "inputs": {"noise_seed": int(os.environ.get("H3_CAPTURE_SEED", "440202"))}},
+    "11": {"class_type": "KSamplerSelect", "inputs": {"sampler_name": "res_multistep"}},
+    "12": {"class_type": "BasicScheduler", "inputs": {"model": ["1", 0], "scheduler": "beta", "steps": 12, "denoise": 1.0}},
+    "20": {"class_type": "LoadImage", "inputs": {"image": "episode-bees/01-desk_00001_.png"}},
+    "21": {"class_type": "LoadImage", "inputs": {"image": "episode-bees/02-launch_00001_.png"}},
+    "22": {"class_type": "LoadImage", "inputs": {"image": "episode-bees/03-flight_00001_.png"}},
+    "8": {
+        "class_type": "MiniMaxH3ReferenceToVideo",
+        "inputs": {
+            "clip": ["3", 0], "vae": ["4", 0], "audio_vae": ["5", 0],
+            "width": 960, "height": 544, "length": 124, "ref_image_size": "match",
+            "prompt": "Generate a continuous five-second storybook transition of a brass-and-paper dragon from a rainy clockmaker workshop to flight above a blue-hour city. Preserve the dragon identity and use clockwork, rain, wing flutter, and open night wind as the soundscape.",
+            "ref_images.ref_image_0": ["20", 0], "ref_images.ref_image_1": ["21", 0], "ref_images.ref_image_2": ["22", 0],
+        },
+    },
+    "9": {"class_type": "BasicGuider", "inputs": {"model": ["1", 0], "conditioning": ["8", 0]}},
+    "13": {"class_type": "SamplerCustomAdvanced", "inputs": {"noise": ["10", 0], "guider": ["9", 0], "sampler": ["11", 0], "sigmas": ["12", 0], "latent_image": ["8", 1]}},
+    "14": {"class_type": "VAEDecode", "inputs": {"samples": ["13", 0], "vae": ["4", 0]}},
+    "15": {"class_type": "VAEDecodeAudio", "inputs": {"samples": ["13", 0], "vae": ["5", 0]}},
+    "16": {"class_type": "CreateVideo", "inputs": {"images": ["14", 0], "audio": ["15", 0], "bit_depth": 8, "fps": 24.0}},
+    "17": {"class_type": "SaveVideo", "inputs": {"video": ["16", 0], "filename_prefix": "h3-blackwell-runtime/capture", "format": "mp4", "codec": "auto"}},
+}
+
+request = Request(
+    "http://localhost:8188/prompt",
+    data=json.dumps({"prompt": prompt}).encode(),
+    headers={"Content-Type": "application/json"},
+)
+with urlopen(request) as response:
+    print(response.read().decode())

tools/sweep_adaln_linear_precision.py

@@ -0,0 +1,27 @@
+"""Find the AdaLN linear precision boundary that matches ComfyUI."""
+
+import torch
+import torch.nn.functional as functional
+
+
+capture_dir = "/artifacts/capture"
+inputs = torch.load(f"{capture_dir}/input.pt", map_location="cuda", weights_only=False)
+expected = torch.load(f"{capture_dir}/block0_norm1_adaln.pt", map_location="cuda", weights_only=False)
+embedding = inputs["t_emb"]
+weight = expected["weight"]
+bias = expected["bias"]
+hidden = expected["shift"].shape[-1]
+
+variants = {
+    "fp32": (embedding.float(), weight.float(), bias.float()),
+    "bf16_weight_fp32": (embedding.float(), weight.bfloat16().float(), bias.float()),
+    "bf16_all": (embedding.bfloat16(), weight.bfloat16(), bias.bfloat16()),
+    "bf16_input_fp32": (embedding.bfloat16().float(), weight.float(), bias.float()),
+}
+
+for name, (x, w, b) in variants.items():
+    values = functional.linear(x, w, b).float().view(x.shape[0] * 3, 6 * hidden)
+    shift, scale, *_ = values.chunk(6, dim=-1)
+    shift_delta = (shift - expected["shift"].float()).abs()
+    scale_delta = (scale - expected["scale"].float()).abs()
+    print(f"{name} shift_mean={shift_delta.mean().item():.6g} shift_max={shift_delta.max().item():.6g} scale_mean={scale_delta.mean().item():.6g} scale_max={scale_delta.max().item():.6g}")

tools/validate_captured_denoiser.py

@@ -0,0 +1,33 @@
+"""Run the direct H3 core against one matched ComfyUI capture."""
+
+import time
+
+import torch
+
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.denoiser import H3PackedDenoiser
+
+
+capture_dir = "/artifacts/capture"
+inputs = torch.load(f"{capture_dir}/input.pt", map_location="cuda", weights_only=False)
+expected = torch.load(f"{capture_dir}/output.pt", map_location="cuda", weights_only=False)
+checkpoint = H3Checkpoint("/models/minimax_h3_ref2va_pruned_nvfp4.safetensors")
+
+start = time.perf_counter()
+model = H3PackedDenoiser.from_checkpoint(checkpoint).eval()
+with torch.inference_mode():
+    video, audio = model(
+        inputs["hidden"],
+        inputs["timesteps"],
+        inputs["position_ids"],
+        inputs["segments"],
+        expected["video_segment"],
+        expected["audio_segment"],
+    )
+torch.cuda.synchronize()
+
+for name, actual, reference in (("video", video, expected["video"]), ("audio", audio, expected["audio"])):
+    reference = reference.reshape_as(actual)
+    delta = (actual.float() - reference.float()).abs()
+    print(f"{name} shape={tuple(actual.shape)} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
+print(f"elapsed_s={time.perf_counter() - start:.3f}")