Cache GB10 SageAttention3 wheel

Dockerfile.spark

@@ -4,12 +4,10 @@ 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
+# Prebuilt against the base image's CUDA 13 / Torch ABI for GB10 (sm_121a).
+COPY wheels/sageattn3-*.whl /tmp/wheels/
+RUN python -m pip install --no-cache-dir --no-deps /tmp/wheels/sageattn3-*.whl \
+    && rm -rf /tmp/wheels
 
 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)"

src/h3_blackwell_runtime/conditioning.py

@@ -11,7 +11,11 @@ class H3PromptTokenizer:
     def __init__(self, tokenizer_dir: str | Path):
         from transformers import AutoTokenizer
 
-        self.tokenizer = AutoTokenizer.from_pretrained(str(Path(tokenizer_dir)), local_files_only=True)
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            str(Path(tokenizer_dir)),
+            local_files_only=True,
+            fix_mistral_regex=True,
+        )
 
     def __call__(self, prompt: str, *, device: torch.device | str = "cuda") -> torch.Tensor:
         if not prompt: