Compare commits
No commits in common. "3d0c093168246038ef287dab00e38023e226bacb" and "ea7712b4b2df960c909685b78c88c0209b2f5a8e" have entirely different histories.
3d0c093168
...
ea7712b4b2
17 changed files with 66 additions and 723 deletions
|
|
@ -1,8 +1,6 @@
|
|||
# GB10/Grace Blackwell development image. It does not start inference by default.
|
||||
FROM ghcr.io/aeon-7/comfyui-aeon-spark:slim
|
||||
|
||||
ARG SOL_ATTN_COMMIT=930a4d6e432ff8b8ed5e30ff2f72519b92d69bdf
|
||||
|
||||
WORKDIR /opt/h3-blackwell-runtime
|
||||
COPY . .
|
||||
|
||||
|
|
@ -15,18 +13,10 @@ RUN python -m pip install --no-cache-dir --no-deps comfy-kitchen==0.2.28
|
|||
|
||||
RUN python -m pip install --no-cache-dir "fastsafetensors>=0.1.10"
|
||||
|
||||
RUN python -m pip uninstall -y pynvml \
|
||||
&& python -m pip install --no-cache-dir nvidia-ml-py
|
||||
|
||||
RUN git clone https://github.com/Saganaki22/ComfyUI-sol-attn.git /opt/ComfyUI-sol-attn \
|
||||
&& cd /opt/ComfyUI-sol-attn \
|
||||
&& git checkout ${SOL_ATTN_COMMIT}
|
||||
|
||||
RUN python -m pip install --no-cache-dir --no-deps -e . \
|
||||
&& python -c "import comfy_kitchen, torch; from sageattn3 import sageattn3_blackwell; assert hasattr(torch.ops.comfy_kitchen, 'rms_rope_split_half_'); print(torch.__version__, torch.version.cuda)"
|
||||
|
||||
ENV H3_MODEL_PATH=/models/minimax_h3_ref2va_pruned_nvfp4.safetensors
|
||||
ENV PYTHONPATH=/opt/ComfyUI-sol-attn
|
||||
ENV TORCH_COMPILE_DISABLE=0 TORCHDYNAMO_DISABLE=0
|
||||
ENTRYPOINT []
|
||||
CMD ["bash"]
|
||||
|
|
|
|||
19
PARITY.md
19
PARITY.md
|
|
@ -187,9 +187,9 @@ gate and is exact.
|
|||
| H3 final layer | Yes | Bit-exact final AdaLN, target-row modulation, and video/audio patch rows; Comfy materializes the AdaLN and output-head biases through BF16 |
|
||||
| H3 DiT | Yes | Strict all-block numeric parity not achieved |
|
||||
| Beta/RES sampler | Yes | H3 replay is bit-exact for all 12 captured sampler calls and RES update is bit-exact from captured denoised; direct raw-output to model-wrapper denoised conversion retains small BF16/state boundary deltas |
|
||||
| Video VAE decoder | Yes | Direct VAE temporal assembly matches upstream after overlap fix. FP16 is the default Comfy-equivalent runtime path; cat benchmark VAE decode is `25.085s`. Use FP32 only for exact direct diagnostics. |
|
||||
| Audio VAE/decode/mux | Yes | Direct decoder-only MiniMax H3 audio VAE returns stereo `32000 Hz` waveform and muxes with generated video. Native audio latent scaling is fixed. |
|
||||
| End-to-end prompt-only FL2VA preview | Yes | Apples-to-apples warm cat benchmark is at Comfy parity: Comfy warm `150.26s`; direct warm after text conditioning `149.304s`; direct warm including text conditioning `151.465s`. |
|
||||
| Video VAE decoder | Yes | Same final sampled latent decoded through direct VAE versus Comfy PNG frames: mean abs pixel delta `0.0833898`, max `6` |
|
||||
| Audio VAE/decode/mux | No | Preview intentionally produces video only |
|
||||
| End-to-end prompt-only FL2VA preview | Yes | Direct 320x192x22 dragon smoke videos generated with Sage2. Comfy-equivalent CPU joint AV seed initialization matches captured `initial.pt` exactly. Video-side wrapper denoised conversion is exact after BF16 raw-output materialization; remaining seed-exact final-latent delta is localized to regenerated text conditioning plus audio-side wrapper conversion |
|
||||
| Full Ref2VA | No | References, vision conditioning, VAE encode, audio, and muxing are unimplemented |
|
||||
|
||||
## Remaining Gates, In Dependency Order
|
||||
|
|
@ -197,11 +197,14 @@ gate and is exact.
|
|||
Only these are outstanding. Do not recapture or revisit rows marked complete
|
||||
unless the checkpoint, Comfy version, prompt, or backend changes.
|
||||
|
||||
1. **Feature/performance work.** Prompt-only FL2VA is now closed against the
|
||||
warm Comfy baseline. Optimize load/caching/sampling, then evaluate Sage3,
|
||||
CUDA graphs, and multi-GPU execution.
|
||||
2. **Full Ref2VA support.** Add references, vision conditioning, VAE encode,
|
||||
and reference-path validation gates.
|
||||
1. **Final video latent layout.** Compare direct `unpatchify_video` output
|
||||
with Comfy before entering the sampler.
|
||||
2. **Text/audio end-to-end closure.** Direct H3, video wrapper conversion,
|
||||
RES update, VAE decode, and initial AV noise now match their captures.
|
||||
Close regenerated text conditioning and audio-side wrapper conversion before
|
||||
claiming identical seed output.
|
||||
3. **Feature/performance work.** Only then add audio, Ref2VA/reference paths,
|
||||
Sage3, CUDA graphs, and multi-GPU execution.
|
||||
|
||||
## Existing Tools And Their Intended Gate
|
||||
|
||||
|
|
|
|||
17
PLAN.md
17
PLAN.md
|
|
@ -48,21 +48,6 @@ The direct runner must first match the model contract and output quality. Beatin
|
|||
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.
|
||||
|
||||
## Performance Backend Plan
|
||||
|
||||
Prompt-only FL2VA is now at warm Comfy parity with the direct Sage2 baseline. Feature and performance work should proceed in this order:
|
||||
|
||||
1. Validate and benchmark the existing `sage3` backend against the same cat prompt, seed, dimensions, and FP16 VAE runtime path used for Sage2 parity.
|
||||
- First cat benchmark result: Sage3 runs successfully but is slower than Sage2 in this direct path. Sampling was `123.675s` versus Sage2 `114.414s`; warm after text conditioning was `158.111s` versus Sage2 `149.304s`. Same-seed MP4 frame diff versus Sage2 was mean `46.563`, max `255`, so keep Sage3 experimental pending human visual review and stricter tensor gates.
|
||||
2. Build a persistent hot runtime service instead of measuring only process-per-run CLIs. A warm container must preload and retain Qwen, H3, video VAE, and audio VAE in GPU memory, then accept video jobs without model-load latency. Add explicit startup warmup, readiness reporting, request-level attention selection (`sage2`, `sdpa`, `sage3` initially), and timing fields that separate resident-model request latency from cold startup.
|
||||
3. Add exact memory/lifetime optimizations next: `kj_head_sliced` and `kj_chunked_ffn`. These must preserve the validated direct outputs before being kept.
|
||||
4. Evaluate prior H3-tested attention candidates as standalone adapters: `sol_attn` and `kj_sage`.
|
||||
- `kj_sage` is implemented as explicit SageAttention mode backends: `kj_sage_cuda`, `kj_sage_triton`, `kj_sage_fp8`, and `kj_sage_fp8pp`; all passed hot-runtime smoke tests.
|
||||
- `sol_attn` is still blocked on locating/adding the standalone Sol-Attn source or package. It is not installed in the Spark image and is not present in this repository.
|
||||
5. Evaluate approximate denoiser caches only after exact baselines are recorded: `easycache` and `h3_cache`.
|
||||
- Initial direct cache modes are implemented as opt-in approximate sampler modes. They reuse cached denoised deltas and report skipped-step stats; full-size quality/threshold sweeps are still required before using them for production output.
|
||||
6. Keep every backend explicit per run, with separate quality and timing records for sampling, VAE, audio, and end-to-end output.
|
||||
|
||||
## Non-Negotiable Validation
|
||||
|
||||
- Never silently pad semantic H3 tokens for unmasked attention.
|
||||
|
|
@ -73,8 +58,6 @@ Prompt-only FL2VA is now at warm Comfy parity with the direct Sage2 baseline. Fe
|
|||
|
||||
## 2026-08-13 VAE Debug Handoff
|
||||
|
||||
Resolved on `vae-decode-optimization`: direct VAE temporal overlap constants now match upstream, audio decode/mux is implemented, and Comfy-equivalent FP16 video VAE is the default runtime path. The 960x544x124 cat benchmark now matches warm Comfy performance: Comfy `150.26s`, direct `149.304s` after text conditioning, direct VAE decode `25.085s`.
|
||||
|
||||
Current saved latent and comparison assets live under:
|
||||
|
||||
- Spark: `/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime`
|
||||
|
|
|
|||
24
README.md
24
README.md
|
|
@ -45,29 +45,5 @@ Generation and latent-decode tools are quiet by default: they suppress ffmpeg ba
|
|||
- `--profile-memory`: print memory checkpoints in `tools/direct_t2v_preview.py`.
|
||||
- `--ffmpeg-loglevel info`: show ffmpeg details instead of the default `error` level.
|
||||
- `--quiet`: suppress JSON summary lines.
|
||||
- `--vae-dtype float16`: use Comfy-style FP16 video VAE decode in `tools/direct_t2v_preview.py` or `tools/decode_video_latent.py`; this is the default runtime path. Use `--vae-dtype float32` only for exact direct-path diagnostics. `tools/direct_t2v_preview.py` also accepts `H3_VAE_DTYPE`.
|
||||
- `--vae-tile-size 256`: set the direct video VAE spatial tile size. `tools/direct_t2v_preview.py` also accepts `H3_VAE_TILE_SIZE`.
|
||||
|
||||
Standalone `tools/compare_*`, `tools/trace_*`, `tools/inspect_*`, and `tools/patch_comfy_*` scripts are debugging utilities and remain opt-in by being separate commands.
|
||||
|
||||
## Hot Runtime Service
|
||||
|
||||
`tools/serve_hot_runtime.py` keeps Qwen, H3, video VAE, and audio VAE resident in one process. Start the optional Spark service with:
|
||||
|
||||
```bash
|
||||
docker compose -f compose.spark.yml up -d h3-hot-runtime
|
||||
```
|
||||
|
||||
Use `GET /ready` to confirm resident model readiness. Use `POST /generate` with JSON fields like `prompt`, `output`, `width`, `height`, `frames`, `steps`, `seed`, and optional `attention`. Supported request-level attention values are reported by `/ready`; switching attention does not reload model weights.
|
||||
|
||||
Exact memory/lifetime options:
|
||||
|
||||
- `attention: "kj_head_sliced"` slices attention heads and runs the slice backend from `H3_HEAD_SLICE_BACKEND` (`sage2` by default) with `H3_HEAD_SLICE_SIZE` heads per slice (`8` by default).
|
||||
- `attention: "sol_attn"` routes eligible H3 attention calls through the pinned ComfyUI Sol-Attn Triton kernel vendored into the Spark image. Configure with `H3_SOL_TAU` (`1.3`), `H3_SOL_MIN_TOKENS` (`4096`), `H3_SOL_THRESH_TYPE` (`diag`), `H3_SOL_INT8_QK`, `H3_SOL_INT8_PV`, `H3_SOL_FALLBACK` (`sage2`), and `H3_SOL_STRICT`.
|
||||
- `--mlp-chunks N` on `tools/serve_hot_runtime.py` or `tools/direct_t2v_preview.py` chunks H3 SwiGLU rows exactly to reduce peak activation memory. Default is `1` (disabled).
|
||||
|
||||
Approximate cache options are opt-in and must be quality-gated per prompt:
|
||||
|
||||
- `cache_mode: "easycache"` reuses cached denoised deltas while cumulative latent input change stays below `cache_threshold`.
|
||||
- `cache_mode: "h3_cache"` reuses cached denoised deltas when the current per-step latent input change is below `cache_threshold`.
|
||||
- Both modes accept `cache_start_percent`, `cache_end_percent`, and `cache_subsample_factor` in `POST /generate`; the CLI exposes equivalent `--cache-*` flags.
|
||||
|
|
|
|||
|
|
@ -18,35 +18,5 @@
|
|||
"sampler": "res_multistep",
|
||||
"seed": 440202,
|
||||
"reference_comfy_sage3_seconds": 49.893,
|
||||
"measurement": "ComfyUI prompt execution time after warm-up",
|
||||
"prompt_only_fl2va_cat_benchmark": {
|
||||
"seed": 440407,
|
||||
"prompt": "A playful orange tabby cat starts in an ordinary cozy living room in a normal house, afternoon light, sofa and rug. The cat crouches, jumps, and does one clean athletic backflip in slow motion. As the backflip completes there is a sharp cinematic cut: the cat lands perfectly on a glowing neon disco dance floor wearing oversized black sunglasses. Mirror ball reflections, colorful lights, joyful party energy, stylish and funny, clear before-and-after transformation.",
|
||||
"comfy_warm_seconds": 150.26,
|
||||
"direct_warm_after_text_conditioned_seconds": 149.304,
|
||||
"direct_warm_including_text_conditioning_seconds": 151.465,
|
||||
"direct_cold_through_audio_decode_seconds": 180.494,
|
||||
"direct_video_vae_dtype": "float16",
|
||||
"direct_video_vae_tile_size": 256,
|
||||
"direct_video_vae_decode_seconds": 25.085,
|
||||
"direct_output": "\\\\192.168.1.162\\StoryStudioAssets\\H3-output\\h3-blackwell-runtime\\direct-cat-benchmark-960x544-124f-seed440407-fp16-vae.mp4",
|
||||
"direct_log": "\\\\192.168.1.162\\StoryStudioAssets\\H3-output\\h3-blackwell-runtime\\benchmarks\\direct-cat-benchmark-960x544-124f-seed440407-fp16-vae.log",
|
||||
"comfy_output": "\\\\192.168.1.162\\StoryStudioAssets\\H3-output\\h3-blackwell-runtime\\comfy-cat-benchmark-960x544-124f-seed440407_00001_.mp4",
|
||||
"comfy_metrics": "\\\\192.168.1.162\\StoryStudioAssets\\H3-output\\h3-blackwell-runtime\\benchmarks\\comfy-cat-benchmark-960x544-124f-seed440407.json"
|
||||
},
|
||||
"prompt_only_fl2va_cat_sage3_benchmark": {
|
||||
"seed": 440407,
|
||||
"attention_backend": "sage3",
|
||||
"direct_cold_through_audio_decode_seconds": 184.646,
|
||||
"direct_warm_after_text_conditioned_seconds": 158.111,
|
||||
"direct_warm_including_text_conditioning_seconds": 160.38,
|
||||
"direct_sample_seconds": 123.675,
|
||||
"direct_video_vae_dtype": "float16",
|
||||
"direct_video_vae_tile_size": 256,
|
||||
"direct_video_vae_decode_seconds": 25.024,
|
||||
"sage2_mp4_frame_diff_max": 255,
|
||||
"sage2_mp4_frame_diff_mean": 46.56300230273561,
|
||||
"direct_output": "\\\\192.168.1.162\\StoryStudioAssets\\H3-output\\h3-blackwell-runtime\\direct-cat-benchmark-960x544-124f-seed440407-sage3-fp16-vae.mp4",
|
||||
"direct_log": "\\\\192.168.1.162\\StoryStudioAssets\\H3-output\\h3-blackwell-runtime\\benchmarks\\direct-cat-benchmark-960x544-124f-seed440407-sage3-fp16-vae.log"
|
||||
}
|
||||
"measurement": "ComfyUI prompt execution time after warm-up"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -12,20 +12,3 @@ services:
|
|||
- /home/daniel/aeon-spark-test/h3/h3-blackwell-runtime/artifacts:/artifacts:ro
|
||||
- /home/daniel/StoryStudioAssets/H3-output:/output
|
||||
command: ["sleep", "infinity"]
|
||||
h3-hot-runtime:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile.spark
|
||||
image: h3-blackwell-runtime:dev
|
||||
gpus: all
|
||||
volumes:
|
||||
- /home/daniel/aeon-spark-test/h3/comfy-models/diffusion_models:/models:ro
|
||||
- /home/daniel/aeon-spark-test/h3/comfy-models/text_encoders:/text-encoders:ro
|
||||
- /home/daniel/aeon-spark-test/h3/comfy-models/vae:/vae:ro
|
||||
- /home/daniel/aeon-spark-test/h3/h3-blackwell-runtime/artifacts:/artifacts:ro
|
||||
- /home/daniel/StoryStudioAssets/H3-output:/output
|
||||
ports:
|
||||
- "8001:8000"
|
||||
environment:
|
||||
H3_DISABLE_MMAP: "1"
|
||||
command: ["python", "/opt/h3-blackwell-runtime/tools/serve_hot_runtime.py", "--host", "0.0.0.0", "--port", "8000", "--attention", "sage2", "--warmup"]
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
"""Direct H3 self-attention using packed NVFP4 linears and SageAttention3."""
|
||||
|
||||
import os
|
||||
import torch
|
||||
import torch.nn.functional as functional
|
||||
from torch import nn
|
||||
|
|
@ -9,18 +8,20 @@ from .checkpoint import H3Checkpoint
|
|||
from .nvfp4 import Nvfp4Linear
|
||||
|
||||
|
||||
AVAILABLE_BACKENDS = ("sage2", "sdpa", "sage3", "sage3_mean", "kj_sage_cuda", "kj_sage_triton", "kj_sage_fp8", "kj_sage_fp8pp", "kj_head_sliced", "sol_attn")
|
||||
PLANNED_BACKENDS = ("flash4", "easycache", "h3_cache", "kj_chunked_ffn")
|
||||
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({"sol_attn": "experimental: sparse Triton attention for eligible non-causal H3 attention calls; falls back below H3_SOL_MIN_TOKENS unless H3_SOL_STRICT=1"})
|
||||
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({"kj_chunked_ffn": "available: exact H3 MLP row chunking via H3_MLP_CHUNKS or runtime args"})
|
||||
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
|
||||
|
||||
|
||||
|
|
@ -30,45 +31,6 @@ def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
|
|||
|
||||
def run_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, backend: str, is_causal: bool) -> torch.Tensor:
|
||||
"""Run one `[batch, heads, sequence, dim]` attention operation."""
|
||||
if backend == "sol_attn":
|
||||
try:
|
||||
from sol_kernel import sol_attn
|
||||
|
||||
tau = float(os.getenv("H3_SOL_TAU", "1.3"))
|
||||
min_tokens = int(os.getenv("H3_SOL_MIN_TOKENS", "4096"))
|
||||
thresh_type = os.getenv("H3_SOL_THRESH_TYPE", "diag")
|
||||
int8_qk = os.getenv("H3_SOL_INT8_QK", "").lower() in {"1", "true", "yes", "on"}
|
||||
int8_pv = os.getenv("H3_SOL_INT8_PV", "").lower() in {"1", "true", "yes", "on"}
|
||||
if is_causal:
|
||||
raise ValueError("Sol-Attn backend only supports non-causal H3 attention")
|
||||
if q.shape[-1] != 128:
|
||||
raise ValueError(f"Sol-Attn requires head dim 128, got {q.shape[-1]}")
|
||||
if q.shape[2] < min_tokens:
|
||||
raise ValueError(f"{q.shape[2]} tokens < H3_SOL_MIN_TOKENS={min_tokens}")
|
||||
out = sol_attn(
|
||||
q.transpose(1, 2).contiguous(),
|
||||
k.transpose(1, 2).contiguous(),
|
||||
v.transpose(1, 2).contiguous(),
|
||||
tau=tau,
|
||||
thresh_type=thresh_type,
|
||||
int8_qk=int8_qk,
|
||||
int8_pv=int8_pv,
|
||||
)
|
||||
return out.transpose(1, 2)
|
||||
except Exception:
|
||||
if os.getenv("H3_SOL_STRICT", "").lower() in {"1", "true", "yes", "on"}:
|
||||
raise
|
||||
return run_attention(q, k, v, backend=os.getenv("H3_SOL_FALLBACK", "sage2"), is_causal=is_causal)
|
||||
if backend == "kj_head_sliced":
|
||||
head_slice_size = int(os.getenv("H3_HEAD_SLICE_SIZE", "8"))
|
||||
base_backend = os.getenv("H3_HEAD_SLICE_BACKEND", "sage2")
|
||||
if head_slice_size <= 0:
|
||||
raise ValueError("H3_HEAD_SLICE_SIZE must be positive")
|
||||
outputs = [
|
||||
run_attention(q[:, start:start + head_slice_size], k[:, start:start + head_slice_size], v[:, start:start + head_slice_size], backend=base_backend, is_causal=is_causal)
|
||||
for start in range(0, q.shape[1], head_slice_size)
|
||||
]
|
||||
return torch.cat(outputs, dim=1)
|
||||
if backend == "sage2":
|
||||
from sageattention import sageattn
|
||||
|
||||
|
|
@ -77,26 +39,6 @@ def run_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, backend:
|
|||
from sageattn3 import sageattn3_blackwell
|
||||
|
||||
return sageattn3_blackwell(q, k, v, is_causal=is_causal)
|
||||
if backend == "sage3_mean":
|
||||
from sageattn3 import sageattn3_blackwell
|
||||
|
||||
return sageattn3_blackwell(q, k, v, is_causal=is_causal, per_block_mean=True)
|
||||
if backend == "kj_sage_cuda":
|
||||
from sageattention import sageattn_qk_int8_pv_fp16_cuda
|
||||
|
||||
return sageattn_qk_int8_pv_fp16_cuda(q, k, v, is_causal=is_causal, pv_accum_dtype="fp32", tensor_layout="HND")
|
||||
if backend == "kj_sage_triton":
|
||||
from sageattention import sageattn_qk_int8_pv_fp16_triton
|
||||
|
||||
return sageattn_qk_int8_pv_fp16_triton(q, k, v, is_causal=is_causal, tensor_layout="HND")
|
||||
if backend == "kj_sage_fp8":
|
||||
from sageattention import sageattn_qk_int8_pv_fp8_cuda
|
||||
|
||||
return sageattn_qk_int8_pv_fp8_cuda(q, k, v, is_causal=is_causal, pv_accum_dtype="fp32+fp32", tensor_layout="HND")
|
||||
if backend == "kj_sage_fp8pp":
|
||||
from sageattention import sageattn_qk_int8_pv_fp8_cuda
|
||||
|
||||
return sageattn_qk_int8_pv_fp8_cuda(q, k, v, is_causal=is_causal, pv_accum_dtype="fp32+fp16", tensor_layout="HND")
|
||||
if backend == "sdpa":
|
||||
return functional.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
|
||||
raise ValueError(f"Unsupported H3 attention backend: {backend}")
|
||||
|
|
|
|||
|
|
@ -39,8 +39,6 @@ class H3SwiGLU(nn.Module):
|
|||
super().__init__()
|
||||
self.fc1 = fc1
|
||||
self.fc2 = fc2
|
||||
self.chunks = 1
|
||||
self.chunk_threshold = 4096
|
||||
|
||||
@classmethod
|
||||
def from_checkpoint(cls, checkpoint: H3Checkpoint, prefix: str, *, output_dtype=torch.bfloat16):
|
||||
|
|
@ -50,25 +48,10 @@ class H3SwiGLU(nn.Module):
|
|||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if self.chunks > 1 and x.shape[0] >= self.chunk_threshold:
|
||||
return torch.cat([self._forward_chunk(chunk) for chunk in x.chunk(self.chunks, dim=0)], dim=0)
|
||||
return self._forward_chunk(x)
|
||||
|
||||
def _forward_chunk(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))
|
||||
|
||||
|
||||
def configure_mlp_chunking(model: nn.Module, chunks: int, threshold: int = 4096) -> None:
|
||||
"""Configure exact row-chunked H3 SwiGLU execution to reduce peak activation memory."""
|
||||
if chunks < 1:
|
||||
raise ValueError("MLP chunks must be >= 1")
|
||||
for module in model.modules():
|
||||
if isinstance(module, H3SwiGLU):
|
||||
module.chunks = chunks
|
||||
module.chunk_threshold = threshold
|
||||
|
||||
|
||||
class H3DiTBlock(nn.Module):
|
||||
"""One H3 transformer block with externally supplied AdaLN tensors."""
|
||||
|
||||
|
|
|
|||
|
|
@ -1,266 +0,0 @@
|
|||
"""Resident prompt-only H3 runtime used by the hot service and tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from .audio_vae_decoder import MiniMaxH3AudioVAE
|
||||
from .attention import AVAILABLE_BACKENDS
|
||||
from .block import configure_mlp_chunking
|
||||
from .checkpoint import H3Checkpoint
|
||||
from .denoiser import H3PackedDenoiser
|
||||
from .packing import H3PromptPacker
|
||||
from .qwen3vl_text import Qwen3VLPromptConditioner
|
||||
from .sampler import sample_video_res_multistep
|
||||
from .t2v import random_av_latents
|
||||
from .token_refiner import H3TokenRefiner
|
||||
from .vae_decoder import MiniMaxH3VideoVAE, dtype_from_name
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RuntimeConfig:
|
||||
model_path: str = "/models/minimax_h3_fl2va_pruned_nvfp4.safetensors"
|
||||
text_encoder_path: str = "/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors"
|
||||
tokenizer_path: str = "/opt/h3-blackwell-runtime/src/h3_blackwell_runtime/qwen25_tokenizer"
|
||||
video_vae_path: str = "/vae/minimax_h3_video_vae_fp16.safetensors"
|
||||
audio_vae_path: str = "/vae/minimax_h3_audio_vae_fp32.safetensors"
|
||||
attention: str = "sage2"
|
||||
vae_dtype: str = "float16"
|
||||
vae_tile_size: int = 256
|
||||
mlp_chunks: int = 1
|
||||
mlp_chunk_threshold: int = 4096
|
||||
device: str = "cuda"
|
||||
|
||||
|
||||
def _sync() -> None:
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
|
||||
|
||||
def _ffmpeg_command(loglevel: str, *parts: str) -> list[str]:
|
||||
return ["ffmpeg", "-hide_banner", "-loglevel", loglevel, *parts]
|
||||
|
||||
|
||||
def _refiner_attention_backend(attention: str) -> str:
|
||||
if attention != "sol_attn":
|
||||
return attention
|
||||
fallback = os.getenv("H3_SOL_FALLBACK", "sage2")
|
||||
return fallback if fallback in AVAILABLE_BACKENDS and fallback != "sol_attn" else "sage2"
|
||||
|
||||
|
||||
class H3HotRuntime:
|
||||
"""Keep all prompt-only H3 models resident for repeated requests."""
|
||||
|
||||
def __init__(self, config: RuntimeConfig):
|
||||
self.config = config
|
||||
self.attention = config.attention
|
||||
self.loaded_at = time.time()
|
||||
self.load_stages: list[dict] = []
|
||||
|
||||
self.checkpoint = H3Checkpoint(config.model_path, device=config.device)
|
||||
self.conditioner = self._timed_load(
|
||||
"qwen_loaded",
|
||||
lambda: Qwen3VLPromptConditioner(config.text_encoder_path, config.tokenizer_path),
|
||||
)
|
||||
self.model = self._timed_load(
|
||||
"h3_loaded",
|
||||
self._load_h3,
|
||||
)
|
||||
self.refiner = self._timed_load(
|
||||
"token_refiner_loaded",
|
||||
lambda: H3TokenRefiner(self.checkpoint, attention_backend=_refiner_attention_backend(config.attention)).eval(),
|
||||
)
|
||||
self.packer = H3PromptPacker(self.checkpoint)
|
||||
self.video_vae = self._timed_load(
|
||||
"video_vae_loaded",
|
||||
self._load_video_vae,
|
||||
)
|
||||
self.audio_vae = self._timed_load(
|
||||
"audio_vae_loaded",
|
||||
lambda: MiniMaxH3AudioVAE.from_safetensors(config.audio_vae_path, device=config.device).eval(),
|
||||
)
|
||||
|
||||
def _timed_load(self, stage: str, fn):
|
||||
_sync()
|
||||
start = time.perf_counter()
|
||||
value = fn()
|
||||
_sync()
|
||||
self.load_stages.append({"stage": stage, "seconds": time.perf_counter() - start})
|
||||
return value
|
||||
|
||||
def _load_video_vae(self) -> MiniMaxH3VideoVAE:
|
||||
vae = MiniMaxH3VideoVAE.from_safetensors(
|
||||
self.config.video_vae_path,
|
||||
device=self.config.device,
|
||||
dtype=dtype_from_name(self.config.vae_dtype),
|
||||
).eval()
|
||||
vae.tile_size = self.config.vae_tile_size
|
||||
return vae
|
||||
|
||||
def _load_h3(self) -> H3PackedDenoiser:
|
||||
model = H3PackedDenoiser.from_checkpoint(self.checkpoint, attention_backend=self.config.attention).eval()
|
||||
configure_mlp_chunking(model, self.config.mlp_chunks, self.config.mlp_chunk_threshold)
|
||||
return model
|
||||
|
||||
def status(self) -> dict:
|
||||
return {
|
||||
"ready": True,
|
||||
"initial_attention": self.config.attention,
|
||||
"current_attention": self.attention,
|
||||
"vae_dtype": self.config.vae_dtype,
|
||||
"vae_tile_size": self.config.vae_tile_size,
|
||||
"mlp_chunks": self.config.mlp_chunks,
|
||||
"mlp_chunk_threshold": self.config.mlp_chunk_threshold,
|
||||
"loaded_at": self.loaded_at,
|
||||
"load_stages": self.load_stages,
|
||||
}
|
||||
|
||||
@torch.inference_mode()
|
||||
def set_attention(self, attention: str) -> None:
|
||||
if attention not in AVAILABLE_BACKENDS:
|
||||
raise ValueError(f"Unsupported attention backend: {attention}")
|
||||
if attention == self.attention:
|
||||
return
|
||||
for module in self.model.modules():
|
||||
if hasattr(module, "backend"):
|
||||
module.backend = attention
|
||||
for block in self.refiner.blocks:
|
||||
block.attention_backend = _refiner_attention_backend(attention)
|
||||
self.attention = attention
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate(
|
||||
self,
|
||||
*,
|
||||
prompt: str,
|
||||
output: str | Path,
|
||||
width: int,
|
||||
height: int,
|
||||
frames: int,
|
||||
steps: int,
|
||||
seed: int,
|
||||
attention: str | None = None,
|
||||
mux_audio: bool = True,
|
||||
ffmpeg_loglevel: str = "error",
|
||||
save_latent: str | Path | None = None,
|
||||
cache_mode: str | None = None,
|
||||
cache_threshold: float = 0.0,
|
||||
cache_start_percent: float = 0.0,
|
||||
cache_end_percent: float = 1.0,
|
||||
cache_subsample_factor: int = 2,
|
||||
) -> dict:
|
||||
stages: list[dict] = []
|
||||
cache_stats: dict = {}
|
||||
|
||||
def timed(stage: str, fn):
|
||||
_sync()
|
||||
start = time.perf_counter()
|
||||
value = fn()
|
||||
_sync()
|
||||
stages.append({"stage": stage, "seconds": time.perf_counter() - start})
|
||||
return value
|
||||
|
||||
output = Path(output)
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
if attention is not None:
|
||||
self.set_attention(attention)
|
||||
|
||||
video, audio, aligned_frames = timed(
|
||||
"latents_initialized",
|
||||
lambda: random_av_latents(width, height, frames, seed, device=self.config.device),
|
||||
)
|
||||
text = timed("text_conditioned", lambda: self.refiner(self.conditioner(prompt)))
|
||||
sampled = timed(
|
||||
"sampled",
|
||||
lambda: sample_video_res_multistep(
|
||||
self.model,
|
||||
self.packer,
|
||||
text,
|
||||
video,
|
||||
audio,
|
||||
steps=steps,
|
||||
return_audio=mux_audio,
|
||||
cache_mode=cache_mode,
|
||||
cache_threshold=cache_threshold,
|
||||
cache_start_percent=cache_start_percent,
|
||||
cache_end_percent=cache_end_percent,
|
||||
cache_subsample_factor=cache_subsample_factor,
|
||||
cache_stats=cache_stats,
|
||||
),
|
||||
)
|
||||
if mux_audio:
|
||||
latent, audio_latent = sampled
|
||||
else:
|
||||
latent, audio_latent = sampled, None
|
||||
if save_latent is not None:
|
||||
latent_path = Path(save_latent)
|
||||
latent_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
state = {"latent": latent.detach().cpu(), "frames": aligned_frames, "width": width, "height": height, "prompt": prompt, "seed": seed}
|
||||
if audio_latent is not None:
|
||||
state["audio_latent"] = audio_latent.detach().cpu()
|
||||
torch.save(state, latent_path)
|
||||
pixels = timed("vae_decoded", lambda: self.video_vae.decode(latent.to(next(self.video_vae.parameters()).dtype))[:, :, :aligned_frames])
|
||||
pixels = timed("pixels_cpu", lambda: ((pixels[0].permute(1, 2, 3, 0).clamp(-1, 1) + 1) * 127.5).to(torch.uint8).cpu())
|
||||
|
||||
raw = output.with_suffix(".rgb")
|
||||
video_output = output.with_name(output.stem + ".video.mp4") if mux_audio else output
|
||||
timed("raw_write", lambda: pixels.numpy().tofile(raw))
|
||||
timed(
|
||||
"video_encode",
|
||||
lambda: subprocess.run(
|
||||
_ffmpeg_command(
|
||||
ffmpeg_loglevel,
|
||||
"-y",
|
||||
"-f",
|
||||
"rawvideo",
|
||||
"-pixel_format",
|
||||
"rgb24",
|
||||
"-video_size",
|
||||
f"{pixels.shape[2]}x{pixels.shape[1]}",
|
||||
"-framerate",
|
||||
"24",
|
||||
"-i",
|
||||
str(raw),
|
||||
"-an",
|
||||
"-c:v",
|
||||
"libx264",
|
||||
"-pix_fmt",
|
||||
"yuv420p",
|
||||
str(video_output),
|
||||
),
|
||||
check=True,
|
||||
),
|
||||
)
|
||||
raw.unlink()
|
||||
|
||||
audio_output = None
|
||||
if mux_audio:
|
||||
audio_output = output.with_suffix(".wav")
|
||||
waveform = timed("audio_decoded", lambda: self.audio_vae.decode(audio_latent.to(next(self.audio_vae.parameters()).dtype)).clamp(-1, 1).cpu()[0])
|
||||
audio_raw = audio_output.with_suffix(".f32le")
|
||||
timed("audio_raw_write", lambda: waveform.transpose(0, 1).contiguous().numpy().tofile(audio_raw))
|
||||
timed("audio_encode", lambda: subprocess.run(_ffmpeg_command(ffmpeg_loglevel, "-y", "-f", "f32le", "-ar", "32000", "-ac", "2", "-i", str(audio_raw), str(audio_output)), check=True))
|
||||
audio_raw.unlink()
|
||||
timed("mux", lambda: subprocess.run(_ffmpeg_command(ffmpeg_loglevel, "-y", "-i", str(video_output), "-i", str(audio_output), "-c:v", "copy", "-c:a", "aac", "-shortest", str(output)), check=True))
|
||||
video_output.unlink()
|
||||
|
||||
return {
|
||||
"output": str(output),
|
||||
"audio_output": str(audio_output) if audio_output is not None else None,
|
||||
"frames": aligned_frames,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"seed": seed,
|
||||
"attention": self.attention,
|
||||
"vae_dtype": self.config.vae_dtype,
|
||||
"vae_tile_size": self.config.vae_tile_size,
|
||||
"stages": stages,
|
||||
"cache": cache_stats,
|
||||
"request_seconds": sum(stage["seconds"] for stage in stages),
|
||||
}
|
||||
|
|
@ -57,14 +57,6 @@ def _decode_audio_latent(audio_carried: torch.Tensor, *, shift_video: float = 12
|
|||
return audio_carried * (shift_audio / shift_video)
|
||||
|
||||
|
||||
def _cache_sample(x: torch.Tensor, factor: int) -> torch.Tensor:
|
||||
if factor <= 1:
|
||||
return x
|
||||
if x.ndim == 5:
|
||||
return x[..., ::factor, ::factor]
|
||||
return x[..., ::factor]
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def sample_video_res_multistep(
|
||||
model,
|
||||
|
|
@ -77,12 +69,6 @@ def sample_video_res_multistep(
|
|||
model_timesteps: list[torch.Tensor] | tuple[torch.Tensor, ...] | None = None,
|
||||
return_audio: bool = False,
|
||||
progress: bool = False,
|
||||
cache_mode: str | None = None,
|
||||
cache_threshold: float = 0.0,
|
||||
cache_start_percent: float = 0.0,
|
||||
cache_end_percent: float = 1.0,
|
||||
cache_subsample_factor: int = 2,
|
||||
cache_stats: dict | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Direct H3 beta/RES sampling with Comfy-equivalent joint AV carry semantics."""
|
||||
sigmas = beta_sigmas(steps, device=video.device)
|
||||
|
|
@ -91,47 +77,10 @@ def sample_video_res_multistep(
|
|||
video_history_sigma = audio_history_sigma = None
|
||||
total_steps = len(sigmas) - 1
|
||||
started = time.perf_counter()
|
||||
cache_mode = None if cache_mode in {None, "", "disabled", "none"} else cache_mode
|
||||
if cache_mode not in {None, "easycache", "h3_cache"}:
|
||||
raise ValueError(f"Unsupported cache mode: {cache_mode}")
|
||||
if cache_stats is not None:
|
||||
cache_stats.update({"mode": cache_mode, "threshold": cache_threshold, "skipped_steps": 0, "rates": []})
|
||||
cache = {
|
||||
"video_diff": None,
|
||||
"audio_diff": None,
|
||||
"video_prev": None,
|
||||
"audio_prev": None,
|
||||
"prev_norm": None,
|
||||
"cumulative_rate": 0.0,
|
||||
}
|
||||
for index, sigma in enumerate(sigmas[:-1], start=1):
|
||||
step_started = time.perf_counter()
|
||||
previous_index = index - 1
|
||||
sigma_down = sigmas[index]
|
||||
current_percent = previous_index / total_steps
|
||||
can_cache = cache_mode is not None and cache_threshold > 0 and cache_start_percent <= current_percent <= cache_end_percent and cache["video_diff"] is not None
|
||||
skipped = False
|
||||
if can_cache:
|
||||
video_now = _cache_sample(video, cache_subsample_factor)
|
||||
audio_now = _cache_sample(audio_carried, cache_subsample_factor)
|
||||
input_change = (video_now - cache["video_prev"]).flatten().abs().mean() + (audio_now - cache["audio_prev"]).flatten().abs().mean()
|
||||
input_norm = cache["prev_norm"].clamp_min(1e-8)
|
||||
rate = (input_change / input_norm).item()
|
||||
if cache_mode == "easycache":
|
||||
cache["cumulative_rate"] += rate
|
||||
skipped = cache["cumulative_rate"] < cache_threshold
|
||||
if not skipped:
|
||||
cache["cumulative_rate"] = 0.0
|
||||
else:
|
||||
skipped = rate < cache_threshold
|
||||
if cache_stats is not None:
|
||||
cache_stats["rates"].append({"step": previous_index, "rate": rate, "skipped": skipped})
|
||||
if skipped:
|
||||
video_denoised = video + cache["video_diff"]
|
||||
audio_denoised = audio_carried + cache["audio_diff"]
|
||||
if cache_stats is not None:
|
||||
cache_stats["skipped_steps"] += 1
|
||||
else:
|
||||
sigma_audio = _audio_sigma(sigma)
|
||||
carry = sigma_audio / sigma
|
||||
native_audio = audio_carried.to(torch.bfloat16) * carry
|
||||
|
|
@ -147,12 +96,6 @@ def sample_video_res_multistep(
|
|||
).float()
|
||||
video_denoised = video - sigma * velocity_video
|
||||
audio_denoised = audio_carried - sigma * velocity_audio
|
||||
if cache_mode is not None:
|
||||
cache["video_diff"] = (video_denoised - video).detach()
|
||||
cache["audio_diff"] = (audio_denoised - audio_carried).detach()
|
||||
cache["video_prev"] = _cache_sample(video, cache_subsample_factor).detach().clone()
|
||||
cache["audio_prev"] = _cache_sample(audio_carried, cache_subsample_factor).detach().clone()
|
||||
cache["prev_norm"] = video.flatten().abs().mean() + audio_carried.flatten().abs().mean()
|
||||
previous_sigma = sigmas[previous_index - 1] if previous_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)
|
||||
|
|
|
|||
|
|
@ -23,14 +23,6 @@ LATENTS_MEAN = (0.858090341091156, -0.9606591463088989, 1.0661640167236328, -0.5
|
|||
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 dtype_from_name(name: str) -> torch.dtype:
|
||||
return {
|
||||
"float32": torch.float32,
|
||||
"float16": torch.float16,
|
||||
"bfloat16": torch.bfloat16,
|
||||
}[name]
|
||||
|
||||
|
||||
def _rms_norm(x: torch.Tensor, weight: torch.Tensor | None, eps: float) -> torch.Tensor:
|
||||
if os.getenv("H3_VAE_FAST_OPS", "").lower() in {"1", "true", "yes", "on"}:
|
||||
if weight is None:
|
||||
|
|
@ -206,7 +198,7 @@ class MiniMaxH3VideoVAE(nn.Module):
|
|||
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, dtype: torch.dtype = torch.float16) -> "MiniMaxH3VideoVAE":
|
||||
def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", tiling: bool = True, dtype: torch.dtype = torch.float32) -> "MiniMaxH3VideoVAE":
|
||||
model = cls(device="meta", tiling=tiling)
|
||||
expected = model.state_dict()
|
||||
if os.getenv("H3_FAST_SAFETENSORS", "").lower() in {"1", "true", "yes", "on"}:
|
||||
|
|
|
|||
|
|
@ -10,8 +10,6 @@ from pathlib import Path
|
|||
|
||||
import torch
|
||||
|
||||
from h3_blackwell_runtime.vae_decoder import dtype_from_name
|
||||
|
||||
|
||||
def sync() -> None:
|
||||
if torch.cuda.is_available():
|
||||
|
|
@ -28,6 +26,14 @@ def timed(stage: str, rows: list[dict], fn):
|
|||
return value
|
||||
|
||||
|
||||
def dtype_from_name(name: str) -> torch.dtype:
|
||||
return {
|
||||
"float32": torch.float32,
|
||||
"float16": torch.float16,
|
||||
"bfloat16": torch.bfloat16,
|
||||
}[name]
|
||||
|
||||
|
||||
def ffmpeg_command(loglevel: str, *parts: str) -> list[str]:
|
||||
return ["ffmpeg", "-hide_banner", "-loglevel", loglevel, *parts]
|
||||
|
||||
|
|
@ -38,7 +44,7 @@ def main() -> None:
|
|||
parser.add_argument("--output", type=Path)
|
||||
parser.add_argument("--metrics", type=Path, required=True)
|
||||
parser.add_argument("--implementation", choices=("direct", "upstream"), default="direct")
|
||||
parser.add_argument("--dtype", choices=("float32", "float16", "bfloat16"), default="float16")
|
||||
parser.add_argument("--dtype", choices=("float32", "float16", "bfloat16"), default="float32")
|
||||
parser.add_argument("--no-tiling", action="store_true")
|
||||
parser.add_argument("--tile-size", type=int)
|
||||
parser.add_argument("--tile-overlap", type=int)
|
||||
|
|
|
|||
|
|
@ -4,8 +4,8 @@ import argparse
|
|||
import time
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as functional
|
||||
|
||||
from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, run_attention
|
||||
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
||||
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
|
||||
|
||||
|
|
@ -21,14 +21,20 @@ q, k, v = payload["q"], payload["k"], payload["v"]
|
|||
model = H3PackedDenoiser.from_checkpoint(H3Checkpoint(args.model)).eval()
|
||||
out_proj = model.backbone.blocks[0].attention.out_proj
|
||||
|
||||
for name in AVAILABLE_BACKENDS:
|
||||
for name in ("sdpa", "sage2", "sage3"):
|
||||
torch.cuda.synchronize()
|
||||
start = time.perf_counter()
|
||||
try:
|
||||
output = run_attention(q, k, v, backend=name, is_causal=False)
|
||||
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}")
|
||||
except Exception as exc:
|
||||
print(f"{name} error={type(exc).__name__}: {exc}")
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@ import argparse
|
|||
import torch
|
||||
|
||||
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
||||
from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS
|
||||
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
|
||||
from h3_blackwell_runtime.packing import H3PromptPacker, unpatchify_video
|
||||
from h3_blackwell_runtime.sampler import _audio_sigma, _model_sigma, _unpack_audio, res_multistep_update
|
||||
|
|
@ -19,7 +18,7 @@ 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("--seed", type=int, default=440204)
|
||||
parser.add_argument("--attention", choices=AVAILABLE_BACKENDS, default="sage2")
|
||||
parser.add_argument("--attention", choices=("sage2", "sdpa", "sage3"), default="sage2")
|
||||
parser.add_argument("--oracle-timesteps", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from pathlib import Path
|
|||
|
||||
import torch
|
||||
|
||||
from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE, dtype_from_name
|
||||
from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
|
|
@ -17,8 +17,6 @@ parser.add_argument("--implementation", choices=("direct", "upstream"), default=
|
|||
parser.add_argument("--frames-dir", type=Path)
|
||||
parser.add_argument("--ffmpeg-loglevel", default="error")
|
||||
parser.add_argument("--quiet", action="store_true")
|
||||
parser.add_argument("--vae-dtype", choices=("float32", "float16", "bfloat16"), default="float16")
|
||||
parser.add_argument("--vae-tile-size", type=int, default=256)
|
||||
args = parser.parse_args()
|
||||
|
||||
state = torch.load(args.latent, map_location="cuda", weights_only=False)
|
||||
|
|
@ -30,9 +28,7 @@ if args.implementation == "direct":
|
|||
"/vae/minimax_h3_video_vae_fp16.safetensors",
|
||||
device="cuda",
|
||||
tiling=not args.no_tiling,
|
||||
dtype=dtype_from_name(args.vae_dtype),
|
||||
).eval()
|
||||
vae.tile_size = args.vae_tile_size
|
||||
else:
|
||||
from safetensors.torch import load_file
|
||||
|
||||
|
|
@ -40,10 +36,7 @@ else:
|
|||
|
||||
vae = UpstreamMiniMaxH3VideoVAE(tiling=not args.no_tiling).to("cuda").eval()
|
||||
checkpoint = load_file("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda")
|
||||
checkpoint = {key: value.to(dtype=dtype_from_name(args.vae_dtype)) for key, value in checkpoint.items()}
|
||||
missing, unexpected = vae.load_state_dict(checkpoint, strict=False)
|
||||
vae = vae.to(dtype=dtype_from_name(args.vae_dtype))
|
||||
vae.tile_size = args.vae_tile_size
|
||||
if not args.quiet:
|
||||
print({"upstream_missing": len(missing), "upstream_unexpected": len(unexpected)}, flush=True)
|
||||
with torch.inference_mode():
|
||||
|
|
@ -68,4 +61,4 @@ subprocess.run([
|
|||
], check=True)
|
||||
raw.unlink()
|
||||
if not args.quiet:
|
||||
print({"output": str(args.output), "frames": frames, "shape": tuple(pixels.shape), "tiling": not args.no_tiling, "implementation": args.implementation, "vae_dtype": args.vae_dtype, "vae_tile_size": vae.tile_size})
|
||||
print({"output": str(args.output), "frames": frames, "shape": tuple(pixels.shape), "tiling": not args.no_tiling, "implementation": args.implementation})
|
||||
|
|
|
|||
|
|
@ -14,15 +14,13 @@ import torch
|
|||
|
||||
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
||||
from h3_blackwell_runtime.audio_vae_decoder import MiniMaxH3AudioVAE
|
||||
from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS
|
||||
from h3_blackwell_runtime.block import configure_mlp_chunking
|
||||
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 random_av_latents
|
||||
from h3_blackwell_runtime.token_refiner import H3TokenRefiner
|
||||
from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE, dtype_from_name
|
||||
from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
|
|
@ -33,7 +31,7 @@ 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=AVAILABLE_BACKENDS, default="sage2")
|
||||
parser.add_argument("--attention", choices=("sage2", "sdpa", "sage3"), default="sage2")
|
||||
parser.add_argument("--model-timesteps-capture", type=Path, help="Directory containing captured input_XX.pt H3 timesteps for strict parity checks.")
|
||||
parser.add_argument("--progress", action="store_true", help="Print per-step sampler progress.")
|
||||
parser.add_argument("--profile-memory", action="store_true")
|
||||
|
|
@ -44,15 +42,6 @@ parser.add_argument("--save-audio-latent", type=Path)
|
|||
parser.add_argument("--audio-output", type=Path)
|
||||
parser.add_argument("--mux-audio", action="store_true")
|
||||
parser.add_argument("--skip-decode", action="store_true")
|
||||
parser.add_argument("--vae-dtype", choices=("float32", "float16", "bfloat16"), default=os.getenv("H3_VAE_DTYPE", "float16"))
|
||||
parser.add_argument("--vae-tile-size", type=int, default=int(os.getenv("H3_VAE_TILE_SIZE", "256")))
|
||||
parser.add_argument("--mlp-chunks", type=int, default=int(os.getenv("H3_MLP_CHUNKS", "1")))
|
||||
parser.add_argument("--mlp-chunk-threshold", type=int, default=int(os.getenv("H3_MLP_CHUNK_THRESHOLD", "4096")))
|
||||
parser.add_argument("--cache-mode", choices=("disabled", "easycache", "h3_cache"), default="disabled")
|
||||
parser.add_argument("--cache-threshold", type=float, default=0.0)
|
||||
parser.add_argument("--cache-start-percent", type=float, default=0.0)
|
||||
parser.add_argument("--cache-end-percent", type=float, default=1.0)
|
||||
parser.add_argument("--cache-subsample-factor", type=int, default=2)
|
||||
args = parser.parse_args()
|
||||
started = time.perf_counter()
|
||||
last_report = started
|
||||
|
|
@ -95,7 +84,6 @@ conditioner = Qwen3VLPromptConditioner(
|
|||
report_memory("qwen_loaded")
|
||||
video, audio, frames = random_av_latents(args.width, args.height, args.frames, args.seed)
|
||||
model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend=args.attention).eval()
|
||||
configure_mlp_chunking(model, args.mlp_chunks, args.mlp_chunk_threshold)
|
||||
report_memory("h3_loaded")
|
||||
text = H3TokenRefiner(checkpoint, attention_backend=args.attention)(conditioner(args.prompt))
|
||||
report_memory("text_conditioned")
|
||||
|
|
@ -106,26 +94,7 @@ if args.model_timesteps_capture is not None:
|
|||
for index in range(args.steps)
|
||||
]
|
||||
want_audio = args.save_audio_latent is not None or args.audio_output is not None or args.mux_audio
|
||||
cache_stats = {}
|
||||
sampled = sample_video_res_multistep(
|
||||
model,
|
||||
H3PromptPacker(checkpoint),
|
||||
text,
|
||||
video,
|
||||
audio,
|
||||
steps=args.steps,
|
||||
model_timesteps=model_timesteps,
|
||||
return_audio=want_audio,
|
||||
progress=args.progress,
|
||||
cache_mode=args.cache_mode,
|
||||
cache_threshold=args.cache_threshold,
|
||||
cache_start_percent=args.cache_start_percent,
|
||||
cache_end_percent=args.cache_end_percent,
|
||||
cache_subsample_factor=args.cache_subsample_factor,
|
||||
cache_stats=cache_stats,
|
||||
)
|
||||
if cache_stats:
|
||||
report({"cache": cache_stats})
|
||||
sampled = sample_video_res_multistep(model, H3PromptPacker(checkpoint), text, video, audio, steps=args.steps, model_timesteps=model_timesteps, return_audio=want_audio, progress=args.progress)
|
||||
if want_audio:
|
||||
latent, audio_latent = sampled
|
||||
else:
|
||||
|
|
@ -146,8 +115,7 @@ if args.save_audio_latent is not None:
|
|||
report({"audio_latent": str(args.save_audio_latent)})
|
||||
if args.skip_decode:
|
||||
raise SystemExit(0)
|
||||
vae = MiniMaxH3VideoVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda", dtype=dtype_from_name(args.vae_dtype)).eval()
|
||||
vae.tile_size = args.vae_tile_size
|
||||
vae = MiniMaxH3VideoVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda").eval()
|
||||
report_memory("vae_loaded")
|
||||
with torch.inference_mode():
|
||||
pixels = vae.decode(latent.to(next(vae.parameters()).dtype))[:, :, :frames]
|
||||
|
|
|
|||
|
|
@ -1,128 +0,0 @@
|
|||
"""Serve a resident prompt-only H3 runtime over a small JSON HTTP API."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS
|
||||
from h3_blackwell_runtime.runtime import H3HotRuntime, RuntimeConfig
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", default="0.0.0.0")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--attention", choices=AVAILABLE_BACKENDS, default="sage2", help="Initial attention backend. Requests can switch with the JSON attention field.")
|
||||
parser.add_argument("--vae-dtype", choices=("float32", "float16", "bfloat16"), default="float16")
|
||||
parser.add_argument("--vae-tile-size", type=int, default=256)
|
||||
parser.add_argument("--mlp-chunks", type=int, default=1)
|
||||
parser.add_argument("--mlp-chunk-threshold", type=int, default=4096)
|
||||
parser.add_argument("--warmup", action="store_true", help="Run a tiny generation before accepting traffic.")
|
||||
parser.add_argument("--warmup-output", type=Path, default=Path("/output/h3-blackwell-runtime/hot-runtime-warmup.mp4"))
|
||||
args = parser.parse_args()
|
||||
|
||||
runtime = H3HotRuntime(RuntimeConfig(attention=args.attention, vae_dtype=args.vae_dtype, vae_tile_size=args.vae_tile_size, mlp_chunks=args.mlp_chunks, mlp_chunk_threshold=args.mlp_chunk_threshold))
|
||||
runtime_lock = threading.Lock()
|
||||
warmup_result = None
|
||||
if args.warmup:
|
||||
warmup_result = runtime.generate(
|
||||
prompt="A small warmup cat blinks in soft light.",
|
||||
output=args.warmup_output,
|
||||
width=320,
|
||||
height=192,
|
||||
frames=22,
|
||||
steps=2,
|
||||
seed=440501,
|
||||
mux_audio=True,
|
||||
)
|
||||
|
||||
|
||||
def service_status() -> dict:
|
||||
return {
|
||||
"ready": True,
|
||||
"attention_backends": list(AVAILABLE_BACKENDS),
|
||||
"runtime": runtime.status(),
|
||||
}
|
||||
|
||||
|
||||
def write_json(handler: BaseHTTPRequestHandler, status: int, payload: dict) -> None:
|
||||
body = json.dumps(payload, indent=2).encode("utf-8")
|
||||
handler.send_response(status)
|
||||
handler.send_header("Content-Type", "application/json")
|
||||
handler.send_header("Content-Length", str(len(body)))
|
||||
handler.end_headers()
|
||||
handler.wfile.write(body)
|
||||
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def log_message(self, format: str, *args) -> None:
|
||||
return
|
||||
|
||||
def do_GET(self) -> None:
|
||||
path = urlparse(self.path).path
|
||||
if path in {"/health", "/ready"}:
|
||||
write_json(self, 200, {**service_status(), "warmup_result": warmup_result})
|
||||
return
|
||||
write_json(self, 404, {"error": "not found"})
|
||||
|
||||
def do_POST(self) -> None:
|
||||
path = urlparse(self.path).path
|
||||
if path != "/generate":
|
||||
write_json(self, 404, {"error": "not found"})
|
||||
return
|
||||
try:
|
||||
length = int(self.headers.get("Content-Length", "0"))
|
||||
payload = json.loads(self.rfile.read(length).decode("utf-8")) if length else {}
|
||||
prompt = payload["prompt"]
|
||||
output = payload["output"]
|
||||
width = int(payload.get("width", 960))
|
||||
height = int(payload.get("height", 544))
|
||||
frames = int(payload.get("frames", 124))
|
||||
steps = int(payload.get("steps", 12))
|
||||
seed = int(payload.get("seed", 440407))
|
||||
attention = payload.get("attention")
|
||||
if attention is not None and attention not in AVAILABLE_BACKENDS:
|
||||
write_json(self, 400, {"error": "unsupported attention", "attention": attention, "available": list(AVAILABLE_BACKENDS)})
|
||||
return
|
||||
mux_audio = bool(payload.get("mux_audio", True))
|
||||
ffmpeg_loglevel = payload.get("ffmpeg_loglevel", "error")
|
||||
save_latent = payload.get("save_latent")
|
||||
cache_mode = payload.get("cache_mode")
|
||||
cache_threshold = float(payload.get("cache_threshold", 0.0))
|
||||
cache_start_percent = float(payload.get("cache_start_percent", 0.0))
|
||||
cache_end_percent = float(payload.get("cache_end_percent", 1.0))
|
||||
cache_subsample_factor = int(payload.get("cache_subsample_factor", 2))
|
||||
started = time.perf_counter()
|
||||
with runtime_lock:
|
||||
result = runtime.generate(
|
||||
prompt=prompt,
|
||||
output=output,
|
||||
width=width,
|
||||
height=height,
|
||||
frames=frames,
|
||||
steps=steps,
|
||||
seed=seed,
|
||||
attention=attention,
|
||||
mux_audio=mux_audio,
|
||||
ffmpeg_loglevel=ffmpeg_loglevel,
|
||||
save_latent=save_latent,
|
||||
cache_mode=cache_mode,
|
||||
cache_threshold=cache_threshold,
|
||||
cache_start_percent=cache_start_percent,
|
||||
cache_end_percent=cache_end_percent,
|
||||
cache_subsample_factor=cache_subsample_factor,
|
||||
)
|
||||
result["wall_seconds"] = time.perf_counter() - started
|
||||
write_json(self, 200, result)
|
||||
except Exception as exc:
|
||||
write_json(self, 500, {"error": type(exc).__name__, "message": str(exc)})
|
||||
|
||||
|
||||
server = ThreadingHTTPServer((args.host, args.port), Handler)
|
||||
print(json.dumps({"serving": True, "host": args.host, "port": args.port, **service_status(), "warmup_result": warmup_result}, indent=2), flush=True)
|
||||
server.serve_forever()
|
||||
Loading…
Add table
Reference in a new issue